diff --git a/.cargo/mutants.toml b/.cargo/mutants.toml new file mode 100644 index 000000000..c3a7e7a40 --- /dev/null +++ b/.cargo/mutants.toml @@ -0,0 +1,11 @@ +# Kernel mutations that cannot be distinguished by native correctness tests. +# V4 and Neon wrappers are exercised by SDE and AArch64 CI respectively. +# The remaining transformations preserve outputs and only switch between +# equivalent scalar-tail or zero-clamp expressions. +exclude_re = [ + "diskann-pipnn/src/(leaf|partition)_kernel\\.rs:.*Target<.*V4", + "diskann-pipnn/src/(leaf|partition)_kernel\\.rs:.*Target<.*Neon", + "diskann-pipnn/src/leaf_kernel\\.rs:.*replace < with <= in pair_distance", + "diskann-pipnn/src/partition_kernel\\.rs:.*replace \\* with / in process_(unary|binary)", + "diskann-pipnn/src/leaf_kernel\\.rs:.*replace > with >= in .*run_simd", +] diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a33b27a63..f435e3752 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -355,6 +355,7 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ + --package diskann-pipnn \ -- --skip compile_tests \ --skip pivots::tests::run_test_happy_path @@ -415,6 +416,7 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ + --package diskann-pipnn \ -- --skip compile_tests test-workspace: diff --git a/Cargo.lock b/Cargo.lock index 50ed17b11..0c9926781 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -680,6 +680,17 @@ dependencies = [ "thiserror 2.0.17", ] +[[package]] +name = "diskann-pipnn" +version = "0.55.0" +dependencies = [ + "criterion", + "diskann-linalg", + "diskann-vector", + "diskann-wide", + "thiserror 2.0.17", +] + [[package]] name = "diskann-providers" version = "0.55.0" diff --git a/Cargo.toml b/Cargo.toml index 6394a7d31..ee1ece101 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ members = [ "diskann-quantization", # Algorithm "diskann", + "diskann-pipnn", # Providers "diskann-providers", "diskann-disk", @@ -59,6 +60,7 @@ diskann-utils = { path = "diskann-utils", default-features = false, version = "0 diskann-quantization = { path = "diskann-quantization", default-features = false, version = "0.55.0" } # Algorithm diskann = { path = "diskann", version = "0.55.0" } +diskann-pipnn = { path = "diskann-pipnn", version = "0.55.0" } # Providers diskann-providers = { path = "diskann-providers", default-features = false, version = "0.55.0" } diskann-inmem = { path = "diskann-inmem", default-features = false, version = "0.55.0" } diff --git a/diskann-linalg/src/faer.rs b/diskann-linalg/src/faer.rs index 700396de4..1e7feeb9d 100644 --- a/diskann-linalg/src/faer.rs +++ b/diskann-linalg/src/faer.rs @@ -53,6 +53,29 @@ pub(super) fn sgemm_impl( faer::linalg::matmul::matmul(c, beta, a, b, alpha, Par::Seq) } +/// Implements the public lower-triangular AAT operation. +/// +/// The caller has already validated the matrix dimensions. +pub(super) fn sgemm_aat_lower_impl(m: usize, k: usize, a: &[f32], c: &mut [f32]) { + use faer::linalg::matmul::triangular::{matmul, BlockStructure}; + + let a = faer::mat::MatRef::from_row_major_slice(a, m, k); + let at = a.transpose(); + let c = faer::mat::MatMut::from_row_major_slice_mut(c, m, m); + + matmul( + c, + BlockStructure::TriangularLower, + faer::Accum::Replace, + a, + BlockStructure::Rectangular, + at, + BlockStructure::Rectangular, + 1.0, + Par::Seq, + ); +} + /// See the documentation for `svd_into`. /// /// The implementation may assume the the specified invariants hold for the sizes of the diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index aa24bb560..a6b8ec798 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -205,6 +205,49 @@ pub fn sgemm( Ok(()) } +/// Computes the lower triangle of $C = A A^\mathsf{T}$ for a dense row-major +/// $m \times k$ matrix $A$. +/// +/// The lower triangle, including the diagonal, is overwritten. The upper +/// triangle of the $m \times m$ destination is left unchanged. +/// +/// # Errors +/// +/// Returns an error if a matrix-size calculation overflows or either slice does +/// not match its declared dimensions. +pub fn sgemm_aat_lower(a: &[f32], m: usize, k: usize, c: &mut [f32]) -> Result<(), SgemmError> { + let expected_a_len = m.checked_mul(k).ok_or(SgemmError::DimensionOverflow { + matrix_name: MatrixName::A, + rows: m, + cols: k, + })?; + if a.len() != expected_a_len { + return Err(SgemmError::InvalidMatrixDimensions { + matrix_name: MatrixName::A, + expected_rows: m, + expected_cols: k, + actual_len: a.len(), + }); + } + + let expected_c_len = m.checked_mul(m).ok_or(SgemmError::DimensionOverflow { + matrix_name: MatrixName::C, + rows: m, + cols: m, + })?; + if c.len() != expected_c_len { + return Err(SgemmError::InvalidMatrixDimensions { + matrix_name: MatrixName::C, + expected_rows: m, + expected_cols: m, + actual_len: c.len(), + }); + } + + faer::sgemm_aat_lower_impl(m, k, a, c); + Ok(()) +} + /// Compute the SVD of the provided matrix implicit row-major matrix `data`. /// /// * `m`: The number of rows in `a`. diff --git a/diskann-linalg/tests/sgemm_aat_lower.rs b/diskann-linalg/tests/sgemm_aat_lower.rs new file mode 100644 index 000000000..e84d600d3 --- /dev/null +++ b/diskann-linalg/tests/sgemm_aat_lower.rs @@ -0,0 +1,109 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_linalg::{sgemm_aat_lower, MatrixName, SgemmError}; + +#[test] +fn computes_lower_triangle_and_preserves_upper_triangle() { + #[rustfmt::skip] + let a = [ + 1.0, 2.0, + 3.0, 4.0, + 5.0, 6.0, + ]; + let untouched = -123.0; + let mut c = [untouched; 9]; + + sgemm_aat_lower(&a, 3, 2, &mut c).unwrap(); + + #[rustfmt::skip] + assert_eq!(c, [ + 5.0, untouched, untouched, + 11.0, 25.0, untouched, + 17.0, 39.0, 61.0, + ]); +} + +#[test] +fn accepts_a_matrix_with_no_rows() { + sgemm_aat_lower(&[], 0, 3, &mut []).unwrap(); +} + +#[test] +fn zero_inner_dimension_zeros_only_the_lower_triangle() { + let untouched = -123.0; + let mut c = [untouched; 9]; + + sgemm_aat_lower(&[], 3, 0, &mut c).unwrap(); + + #[rustfmt::skip] + assert_eq!(c, [ + 0.0, untouched, untouched, + 0.0, 0.0, untouched, + 0.0, 0.0, 0.0, + ]); +} + +#[test] +fn rejects_invalid_input_dimensions() { + let mut c = [0.0; 4]; + + let error = sgemm_aat_lower(&[0.0; 3], 2, 2, &mut c).unwrap_err(); + + assert_eq!( + error, + SgemmError::InvalidMatrixDimensions { + matrix_name: MatrixName::A, + expected_rows: 2, + expected_cols: 2, + actual_len: 3, + } + ); +} + +#[test] +fn rejects_invalid_output_dimensions() { + let mut c = [0.0; 3]; + + let error = sgemm_aat_lower(&[0.0; 4], 2, 2, &mut c).unwrap_err(); + + assert_eq!( + error, + SgemmError::InvalidMatrixDimensions { + matrix_name: MatrixName::C, + expected_rows: 2, + expected_cols: 2, + actual_len: 3, + } + ); +} + +#[test] +fn rejects_input_size_overflow() { + let error = sgemm_aat_lower(&[], usize::MAX, 2, &mut []).unwrap_err(); + + assert_eq!( + error, + SgemmError::DimensionOverflow { + matrix_name: MatrixName::A, + rows: usize::MAX, + cols: 2, + } + ); +} + +#[test] +fn rejects_output_size_overflow() { + let error = sgemm_aat_lower(&[], usize::MAX, 0, &mut []).unwrap_err(); + + assert_eq!( + error, + SgemmError::DimensionOverflow { + matrix_name: MatrixName::C, + rows: usize::MAX, + cols: usize::MAX, + } + ); +} diff --git a/diskann-pipnn/Cargo.toml b/diskann-pipnn/Cargo.toml new file mode 100644 index 000000000..1ff7c3bfd --- /dev/null +++ b/diskann-pipnn/Cargo.toml @@ -0,0 +1,27 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT license. + +[package] +name = "diskann-pipnn" +version.workspace = true +description = "PiPNN graph construction for DiskANN" +authors.workspace = true +repository.workspace = true +license.workspace = true +edition.workspace = true + +[dependencies] +diskann-vector.workspace = true +diskann-wide.workspace = true +thiserror.workspace = true + +[dev-dependencies] +criterion.workspace = true +diskann-linalg.workspace = true + +[[bench]] +name = "kernels" +harness = false + +[lints] +workspace = true diff --git a/diskann-pipnn/benches/kernels.rs b/diskann-pipnn/benches/kernels.rs new file mode 100644 index 000000000..22790c554 --- /dev/null +++ b/diskann-pipnn/benches/kernels.rs @@ -0,0 +1,224 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{hint::black_box, time::Duration}; + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use diskann_linalg::{sgemm, sgemm_aat_lower, Transpose}; +use diskann_pipnn::{ + leaf_kernel::{nearest_leaf_neighbors, LeafNeighbor, LeafTopK, LeafTopKWorkspace}, + partition_kernel::{nearest_leaders, PartitionTopK}, +}; +use diskann_vector::distance::Metric; + +const BIGANN_DIMENSIONS: usize = 128; +const PARTITION_FANOUT: usize = 10; +const LEAF_KS: [usize; 2] = [2, 3]; +const LEAF_SIZES: [usize; 3] = [64, 256, 512]; +const METRICS: [Metric; 4] = [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, +]; + +fn fixed_data(rows: usize, columns: usize, sequence: usize) -> Vec { + (0..rows * columns) + .map(|index| { + let value = index + .wrapping_mul(1_664_525) + .wrapping_add(sequence.wrapping_mul(1_013_904_223)) + % 2_003; + (value as f32 - 1_001.0) / 1_001.0 + }) + .collect() +} + +fn normalize_rows(data: &mut [f32], columns: usize) { + for row in data.chunks_exact_mut(columns) { + let inverse_norm = row + .iter() + .map(|value| value * value) + .sum::() + .sqrt() + .recip(); + row.iter_mut().for_each(|value| *value *= inverse_norm); + } +} + +fn lower_dots(points: usize, metric: Metric) -> Vec { + let mut data = fixed_data(points, BIGANN_DIMENSIONS, points); + if metric == Metric::CosineNormalized { + normalize_rows(&mut data, BIGANN_DIMENSIONS); + } + let mut dots = vec![0.0; points * points]; + sgemm_aat_lower(&data, points, BIGANN_DIMENSIONS, &mut dots).unwrap(); + dots +} + +fn benchmark_partition_topk(c: &mut Criterion) { + let mut group = c.benchmark_group("pipnn/partition-topk"); + for (rows, leaders) in [(1_024, 64), (512, 256), (128, 1_000)] { + let points = fixed_data(rows, BIGANN_DIMENSIONS, rows); + let leader_data = fixed_data(leaders, BIGANN_DIMENSIONS, leaders); + let mut dots = vec![0.0; rows * leaders]; + sgemm( + Transpose::None, + Transpose::Ordinary, + rows, + leaders, + BIGANN_DIMENSIONS, + 1.0, + &points, + &leader_data, + None, + &mut dots, + ) + .unwrap(); + let leader_scales = leader_data + .chunks_exact(BIGANN_DIMENSIONS) + .map(|row| row.iter().map(|value| value * value).sum()) + .collect::>(); + let input = PartitionTopK { + dots: &dots, + rows, + leaders, + row_scales: &[], + leader_scales: &leader_scales, + metric: Metric::L2, + }; + let mut output = vec![0; rows * PARTITION_FANOUT]; + + group.throughput(Throughput::Elements(rows as u64)); + group.bench_with_input( + BenchmarkId::new( + "l2", + format!("{BIGANN_DIMENSIONS}d/{rows}x{leaders}/k{PARTITION_FANOUT}"), + ), + &input, + |bencher, input| { + bencher.iter(|| { + nearest_leaders(*input, PARTITION_FANOUT, &mut output).unwrap(); + black_box(&output); + }); + }, + ); + } + group.finish(); +} + +fn benchmark_lower_aat(c: &mut Criterion) { + let mut group = c.benchmark_group("pipnn/lower-aat"); + for points in LEAF_SIZES { + let data = fixed_data(points, BIGANN_DIMENSIONS, points); + let mut dots = vec![0.0; points * points]; + + group.throughput(Throughput::Elements((points * (points + 1) / 2) as u64)); + group.bench_function( + BenchmarkId::new("f32", format!("{points}x{BIGANN_DIMENSIONS}")), + |bencher| { + bencher.iter(|| { + sgemm_aat_lower(&data, points, BIGANN_DIMENSIONS, &mut dots).unwrap(); + black_box(&dots); + }); + }, + ); + } + group.finish(); +} + +fn benchmark_leaf_topk(c: &mut Criterion) { + let mut group = c.benchmark_group("pipnn/leaf-topk"); + for points in LEAF_SIZES { + for metric in METRICS { + for leaf_k in LEAF_KS { + let dots = lower_dots(points, metric); + let input = LeafTopK { + dots: &dots, + points, + metric, + }; + let mut output = vec![LeafNeighbor::default(); points * leaf_k]; + let mut workspace = LeafTopKWorkspace::new(); + nearest_leaf_neighbors(input, leaf_k, &mut output, &mut workspace).unwrap(); + + group.throughput(Throughput::Elements((points * (points - 1) / 2) as u64)); + group.bench_with_input( + BenchmarkId::new(metric.as_str(), format!("{points}/k{leaf_k}")), + &input, + |bencher, input| { + bencher.iter(|| { + nearest_leaf_neighbors(*input, leaf_k, &mut output, &mut workspace) + .unwrap(); + black_box(&output); + }); + }, + ); + } + } + } + group.finish(); +} + +fn benchmark_full_leaf(c: &mut Criterion) { + let mut group = c.benchmark_group("pipnn/full-leaf-numerical"); + for points in LEAF_SIZES { + for leaf_k in LEAF_KS { + let data = fixed_data(points, BIGANN_DIMENSIONS, points); + let mut dots = vec![0.0; points * points]; + let mut output = vec![LeafNeighbor::default(); points * leaf_k]; + let mut workspace = LeafTopKWorkspace::new(); + sgemm_aat_lower(&data, points, BIGANN_DIMENSIONS, &mut dots).unwrap(); + nearest_leaf_neighbors( + LeafTopK { + dots: &dots, + points, + metric: Metric::L2, + }, + leaf_k, + &mut output, + &mut workspace, + ) + .unwrap(); + + group.throughput(Throughput::Elements(points as u64)); + group.bench_function( + BenchmarkId::new("l2", format!("{points}x{BIGANN_DIMENSIONS}/k{leaf_k}")), + |bencher| { + bencher.iter(|| { + sgemm_aat_lower(&data, points, BIGANN_DIMENSIONS, &mut dots).unwrap(); + nearest_leaf_neighbors( + LeafTopK { + dots: &dots, + points, + metric: Metric::L2, + }, + leaf_k, + &mut output, + &mut workspace, + ) + .unwrap(); + black_box(&output); + }); + }, + ); + } + } + group.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(30) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(3)); + targets = + benchmark_partition_topk, + benchmark_lower_aat, + benchmark_leaf_topk, + benchmark_full_leaf +} +criterion_main!(benches); diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs new file mode 100644 index 000000000..f8531d429 --- /dev/null +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -0,0 +1,735 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Fused nearest-neighbor kernel for a leaf's lower dot-product matrix. + +use diskann_vector::distance::Metric; +use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; + +/// Widest f32 SIMD lane count DiskANN dispatches to, used to size lane scratch. +const MAX_LANES: usize = 16; + +const L2: u8 = 0; +const COSINE_NORMALIZED: u8 = 1; +const INNER_PRODUCT: u8 = 2; +const COSINE: u8 = 3; + +/// One leaf-local neighbor and its metric distance. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct LeafNeighbor { + /// Position in the leaf, not a dataset ID. + pub position: u32, + /// Distance from the row point to `position`. + pub distance: f32, +} + +impl LeafNeighbor { + /// Construct a leaf-local neighbor. + pub const fn new(position: u32, distance: f32) -> Self { + Self { position, distance } + } +} + +impl Default for LeafNeighbor { + fn default() -> Self { + Self::new(u32::MAX, f32::INFINITY) + } +} + +/// Lower-triangular dot products consumed by [`nearest_leaf_neighbors`]. +#[derive(Clone, Copy, Debug)] +pub struct LeafTopK<'a> { + /// Row-major `points * points` matrix. Only entries with `column <= row` are read. + pub dots: &'a [f32], + /// Number of points represented by the matrix. + pub points: usize, + /// Metric used to rank pairs. + pub metric: Metric, +} + +/// Reusable temporary storage for leaf top-k selection. +#[derive(Debug, Default)] +pub struct LeafTopKWorkspace { + norms: Vec, + worst: Vec, +} + +impl LeafTopKWorkspace { + /// Construct an empty workspace. + pub const fn new() -> Self { + Self { + norms: Vec::new(), + worst: Vec::new(), + } + } +} + +/// Validation or allocation error returned by [`nearest_leaf_neighbors`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum LeafKernelError { + /// The point count cannot be represented in leaf-local `u32` positions. + #[error("point count {0} exceeds the u32 position limit")] + TooManyPoints(usize), + /// A declared shape overflowed `usize`. + #[error("{buffer} shape {rows} x {cols} overflows usize")] + ShapeOverflow { + /// Name of the buffer whose shape overflowed. + buffer: &'static str, + /// Declared row count. + rows: usize, + /// Declared column count. + cols: usize, + }, + /// A supplied slice did not match its declared shape. + #[error("invalid {buffer} length: expected {expected}, got {actual}")] + InvalidBufferLength { + /// Name of the invalid buffer. + buffer: &'static str, + /// Required length. + expected: usize, + /// Supplied length. + actual: usize, + }, + /// Temporary storage could not be reserved. + #[error("failed to reserve {additional} values for {buffer}")] + Allocation { + /// Name of the temporary buffer. + buffer: &'static str, + /// Additional element capacity requested. + additional: usize, + }, + /// A row did not contain enough rankable pair distances to fill its output. + #[error("row {row} has fewer than {neighbors} rankable leaf neighbors")] + InsufficientRankableNeighbors { + /// Zero-based row position in the leaf. + row: usize, + /// Required number of non-self neighbors. + neighbors: usize, + }, +} + +/// Select the nearest non-self leaf positions for every row. +/// +/// The strictly lower triangle is scanned once. Each pair updates both row +/// trackers, so the upper triangle is neither read nor materialized. The +/// returned value is `min(k, points - 1)`, and `output` contains exactly +/// `points * returned_k` entries grouped by row and ordered by ascending +/// distance. Equal distances retain pair scan order. +pub fn nearest_leaf_neighbors( + input: LeafTopK<'_>, + k: usize, + output: &mut [LeafNeighbor], + workspace: &mut LeafTopKWorkspace, +) -> Result { + let actual_k = validate(input, k, output)?; + if actual_k == 0 { + return Ok(0); + } + + resize("norms", &mut workspace.norms, input.points, 0.0)?; + resize( + "worst distances", + &mut workspace.worst, + input.points, + f32::INFINITY, + )?; + for (row, norm) in workspace.norms.iter_mut().enumerate() { + let squared_norm = input.dots[row * input.points + row]; + *norm = if input.metric == Metric::Cosine { + // Match diskann-vector: a finite/subnormal squared norm below this + // threshold is a zero vector, while NaN continues through the + // distance calculation as non-rankable. + if squared_norm < f32::MIN_POSITIVE { + 0.0 + } else { + squared_norm.sqrt() + } + } else { + squared_norm + }; + } + output.fill(LeafNeighbor::default()); + workspace.worst.fill(f32::INFINITY); + + diskann_wide::arch::dispatch(LeafKernel { + input, + k: actual_k, + output, + norms: &workspace.norms, + worst: &mut workspace.worst, + }); + if let Some(row) = output + .chunks_exact(actual_k) + .position(|neighbors| neighbors[actual_k - 1].position == u32::MAX) + { + return Err(LeafKernelError::InsufficientRankableNeighbors { + row, + neighbors: actual_k, + }); + } + Ok(actual_k) +} + +fn validate( + input: LeafTopK<'_>, + k: usize, + output: &[LeafNeighbor], +) -> Result { + if input.points > u32::MAX as usize { + return Err(LeafKernelError::TooManyPoints(input.points)); + } + let matrix_len = checked_area("lower dot-product matrix", input.points, input.points)?; + check_length("lower dot-product matrix", input.dots.len(), matrix_len)?; + let actual_k = k.min(input.points.saturating_sub(1)); + let output_len = checked_area("output", input.points, actual_k)?; + check_length("output", output.len(), output_len)?; + Ok(actual_k) +} + +fn resize( + buffer: &'static str, + values: &mut Vec, + len: usize, + value: T, +) -> Result<(), LeafKernelError> { + let additional = len.saturating_sub(values.len()); + values + .try_reserve(additional) + .map_err(|_| LeafKernelError::Allocation { buffer, additional })?; + values.resize(len, value); + Ok(()) +} + +fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> Result { + rows.checked_mul(cols) + .ok_or(LeafKernelError::ShapeOverflow { buffer, rows, cols }) +} + +fn check_length( + buffer: &'static str, + actual: usize, + expected: usize, +) -> Result<(), LeafKernelError> { + if actual == expected { + Ok(()) + } else { + Err(LeafKernelError::InvalidBufferLength { + buffer, + expected, + actual, + }) + } +} + +struct LeafKernel<'a, 'o, 'w> { + input: LeafTopK<'a>, + k: usize, + output: &'o mut [LeafNeighbor], + norms: &'w [f32], + worst: &'w mut [f32], +} + +impl LeafKernel<'_, '_, '_> { + fn run_simd(self, arch: F::Arch) + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, + { + if self.k > 3 { + process_pairs_simd_dynamic::( + arch, + self.input, + self.k, + self.output, + self.norms, + self.worst, + ); + return; + } + match self.k { + 1 => self.run_fused::(arch), + 2 => self.run_fused::(arch), + 3 => self.run_fused::(arch), + _ => unreachable!("validated non-zero leaf width"), + } + } + + fn run_fused(self, arch: F::Arch) + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, + { + match self.input.metric { + Metric::L2 => process_pairs_simd_fused::( + arch, + self.input, + self.output, + self.norms, + self.worst, + ), + Metric::CosineNormalized => process_pairs_simd_fused::( + arch, + self.input, + self.output, + self.norms, + self.worst, + ), + Metric::InnerProduct => process_pairs_simd_fused::( + arch, + self.input, + self.output, + self.norms, + self.worst, + ), + Metric::Cosine => process_pairs_simd_fused::( + arch, + self.input, + self.output, + self.norms, + self.worst, + ), + } + } +} + +impl diskann_wide::arch::Target for LeafKernel<'_, '_, '_> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + #[inline(always)] + fn run(self, arch: A) { + self.run_simd::(arch); + } +} + +#[cfg(test)] +fn process_pairs_scalar( + input: LeafTopK<'_>, + k: usize, + output: &mut [LeafNeighbor], + norms: &[f32], + worst: &mut [f32], +) { + for row in 1..input.points { + for column in 0..row { + let dot = input.dots[row * input.points + column]; + let distance = pair_distance(input.metric, dot, norms[row], norms[column]); + insert_row(output, worst, k, row, column as u32, distance); + insert_row(output, worst, k, column, row as u32, distance); + } + } +} + +/// Fused dual-endpoint scan for row widths without a specialized arm. +/// +/// Identical structure to [`process_pairs_simd_fused`], with the slot count +/// read at run time. Wider leaves are rare, so the extra indirection is +/// cheaper than instantiating an arm per width. +fn process_pairs_simd_dynamic( + arch: F::Arch, + input: LeafTopK<'_>, + k: usize, + output: &mut [LeafNeighbor], + norms: &[f32], + worst: &mut [f32], +) where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + let output_ptr = output.as_mut_ptr(); + let worst_ptr = worst.as_mut_ptr(); + for row in 1..input.points { + let row_start = row * input.points; + let row_norm = F::splat(arch, norms[row]); + // SAFETY: `row < input.points == worst.len()`. + let mut row_worst = unsafe { *worst_ptr.add(row) }; + let mut column = 0; + while column + F::LANES <= row { + // SAFETY: the full chunk is contained in the strict lower row prefix. + let dots = unsafe { F::load_simd(arch, input.dots.as_ptr().add(row_start + column)) }; + // SAFETY: `column + F::LANES <= row < input.points == norms.len()`. + let column_norms = unsafe { F::load_simd(arch, norms.as_ptr().add(column)) }; + let distances = pair_distances::(arch, input.metric, dots, row_norm, column_norms); + let row_eligible = distances.lt_simd(F::splat(arch, row_worst)); + // SAFETY: the full chunk lies below `row`, so it is within `worst`. + let column_worst = unsafe { F::load_simd(arch, worst_ptr.add(column)) }; + let column_eligible = distances.lt_simd(column_worst); + let row_bits = u64::from(row_eligible.bitmask().to_underlying()); + let column_bits = u64::from(column_eligible.bitmask().to_underlying()); + if row_bits | column_bits != 0 { + let mut values = [0.0f32; MAX_LANES]; + // SAFETY: the array covers every f32 SIMD width DiskANN exposes. + unsafe { distances.store_simd(values.as_mut_ptr()) }; + let mut row_bits = row_bits; + while row_bits != 0 { + let lane = row_bits.trailing_zeros() as usize; + row_bits &= row_bits - 1; + let distance = values[lane]; + if distance < row_worst { + // SAFETY: `row * k + k` is inside the validated output. + row_worst = unsafe { + insert_slots(output_ptr, row * k, k, (column + lane) as u32, distance) + }; + } + } + let mut column_bits = column_bits; + while column_bits != 0 { + let lane = column_bits.trailing_zeros() as usize; + column_bits &= column_bits - 1; + let target = column + lane; + // SAFETY: `target < row`, so its slots are inside the output. + let new_worst = unsafe { + insert_slots(output_ptr, target * k, k, row as u32, values[lane]) + }; + // SAFETY: `target < row < worst.len()`. + unsafe { *worst_ptr.add(target) = new_worst }; + } + } + column += F::LANES; + } + while column < row { + // SAFETY: the scalar tail remains in the strict lower triangle. + let dot = unsafe { *input.dots.get_unchecked(row_start + column) }; + // SAFETY: `column < row < input.points == norms.len()`. + let column_norm = unsafe { *norms.get_unchecked(column) }; + let distance = pair_distance(input.metric, dot, norms[row], column_norm); + if distance < row_worst { + // SAFETY: `row * k + k` is inside the validated output. + row_worst = + unsafe { insert_slots(output_ptr, row * k, k, column as u32, distance) }; + } + // SAFETY: `column < row < worst.len()`. + let column_worst = unsafe { *worst_ptr.add(column) }; + if distance < column_worst { + // SAFETY: `column < row`, so its slots are inside the output. + let new_worst = + unsafe { insert_slots(output_ptr, column * k, k, row as u32, distance) }; + // SAFETY: `column < row < worst.len()`. + unsafe { *worst_ptr.add(column) = new_worst }; + } + column += 1; + } + // SAFETY: `row < worst.len()`. + unsafe { *worst_ptr.add(row) = row_worst }; + } +} + +/// Fused dual-endpoint scan of the strict lower triangle. +/// +/// The row's current worst distance stays in a register for the whole row, and +/// each chunk derives both endpoint candidate masks before touching memory, so +/// a chunk where neither endpoint can accept costs one branch. `SLOTS` is the +/// per-row neighbor count, threaded as a const so the insert arm is selected at +/// compile time. +#[inline(never)] +fn process_pairs_simd_fused( + arch: F::Arch, + input: LeafTopK<'_>, + output: &mut [LeafNeighbor], + norms: &[f32], + worst: &mut [f32], +) where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + let output_ptr = output.as_mut_ptr(); + let worst_ptr = worst.as_mut_ptr(); + for row in 1..input.points { + let row_start = row * input.points; + let row_norm = F::splat(arch, norms[row]); + // SAFETY: `row < input.points == worst.len()`. + let mut row_worst = unsafe { *worst_ptr.add(row) }; + let mut column = 0; + while column + F::LANES <= row { + // SAFETY: the full chunks are inside the validated matrix and norms. + let dots = unsafe { F::load_simd(arch, input.dots.as_ptr().add(row_start + column)) }; + // SAFETY: `column + F::LANES <= row < input.points == norms.len()`. + let column_norms = unsafe { F::load_simd(arch, norms.as_ptr().add(column)) }; + let distances = + pair_distances::(arch, metric::(), dots, row_norm, column_norms); + let row_eligible = distances.lt_simd(F::splat(arch, row_worst)); + // SAFETY: the full chunk lies below `row`, so it is within `worst`. + let column_worst = unsafe { F::load_simd(arch, worst_ptr.add(column)) }; + let column_eligible = distances.lt_simd(column_worst); + // Test both candidate masks with a single reduction. Reducing each + // mask separately costs an extra cross-lane extraction per chunk, + // and the overwhelmingly common case is that neither end accepts. + let row_bits = u64::from(row_eligible.bitmask().to_underlying()); + let column_bits = u64::from(column_eligible.bitmask().to_underlying()); + if row_bits | column_bits != 0 { + let mut values = [0.0f32; MAX_LANES]; + // SAFETY: the array covers every f32 SIMD width DiskANN exposes. + unsafe { distances.store_simd(values.as_mut_ptr()) }; + let mut row_bits = row_bits; + while row_bits != 0 { + let lane = row_bits.trailing_zeros() as usize; + row_bits &= row_bits - 1; + let distance = values[lane]; + // Earlier lanes in this chunk may already have tightened the + // threshold, so re-check against the live value. + if distance < row_worst { + // SAFETY: `row * SLOTS + SLOTS` is inside the validated output. + row_worst = unsafe { + insert_slots( + output_ptr, + row * SLOTS, + SLOTS, + (column + lane) as u32, + distance, + ) + }; + } + } + let mut column_bits = column_bits; + while column_bits != 0 { + let lane = column_bits.trailing_zeros() as usize; + column_bits &= column_bits - 1; + let target = column + lane; + // SAFETY: `target < row`, so its slots are inside the output. + let new_worst = unsafe { + insert_slots(output_ptr, target * SLOTS, SLOTS, row as u32, values[lane]) + }; + // SAFETY: `target < row < worst.len()`. + unsafe { *worst_ptr.add(target) = new_worst }; + } + } + column += F::LANES; + } + while column < row { + // SAFETY: the scalar tail remains in the strict lower triangle. + let dot = unsafe { *input.dots.get_unchecked(row_start + column) }; + // SAFETY: `column < row < input.points == norms.len()`. + let column_norm = unsafe { *norms.get_unchecked(column) }; + let distance = pair_distance(metric::(), dot, norms[row], column_norm); + if distance < row_worst { + // SAFETY: `row * SLOTS + SLOTS` is inside the validated output. + row_worst = unsafe { + insert_slots(output_ptr, row * SLOTS, SLOTS, column as u32, distance) + }; + } + // SAFETY: `column < row < worst.len()`. + let column_worst = unsafe { *worst_ptr.add(column) }; + if distance < column_worst { + // SAFETY: `column < row`, so its slots are inside the output. + let new_worst = unsafe { + insert_slots(output_ptr, column * SLOTS, SLOTS, row as u32, distance) + }; + // SAFETY: `column < row < worst.len()`. + unsafe { *worst_ptr.add(column) = new_worst }; + } + column += 1; + } + // SAFETY: `row < worst.len()`. + unsafe { *worst_ptr.add(row) = row_worst }; + } +} + +const fn metric() -> Metric { + match METRIC { + L2 => Metric::L2, + COSINE_NORMALIZED => Metric::CosineNormalized, + INNER_PRODUCT => Metric::InnerProduct, + COSINE => Metric::Cosine, + _ => unreachable!(), + } +} + +/// Insert one candidate into a row's ascending-distance slots and return the +/// row's new worst distance. +/// +/// Slot counts of one, two, and three are the production leaf widths and get +/// straight-line arms. Wider rows fall back to a bubble-up over the same +/// layout, which produces identical results at a lower instruction count than +/// specializing further would justify. +/// +/// # Safety +/// +/// `base + slots` must be within the allocation behind `output`. +#[inline(always)] +unsafe fn insert_slots( + output: *mut LeafNeighbor, + base: usize, + slots: usize, + position: u32, + distance: f32, +) -> f32 { + let entry = LeafNeighbor::new(position, distance); + match slots { + 1 => { + // SAFETY: the caller guarantees `base` is in bounds. + unsafe { *output.add(base) = entry }; + distance + } + 2 => { + // SAFETY: the caller guarantees `base` and `base + 1` are in bounds. + let first = unsafe { *output.add(base) }; + if distance < first.distance { + // SAFETY: as above. + unsafe { + *output.add(base) = entry; + *output.add(base + 1) = first; + } + first.distance + } else { + // SAFETY: as above. + unsafe { *output.add(base + 1) = entry }; + distance + } + } + 3 => { + // SAFETY: the caller guarantees `base..base + 3` is in bounds. + let (first, second) = unsafe { (*output.add(base), *output.add(base + 1)) }; + if distance < first.distance { + // SAFETY: as above. + unsafe { + *output.add(base) = entry; + *output.add(base + 1) = first; + *output.add(base + 2) = second; + } + } else if distance < second.distance { + // SAFETY: as above. + unsafe { + *output.add(base + 1) = entry; + *output.add(base + 2) = second; + } + } else { + // SAFETY: as above. + unsafe { *output.add(base + 2) = entry }; + return distance; + } + second.distance + } + _ => { + let last = base + slots - 1; + // SAFETY: the caller guarantees `base..base + slots` is in bounds. + unsafe { *output.add(last) = entry }; + let mut position = last; + while position > base { + // SAFETY: `base < position <= last` stays inside the row. + let (current, previous) = + unsafe { (*output.add(position), *output.add(position - 1)) }; + if current.distance >= previous.distance { + break; + } + // SAFETY: as above. + unsafe { + *output.add(position) = previous; + *output.add(position - 1) = current; + } + position -= 1; + } + // SAFETY: `last` is in bounds. + unsafe { (*output.add(last)).distance } + } + } +} + +#[inline(always)] +fn pair_distances(arch: F::Arch, metric: Metric, dot: F, row_norm: F, column_norm: F) -> F +where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, +{ + let zero = F::default(arch); + let clamp_nonnegative = |distance: F| { + // SIMD max has ISA-specific NaN behavior. Select the original NaN + // explicitly so it remains non-rankable on every backend. + distance + .eq_simd(distance) + .select(zero.max_simd(distance), distance) + }; + match metric { + Metric::L2 => { + let distance = row_norm + column_norm - F::splat(arch, 2.0) * dot; + clamp_nonnegative(distance) + } + Metric::CosineNormalized => { + let distance = F::splat(arch, 1.0) - dot; + clamp_nonnegative(distance) + } + Metric::InnerProduct => zero - dot, + Metric::Cosine => { + let one = F::splat(arch, 1.0); + let row_zero = row_norm.eq_simd(zero); + let column_zero = column_norm.eq_simd(zero); + let denominator = row_norm * column_norm; + let safe_denominator = row_zero.select(one, column_zero.select(one, denominator)); + let cosine = row_zero.select(zero, column_zero.select(zero, dot / safe_denominator)); + clamp_nonnegative(one - cosine) + } + } +} + +#[inline(always)] +fn pair_distance(metric: Metric, dot: f32, row_norm: f32, column_norm: f32) -> f32 { + match metric { + Metric::L2 => { + let distance = row_norm + column_norm - 2.0 * dot; + if distance < 0.0 { + 0.0 + } else { + distance + } + } + Metric::CosineNormalized => { + let distance = 1.0 - dot; + if distance < 0.0 { + 0.0 + } else { + distance + } + } + Metric::InnerProduct => -dot, + Metric::Cosine => { + let denominator = row_norm * column_norm; + let cosine = if row_norm != 0.0 && column_norm != 0.0 { + dot / denominator + } else { + 0.0 + }; + let distance = 1.0 - cosine; + if distance < 0.0 { + 0.0 + } else { + distance + } + } + } +} + +#[inline(always)] +#[cfg(test)] +fn insert_row( + output: &mut [LeafNeighbor], + worst: &mut [f32], + k: usize, + row: usize, + position: u32, + distance: f32, +) { + if distance.partial_cmp(&worst[row]) != Some(std::cmp::Ordering::Less) { + return; + } + + let start = row * k; + let row_output = &mut output[start..start + k]; + row_output[k - 1] = LeafNeighbor::new(position, distance); + let mut index = k - 1; + while index > 0 && row_output[index].distance < row_output[index - 1].distance { + row_output.swap(index, index - 1); + index -= 1; + } + worst[row] = row_output[k - 1].distance; +} + +#[cfg(test)] +mod tests; diff --git a/diskann-pipnn/src/leaf_kernel/tests.rs b/diskann-pipnn/src/leaf_kernel/tests.rs new file mode 100644 index 000000000..90381bc1d --- /dev/null +++ b/diskann-pipnn/src/leaf_kernel/tests.rs @@ -0,0 +1,112 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; + +fn dots(metric: Metric, points: usize) -> Vec { + let mut dots = vec![f32::NAN; points * points]; + for row in 0..points { + dots[row * points + row] = if metric == Metric::Cosine && row == 0 { + 0.0 + } else { + 1.0 + (row % 5) as f32 + }; + for column in 0..row { + dots[row * points + column] = (((row * 17 + column * 11) % 23) as f32 - 11.0) * 0.03125; + } + } + dots +} + +fn norms(input: LeafTopK<'_>) -> Vec { + (0..input.points) + .map(|row| { + let squared = input.dots[row * input.points + row]; + if input.metric == Metric::Cosine { + if squared < f32::MIN_POSITIVE { + 0.0 + } else { + squared.sqrt() + } + } else { + squared + } + }) + .collect() +} + +#[test] +fn scalar_reference_matches_runtime_dispatch() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + for points in [7, 17] { + let dots = dots(metric, points); + let input = LeafTopK { + dots: &dots, + points, + metric, + }; + for k in [1, 2, 3, 4] { + let mut expected = vec![LeafNeighbor::default(); points * k]; + nearest_leaf_neighbors(input, k, &mut expected, &mut LeafTopKWorkspace::new()) + .unwrap(); + + let mut actual = vec![LeafNeighbor::default(); points * k]; + let mut worst = vec![f32::INFINITY; points]; + let norms = norms(input); + process_pairs_scalar(input, k, &mut actual, &norms, &mut worst); + + assert_eq!(actual, expected, "{metric:?}, n={points}, k={k}"); + } + } + } +} + +#[test] +fn scalar_insertion_orders_candidates_and_rejects_nan() { + let mut output = [LeafNeighbor::default(); 4]; + let mut worst = [f32::INFINITY]; + + for (position, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 0.5)] { + insert_row(&mut output, &mut worst, 4, 0, position, distance); + } + insert_row(&mut output, &mut worst, 4, 0, 5, f32::NAN); + + assert_eq!( + output, + [ + LeafNeighbor::new(4, 0.5), + LeafNeighbor::new(1, 1.0), + LeafNeighbor::new(3, 2.0), + LeafNeighbor::new(2, 3.0), + ] + ); + assert_eq!(worst, [3.0]); +} + +#[test] +fn workspace_can_shrink_and_grow_between_calls() { + let mut workspace = LeafTopKWorkspace::new(); + for points in [17, 7, 17] { + let dots = dots(Metric::L2, points); + let mut output = vec![LeafNeighbor::default(); points * 2]; + nearest_leaf_neighbors( + LeafTopK { + dots: &dots, + points, + metric: Metric::L2, + }, + 2, + &mut output, + &mut workspace, + ) + .unwrap(); + assert!(output.iter().all(|neighbor| neighbor.position != u32::MAX)); + } +} diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs new file mode 100644 index 000000000..198434b72 --- /dev/null +++ b/diskann-pipnn/src/lib.rs @@ -0,0 +1,9 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! PiPNN graph construction. + +pub mod leaf_kernel; +pub mod partition_kernel; diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs new file mode 100644 index 000000000..b2942fa65 --- /dev/null +++ b/diskann-pipnn/src/partition_kernel.rs @@ -0,0 +1,411 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Distance and top-k kernel for partition assignment. +//! +//! The kernel consumes a row-major tile of point-to-leader dot products. It +//! converts those products to metric distances while retaining only leader +//! positions; partition recursion and cluster ownership stay with the caller. + +use diskann_vector::distance::Metric; +use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector}; + +/// Maximum number of leaders retained for one point. +pub const MAX_PARTITION_FANOUT: usize = 16; + +type TopK = [(u32, f32); MAX_PARTITION_FANOUT]; + +/// Input tile and metric-specific normalization terms for partition top-k. +#[derive(Clone, Copy, Debug)] +pub struct PartitionTopK<'a> { + /// Row-major `rows * leaders` point-to-leader dot products. + pub dots: &'a [f32], + /// Number of points represented by `dots`. + pub rows: usize, + /// Number of leaders represented by each row. + pub leaders: usize, + /// Squared point norms for cosine, otherwise empty. + pub row_scales: &'a [f32], + /// Leader norms for cosine, squared leader norms for L2, otherwise empty. + pub leader_scales: &'a [f32], + /// Distance metric used to rank leaders. + pub metric: Metric, +} + +/// Validation error returned by [`nearest_leaders`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] +pub enum PartitionKernelError { + /// A declared matrix or output shape overflowed `usize`. + #[error("{buffer} shape {rows} x {cols} overflows usize")] + ShapeOverflow { + /// Name of the buffer whose shape overflowed. + buffer: &'static str, + /// Declared row count. + rows: usize, + /// Declared column count. + cols: usize, + }, + /// A supplied slice did not match its declared shape. + #[error("invalid {buffer} length: expected {expected}, got {actual}")] + InvalidBufferLength { + /// Name of the invalid buffer. + buffer: &'static str, + /// Required length. + expected: usize, + /// Supplied length. + actual: usize, + }, + /// The requested fanout cannot be represented by the fixed top-k tracker. + #[error("invalid fanout {fanout} for {leaders} leaders; maximum is {maximum}")] + InvalidFanout { + /// Requested number of leaders per row. + fanout: usize, + /// Available leader count. + leaders: usize, + /// Kernel maximum. + maximum: usize, + }, + /// Leader positions cannot be represented as `u32`. + #[error("leader count {0} exceeds the u32 position limit")] + TooManyLeaders(usize), + /// A row did not contain enough rankable distances to fill its output. + #[error("row {row} has fewer than {fanout} rankable leader distances")] + InsufficientRankableDistances { + /// Zero-based row position in the input tile. + row: usize, + /// Requested number of leader positions. + fanout: usize, + }, +} + +/// Select the nearest `fanout` leader positions for every input row. +/// +/// Results for each row are ordered by ascending distance. Equal distances do +/// not replace or move an already retained entry, so leader scan order breaks +/// ties. A zero fanout is a validated no-op. +/// +/// For L2, the point's squared norm is omitted because it is constant across +/// every leader in a row and cannot change the ranking. +pub fn nearest_leaders( + input: PartitionTopK<'_>, + fanout: usize, + output: &mut [u32], +) -> Result<(), PartitionKernelError> { + validate(input, fanout, output)?; + if fanout == 0 || input.rows == 0 { + return Ok(()); + } + + diskann_wide::arch::dispatch(PartitionKernel { + input, + fanout, + output, + }); + if let Some(row) = output + .chunks_exact(fanout) + .position(|leaders| leaders.contains(&u32::MAX)) + { + return Err(PartitionKernelError::InsufficientRankableDistances { row, fanout }); + } + Ok(()) +} + +fn validate( + input: PartitionTopK<'_>, + fanout: usize, + output: &[u32], +) -> Result<(), PartitionKernelError> { + if input.leaders > u32::MAX as usize { + return Err(PartitionKernelError::TooManyLeaders(input.leaders)); + } + if fanout > MAX_PARTITION_FANOUT || fanout > input.leaders { + return Err(PartitionKernelError::InvalidFanout { + fanout, + leaders: input.leaders, + maximum: MAX_PARTITION_FANOUT, + }); + } + + let expected_dots = checked_area("dot-product tile", input.rows, input.leaders)?; + check_length("dot-product tile", input.dots.len(), expected_dots)?; + let expected_output = checked_area("output", input.rows, fanout)?; + check_length("output", output.len(), expected_output)?; + + let (row_scales, leader_scales) = match input.metric { + Metric::Cosine => (input.rows, input.leaders), + Metric::L2 => (0, input.leaders), + Metric::CosineNormalized | Metric::InnerProduct => (0, 0), + }; + check_length("row scales", input.row_scales.len(), row_scales)?; + check_length("leader scales", input.leader_scales.len(), leader_scales) +} + +fn checked_area( + buffer: &'static str, + rows: usize, + cols: usize, +) -> Result { + rows.checked_mul(cols) + .ok_or(PartitionKernelError::ShapeOverflow { buffer, rows, cols }) +} + +fn check_length( + buffer: &'static str, + actual: usize, + expected: usize, +) -> Result<(), PartitionKernelError> { + if actual == expected { + Ok(()) + } else { + Err(PartitionKernelError::InvalidBufferLength { + buffer, + expected, + actual, + }) + } +} + +struct PartitionKernel<'a, 'o> { + input: PartitionTopK<'a>, + fanout: usize, + output: &'o mut [u32], +} + +impl PartitionKernel<'_, '_> { + fn run_simd(self, arch: F::Arch) + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, + { + process_rows_simd::(arch, self.input, self.fanout, self.output); + } +} + +impl diskann_wide::arch::Target for PartitionKernel<'_, '_> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + #[inline(always)] + fn run(self, arch: A) { + self.run_simd::(arch); + } +} + +#[cfg(test)] +fn process_rows_scalar(input: PartitionTopK<'_>, fanout: usize, output: &mut [u32]) { + for (row_index, (dot_row, output_row)) in input + .dots + .chunks_exact(input.leaders) + .zip(output.chunks_exact_mut(fanout)) + .enumerate() + { + let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + let row_scale = input.row_scales.get(row_index).copied().unwrap_or(0.0); + for (leader, &dot) in dot_row.iter().enumerate() { + let leader_scale = input.leader_scales.get(leader).copied().unwrap_or(0.0); + insert_topk( + &mut top, + fanout, + leader as u32, + distance(input.metric, dot, row_scale, leader_scale), + ); + } + copy_ids(&top, output_row); + } +} + +fn process_rows_simd(arch: F::Arch, input: PartitionTopK<'_>, fanout: usize, output: &mut [u32]) +where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + for (row_index, (dot_row, output_row)) in input + .dots + .chunks_exact(input.leaders) + .zip(output.chunks_exact_mut(fanout)) + .enumerate() + { + let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + match input.metric { + Metric::L2 => process_binary::( + arch, + dot_row, + input.leader_scales, + &mut top, + fanout, + |dot, norm| F::splat(arch, -2.0).mul_add_simd(dot, norm), + ), + Metric::CosineNormalized => { + process_unary::(arch, dot_row, &mut top, fanout, |dot| { + F::splat(arch, 1.0) - dot + }) + } + Metric::InnerProduct => process_unary::(arch, dot_row, &mut top, fanout, |dot| { + F::default(arch) - dot + }), + Metric::Cosine => process_cosine::( + arch, + dot_row, + input.row_scales[row_index], + input.leader_scales, + &mut top, + fanout, + ), + } + copy_ids(&top, output_row); + } +} + +fn process_cosine( + arch: F::Arch, + dots: &[f32], + row_norm_squared: f32, + leader_norms: &[f32], + top: &mut TopK, + fanout: usize, +) where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + let row_norm = F::splat(arch, row_norm_squared.sqrt()); + let one = F::splat(arch, 1.0); + let zero = F::default(arch); + process_binary::(arch, dots, leader_norms, top, fanout, |dot, leader_norm| { + let denominator = row_norm * leader_norm; + let valid = denominator.gt_simd(zero); + let safe_denominator = valid.select(denominator, one); + let cosine = valid.select(dot / safe_denominator, zero); + one - cosine + }); +} + +fn process_unary( + arch: F::Arch, + dots: &[f32], + top: &mut TopK, + fanout: usize, + transform: Transform, +) where + F: SIMDVector + SIMDFloat, + Transform: Fn(F) -> F, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + let full = dots.len() / F::LANES * F::LANES; + for base in (0..full).step_by(F::LANES) { + // SAFETY: `base + F::LANES <= full <= dots.len()`. + let dots = unsafe { F::load_simd(arch, dots.as_ptr().add(base)) }; + insert_lanes(transform(dots), base, top, fanout); + } + for (offset, &dot) in dots[full..].iter().enumerate() { + let mut lane = [0.0f32; 16]; + let value = transform(F::splat(arch, dot)); + // SAFETY: `lane` has capacity for every supported `F`. + unsafe { value.store_simd(lane.as_mut_ptr()) }; + insert_topk(top, fanout, (full + offset) as u32, lane[0]); + } +} + +fn process_binary( + arch: F::Arch, + dots: &[f32], + scales: &[f32], + top: &mut TopK, + fanout: usize, + transform: Transform, +) where + F: SIMDVector + SIMDFloat, + Transform: Fn(F, F) -> F, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + let full = dots.len() / F::LANES * F::LANES; + for base in (0..full).step_by(F::LANES) { + // SAFETY: both slices contain the full SIMD chunk at `base`. + let dots = unsafe { F::load_simd(arch, dots.as_ptr().add(base)) }; + // SAFETY: shape validation guarantees `scales.len() == dots.len()`. + let scales = unsafe { F::load_simd(arch, scales.as_ptr().add(base)) }; + insert_lanes(transform(dots, scales), base, top, fanout); + } + for offset in 0..dots.len() - full { + let mut lane = [0.0f32; 16]; + let value = transform( + F::splat(arch, dots[full + offset]), + F::splat(arch, scales[full + offset]), + ); + // SAFETY: `lane` has capacity for every supported `F`. + unsafe { value.store_simd(lane.as_mut_ptr()) }; + insert_topk(top, fanout, (full + offset) as u32, lane[0]); + } +} + +fn insert_lanes(distances: F, base: usize, top: &mut TopK, fanout: usize) +where + F: SIMDVector + SIMDPartialOrd, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + let threshold = F::splat(distances.arch(), top[fanout - 1].1); + let eligible = distances.lt_simd(threshold); + if eligible.none() { + return; + } + + let mut values = [0.0f32; 16]; + // SAFETY: `values` has capacity for every f32 SIMD width DiskANN exposes. + unsafe { distances.store_simd(values.as_mut_ptr()) }; + let mut lanes = u64::from(eligible.bitmask().to_underlying()); + while lanes != 0 { + let lane = lanes.trailing_zeros() as usize; + lanes &= lanes - 1; + insert_topk(top, fanout, (base + lane) as u32, values[lane]); + } +} + +#[inline(always)] +#[cfg(test)] +fn distance(metric: Metric, dot: f32, row_scale: f32, leader_scale: f32) -> f32 { + match metric { + Metric::L2 => (-2.0f32).mul_add(dot, leader_scale), + Metric::CosineNormalized => 1.0 - dot, + Metric::InnerProduct => -dot, + Metric::Cosine => { + let denominator = row_scale.sqrt() * leader_scale; + let cosine = if denominator > 0.0 { + dot / denominator + } else { + 0.0 + }; + 1.0 - cosine + } + } +} + +#[inline(always)] +fn insert_topk(top: &mut TopK, fanout: usize, leader: u32, distance: f32) { + let threshold = fanout - 1; + if distance.partial_cmp(&top[threshold].1) != Some(std::cmp::Ordering::Less) { + return; + } + + top[threshold] = (leader, distance); + let mut position = threshold; + while position > 0 && top[position].1 < top[position - 1].1 { + top.swap(position, position - 1); + position -= 1; + } +} + +fn copy_ids(top: &TopK, output: &mut [u32]) { + for (destination, &(leader, _)) in output.iter_mut().zip(top) { + *destination = leader; + } +} + +#[cfg(test)] +mod tests; diff --git a/diskann-pipnn/src/partition_kernel/tests.rs b/diskann-pipnn/src/partition_kernel/tests.rs new file mode 100644 index 000000000..aaa03917b --- /dev/null +++ b/diskann-pipnn/src/partition_kernel/tests.rs @@ -0,0 +1,85 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; + +fn input(metric: Metric, leaders: usize) -> (Vec, Vec, Vec) { + let dots = (0..2 * leaders) + .map(|index| (((index * 13 + 7) % 29) as f32 - 14.0) * 0.125) + .collect(); + let row_scales = if metric == Metric::Cosine { + vec![0.0, 16.0] + } else { + Vec::new() + }; + let leader_scales = match metric { + Metric::L2 => (0..leaders).map(|leader| (leader + 1) as f32).collect(), + Metric::Cosine => (0..leaders) + .map(|leader| { + if leader == 0 { + 0.0 + } else { + (leader + 1) as f32 + } + }) + .collect(), + Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), + }; + (dots, row_scales, leader_scales) +} + +#[test] +fn scalar_reference_matches_runtime_dispatch() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + for leaders in [7, 17] { + let (dots, row_scales, leader_scales) = input(metric, leaders); + let input = PartitionTopK { + dots: &dots, + rows: 2, + leaders, + row_scales: &row_scales, + leader_scales: &leader_scales, + metric, + }; + for fanout in [1, 2, 6] { + let mut expected = vec![u32::MAX; input.rows * fanout]; + nearest_leaders(input, fanout, &mut expected).unwrap(); + + let mut actual = vec![u32::MAX; input.rows * fanout]; + process_rows_scalar(input, fanout, &mut actual); + + assert_eq!( + actual, expected, + "{metric:?}, leaders={leaders}, k={fanout}" + ); + } + } + } +} + +#[test] +fn scalar_distance_matches_metric_contract() { + assert_eq!(distance(Metric::L2, 2.0, 99.0, 9.0), 5.0); + assert_eq!(distance(Metric::CosineNormalized, 0.25, 99.0, 99.0), 0.75); + assert_eq!(distance(Metric::InnerProduct, 3.0, 99.0, 99.0), -3.0); + assert_eq!(distance(Metric::Cosine, 4.0, 4.0, 4.0), 0.5); + assert_eq!(distance(Metric::Cosine, 4.0, 0.0, 4.0), 1.0); +} + +#[test] +fn scalar_topk_orders_candidates_and_preserves_ties() { + let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + for (leader, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 1.0)] { + insert_topk(&mut top, 4, leader, distance); + } + insert_topk(&mut top, 4, 5, f32::NAN); + + assert_eq!(top[..4], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); +} diff --git a/diskann-pipnn/tests/leaf_kernel.rs b/diskann-pipnn/tests/leaf_kernel.rs new file mode 100644 index 000000000..5cead7e96 --- /dev/null +++ b/diskann-pipnn/tests/leaf_kernel.rs @@ -0,0 +1,489 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_pipnn::leaf_kernel::{ + nearest_leaf_neighbors, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, +}; +use diskann_vector::distance::Metric; +use std::cmp::Ordering; + +fn differential_input(metric: Metric, points: usize) -> Vec { + let mut dots = vec![f32::NAN; points * points]; + for row in 0..points { + dots[row * points + row] = if metric == Metric::Cosine && row == 0 { + 0.0 + } else if row == 2 { + 2.0 + } else { + 1.0 + (row % 5) as f32 + }; + for column in 0..row { + let pair = ((row * 17 + column * 11) % 23) as f32 - 11.0; + dots[row * points + column] = if row == points - 1 && column == 0 { + f32::NAN + } else if column == 1 || column == 2 { + 0.5 + } else { + pair * 0.03125 + }; + } + } + dots +} + +fn reference(input: LeafTopK<'_>, requested_k: usize) -> Vec { + let k = requested_k.min(input.points.saturating_sub(1)); + let mut output = vec![LeafNeighbor::default(); input.points * k]; + if k == 0 { + return output; + } + + let norms: Vec<_> = (0..input.points) + .map(|row| { + let diagonal = input.dots[row * input.points + row]; + if input.metric == Metric::Cosine { + if diagonal < f32::MIN_POSITIVE { + 0.0 + } else { + diagonal.sqrt() + } + } else { + diagonal + } + }) + .collect(); + + for row in 0..input.points { + let mut candidates = Vec::with_capacity(input.points - 1); + for position in 0..input.points { + if position == row { + continue; + } + let (lower_row, lower_column) = if row > position { + (row, position) + } else { + (position, row) + }; + let dot = input.dots[lower_row * input.points + lower_column]; + let clamp = |distance: f32| { + if distance < 0.0 { + 0.0 + } else { + distance + } + }; + let distance = match input.metric { + Metric::L2 => clamp(norms[row] + norms[position] - 2.0 * dot), + Metric::CosineNormalized => clamp(1.0 - dot), + Metric::InnerProduct => -dot, + Metric::Cosine => { + let denominator = norms[row] * norms[position]; + let similarity = if denominator == 0.0 { + 0.0 + } else { + dot / denominator + }; + clamp(1.0 - similarity) + } + }; + if distance.partial_cmp(&f32::INFINITY) == Some(Ordering::Less) { + candidates.push(LeafNeighbor::new(position as u32, distance)); + } + } + candidates.sort_by(|left, right| { + left.distance + .partial_cmp(&right.distance) + .expect("NaN distances were filtered") + }); + let count = candidates.len().min(k); + output[row * k..row * k + count].copy_from_slice(&candidates[..count]); + } + output +} + +#[test] +fn dispatch_matches_reference_across_simd_width_boundaries() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + for points in [7, 8, 9, 15, 16, 17, 64, 256, 512] { + let dots = differential_input(metric, points); + let input = LeafTopK { + dots: &dots, + points, + metric, + }; + // Covers every specialized insertion arm (1, 2, 3), the first width + // that falls back to the general bubble-up (4), and a wider row (5). + for requested_k in [1, 2, 3, 4, 5] { + let expected = reference(input, requested_k); + let mut actual = vec![LeafNeighbor::default(); expected.len()]; + let mut workspace = LeafTopKWorkspace::new(); + nearest_leaf_neighbors(input, requested_k, &mut actual, &mut workspace).unwrap(); + assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); + } + } + } +} + +fn run(dots: &[f32], points: usize, k: usize, metric: Metric) -> (usize, Vec) { + let actual_k = k.min(points.saturating_sub(1)); + let mut output = vec![LeafNeighbor::default(); points * actual_k]; + let mut workspace = LeafTopKWorkspace::new(); + let returned_k = nearest_leaf_neighbors( + LeafTopK { + dots, + points, + metric, + }, + k, + &mut output, + &mut workspace, + ) + .unwrap(); + assert_eq!(returned_k, actual_k); + (returned_k, output) +} + +#[test] +fn l2_scans_only_the_lower_triangle_and_breaks_ties_by_position() { + #[rustfmt::skip] + let dots = [ + 0.0, 999.0, 999.0, 999.0, + 0.0, 1.0, 999.0, 999.0, + 0.0, 0.0, 1.0, 999.0, + 0.0, 1.0, 1.0, 2.0, + ]; + + let (_, output) = run(&dots, 4, 2, Metric::L2); + + assert_eq!( + output, + [ + LeafNeighbor::new(1, 1.0), + LeafNeighbor::new(2, 1.0), + LeafNeighbor::new(0, 1.0), + LeafNeighbor::new(3, 1.0), + LeafNeighbor::new(0, 1.0), + LeafNeighbor::new(3, 1.0), + LeafNeighbor::new(1, 1.0), + LeafNeighbor::new(2, 1.0), + ] + ); +} + +#[test] +fn supports_every_leaf_metric() { + #[rustfmt::skip] + let dots = [ + 1.0, 77.0, 77.0, + 0.0, 1.0, 77.0, + -1.0, 0.5, 1.0, + ]; + + let cases = [ + (Metric::L2, [1, 2, 1]), + (Metric::Cosine, [1, 2, 1]), + (Metric::CosineNormalized, [1, 2, 1]), + (Metric::InnerProduct, [1, 2, 1]), + ]; + + for (metric, expected) in cases { + let (_, output) = run(&dots, 3, 1, metric); + let positions: Vec<_> = output.iter().map(|neighbor| neighbor.position).collect(); + assert_eq!(positions, expected, "metric {metric:?}"); + } +} + +#[test] +fn cosine_treats_zero_norm_as_zero_similarity() { + #[rustfmt::skip] + let dots = [ + 0.0, 11.0, 11.0, + 0.0, 1.0, 11.0, + 0.0, 0.0, 1.0, + ]; + + let (_, output) = run(&dots, 3, 2, Metric::Cosine); + + assert_eq!(output[0], LeafNeighbor::new(1, 1.0)); + assert_eq!(output[1], LeafNeighbor::new(2, 1.0)); +} + +#[test] +fn preserves_pipnn_metric_edge_semantics() { + #[rustfmt::skip] + let out_of_range = [ + 1.0, 0.0, + 2.0, 1.0, + ]; + assert_eq!(run(&out_of_range, 2, 1, Metric::L2).1[0].distance, 0.0); + assert_eq!( + run(&out_of_range, 2, 1, Metric::CosineNormalized).1[0].distance, + 0.0 + ); + assert_eq!(run(&out_of_range, 2, 1, Metric::Cosine).1[0].distance, 0.0); + + #[rustfmt::skip] + let opposite = [ + 1.0, 0.0, + -2.0, 1.0, + ]; + assert_eq!(run(&opposite, 2, 1, Metric::Cosine).1[0].distance, 3.0); + + let subnormal_squared_norm = f32::MIN_POSITIVE / 2.0; + #[rustfmt::skip] + let subnormal = [ + subnormal_squared_norm, 0.0, + 1.0, 1.0, + ]; + assert_eq!(run(&subnormal, 2, 1, Metric::Cosine).1[0].distance, 1.0); + + let minimum_normal_squared_norm = f32::MIN_POSITIVE; + #[rustfmt::skip] + let minimum_normal = [ + minimum_normal_squared_norm, 0.0, + minimum_normal_squared_norm.sqrt(), 1.0, + ]; + assert_eq!( + run(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, + 0.0 + ); +} + +#[test] +fn finite_max_distance_fills_the_final_simd_slot() { + let points = 9; + let mut dots = vec![0.0; points * points]; + dots[8 * points] = -f32::MAX; + + let (actual_k, output) = run(&dots, points, points - 1, Metric::InnerProduct); + + assert_eq!(actual_k, 8); + assert_eq!( + output[8 * actual_k + actual_k - 1], + LeafNeighbor::new(0, f32::MAX) + ); +} + +#[test] +fn every_metric_ignores_nan_pairs() { + #[rustfmt::skip] + let dots = [ + 1.0, 0.0, 0.0, + f32::NAN, 1.0, 0.0, + 0.5, 0.25, 1.0, + ]; + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let (_, output) = run(&dots, 3, 1, metric); + assert_eq!(output[0].position, 2, "metric {metric:?}"); + assert_eq!(output[1].position, 2, "metric {metric:?}"); + } +} + +#[test] +fn rejects_incomplete_neighbor_rows() { + #[rustfmt::skip] + let dots = [ + 1.0, 0.0, + f32::NAN, 1.0, + ]; + let mut output = [LeafNeighbor::default(); 2]; + let mut workspace = LeafTopKWorkspace::new(); + + let error = nearest_leaf_neighbors( + LeafTopK { + dots: &dots, + points: 2, + metric: Metric::L2, + }, + 1, + &mut output, + &mut workspace, + ) + .unwrap_err(); + + assert_eq!( + error, + LeafKernelError::InsufficientRankableNeighbors { + row: 0, + neighbors: 1, + } + ); +} + +#[test] +fn clamps_k_to_available_non_self_neighbors() { + #[rustfmt::skip] + let dots = [ + 1.0, 3.0, 3.0, + 0.0, 1.0, 3.0, + 0.0, 0.0, 1.0, + ]; + + let (actual_k, output) = run(&dots, 3, 99, Metric::L2); + + assert_eq!(actual_k, 2); + assert_eq!(output.len(), 6); + for (row, neighbors) in output.chunks_exact(actual_k).enumerate() { + assert!(neighbors + .iter() + .all(|neighbor| neighbor.position as usize != row)); + } +} + +#[test] +fn accepts_empty_singleton_and_zero_k_inputs() { + let mut workspace = LeafTopKWorkspace::new(); + let empty = LeafTopK { + dots: &[], + points: 0, + metric: Metric::L2, + }; + assert_eq!( + nearest_leaf_neighbors(empty, 2, &mut [], &mut workspace).unwrap(), + 0 + ); + + let singleton = LeafTopK { + dots: &[4.0], + points: 1, + metric: Metric::Cosine, + }; + assert_eq!( + nearest_leaf_neighbors(singleton, 2, &mut [], &mut workspace).unwrap(), + 0 + ); + + let pair = LeafTopK { + dots: &[1.0, 0.0, 0.0, 1.0], + points: 2, + metric: Metric::InnerProduct, + }; + assert_eq!( + nearest_leaf_neighbors(pair, 0, &mut [], &mut workspace).unwrap(), + 0 + ); +} + +#[test] +fn rejects_invalid_shapes_before_dispatch() { + let mut workspace = LeafTopKWorkspace::new(); + let error = nearest_leaf_neighbors( + LeafTopK { + dots: &[0.0; 8], + points: 3, + metric: Metric::L2, + }, + 1, + &mut [LeafNeighbor::default(); 3], + &mut workspace, + ) + .unwrap_err(); + assert_eq!( + error, + LeafKernelError::InvalidBufferLength { + buffer: "lower dot-product matrix", + expected: 9, + actual: 8, + } + ); + + let error = nearest_leaf_neighbors( + LeafTopK { + dots: &[0.0; 9], + points: 3, + metric: Metric::L2, + }, + 2, + &mut [LeafNeighbor::default(); 5], + &mut workspace, + ) + .unwrap_err(); + assert_eq!( + error, + LeafKernelError::InvalidBufferLength { + buffer: "output", + expected: 6, + actual: 5, + } + ); +} + +#[test] +fn rejects_shape_overflow_before_reading_buffers() { + let mut workspace = LeafTopKWorkspace::new(); + let error = nearest_leaf_neighbors( + LeafTopK { + dots: &[], + points: usize::MAX, + metric: Metric::L2, + }, + 1, + &mut [], + &mut workspace, + ) + .unwrap_err(); + + assert_eq!(error, LeafKernelError::TooManyPoints(usize::MAX)); +} + +#[test] +fn cosine_zero_norm_masks_nan_norm_at_simd_boundaries() { + for points in [9, 17] { + let mut dots = vec![0.0; points * points]; + dots[0] = 0.0; + for row in 1..points { + dots[row * points + row] = f32::NAN; + } + + let (_, output) = run(&dots, points, 1, Metric::Cosine); + for (row, neighbor) in output.iter().enumerate().skip(1) { + assert_eq!( + *neighbor, + LeafNeighbor::new(0, 1.0), + "n={points}, row={row}" + ); + } + } +} + +#[cfg(target_pointer_width = "64")] +#[test] +fn accepts_the_largest_representable_point_count_before_shape_validation() { + let points = u32::MAX as usize; + let expected = points.checked_mul(points).unwrap(); + let mut workspace = LeafTopKWorkspace::new(); + + let error = nearest_leaf_neighbors( + LeafTopK { + dots: &[], + points, + metric: Metric::InnerProduct, + }, + 0, + &mut [], + &mut workspace, + ) + .unwrap_err(); + + assert_eq!( + error, + LeafKernelError::InvalidBufferLength { + buffer: "lower dot-product matrix", + expected, + actual: 0, + } + ); +} diff --git a/diskann-pipnn/tests/partition_kernel.rs b/diskann-pipnn/tests/partition_kernel.rs new file mode 100644 index 000000000..bb90a8b9f --- /dev/null +++ b/diskann-pipnn/tests/partition_kernel.rs @@ -0,0 +1,442 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_pipnn::partition_kernel::{ + nearest_leaders, PartitionKernelError, PartitionTopK, MAX_PARTITION_FANOUT, +}; +use diskann_vector::distance::Metric; + +fn reference(input: PartitionTopK<'_>, fanout: usize) -> Vec { + let mut output = vec![u32::MAX; input.rows * fanout]; + for (row_index, (dots, output)) in input + .dots + .chunks_exact(input.leaders) + .zip(output.chunks_exact_mut(fanout)) + .enumerate() + { + let row_scale = input.row_scales.get(row_index).copied().unwrap_or(0.0); + let mut candidates: Vec<_> = dots + .iter() + .enumerate() + .filter_map(|(leader, &dot)| { + let leader_scale = input.leader_scales.get(leader).copied().unwrap_or(0.0); + let distance = match input.metric { + Metric::L2 => leader_scale - 2.0 * dot, + Metric::CosineNormalized => 1.0 - dot, + Metric::InnerProduct => -dot, + Metric::Cosine => { + let denominator = row_scale.sqrt() * leader_scale; + 1.0 - if denominator > 0.0 { + dot / denominator + } else { + 0.0 + } + } + }; + (distance.partial_cmp(&f32::INFINITY) == Some(std::cmp::Ordering::Less)) + .then_some((leader as u32, distance)) + }) + .collect(); + candidates.sort_by(|left, right| left.1.partial_cmp(&right.1).unwrap()); + for (destination, (leader, _)) in output.iter_mut().zip(candidates) { + *destination = leader; + } + } + output +} + +fn differential_input(metric: Metric, leaders: usize) -> (Vec, Vec, Vec) { + let dots = (0..2 * leaders) + .map(|index| { + let leader = index % leaders; + let row = index / leaders; + let base = ((leader * 13 + row * 7) % 19) as f32 - 9.0; + if leader == 2 || leader == 3 { + 1.0 + } else if leader + 1 == leaders { + f32::NAN + } else { + base * 0.25 + } + }) + .collect(); + let row_scales = if metric == Metric::Cosine { + vec![0.0, 16.0] + } else { + Vec::new() + }; + let leader_scales = match metric { + Metric::Cosine => (0..leaders) + .map(|leader| { + if leader == 1 { + 0.0 + } else if leader == 2 || leader == 3 { + 3.0 + } else { + 1.0 + leader as f32 + } + }) + .collect(), + Metric::L2 => (0..leaders) + .map(|leader| { + let norm = if leader == 2 || leader == 3 { + 3.0 + } else { + leader as f32 + 1.0 + }; + norm * norm + }) + .collect(), + Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), + }; + (dots, row_scales, leader_scales) +} + +#[test] +fn dispatch_matches_reference_across_simd_width_boundaries() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + for leaders in [7, 8, 9, 15, 16, 17] { + let (dots, row_scales, leader_scales) = differential_input(metric, leaders); + for fanout in [1, 2, 16] { + if fanout >= leaders { + continue; + } + let input = PartitionTopK { + dots: &dots, + rows: 2, + leaders, + row_scales: &row_scales, + leader_scales: &leader_scales, + metric, + }; + let expected = reference(input, fanout); + let mut actual = vec![u32::MAX; expected.len()]; + nearest_leaders(input, fanout, &mut actual).unwrap(); + assert_eq!( + actual, expected, + "{metric:?}, leaders={leaders}, k={fanout}" + ); + } + } + } +} + +#[test] +fn l2_keeps_the_first_leader_when_boundary_distances_tie() { + #[rustfmt::skip] + let dots = [ + 0.0, 0.0, 0.0, 0.0, + 0.0, 2.0, 4.0, 6.0, + ]; + let leader_squared_norms = [0.0, 1.0, 4.0, 9.0]; + let mut assignments = [u32::MAX; 4]; + + let input = PartitionTopK { + dots: &dots, + rows: 2, + leaders: 4, + row_scales: &[], + leader_scales: &leader_squared_norms, + metric: Metric::L2, + }; + + nearest_leaders(input, 2, &mut assignments).unwrap(); + + assert_eq!(assignments, [0, 1, 2, 1]); +} + +#[test] +fn supports_every_partition_metric() { + #[rustfmt::skip] + let dots = [ + 1.0, 0.0, -1.0, + 2.0, 6.0, 0.0, + ]; + + let cases = [ + (Metric::L2, &[][..], &[1.0, 4.0, 9.0][..], [0, 1, 1, 0]), + ( + Metric::Cosine, + &[1.0, 4.0][..], + &[1.0, 2.0, 3.0][..], + [0, 1, 1, 0], + ), + (Metric::CosineNormalized, &[][..], &[][..], [0, 1, 1, 0]), + (Metric::InnerProduct, &[][..], &[][..], [0, 1, 1, 0]), + ]; + + for (metric, row_scales, leader_scales, expected) in cases { + let mut assignments = [u32::MAX; 4]; + nearest_leaders( + PartitionTopK { + dots: &dots, + rows: 2, + leaders: 3, + row_scales, + leader_scales, + metric, + }, + 2, + &mut assignments, + ) + .unwrap(); + + assert_eq!(assignments, expected, "metric {metric:?}"); + } +} + +#[test] +fn cosine_treats_a_zero_norm_as_zero_similarity() { + let mut assignments = [u32::MAX; 2]; + + nearest_leaders( + PartitionTopK { + dots: &[100.0, -100.0], + rows: 1, + leaders: 2, + row_scales: &[0.0], + leader_scales: &[1.0, 1.0], + metric: Metric::Cosine, + }, + 2, + &mut assignments, + ) + .unwrap(); + + assert_eq!(assignments, [0, 1]); +} + +#[test] +fn finite_max_distance_fills_the_final_simd_slot() { + let mut assignments = [u32::MAX; 8]; + let mut dots = [0.0; 8]; + dots[7] = -f32::MAX; + + nearest_leaders( + PartitionTopK { + dots: &dots, + rows: 1, + leaders: 8, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }, + 8, + &mut assignments, + ) + .unwrap(); + + assert_eq!(assignments, [0, 1, 2, 3, 4, 5, 6, 7]); +} + +#[test] +fn ignores_nan_distances_without_displacing_finite_leaders() { + let mut assignments = [u32::MAX; 2]; + + nearest_leaders( + PartitionTopK { + dots: &[f32::NAN, 3.0, 2.0], + rows: 1, + leaders: 3, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }, + 2, + &mut assignments, + ) + .unwrap(); + + assert_eq!(assignments, [1, 2]); +} + +#[test] +fn rejects_rows_with_too_few_rankable_distances() { + let error = nearest_leaders( + PartitionTopK { + dots: &[f32::NAN, 3.0], + rows: 1, + leaders: 2, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }, + 2, + &mut [u32::MAX; 2], + ) + .unwrap_err(); + + assert_eq!( + error, + PartitionKernelError::InsufficientRankableDistances { row: 0, fanout: 2 } + ); +} + +#[test] +fn accepts_empty_rows_and_zero_fanout() { + nearest_leaders( + PartitionTopK { + dots: &[], + rows: 0, + leaders: 3, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }, + 2, + &mut [], + ) + .unwrap(); + + nearest_leaders( + PartitionTopK { + dots: &[1.0, 2.0, 3.0], + rows: 1, + leaders: 3, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }, + 0, + &mut [], + ) + .unwrap(); + + // `u32::MAX` leaders still have positions representable by `u32`: the + // largest position is `u32::MAX - 1`. An empty batch lets us exercise the + // validation boundary without allocating the declared tile. + nearest_leaders( + PartitionTopK { + dots: &[], + rows: 0, + leaders: u32::MAX as usize, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }, + 0, + &mut [], + ) + .unwrap(); + + #[cfg(target_pointer_width = "64")] + assert_eq!( + nearest_leaders( + PartitionTopK { + dots: &[], + rows: 0, + leaders: u32::MAX as usize + 1, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }, + 0, + &mut [], + ), + Err(PartitionKernelError::TooManyLeaders(u32::MAX as usize + 1)) + ); +} + +#[test] +fn rejects_inconsistent_shapes_and_fanout() { + let base = PartitionTopK { + dots: &[0.0; 6], + rows: 2, + leaders: 3, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }; + + assert_eq!( + nearest_leaders( + PartitionTopK { + dots: &[0.0; 5], + ..base + }, + 2, + &mut [0; 4], + ), + Err(PartitionKernelError::InvalidBufferLength { + buffer: "dot-product tile", + expected: 6, + actual: 5, + }) + ); + assert_eq!( + nearest_leaders(base, 2, &mut [0; 3]), + Err(PartitionKernelError::InvalidBufferLength { + buffer: "output", + expected: 4, + actual: 3, + }) + ); + assert_eq!( + nearest_leaders(base, MAX_PARTITION_FANOUT + 1, &mut []), + Err(PartitionKernelError::InvalidFanout { + fanout: MAX_PARTITION_FANOUT + 1, + leaders: 3, + maximum: MAX_PARTITION_FANOUT, + }) + ); + + let one_leader = PartitionTopK { + dots: &[0.0], + rows: 1, + leaders: 1, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }; + assert_eq!( + nearest_leaders(one_leader, 2, &mut []), + Err(PartitionKernelError::InvalidFanout { + fanout: 2, + leaders: 1, + maximum: MAX_PARTITION_FANOUT, + }) + ); + + let exact_maximum = PartitionTopK { + dots: &[], + rows: 0, + leaders: MAX_PARTITION_FANOUT, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }; + nearest_leaders(exact_maximum, MAX_PARTITION_FANOUT, &mut []).unwrap(); +} + +#[test] +fn rejects_shape_overflow_before_reading_buffers() { + let error = nearest_leaders( + PartitionTopK { + dots: &[], + rows: usize::MAX, + leaders: 2, + row_scales: &[], + leader_scales: &[], + metric: Metric::InnerProduct, + }, + 1, + &mut [], + ) + .unwrap_err(); + + assert_eq!( + error, + PartitionKernelError::ShapeOverflow { + buffer: "dot-product tile", + rows: usize::MAX, + cols: 2, + } + ); +} diff --git a/diskann-wide/src/arch/aarch64/f32x2_.rs b/diskann-wide/src/arch/aarch64/f32x2_.rs index 318227ca5..f0b7431a8 100644 --- a/diskann-wide/src/arch/aarch64/f32x2_.rs +++ b/diskann-wide/src/arch/aarch64/f32x2_.rs @@ -31,6 +31,7 @@ macros::aarch64_define_loadstore!(f32x2, vld1_f32, internal::load_first::f32x2, helpers::unsafe_map_binary_op!(f32x2, std::ops::Add, add, vadd_f32, "neon"); helpers::unsafe_map_binary_op!(f32x2, std::ops::Sub, sub, vsub_f32, "neon"); helpers::unsafe_map_binary_op!(f32x2, std::ops::Mul, mul, vmul_f32, "neon"); +helpers::unsafe_map_binary_op!(f32x2, std::ops::Div, div, vdiv_f32, "neon"); macros::aarch64_define_fma!(f32x2, vfma_f32); macros::aarch64_define_cmp!( @@ -90,6 +91,7 @@ mod tests { test_utils::ops::test_add!(f32x2, 0xcd7a8fea9a3fb727, test_neon()); test_utils::ops::test_sub!(f32x2, 0x3f6562c94c923238, test_neon()); test_utils::ops::test_mul!(f32x2, 0x07e48666c0fc564c, test_neon()); + test_utils::ops::test_div!(f32x2, 0xa0352efeb9bc5ca5, test_neon()); test_utils::ops::test_fma!(f32x2, 0xcfde9d031302cf2c, test_neon()); test_utils::ops::test_cmp!(f32x2, 0xc4f468b224622326, test_neon()); diff --git a/diskann-wide/src/arch/aarch64/f32x4_.rs b/diskann-wide/src/arch/aarch64/f32x4_.rs index 82cf39107..83779dacf 100644 --- a/diskann-wide/src/arch/aarch64/f32x4_.rs +++ b/diskann-wide/src/arch/aarch64/f32x4_.rs @@ -32,6 +32,7 @@ macros::aarch64_splitjoin!(f32x4, f32x2, vget_low_f32, vget_high_f32, vcombine_f helpers::unsafe_map_binary_op!(f32x4, std::ops::Add, add, vaddq_f32, "neon"); helpers::unsafe_map_binary_op!(f32x4, std::ops::Sub, sub, vsubq_f32, "neon"); helpers::unsafe_map_binary_op!(f32x4, std::ops::Mul, mul, vmulq_f32, "neon"); +helpers::unsafe_map_binary_op!(f32x4, std::ops::Div, div, vdivq_f32, "neon"); helpers::unsafe_map_unary_op!(f32x4, SIMDAbs, abs_simd, vabsq_f32, "neon"); macros::aarch64_define_fma!(f32x4, vfmaq_f32); @@ -187,6 +188,7 @@ mod tests { test_utils::ops::test_add!(f32x4, 0xcd7a8fea9a3fb727, test_neon()); test_utils::ops::test_sub!(f32x4, 0x3f6562c94c923238, test_neon()); test_utils::ops::test_mul!(f32x4, 0x07e48666c0fc564c, test_neon()); + test_utils::ops::test_div!(f32x4, 0xa0352efeb9bc5ca5, test_neon()); test_utils::ops::test_fma!(f32x4, 0xcfde9d031302cf2c, test_neon()); test_utils::ops::test_abs!(f32x4, 0xb8f702ba85375041, test_neon()); test_utils::ops::test_minmax!(f32x4, 0x6d7fc8ed6d852187, test_neon()); diff --git a/diskann-wide/src/arch/x86_64/v3/f32x16_.rs b/diskann-wide/src/arch/x86_64/v3/f32x16_.rs index 836193a45..b93c861e7 100644 --- a/diskann-wide/src/arch/x86_64/v3/f32x16_.rs +++ b/diskann-wide/src/arch/x86_64/v3/f32x16_.rs @@ -54,6 +54,7 @@ mod test_x86_f32 { test_utils::ops::test_add!(f32x16, 0xa8989b97ca888d11, V3::new_checked_uncached()); test_utils::ops::test_sub!(f32x16, 0xb2554fc13fdc1182, V3::new_checked_uncached()); test_utils::ops::test_mul!(f32x16, 0x23becaa968b0cd71, V3::new_checked_uncached()); + test_utils::ops::test_div!(f32x16, 0x6fd16af08fa1f498, V3::new_checked_uncached()); test_utils::ops::test_fma!(f32x16, 0x32a814070a93df4e, V3::new_checked_uncached()); test_utils::ops::test_minmax!(f32x16, 0x6d7fc8ed6d852187, V3::new_checked_uncached()); test_utils::ops::test_abs!(f32x16, 0x6799e60873a2efe2, V3::new_checked_uncached()); diff --git a/diskann-wide/src/arch/x86_64/v3/f32x4_.rs b/diskann-wide/src/arch/x86_64/v3/f32x4_.rs index 60ffa4477..45cc64a4a 100644 --- a/diskann-wide/src/arch/x86_64/v3/f32x4_.rs +++ b/diskann-wide/src/arch/x86_64/v3/f32x4_.rs @@ -31,6 +31,7 @@ macros::x86_define_default!(f32x4, _mm_setzero_ps, "sse"); helpers::unsafe_map_binary_op!(f32x4, std::ops::Add, add, _mm_add_ps, "sse"); helpers::unsafe_map_binary_op!(f32x4, std::ops::Sub, sub, _mm_sub_ps, "sse"); helpers::unsafe_map_binary_op!(f32x4, std::ops::Mul, mul, _mm_mul_ps, "sse"); +helpers::unsafe_map_binary_op!(f32x4, std::ops::Div, div, _mm_div_ps, "sse"); impl f32x4 { #[inline(always)] @@ -253,6 +254,7 @@ mod test_x86_f32 { test_utils::ops::test_add!(f32x4, 0xcd7a8fea9a3fb727, V3::new_checked_uncached()); test_utils::ops::test_sub!(f32x4, 0x3f6562c94c923238, V3::new_checked_uncached()); test_utils::ops::test_mul!(f32x4, 0x07e48666c0fc564c, V3::new_checked_uncached()); + test_utils::ops::test_div!(f32x4, 0xa0352efeb9bc5ca5, V3::new_checked_uncached()); test_utils::ops::test_fma!(f32x4, 0xcfde9d031302cf2c, V3::new_checked_uncached()); test_utils::ops::test_minmax!(f32x4, 0x6d7fc8ed6d852187, V3::new_checked_uncached()); test_utils::ops::test_abs!(f32x4, 0x8e6d9944c9c43a74, V3::new_checked_uncached()); diff --git a/diskann-wide/src/arch/x86_64/v3/f32x8_.rs b/diskann-wide/src/arch/x86_64/v3/f32x8_.rs index 054b249e8..48ecea7ae 100644 --- a/diskann-wide/src/arch/x86_64/v3/f32x8_.rs +++ b/diskann-wide/src/arch/x86_64/v3/f32x8_.rs @@ -33,6 +33,7 @@ macros::x86_splitjoin!(f32x8, f32x4, _mm256_extractf128_ps, _mm256_set_m128, "av helpers::unsafe_map_binary_op!(f32x8, std::ops::Add, add, _mm256_add_ps, "avx"); helpers::unsafe_map_binary_op!(f32x8, std::ops::Sub, sub, _mm256_sub_ps, "avx"); helpers::unsafe_map_binary_op!(f32x8, std::ops::Mul, mul, _mm256_mul_ps, "avx"); +helpers::unsafe_map_binary_op!(f32x8, std::ops::Div, div, _mm256_div_ps, "avx"); impl f32x8 { #[inline(always)] @@ -266,6 +267,7 @@ mod test_x86_f32 { test_utils::ops::test_add!(f32x8, 0x3824379d4a43a416, V3::new_checked_uncached()); test_utils::ops::test_sub!(f32x8, 0x548fc74c07ba425d, V3::new_checked_uncached()); test_utils::ops::test_mul!(f32x8, 0x6d340672ff91b256, V3::new_checked_uncached()); + test_utils::ops::test_div!(f32x8, 0x776f54898c62dd0b, V3::new_checked_uncached()); test_utils::ops::test_fma!(f32x8, 0x5f566d8968d4d201, V3::new_checked_uncached()); test_utils::ops::test_minmax!(f32x8, 0x6d7fc8ed6d852187, V3::new_checked_uncached()); test_utils::ops::test_abs!(f32x8, 0x2a4a9651d8ebe912, V3::new_checked_uncached()); diff --git a/diskann-wide/src/arch/x86_64/v4/f32x16_.rs b/diskann-wide/src/arch/x86_64/v4/f32x16_.rs index d38465f90..ff5911999 100644 --- a/diskann-wide/src/arch/x86_64/v4/f32x16_.rs +++ b/diskann-wide/src/arch/x86_64/v4/f32x16_.rs @@ -57,6 +57,7 @@ impl crate::SplitJoin for f32x16 { helpers::unsafe_map_binary_op!(f32x16, std::ops::Add, add, _mm512_add_ps, "avx512f"); helpers::unsafe_map_binary_op!(f32x16, std::ops::Sub, sub, _mm512_sub_ps, "avx512f"); helpers::unsafe_map_binary_op!(f32x16, std::ops::Mul, mul, _mm512_mul_ps, "avx512f"); +helpers::unsafe_map_binary_op!(f32x16, std::ops::Div, div, _mm512_div_ps, "avx512f"); impl f32x16 { #[inline(always)] @@ -240,6 +241,7 @@ mod test_x86_f32 { test_utils::ops::test_add!(f32x16, 0xa8989b97ca888d11, V4::new_checked_uncached()); test_utils::ops::test_sub!(f32x16, 0xb2554fc13fdc1182, V4::new_checked_uncached()); test_utils::ops::test_mul!(f32x16, 0x23becaa968b0cd71, V4::new_checked_uncached()); + test_utils::ops::test_div!(f32x16, 0x6fd16af08fa1f498, V4::new_checked_uncached()); test_utils::ops::test_fma!(f32x16, 0x32a814070a93df4e, V4::new_checked_uncached()); test_utils::ops::test_minmax!(f32x16, 0x6d7fc8ed6d852187, V4::new_checked_uncached()); test_utils::ops::test_abs!(f32x16, 0x6799e60873a2efe2, V4::new_checked_uncached()); diff --git a/diskann-wide/src/arch/x86_64/v4/f32x4_.rs b/diskann-wide/src/arch/x86_64/v4/f32x4_.rs index 328dba4d2..7028b2fdc 100644 --- a/diskann-wide/src/arch/x86_64/v4/f32x4_.rs +++ b/diskann-wide/src/arch/x86_64/v4/f32x4_.rs @@ -33,6 +33,7 @@ macros::x86_retarget!(f32x4 => v3::f32x4); helpers::unsafe_map_binary_op!(f32x4, std::ops::Add, add, _mm_add_ps, "sse"); helpers::unsafe_map_binary_op!(f32x4, std::ops::Sub, sub, _mm_sub_ps, "sse"); helpers::unsafe_map_binary_op!(f32x4, std::ops::Mul, mul, _mm_mul_ps, "sse"); +helpers::unsafe_map_binary_op!(f32x4, std::ops::Div, div, _mm_div_ps, "sse"); impl f32x4 { #[inline(always)] @@ -210,6 +211,7 @@ mod test_x86_f32 { test_utils::ops::test_add!(f32x4, 0xcd7a8fea9a3fb727, V4::new_checked_uncached()); test_utils::ops::test_sub!(f32x4, 0x3f6562c94c923238, V4::new_checked_uncached()); test_utils::ops::test_mul!(f32x4, 0x07e48666c0fc564c, V4::new_checked_uncached()); + test_utils::ops::test_div!(f32x4, 0xa0352efeb9bc5ca5, V4::new_checked_uncached()); test_utils::ops::test_fma!(f32x4, 0xcfde9d031302cf2c, V4::new_checked_uncached()); test_utils::ops::test_minmax!(f32x4, 0x6d7fc8ed6d852187, V4::new_checked_uncached()); test_utils::ops::test_abs!(f32x4, 0x8e6d9944c9c43a74, V4::new_checked_uncached()); diff --git a/diskann-wide/src/arch/x86_64/v4/f32x8_.rs b/diskann-wide/src/arch/x86_64/v4/f32x8_.rs index 3158ffc1d..d38de49de 100644 --- a/diskann-wide/src/arch/x86_64/v4/f32x8_.rs +++ b/diskann-wide/src/arch/x86_64/v4/f32x8_.rs @@ -36,6 +36,7 @@ macros::x86_retarget!(f32x8 => v3::f32x8); helpers::unsafe_map_binary_op!(f32x8, std::ops::Add, add, _mm256_add_ps, "avx"); helpers::unsafe_map_binary_op!(f32x8, std::ops::Sub, sub, _mm256_sub_ps, "avx"); helpers::unsafe_map_binary_op!(f32x8, std::ops::Mul, mul, _mm256_mul_ps, "avx"); +helpers::unsafe_map_binary_op!(f32x8, std::ops::Div, div, _mm256_div_ps, "avx"); impl f32x8 { #[inline(always)] @@ -206,6 +207,7 @@ mod test_x86_f32 { test_utils::ops::test_add!(f32x8, 0x3824379d4a43a416, V4::new_checked_uncached()); test_utils::ops::test_sub!(f32x8, 0x548fc74c07ba425d, V4::new_checked_uncached()); test_utils::ops::test_mul!(f32x8, 0x6d340672ff91b256, V4::new_checked_uncached()); + test_utils::ops::test_div!(f32x8, 0x776f54898c62dd0b, V4::new_checked_uncached()); test_utils::ops::test_fma!(f32x8, 0x5f566d8968d4d201, V4::new_checked_uncached()); test_utils::ops::test_minmax!(f32x8, 0x6d7fc8ed6d852187, V4::new_checked_uncached()); test_utils::ops::test_abs!(f32x8, 0x2a4a9651d8ebe912, V4::new_checked_uncached()); diff --git a/diskann-wide/src/doubled.rs b/diskann-wide/src/doubled.rs index 4f0ecd681..0a1ca5015 100644 --- a/diskann-wide/src/doubled.rs +++ b/diskann-wide/src/doubled.rs @@ -205,6 +205,15 @@ impl> std::ops::Mul for Doubled { } } +impl> std::ops::Div for Doubled { + type Output = Self; + + #[inline(always)] + fn div(self, rhs: Self) -> Self { + Self(self.0 / rhs.0, self.1 / rhs.1) + } +} + impl> std::ops::BitAnd for Doubled { type Output = Self; #[inline(always)] diff --git a/diskann-wide/src/emulated.rs b/diskann-wide/src/emulated.rs index a91d5d767..6b444bd45 100644 --- a/diskann-wide/src/emulated.rs +++ b/diskann-wide/src/emulated.rs @@ -198,6 +198,15 @@ where } } +impl std::ops::Div for Emulated { + type Output = Self; + + #[inline(always)] + fn div(self, rhs: Self) -> Self { + Self::from_arch_fn(self.1, |i| self.0[i] / rhs.0[i]) + } +} + /// MulAdd impl SIMDMulAdd for Emulated where @@ -886,6 +895,10 @@ mod test_emulated { test_emulated!(f32, 4); test_emulated!(f32, 8); test_emulated!(f32, 16); + test_utils::ops::test_div!(Emulated, 0x32f0d2991be50f13, SC); + test_utils::ops::test_div!(Emulated, 0xf65f08475f5e30c9, SC); + test_utils::ops::test_div!(Emulated, 0x31e044b2369bf812, SC); + test_utils::ops::test_div!(Emulated, 0x87f74cf00a528a2d, SC); // test_emulated!(f64, 8); // unsigned integer diff --git a/diskann-wide/src/test_utils/ops.rs b/diskann-wide/src/test_utils/ops.rs index 24da31892..8c691f001 100644 --- a/diskann-wide/src/test_utils/ops.rs +++ b/diskann-wide/src/test_utils/ops.rs @@ -425,6 +425,38 @@ macro_rules! test_mul { }; } +macro_rules! test_div { + ($wide:ident $(< $($ps:tt),+ >)?, $seed:literal, $arch:expr) => { + paste::paste! { + #[test] + fn []() { + use $crate::SIMDVector; + type T = $wide $(< $($ps),+>)?; + type Scalar = ::Scalar; + + if let Some(arch) = $arch { + let f = move |a: &[Scalar], b: &[Scalar]| { + let got = ( + ::from_array(arch, a.try_into().unwrap()) / + ::from_array(arch, b.try_into().unwrap()) + ).to_array(); + test_utils::test_binary_op( + &a, + &b, + &got, + &|l: Scalar, r: Scalar| { l / r }, + "binary division", + ) + }; + + let n = T::LANES; + $crate::test_utils::driver::drive_binary(&f, (n, n), $seed); + } + } + } + }; +} + macro_rules! test_fma { ($wide:ident $(< $($ps:tt),+ >)?, $seed:literal, $arch:expr) => { paste::paste! { @@ -1110,6 +1142,7 @@ pub(crate) use test_add; pub(crate) use test_bitops; pub(crate) use test_cast; pub(crate) use test_cmp; +pub(crate) use test_div; pub(crate) use test_fma; pub(crate) use test_lossless_convert; pub(crate) use test_minmax;