From d680d5d9a8b1a0e57aab0997144bfeafd33ba67d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:02:36 +0000 Subject: [PATCH 01/26] linalg: add fallible lower-triangular AAT kernel --- diskann-linalg/src/faer.rs | 23 +++++ diskann-linalg/src/lib.rs | 43 ++++++++++ diskann-linalg/tests/sgemm_aat_lower.rs | 109 ++++++++++++++++++++++++ 3 files changed, 175 insertions(+) create mode 100644 diskann-linalg/tests/sgemm_aat_lower.rs 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, + } + ); +} From bf6a5eb0f2c48034b70f1a3eecb221e972ddb08f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:02:37 +0000 Subject: [PATCH 02/26] pipnn: add dispatched numerical kernels --- Cargo.lock | 11 + Cargo.toml | 2 + diskann-pipnn/Cargo.toml | 27 + diskann-pipnn/src/leaf_kernel.rs | 764 +++++++++++++++++++++ diskann-pipnn/src/lib.rs | 9 + diskann-pipnn/src/partition_kernel.rs | 438 ++++++++++++ diskann-pipnn/tests/leaf_kernel.rs | 454 ++++++++++++ diskann-pipnn/tests/partition_kernel.rs | 419 +++++++++++ diskann-wide/src/arch/aarch64/f32x2_.rs | 2 + diskann-wide/src/arch/aarch64/f32x4_.rs | 2 + diskann-wide/src/arch/x86_64/v3/f32x16_.rs | 1 + diskann-wide/src/arch/x86_64/v3/f32x4_.rs | 2 + diskann-wide/src/arch/x86_64/v3/f32x8_.rs | 2 + diskann-wide/src/arch/x86_64/v4/f32x16_.rs | 2 + diskann-wide/src/arch/x86_64/v4/f32x4_.rs | 2 + diskann-wide/src/arch/x86_64/v4/f32x8_.rs | 2 + diskann-wide/src/doubled.rs | 9 + diskann-wide/src/emulated.rs | 13 + diskann-wide/src/test_utils/ops.rs | 33 + 19 files changed, 2194 insertions(+) create mode 100644 diskann-pipnn/Cargo.toml create mode 100644 diskann-pipnn/src/leaf_kernel.rs create mode 100644 diskann-pipnn/src/lib.rs create mode 100644 diskann-pipnn/src/partition_kernel.rs create mode 100644 diskann-pipnn/tests/leaf_kernel.rs create mode 100644 diskann-pipnn/tests/partition_kernel.rs 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-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/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs new file mode 100644 index 000000000..a913ba138 --- /dev/null +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -0,0 +1,764 @@ +/* + * 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; +#[cfg(target_arch = "x86_64")] +use diskann_wide::{SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; + +/// Widest f32 SIMD lane count DiskANN dispatches to, used to size lane scratch. +#[cfg(target_arch = "x86_64")] +const MAX_LANES: usize = 16; + +#[cfg(target_arch = "x86_64")] +const L2: u8 = 0; +#[cfg(target_arch = "x86_64")] +const COSINE_NORMALIZED: u8 = 1; +#[cfg(target_arch = "x86_64")] +const INNER_PRODUCT: u8 = 2; +#[cfg(target_arch = "x86_64")] +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::MAX) + } +} + +/// 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::MAX, + )?; + 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::MAX); + + 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> { + if len > values.len() { + let additional = len - values.len(); + values + .try_reserve(additional) + .map_err(|_| LeafKernelError::Allocation { buffer, additional })?; + values.resize(len, value); + } else { + values.truncate(len); + } + 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_scalar(self) { + process_pairs_scalar(self.input, self.k, self.output, self.norms, self.worst); + } + + #[cfg(target_arch = "x86_64")] + fn run_simd(self, arch: F::Arch) + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, + { + match self.k { + 1 => self.run_fused::(arch), + 2 => self.run_fused::(arch), + 3 => self.run_fused::(arch), + _ => process_pairs_simd_dynamic::( + arch, + self.input, + self.k, + self.output, + self.norms, + self.worst, + ), + } + } + + #[cfg(target_arch = "x86_64")] + 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<'_, '_, '_> { + #[inline(always)] + fn run(self, _: diskann_wide::arch::Scalar) { + self.run_scalar(); + } +} + +#[cfg(target_arch = "x86_64")] +impl diskann_wide::arch::Target for LeafKernel<'_, '_, '_> { + #[inline(always)] + fn run(self, arch: diskann_wide::arch::x86_64::V3) { + diskann_wide::alias!(F32x8 = ::f32x8); + self.run_simd::(arch); + } +} + +#[cfg(target_arch = "x86_64")] +impl diskann_wide::arch::Target for LeafKernel<'_, '_, '_> { + #[inline(always)] + fn run(self, arch: diskann_wide::arch::x86_64::V4) { + diskann_wide::alias!(F32x16 = ::f32x16); + self.run_simd::(arch); + } +} + +#[cfg(target_arch = "aarch64")] +impl diskann_wide::arch::Target for LeafKernel<'_, '_, '_> { + #[inline(always)] + fn run(self, arch: diskann_wide::arch::aarch64::Neon) { + let _scalar = arch.retarget(); + self.run_scalar(); + } +} + +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); + } + } +} + +#[cfg(target_arch = "x86_64")] +/// 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. +#[cfg(target_arch = "x86_64")] +#[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 }; + } +} + +#[cfg(target_arch = "x86_64")] +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`. +#[cfg(target_arch = "x86_64")] +#[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 } + } + } +} + +#[cfg(target_arch = "x86_64")] +#[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); + match metric { + Metric::L2 => { + let distance = row_norm + column_norm - F::splat(arch, 2.0) * dot; + zero.max_simd(distance) + } + Metric::CosineNormalized => { + let distance = F::splat(arch, 1.0) - dot; + zero.max_simd(distance) + } + Metric::InnerProduct => zero - dot, + Metric::Cosine => { + let one = F::splat(arch, 1.0); + let denominator = row_norm * column_norm; + let zero_denominator = denominator.eq_simd(zero); + let safe_denominator = zero_denominator.select(one, denominator); + let cosine = zero_denominator.select(zero, dot / safe_denominator); + let distance = one - cosine; + // Comparisons with NaN are false, so this explicit lower clamp + // preserves non-rankable NaNs while matching the existing PiPNN + // distance formulas for finite values. + zero.max_simd(distance) + } + } +} + +#[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)] +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; +} 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..b39c05aa8 --- /dev/null +++ b/diskann-pipnn/src/partition_kernel.rs @@ -0,0 +1,438 @@ +/* + * 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; +#[cfg(target_arch = "x86_64")] +use diskann_wide::{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_scalar(self) { + process_rows_scalar(self.input, self.fanout, self.output); + } + + #[cfg(target_arch = "x86_64")] + 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<'_, '_> { + #[inline(always)] + fn run(self, _: diskann_wide::arch::Scalar) { + self.run_scalar(); + } +} + +#[cfg(target_arch = "x86_64")] +impl diskann_wide::arch::Target for PartitionKernel<'_, '_> { + #[inline(always)] + fn run(self, arch: diskann_wide::arch::x86_64::V3) { + diskann_wide::alias!(F32x8 = ::f32x8); + self.run_simd::(arch); + } +} + +#[cfg(target_arch = "x86_64")] +impl diskann_wide::arch::Target for PartitionKernel<'_, '_> { + #[inline(always)] + fn run(self, arch: diskann_wide::arch::x86_64::V4) { + diskann_wide::alias!(F32x16 = ::f32x16); + self.run_simd::(arch); + } +} + +#[cfg(target_arch = "aarch64")] +impl diskann_wide::arch::Target for PartitionKernel<'_, '_> { + #[inline(always)] + fn run(self, arch: diskann_wide::arch::aarch64::Neon) { + let _scalar = arch.retarget(); + self.run_scalar(); + } +} + +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::MAX); 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); + } +} + +#[cfg(target_arch = "x86_64")] +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::MAX); 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); + } +} + +#[cfg(target_arch = "x86_64")] +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 + }); +} + +#[cfg(target_arch = "x86_64")] +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]); + } +} + +#[cfg(target_arch = "x86_64")] +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]); + } +} + +#[cfg(target_arch = "x86_64")] +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)] +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; + } +} diff --git a/diskann-pipnn/tests/leaf_kernel.rs b/diskann-pipnn/tests/leaf_kernel.rs new file mode 100644 index 000000000..4028304b4 --- /dev/null +++ b/diskann-pipnn/tests/leaf_kernel.rs @@ -0,0 +1,454 @@ +/* + * 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::MAX) == 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 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)); +} + +#[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..f3082f3d3 --- /dev/null +++ b/diskann-pipnn/tests/partition_kernel.rs @@ -0,0 +1,419 @@ +/* + * 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::MAX) == 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 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; From 4b9116c07ead50e0e11e9b7c030ef1c30541c015 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:02:37 +0000 Subject: [PATCH 03/26] pipnn: benchmark and exercise numerical kernels --- .github/workflows/ci.yml | 2 + diskann-pipnn/benches/kernels.rs | 220 +++++++++++++++++++++++++++++++ 2 files changed, 222 insertions(+) create mode 100644 diskann-pipnn/benches/kernels.rs 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/diskann-pipnn/benches/kernels.rs b/diskann-pipnn/benches/kernels.rs new file mode 100644 index 000000000..333233922 --- /dev/null +++ b/diskann-pipnn/benches/kernels.rs @@ -0,0 +1,220 @@ +/* + * 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_K: usize = 2; +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 { + 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 { + 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); From 7baade7f801bf077c08dbf5e0b3797698573afea Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:02:46 +0000 Subject: [PATCH 04/26] pipnn: harden kernel coverage and benchmarks --- diskann-pipnn/benches/kernels.rs | 122 ++++++++++---------- diskann-pipnn/src/leaf_kernel.rs | 40 ++++--- diskann-pipnn/src/leaf_kernel/tests.rs | 122 ++++++++++++++++++++ diskann-pipnn/src/partition_kernel.rs | 3 + diskann-pipnn/src/partition_kernel/tests.rs | 93 +++++++++++++++ diskann-pipnn/tests/leaf_kernel.rs | 20 ++++ 6 files changed, 323 insertions(+), 77 deletions(-) create mode 100644 diskann-pipnn/src/leaf_kernel/tests.rs create mode 100644 diskann-pipnn/src/partition_kernel/tests.rs diff --git a/diskann-pipnn/benches/kernels.rs b/diskann-pipnn/benches/kernels.rs index 333233922..22790c554 100644 --- a/diskann-pipnn/benches/kernels.rs +++ b/diskann-pipnn/benches/kernels.rs @@ -15,7 +15,7 @@ use diskann_vector::distance::Metric; const BIGANN_DIMENSIONS: usize = 128; const PARTITION_FANOUT: usize = 10; -const LEAF_K: usize = 2; +const LEAF_KS: [usize; 2] = [2, 3]; const LEAF_SIZES: [usize; 3] = [64, 256, 512]; const METRICS: [Metric; 4] = [ Metric::L2, @@ -133,28 +133,30 @@ 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 { - 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(); + 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.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(); @@ -163,44 +165,46 @@ fn benchmark_leaf_topk(c: &mut Criterion) { fn benchmark_full_leaf(c: &mut Criterion) { let mut group = c.benchmark_group("pipnn/full-leaf-numerical"); for points in LEAF_SIZES { - 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(); + 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.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(); } diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index a913ba138..0953bee21 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -200,15 +200,11 @@ fn resize( len: usize, value: T, ) -> Result<(), LeafKernelError> { - if len > values.len() { - let additional = len - values.len(); - values - .try_reserve(additional) - .map_err(|_| LeafKernelError::Allocation { buffer, additional })?; - values.resize(len, value); - } else { - values.truncate(len); - } + let additional = len.saturating_sub(values.len()); + values + .try_reserve(additional) + .map_err(|_| LeafKernelError::Allocation { buffer, additional })?; + values.resize(len, value); Ok(()) } @@ -253,18 +249,22 @@ impl LeafKernel<'_, '_, '_> { F::Mask: SIMDSelect, u64: From<<::BitMask as SIMDMask>::Underlying>, { - match self.k { - 1 => self.run_fused::(arch), - 2 => self.run_fused::(arch), - 3 => self.run_fused::(arch), - _ => process_pairs_simd_dynamic::( + 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"), } } @@ -689,10 +689,11 @@ where 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 zero_denominator = denominator.eq_simd(zero); - let safe_denominator = zero_denominator.select(one, denominator); - let cosine = zero_denominator.select(zero, dot / safe_denominator); + 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)); let distance = one - cosine; // Comparisons with NaN are false, so this explicit lower clamp // preserves non-rankable NaNs while matching the existing PiPNN @@ -762,3 +763,6 @@ fn insert_row( } 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..661f0f558 --- /dev/null +++ b/diskann-pipnn/src/leaf_kernel/tests.rs @@ -0,0 +1,122 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; +use diskann_wide::arch::{Scalar, Target}; + +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_target_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::MAX; points]; + let norms = norms(input); + as Target>::run( + LeafKernel { + input, + k, + output: &mut actual, + norms: &norms, + worst: &mut worst, + }, + Scalar::new(), + ); + + 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::MAX]; + + 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/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index b39c05aa8..3525ae157 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -436,3 +436,6 @@ fn copy_ids(top: &TopK, output: &mut [u32]) { *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..df51bfc66 --- /dev/null +++ b/diskann-pipnn/src/partition_kernel/tests.rs @@ -0,0 +1,93 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; +use diskann_wide::arch::{Scalar, Target}; + +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_target_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]; + as Target>::run( + PartitionKernel { + input, + fanout, + output: &mut actual, + }, + Scalar::new(), + ); + + 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::MAX); 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 index 4028304b4..2110ddffa 100644 --- a/diskann-pipnn/tests/leaf_kernel.rs +++ b/diskann-pipnn/tests/leaf_kernel.rs @@ -424,6 +424,26 @@ fn rejects_shape_overflow_before_reading_buffers() { 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() { From e6f6ad979d08922ba628267b1236cfa801934ed3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:04:07 +0000 Subject: [PATCH 05/26] fix(pipnn): accept finite maximum distances --- diskann-pipnn/src/leaf_kernel.rs | 6 ++--- diskann-pipnn/src/leaf_kernel/tests.rs | 4 ++-- diskann-pipnn/src/partition_kernel.rs | 4 ++-- diskann-pipnn/src/partition_kernel/tests.rs | 2 +- diskann-pipnn/tests/leaf_kernel.rs | 17 +++++++++++++- diskann-pipnn/tests/partition_kernel.rs | 25 ++++++++++++++++++++- 6 files changed, 48 insertions(+), 10 deletions(-) diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 0953bee21..b1eeb6170 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -40,7 +40,7 @@ impl LeafNeighbor { impl Default for LeafNeighbor { fn default() -> Self { - Self::new(u32::MAX, f32::MAX) + Self::new(u32::MAX, f32::INFINITY) } } @@ -139,7 +139,7 @@ pub fn nearest_leaf_neighbors( "worst distances", &mut workspace.worst, input.points, - f32::MAX, + f32::INFINITY, )?; for (row, norm) in workspace.norms.iter_mut().enumerate() { let squared_norm = input.dots[row * input.points + row]; @@ -157,7 +157,7 @@ pub fn nearest_leaf_neighbors( }; } output.fill(LeafNeighbor::default()); - workspace.worst.fill(f32::MAX); + workspace.worst.fill(f32::INFINITY); diskann_wide::arch::dispatch(LeafKernel { input, diff --git a/diskann-pipnn/src/leaf_kernel/tests.rs b/diskann-pipnn/src/leaf_kernel/tests.rs index 661f0f558..becef5039 100644 --- a/diskann-pipnn/src/leaf_kernel/tests.rs +++ b/diskann-pipnn/src/leaf_kernel/tests.rs @@ -59,7 +59,7 @@ fn scalar_target_matches_runtime_dispatch() { .unwrap(); let mut actual = vec![LeafNeighbor::default(); points * k]; - let mut worst = vec![f32::MAX; points]; + let mut worst = vec![f32::INFINITY; points]; let norms = norms(input); as Target>::run( LeafKernel { @@ -81,7 +81,7 @@ fn scalar_target_matches_runtime_dispatch() { #[test] fn scalar_insertion_orders_candidates_and_rejects_nan() { let mut output = [LeafNeighbor::default(); 4]; - let mut worst = [f32::MAX]; + 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); diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index 3525ae157..4865005c4 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -231,7 +231,7 @@ fn process_rows_scalar(input: PartitionTopK<'_>, fanout: usize, output: &mut [u3 .zip(output.chunks_exact_mut(fanout)) .enumerate() { - let mut top = [(u32::MAX, f32::MAX); MAX_PARTITION_FANOUT]; + 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); @@ -259,7 +259,7 @@ where .zip(output.chunks_exact_mut(fanout)) .enumerate() { - let mut top = [(u32::MAX, f32::MAX); MAX_PARTITION_FANOUT]; + let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; match input.metric { Metric::L2 => process_binary::( arch, diff --git a/diskann-pipnn/src/partition_kernel/tests.rs b/diskann-pipnn/src/partition_kernel/tests.rs index df51bfc66..feebc5348 100644 --- a/diskann-pipnn/src/partition_kernel/tests.rs +++ b/diskann-pipnn/src/partition_kernel/tests.rs @@ -83,7 +83,7 @@ fn scalar_distance_matches_metric_contract() { #[test] fn scalar_topk_orders_candidates_and_preserves_ties() { - let mut top = [(u32::MAX, f32::MAX); MAX_PARTITION_FANOUT]; + 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); } diff --git a/diskann-pipnn/tests/leaf_kernel.rs b/diskann-pipnn/tests/leaf_kernel.rs index 2110ddffa..5cead7e96 100644 --- a/diskann-pipnn/tests/leaf_kernel.rs +++ b/diskann-pipnn/tests/leaf_kernel.rs @@ -88,7 +88,7 @@ fn reference(input: LeafTopK<'_>, requested_k: usize) -> Vec { clamp(1.0 - similarity) } }; - if distance.partial_cmp(&f32::MAX) == Some(Ordering::Less) { + if distance.partial_cmp(&f32::INFINITY) == Some(Ordering::Less) { candidates.push(LeafNeighbor::new(position as u32, distance)); } } @@ -256,6 +256,21 @@ fn preserves_pipnn_metric_edge_semantics() { ); } +#[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] diff --git a/diskann-pipnn/tests/partition_kernel.rs b/diskann-pipnn/tests/partition_kernel.rs index f3082f3d3..bb90a8b9f 100644 --- a/diskann-pipnn/tests/partition_kernel.rs +++ b/diskann-pipnn/tests/partition_kernel.rs @@ -35,7 +35,7 @@ fn reference(input: PartitionTopK<'_>, fanout: usize) -> Vec { } } }; - (distance.partial_cmp(&f32::MAX) == Some(std::cmp::Ordering::Less)) + (distance.partial_cmp(&f32::INFINITY) == Some(std::cmp::Ordering::Less)) .then_some((leader as u32, distance)) }) .collect(); @@ -213,6 +213,29 @@ fn cosine_treats_a_zero_norm_as_zero_similarity() { 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]; From aa6347f2c8443c456541fdde34a08a757498c7e2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 07:48:17 +0000 Subject: [PATCH 06/26] refactor(pipnn): make kernels architecture-neutral --- diskann-pipnn/src/leaf_kernel.rs | 60 +++++---------------- diskann-pipnn/src/leaf_kernel/tests.rs | 14 +---- diskann-pipnn/src/partition_kernel.rs | 54 +++++-------------- diskann-pipnn/src/partition_kernel/tests.rs | 12 +---- 4 files changed, 28 insertions(+), 112 deletions(-) diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index b1eeb6170..34e78a02b 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -6,20 +6,14 @@ //! Fused nearest-neighbor kernel for a leaf's lower dot-product matrix. use diskann_vector::distance::Metric; -#[cfg(target_arch = "x86_64")] -use diskann_wide::{SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; +use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; /// Widest f32 SIMD lane count DiskANN dispatches to, used to size lane scratch. -#[cfg(target_arch = "x86_64")] const MAX_LANES: usize = 16; -#[cfg(target_arch = "x86_64")] const L2: u8 = 0; -#[cfg(target_arch = "x86_64")] const COSINE_NORMALIZED: u8 = 1; -#[cfg(target_arch = "x86_64")] const INNER_PRODUCT: u8 = 2; -#[cfg(target_arch = "x86_64")] const COSINE: u8 = 3; /// One leaf-local neighbor and its metric distance. @@ -238,11 +232,6 @@ struct LeafKernel<'a, 'o, 'w> { } impl LeafKernel<'_, '_, '_> { - fn run_scalar(self) { - process_pairs_scalar(self.input, self.k, self.output, self.norms, self.worst); - } - - #[cfg(target_arch = "x86_64")] fn run_simd(self, arch: F::Arch) where F: SIMDVector + SIMDFloat + std::ops::Div, @@ -268,7 +257,6 @@ impl LeafKernel<'_, '_, '_> { } } - #[cfg(target_arch = "x86_64")] fn run_fused(self, arch: F::Arch) where F: SIMDVector + SIMDFloat + std::ops::Div, @@ -308,40 +296,20 @@ impl LeafKernel<'_, '_, '_> { } } -impl diskann_wide::arch::Target for LeafKernel<'_, '_, '_> { - #[inline(always)] - fn run(self, _: diskann_wide::arch::Scalar) { - self.run_scalar(); - } -} - -#[cfg(target_arch = "x86_64")] -impl diskann_wide::arch::Target for LeafKernel<'_, '_, '_> { - #[inline(always)] - fn run(self, arch: diskann_wide::arch::x86_64::V3) { - diskann_wide::alias!(F32x8 = ::f32x8); - self.run_simd::(arch); - } -} - -#[cfg(target_arch = "x86_64")] -impl diskann_wide::arch::Target for LeafKernel<'_, '_, '_> { - #[inline(always)] - fn run(self, arch: diskann_wide::arch::x86_64::V4) { - diskann_wide::alias!(F32x16 = ::f32x16); - self.run_simd::(arch); - } -} - -#[cfg(target_arch = "aarch64")] -impl diskann_wide::arch::Target for LeafKernel<'_, '_, '_> { +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: diskann_wide::arch::aarch64::Neon) { - let _scalar = arch.retarget(); - self.run_scalar(); + fn run(self, arch: A) { + self.run_simd::(arch); } } +#[cfg(test)] fn process_pairs_scalar( input: LeafTopK<'_>, k: usize, @@ -359,7 +327,6 @@ fn process_pairs_scalar( } } -#[cfg(target_arch = "x86_64")] /// Fused dual-endpoint scan for row widths without a specialized arm. /// /// Identical structure to [`process_pairs_simd_fused`], with the slot count @@ -462,7 +429,6 @@ fn process_pairs_simd_dynamic( /// 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. -#[cfg(target_arch = "x86_64")] #[inline(never)] fn process_pairs_simd_fused( arch: F::Arch, @@ -567,7 +533,6 @@ fn process_pairs_simd_fused( } } -#[cfg(target_arch = "x86_64")] const fn metric() -> Metric { match METRIC { L2 => Metric::L2, @@ -589,7 +554,6 @@ const fn metric() -> Metric { /// # Safety /// /// `base + slots` must be within the allocation behind `output`. -#[cfg(target_arch = "x86_64")] #[inline(always)] unsafe fn insert_slots( output: *mut LeafNeighbor, @@ -669,7 +633,6 @@ unsafe fn insert_slots( } } -#[cfg(target_arch = "x86_64")] #[inline(always)] fn pair_distances(arch: F::Arch, metric: Metric, dot: F, row_norm: F, column_norm: F) -> F where @@ -741,6 +704,7 @@ fn pair_distance(metric: Metric, dot: f32, row_norm: f32, column_norm: f32) -> f } #[inline(always)] +#[cfg(test)] fn insert_row( output: &mut [LeafNeighbor], worst: &mut [f32], diff --git a/diskann-pipnn/src/leaf_kernel/tests.rs b/diskann-pipnn/src/leaf_kernel/tests.rs index becef5039..90381bc1d 100644 --- a/diskann-pipnn/src/leaf_kernel/tests.rs +++ b/diskann-pipnn/src/leaf_kernel/tests.rs @@ -4,7 +4,6 @@ */ use super::*; -use diskann_wide::arch::{Scalar, Target}; fn dots(metric: Metric, points: usize) -> Vec { let mut dots = vec![f32::NAN; points * points]; @@ -39,7 +38,7 @@ fn norms(input: LeafTopK<'_>) -> Vec { } #[test] -fn scalar_target_matches_runtime_dispatch() { +fn scalar_reference_matches_runtime_dispatch() { for metric in [ Metric::L2, Metric::Cosine, @@ -61,16 +60,7 @@ fn scalar_target_matches_runtime_dispatch() { let mut actual = vec![LeafNeighbor::default(); points * k]; let mut worst = vec![f32::INFINITY; points]; let norms = norms(input); - as Target>::run( - LeafKernel { - input, - k, - output: &mut actual, - norms: &norms, - worst: &mut worst, - }, - Scalar::new(), - ); + process_pairs_scalar(input, k, &mut actual, &norms, &mut worst); assert_eq!(actual, expected, "{metric:?}, n={points}, k={k}"); } diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index 4865005c4..b2942fa65 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -10,8 +10,7 @@ //! positions; partition recursion and cluster ownership stay with the caller. use diskann_vector::distance::Metric; -#[cfg(target_arch = "x86_64")] -use diskann_wide::{SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector}; +use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector}; /// Maximum number of leaders retained for one point. pub const MAX_PARTITION_FANOUT: usize = 16; @@ -175,11 +174,6 @@ struct PartitionKernel<'a, 'o> { } impl PartitionKernel<'_, '_> { - fn run_scalar(self) { - process_rows_scalar(self.input, self.fanout, self.output); - } - - #[cfg(target_arch = "x86_64")] fn run_simd(self, arch: F::Arch) where F: SIMDVector + SIMDFloat + std::ops::Div, @@ -190,40 +184,20 @@ impl PartitionKernel<'_, '_> { } } -impl diskann_wide::arch::Target for PartitionKernel<'_, '_> { - #[inline(always)] - fn run(self, _: diskann_wide::arch::Scalar) { - self.run_scalar(); - } -} - -#[cfg(target_arch = "x86_64")] -impl diskann_wide::arch::Target for PartitionKernel<'_, '_> { - #[inline(always)] - fn run(self, arch: diskann_wide::arch::x86_64::V3) { - diskann_wide::alias!(F32x8 = ::f32x8); - self.run_simd::(arch); - } -} - -#[cfg(target_arch = "x86_64")] -impl diskann_wide::arch::Target for PartitionKernel<'_, '_> { - #[inline(always)] - fn run(self, arch: diskann_wide::arch::x86_64::V4) { - diskann_wide::alias!(F32x16 = ::f32x16); - self.run_simd::(arch); - } -} - -#[cfg(target_arch = "aarch64")] -impl diskann_wide::arch::Target for PartitionKernel<'_, '_> { +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: diskann_wide::arch::aarch64::Neon) { - let _scalar = arch.retarget(); - self.run_scalar(); + 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 @@ -246,7 +220,6 @@ fn process_rows_scalar(input: PartitionTopK<'_>, fanout: usize, output: &mut [u3 } } -#[cfg(target_arch = "x86_64")] fn process_rows_simd(arch: F::Arch, input: PartitionTopK<'_>, fanout: usize, output: &mut [u32]) where F: SIMDVector + SIMDFloat + std::ops::Div, @@ -290,7 +263,6 @@ where } } -#[cfg(target_arch = "x86_64")] fn process_cosine( arch: F::Arch, dots: &[f32], @@ -315,7 +287,6 @@ fn process_cosine( }); } -#[cfg(target_arch = "x86_64")] fn process_unary( arch: F::Arch, dots: &[f32], @@ -342,7 +313,6 @@ fn process_unary( } } -#[cfg(target_arch = "x86_64")] fn process_binary( arch: F::Arch, dots: &[f32], @@ -375,7 +345,6 @@ fn process_binary( } } -#[cfg(target_arch = "x86_64")] fn insert_lanes(distances: F, base: usize, top: &mut TopK, fanout: usize) where F: SIMDVector + SIMDPartialOrd, @@ -399,6 +368,7 @@ where } #[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), diff --git a/diskann-pipnn/src/partition_kernel/tests.rs b/diskann-pipnn/src/partition_kernel/tests.rs index feebc5348..aaa03917b 100644 --- a/diskann-pipnn/src/partition_kernel/tests.rs +++ b/diskann-pipnn/src/partition_kernel/tests.rs @@ -4,7 +4,6 @@ */ use super::*; -use diskann_wide::arch::{Scalar, Target}; fn input(metric: Metric, leaders: usize) -> (Vec, Vec, Vec) { let dots = (0..2 * leaders) @@ -32,7 +31,7 @@ fn input(metric: Metric, leaders: usize) -> (Vec, Vec, Vec) { } #[test] -fn scalar_target_matches_runtime_dispatch() { +fn scalar_reference_matches_runtime_dispatch() { for metric in [ Metric::L2, Metric::Cosine, @@ -54,14 +53,7 @@ fn scalar_target_matches_runtime_dispatch() { nearest_leaders(input, fanout, &mut expected).unwrap(); let mut actual = vec![u32::MAX; input.rows * fanout]; - as Target>::run( - PartitionKernel { - input, - fanout, - output: &mut actual, - }, - Scalar::new(), - ); + process_rows_scalar(input, fanout, &mut actual); assert_eq!( actual, expected, From 152bcf934af3552fbaf3ce79cbdec3acb67e07c1 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:53:48 +0000 Subject: [PATCH 07/26] fix(pipnn): preserve NaN distances across SIMD backends --- diskann-pipnn/src/leaf_kernel.rs | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 34e78a02b..f8531d429 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -640,14 +640,21 @@ where 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; - zero.max_simd(distance) + clamp_nonnegative(distance) } Metric::CosineNormalized => { let distance = F::splat(arch, 1.0) - dot; - zero.max_simd(distance) + clamp_nonnegative(distance) } Metric::InnerProduct => zero - dot, Metric::Cosine => { @@ -657,11 +664,7 @@ where 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)); - let distance = one - cosine; - // Comparisons with NaN are false, so this explicit lower clamp - // preserves non-rankable NaNs while matching the existing PiPNN - // distance formulas for finite values. - zero.max_simd(distance) + clamp_nonnegative(one - cosine) } } } From 4354948110664893483a15cb3c20079d36c7c48d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 02:48:33 +0000 Subject: [PATCH 08/26] docs(pipnn): explain numerical kernel invariants --- diskann-linalg/src/faer.rs | 7 ++++++- diskann-pipnn/src/leaf_kernel.rs | 15 ++++++++++++++- diskann-pipnn/src/leaf_kernel/tests.rs | 20 +++++++++++++++----- diskann-pipnn/src/partition_kernel.rs | 14 +++++++++++--- diskann-pipnn/src/partition_kernel/tests.rs | 9 +++++++-- 5 files changed, 53 insertions(+), 12 deletions(-) diff --git a/diskann-linalg/src/faer.rs b/diskann-linalg/src/faer.rs index 1e7feeb9d..55ca5b2df 100644 --- a/diskann-linalg/src/faer.rs +++ b/diskann-linalg/src/faer.rs @@ -55,7 +55,12 @@ pub(super) fn sgemm_impl( /// Implements the public lower-triangular AAT operation. /// -/// The caller has already validated the matrix dimensions. +/// Leaf selection consumes each symmetric pair once and updates both endpoints, +/// so computing or initializing the upper triangle would be wasted bandwidth. +/// Faer's triangular block structure is the contract that prevents those stores; +/// callers may keep unrelated values in the upper triangle. The public wrapper +/// has already checked `a.len() == m * k`, `c.len() == m * m`, and overflow, so +/// the unchecked matrix views below cannot escape their backing slices. pub(super) fn sgemm_aat_lower_impl(m: usize, k: usize, a: &[f32], c: &mut [f32]) { use faer::linalg::matmul::triangular::{matmul, BlockStructure}; diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index f8531d429..7f75f228b 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -3,7 +3,20 @@ * Licensed under the MIT license. */ -//! Fused nearest-neighbor kernel for a leaf's lower dot-product matrix. +//! Fused nearest-neighbor selection over a leaf's lower dot-product matrix. +//! +//! `sgemm_aat_lower` writes only pair `(row, column)` with `column <= row`. +//! This kernel therefore walks the strict lower triangle once and offers each +//! computed distance to both endpoint rows. Keeping one top-k tracker per row +//! avoids materializing the upper triangle or computing a symmetric distance +//! twice. +//! +//! The public entry point validates every shape before dispatch. The dispatched +//! path processes complete SIMD chunks, then a scalar tail. For `k <= 3`, const +//! slot counts remove the dynamic insertion loop from the hot path; larger `k` +//! uses the same ordering rules through the dynamic fallback. NaN distances are +//! never rankable, and ties retain scan order so scalar and SIMD backends produce +//! the same graph. use diskann_vector::distance::Metric; use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; diff --git a/diskann-pipnn/src/leaf_kernel/tests.rs b/diskann-pipnn/src/leaf_kernel/tests.rs index 90381bc1d..67878424f 100644 --- a/diskann-pipnn/src/leaf_kernel/tests.rs +++ b/diskann-pipnn/src/leaf_kernel/tests.rs @@ -45,24 +45,34 @@ fn scalar_reference_matches_runtime_dispatch() { Metric::CosineNormalized, Metric::InnerProduct, ] { - for points in [7, 17] { + // Point count, rather than source-vector dimension, controls this + // kernel's SIMD boundaries. These values cover short rows plus the + // lane-1/lane/lane+1 boundaries for 4-, 8-, and 16-lane backends, + // then the boundary around a second 16-lane chunk. + for points in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let dots = dots(metric, points); let input = LeafTopK { dots: &dots, points, metric, }; - for k in [1, 2, 3, 4] { + for requested_k in [1, 2, 3, 4] { + let k = requested_k.min(points - 1); let mut expected = vec![LeafNeighbor::default(); points * k]; - nearest_leaf_neighbors(input, k, &mut expected, &mut LeafTopKWorkspace::new()) - .unwrap(); + nearest_leaf_neighbors( + input, + requested_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}"); + assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); } } } diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index b2942fa65..bb193eb9b 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -5,9 +5,17 @@ //! 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. +//! The caller gathers a point stripe and a leader matrix, then computes the +//! row-major `points · leadersᵀ` tile with GEMM. This module performs the second +//! half of assignment: convert each dot product to the configured metric and +//! retain only the nearest leader positions. +//! +//! L2 deliberately omits the point norm because it adds the same constant to +//! every leader in one row and cannot change their order. Cosine still needs a +//! point scale because it divides each dot product. The fixed 16-entry tracker +//! bounds stack use and matches the configuration fanout limit. SIMD chunks and +//! scalar tails feed the same insertion routine; NaNs are ignored and equal +//! distances keep the first leader encountered. use diskann_vector::distance::Metric; use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector}; diff --git a/diskann-pipnn/src/partition_kernel/tests.rs b/diskann-pipnn/src/partition_kernel/tests.rs index aaa03917b..e7b33ead1 100644 --- a/diskann-pipnn/src/partition_kernel/tests.rs +++ b/diskann-pipnn/src/partition_kernel/tests.rs @@ -38,7 +38,9 @@ fn scalar_reference_matches_runtime_dispatch() { Metric::CosineNormalized, Metric::InnerProduct, ] { - for leaders in [7, 17] { + // Leader count controls SIMD chunking. Exercise the tail on both sides + // of 4-, 8-, and 16-lane boundaries, then a second 16-lane chunk. + for leaders in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let (dots, row_scales, leader_scales) = input(metric, leaders); let input = PartitionTopK { dots: &dots, @@ -48,7 +50,10 @@ fn scalar_reference_matches_runtime_dispatch() { leader_scales: &leader_scales, metric, }; - for fanout in [1, 2, 6] { + for fanout in [1, 2, 6, MAX_PARTITION_FANOUT] { + if fanout > leaders { + continue; + } let mut expected = vec![u32::MAX; input.rows * fanout]; nearest_leaders(input, fanout, &mut expected).unwrap(); From 20ab8a0556fa9e687c84eb21df0fb6864b4507fb Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:39:16 +0000 Subject: [PATCH 09/26] test(pipnn): name SIMD metric boundary matrices --- diskann-pipnn/src/leaf_kernel/tests.rs | 81 ++++++++++++--------- diskann-pipnn/src/partition_kernel/tests.rs | 78 +++++++++++--------- 2 files changed, 91 insertions(+), 68 deletions(-) diff --git a/diskann-pipnn/src/leaf_kernel/tests.rs b/diskann-pipnn/src/leaf_kernel/tests.rs index 67878424f..1dba1805c 100644 --- a/diskann-pipnn/src/leaf_kernel/tests.rs +++ b/diskann-pipnn/src/leaf_kernel/tests.rs @@ -37,47 +37,58 @@ fn norms(input: LeafTopK<'_>) -> Vec { .collect() } -#[test] -fn scalar_reference_matches_runtime_dispatch() { - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - // Point count, rather than source-vector dimension, controls this - // kernel's SIMD boundaries. These values cover short rows plus the - // lane-1/lane/lane+1 boundaries for 4-, 8-, and 16-lane backends, - // then the boundary around a second 16-lane chunk. - for points in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let dots = dots(metric, points); - let input = LeafTopK { - dots: &dots, - points, - metric, - }; - for requested_k in [1, 2, 3, 4] { - let k = requested_k.min(points - 1); - let mut expected = vec![LeafNeighbor::default(); points * k]; - nearest_leaf_neighbors( - input, - requested_k, - &mut expected, - &mut LeafTopKWorkspace::new(), - ) - .unwrap(); +fn assert_scalar_reference_matches_runtime_dispatch(metric: Metric) { + // Point count, rather than source-vector dimension, controls this kernel's + // SIMD boundaries. Cover lane-1/lane/lane+1 for 4-, 8-, and 16-lane + // backends, then the boundary around a second 16-lane chunk. + for points in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let dots = dots(metric, points); + let input = LeafTopK { + dots: &dots, + points, + metric, + }; + for requested_k in [1, 2, 3, 4] { + let k = requested_k.min(points - 1); + let mut expected = vec![LeafNeighbor::default(); points * k]; + nearest_leaf_neighbors( + input, + requested_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); + 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={requested_k}"); - } + assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); } } } +#[test] +fn l2_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_runtime_dispatch(Metric::L2); +} + +#[test] +fn cosine_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_runtime_dispatch(Metric::Cosine); +} + +#[test] +fn normalized_cosine_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_runtime_dispatch(Metric::CosineNormalized); +} + +#[test] +fn inner_product_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_runtime_dispatch(Metric::InnerProduct); +} + #[test] fn scalar_insertion_orders_candidates_and_rejects_nan() { let mut output = [LeafNeighbor::default(); 4]; diff --git a/diskann-pipnn/src/partition_kernel/tests.rs b/diskann-pipnn/src/partition_kernel/tests.rs index e7b33ead1..ad3e048a0 100644 --- a/diskann-pipnn/src/partition_kernel/tests.rs +++ b/diskann-pipnn/src/partition_kernel/tests.rs @@ -30,45 +30,57 @@ fn input(metric: Metric, leaders: usize) -> (Vec, Vec, Vec) { (dots, row_scales, leader_scales) } -#[test] -fn scalar_reference_matches_runtime_dispatch() { - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - // Leader count controls SIMD chunking. Exercise the tail on both sides - // of 4-, 8-, and 16-lane boundaries, then a second 16-lane chunk. - for leaders in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - 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, MAX_PARTITION_FANOUT] { - if fanout > leaders { - continue; - } - let mut expected = vec![u32::MAX; input.rows * fanout]; - nearest_leaders(input, fanout, &mut expected).unwrap(); +fn assert_scalar_reference_matches_runtime_dispatch(metric: Metric) { + // Leader count controls SIMD chunking. Exercise the tail on both sides of + // 4-, 8-, and 16-lane boundaries, then a second 16-lane chunk. + for leaders in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + 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, MAX_PARTITION_FANOUT] { + if fanout > leaders { + continue; + } + 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); + 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}" - ); - } + assert_eq!( + actual, expected, + "{metric:?}, leaders={leaders}, k={fanout}" + ); } } } +#[test] +fn l2_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_runtime_dispatch(Metric::L2); +} + +#[test] +fn cosine_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_runtime_dispatch(Metric::Cosine); +} + +#[test] +fn normalized_cosine_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_runtime_dispatch(Metric::CosineNormalized); +} + +#[test] +fn inner_product_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_runtime_dispatch(Metric::InnerProduct); +} + #[test] fn scalar_distance_matches_metric_contract() { assert_eq!(distance(Metric::L2, 2.0, 99.0, 9.0), 5.0); From 943ffc38f3c90e5e2df67245953d0c176eea1672 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 02:08:31 +0000 Subject: [PATCH 10/26] fix(pipnn): preserve dispatched kernel semantics --- diskann-pipnn/src/leaf_kernel.rs | 13 ++---- diskann-pipnn/src/partition_kernel.rs | 59 +++++++++++++++------------ diskann-wide/src/traits.rs | 10 ++--- 3 files changed, 43 insertions(+), 39 deletions(-) diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 7f75f228b..2aee917f3 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -21,9 +21,6 @@ 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; @@ -378,9 +375,8 @@ fn process_pairs_simd_dynamic( 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 values = distances.to_array(); + let values = values.as_ref(); let mut row_bits = row_bits; while row_bits != 0 { let lane = row_bits.trailing_zeros() as usize; @@ -479,9 +475,8 @@ fn process_pairs_simd_fused( 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 values = distances.to_array(); + let values = values.as_ref(); let mut row_bits = row_bits; while row_bits != 0 { let lane = row_bits.trailing_zeros() as usize; diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index bb193eb9b..ab4d02255 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -242,13 +242,14 @@ where { let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; match input.metric { - Metric::L2 => process_binary::( + 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), + |dot, norm| norm - 2.0 * dot, ), Metric::CosineNormalized => { process_unary::(arch, dot_row, &mut top, fanout, |dot| { @@ -286,13 +287,29 @@ fn process_cosine( 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 - }); + 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 + }, + |dot, leader_norm| { + let denominator = row_norm_squared.sqrt() * leader_norm; + let cosine = if denominator > 0.0 { + dot / denominator + } else { + 0.0 + }; + 1.0 - cosine + }, + ); } fn process_unary( @@ -313,24 +330,23 @@ fn process_unary( 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]); + let value = transform(F::splat(arch, dot)).to_array(); + insert_topk(top, fanout, (full + offset) as u32, value.as_ref()[0]); } } -fn process_binary( +fn process_binary( arch: F::Arch, dots: &[f32], scales: &[f32], top: &mut TopK, fanout: usize, transform: Transform, + scalar_transform: ScalarTransform, ) where F: SIMDVector + SIMDFloat, Transform: Fn(F, F) -> F, + ScalarTransform: Fn(f32, f32) -> f32, u64: From<<::BitMask as SIMDMask>::Underlying>, { let full = dots.len() / F::LANES * F::LANES; @@ -342,14 +358,8 @@ fn process_binary( 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]); + let value = scalar_transform(dots[full + offset], scales[full + offset]); + insert_topk(top, fanout, (full + offset) as u32, value); } } @@ -364,9 +374,8 @@ where 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 values = distances.to_array(); + let values = values.as_ref(); let mut lanes = u64::from(eligible.bitmask().to_underlying()); while lanes != 0 { let lane = lanes.trailing_zeros() as usize; diff --git a/diskann-wide/src/traits.rs b/diskann-wide/src/traits.rs index d581406c9..5eb002f9a 100644 --- a/diskann-wide/src/traits.rs +++ b/diskann-wide/src/traits.rs @@ -28,7 +28,7 @@ use super::{ /// - /// - pub trait ArrayType: SupportedLaneCount { - type Type; + type Type: AsRef<[T]> + AsMut<[T]>; } /// Map scalar + lengths to arrays. @@ -262,7 +262,7 @@ pub trait SIMDVector: Copy + std::fmt::Debug { /// The argument `arch` provides a "proof of compatibility" as `A` can only be safely /// instantiated when all the requirements for the architecture are met. fn from_array(arch: Self::Arch, x: >::Type) - -> Self; + -> Self; /// Broadcast the provided scalar across all lanes. /// @@ -897,16 +897,16 @@ impl_simd_mask_for_bitmask!(64, u64, u64::MAX); #[cfg(test)] mod test_traits { use rand::{ - SeedableRng, distr::{Distribution, StandardUniform}, rngs::StdRng, + SeedableRng, }; use super::*; use crate::{ - ARCH, arch, + arch, splitjoin::{LoHi, SplitJoin}, - test_utils, + test_utils, ARCH, }; // Allow unsigned 128-bit integers to be converted to narrow types. From 4d8cfb3358efe97621c829d719ce8e92b1916579 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 02:15:54 +0000 Subject: [PATCH 11/26] style(wide): format array trait bounds --- diskann-wide/src/traits.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/diskann-wide/src/traits.rs b/diskann-wide/src/traits.rs index 5eb002f9a..caea7ce6c 100644 --- a/diskann-wide/src/traits.rs +++ b/diskann-wide/src/traits.rs @@ -262,7 +262,7 @@ pub trait SIMDVector: Copy + std::fmt::Debug { /// The argument `arch` provides a "proof of compatibility" as `A` can only be safely /// instantiated when all the requirements for the architecture are met. fn from_array(arch: Self::Arch, x: >::Type) - -> Self; + -> Self; /// Broadcast the provided scalar across all lanes. /// @@ -897,16 +897,16 @@ impl_simd_mask_for_bitmask!(64, u64, u64::MAX); #[cfg(test)] mod test_traits { use rand::{ + SeedableRng, distr::{Distribution, StandardUniform}, rngs::StdRng, - SeedableRng, }; use super::*; use crate::{ - arch, + ARCH, arch, splitjoin::{LoHi, SplitJoin}, - test_utils, ARCH, + test_utils, }; // Allow unsigned 128-bit integers to be converted to narrow types. From 437617d6f2919773259b799c5ec8b2c40ec03498 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 07:30:53 +0000 Subject: [PATCH 12/26] fix(pipnn): address kernel review feedback --- .github/workflows/nightly.yml | 19 +- diskann-linalg/src/lib.rs | 103 +++----- diskann-linalg/tests/sgemm_aat_lower.rs | 14 +- diskann-pipnn/benches/kernels.rs | 8 +- diskann-pipnn/src/leaf_kernel.rs | 276 ++++++++++---------- diskann-pipnn/src/leaf_kernel/tests.rs | 11 + diskann-pipnn/src/lib.rs | 15 +- diskann-pipnn/src/partition_kernel.rs | 144 ++++++---- diskann-pipnn/src/partition_kernel/tests.rs | 48 +++- diskann-pipnn/tests/leaf_kernel.rs | 27 +- 10 files changed, 377 insertions(+), 288 deletions(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index fab755b6a..ccf910024 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -17,7 +17,10 @@ env: RUST_CONFIG: 'build.rustflags=["-Dwarnings"]' RUST_BACKTRACE: 1 CARGO_TERM_COLOR: always - DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test" + DISKANN_FEATURES: >- + virtual_storage,spherical-quantization,product-quantization,tracing, + experimental_diversity_search,disk-index,flatbuffers,linalg,codegen, + multi-vector,bftree,inmem2,integration-test defaults: run: @@ -36,7 +39,9 @@ jobs: run: rustup show && rustup component add clippy - uses: Swatinem/rust-cache@v2 - name: "clippy --workspace --all-targets" - run: cargo clippy --locked --workspace --all-targets --no-deps --config "$RUST_CONFIG" -- -Dwarnings + run: | + cargo clippy --locked --workspace --all-targets --no-deps \ + --config "$RUST_CONFIG" -- -Dwarnings clippy-features: name: clippy-features (macos) @@ -52,7 +57,7 @@ jobs: cargo clippy --locked --workspace \ --all-targets \ --no-deps \ - --features ${{ env.DISKANN_FEATURES }} \ + --features "${{ env.DISKANN_FEATURES }}" \ --config "$RUST_CONFIG" \ -- -Dwarnings @@ -109,7 +114,7 @@ jobs: set -euxo pipefail cargo nextest run --locked --workspace \ --config "$RUST_CONFIG" \ - --features ${{ env.DISKANN_FEATURES }} + --features "${{ env.DISKANN_FEATURES }}" cargo test --locked --doc --workspace --config "$RUST_CONFIG" @@ -133,6 +138,10 @@ jobs: - uses: Swatinem/rust-cache@v2 - name: miri - run: cargo +nightly miri nextest run --locked --package diskann-quantization + run: | + cargo +nightly miri nextest run --locked \ + --package diskann-quantization + cargo +nightly miri test --locked \ + --package diskann-pipnn --lib env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index a6b8ec798..638c075c5 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -80,6 +80,30 @@ impl fmt::Display for SgemmError { impl std::error::Error for SgemmError {} +fn check_matrix( + matrix_name: MatrixName, + actual_len: usize, + rows: usize, + cols: usize, +) -> Result<(), SgemmError> { + let expected_len = rows + .checked_mul(cols) + .ok_or(SgemmError::DimensionOverflow { + matrix_name, + rows, + cols, + })?; + if actual_len != expected_len { + return Err(SgemmError::InvalidMatrixDimensions { + matrix_name, + expected_rows: rows, + expected_cols: cols, + actual_len, + }); + } + Ok(()) +} + // Make the reference implementation available for internal testing. #[cfg(test)] mod reference; @@ -154,51 +178,9 @@ pub fn sgemm( beta: Option, c: &mut [f32], ) -> Result<(), SgemmError> { - // Check size requirements with overflow protection. - 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_b_len = k.checked_mul(n).ok_or(SgemmError::DimensionOverflow { - matrix_name: MatrixName::B, - rows: k, - cols: n, - })?; - - if b.len() != expected_b_len { - return Err(SgemmError::InvalidMatrixDimensions { - matrix_name: MatrixName::B, - expected_rows: k, - expected_cols: n, - actual_len: b.len(), - }); - } - - let expected_c_len = m.checked_mul(n).ok_or(SgemmError::DimensionOverflow { - matrix_name: MatrixName::C, - rows: m, - cols: n, - })?; - - if c.len() != expected_c_len { - return Err(SgemmError::InvalidMatrixDimensions { - matrix_name: MatrixName::C, - expected_rows: m, - expected_cols: n, - actual_len: c.len(), - }); - } + check_matrix(MatrixName::A, a.len(), m, k)?; + check_matrix(MatrixName::B, b.len(), k, n)?; + check_matrix(MatrixName::C, c.len(), m, n)?; // Invoke the actual implementation. sgemm_impl(atranspose, btranspose, m, n, k, alpha, a, b, beta, c); @@ -215,34 +197,9 @@ pub fn sgemm( /// /// 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(), - }); - } +pub fn sgemm_aat_lower(m: usize, k: usize, a: &[f32], c: &mut [f32]) -> Result<(), SgemmError> { + check_matrix(MatrixName::A, a.len(), m, k)?; + check_matrix(MatrixName::C, c.len(), m, m)?; faer::sgemm_aat_lower_impl(m, k, a, c); Ok(()) diff --git a/diskann-linalg/tests/sgemm_aat_lower.rs b/diskann-linalg/tests/sgemm_aat_lower.rs index e84d600d3..d19c92c76 100644 --- a/diskann-linalg/tests/sgemm_aat_lower.rs +++ b/diskann-linalg/tests/sgemm_aat_lower.rs @@ -16,7 +16,7 @@ fn computes_lower_triangle_and_preserves_upper_triangle() { let untouched = -123.0; let mut c = [untouched; 9]; - sgemm_aat_lower(&a, 3, 2, &mut c).unwrap(); + sgemm_aat_lower(3, 2, &a, &mut c).unwrap(); #[rustfmt::skip] assert_eq!(c, [ @@ -28,7 +28,7 @@ fn computes_lower_triangle_and_preserves_upper_triangle() { #[test] fn accepts_a_matrix_with_no_rows() { - sgemm_aat_lower(&[], 0, 3, &mut []).unwrap(); + sgemm_aat_lower(0, 3, &[], &mut []).unwrap(); } #[test] @@ -36,7 +36,7 @@ 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(); + sgemm_aat_lower(3, 0, &[], &mut c).unwrap(); #[rustfmt::skip] assert_eq!(c, [ @@ -50,7 +50,7 @@ fn zero_inner_dimension_zeros_only_the_lower_triangle() { 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(); + let error = sgemm_aat_lower(2, 2, &[0.0; 3], &mut c).unwrap_err(); assert_eq!( error, @@ -67,7 +67,7 @@ fn rejects_invalid_input_dimensions() { 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(); + let error = sgemm_aat_lower(2, 2, &[0.0; 4], &mut c).unwrap_err(); assert_eq!( error, @@ -82,7 +82,7 @@ fn rejects_invalid_output_dimensions() { #[test] fn rejects_input_size_overflow() { - let error = sgemm_aat_lower(&[], usize::MAX, 2, &mut []).unwrap_err(); + let error = sgemm_aat_lower(usize::MAX, 2, &[], &mut []).unwrap_err(); assert_eq!( error, @@ -96,7 +96,7 @@ fn rejects_input_size_overflow() { #[test] fn rejects_output_size_overflow() { - let error = sgemm_aat_lower(&[], usize::MAX, 0, &mut []).unwrap_err(); + let error = sgemm_aat_lower(usize::MAX, 0, &[], &mut []).unwrap_err(); assert_eq!( error, diff --git a/diskann-pipnn/benches/kernels.rs b/diskann-pipnn/benches/kernels.rs index 22790c554..52661f519 100644 --- a/diskann-pipnn/benches/kernels.rs +++ b/diskann-pipnn/benches/kernels.rs @@ -54,7 +54,7 @@ fn lower_dots(points: usize, metric: Metric) -> Vec { normalize_rows(&mut data, BIGANN_DIMENSIONS); } let mut dots = vec![0.0; points * points]; - sgemm_aat_lower(&data, points, BIGANN_DIMENSIONS, &mut dots).unwrap(); + sgemm_aat_lower(points, BIGANN_DIMENSIONS, &data, &mut dots).unwrap(); dots } @@ -120,7 +120,7 @@ fn benchmark_lower_aat(c: &mut Criterion) { BenchmarkId::new("f32", format!("{points}x{BIGANN_DIMENSIONS}")), |bencher| { bencher.iter(|| { - sgemm_aat_lower(&data, points, BIGANN_DIMENSIONS, &mut dots).unwrap(); + sgemm_aat_lower(points, BIGANN_DIMENSIONS, &data, &mut dots).unwrap(); black_box(&dots); }); }, @@ -170,7 +170,7 @@ fn benchmark_full_leaf(c: &mut Criterion) { 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(); + sgemm_aat_lower(points, BIGANN_DIMENSIONS, &data, &mut dots).unwrap(); nearest_leaf_neighbors( LeafTopK { dots: &dots, @@ -188,7 +188,7 @@ fn benchmark_full_leaf(c: &mut Criterion) { BenchmarkId::new("l2", format!("{points}x{BIGANN_DIMENSIONS}/k{leaf_k}")), |bencher| { bencher.iter(|| { - sgemm_aat_lower(&data, points, BIGANN_DIMENSIONS, &mut dots).unwrap(); + sgemm_aat_lower(points, BIGANN_DIMENSIONS, &data, &mut dots).unwrap(); nearest_leaf_neighbors( LeafTopK { dots: &dots, diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 2aee917f3..c514de43e 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -120,13 +120,21 @@ pub enum LeafKernelError { }, } +/// Return the required output length for [`nearest_leaf_neighbors`]. +pub fn leaf_output_len(points: usize, k: usize) -> Result { + if points > u32::MAX as usize { + return Err(LeafKernelError::TooManyPoints(points)); + } + checked_area("output", points, k.min(points.saturating_sub(1))) +} + /// 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. +/// [`leaf_output_len`] entries grouped by row and ordered by ascending distance. +/// Equal distances retain pair scan order. pub fn nearest_leaf_neighbors( input: LeafTopK<'_>, k: usize, @@ -138,28 +146,33 @@ pub fn nearest_leaf_neighbors( return Ok(0); } - resize("norms", &mut workspace.norms, input.points, 0.0)?; + let uses_norms = matches!(input.metric, Metric::L2 | Metric::Cosine); + if uses_norms { + resize("norms", &mut workspace.norms, input.points, 0.0)?; + 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 + }; + } + } else { + workspace.norms.clear(); + } 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); @@ -187,13 +200,10 @@ fn validate( k: usize, output: &[LeafNeighbor], ) -> Result { - if input.points > u32::MAX as usize { - return Err(LeafKernelError::TooManyPoints(input.points)); - } + let output_len = leaf_output_len(input.points, k)?; 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) } @@ -354,19 +364,27 @@ fn process_pairs_simd_dynamic( F::Mask: SIMDSelect, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let output_ptr = output.as_mut_ptr(); let worst_ptr = worst.as_mut_ptr(); + let uses_norms = matches!(input.metric, Metric::L2 | Metric::Cosine); for row in 1..input.points { let row_start = row * input.points; - let row_norm = F::splat(arch, norms[row]); + let row_norm = if uses_norms { + F::splat(arch, norms[row]) + } else { + F::default(arch) + }; // 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 column_norms = if uses_norms { + // SAFETY: `column + F::LANES <= row < input.points == norms.len()`. + unsafe { F::load_simd(arch, norms.as_ptr().add(column)) } + } else { + F::default(arch) + }; 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`. @@ -383,10 +401,11 @@ fn process_pairs_simd_dynamic( 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) - }; + row_worst = insert_slots( + &mut output[row * k..(row + 1) * k], + (column + lane) as u32, + distance, + ); } } let mut column_bits = column_bits; @@ -394,10 +413,11 @@ fn process_pairs_simd_dynamic( 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]) - }; + let new_worst = insert_slots( + &mut output[target * k..(target + 1) * k], + row as u32, + values[lane], + ); // SAFETY: `target < row < worst.len()`. unsafe { *worst_ptr.add(target) = new_worst }; } @@ -407,20 +427,25 @@ fn process_pairs_simd_dynamic( 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); + let (row_norm, column_norm) = if uses_norms { + // SAFETY: `column < row < input.points == norms.len()`. + (norms[row], unsafe { *norms.get_unchecked(column) }) + } else { + (0.0, 0.0) + }; + let distance = pair_distance(input.metric, dot, row_norm, 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) }; + insert_slots(&mut output[row * k..(row + 1) * 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) }; + let new_worst = insert_slots( + &mut output[column * k..(column + 1) * k], + row as u32, + distance, + ); // SAFETY: `column < row < worst.len()`. unsafe { *worst_ptr.add(column) = new_worst }; } @@ -450,19 +475,27 @@ fn process_pairs_simd_fused( F::Mask: SIMDSelect, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let output_ptr = output.as_mut_ptr(); let worst_ptr = worst.as_mut_ptr(); + let uses_norms = METRIC == L2 || METRIC == COSINE; for row in 1..input.points { let row_start = row * input.points; - let row_norm = F::splat(arch, norms[row]); + let row_norm = if uses_norms { + F::splat(arch, norms[row]) + } else { + F::default(arch) + }; // 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 column_norms = if uses_norms { + // SAFETY: `column + F::LANES <= row < input.points == norms.len()`. + unsafe { F::load_simd(arch, norms.as_ptr().add(column)) } + } else { + F::default(arch) + }; let distances = pair_distances::(arch, metric::(), dots, row_norm, column_norms); let row_eligible = distances.lt_simd(F::splat(arch, row_worst)); @@ -485,16 +518,11 @@ fn process_pairs_simd_fused( // 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, - ) - }; + row_worst = insert_fixed::( + &mut output[row * SLOTS..(row + 1) * SLOTS], + (column + lane) as u32, + distance, + ); } } let mut column_bits = column_bits; @@ -502,10 +530,11 @@ fn process_pairs_simd_fused( 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]) - }; + let new_worst = insert_fixed::( + &mut output[target * SLOTS..(target + 1) * SLOTS], + row as u32, + values[lane], + ); // SAFETY: `target < row < worst.len()`. unsafe { *worst_ptr.add(target) = new_worst }; } @@ -515,22 +544,28 @@ fn process_pairs_simd_fused( 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); + let (row_norm, column_norm) = if uses_norms { + // SAFETY: `column < row < input.points == norms.len()`. + (norms[row], unsafe { *norms.get_unchecked(column) }) + } else { + (0.0, 0.0) + }; + let distance = pair_distance(metric::(), dot, row_norm, 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) - }; + row_worst = insert_fixed::( + &mut output[row * SLOTS..(row + 1) * 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) - }; + let new_worst = insert_fixed::( + &mut output[column * SLOTS..(column + 1) * SLOTS], + row as u32, + distance, + ); // SAFETY: `column < row < worst.len()`. unsafe { *worst_ptr.add(column) = new_worst }; } @@ -551,94 +586,59 @@ const fn metric() -> Metric { } } -/// 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`. +/// Insert into a production row whose width is known at dispatch. #[inline(always)] -unsafe fn insert_slots( - output: *mut LeafNeighbor, - base: usize, - slots: usize, - position: u32, - distance: f32, -) -> f32 { +fn insert_fixed(row: &mut [LeafNeighbor], position: u32, distance: f32) -> f32 { + let row: &mut [LeafNeighbor; N] = row + .try_into() + .expect("validated fixed-width leaf output row"); let entry = LeafNeighbor::new(position, distance); - match slots { + match N { 1 => { - // SAFETY: the caller guarantees `base` is in bounds. - unsafe { *output.add(base) = entry }; + row[0] = entry; distance } 2 => { - // SAFETY: the caller guarantees `base` and `base + 1` are in bounds. - let first = unsafe { *output.add(base) }; + let first = row[0]; if distance < first.distance { - // SAFETY: as above. - unsafe { - *output.add(base) = entry; - *output.add(base + 1) = first; - } + row[0] = entry; + row[1] = first; first.distance } else { - // SAFETY: as above. - unsafe { *output.add(base + 1) = entry }; + row[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)) }; + let (first, second) = (row[0], row[1]); if distance < first.distance { - // SAFETY: as above. - unsafe { - *output.add(base) = entry; - *output.add(base + 1) = first; - *output.add(base + 2) = second; - } + row[0] = entry; + row[1] = first; + row[2] = second; } else if distance < second.distance { - // SAFETY: as above. - unsafe { - *output.add(base + 1) = entry; - *output.add(base + 2) = second; - } + row[1] = entry; + row[2] = second; } else { - // SAFETY: as above. - unsafe { *output.add(base + 2) = entry }; + row[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 } - } + _ => unreachable!("fixed leaf widths are one through three"), + } +} + +/// Insert into the uncommon run-time-width row (`k > 3`). +#[inline(always)] +fn insert_slots(row: &mut [LeafNeighbor], position: u32, distance: f32) -> f32 { + let last = row.len() - 1; + row[last] = LeafNeighbor::new(position, distance); + let mut index = last; + while index > 0 && row[index].distance < row[index - 1].distance { + row.swap(index, index - 1); + index -= 1; } + row[last].distance } #[inline(always)] diff --git a/diskann-pipnn/src/leaf_kernel/tests.rs b/diskann-pipnn/src/leaf_kernel/tests.rs index 1dba1805c..1bef128ab 100644 --- a/diskann-pipnn/src/leaf_kernel/tests.rs +++ b/diskann-pipnn/src/leaf_kernel/tests.rs @@ -111,6 +111,17 @@ fn scalar_insertion_orders_candidates_and_rejects_nan() { assert_eq!(worst, [3.0]); } +#[test] +fn output_length_clamps_to_non_self_neighbors() { + assert_eq!(leaf_output_len(0, 3).unwrap(), 0); + assert_eq!(leaf_output_len(1, 3).unwrap(), 0); + assert_eq!(leaf_output_len(4, 9).unwrap(), 12); + assert_eq!( + leaf_output_len(u32::MAX as usize + 1, 1), + Err(LeafKernelError::TooManyPoints(u32::MAX as usize + 1)) + ); +} + #[test] fn workspace_can_shrink_and_grow_between_calls() { let mut workspace = LeafTopKWorkspace::new(); diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index 198434b72..17a5dd708 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -3,7 +3,20 @@ * Licensed under the MIT license. */ -//! PiPNN graph construction. +//! Numerical kernels used by PiPNN graph construction. +//! +//! PiPNN first partitions points around sampled leaders, then builds local +//! neighbor candidates inside each leaf. This crate owns the numerical seams +//! of those stages while callers retain dataset storage, GEMM workspaces, graph +//! policy, and scheduling: +//! +//! - [`partition_kernel`] converts a point-by-leader dot-product tile into the +//! nearest leader positions for each point. +//! - [`leaf_kernel`] scans a leaf's lower-triangular dot-product matrix once and +//! retains nearest non-self neighbors for both endpoints. +//! +//! Both modules validate slice shapes before dispatch and use `diskann-wide` for +//! architecture selection; PiPNN does not detect or name instruction sets. pub mod leaf_kernel; pub mod partition_kernel; diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index ab4d02255..8122244f9 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -21,11 +21,25 @@ use diskann_vector::distance::Metric; use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector}; /// Maximum number of leaders retained for one point. +/// +/// Supported PiPNN partition fanouts fit within 16. Keeping this as a fixed +/// stack tracker bounds per-row stack use and code size; larger requests are +/// rejected rather than silently truncated. 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. +/// One row-major point-by-leader dot-product tile and its normalization terms. +/// +/// The scale slices are deliberately metric-specific: +/// +/// | metric | `row_scales` | `leader_scales` | +/// |---|---|---| +/// | [`Metric::L2`] | empty | squared leader norms | +/// | [`Metric::Cosine`] | squared point norms | leader norms | +/// | [`Metric::CosineNormalized`] / [`Metric::InnerProduct`] | empty | empty | +/// +/// [`nearest_leaders`] validates every declared shape before dispatch. #[derive(Clone, Copy, Debug)] pub struct PartitionTopK<'a> { /// Row-major `rows * leaders` point-to-leader dot products. @@ -34,9 +48,9 @@ pub struct PartitionTopK<'a> { pub rows: usize, /// Number of leaders represented by each row. pub leaders: usize, - /// Squared point norms for cosine, otherwise empty. + /// Metric-specific point normalization terms described in the type table. pub row_scales: &'a [f32], - /// Leader norms for cosine, squared leader norms for L2, otherwise empty. + /// Metric-specific leader normalization terms described in the type table. pub leader_scales: &'a [f32], /// Distance metric used to rank leaders. pub metric: Metric, @@ -66,7 +80,9 @@ pub enum PartitionKernelError { actual: usize, }, /// The requested fanout cannot be represented by the fixed top-k tracker. - #[error("invalid fanout {fanout} for {leaders} leaders; maximum is {maximum}")] + #[error( + "invalid fanout {fanout}: must not exceed {leaders} leaders or kernel maximum {maximum}" + )] InvalidFanout { /// Requested number of leaders per row. fanout: usize, @@ -113,7 +129,7 @@ pub fn nearest_leaders( }); if let Some(row) = output .chunks_exact(fanout) - .position(|leaders| leaders.contains(&u32::MAX)) + .position(|leaders| leaders[fanout - 1] == u32::MAX) { return Err(PartitionKernelError::InsufficientRankableDistances { row, fanout }); } @@ -234,44 +250,75 @@ where 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::( + match input.metric { + Metric::L2 => process_rows(input, fanout, output, |_, dot_row, top| { + process_binary::( arch, dot_row, input.leader_scales, - &mut top, + top, fanout, |dot, norm| F::splat(arch, -2.0).mul_add_simd(dot, norm), |dot, norm| norm - 2.0 * dot, - ), - 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::( + ); + }), + Metric::CosineNormalized => process_rows(input, fanout, output, |_, dot_row, top| { + process_unary::(arch, dot_row, top, fanout, |dot| F::splat(arch, 1.0) - dot); + }), + Metric::InnerProduct => process_rows(input, fanout, output, |_, dot_row, top| { + process_unary::(arch, dot_row, top, fanout, |dot| F::default(arch) - dot); + }), + Metric::Cosine => process_rows(input, fanout, output, |row, dot_row, top| { + process_cosine::( arch, dot_row, - input.row_scales[row_index], + input.row_scales[row], input.leader_scales, - &mut top, + top, fanout, - ), - } + ); + }), + } +} + +#[inline(always)] +fn process_rows( + input: PartitionTopK<'_>, + fanout: usize, + output: &mut [u32], + mut process: impl FnMut(usize, &[f32], &mut TopK), +) { + for (row, (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]; + process(row, dot_row, &mut top); copy_ids(&top, output_row); } } +#[inline(always)] +fn cosine_distance(row_norm_squared: f32, leader_norm: f32, dot: f32) -> f32 { + let row_norm = if row_norm_squared < f32::MIN_POSITIVE { + 0.0 + } else { + row_norm_squared.sqrt() + }; + let leader_norm = if leader_norm < f32::MIN_POSITIVE.sqrt() { + 0.0 + } else { + leader_norm + }; + if row_norm == 0.0 || leader_norm == 0.0 { + 1.0 + } else { + 1.0 - dot / (row_norm * leader_norm) + } +} + fn process_cosine( arch: F::Arch, dots: &[f32], @@ -284,9 +331,14 @@ fn process_cosine( F::Mask: SIMDSelect, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let row_norm = F::splat(arch, row_norm_squared.sqrt()); + let row_norm = if row_norm_squared < f32::MIN_POSITIVE { + 0.0 + } else { + row_norm_squared.sqrt() + }; + let row_norm = F::splat(arch, row_norm); let one = F::splat(arch, 1.0); - let zero = F::default(arch); + let minimum_norm = F::splat(arch, f32::MIN_POSITIVE.sqrt()); process_binary::( arch, dots, @@ -294,21 +346,17 @@ fn process_cosine( top, fanout, |dot, leader_norm| { + let row_zero = row_norm.lt_simd(minimum_norm); + let leader_zero = leader_norm.lt_simd(minimum_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); + let safe_denominator = row_zero.select(one, leader_zero.select(one, denominator)); + let cosine = row_zero.select( + F::default(arch), + leader_zero.select(F::default(arch), dot / safe_denominator), + ); one - cosine }, - |dot, leader_norm| { - let denominator = row_norm_squared.sqrt() * leader_norm; - let cosine = if denominator > 0.0 { - dot / denominator - } else { - 0.0 - }; - 1.0 - cosine - }, + |dot, leader_norm| cosine_distance(row_norm_squared, leader_norm, dot), ); } @@ -391,15 +439,7 @@ fn distance(metric: Metric, dot: f32, row_scale: f32, leader_scale: f32) -> f32 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 - } + Metric::Cosine => cosine_distance(row_scale, leader_scale, dot), } } diff --git a/diskann-pipnn/src/partition_kernel/tests.rs b/diskann-pipnn/src/partition_kernel/tests.rs index ad3e048a0..7da20fd58 100644 --- a/diskann-pipnn/src/partition_kernel/tests.rs +++ b/diskann-pipnn/src/partition_kernel/tests.rs @@ -15,7 +15,9 @@ fn input(metric: Metric, leaders: usize) -> (Vec, Vec, Vec) { Vec::new() }; let leader_scales = match metric { - Metric::L2 => (0..leaders).map(|leader| (leader + 1) as f32).collect(), + Metric::L2 => (0..leaders) + .map(|leader| ((leader + 1) as f32).powi(2)) + .collect(), Metric::Cosine => (0..leaders) .map(|leader| { if leader == 0 { @@ -88,6 +90,50 @@ fn scalar_distance_matches_metric_contract() { 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); + assert_eq!( + distance(Metric::Cosine, 1.0, f32::MIN_POSITIVE / 2.0, 1.0), + 1.0 + ); + assert_eq!( + distance( + Metric::Cosine, + f32::MIN_POSITIVE, + f32::MIN_POSITIVE, + f32::MIN_POSITIVE.sqrt() + ), + 0.0 + ); + assert!(distance(Metric::Cosine, 1.0, f32::NAN, 1.0).is_nan()); +} + +#[test] +fn cosine_special_norms_match_scalar_and_runtime_dispatch() { + let leaders = 17; + let dots = vec![1.0; 4 * leaders]; + let row_scales = [0.0, f32::MIN_POSITIVE / 2.0, f32::MIN_POSITIVE, f32::NAN]; + let mut leader_scales = vec![1.0; leaders]; + leader_scales[..4].copy_from_slice(&[ + 0.0, + f32::MIN_POSITIVE.sqrt() / 2.0, + f32::MIN_POSITIVE.sqrt(), + f32::NAN, + ]); + let input = PartitionTopK { + dots: &dots, + rows: row_scales.len(), + leaders, + row_scales: &row_scales, + leader_scales: &leader_scales, + metric: Metric::Cosine, + }; + let mut expected = vec![u32::MAX; input.rows * 2]; + process_rows_scalar(input, 2, &mut expected); + let mut actual = vec![u32::MAX; input.rows * 2]; + nearest_leaders(input, 2, &mut actual).unwrap(); + + assert_eq!(actual, expected); + assert_eq!(&actual[..4], &[0, 1, 0, 1]); + assert_eq!(&actual[6..], &[0, 1]); } #[test] diff --git a/diskann-pipnn/tests/leaf_kernel.rs b/diskann-pipnn/tests/leaf_kernel.rs index 5cead7e96..5971f4cb6 100644 --- a/diskann-pipnn/tests/leaf_kernel.rs +++ b/diskann-pipnn/tests/leaf_kernel.rs @@ -9,24 +9,36 @@ use diskann_pipnn::leaf_kernel::{ use diskann_vector::distance::Metric; use std::cmp::Ordering; +const SIMD_BOUNDARY_POINTS: [usize; 9] = [7, 8, 9, 15, 16, 17, 64, 256, 512]; +const ZERO_NORM_POSITION: usize = 0; +const DISTINCT_NORM_POSITION: usize = 2; +const NORM_PERIOD: usize = 5; +const ROW_MIXER: usize = 17; +const COLUMN_MIXER: usize = 11; +const MIX_MODULUS: usize = 23; +const MIX_CENTER: f32 = 11.0; +const DOT_SCALE: f32 = 1.0 / 32.0; +const TIED_COLUMNS: [usize; 2] = [1, 2]; + 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 { + dots[row * points + row] = if metric == Metric::Cosine && row == ZERO_NORM_POSITION { 0.0 - } else if row == 2 { + } else if row == DISTINCT_NORM_POSITION { 2.0 } else { - 1.0 + (row % 5) as f32 + 1.0 + (row % NORM_PERIOD) as f32 }; for column in 0..row { - let pair = ((row * 17 + column * 11) % 23) as f32 - 11.0; + let pair = + ((row * ROW_MIXER + column * COLUMN_MIXER) % MIX_MODULUS) as f32 - MIX_CENTER; dots[row * points + column] = if row == points - 1 && column == 0 { f32::NAN - } else if column == 1 || column == 2 { + } else if TIED_COLUMNS.contains(&column) { 0.5 } else { - pair * 0.03125 + pair * DOT_SCALE }; } } @@ -111,7 +123,8 @@ fn dispatch_matches_reference_across_simd_width_boundaries() { Metric::CosineNormalized, Metric::InnerProduct, ] { - for points in [7, 8, 9, 15, 16, 17, 64, 256, 512] { + // Straddle the 8- and 16-lane boundaries, then cover production leaf sizes. + for points in SIMD_BOUNDARY_POINTS { let dots = differential_input(metric, points); let input = LeafTopK { dots: &dots, From b4a73e98242f9278fd078d987a966877e023fae3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 09:55:57 +0000 Subject: [PATCH 13/26] refactor(pipnn)!: prepare kernel dispatch Select architecture, metric, and leaf-width implementations once, then reuse direct diskann-wide function pointers across stripes and leaves. BREAKING CHANGE: callers construct LeafKernel or PartitionKernel and pass MatrixView-backed inputs and outputs. --- Cargo.lock | 3 +- diskann-pipnn/Cargo.toml | 9 +- diskann-pipnn/benches/kernels.rs | 224 ---- diskann-pipnn/src/kernel_metric.rs | 314 ++++++ diskann-pipnn/src/leaf_kernel.rs | 996 ++++++++++-------- diskann-pipnn/src/leaf_kernel/tests.rs | 144 --- diskann-pipnn/src/lib.rs | 17 +- diskann-pipnn/src/partition_kernel.rs | 839 ++++++++++----- diskann-pipnn/src/partition_kernel/tests.rs | 148 --- .../{leaf_kernel.rs => leaf_kernel_api.rs} | 316 ++---- diskann-pipnn/tests/partition_kernel.rs | 442 -------- diskann-pipnn/tests/partition_kernel_api.rs | 358 +++++++ 12 files changed, 1917 insertions(+), 1893 deletions(-) delete mode 100644 diskann-pipnn/benches/kernels.rs create mode 100644 diskann-pipnn/src/kernel_metric.rs delete mode 100644 diskann-pipnn/src/leaf_kernel/tests.rs delete mode 100644 diskann-pipnn/src/partition_kernel/tests.rs rename diskann-pipnn/tests/{leaf_kernel.rs => leaf_kernel_api.rs} (56%) delete mode 100644 diskann-pipnn/tests/partition_kernel.rs create mode 100644 diskann-pipnn/tests/partition_kernel_api.rs diff --git a/Cargo.lock b/Cargo.lock index 0c9926781..059a0b2c6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -684,8 +684,7 @@ dependencies = [ name = "diskann-pipnn" version = "0.55.0" dependencies = [ - "criterion", - "diskann-linalg", + "diskann-utils", "diskann-vector", "diskann-wide", "thiserror 2.0.17", diff --git a/diskann-pipnn/Cargo.toml b/diskann-pipnn/Cargo.toml index 1ff7c3bfd..848fbce2a 100644 --- a/diskann-pipnn/Cargo.toml +++ b/diskann-pipnn/Cargo.toml @@ -11,17 +11,10 @@ license.workspace = true edition.workspace = true [dependencies] +diskann-utils.workspace = true 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 deleted file mode 100644 index 52661f519..000000000 --- a/diskann-pipnn/benches/kernels.rs +++ /dev/null @@ -1,224 +0,0 @@ -/* - * 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(points, BIGANN_DIMENSIONS, &data, &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(points, BIGANN_DIMENSIONS, &data, &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(points, BIGANN_DIMENSIONS, &data, &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(points, BIGANN_DIMENSIONS, &data, &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/kernel_metric.rs b/diskann-pipnn/src/kernel_metric.rs new file mode 100644 index 000000000..f7e61d43b --- /dev/null +++ b/diskann-pipnn/src/kernel_metric.rs @@ -0,0 +1,314 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Metric marker types shared by the partition and leaf kernels. +//! +//! Runtime metric selection happens only while preparing a dispatched kernel. +//! The hot loops receive a concrete marker type, allowing metric arithmetic and +//! scale handling to inline without a per-row or per-chunk enum match. + +use diskann_vector::distance::Metric; +use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum ScaleKind { + None, + SquaredNorm, + NormFromSquared, + Norm, +} + +impl ScaleKind { + #[inline(always)] + pub(crate) fn transform(self, stored: f32) -> f32 { + match self { + Self::None => 0.0, + Self::SquaredNorm => stored, + Self::Norm => { + if stored < f32::MIN_POSITIVE.sqrt() { + 0.0 + } else { + stored + } + } + Self::NormFromSquared => { + if stored < f32::MIN_POSITIVE { + 0.0 + } else { + stored.sqrt() + } + } + } + } + + pub(crate) const fn is_some(self) -> bool { + !matches!(self, Self::None) + } +} + +pub(crate) trait KernelMetric: Send + Sync + 'static { + const METRIC: Metric; + const LEAF_SCALE: ScaleKind; + const PARTITION_ROW_SCALE: ScaleKind; + const PARTITION_LEADER_SCALE: ScaleKind; + + fn leaf_distance(arch: F::Arch, dot: F, row_scale: F, column_scale: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect; + + fn leaf_distance_scalar(dot: f32, row_scale: f32, column_scale: f32) -> f32; + + fn partition_distance(arch: F::Arch, dot: F, row_scale: F, leader_scale: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect; + + fn partition_distance_scalar(dot: f32, row_scale: f32, leader_scale: f32) -> f32; +} + +pub(crate) struct L2; +pub(crate) struct Cosine; +pub(crate) struct CosineNormalized; +pub(crate) struct InnerProduct; + +#[inline(always)] +fn clamp_nonnegative(arch: F::Arch, distance: F) -> F +where + F: SIMDVector + SIMDFloat, + F::Mask: SIMDSelect, +{ + let zero = F::default(arch); + // SIMD max has ISA-specific NaN behavior. Select the original NaN so it + // remains non-rankable on every backend. + distance + .eq_simd(distance) + .select(zero.max_simd(distance), distance) +} + +#[inline(always)] +fn clamp_nonnegative_scalar(distance: f32) -> f32 { + if distance < 0.0 { + 0.0 + } else { + distance + } +} + +#[inline(always)] +fn cosine_distance(arch: F::Arch, 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 one = F::splat(arch, 1.0); + let minimum_norm = F::splat(arch, f32::MIN_POSITIVE.sqrt()); + let row_zero = row_norm.lt_simd(minimum_norm); + let column_zero = column_norm.lt_simd(minimum_norm); + 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)); + one - cosine +} + +#[inline(always)] +fn cosine_distance_scalar(dot: f32, row_norm: f32, column_norm: f32) -> f32 { + if row_norm < f32::MIN_POSITIVE.sqrt() || column_norm < f32::MIN_POSITIVE.sqrt() { + 1.0 + } else { + 1.0 - dot / (row_norm * column_norm) + } +} + +impl KernelMetric for L2 { + const METRIC: Metric = Metric::L2; + const LEAF_SCALE: ScaleKind = ScaleKind::SquaredNorm; + const PARTITION_ROW_SCALE: ScaleKind = ScaleKind::None; + const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::SquaredNorm; + + #[inline(always)] + fn leaf_distance(arch: F::Arch, dot: F, row_scale: F, column_scale: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + clamp_nonnegative(arch, row_scale + column_scale - F::splat(arch, 2.0) * dot) + } + + #[inline(always)] + fn leaf_distance_scalar(dot: f32, row_scale: f32, column_scale: f32) -> f32 { + clamp_nonnegative_scalar(row_scale + column_scale - 2.0 * dot) + } + + #[inline(always)] + fn partition_distance(arch: F::Arch, dot: F, _: F, leader_scale: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + F::splat(arch, -2.0).mul_add_simd(dot, leader_scale) + } + + #[inline(always)] + fn partition_distance_scalar(dot: f32, _: f32, leader_scale: f32) -> f32 { + // Preserve the scalar reduction shape used by the original partition + // kernel; changing this rounding can change leader tie order. + leader_scale - 2.0 * dot + } +} + +impl KernelMetric for Cosine { + const METRIC: Metric = Metric::Cosine; + const LEAF_SCALE: ScaleKind = ScaleKind::NormFromSquared; + const PARTITION_ROW_SCALE: ScaleKind = ScaleKind::NormFromSquared; + const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::Norm; + + #[inline(always)] + fn leaf_distance(arch: F::Arch, dot: F, row_scale: F, column_scale: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + clamp_nonnegative(arch, cosine_distance(arch, dot, row_scale, column_scale)) + } + + #[inline(always)] + fn leaf_distance_scalar(dot: f32, row_scale: f32, column_scale: f32) -> f32 { + clamp_nonnegative_scalar(cosine_distance_scalar(dot, row_scale, column_scale)) + } + + #[inline(always)] + fn partition_distance(arch: F::Arch, dot: F, row_scale: F, leader_scale: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + cosine_distance(arch, dot, row_scale, leader_scale) + } + + #[inline(always)] + fn partition_distance_scalar(dot: f32, row_scale: f32, leader_scale: f32) -> f32 { + cosine_distance_scalar(dot, row_scale, leader_scale) + } +} + +impl KernelMetric for CosineNormalized { + const METRIC: Metric = Metric::CosineNormalized; + const LEAF_SCALE: ScaleKind = ScaleKind::None; + const PARTITION_ROW_SCALE: ScaleKind = ScaleKind::None; + const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::None; + + #[inline(always)] + fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + clamp_nonnegative(arch, F::splat(arch, 1.0) - dot) + } + + #[inline(always)] + fn leaf_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { + clamp_nonnegative_scalar(1.0 - dot) + } + + #[inline(always)] + fn partition_distance(arch: F::Arch, dot: F, _: F, _: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + F::splat(arch, 1.0) - dot + } + + #[inline(always)] + fn partition_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { + 1.0 - dot + } +} + +impl KernelMetric for InnerProduct { + const METRIC: Metric = Metric::InnerProduct; + const LEAF_SCALE: ScaleKind = ScaleKind::None; + const PARTITION_ROW_SCALE: ScaleKind = ScaleKind::None; + const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::None; + + #[inline(always)] + fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + F::default(arch) - dot + } + + #[inline(always)] + fn leaf_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { + -dot + } + + #[inline(always)] + fn partition_distance(arch: F::Arch, dot: F, _: F, _: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + F::default(arch) - dot + } + + #[inline(always)] + fn partition_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { + -dot + } +} + +pub(crate) trait EraseMetric { + type Output; + + fn erase(self) -> Self::Output; +} + +pub(crate) fn erase_metric(metric: Metric, erase: E) -> E::Output { + match metric { + Metric::L2 => erase.erase::(), + Metric::Cosine => erase.erase::(), + Metric::CosineNormalized => erase.erase::(), + Metric::InnerProduct => erase.erase::(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn norm_scales_apply_zero_threshold_without_erasing_nan() { + assert_eq!(ScaleKind::Norm.transform(-0.0).to_bits(), 0.0f32.to_bits()); + assert_eq!( + ScaleKind::Norm.transform(f32::MIN_POSITIVE.sqrt() / 2.0), + 0.0 + ); + assert_eq!( + ScaleKind::NormFromSquared.transform(f32::MIN_POSITIVE / 2.0), + 0.0 + ); + assert_eq!( + ScaleKind::NormFromSquared.transform(f32::MIN_POSITIVE), + f32::MIN_POSITIVE.sqrt() + ); + assert!(ScaleKind::Norm.transform(f32::NAN).is_nan()); + assert!(ScaleKind::NormFromSquared.transform(f32::NAN).is_nan()); + } + + #[test] + fn l2_partition_scalar_tail_preserves_non_fused_rounding() { + let scalar = L2::partition_distance_scalar(f32::MAX, 0.0, f32::MAX); + let fused = (-2.0f32).mul_add(f32::MAX, f32::MAX); + + assert_eq!(scalar, f32::NEG_INFINITY); + assert_eq!(fused, -f32::MAX); + } +} diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index c514de43e..d0869bf61 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -3,28 +3,26 @@ * Licensed under the MIT license. */ -//! Fused nearest-neighbor selection over a leaf's lower dot-product matrix. +//! Prepared nearest-neighbor kernels over a leaf's lower dot-product matrix. //! -//! `sgemm_aat_lower` writes only pair `(row, column)` with `column <= row`. -//! This kernel therefore walks the strict lower triangle once and offers each -//! computed distance to both endpoint rows. Keeping one top-k tracker per row -//! avoids materializing the upper triangle or computing a symmetric distance -//! twice. -//! -//! The public entry point validates every shape before dispatch. The dispatched -//! path processes complete SIMD chunks, then a scalar tail. For `k <= 3`, const -//! slot counts remove the dynamic insertion loop from the hot path; larger `k` -//! uses the same ordering rules through the dynamic fallback. NaN distances are -//! never rankable, and ties retain scan order so scalar and SIMD backends produce -//! the same graph. +//! `sgemm_aat_lower` writes pair `(row, column)` only when `column <= row`. +//! The kernel scans that strict lower triangle once and offers each distance to +//! both endpoint rows. A [`LeafKernel`] is prepared once for the build metric, +//! requested neighbor count, and runtime CPU; repeated leaves call a direct +//! `diskann-wide` function pointer without ISA or metric dispatch in the loop. +//! NaN distances are not rankable, and equal distances retain pair scan order. + +use std::marker::PhantomData; +use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; -use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; +use diskann_wide::{ + arch::{self, Dispatched1, FTarget1}, + lifetime::AddLifetime, + Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, +}; -const L2: u8 = 0; -const COSINE_NORMALIZED: u8 = 1; -const INNER_PRODUCT: u8 = 2; -const COSINE: u8 = 3; +use crate::kernel_metric::{erase_metric, EraseMetric, KernelMetric}; /// One leaf-local neighbor and its metric distance. #[derive(Clone, Copy, Debug, PartialEq)] @@ -48,15 +46,11 @@ impl Default for LeafNeighbor { } } -/// Lower-triangular dot products consumed by [`nearest_leaf_neighbors`]. +/// Square lower-triangular dot-product matrix for one leaf. #[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, + /// Point-by-point matrix. Only entries with `column <= row` are read. + pub dots: MatrixView<'a, f32>, } /// Reusable temporary storage for leaf top-k selection. @@ -76,13 +70,21 @@ impl LeafTopKWorkspace { } } -/// Validation or allocation error returned by [`nearest_leaf_neighbors`]. +/// Validation or allocation error returned by [`LeafKernel::nearest_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`. + /// The dot-product matrix is not square. + #[error("leaf dot-product matrix must be square, got {rows} x {cols}")] + NonSquareDots { + /// Supplied row count. + rows: usize, + /// Supplied column count. + cols: usize, + }, + /// A declared output shape overflowed `usize`. #[error("{buffer} shape {rows} x {cols} overflows usize")] ShapeOverflow { /// Name of the buffer whose shape overflowed. @@ -92,7 +94,7 @@ pub enum LeafKernelError { /// Declared column count. cols: usize, }, - /// A supplied slice did not match its declared shape. + /// A view's backing slice does not match its declared shape. #[error("invalid {buffer} length: expected {expected}, got {actual}")] InvalidBufferLength { /// Name of the invalid buffer. @@ -102,6 +104,20 @@ pub enum LeafKernelError { /// Supplied length. actual: usize, }, + /// The output matrix does not match the requested neighbor shape. + #[error( + "invalid output shape: expected {expected_rows} x {expected_cols}, got {actual_rows} x {actual_cols}" + )] + InvalidOutputShape { + /// Required row count. + expected_rows: usize, + /// Required column count. + expected_cols: usize, + /// Supplied row count. + actual_rows: usize, + /// Supplied column count. + actual_cols: usize, + }, /// Temporary storage could not be reserved. #[error("failed to reserve {additional} values for {buffer}")] Allocation { @@ -120,7 +136,7 @@ pub enum LeafKernelError { }, } -/// Return the required output length for [`nearest_leaf_neighbors`]. +/// Return the required output length for [`LeafKernel::nearest_neighbors`]. pub fn leaf_output_len(points: usize, k: usize) -> Result { if points > u32::MAX as usize { return Err(LeafKernelError::TooManyPoints(points)); @@ -128,41 +144,228 @@ pub fn leaf_output_len(points: usize, k: usize) -> Result { + input: LeafTopK<'a>, + output: MutMatrixView<'a, LeafNeighbor>, + workspace: &'a mut LeafTopKWorkspace, + requested_k: usize, +} + +#[derive(Debug)] +struct LeafCallArg; + +impl AddLifetime for LeafCallArg { + type Of<'a> = LeafCall<'a>; +} + +type LeafFn = Dispatched1, LeafCallArg>; + +/// A leaf kernel prepared for one metric, neighbor count, and the current CPU. /// -/// 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 -/// [`leaf_output_len`] entries grouped by row and ordered by ascending distance. -/// Equal distances retain pair scan order. -pub fn nearest_leaf_neighbors( +/// Construct this once with [`LeafKernel::new`] and share it across leaf workers. +/// The handle stores only a direct function pointer and the requested `k`. +#[derive(Clone, Copy, Debug)] +pub struct LeafKernel { + run: LeafFn, + requested_k: usize, +} + +impl LeafKernel { + /// Prepare a leaf kernel for `metric`, `k`, and the current CPU. + pub fn new(metric: Metric, k: usize) -> Self { + diskann_wide::arch::dispatch1_no_features(PrepareLeaf { requested_k: k }, metric) + } + + /// Select the nearest non-self leaf positions for every row. + /// + /// `output` must have `input.dots.nrows()` rows and + /// `min(k, rows - 1)` columns. The returned value is that effective column + /// count. Equal distances retain pair scan order. + pub fn nearest_neighbors( + &self, + input: LeafTopK<'_>, + output: MutMatrixView<'_, LeafNeighbor>, + workspace: &mut LeafTopKWorkspace, + ) -> Result { + self.run.call(LeafCall { + input, + output, + workspace, + requested_k: self.requested_k, + }) + } +} + +#[derive(Clone, Copy, Debug)] +enum KValue { + One, + Two, + Three, + Large, +} + +impl KValue { + const fn from_requested(k: usize) -> Self { + match k { + 1 => Self::One, + 2 => Self::Two, + 3 => Self::Three, + _ => Self::Large, + } + } +} + +struct PrepareLeaf { + requested_k: usize, +} + +impl arch::Target1 for PrepareLeaf +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + fn run(self, arch: A, metric: Metric) -> LeafKernel { + erase_metric( + metric, + BuildLeaf { + arch, + requested_k: self.requested_k, + }, + ) + } +} + +struct BuildLeaf { + arch: A, + requested_k: usize, +} + +impl BuildLeaf +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + fn build(self) -> LeafKernel { + LeafKernel { + run: self + .arch + .dispatch1::, Result, LeafCallArg>(), + requested_k: self.requested_k, + } + } +} + +impl EraseMetric for BuildLeaf +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + type Output = LeafKernel; + + fn erase(self) -> Self::Output { + match KValue::from_requested(self.requested_k) { + KValue::One => self.build::>(), + KValue::Two => self.build::>(), + KValue::Three => self.build::>(), + KValue::Large => self.build::(), + } + } +} + +struct LeafEntry(PhantomData<(M, S)>); + +impl FTarget1, LeafCall<'_>> for LeafEntry +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + M: KernelMetric, + S: SlotSelection, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + fn run(arch: A, mut call: LeafCall<'_>) -> Result { + let actual_k = validate(call.input, call.requested_k, &call.output)?; + if actual_k == 0 { + return Ok(0); + } + + prepare_workspace::(call.input, call.workspace)?; + call.output.as_mut_slice().fill(LeafNeighbor::default()); + call.workspace.worst.fill(f32::INFINITY); + + S::process::( + arch, + call.input, + actual_k, + call.output.as_mut_slice(), + &call.workspace.norms, + &mut call.workspace.worst, + ); + if let Some(row) = call + .output + .as_slice() + .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: &mut [LeafNeighbor], - workspace: &mut LeafTopKWorkspace, + output: &MutMatrixView<'_, LeafNeighbor>, ) -> Result { - let actual_k = validate(input, k, output)?; - if actual_k == 0 { - return Ok(0); + let rows = input.dots.nrows(); + let columns = input.dots.ncols(); + if rows != columns { + return Err(LeafKernelError::NonSquareDots { + rows, + cols: columns, + }); + } + let output_len = leaf_output_len(rows, k)?; + let dots_len = checked_area("leaf dot-product matrix", rows, columns)?; + check_length( + "leaf dot-product matrix", + input.dots.as_slice().len(), + dots_len, + )?; + + let actual_k = k.min(rows.saturating_sub(1)); + if output.nrows() != rows || output.ncols() != actual_k { + return Err(LeafKernelError::InvalidOutputShape { + expected_rows: rows, + expected_cols: actual_k, + actual_rows: output.nrows(), + actual_cols: output.ncols(), + }); } + check_length("output", output.as_slice().len(), output_len)?; + Ok(actual_k) +} - let uses_norms = matches!(input.metric, Metric::L2 | Metric::Cosine); - if uses_norms { - resize("norms", &mut workspace.norms, input.points, 0.0)?; +fn prepare_workspace( + input: LeafTopK<'_>, + workspace: &mut LeafTopKWorkspace, +) -> Result<(), LeafKernelError> { + let points = input.dots.nrows(); + if M::LEAF_SCALE.is_some() { + resize("norms", &mut workspace.norms, points, 0.0)?; 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 - }; + *norm = M::LEAF_SCALE.transform(input.dots[(row, row)]); } } else { workspace.norms.clear(); @@ -170,42 +373,9 @@ pub fn nearest_leaf_neighbors( resize( "worst distances", &mut workspace.worst, - input.points, + points, f32::INFINITY, - )?; - 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 { - let output_len = leaf_output_len(input.points, k)?; - 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)); - check_length("output", output.len(), output_len)?; - Ok(actual_k) + ) } fn resize( @@ -243,155 +413,199 @@ fn check_length( } } -struct LeafKernel<'a, 'o, 'w> { - input: LeafTopK<'a>, - k: usize, - output: &'o mut [LeafNeighbor], - norms: &'w [f32], - worst: &'w mut [f32], +trait SlotSelection: Send + Sync + 'static { + fn process( + arch: F::Arch, + input: LeafTopK<'_>, + actual_k: usize, + output: &mut [LeafNeighbor], + norms: &[f32], + worst: &mut [f32], + ) where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + M: KernelMetric, + u64: From<<::BitMask as SIMDMask>::Underlying>; } -impl LeafKernel<'_, '_, '_> { - fn run_simd(self, arch: F::Arch) - where +struct FixedSelection; +struct DynamicSelection; + +impl SlotSelection for FixedSelection { + fn process( + arch: F::Arch, + input: LeafTopK<'_>, + actual_k: usize, + output: &mut [LeafNeighbor], + norms: &[f32], + worst: &mut [f32], + ) where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, + M: KernelMetric, 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"), - } + debug_assert!(actual_k <= N); + process_selected::(arch, input, actual_k, output, norms, worst); } +} - fn run_fused(self, arch: F::Arch) - where +impl SlotSelection for DynamicSelection { + fn process( + arch: F::Arch, + input: LeafTopK<'_>, + actual_k: usize, + output: &mut [LeafNeighbor], + norms: &[f32], + worst: &mut [f32], + ) where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, + M: KernelMetric, 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, - ), - } + process_selected::(arch, input, actual_k, output, norms, 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>, +fn process_selected( + arch: F::Arch, + input: LeafTopK<'_>, + actual_k: usize, + output: &mut [LeafNeighbor], + norms: &[f32], + worst: &mut [f32], +) where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + M: KernelMetric, + u64: From<<::BitMask as SIMDMask>::Underlying>, { - #[inline(always)] - fn run(self, arch: A) { - self.run_simd::(arch); + match actual_k { + 1 => process_fixed::(arch, input, output, norms, worst), + 2 => process_fixed::(arch, input, output, norms, worst), + 3 => process_fixed::(arch, input, output, norms, worst), + width => process_pairs::( + arch, + input, + DynamicRows { + values: output, + width, + }, + norms, + worst, + ), } } -#[cfg(test)] -fn process_pairs_scalar( +fn process_fixed( + arch: F::Arch, 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); - } +) where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + M: KernelMetric, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + let (rows, remainder) = output.as_chunks_mut::(); + debug_assert!(remainder.is_empty()); + process_pairs::(arch, input, FixedRows(rows), norms, worst); +} + +trait NeighborRows { + fn len(&self) -> usize; + fn insert(&mut self, row: usize, position: u32, distance: f32) -> f32; +} + +struct FixedRows<'a, const N: usize>(&'a mut [[LeafNeighbor; N]]); + +impl NeighborRows for FixedRows<'_, N> { + #[inline(always)] + fn len(&self) -> usize { + self.0.len() + } + + #[inline(always)] + fn insert(&mut self, row: usize, position: u32, distance: f32) -> f32 { + insert_fixed(&mut self.0[row], position, distance) } } -/// Fused dual-endpoint scan for row widths without a specialized arm. +struct DynamicRows<'a> { + values: &'a mut [LeafNeighbor], + width: usize, +} + +impl NeighborRows for DynamicRows<'_> { + #[inline(always)] + fn len(&self) -> usize { + self.values.len() / self.width + } + + #[inline(always)] + fn insert(&mut self, row: usize, position: u32, distance: f32) -> f32 { + insert_dynamic( + &mut self.values[row * self.width..(row + 1) * self.width], + position, + distance, + ) + } +} + +/// Scan the strict lower triangle and update both endpoint rows. /// -/// 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( +/// `M` fixes metric arithmetic before type erasure. `R` presents either +/// fixed-width array rows or the uncommon run-time-width rows. +#[inline(never)] +fn process_pairs( arch: F::Arch, input: LeafTopK<'_>, - k: usize, - output: &mut [LeafNeighbor], + mut output: R, norms: &[f32], worst: &mut [f32], ) where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, + M: KernelMetric, + R: NeighborRows, u64: From<<::BitMask as SIMDMask>::Underlying>, { + let points = input.dots.nrows(); + let dots = input.dots.as_slice(); + let uses_norms = M::LEAF_SCALE.is_some(); let worst_ptr = worst.as_mut_ptr(); - let uses_norms = matches!(input.metric, Metric::L2 | Metric::Cosine); - for row in 1..input.points { - let row_start = row * input.points; + + for row in 1..points { + let row_start = row * points; let row_norm = if uses_norms { F::splat(arch, norms[row]) } else { F::default(arch) }; - // SAFETY: `row < input.points == worst.len()`. + // SAFETY: `row < points == worst.len()` after validation. 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)) }; + let pair_dots = unsafe { F::load_simd(arch, dots.as_ptr().add(row_start + column)) }; let column_norms = if uses_norms { - // SAFETY: `column + F::LANES <= row < input.points == norms.len()`. + // SAFETY: the full chunk lies below `row <= norms.len()`. unsafe { F::load_simd(arch, norms.as_ptr().add(column)) } } else { F::default(arch) }; - let distances = pair_distances::(arch, input.metric, dots, row_norm, column_norms); + let distances = M::leaf_distance(arch, pair_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`. + // SAFETY: the full chunk lies below `row`, so it is inside `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 values = distances.to_array(); let values = values.as_ref(); @@ -401,51 +615,40 @@ fn process_pairs_simd_dynamic( row_bits &= row_bits - 1; let distance = values[lane]; if distance < row_worst { - row_worst = insert_slots( - &mut output[row * k..(row + 1) * k], - (column + lane) as u32, - distance, - ); + row_worst = output.insert(row, (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; - let new_worst = insert_slots( - &mut output[target * k..(target + 1) * k], - row as u32, - values[lane], - ); + let new_worst = output.insert(target, 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) }; + let dot = unsafe { *dots.get_unchecked(row_start + column) }; let (row_norm, column_norm) = if uses_norms { - // SAFETY: `column < row < input.points == norms.len()`. + // SAFETY: `column < row < points == norms.len()`. (norms[row], unsafe { *norms.get_unchecked(column) }) } else { (0.0, 0.0) }; - let distance = pair_distance(input.metric, dot, row_norm, column_norm); + let distance = M::leaf_distance_scalar(dot, row_norm, column_norm); if distance < row_worst { - row_worst = - insert_slots(&mut output[row * k..(row + 1) * k], column as u32, distance); + row_worst = output.insert(row, column as u32, distance); } // SAFETY: `column < row < worst.len()`. let column_worst = unsafe { *worst_ptr.add(column) }; if distance < column_worst { - let new_worst = insert_slots( - &mut output[column * k..(column + 1) * k], - row as u32, - distance, - ); + let new_worst = output.insert(column, row as u32, distance); // SAFETY: `column < row < worst.len()`. unsafe { *worst_ptr.add(column) = new_worst }; } @@ -454,144 +657,53 @@ fn process_pairs_simd_dynamic( // SAFETY: `row < worst.len()`. unsafe { *worst_ptr.add(row) = row_worst }; } + + debug_assert_eq!(output.len(), points); } -/// 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, +#[cfg(test)] +fn process_pairs_scalar( 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 worst_ptr = worst.as_mut_ptr(); - let uses_norms = METRIC == L2 || METRIC == COSINE; - for row in 1..input.points { - let row_start = row * input.points; - let row_norm = if uses_norms { - F::splat(arch, norms[row]) - } else { - F::default(arch) - }; - // 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)) }; - let column_norms = if uses_norms { - // SAFETY: `column + F::LANES <= row < input.points == norms.len()`. - unsafe { F::load_simd(arch, norms.as_ptr().add(column)) } - } else { - F::default(arch) - }; - 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 values = distances.to_array(); - let values = values.as_ref(); - 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 { - row_worst = insert_fixed::( - &mut output[row * SLOTS..(row + 1) * 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; - let new_worst = insert_fixed::( - &mut output[target * SLOTS..(target + 1) * 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) }; +) { + let points = input.dots.nrows(); + let uses_norms = M::LEAF_SCALE.is_some(); + for row in 1..points { + for column in 0..row { let (row_norm, column_norm) = if uses_norms { - // SAFETY: `column < row < input.points == norms.len()`. - (norms[row], unsafe { *norms.get_unchecked(column) }) + (norms[row], norms[column]) } else { (0.0, 0.0) }; - let distance = pair_distance(metric::(), dot, row_norm, column_norm); - if distance < row_worst { - row_worst = insert_fixed::( - &mut output[row * SLOTS..(row + 1) * SLOTS], - column as u32, - distance, - ); - } - // SAFETY: `column < row < worst.len()`. - let column_worst = unsafe { *worst_ptr.add(column) }; - if distance < column_worst { - let new_worst = insert_fixed::( - &mut output[column * SLOTS..(column + 1) * SLOTS], - row as u32, - distance, - ); - // SAFETY: `column < row < worst.len()`. - unsafe { *worst_ptr.add(column) = new_worst }; - } - column += 1; + let distance = + M::leaf_distance_scalar(input.dots[(row, column)], row_norm, column_norm); + insert_scalar(output, worst, k, row, column as u32, distance); + insert_scalar(output, worst, k, column, row as u32, distance); } - // 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!(), +#[cfg(test)] +fn insert_scalar( + 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; } + worst[row] = insert_dynamic(&mut output[row * k..(row + 1) * k], position, distance); } /// Insert into a production row whose width is known at dispatch. #[inline(always)] -fn insert_fixed(row: &mut [LeafNeighbor], position: u32, distance: f32) -> f32 { - let row: &mut [LeafNeighbor; N] = row - .try_into() - .expect("validated fixed-width leaf output row"); +fn insert_fixed(row: &mut [LeafNeighbor; N], position: u32, distance: f32) -> f32 { let entry = LeafNeighbor::new(position, distance); match N { 1 => { @@ -628,9 +740,8 @@ fn insert_fixed(row: &mut [LeafNeighbor], position: u32, distanc } } -/// Insert into the uncommon run-time-width row (`k > 3`). #[inline(always)] -fn insert_slots(row: &mut [LeafNeighbor], position: u32, distance: f32) -> f32 { +fn insert_dynamic(row: &mut [LeafNeighbor], position: u32, distance: f32) -> f32 { let last = row.len() - 1; row[last] = LeafNeighbor::new(position, distance); let mut index = last; @@ -641,103 +752,164 @@ fn insert_slots(row: &mut [LeafNeighbor], position: u32, distance: f32) -> f32 { row[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) - } - } -} +#[cfg(test)] +mod tests { + use crate::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; -#[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 { + 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 { - distance + 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; } } - Metric::CosineNormalized => { - let distance = 1.0 - dot; - if distance < 0.0 { - 0.0 - } else { - distance - } + dots + } + + fn input(dots: &[f32], points: usize) -> LeafTopK<'_> { + LeafTopK { + dots: MatrixView::try_from(dots, points, points).unwrap(), } - 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 + } + + fn scalar(input: LeafTopK<'_>, k: usize, output: &mut [LeafNeighbor]) { + let points = input.dots.nrows(); + let norms: Vec<_> = (0..points) + .map(|row| M::LEAF_SCALE.transform(input.dots[(row, row)])) + .collect(); + let mut worst = vec![f32::INFINITY; points]; + process_pairs_scalar::(input, k, output, &norms, &mut worst); + } + + fn scalar_for_metric( + metric: Metric, + input: LeafTopK<'_>, + k: usize, + output: &mut [LeafNeighbor], + ) { + match metric { + Metric::L2 => scalar::(input, k, output), + Metric::Cosine => scalar::(input, k, output), + Metric::CosineNormalized => scalar::(input, k, output), + Metric::InnerProduct => scalar::(input, k, output), + } + } + + fn assert_scalar_reference_matches_prepared_dispatch(metric: Metric) { + // Point count controls SIMD chunking. Cover both sides of 4-, 8-, and + // 16-lane boundaries, then the boundary around a second 16-lane chunk. + for points in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let dots = dots(metric, points); + let input = input(&dots, points); + for requested_k in [1, 2, 3, 4] { + let k = requested_k.min(points - 1); + let kernel = LeafKernel::new(metric, requested_k); + let mut expected = vec![LeafNeighbor::default(); points * k]; + kernel + .nearest_neighbors( + input, + MutMatrixView::try_from(expected.as_mut_slice(), points, k).unwrap(), + &mut LeafTopKWorkspace::new(), + ) + .unwrap(); + + let mut actual = vec![LeafNeighbor::default(); points * k]; + scalar_for_metric(metric, input, k, &mut actual); + + assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); } } } -} -#[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; + #[test] + fn l2_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_prepared_dispatch(Metric::L2); } - 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; + #[test] + fn cosine_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_prepared_dispatch(Metric::Cosine); } - worst[row] = row_output[k - 1].distance; -} -#[cfg(test)] -mod tests; + #[test] + fn normalized_cosine_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_prepared_dispatch(Metric::CosineNormalized); + } + + #[test] + fn inner_product_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_prepared_dispatch(Metric::InnerProduct); + } + + #[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_scalar(&mut output, &mut worst, 4, 0, position, distance); + } + insert_scalar(&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 output_length_clamps_to_non_self_neighbors() { + assert_eq!(leaf_output_len(0, 3).unwrap(), 0); + assert_eq!(leaf_output_len(1, 3).unwrap(), 0); + assert_eq!(leaf_output_len(4, 9).unwrap(), 12); + #[cfg(target_pointer_width = "64")] + assert_eq!( + leaf_output_len(u32::MAX as usize + 1, 1), + Err(LeafKernelError::TooManyPoints(u32::MAX as usize + 1)) + ); + } + + #[test] + fn matrix_area_overflow_is_rejected_before_kernel_access() { + assert_eq!( + checked_area("leaf dot-product matrix", usize::MAX, 2), + Err(LeafKernelError::ShapeOverflow { + buffer: "leaf dot-product matrix", + rows: usize::MAX, + cols: 2, + }) + ); + } + + #[test] + fn workspace_can_shrink_and_grow_between_calls() { + let kernel = LeafKernel::new(Metric::L2, 2); + 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]; + kernel + .nearest_neighbors( + input(&dots, points), + MutMatrixView::try_from(output.as_mut_slice(), points, 2).unwrap(), + &mut workspace, + ) + .unwrap(); + assert!(output.iter().all(|neighbor| neighbor.position != u32::MAX)); + } + } +} diff --git a/diskann-pipnn/src/leaf_kernel/tests.rs b/diskann-pipnn/src/leaf_kernel/tests.rs deleted file mode 100644 index 1bef128ab..000000000 --- a/diskann-pipnn/src/leaf_kernel/tests.rs +++ /dev/null @@ -1,144 +0,0 @@ -/* - * 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() -} - -fn assert_scalar_reference_matches_runtime_dispatch(metric: Metric) { - // Point count, rather than source-vector dimension, controls this kernel's - // SIMD boundaries. Cover lane-1/lane/lane+1 for 4-, 8-, and 16-lane - // backends, then the boundary around a second 16-lane chunk. - for points in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let dots = dots(metric, points); - let input = LeafTopK { - dots: &dots, - points, - metric, - }; - for requested_k in [1, 2, 3, 4] { - let k = requested_k.min(points - 1); - let mut expected = vec![LeafNeighbor::default(); points * k]; - nearest_leaf_neighbors( - input, - requested_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={requested_k}"); - } - } -} - -#[test] -fn l2_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_runtime_dispatch(Metric::L2); -} - -#[test] -fn cosine_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_runtime_dispatch(Metric::Cosine); -} - -#[test] -fn normalized_cosine_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_runtime_dispatch(Metric::CosineNormalized); -} - -#[test] -fn inner_product_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_runtime_dispatch(Metric::InnerProduct); -} - -#[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 output_length_clamps_to_non_self_neighbors() { - assert_eq!(leaf_output_len(0, 3).unwrap(), 0); - assert_eq!(leaf_output_len(1, 3).unwrap(), 0); - assert_eq!(leaf_output_len(4, 9).unwrap(), 12); - assert_eq!( - leaf_output_len(u32::MAX as usize + 1, 1), - Err(LeafKernelError::TooManyPoints(u32::MAX as usize + 1)) - ); -} - -#[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 index 17a5dd708..ad08807e0 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -10,13 +10,18 @@ //! of those stages while callers retain dataset storage, GEMM workspaces, graph //! policy, and scheduling: //! -//! - [`partition_kernel`] converts a point-by-leader dot-product tile into the -//! nearest leader positions for each point. -//! - [`leaf_kernel`] scans a leaf's lower-triangular dot-product matrix once and -//! retains nearest non-self neighbors for both endpoints. +//! - [`partition_kernel::PartitionKernel`] converts point-by-leader dot-product +//! tiles into nearest leader positions. +//! - [`leaf_kernel::LeafKernel`] scans each leaf's lower-triangular dot-product +//! matrix once and retains nearest non-self neighbors for both endpoints. //! -//! Both modules validate slice shapes before dispatch and use `diskann-wide` for -//! architecture selection; PiPNN does not detect or name instruction sets. +//! Callers prepare these small handles once per build metric (and leaf `k`) and +//! reuse them across stripes or leaves. Preparation uses `diskann-wide` to select +//! the runtime architecture and returns a direct function pointer; repeated calls +//! do not repeat ISA or metric dispatch. PiPNN itself never names instruction +//! sets. + +mod kernel_metric; pub mod leaf_kernel; pub mod partition_kernel; diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index 8122244f9..80d2a02e8 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -3,22 +3,29 @@ * Licensed under the MIT license. */ -//! Distance and top-k kernel for partition assignment. +//! Prepared distance and top-k kernels for partition assignment. //! -//! The caller gathers a point stripe and a leader matrix, then computes the -//! row-major `points · leadersᵀ` tile with GEMM. This module performs the second -//! half of assignment: convert each dot product to the configured metric and -//! retain only the nearest leader positions. +//! The caller computes a row-major `points · leadersᵀ` tile with GEMM, then +//! passes it to a [`PartitionKernel`] prepared once for the build metric. Kernel +//! preparation selects the runtime architecture and concrete metric type once; +//! repeated stripes call a direct `diskann-wide` function pointer with no ISA or +//! metric branch in the row loop. //! -//! L2 deliberately omits the point norm because it adds the same constant to -//! every leader in one row and cannot change their order. Cosine still needs a -//! point scale because it divides each dot product. The fixed 16-entry tracker -//! bounds stack use and matches the configuration fanout limit. SIMD chunks and -//! scalar tails feed the same insertion routine; NaNs are ignored and equal -//! distances keep the first leader encountered. +//! L2 deliberately omits the point norm because it is constant across every +//! leader in one row. Cosine consumes squared point norms and leader norms. NaN +//! distances are not rankable, and equal distances retain leader scan order. +use std::marker::PhantomData; + +use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; -use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector}; +use diskann_wide::{ + arch::{self, Dispatched2, FTarget2}, + lifetime::AddLifetime, + Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, +}; + +use crate::kernel_metric::{erase_metric, EraseMetric, KernelMetric, ScaleKind}; /// Maximum number of leaders retained for one point. /// @@ -29,37 +36,38 @@ pub const MAX_PARTITION_FANOUT: usize = 16; type TopK = [(u32, f32); MAX_PARTITION_FANOUT]; -/// One row-major point-by-leader dot-product tile and its normalization terms. -/// -/// The scale slices are deliberately metric-specific: -/// -/// | metric | `row_scales` | `leader_scales` | -/// |---|---|---| -/// | [`Metric::L2`] | empty | squared leader norms | -/// | [`Metric::Cosine`] | squared point norms | leader norms | -/// | [`Metric::CosineNormalized`] / [`Metric::InnerProduct`] | empty | empty | -/// -/// [`nearest_leaders`] validates every declared shape before dispatch. +/// Metric-specific normalization inputs for one partition tile. +#[derive(Clone, Copy, Debug)] +pub enum PartitionScales<'a> { + /// L2 needs only squared leader norms; the point norm cannot affect ranking. + L2 { + /// Squared norm for every leader column. + leader_squared_norms: &'a [f32], + }, + /// Unnormalized cosine needs squared point norms and leader norms. + Cosine { + /// Squared norm for every point row. + row_squared_norms: &'a [f32], + /// Norm for every leader column. + leader_norms: &'a [f32], + }, + /// Normalized cosine and inner product need no normalization inputs. + None, +} + +/// One row-major point-by-leader dot-product tile. #[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, - /// Metric-specific point normalization terms described in the type table. - pub row_scales: &'a [f32], - /// Metric-specific leader normalization terms described in the type table. - pub leader_scales: &'a [f32], - /// Distance metric used to rank leaders. - pub metric: Metric, + /// Point rows by leader columns. + pub dots: MatrixView<'a, f32>, + /// Normalization inputs matching the prepared metric. + pub scales: PartitionScales<'a>, } -/// Validation error returned by [`nearest_leaders`]. +/// Validation error returned by [`PartitionKernel::nearest_leaders`]. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub enum PartitionKernelError { - /// A declared matrix or output shape overflowed `usize`. + /// A declared matrix shape overflowed `usize`. #[error("{buffer} shape {rows} x {cols} overflows usize")] ShapeOverflow { /// Name of the buffer whose shape overflowed. @@ -69,16 +77,34 @@ pub enum PartitionKernelError { /// Declared column count. cols: usize, }, - /// A supplied slice did not match its declared shape. + /// The output matrix does not match the input row count. + #[error( + "invalid output shape: expected {expected_rows} rows, got {actual_rows} rows and {actual_cols} columns" + )] + InvalidOutputShape { + /// Required row count. + expected_rows: usize, + /// Supplied row count. + actual_rows: usize, + /// Supplied column count. + actual_cols: usize, + }, + /// A metric-specific scale slice has the wrong length. #[error("invalid {buffer} length: expected {expected}, got {actual}")] InvalidBufferLength { - /// Name of the invalid buffer. + /// Name of the invalid scale buffer. buffer: &'static str, /// Required length. expected: usize, /// Supplied length. actual: usize, }, + /// Scale inputs do not match the metric used to prepare the kernel. + #[error("partition scales do not match prepared {expected} metric")] + InvalidScales { + /// Expected scale layout. + expected: &'static str, + }, /// The requested fanout cannot be represented by the fixed top-k tracker. #[error( "invalid fanout {fanout}: must not exceed {leaders} leaders or kernel maximum {maximum}" @@ -104,66 +130,219 @@ pub enum PartitionKernelError { }, } -/// 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. +#[derive(Debug)] +struct PartitionInput; + +impl AddLifetime for PartitionInput { + type Of<'a> = PartitionTopK<'a>; +} + +#[derive(Debug)] +struct PartitionOutput; + +impl AddLifetime for PartitionOutput { + type Of<'a> = MutMatrixView<'a, u32>; +} + +type PartitionFn = Dispatched2, PartitionInput, PartitionOutput>; + +/// A partition kernel prepared for one metric and the current CPU. /// -/// 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(()); +/// Construct this once with [`PartitionKernel::new`] and reuse it for every +/// point stripe. The handle is a direct function pointer and is `Copy`, `Send`, +/// and `Sync`. +#[derive(Clone, Copy, Debug)] +pub struct PartitionKernel { + run: PartitionFn, +} + +impl PartitionKernel { + /// Prepare a partition kernel for `metric` and the current CPU. + pub fn new(metric: Metric) -> Self { + diskann_wide::arch::dispatch1_no_features(PreparePartition, metric) } - diskann_wide::arch::dispatch(PartitionKernel { - input, - fanout, - output, - }); - if let Some(row) = output - .chunks_exact(fanout) - .position(|leaders| leaders[fanout - 1] == u32::MAX) - { - return Err(PartitionKernelError::InsufficientRankableDistances { row, fanout }); + /// Select the nearest leader positions for every input row. + /// + /// `output.nrows()` must equal `input.dots.nrows()`; its column count is the + /// requested fanout. Results are ordered by ascending distance. For L2, the + /// score omits the point norm because it cannot affect within-row ranking. + pub fn nearest_leaders( + &self, + input: PartitionTopK<'_>, + output: MutMatrixView<'_, u32>, + ) -> Result<(), PartitionKernelError> { + self.run.call(input, output) } - Ok(()) } -fn validate( - input: PartitionTopK<'_>, - fanout: usize, - output: &[u32], -) -> Result<(), PartitionKernelError> { - if input.leaders > u32::MAX as usize { - return Err(PartitionKernelError::TooManyLeaders(input.leaders)); +struct PreparePartition; + +impl arch::Target1 for PreparePartition +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + fn run(self, arch: A, metric: Metric) -> PartitionKernel { + erase_metric(metric, BuildPartition(arch)) + } +} + +struct BuildPartition(A); + +impl EraseMetric for BuildPartition +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + type Output = PartitionKernel; + + fn erase(self) -> Self::Output { + PartitionKernel { + run: self.0.dispatch2::< + PartitionEntry, + Result<(), PartitionKernelError>, + PartitionInput, + PartitionOutput, + >(), + } + } +} + +struct PartitionEntry(PhantomData); + +impl FTarget2, PartitionTopK<'_>, MutMatrixView<'_, u32>> + for PartitionEntry +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, +{ + fn run( + arch: A, + input: PartitionTopK<'_>, + mut output: MutMatrixView<'_, u32>, + ) -> Result<(), PartitionKernelError> { + let scales = validate::(input, &output)?; + let fanout = output.ncols(); + if fanout == 0 || input.dots.nrows() == 0 { + return Ok(()); + } + + process_rows::(arch, input.dots, scales, fanout, output.as_mut_slice()); + if let Some(row) = output + .as_slice() + .chunks_exact(fanout) + .position(|leaders| leaders[fanout - 1] == u32::MAX) + { + return Err(PartitionKernelError::InsufficientRankableDistances { row, fanout }); + } + Ok(()) + } +} + +#[derive(Clone, Copy)] +struct ScaleSlices<'a> { + rows: &'a [f32], + leaders: &'a [f32], +} + +fn validate<'a, M: KernelMetric>( + input: PartitionTopK<'a>, + output: &MutMatrixView<'_, u32>, +) -> Result, PartitionKernelError> { + let rows = input.dots.nrows(); + let leaders = input.dots.ncols(); + let fanout = output.ncols(); + + let dots_len = checked_area("dot-product tile", rows, leaders)?; + check_length("dot-product tile", input.dots.as_slice().len(), dots_len)?; + let output_len = checked_area("output", output.nrows(), fanout)?; + check_length("output", output.as_slice().len(), output_len)?; + + if output.nrows() != rows { + return Err(PartitionKernelError::InvalidOutputShape { + expected_rows: rows, + actual_rows: output.nrows(), + actual_cols: output.ncols(), + }); } - if fanout > MAX_PARTITION_FANOUT || fanout > input.leaders { + if leaders > u32::MAX as usize { + return Err(PartitionKernelError::TooManyLeaders(leaders)); + } + if fanout > MAX_PARTITION_FANOUT || fanout > leaders { return Err(PartitionKernelError::InvalidFanout { fanout, - leaders: input.leaders, + 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), + let scales = match (M::METRIC, input.scales) { + ( + Metric::L2, + PartitionScales::L2 { + leader_squared_norms, + }, + ) => ScaleSlices { + rows: &[], + leaders: leader_squared_norms, + }, + ( + Metric::Cosine, + PartitionScales::Cosine { + row_squared_norms, + leader_norms, + }, + ) => ScaleSlices { + rows: row_squared_norms, + leaders: leader_norms, + }, + (Metric::CosineNormalized | Metric::InnerProduct, PartitionScales::None) => ScaleSlices { + rows: &[], + leaders: &[], + }, + (Metric::L2, _) => return Err(PartitionKernelError::InvalidScales { expected: "L2" }), + (Metric::Cosine, _) => { + return Err(PartitionKernelError::InvalidScales { expected: "cosine" }); + } + (Metric::CosineNormalized, _) => { + return Err(PartitionKernelError::InvalidScales { + expected: "normalized cosine", + }); + } + (Metric::InnerProduct, _) => { + return Err(PartitionKernelError::InvalidScales { + expected: "inner product", + }); + } }; - check_length("row scales", input.row_scales.len(), row_scales)?; - check_length("leader scales", input.leader_scales.len(), leader_scales) + + check_length( + "row scales", + scales.rows.len(), + expected_scale_len(M::PARTITION_ROW_SCALE, rows), + )?; + check_length( + "leader scales", + scales.leaders.len(), + expected_scale_len(M::PARTITION_LEADER_SCALE, leaders), + )?; + Ok(scales) +} + +const fn expected_scale_len(kind: ScaleKind, count: usize) -> usize { + if kind.is_some() { + count + } else { + 0 + } } fn checked_area( @@ -191,223 +370,102 @@ fn check_length( } } -struct PartitionKernel<'a, 'o> { - input: PartitionTopK<'a>, +fn process_rows( + arch: F::Arch, + dots: MatrixView<'_, f32>, + scales: ScaleSlices<'_>, 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>, + output: &mut [u32], +) where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + M: KernelMetric, + u64: From<<::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) + let leaders = dots.ncols(); + for (row, (dot_row, output_row)) in dots + .as_slice() + .chunks_exact(leaders) .zip(output.chunks_exact_mut(fanout)) .enumerate() { + let row_scale = if M::PARTITION_ROW_SCALE.is_some() { + M::PARTITION_ROW_SCALE.transform(scales.rows[row]) + } else { + 0.0 + }; + let row_scale_vector = F::splat(arch, row_scale); 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( + let full = leaders / F::LANES * F::LANES; + + for base in (0..full).step_by(F::LANES) { + // SAFETY: `base + F::LANES <= full <= dot_row.len()`. + let dots = unsafe { F::load_simd(arch, dot_row.as_ptr().add(base)) }; + let leader_scales = if M::PARTITION_LEADER_SCALE.is_some() { + // SAFETY: validation requires one leader scale per dot-product column. + unsafe { F::load_simd(arch, scales.leaders.as_ptr().add(base)) } + } else { + F::default(arch) + }; + insert_lanes( + M::partition_distance(arch, dots, row_scale_vector, leader_scales), + base, &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>, -{ - match input.metric { - Metric::L2 => process_rows(input, fanout, output, |_, dot_row, top| { - process_binary::( - arch, - dot_row, - input.leader_scales, - top, - fanout, - |dot, norm| F::splat(arch, -2.0).mul_add_simd(dot, norm), - |dot, norm| norm - 2.0 * dot, - ); - }), - Metric::CosineNormalized => process_rows(input, fanout, output, |_, dot_row, top| { - process_unary::(arch, dot_row, top, fanout, |dot| F::splat(arch, 1.0) - dot); - }), - Metric::InnerProduct => process_rows(input, fanout, output, |_, dot_row, top| { - process_unary::(arch, dot_row, top, fanout, |dot| F::default(arch) - dot); - }), - Metric::Cosine => process_rows(input, fanout, output, |row, dot_row, top| { - process_cosine::( - arch, - dot_row, - input.row_scales[row], - input.leader_scales, - top, + for (leader, &dot) in dot_row.iter().enumerate().skip(full) { + let leader_scale = if M::PARTITION_LEADER_SCALE.is_some() { + M::PARTITION_LEADER_SCALE.transform(scales.leaders[leader]) + } else { + 0.0 + }; + insert_topk( + &mut top, fanout, + leader as u32, + M::partition_distance_scalar(dot, row_scale, leader_scale), ); - }), + } + copy_ids(&top, output_row); } } -#[inline(always)] -fn process_rows( - input: PartitionTopK<'_>, +#[cfg(test)] +fn process_rows_scalar( + dots: MatrixView<'_, f32>, + scales: ScaleSlices<'_>, fanout: usize, output: &mut [u32], - mut process: impl FnMut(usize, &[f32], &mut TopK), ) { - for (row, (dot_row, output_row)) in input - .dots - .chunks_exact(input.leaders) + let leaders = dots.ncols(); + for (row, (dot_row, output_row)) in dots + .as_slice() + .chunks_exact(leaders) .zip(output.chunks_exact_mut(fanout)) .enumerate() { + let row_scale = if M::PARTITION_ROW_SCALE.is_some() { + M::PARTITION_ROW_SCALE.transform(scales.rows[row]) + } else { + 0.0 + }; let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; - process(row, dot_row, &mut top); - copy_ids(&top, output_row); - } -} - -#[inline(always)] -fn cosine_distance(row_norm_squared: f32, leader_norm: f32, dot: f32) -> f32 { - let row_norm = if row_norm_squared < f32::MIN_POSITIVE { - 0.0 - } else { - row_norm_squared.sqrt() - }; - let leader_norm = if leader_norm < f32::MIN_POSITIVE.sqrt() { - 0.0 - } else { - leader_norm - }; - if row_norm == 0.0 || leader_norm == 0.0 { - 1.0 - } else { - 1.0 - dot / (row_norm * leader_norm) - } -} - -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 = if row_norm_squared < f32::MIN_POSITIVE { - 0.0 - } else { - row_norm_squared.sqrt() - }; - let row_norm = F::splat(arch, row_norm); - let one = F::splat(arch, 1.0); - let minimum_norm = F::splat(arch, f32::MIN_POSITIVE.sqrt()); - process_binary::( - arch, - dots, - leader_norms, - top, - fanout, - |dot, leader_norm| { - let row_zero = row_norm.lt_simd(minimum_norm); - let leader_zero = leader_norm.lt_simd(minimum_norm); - let denominator = row_norm * leader_norm; - let safe_denominator = row_zero.select(one, leader_zero.select(one, denominator)); - let cosine = row_zero.select( - F::default(arch), - leader_zero.select(F::default(arch), dot / safe_denominator), + for (leader, &dot) in dot_row.iter().enumerate() { + let leader_scale = if M::PARTITION_LEADER_SCALE.is_some() { + M::PARTITION_LEADER_SCALE.transform(scales.leaders[leader]) + } else { + 0.0 + }; + insert_topk( + &mut top, + fanout, + leader as u32, + M::partition_distance_scalar(dot, row_scale, leader_scale), ); - one - cosine - }, - |dot, leader_norm| cosine_distance(row_norm_squared, leader_norm, dot), - ); -} - -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 value = transform(F::splat(arch, dot)).to_array(); - insert_topk(top, fanout, (full + offset) as u32, value.as_ref()[0]); - } -} - -fn process_binary( - arch: F::Arch, - dots: &[f32], - scales: &[f32], - top: &mut TopK, - fanout: usize, - transform: Transform, - scalar_transform: ScalarTransform, -) where - F: SIMDVector + SIMDFloat, - Transform: Fn(F, F) -> F, - ScalarTransform: Fn(f32, f32) -> f32, - 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 value = scalar_transform(dots[full + offset], scales[full + offset]); - insert_topk(top, fanout, (full + offset) as u32, value); + } + copy_ids(&top, output_row); } } @@ -432,17 +490,6 @@ where } } -#[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 => cosine_distance(row_scale, leader_scale, dot), - } -} - #[inline(always)] fn insert_topk(top: &mut TopK, fanout: usize, leader: u32, distance: f32) { let threshold = fanout - 1; @@ -465,4 +512,216 @@ fn copy_ids(top: &TopK, output: &mut [u32]) { } #[cfg(test)] -mod tests; +mod tests { + use crate::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; + + use super::*; + + fn data(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).powi(2)) + .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) + } + + fn input<'a>( + metric: Metric, + dots: &'a [f32], + rows: usize, + leaders: usize, + row_scales: &'a [f32], + leader_scales: &'a [f32], + ) -> PartitionTopK<'a> { + let scales = match metric { + Metric::L2 => PartitionScales::L2 { + leader_squared_norms: leader_scales, + }, + Metric::Cosine => PartitionScales::Cosine { + row_squared_norms: row_scales, + leader_norms: leader_scales, + }, + Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, + }; + PartitionTopK { + dots: MatrixView::try_from(dots, rows, leaders).unwrap(), + scales, + } + } + + fn scalar(input: PartitionTopK<'_>, fanout: usize, output: &mut [u32]) { + let scales = match input.scales { + PartitionScales::L2 { + leader_squared_norms, + } => ScaleSlices { + rows: &[], + leaders: leader_squared_norms, + }, + PartitionScales::Cosine { + row_squared_norms, + leader_norms, + } => ScaleSlices { + rows: row_squared_norms, + leaders: leader_norms, + }, + PartitionScales::None => ScaleSlices { + rows: &[], + leaders: &[], + }, + }; + process_rows_scalar::(input.dots, scales, fanout, output); + } + + fn scalar_for_metric( + metric: Metric, + input: PartitionTopK<'_>, + fanout: usize, + output: &mut [u32], + ) { + match metric { + Metric::L2 => scalar::(input, fanout, output), + Metric::Cosine => scalar::(input, fanout, output), + Metric::CosineNormalized => scalar::(input, fanout, output), + Metric::InnerProduct => scalar::(input, fanout, output), + } + } + + fn assert_scalar_reference_matches_prepared_dispatch(metric: Metric) { + // Leader count controls SIMD chunking. Exercise both sides of 4-, 8-, and + // 16-lane boundaries, then a second 16-lane chunk. + for leaders in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let (dots, row_scales, leader_scales) = data(metric, leaders); + let input = input(metric, &dots, 2, leaders, &row_scales, &leader_scales); + let kernel = PartitionKernel::new(metric); + for fanout in [1, 2, 6, MAX_PARTITION_FANOUT] { + if fanout > leaders { + continue; + } + let mut expected = vec![u32::MAX; 2 * fanout]; + kernel + .nearest_leaders( + input, + MutMatrixView::try_from(expected.as_mut_slice(), 2, fanout).unwrap(), + ) + .unwrap(); + + let mut actual = vec![u32::MAX; 2 * fanout]; + scalar_for_metric(metric, input, fanout, &mut actual); + assert_eq!( + actual, expected, + "{metric:?}, leaders={leaders}, k={fanout}" + ); + } + } + } + + #[test] + fn l2_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_prepared_dispatch(Metric::L2); + } + + #[test] + fn cosine_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_prepared_dispatch(Metric::Cosine); + } + + #[test] + fn normalized_cosine_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_prepared_dispatch(Metric::CosineNormalized); + } + + #[test] + fn inner_product_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { + assert_scalar_reference_matches_prepared_dispatch(Metric::InnerProduct); + } + + #[test] + fn scalar_distance_matches_metric_contract() { + assert_eq!(L2::partition_distance_scalar(2.0, 0.0, 9.0), 5.0); + assert_eq!( + CosineNormalized::partition_distance_scalar(0.25, 0.0, 0.0), + 0.75 + ); + assert_eq!(InnerProduct::partition_distance_scalar(3.0, 0.0, 0.0), -3.0); + assert_eq!(Cosine::partition_distance_scalar(4.0, 2.0, 4.0), 0.5); + assert_eq!(Cosine::partition_distance_scalar(4.0, 0.0, 4.0), 1.0); + assert!(Cosine::partition_distance_scalar(1.0, f32::NAN, 1.0).is_nan()); + } + + #[test] + fn cosine_special_norms_match_scalar_and_prepared_dispatch() { + let leaders = 17; + let dots = vec![1.0; 4 * leaders]; + let row_scales = [0.0, f32::MIN_POSITIVE / 2.0, f32::MIN_POSITIVE, f32::NAN]; + let mut leader_scales = vec![1.0; leaders]; + leader_scales[..4].copy_from_slice(&[ + 0.0, + f32::MIN_POSITIVE.sqrt() / 2.0, + f32::MIN_POSITIVE.sqrt(), + f32::NAN, + ]); + let input = input( + Metric::Cosine, + &dots, + row_scales.len(), + leaders, + &row_scales, + &leader_scales, + ); + let mut expected = vec![u32::MAX; row_scales.len() * 2]; + scalar::(input, 2, &mut expected); + let mut actual = vec![u32::MAX; row_scales.len() * 2]; + PartitionKernel::new(Metric::Cosine) + .nearest_leaders( + input, + MutMatrixView::try_from(actual.as_mut_slice(), row_scales.len(), 2).unwrap(), + ) + .unwrap(); + + assert_eq!(actual, expected); + assert_eq!(&actual[..4], &[0, 1, 0, 1]); + assert_eq!(&actual[6..], &[0, 1]); + } + + #[test] + fn matrix_area_overflow_is_rejected_before_kernel_access() { + assert_eq!( + checked_area("dot-product tile", usize::MAX, 2), + Err(PartitionKernelError::ShapeOverflow { + buffer: "dot-product tile", + rows: usize::MAX, + cols: 2, + }) + ); + } + + #[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/src/partition_kernel/tests.rs b/diskann-pipnn/src/partition_kernel/tests.rs deleted file mode 100644 index 7da20fd58..000000000 --- a/diskann-pipnn/src/partition_kernel/tests.rs +++ /dev/null @@ -1,148 +0,0 @@ -/* - * 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).powi(2)) - .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) -} - -fn assert_scalar_reference_matches_runtime_dispatch(metric: Metric) { - // Leader count controls SIMD chunking. Exercise the tail on both sides of - // 4-, 8-, and 16-lane boundaries, then a second 16-lane chunk. - for leaders in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - 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, MAX_PARTITION_FANOUT] { - if fanout > leaders { - continue; - } - 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 l2_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_runtime_dispatch(Metric::L2); -} - -#[test] -fn cosine_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_runtime_dispatch(Metric::Cosine); -} - -#[test] -fn normalized_cosine_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_runtime_dispatch(Metric::CosineNormalized); -} - -#[test] -fn inner_product_scalar_reference_matches_runtime_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_runtime_dispatch(Metric::InnerProduct); -} - -#[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); - assert_eq!( - distance(Metric::Cosine, 1.0, f32::MIN_POSITIVE / 2.0, 1.0), - 1.0 - ); - assert_eq!( - distance( - Metric::Cosine, - f32::MIN_POSITIVE, - f32::MIN_POSITIVE, - f32::MIN_POSITIVE.sqrt() - ), - 0.0 - ); - assert!(distance(Metric::Cosine, 1.0, f32::NAN, 1.0).is_nan()); -} - -#[test] -fn cosine_special_norms_match_scalar_and_runtime_dispatch() { - let leaders = 17; - let dots = vec![1.0; 4 * leaders]; - let row_scales = [0.0, f32::MIN_POSITIVE / 2.0, f32::MIN_POSITIVE, f32::NAN]; - let mut leader_scales = vec![1.0; leaders]; - leader_scales[..4].copy_from_slice(&[ - 0.0, - f32::MIN_POSITIVE.sqrt() / 2.0, - f32::MIN_POSITIVE.sqrt(), - f32::NAN, - ]); - let input = PartitionTopK { - dots: &dots, - rows: row_scales.len(), - leaders, - row_scales: &row_scales, - leader_scales: &leader_scales, - metric: Metric::Cosine, - }; - let mut expected = vec![u32::MAX; input.rows * 2]; - process_rows_scalar(input, 2, &mut expected); - let mut actual = vec![u32::MAX; input.rows * 2]; - nearest_leaders(input, 2, &mut actual).unwrap(); - - assert_eq!(actual, expected); - assert_eq!(&actual[..4], &[0, 1, 0, 1]); - assert_eq!(&actual[6..], &[0, 1]); -} - -#[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_api.rs similarity index 56% rename from diskann-pipnn/tests/leaf_kernel.rs rename to diskann-pipnn/tests/leaf_kernel_api.rs index 5971f4cb6..f82bfe11d 100644 --- a/diskann-pipnn/tests/leaf_kernel.rs +++ b/diskann-pipnn/tests/leaf_kernel_api.rs @@ -3,11 +3,13 @@ * Licensed under the MIT license. */ +use std::cmp::Ordering; + use diskann_pipnn::leaf_kernel::{ - nearest_leaf_neighbors, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, + leaf_output_len, LeafKernel, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, }; +use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; -use std::cmp::Ordering; const SIMD_BOUNDARY_POINTS: [usize; 9] = [7, 8, 9, 15, 16, 17, 64, 256, 512]; const ZERO_NORM_POSITION: usize = 0; @@ -45,17 +47,23 @@ fn differential_input(metric: Metric, points: usize) -> Vec { 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]; +fn input(dots: &[f32], points: usize) -> LeafTopK<'_> { + LeafTopK { + dots: MatrixView::try_from(dots, points, points).unwrap(), + } +} + +fn reference(dots: &[f32], points: usize, requested_k: usize, metric: Metric) -> Vec { + let k = requested_k.min(points.saturating_sub(1)); + let mut output = vec![LeafNeighbor::default(); points * k]; if k == 0 { return output; } - let norms: Vec<_> = (0..input.points) + let norms: Vec<_> = (0..points) .map(|row| { - let diagonal = input.dots[row * input.points + row]; - if input.metric == Metric::Cosine { + let diagonal = dots[row * points + row]; + if metric == Metric::Cosine { if diagonal < f32::MIN_POSITIVE { 0.0 } else { @@ -67,9 +75,9 @@ fn reference(input: LeafTopK<'_>, requested_k: usize) -> Vec { }) .collect(); - for row in 0..input.points { - let mut candidates = Vec::with_capacity(input.points - 1); - for position in 0..input.points { + for row in 0..points { + let mut candidates = Vec::with_capacity(points - 1); + for position in 0..points { if position == row { continue; } @@ -78,15 +86,9 @@ fn reference(input: LeafTopK<'_>, requested_k: usize) -> Vec { } 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 { + let dot = dots[lower_row * points + lower_column]; + let clamp = |distance: f32| if distance < 0.0 { 0.0 } else { distance }; + let distance = match metric { Metric::L2 => clamp(norms[row] + norms[position] - 2.0 * dot), Metric::CosineNormalized => clamp(1.0 - dot), Metric::InnerProduct => -dot, @@ -115,54 +117,39 @@ fn reference(input: LeafTopK<'_>, requested_k: usize) -> Vec { output } +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 returned_k = LeafKernel::new(metric, k) + .nearest_neighbors( + input(dots, points), + MutMatrixView::try_from(output.as_mut_slice(), points, actual_k).unwrap(), + &mut LeafTopKWorkspace::new(), + ) + .unwrap(); + assert_eq!(returned_k, actual_k); + (returned_k, output) +} + #[test] -fn dispatch_matches_reference_across_simd_width_boundaries() { +fn prepared_dispatch_matches_reference_across_simd_width_boundaries() { for metric in [ Metric::L2, Metric::Cosine, Metric::CosineNormalized, Metric::InnerProduct, ] { - // Straddle the 8- and 16-lane boundaries, then cover production leaf sizes. for points in SIMD_BOUNDARY_POINTS { 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(); + let expected = reference(&dots, points, requested_k, metric); + let actual = run(&dots, points, requested_k, metric).1; 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] @@ -173,10 +160,8 @@ fn l2_scans_only_the_lower_triangle_and_breaks_ties_by_position() { 0.0, 1.0, 1.0, 2.0, ]; - let (_, output) = run(&dots, 4, 2, Metric::L2); - assert_eq!( - output, + run(&dots, 4, 2, Metric::L2).1, [ LeafNeighbor::new(1, 1.0), LeafNeighbor::new(2, 1.0), @@ -198,17 +183,17 @@ fn supports_every_leaf_metric() { 0.0, 1.0, 77.0, -1.0, 0.5, 1.0, ]; - - let cases = [ + for (metric, expected) in [ (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(); + ] { + let positions: Vec<_> = run(&dots, 3, 1, metric) + .1 + .iter() + .map(|neighbor| neighbor.position) + .collect(); assert_eq!(positions, expected, "metric {metric:?}"); } } @@ -222,8 +207,7 @@ fn cosine_treats_zero_norm_as_zero_similarity() { 0.0, 0.0, 1.0, ]; - let (_, output) = run(&dots, 3, 2, Metric::Cosine); - + let output = run(&dots, 3, 2, Metric::Cosine).1; assert_eq!(output[0], LeafNeighbor::new(1, 1.0)); assert_eq!(output[1], LeafNeighbor::new(2, 1.0)); } @@ -231,10 +215,7 @@ fn cosine_treats_zero_norm_as_zero_similarity() { #[test] fn preserves_pipnn_metric_edge_semantics() { #[rustfmt::skip] - let out_of_range = [ - 1.0, 0.0, - 2.0, 1.0, - ]; + 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, @@ -243,26 +224,13 @@ fn preserves_pipnn_metric_edge_semantics() { 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, - ]; + 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, - ]; + let subnormal = [f32::MIN_POSITIVE / 2.0, 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, - ]; + let minimum_normal = [f32::MIN_POSITIVE, 0.0, f32::MIN_POSITIVE.sqrt(), 1.0]; assert_eq!( run(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, 0.0 @@ -276,7 +244,6 @@ fn finite_max_distance_fills_the_final_simd_slot() { 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], @@ -299,7 +266,7 @@ fn every_metric_ignores_nan_pairs() { Metric::CosineNormalized, Metric::InnerProduct, ] { - let (_, output) = run(&dots, 3, 1, metric); + let output = run(&dots, 3, 1, metric).1; assert_eq!(output[0].position, 2, "metric {metric:?}"); assert_eq!(output[1].position, 2, "metric {metric:?}"); } @@ -307,31 +274,21 @@ fn every_metric_ignores_nan_pairs() { #[test] fn rejects_incomplete_neighbor_rows() { - #[rustfmt::skip] - let dots = [ - 1.0, 0.0, - f32::NAN, 1.0, - ]; + 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(); + let error = LeafKernel::new(Metric::L2, 1) + .nearest_neighbors( + input(&dots, 2), + MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), + &mut LeafTopKWorkspace::new(), + ) + .unwrap_err(); assert_eq!( error, LeafKernelError::InsufficientRankableNeighbors { row: 0, - neighbors: 1, + neighbors: 1 } ); } @@ -344,11 +301,9 @@ fn clamps_k_to_available_non_self_neighbors() { 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() @@ -358,110 +313,58 @@ fn clamps_k_to_available_non_self_neighbors() { #[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 - ); + for (dots, points, k, metric) in [ + (&[][..], 0, 2, Metric::L2), + (&[4.0][..], 1, 2, Metric::Cosine), + (&[1.0, 0.0, 0.0, 1.0][..], 2, 0, Metric::InnerProduct), + ] { + assert_eq!(run(dots, points, k, metric).0, 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(); +fn rejects_non_square_input_and_wrong_output_shape() { + let dots = [0.0; 6]; + let non_square = LeafTopK { + dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), + }; + let mut output = [LeafNeighbor::default(); 2]; + let kernel = LeafKernel::new(Metric::L2, 1); assert_eq!( - error, - LeafKernelError::InvalidBufferLength { - buffer: "lower dot-product matrix", - expected: 9, - actual: 8, - } + kernel.nearest_neighbors( + non_square, + MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), + &mut LeafTopKWorkspace::new(), + ), + Err(LeafKernelError::NonSquareDots { rows: 2, cols: 3 }) ); - let error = nearest_leaf_neighbors( - LeafTopK { - dots: &[0.0; 9], - points: 3, - metric: Metric::L2, - }, - 2, - &mut [LeafNeighbor::default(); 5], - &mut workspace, - ) - .unwrap_err(); + let square = [0.0; 9]; + let mut wrong = [LeafNeighbor::default(); 3]; assert_eq!( - error, - LeafKernelError::InvalidBufferLength { - buffer: "output", - expected: 6, - actual: 5, - } + LeafKernel::new(Metric::L2, 2).nearest_neighbors( + input(&square, 3), + MutMatrixView::try_from(&mut wrong[..], 3, 1).unwrap(), + &mut LeafTopKWorkspace::new(), + ), + Err(LeafKernelError::InvalidOutputShape { + expected_rows: 3, + expected_cols: 2, + actual_rows: 3, + actual_cols: 1, + }) ); } -#[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); + let output = run(&dots, points, 1, Metric::Cosine).1; for (row, neighbor) in output.iter().enumerate().skip(1) { assert_eq!( *neighbor, @@ -472,31 +375,10 @@ fn cosine_zero_norm_masks_nan_norm_at_simd_boundaries() { } } -#[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(); - +fn output_length_rejects_unrepresentable_point_count() { assert_eq!( - error, - LeafKernelError::InvalidBufferLength { - buffer: "lower dot-product matrix", - expected, - actual: 0, - } + leaf_output_len(usize::MAX, 1), + Err(LeafKernelError::TooManyPoints(usize::MAX)) ); } diff --git a/diskann-pipnn/tests/partition_kernel.rs b/diskann-pipnn/tests/partition_kernel.rs deleted file mode 100644 index bb90a8b9f..000000000 --- a/diskann-pipnn/tests/partition_kernel.rs +++ /dev/null @@ -1,442 +0,0 @@ -/* - * 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-pipnn/tests/partition_kernel_api.rs b/diskann-pipnn/tests/partition_kernel_api.rs new file mode 100644 index 000000000..e5318b8d1 --- /dev/null +++ b/diskann-pipnn/tests/partition_kernel_api.rs @@ -0,0 +1,358 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_pipnn::partition_kernel::{ + PartitionKernel, PartitionKernelError, PartitionScales, PartitionTopK, MAX_PARTITION_FANOUT, +}; +use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_vector::distance::Metric; + +fn input<'a>( + metric: Metric, + dots: &'a [f32], + rows: usize, + leaders: usize, + row_scales: &'a [f32], + leader_scales: &'a [f32], +) -> PartitionTopK<'a> { + let scales = match metric { + Metric::L2 => PartitionScales::L2 { + leader_squared_norms: leader_scales, + }, + Metric::Cosine => PartitionScales::Cosine { + row_squared_norms: row_scales, + leader_norms: leader_scales, + }, + Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, + }; + PartitionTopK { + dots: MatrixView::try_from(dots, rows, leaders).unwrap(), + scales, + } +} + +fn reference(input: PartitionTopK<'_>, fanout: usize, metric: Metric) -> Vec { + let rows = input.dots.nrows(); + let leaders = input.dots.ncols(); + let (row_scales, leader_scales) = match input.scales { + PartitionScales::L2 { + leader_squared_norms, + } => (&[][..], leader_squared_norms), + PartitionScales::Cosine { + row_squared_norms, + leader_norms, + } => (row_squared_norms, leader_norms), + PartitionScales::None => (&[][..], &[][..]), + }; + let mut output = vec![u32::MAX; rows * fanout]; + for (row, (dots, output)) in input + .dots + .as_slice() + .chunks_exact(leaders) + .zip(output.chunks_exact_mut(fanout)) + .enumerate() + { + let row_scale = row_scales.get(row).copied().unwrap_or(0.0); + let mut candidates: Vec<_> = dots + .iter() + .enumerate() + .filter_map(|(leader, &dot)| { + let leader_scale = leader_scales.get(leader).copied().unwrap_or(0.0); + let distance = match metric { + Metric::L2 => leader_scale - 2.0 * dot, + Metric::CosineNormalized => 1.0 - dot, + Metric::InnerProduct => -dot, + Metric::Cosine => { + let row_norm = if row_scale < f32::MIN_POSITIVE { + 0.0 + } else { + row_scale.sqrt() + }; + 1.0 - if row_norm == 0.0 || leader_scale == 0.0 { + 0.0 + } else { + dot / (row_norm * leader_scale) + } + } + }; + (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) +} + +fn run( + metric: Metric, + input: PartitionTopK<'_>, + fanout: usize, +) -> Result, PartitionKernelError> { + let mut output = vec![u32::MAX; input.dots.nrows() * fanout]; + PartitionKernel::new(metric).nearest_leaders( + input, + MutMatrixView::try_from(output.as_mut_slice(), input.dots.nrows(), fanout).unwrap(), + )?; + Ok(output) +} + +#[test] +fn prepared_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); + let input = input(metric, &dots, 2, leaders, &row_scales, &leader_scales); + for fanout in [1, 2, 16] { + if fanout >= leaders { + continue; + } + assert_eq!( + run(metric, input, fanout).unwrap(), + reference(input, fanout, metric), + "{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 norms = [0.0, 1.0, 4.0, 9.0]; + + assert_eq!( + run(Metric::L2, input(Metric::L2, &dots, 2, 4, &[], &norms), 2).unwrap(), + [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, + ]; + for (metric, rows, leaders, expected) in [ + (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]), + ] { + assert_eq!( + run(metric, input(metric, &dots, 2, 3, rows, leaders), 2).unwrap(), + expected, + "metric {metric:?}" + ); + } +} + +#[test] +fn cosine_treats_a_zero_norm_as_zero_similarity() { + assert_eq!( + run( + Metric::Cosine, + input(Metric::Cosine, &[100.0, -100.0], 1, 2, &[0.0], &[1.0, 1.0]), + 2, + ) + .unwrap(), + [0, 1] + ); +} + +#[test] +fn finite_max_distance_fills_the_final_simd_slot() { + let mut dots = [0.0; 8]; + dots[7] = -f32::MAX; + assert_eq!( + run( + Metric::InnerProduct, + input(Metric::InnerProduct, &dots, 1, 8, &[], &[]), + 8 + ) + .unwrap(), + [0, 1, 2, 3, 4, 5, 6, 7] + ); +} + +#[test] +fn ignores_nan_distances_without_displacing_finite_leaders() { + assert_eq!( + run( + Metric::InnerProduct, + input(Metric::InnerProduct, &[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), + 2, + ) + .unwrap(), + [1, 2] + ); +} + +#[test] +fn rejects_rows_with_too_few_rankable_distances() { + assert_eq!( + run( + Metric::InnerProduct, + input(Metric::InnerProduct, &[f32::NAN, 3.0], 1, 2, &[], &[]), + 2, + ), + Err(PartitionKernelError::InsufficientRankableDistances { row: 0, fanout: 2 }) + ); +} + +#[test] +fn accepts_empty_rows_zero_fanout_and_largest_leader_id() { + run( + Metric::InnerProduct, + input(Metric::InnerProduct, &[], 0, 3, &[], &[]), + 2, + ) + .unwrap(); + run( + Metric::InnerProduct, + input(Metric::InnerProduct, &[1.0, 2.0, 3.0], 1, 3, &[], &[]), + 0, + ) + .unwrap(); + run( + Metric::InnerProduct, + input(Metric::InnerProduct, &[], 0, u32::MAX as usize, &[], &[]), + 0, + ) + .unwrap(); + + #[cfg(target_pointer_width = "64")] + assert_eq!( + run( + Metric::InnerProduct, + input( + Metric::InnerProduct, + &[], + 0, + u32::MAX as usize + 1, + &[], + &[], + ), + 0, + ), + Err(PartitionKernelError::TooManyLeaders(u32::MAX as usize + 1)) + ); +} + +#[test] +fn rejects_wrong_output_scales_and_fanout() { + let dots = [0.0; 6]; + let valid_input = input(Metric::InnerProduct, &dots, 2, 3, &[], &[]); + let mut wrong_output = [u32::MAX; 3]; + assert_eq!( + PartitionKernel::new(Metric::InnerProduct).nearest_leaders( + valid_input, + MutMatrixView::try_from(&mut wrong_output[..], 1, 3).unwrap(), + ), + Err(PartitionKernelError::InvalidOutputShape { + expected_rows: 2, + actual_rows: 1, + actual_cols: 3, + }) + ); + + let wrong_scales = PartitionTopK { + dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), + scales: PartitionScales::None, + }; + assert_eq!( + run(Metric::L2, wrong_scales, 2), + Err(PartitionKernelError::InvalidScales { expected: "L2" }) + ); + + assert_eq!( + run(Metric::InnerProduct, valid_input, MAX_PARTITION_FANOUT + 1,), + Err(PartitionKernelError::InvalidFanout { + fanout: MAX_PARTITION_FANOUT + 1, + leaders: 3, + maximum: MAX_PARTITION_FANOUT, + }) + ); + + let one = [0.0]; + assert_eq!( + run( + Metric::InnerProduct, + input(Metric::InnerProduct, &one, 1, 1, &[], &[]), + 2, + ), + Err(PartitionKernelError::InvalidFanout { + fanout: 2, + leaders: 1, + maximum: MAX_PARTITION_FANOUT, + }) + ); +} From 4b731c5ed049c42d7fd1688436b26e43d7dd5d6c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:58:42 +0000 Subject: [PATCH 14/26] docs(pipnn): describe full crate scope --- diskann-pipnn/src/lib.rs | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index ad08807e0..9cfc3e25e 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -3,12 +3,15 @@ * Licensed under the MIT license. */ -//! Numerical kernels used by PiPNN graph construction. +//! Provider-independent PiPNN graph construction. //! -//! PiPNN first partitions points around sampled leaders, then builds local -//! neighbor candidates inside each leaf. This crate owns the numerical seams -//! of those stages while callers retain dataset storage, GEMM workspaces, graph -//! policy, and scheduling: +//! The crate owns overlapping partition generation, leaf-local nearest-neighbor +//! construction, candidate merging, and optional graph-degree finalization. The +//! caller supplies contiguous data, DiskANN graph policy, and the Rayon pool. +//! Providers, start/frozen points, quantization, persistence, and search remain +//! outside this algorithm seam. +//! +//! Numerical kernels include: //! //! - [`partition_kernel::PartitionKernel`] converts point-by-leader dot-product //! tiles into nearest leader positions. From 4e425baeea17535ba13f52e6f67a6a8429e19f4f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 11:14:01 +0000 Subject: [PATCH 15/26] docs(pipnn): map dispatch hot paths --- diskann-pipnn/src/kernel_metric.rs | 41 ++++++++++++ diskann-pipnn/src/leaf_kernel.rs | 96 +++++++++++++++++++++++++-- diskann-pipnn/src/partition_kernel.rs | 66 ++++++++++++++++++ 3 files changed, 199 insertions(+), 4 deletions(-) diff --git a/diskann-pipnn/src/kernel_metric.rs b/diskann-pipnn/src/kernel_metric.rs index f7e61d43b..08bc2987d 100644 --- a/diskann-pipnn/src/kernel_metric.rs +++ b/diskann-pipnn/src/kernel_metric.rs @@ -12,15 +12,28 @@ use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; +/// Stored scale representation consumed by one kernel position. +/// +/// Associated constants on `KernelMetric` let the compiler remove unused scale +/// loads and allocations after metric selection. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum ScaleKind { + /// Metric does not read this scale position. None, + /// Stored value is already a squared norm. SquaredNorm, + /// Stored value is a squared norm that must become a norm. NormFromSquared, + /// Stored value is already a norm. Norm, } impl ScaleKind { + /// Convert stored scale to the arithmetic form required by a kernel. + /// + /// DiskANN treats subnormal squared norms, and corresponding subnormal + /// norms, as zero before division. Ordered comparisons intentionally leave + /// NaN unchanged so later distance comparisons keep it non-rankable. #[inline(always)] pub(crate) fn transform(self, stored: f32) -> f32 { match self { @@ -48,27 +61,42 @@ impl ScaleKind { } } +/// Concrete metric contract shared by leaf and partition hot loops. +/// +/// Runtime `Metric` is converted to one implementor before final type erasure. +/// Generic methods then inline metric arithmetic into the architecture-specific +/// function pointer. Leaf and partition operations remain separate because L2 +/// partition ranking deliberately omits the row norm. pub(crate) trait KernelMetric: Send + Sync + 'static { + /// Runtime tag represented by this marker. const METRIC: Metric; + /// Diagonal scale representation used by the leaf kernel. const LEAF_SCALE: ScaleKind; + /// Point-row scale representation used by partition assignment. const PARTITION_ROW_SCALE: ScaleKind; + /// Leader-column scale representation used by partition assignment. const PARTITION_LEADER_SCALE: ScaleKind; + /// SIMD distance for one leaf row against a lane group of earlier points. fn leaf_distance(arch: F::Arch, dot: F, row_scale: F, column_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; + /// Scalar-tail equivalent of `leaf_distance`. fn leaf_distance_scalar(dot: f32, row_scale: f32, column_scale: f32) -> f32; + /// SIMD ranking score for one point row against a lane group of leaders. fn partition_distance(arch: F::Arch, dot: F, row_scale: F, leader_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; + /// Scalar-tail equivalent of `partition_distance`. fn partition_distance_scalar(dot: f32, row_scale: f32, leader_scale: f32) -> f32; } +/// Zero-sized metric markers used only for monomorphization. pub(crate) struct L2; pub(crate) struct Cosine; pub(crate) struct CosineNormalized; @@ -97,6 +125,11 @@ fn clamp_nonnegative_scalar(distance: f32) -> f32 { } } +/// Compute cosine distance while preserving DiskANN zero/NaN semantics. +/// +/// Zero lanes divide by one only to keep the operation defined, then explicitly +/// select zero similarity. NaN norms fail the zero comparison and propagate +/// through division, leaving the final distance non-rankable. #[inline(always)] fn cosine_distance(arch: F::Arch, dot: F, row_norm: F, column_norm: F) -> F where @@ -265,12 +298,20 @@ impl KernelMetric for InnerProduct { } } +/// BYO-type-erasure visitor for runtime metric selection. +/// +/// The visitor receives concrete `M`, allowing architecture and width wrappers +/// to compose with metric arithmetic before producing the final function pointer. +/// This avoids a nested metric trait object inside architecture dispatch. pub(crate) trait EraseMetric { + /// Final caller-selected erased representation. type Output; + /// Consume the visitor with one concrete metric marker. fn erase(self) -> Self::Output; } +/// Visit the concrete marker represented by a runtime metric tag. pub(crate) fn erase_metric(metric: Metric, erase: E) -> E::Output { match metric { Metric::L2 => erase.erase::(), diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index d0869bf61..aac60c93f 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -11,6 +11,20 @@ //! requested neighbor count, and runtime CPU; repeated leaves call a direct //! `diskann-wide` function pointer without ISA or metric dispatch in the loop. //! NaN distances are not rankable, and equal distances retain pair scan order. +//! +//! ```text +//! metric + requested k + runtime architecture +//! │ +//! v +//! prepared Dispatched1 handle +//! │ reused for every leaf +//! v +//! shape validation -> scale scratch -> strict-lower scan -> sorted row slots +//! ``` +//! +//! `workspace.worst[row]` always mirrors the last (worst) retained slot for that +//! row. The SIMD loop may update both endpoints of a pair, so this mirror is the +//! threshold shared by row and column candidate masks. use std::marker::PhantomData; @@ -144,6 +158,11 @@ pub fn leaf_output_len(points: usize, k: usize) -> Result { input: LeafTopK<'a>, @@ -197,6 +216,10 @@ impl LeafKernel { } } +/// Requested-width dispatch selected once while preparing the kernel. +/// +/// Widths one through three receive fixed array rows. Larger widths retain one +/// dynamic implementation instead of multiplying code size by every possible k. #[derive(Clone, Copy, Debug)] enum KValue { One, @@ -216,6 +239,11 @@ impl KValue { } } +/// First dispatch stage: choose the runtime architecture once. +/// +/// The factory itself uses `dispatch1_no_features`; only the returned leaf entry +/// needs target features, so architecture-specific code remains behind the final +/// direct function pointer. struct PrepareLeaf { requested_k: usize, } @@ -238,6 +266,10 @@ where } } +/// BYO-type-erasure visitor holding a concrete architecture. +/// +/// `erase` receives a concrete metric marker, then combines `A`, `M`, and +/// the requested width before erasing the result into exactly one `Dispatched1`. struct BuildLeaf { arch: A, requested_k: usize, @@ -279,6 +311,10 @@ where } } +/// Architecture/metric/width-specialized function-pointer destination. +/// +/// This type is zero-sized. All per-leaf state arrives through `LeafCall`; the +/// entry validates and initializes that state before reaching pointer-based SIMD. struct LeafEntry(PhantomData<(M, S)>); impl FTarget1, LeafCall<'_>> for LeafEntry @@ -291,11 +327,15 @@ where u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { fn run(arch: A, mut call: LeafCall<'_>) -> Result { + // Validation establishes every shape and active-prefix invariant used by + // unchecked loads below. No output or scratch mutation occurs on error. let actual_k = validate(call.input, call.requested_k, &call.output)?; if actual_k == 0 { return Ok(0); } + // Norm and threshold scratch are reset for this leaf, while Vec capacity + // remains reusable by the worker that owns the workspace. prepare_workspace::(call.input, call.workspace)?; call.output.as_mut_slice().fill(LeafNeighbor::default()); call.workspace.worst.fill(f32::INFINITY); @@ -323,6 +363,11 @@ where } } +/// Validate the complete safety contract before dispatched SIMD executes. +/// +/// Matrix views are rechecked with `checked_mul` because the hot loop performs +/// unchecked contiguous loads. Output columns must equal the clamped effective +/// k so fixed-row conversion cannot expose a partial row. fn validate( input: LeafTopK<'_>, k: usize, @@ -357,6 +402,11 @@ fn validate( Ok(actual_k) } +/// Prepare metric-specific scale and threshold scratch. +/// +/// L2 stores diagonal squared norms; cosine converts diagonals to norms using +/// DiskANN's zero threshold. Normalized cosine and inner product skip the norm +/// allocation entirely. `worst` is reset separately after allocation succeeds. fn prepare_workspace( input: LeafTopK<'_>, workspace: &mut LeafTopKWorkspace, @@ -413,6 +463,11 @@ fn check_length( } } +/// Prepared requested-width policy. +/// +/// The actual width can be smaller for singleton/tiny leaves, so each policy +/// performs one pre-loop clamp dispatch while keeping width selection out of the +/// pair scan. trait SlotSelection: Send + Sync + 'static { fn process( arch: F::Arch, @@ -468,6 +523,10 @@ impl SlotSelection for DynamicSelection { } } +/// Convert effective k into one fixed row representation or the dynamic fallback. +/// +/// This branch runs once per leaf. Fixed conversion uses `as_chunks_mut` once, +/// avoiding per-candidate slice-to-array checks while retaining safe insertion. fn process_selected( arch: F::Arch, input: LeafTopK<'_>, @@ -515,6 +574,11 @@ fn process_fixed( process_pairs::(arch, input, FixedRows(rows), norms, worst); } +/// Mutable row adapter used by the shared pair traversal. +/// +/// Implementations own the exclusive output borrow for the whole scan. Each +/// insertion borrows one row briefly, so updates to the current row and earlier +/// endpoint rows cannot alias simultaneously. trait NeighborRows { fn len(&self) -> usize; fn insert(&mut self, row: usize, position: u32, distance: f32) -> f32; @@ -555,10 +619,23 @@ impl NeighborRows for DynamicRows<'_> { } } -/// Scan the strict lower triangle and update both endpoint rows. +/// Scan the strict lower triangle once and update both endpoint rows. +/// +/// Invariants on entry: /// -/// `M` fixes metric arithmetic before type erasure. `R` presents either -/// fixed-width array rows or the uncommon run-time-width rows. +/// - `dots` is a validated square row-major matrix; +/// - `output` has one sorted ascending-distance row per point; +/// - `worst[row]` equals that row's last slot; +/// - `norms` has one value per point exactly when `M` requires scales. +/// +/// Each SIMD chunk computes both endpoint eligibility masks before mutation. +/// Multiple lanes compete for the current row, so row candidates recheck its +/// live cached threshold. Every column lane targets a distinct earlier row and +/// can use the precomputed mask directly. Scalar tails call the matching scalar +/// metric operation to preserve established rounding semantics. +/// +/// `M` is concrete before type erasure. `R` presents fixed array rows for common +/// widths or safe dynamic slices for the uncommon fallback. #[inline(never)] fn process_pairs( arch: F::Arch, @@ -599,6 +676,9 @@ fn process_pairs( F::default(arch) }; let distances = M::leaf_distance(arch, pair_dots, row_norm, column_norms); + // Every pair may improve the current row and its earlier endpoint. + // Derive both masks from the same distance vector before either side + // mutates its threshold. let row_eligible = distances.lt_simd(F::splat(arch, row_worst)); // SAFETY: the full chunk lies below `row`, so it is inside `worst`. let column_worst = unsafe { F::load_simd(arch, worst_ptr.add(column)) }; @@ -701,7 +781,11 @@ fn insert_scalar( worst[row] = insert_dynamic(&mut output[row * k..(row + 1) * k], position, distance); } -/// Insert into a production row whose width is known at dispatch. +/// Insert into a fixed-width row and return its new worst distance. +/// +/// Production widths one through three use straight-line shifts. Strict `<` +/// comparisons preserve scan order for ties; callers already rejected NaN via +/// the eligibility comparison. #[inline(always)] fn insert_fixed(row: &mut [LeafNeighbor; N], position: u32, distance: f32) -> f32 { let entry = LeafNeighbor::new(position, distance); @@ -740,6 +824,10 @@ fn insert_fixed(row: &mut [LeafNeighbor; N], position: u32, dist } } +/// Insert into a run-time-width row using the same stable ordering contract. +/// +/// The candidate replaces the last slot, then bubbles toward the front. This +/// path is used only for k greater than three. #[inline(always)] fn insert_dynamic(row: &mut [LeafNeighbor], position: u32, distance: f32) -> f32 { let last = row.len() - 1; diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index 80d2a02e8..e19560c6b 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -14,6 +14,19 @@ //! L2 deliberately omits the point norm because it is constant across every //! leader in one row. Cosine consumes squared point norms and leader norms. NaN //! distances are not rankable, and equal distances retain leader scan order. +//! +//! ```text +//! metric + runtime architecture +//! │ +//! v +//! prepared Dispatched2 handle +//! │ reused for every point stripe +//! v +//! shape/scale validation -> SIMD chunks + scalar tail -> sorted leader IDs +//! ``` +//! +//! Each row owns a fixed-capacity sorted tracker. Its last retained distance is +//! the rejection threshold, so noncompetitive SIMD chunks avoid lane extraction. use std::marker::PhantomData; @@ -130,6 +143,10 @@ pub enum PartitionKernelError { }, } +/// Lifetime families used by the direct function-pointer interface. +/// +/// Input and output receive independent call lifetimes. The prepared handle +/// stores neither view, so it remains `Copy + Send + Sync` across worker threads. #[derive(Debug)] struct PartitionInput; @@ -176,6 +193,10 @@ impl PartitionKernel { } } +/// First dispatch stage: select runtime architecture once. +/// +/// `dispatch1_no_features` runs only this factory. The returned entry pointer is +/// generated by the selected architecture and carries its required features. struct PreparePartition; impl arch::Target1 for PreparePartition @@ -190,6 +211,10 @@ where } } +/// BYO-type-erasure visitor holding a concrete architecture. +/// +/// `erase` combines architecture `A` and concrete metric `M`, then produces +/// one direct function pointer. No nested metric trait object remains at runtime. struct BuildPartition(A); impl EraseMetric for BuildPartition @@ -213,6 +238,10 @@ where } } +/// Architecture/metric-specialized function-pointer destination. +/// +/// The zero-sized entry receives all stripe state as arguments. Validation must +/// complete before `process_rows` reaches unchecked contiguous SIMD loads. struct PartitionEntry(PhantomData); impl FTarget2, PartitionTopK<'_>, MutMatrixView<'_, u32>> @@ -229,6 +258,8 @@ where input: PartitionTopK<'_>, mut output: MutMatrixView<'_, u32>, ) -> Result<(), PartitionKernelError> { + // Validation establishes matrix areas, backing lengths, scale units, + // and fanout bounds before any output mutation or unchecked load. let scales = validate::(input, &output)?; let fanout = output.ncols(); if fanout == 0 || input.dots.nrows() == 0 { @@ -236,6 +267,8 @@ where } process_rows::(arch, input.dots, scales, fanout, output.as_mut_slice()); + // A sorted tracker can be underfilled only at its last slot. This keeps + // post-validation linear in rows rather than scanning every output ID. if let Some(row) = output .as_slice() .chunks_exact(fanout) @@ -247,12 +280,21 @@ where } } +/// Validated scale slices in the storage form required by `M`. +/// +/// Empty slices are intentional for metrics that omit a scale; consumers branch +/// on associated `ScaleKind` constants that monomorphize out of hot loops. #[derive(Clone, Copy)] struct ScaleSlices<'a> { rows: &'a [f32], leaders: &'a [f32], } +/// Validate the complete partition-kernel safety and metric contract. +/// +/// Matrix areas are recomputed with `checked_mul` before pointer loads. The +/// `PartitionScales` variant must match concrete metric `M`, preventing plausible +/// but incorrect norm units from crossing the interface. fn validate<'a, M: KernelMetric>( input: PartitionTopK<'a>, output: &MutMatrixView<'_, u32>, @@ -370,6 +412,19 @@ fn check_length( } } +/// Convert each point-to-leader dot-product row into sorted top-fanout IDs. +/// +/// Per-row flow: +/// +/// 1. transform the row scale once according to concrete metric `M`; +/// 2. process full SIMD chunks, rejecting lanes against the tracker's last slot; +/// 3. process the tail with the scalar metric operation; +/// 4. copy the sorted tracker prefix to that row's output. +/// +/// `top[..fanout]` remains sorted after every accepted candidate. Strict `<` +/// preserves leader scan order for ties and makes NaNs non-rankable. L2 keeps +/// historical bulk-FMA/scalar-tail rounding because changing it can alter graph +/// assignment at near ties. fn process_rows( arch: F::Arch, dots: MatrixView<'_, f32>, @@ -469,6 +524,11 @@ fn process_rows_scalar( } } +/// Offer competitive SIMD lanes to a row tracker in increasing leader order. +/// +/// The broadcast threshold avoids materializing lanes when none can improve the +/// last slot. Bit iteration follows low-to-high lane order, preserving scalar tie +/// behavior across SIMD widths. fn insert_lanes(distances: F, base: usize, top: &mut TopK, fanout: usize) where F: SIMDVector + SIMDPartialOrd, @@ -490,6 +550,11 @@ where } } +/// Insert one strictly better candidate while preserving sorted-prefix state. +/// +/// The last slot is overwritten, then bubbled left. Equal and NaN distances do +/// not enter, so scan order is the deterministic tie breaker and the last slot +/// remains both rejection threshold and underfill sentinel. #[inline(always)] fn insert_topk(top: &mut TopK, fanout: usize, leader: u32, distance: f32) { let threshold = fanout - 1; @@ -505,6 +570,7 @@ fn insert_topk(top: &mut TopK, fanout: usize, leader: u32, distance: f32) { } } +/// Publish only leader IDs; distances stay private tracker state. fn copy_ids(top: &TopK, output: &mut [u32]) { for (destination, &(leader, _)) in output.iter_mut().zip(top) { *destination = leader; From 928a359830ddf4daf6397f43e92f1e452fc26e88 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 11:28:08 +0000 Subject: [PATCH 16/26] test(pipnn): confine traversal references --- diskann-pipnn/src/leaf_kernel.rs | 92 +++++++++++++-------------- diskann-pipnn/src/partition_kernel.rs | 90 +++++++++++++------------- 2 files changed, 90 insertions(+), 92 deletions(-) diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index aac60c93f..9697b13c3 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -741,46 +741,6 @@ fn process_pairs( debug_assert_eq!(output.len(), points); } -#[cfg(test)] -fn process_pairs_scalar( - input: LeafTopK<'_>, - k: usize, - output: &mut [LeafNeighbor], - norms: &[f32], - worst: &mut [f32], -) { - let points = input.dots.nrows(); - let uses_norms = M::LEAF_SCALE.is_some(); - for row in 1..points { - for column in 0..row { - let (row_norm, column_norm) = if uses_norms { - (norms[row], norms[column]) - } else { - (0.0, 0.0) - }; - let distance = - M::leaf_distance_scalar(input.dots[(row, column)], row_norm, column_norm); - insert_scalar(output, worst, k, row, column as u32, distance); - insert_scalar(output, worst, k, column, row as u32, distance); - } - } -} - -#[cfg(test)] -fn insert_scalar( - 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; - } - worst[row] = insert_dynamic(&mut output[row * k..(row + 1) * k], position, distance); -} - /// Insert into a fixed-width row and return its new worst distance. /// /// Production widths one through three use straight-line shifts. Strict `<` @@ -868,13 +828,47 @@ mod tests { } } - fn scalar(input: LeafTopK<'_>, k: usize, output: &mut [LeafNeighbor]) { + // Differential oracle for traversal and dispatch only. It intentionally + // shares `M::leaf_distance_scalar`; public API tests independently spell + // out metric formulas and full sorting behavior. + fn scalar_traversal_reference( + input: LeafTopK<'_>, + k: usize, + output: &mut [LeafNeighbor], + ) { let points = input.dots.nrows(); let norms: Vec<_> = (0..points) .map(|row| M::LEAF_SCALE.transform(input.dots[(row, row)])) .collect(); let mut worst = vec![f32::INFINITY; points]; - process_pairs_scalar::(input, k, output, &norms, &mut worst); + let uses_norms = M::LEAF_SCALE.is_some(); + for row in 1..points { + for column in 0..row { + let (row_norm, column_norm) = if uses_norms { + (norms[row], norms[column]) + } else { + (0.0, 0.0) + }; + let distance = + M::leaf_distance_scalar(input.dots[(row, column)], row_norm, column_norm); + insert_reference(output, &mut worst, k, row, column as u32, distance); + insert_reference(output, &mut worst, k, column, row as u32, distance); + } + } + } + + fn insert_reference( + 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; + } + worst[row] = insert_dynamic(&mut output[row * k..(row + 1) * k], position, distance); } fn scalar_for_metric( @@ -884,10 +878,12 @@ mod tests { output: &mut [LeafNeighbor], ) { match metric { - Metric::L2 => scalar::(input, k, output), - Metric::Cosine => scalar::(input, k, output), - Metric::CosineNormalized => scalar::(input, k, output), - Metric::InnerProduct => scalar::(input, k, output), + Metric::L2 => scalar_traversal_reference::(input, k, output), + Metric::Cosine => scalar_traversal_reference::(input, k, output), + Metric::CosineNormalized => { + scalar_traversal_reference::(input, k, output) + } + Metric::InnerProduct => scalar_traversal_reference::(input, k, output), } } @@ -943,9 +939,9 @@ mod tests { 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_scalar(&mut output, &mut worst, 4, 0, position, distance); + insert_reference(&mut output, &mut worst, 4, 0, position, distance); } - insert_scalar(&mut output, &mut worst, 4, 0, 5, f32::NAN); + insert_reference(&mut output, &mut worst, 4, 0, 5, f32::NAN); assert_eq!( output, diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index e19560c6b..f01a3fecc 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -487,43 +487,6 @@ fn process_rows( } } -#[cfg(test)] -fn process_rows_scalar( - dots: MatrixView<'_, f32>, - scales: ScaleSlices<'_>, - fanout: usize, - output: &mut [u32], -) { - let leaders = dots.ncols(); - for (row, (dot_row, output_row)) in dots - .as_slice() - .chunks_exact(leaders) - .zip(output.chunks_exact_mut(fanout)) - .enumerate() - { - let row_scale = if M::PARTITION_ROW_SCALE.is_some() { - M::PARTITION_ROW_SCALE.transform(scales.rows[row]) - } else { - 0.0 - }; - let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; - for (leader, &dot) in dot_row.iter().enumerate() { - let leader_scale = if M::PARTITION_LEADER_SCALE.is_some() { - M::PARTITION_LEADER_SCALE.transform(scales.leaders[leader]) - } else { - 0.0 - }; - insert_topk( - &mut top, - fanout, - leader as u32, - M::partition_distance_scalar(dot, row_scale, leader_scale), - ); - } - copy_ids(&top, output_row); - } -} - /// Offer competitive SIMD lanes to a row tracker in increasing leader order. /// /// The broadcast threshold avoids materializing lanes when none can improve the @@ -634,7 +597,14 @@ mod tests { } } - fn scalar(input: PartitionTopK<'_>, fanout: usize, output: &mut [u32]) { + // Differential oracle for SIMD chunking, scalar tails, and tracker order. + // It intentionally shares `M::partition_distance_scalar`; public API tests + // independently spell out ranking formulas and full sorting behavior. + fn scalar_traversal_reference( + input: PartitionTopK<'_>, + fanout: usize, + output: &mut [u32], + ) { let scales = match input.scales { PartitionScales::L2 { leader_squared_norms, @@ -654,7 +624,35 @@ mod tests { leaders: &[], }, }; - process_rows_scalar::(input.dots, scales, fanout, output); + let leaders = input.dots.ncols(); + for (row, (dot_row, output_row)) in input + .dots + .as_slice() + .chunks_exact(leaders) + .zip(output.chunks_exact_mut(fanout)) + .enumerate() + { + let row_scale = if M::PARTITION_ROW_SCALE.is_some() { + M::PARTITION_ROW_SCALE.transform(scales.rows[row]) + } else { + 0.0 + }; + let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + for (leader, &dot) in dot_row.iter().enumerate() { + let leader_scale = if M::PARTITION_LEADER_SCALE.is_some() { + M::PARTITION_LEADER_SCALE.transform(scales.leaders[leader]) + } else { + 0.0 + }; + insert_topk( + &mut top, + fanout, + leader as u32, + M::partition_distance_scalar(dot, row_scale, leader_scale), + ); + } + copy_ids(&top, output_row); + } } fn scalar_for_metric( @@ -664,10 +662,14 @@ mod tests { output: &mut [u32], ) { match metric { - Metric::L2 => scalar::(input, fanout, output), - Metric::Cosine => scalar::(input, fanout, output), - Metric::CosineNormalized => scalar::(input, fanout, output), - Metric::InnerProduct => scalar::(input, fanout, output), + Metric::L2 => scalar_traversal_reference::(input, fanout, output), + Metric::Cosine => scalar_traversal_reference::(input, fanout, output), + Metric::CosineNormalized => { + scalar_traversal_reference::(input, fanout, output) + } + Metric::InnerProduct => { + scalar_traversal_reference::(input, fanout, output) + } } } @@ -754,7 +756,7 @@ mod tests { &leader_scales, ); let mut expected = vec![u32::MAX; row_scales.len() * 2]; - scalar::(input, 2, &mut expected); + scalar_traversal_reference::(input, 2, &mut expected); let mut actual = vec![u32::MAX; row_scales.len() * 2]; PartitionKernel::new(Metric::Cosine) .nearest_leaders( From 531dd2a49718b3ec2d0660a8bb8ab253720fbb22 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 11:31:13 +0000 Subject: [PATCH 17/26] refactor(pipnn): name metric visitor by action --- diskann-pipnn/src/kernel_metric.rs | 14 +++++++------- diskann-pipnn/src/leaf_kernel.rs | 12 ++++++------ diskann-pipnn/src/partition_kernel.rs | 10 +++++----- 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/diskann-pipnn/src/kernel_metric.rs b/diskann-pipnn/src/kernel_metric.rs index 08bc2987d..181b642f7 100644 --- a/diskann-pipnn/src/kernel_metric.rs +++ b/diskann-pipnn/src/kernel_metric.rs @@ -303,21 +303,21 @@ impl KernelMetric for InnerProduct { /// The visitor receives concrete `M`, allowing architecture and width wrappers /// to compose with metric arithmetic before producing the final function pointer. /// This avoids a nested metric trait object inside architecture dispatch. -pub(crate) trait EraseMetric { +pub(crate) trait MetricVisitor { /// Final caller-selected erased representation. type Output; /// Consume the visitor with one concrete metric marker. - fn erase(self) -> Self::Output; + fn visit(self) -> Self::Output; } /// Visit the concrete marker represented by a runtime metric tag. -pub(crate) fn erase_metric(metric: Metric, erase: E) -> E::Output { +pub(crate) fn visit_metric(metric: Metric, visitor: V) -> V::Output { match metric { - Metric::L2 => erase.erase::(), - Metric::Cosine => erase.erase::(), - Metric::CosineNormalized => erase.erase::(), - Metric::InnerProduct => erase.erase::(), + Metric::L2 => visitor.visit::(), + Metric::Cosine => visitor.visit::(), + Metric::CosineNormalized => visitor.visit::(), + Metric::InnerProduct => visitor.visit::(), } } diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 9697b13c3..71276a970 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -36,7 +36,7 @@ use diskann_wide::{ Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, }; -use crate::kernel_metric::{erase_metric, EraseMetric, KernelMetric}; +use crate::kernel_metric::{visit_metric, KernelMetric, MetricVisitor}; /// One leaf-local neighbor and its metric distance. #[derive(Clone, Copy, Debug, PartialEq)] @@ -256,7 +256,7 @@ where u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { fn run(self, arch: A, metric: Metric) -> LeafKernel { - erase_metric( + visit_metric( metric, BuildLeaf { arch, @@ -268,8 +268,8 @@ where /// BYO-type-erasure visitor holding a concrete architecture. /// -/// `erase` receives a concrete metric marker, then combines `A`, `M`, and -/// the requested width before erasing the result into exactly one `Dispatched1`. +/// `visit` receives a concrete metric marker, then combines `A`, `M`, and +/// the requested width into exactly one `Dispatched1`. struct BuildLeaf { arch: A, requested_k: usize, @@ -292,7 +292,7 @@ where } } -impl EraseMetric for BuildLeaf +impl MetricVisitor for BuildLeaf where A: Architecture, A::f32x16: std::ops::Div, @@ -301,7 +301,7 @@ where { type Output = LeafKernel; - fn erase(self) -> Self::Output { + fn visit(self) -> Self::Output { match KValue::from_requested(self.requested_k) { KValue::One => self.build::>(), KValue::Two => self.build::>(), diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index f01a3fecc..992d9fa28 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -38,7 +38,7 @@ use diskann_wide::{ Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, }; -use crate::kernel_metric::{erase_metric, EraseMetric, KernelMetric, ScaleKind}; +use crate::kernel_metric::{visit_metric, KernelMetric, MetricVisitor, ScaleKind}; /// Maximum number of leaders retained for one point. /// @@ -207,17 +207,17 @@ where u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { fn run(self, arch: A, metric: Metric) -> PartitionKernel { - erase_metric(metric, BuildPartition(arch)) + visit_metric(metric, BuildPartition(arch)) } } /// BYO-type-erasure visitor holding a concrete architecture. /// -/// `erase` combines architecture `A` and concrete metric `M`, then produces +/// `visit` combines architecture `A` and concrete metric `M`, then produces /// one direct function pointer. No nested metric trait object remains at runtime. struct BuildPartition(A); -impl EraseMetric for BuildPartition +impl MetricVisitor for BuildPartition where A: Architecture, A::f32x16: std::ops::Div, @@ -226,7 +226,7 @@ where { type Output = PartitionKernel; - fn erase(self) -> Self::Output { + fn visit(self) -> Self::Output { PartitionKernel { run: self.0.dispatch2::< PartitionEntry, From 9209c22b4e5e512c655f38ca2adb20dc9367e2fa Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 11:39:26 +0000 Subject: [PATCH 18/26] refactor(pipnn): carry dynamic leaf width --- diskann-pipnn/src/leaf_kernel.rs | 88 ++++++++++++++++++++++++-------- 1 file changed, 67 insertions(+), 21 deletions(-) diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 71276a970..8efc76718 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -183,17 +183,22 @@ type LeafFn = Dispatched1, LeafCallArg>; /// A leaf kernel prepared for one metric, neighbor count, and the current CPU. /// /// Construct this once with [`LeafKernel::new`] and share it across leaf workers. -/// The handle stores only a direct function pointer and the requested `k`. +/// The handle stores only a direct function pointer and the requested-width mode. #[derive(Clone, Copy, Debug)] pub struct LeafKernel { run: LeafFn, - requested_k: usize, + k: KValue, } impl LeafKernel { /// Prepare a leaf kernel for `metric`, `k`, and the current CPU. pub fn new(metric: Metric, k: usize) -> Self { - diskann_wide::arch::dispatch1_no_features(PrepareLeaf { requested_k: k }, metric) + diskann_wide::arch::dispatch1_no_features( + PrepareLeaf { + k: KValue::from_requested(k), + }, + metric, + ) } /// Select the nearest non-self leaf positions for every row. @@ -211,30 +216,42 @@ impl LeafKernel { input, output, workspace, - requested_k: self.requested_k, + requested_k: self.k.requested(), }) } } /// Requested-width dispatch selected once while preparing the kernel. /// -/// Widths one through three receive fixed array rows. Larger widths retain one -/// dynamic implementation instead of multiplying code size by every possible k. -#[derive(Clone, Copy, Debug)] +/// Widths one through three receive fixed array rows. Zero is a validated no-op; +/// larger values carry their requested width into the dynamic implementation. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] enum KValue { + Zero, One, Two, Three, - Large, + Dynamic(usize), } impl KValue { const fn from_requested(k: usize) -> Self { match k { + 0 => Self::Zero, 1 => Self::One, 2 => Self::Two, 3 => Self::Three, - _ => Self::Large, + width => Self::Dynamic(width), + } + } + + const fn requested(self) -> usize { + match self { + Self::Zero => 0, + Self::One => 1, + Self::Two => 2, + Self::Three => 3, + Self::Dynamic(width) => width, } } } @@ -245,7 +262,7 @@ impl KValue { /// needs target features, so architecture-specific code remains behind the final /// direct function pointer. struct PrepareLeaf { - requested_k: usize, + k: KValue, } impl arch::Target1 for PrepareLeaf @@ -256,13 +273,7 @@ where u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { fn run(self, arch: A, metric: Metric) -> LeafKernel { - visit_metric( - metric, - BuildLeaf { - arch, - requested_k: self.requested_k, - }, - ) + visit_metric(metric, BuildLeaf { arch, k: self.k }) } } @@ -272,7 +283,7 @@ where /// the requested width into exactly one `Dispatched1`. struct BuildLeaf { arch: A, - requested_k: usize, + k: KValue, } impl BuildLeaf @@ -287,7 +298,7 @@ where run: self .arch .dispatch1::, Result, LeafCallArg>(), - requested_k: self.requested_k, + k: self.k, } } } @@ -302,11 +313,12 @@ where type Output = LeafKernel; fn visit(self) -> Self::Output { - match KValue::from_requested(self.requested_k) { + match self.k { + KValue::Zero => self.build::(), KValue::One => self.build::>(), KValue::Two => self.build::>(), KValue::Three => self.build::>(), - KValue::Large => self.build::(), + KValue::Dynamic(_) => self.build::(), } } } @@ -483,9 +495,28 @@ trait SlotSelection: Send + Sync + 'static { u64: From<<::BitMask as SIMDMask>::Underlying>; } +struct ZeroSelection; struct FixedSelection; struct DynamicSelection; +impl SlotSelection for ZeroSelection { + fn process( + _arch: F::Arch, + _input: LeafTopK<'_>, + actual_k: usize, + _output: &mut [LeafNeighbor], + _norms: &[f32], + _worst: &mut [f32], + ) where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + M: KernelMetric, + u64: From<<::BitMask as SIMDMask>::Underlying>, + { + debug_assert_eq!(actual_k, 0); + } +} + impl SlotSelection for FixedSelection { fn process( arch: F::Arch, @@ -828,6 +859,21 @@ mod tests { } } + #[test] + fn k_value_preserves_requested_width() { + for (requested, value) in [ + (0, KValue::Zero), + (1, KValue::One), + (2, KValue::Two), + (3, KValue::Three), + (4, KValue::Dynamic(4)), + (17, KValue::Dynamic(17)), + ] { + assert_eq!(KValue::from_requested(requested), value); + assert_eq!(value.requested(), requested); + } + } + // Differential oracle for traversal and dispatch only. It intentionally // shares `M::leaf_distance_scalar`; public API tests independently spell // out metric formulas and full sorting behavior. From 34de407ab96b3e31010ce7cf430adc533a05d3a7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 12:51:02 +0000 Subject: [PATCH 19/26] refactor(pipnn)!: clarify point and neighbor roles Use output columns as the sole leaf-specific neighbor count and reserve row/column terminology for matrix shapes. BREAKING CHANGE: LeafKernel::new no longer takes k, nearest_neighbors returns (), and kernel input/neighbor/error fields use source-target and point-leader names. --- diskann-pipnn/src/kernel_metric.rs | 73 +- diskann-pipnn/src/leaf_kernel.rs | 815 +++++++++----------- diskann-pipnn/src/partition_kernel.rs | 328 ++++---- diskann-pipnn/tests/leaf_kernel_api.rs | 218 +++--- diskann-pipnn/tests/partition_kernel_api.rs | 146 ++-- 5 files changed, 782 insertions(+), 798 deletions(-) diff --git a/diskann-pipnn/src/kernel_metric.rs b/diskann-pipnn/src/kernel_metric.rs index 181b642f7..01aaca0d3 100644 --- a/diskann-pipnn/src/kernel_metric.rs +++ b/diskann-pipnn/src/kernel_metric.rs @@ -7,7 +7,7 @@ //! //! Runtime metric selection happens only while preparing a dispatched kernel. //! The hot loops receive a concrete marker type, allowing metric arithmetic and -//! scale handling to inline without a per-row or per-chunk enum match. +//! scale handling to inline without a per-point or per-chunk enum match. use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; @@ -66,34 +66,34 @@ impl ScaleKind { /// Runtime `Metric` is converted to one implementor before final type erasure. /// Generic methods then inline metric arithmetic into the architecture-specific /// function pointer. Leaf and partition operations remain separate because L2 -/// partition ranking deliberately omits the row norm. +/// partition ranking deliberately omits the point norm. pub(crate) trait KernelMetric: Send + Sync + 'static { /// Runtime tag represented by this marker. const METRIC: Metric; /// Diagonal scale representation used by the leaf kernel. const LEAF_SCALE: ScaleKind; - /// Point-row scale representation used by partition assignment. - const PARTITION_ROW_SCALE: ScaleKind; + /// Point scale representation used by partition assignment. + const PARTITION_POINT_SCALE: ScaleKind; /// Leader-column scale representation used by partition assignment. const PARTITION_LEADER_SCALE: ScaleKind; - /// SIMD distance for one leaf row against a lane group of earlier points. - fn leaf_distance(arch: F::Arch, dot: F, row_scale: F, column_scale: F) -> F + /// SIMD distance for one leaf source against a lane group of earlier targets. + fn leaf_distance(arch: F::Arch, dot: F, source_scale: F, target_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; /// Scalar-tail equivalent of `leaf_distance`. - fn leaf_distance_scalar(dot: f32, row_scale: f32, column_scale: f32) -> f32; + fn leaf_distance_scalar(dot: f32, source_scale: f32, target_scale: f32) -> f32; - /// SIMD ranking score for one point row against a lane group of leaders. - fn partition_distance(arch: F::Arch, dot: F, row_scale: F, leader_scale: F) -> F + /// SIMD ranking score for one point against a lane group of leaders. + fn partition_distance(arch: F::Arch, dot: F, point_scale: F, leader_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; /// Scalar-tail equivalent of `partition_distance`. - fn partition_distance_scalar(dot: f32, row_scale: f32, leader_scale: f32) -> f32; + fn partition_distance_scalar(dot: f32, point_scale: f32, leader_scale: f32) -> f32; } /// Zero-sized metric markers used only for monomorphization. @@ -131,7 +131,7 @@ fn clamp_nonnegative_scalar(distance: f32) -> f32 { /// select zero similarity. NaN norms fail the zero comparison and propagate /// through division, leaving the final distance non-rankable. #[inline(always)] -fn cosine_distance(arch: F::Arch, dot: F, row_norm: F, column_norm: F) -> F +fn cosine_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, @@ -139,41 +139,44 @@ where let zero = F::default(arch); let one = F::splat(arch, 1.0); let minimum_norm = F::splat(arch, f32::MIN_POSITIVE.sqrt()); - let row_zero = row_norm.lt_simd(minimum_norm); - let column_zero = column_norm.lt_simd(minimum_norm); - 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)); + let source_zero = source_norm.lt_simd(minimum_norm); + let target_zero = target_norm.lt_simd(minimum_norm); + let denominator = source_norm * target_norm; + let safe_denominator = source_zero.select(one, target_zero.select(one, denominator)); + let cosine = source_zero.select(zero, target_zero.select(zero, dot / safe_denominator)); one - cosine } #[inline(always)] -fn cosine_distance_scalar(dot: f32, row_norm: f32, column_norm: f32) -> f32 { - if row_norm < f32::MIN_POSITIVE.sqrt() || column_norm < f32::MIN_POSITIVE.sqrt() { +fn cosine_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32 { + if source_norm < f32::MIN_POSITIVE.sqrt() || target_norm < f32::MIN_POSITIVE.sqrt() { 1.0 } else { - 1.0 - dot / (row_norm * column_norm) + 1.0 - dot / (source_norm * target_norm) } } impl KernelMetric for L2 { const METRIC: Metric = Metric::L2; const LEAF_SCALE: ScaleKind = ScaleKind::SquaredNorm; - const PARTITION_ROW_SCALE: ScaleKind = ScaleKind::None; + const PARTITION_POINT_SCALE: ScaleKind = ScaleKind::None; const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::SquaredNorm; #[inline(always)] - fn leaf_distance(arch: F::Arch, dot: F, row_scale: F, column_scale: F) -> F + fn leaf_distance(arch: F::Arch, dot: F, source_scale: F, target_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - clamp_nonnegative(arch, row_scale + column_scale - F::splat(arch, 2.0) * dot) + clamp_nonnegative( + arch, + source_scale + target_scale - F::splat(arch, 2.0) * dot, + ) } #[inline(always)] - fn leaf_distance_scalar(dot: f32, row_scale: f32, column_scale: f32) -> f32 { - clamp_nonnegative_scalar(row_scale + column_scale - 2.0 * dot) + fn leaf_distance_scalar(dot: f32, source_scale: f32, target_scale: f32) -> f32 { + clamp_nonnegative_scalar(source_scale + target_scale - 2.0 * dot) } #[inline(always)] @@ -196,42 +199,42 @@ impl KernelMetric for L2 { impl KernelMetric for Cosine { const METRIC: Metric = Metric::Cosine; const LEAF_SCALE: ScaleKind = ScaleKind::NormFromSquared; - const PARTITION_ROW_SCALE: ScaleKind = ScaleKind::NormFromSquared; + const PARTITION_POINT_SCALE: ScaleKind = ScaleKind::NormFromSquared; const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::Norm; #[inline(always)] - fn leaf_distance(arch: F::Arch, dot: F, row_scale: F, column_scale: F) -> F + fn leaf_distance(arch: F::Arch, dot: F, source_scale: F, target_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - clamp_nonnegative(arch, cosine_distance(arch, dot, row_scale, column_scale)) + clamp_nonnegative(arch, cosine_distance(arch, dot, source_scale, target_scale)) } #[inline(always)] - fn leaf_distance_scalar(dot: f32, row_scale: f32, column_scale: f32) -> f32 { - clamp_nonnegative_scalar(cosine_distance_scalar(dot, row_scale, column_scale)) + fn leaf_distance_scalar(dot: f32, source_scale: f32, target_scale: f32) -> f32 { + clamp_nonnegative_scalar(cosine_distance_scalar(dot, source_scale, target_scale)) } #[inline(always)] - fn partition_distance(arch: F::Arch, dot: F, row_scale: F, leader_scale: F) -> F + fn partition_distance(arch: F::Arch, dot: F, point_scale: F, leader_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - cosine_distance(arch, dot, row_scale, leader_scale) + cosine_distance(arch, dot, point_scale, leader_scale) } #[inline(always)] - fn partition_distance_scalar(dot: f32, row_scale: f32, leader_scale: f32) -> f32 { - cosine_distance_scalar(dot, row_scale, leader_scale) + fn partition_distance_scalar(dot: f32, point_scale: f32, leader_scale: f32) -> f32 { + cosine_distance_scalar(dot, point_scale, leader_scale) } } impl KernelMetric for CosineNormalized { const METRIC: Metric = Metric::CosineNormalized; const LEAF_SCALE: ScaleKind = ScaleKind::None; - const PARTITION_ROW_SCALE: ScaleKind = ScaleKind::None; + const PARTITION_POINT_SCALE: ScaleKind = ScaleKind::None; const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::None; #[inline(always)] @@ -266,7 +269,7 @@ impl KernelMetric for CosineNormalized { impl KernelMetric for InnerProduct { const METRIC: Metric = Metric::InnerProduct; const LEAF_SCALE: ScaleKind = ScaleKind::None; - const PARTITION_ROW_SCALE: ScaleKind = ScaleKind::None; + const PARTITION_POINT_SCALE: ScaleKind = ScaleKind::None; const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::None; #[inline(always)] diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 8efc76718..25f5e841b 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -5,26 +5,27 @@ //! Prepared nearest-neighbor kernels over a leaf's lower dot-product matrix. //! -//! `sgemm_aat_lower` writes pair `(row, column)` only when `column <= row`. +//! `sgemm_aat_lower` writes pair `(source, target)` only when `target <= source`. //! The kernel scans that strict lower triangle once and offers each distance to -//! both endpoint rows. A [`LeafKernel`] is prepared once for the build metric, -//! requested neighbor count, and runtime CPU; repeated leaves call a direct -//! `diskann-wide` function pointer without ISA or metric dispatch in the loop. +//! both endpoint points. A [`LeafKernel`] is prepared once for the build metric +//! and runtime CPU; each output view supplies its leaf-specific neighbor count. +//! Repeated leaves call a direct `diskann-wide` function pointer without ISA or +//! metric dispatch in the loop. //! NaN distances are not rankable, and equal distances retain pair scan order. //! //! ```text -//! metric + requested k + runtime architecture -//! │ -//! v -//! prepared Dispatched1 handle -//! │ reused for every leaf -//! v -//! shape validation -> scale scratch -> strict-lower scan -> sorted row slots +//! metric + runtime architecture +//! │ +//! v +//! prepared Dispatched1 handle +//! │ reused with input + output.ncols() +//! v +//! shape validation -> scale scratch -> strict-lower scan -> sorted neighbor slots //! ``` //! -//! `workspace.worst[row]` always mirrors the last (worst) retained slot for that -//! row. The SIMD loop may update both endpoints of a pair, so this mirror is the -//! threshold shared by row and column candidate masks. +//! `workspace.worst[source]` always mirrors the last retained slot for that +//! source point. The SIMD loop may update both endpoints of a pair, so this +//! mirror is the threshold shared by source and target candidate masks. use std::marker::PhantomData; @@ -41,16 +42,16 @@ use crate::kernel_metric::{visit_metric, KernelMetric, MetricVisitor}; /// 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`. + /// Target position in the leaf, not a dataset ID. + pub target: u32, + /// Distance from the source point to `target`. pub distance: f32, } impl LeafNeighbor { /// Construct a leaf-local neighbor. - pub const fn new(position: u32, distance: f32) -> Self { - Self { position, distance } + pub const fn new(target: u32, distance: f32) -> Self { + Self { target, distance } } } @@ -62,19 +63,19 @@ impl Default for LeafNeighbor { /// Square lower-triangular dot-product matrix for one leaf. #[derive(Clone, Copy, Debug)] -pub struct LeafTopK<'a> { - /// Point-by-point matrix. Only entries with `column <= row` are read. +pub struct LeafInput<'a> { + /// Point-by-point matrix. Only entries with `target <= source` are read. pub dots: MatrixView<'a, f32>, } /// Reusable temporary storage for leaf top-k selection. #[derive(Debug, Default)] -pub struct LeafTopKWorkspace { +pub struct LeafKernelWorkspace { norms: Vec, worst: Vec, } -impl LeafTopKWorkspace { +impl LeafKernelWorkspace { /// Construct an empty workspace. pub const fn new() -> Self { Self { @@ -118,19 +119,25 @@ pub enum LeafKernelError { /// Supplied length. actual: usize, }, - /// The output matrix does not match the requested neighbor shape. - #[error( - "invalid output shape: expected {expected_rows} x {expected_cols}, got {actual_rows} x {actual_cols}" - )] - InvalidOutputShape { + /// The output matrix does not have one row per input point. + #[error("invalid output row count: expected {expected}, got {actual} with {columns} columns")] + InvalidOutputRows { /// Required row count. - expected_rows: usize, - /// Required column count. - expected_cols: usize, + expected: usize, /// Supplied row count. - actual_rows: usize, - /// Supplied column count. - actual_cols: usize, + actual: usize, + /// Supplied neighbor columns. + columns: usize, + }, + /// A source requests more non-self neighbors than the leaf contains. + #[error("invalid leaf neighbor count {neighbors} for {points} points; maximum is {maximum}")] + InvalidNeighborCount { + /// Point count in the leaf. + points: usize, + /// Supplied output-column count. + neighbors: usize, + /// Maximum non-self neighbors per point. + maximum: usize, }, /// Temporary storage could not be reserved. #[error("failed to reserve {additional} values for {buffer}")] @@ -140,22 +147,27 @@ pub enum LeafKernelError { /// 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")] + /// A source did not contain enough rankable targets to fill its output. + #[error("source {source_index} has fewer than {neighbors} rankable leaf neighbors")] InsufficientRankableNeighbors { - /// Zero-based row position in the leaf. - row: usize, + /// Zero-based source position in the leaf. + source_index: usize, /// Required number of non-self neighbors. neighbors: usize, }, } -/// Return the required output length for [`LeafKernel::nearest_neighbors`]. -pub fn leaf_output_len(points: usize, k: usize) -> Result { +/// Return the usable non-self neighbor count for one leaf. +pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result { if points > u32::MAX as usize { return Err(LeafKernelError::TooManyPoints(points)); } - checked_area("output", points, k.min(points.saturating_sub(1))) + Ok(requested_k.min(points.saturating_sub(1))) +} + +/// Return the required output length for [`LeafKernel::nearest_neighbors`]. +pub fn leaf_output_len(points: usize, requested_k: usize) -> Result { + checked_area("output", points, leaf_neighbor_count(points, requested_k)?) } /// One invocation bundled for `Dispatched1`. @@ -165,10 +177,9 @@ pub fn leaf_output_len(points: usize, k: usize) -> Result { - input: LeafTopK<'a>, + input: LeafInput<'a>, output: MutMatrixView<'a, LeafNeighbor>, - workspace: &'a mut LeafTopKWorkspace, - requested_k: usize, + workspace: &'a mut LeafKernelWorkspace, } #[derive(Debug)] @@ -178,92 +189,48 @@ impl AddLifetime for LeafCallArg { type Of<'a> = LeafCall<'a>; } -type LeafFn = Dispatched1, LeafCallArg>; +type LeafFn = Dispatched1, LeafCallArg>; -/// A leaf kernel prepared for one metric, neighbor count, and the current CPU. +/// A leaf kernel prepared for one metric and the current CPU. /// /// Construct this once with [`LeafKernel::new`] and share it across leaf workers. -/// The handle stores only a direct function pointer and the requested-width mode. +/// Each output view carries its leaf-specific neighbor width. #[derive(Clone, Copy, Debug)] pub struct LeafKernel { run: LeafFn, - k: KValue, } impl LeafKernel { - /// Prepare a leaf kernel for `metric`, `k`, and the current CPU. - pub fn new(metric: Metric, k: usize) -> Self { - diskann_wide::arch::dispatch1_no_features( - PrepareLeaf { - k: KValue::from_requested(k), - }, - metric, - ) + /// Prepare a leaf kernel for `metric` and the current CPU. + pub fn new(metric: Metric) -> Self { + diskann_wide::arch::dispatch1_no_features(PrepareLeaf, metric) } - /// Select the nearest non-self leaf positions for every row. + /// Select the nearest non-self leaf positions for every source point. /// - /// `output` must have `input.dots.nrows()` rows and - /// `min(k, rows - 1)` columns. The returned value is that effective column - /// count. Equal distances retain pair scan order. + /// `output` must have one row per input point. Its column count is the + /// neighbor count for this leaf and must not exceed `point_count - 1`. + /// Equal distances retain pair scan order. pub fn nearest_neighbors( &self, - input: LeafTopK<'_>, + input: LeafInput<'_>, output: MutMatrixView<'_, LeafNeighbor>, - workspace: &mut LeafTopKWorkspace, - ) -> Result { + workspace: &mut LeafKernelWorkspace, + ) -> Result<(), LeafKernelError> { self.run.call(LeafCall { input, output, workspace, - requested_k: self.k.requested(), }) } } -/// Requested-width dispatch selected once while preparing the kernel. -/// -/// Widths one through three receive fixed array rows. Zero is a validated no-op; -/// larger values carry their requested width into the dynamic implementation. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum KValue { - Zero, - One, - Two, - Three, - Dynamic(usize), -} - -impl KValue { - const fn from_requested(k: usize) -> Self { - match k { - 0 => Self::Zero, - 1 => Self::One, - 2 => Self::Two, - 3 => Self::Three, - width => Self::Dynamic(width), - } - } - - const fn requested(self) -> usize { - match self { - Self::Zero => 0, - Self::One => 1, - Self::Two => 2, - Self::Three => 3, - Self::Dynamic(width) => width, - } - } -} - /// First dispatch stage: choose the runtime architecture once. /// /// The factory itself uses `dispatch1_no_features`; only the returned leaf entry /// needs target features, so architecture-specific code remains behind the final /// direct function pointer. -struct PrepareLeaf { - k: KValue, -} +struct PrepareLeaf; impl arch::Target1 for PrepareLeaf where @@ -273,35 +240,15 @@ where u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { fn run(self, arch: A, metric: Metric) -> LeafKernel { - visit_metric(metric, BuildLeaf { arch, k: self.k }) + visit_metric(metric, BuildLeaf(arch)) } } -/// BYO-type-erasure visitor holding a concrete architecture. +/// Metric visitor holding a concrete architecture. /// -/// `visit` receives a concrete metric marker, then combines `A`, `M`, and -/// the requested width into exactly one `Dispatched1`. -struct BuildLeaf { - arch: A, - k: KValue, -} - -impl BuildLeaf -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, -{ - fn build(self) -> LeafKernel { - LeafKernel { - run: self - .arch - .dispatch1::, Result, LeafCallArg>(), - k: self.k, - } - } -} +/// `visit` combines architecture `A` and concrete metric `M` into exactly +/// one `Dispatched1`. Leaf width remains call data because it varies by leaf. +struct BuildLeaf(A); impl MetricVisitor for BuildLeaf where @@ -313,37 +260,35 @@ where type Output = LeafKernel; fn visit(self) -> Self::Output { - match self.k { - KValue::Zero => self.build::(), - KValue::One => self.build::>(), - KValue::Two => self.build::>(), - KValue::Three => self.build::>(), - KValue::Dynamic(_) => self.build::(), + LeafKernel { + run: self + .0 + .dispatch1::, Result<(), LeafKernelError>, LeafCallArg>(), } } } -/// Architecture/metric/width-specialized function-pointer destination. +/// Architecture/metric-specialized function-pointer destination. /// -/// This type is zero-sized. All per-leaf state arrives through `LeafCall`; the -/// entry validates and initializes that state before reaching pointer-based SIMD. -struct LeafEntry(PhantomData<(M, S)>); +/// This type is zero-sized. All per-leaf state, including output width, arrives +/// through `LeafCall`; validation completes before pointer-based SIMD executes. +struct LeafEntry(PhantomData); -impl FTarget1, LeafCall<'_>> for LeafEntry +impl FTarget1, LeafCall<'_>> for LeafEntry where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, M: KernelMetric, - S: SlotSelection, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { - fn run(arch: A, mut call: LeafCall<'_>) -> Result { + fn run(arch: A, mut call: LeafCall<'_>) -> Result<(), LeafKernelError> { // Validation establishes every shape and active-prefix invariant used by // unchecked loads below. No output or scratch mutation occurs on error. - let actual_k = validate(call.input, call.requested_k, &call.output)?; - if actual_k == 0 { - return Ok(0); + validate(call.input, &call.output)?; + let neighbor_count = call.output.ncols(); + if neighbor_count == 0 { + return Ok(()); } // Norm and threshold scratch are reset for this leaf, while Vec capacity @@ -352,66 +297,75 @@ where call.output.as_mut_slice().fill(LeafNeighbor::default()); call.workspace.worst.fill(f32::INFINITY); - S::process::( + process_neighbor_width::( arch, call.input, - actual_k, + neighbor_count, call.output.as_mut_slice(), &call.workspace.norms, &mut call.workspace.worst, ); - if let Some(row) = call + if let Some(source) = call .output .as_slice() - .chunks_exact(actual_k) - .position(|neighbors| neighbors[actual_k - 1].position == u32::MAX) + .chunks_exact(neighbor_count) + .position(|neighbors| neighbors[neighbor_count - 1].target == u32::MAX) { return Err(LeafKernelError::InsufficientRankableNeighbors { - row, - neighbors: actual_k, + source_index: source, + neighbors: neighbor_count, }); } - Ok(actual_k) + Ok(()) } } /// Validate the complete safety contract before dispatched SIMD executes. /// /// Matrix views are rechecked with `checked_mul` because the hot loop performs -/// unchecked contiguous loads. Output columns must equal the clamped effective -/// k so fixed-row conversion cannot expose a partial row. +/// unchecked contiguous loads. Output columns are the leaf-specific neighbor +/// width and cannot exceed the number of non-self points. fn validate( - input: LeafTopK<'_>, - k: usize, + input: LeafInput<'_>, output: &MutMatrixView<'_, LeafNeighbor>, -) -> Result { - let rows = input.dots.nrows(); - let columns = input.dots.ncols(); - if rows != columns { +) -> Result<(), LeafKernelError> { + let point_count = input.dots.nrows(); + let dot_columns = input.dots.ncols(); + if point_count > u32::MAX as usize { + return Err(LeafKernelError::TooManyPoints(point_count)); + } + if point_count != dot_columns { return Err(LeafKernelError::NonSquareDots { - rows, - cols: columns, + rows: point_count, + cols: dot_columns, }); } - let output_len = leaf_output_len(rows, k)?; - let dots_len = checked_area("leaf dot-product matrix", rows, columns)?; + let dots_len = checked_area("leaf dot-product matrix", point_count, dot_columns)?; check_length( "leaf dot-product matrix", input.dots.as_slice().len(), dots_len, )?; + let output_len = checked_area("output", output.nrows(), output.ncols())?; + check_length("output", output.as_slice().len(), output_len)?; - let actual_k = k.min(rows.saturating_sub(1)); - if output.nrows() != rows || output.ncols() != actual_k { - return Err(LeafKernelError::InvalidOutputShape { - expected_rows: rows, - expected_cols: actual_k, - actual_rows: output.nrows(), - actual_cols: output.ncols(), + if output.nrows() != point_count { + return Err(LeafKernelError::InvalidOutputRows { + expected: point_count, + actual: output.nrows(), + columns: output.ncols(), }); } - check_length("output", output.as_slice().len(), output_len)?; - Ok(actual_k) + let maximum_neighbors = point_count.saturating_sub(1); + let neighbor_count = output.ncols(); + if neighbor_count > maximum_neighbors { + return Err(LeafKernelError::InvalidNeighborCount { + points: point_count, + neighbors: neighbor_count, + maximum: maximum_neighbors, + }); + } + Ok(()) } /// Prepare metric-specific scale and threshold scratch. @@ -420,14 +374,14 @@ fn validate( /// DiskANN's zero threshold. Normalized cosine and inner product skip the norm /// allocation entirely. `worst` is reset separately after allocation succeeds. fn prepare_workspace( - input: LeafTopK<'_>, - workspace: &mut LeafTopKWorkspace, + input: LeafInput<'_>, + workspace: &mut LeafKernelWorkspace, ) -> Result<(), LeafKernelError> { let points = input.dots.nrows(); if M::LEAF_SCALE.is_some() { resize("norms", &mut workspace.norms, points, 0.0)?; - for (row, norm) in workspace.norms.iter_mut().enumerate() { - *norm = M::LEAF_SCALE.transform(input.dots[(row, row)]); + for (source, norm) in workspace.norms.iter_mut().enumerate() { + *norm = M::LEAF_SCALE.transform(input.dots[(source, source)]); } } else { workspace.norms.clear(); @@ -475,93 +429,14 @@ fn check_length( } } -/// Prepared requested-width policy. -/// -/// The actual width can be smaller for singleton/tiny leaves, so each policy -/// performs one pre-loop clamp dispatch while keeping width selection out of the -/// pair scan. -trait SlotSelection: Send + Sync + 'static { - fn process( - arch: F::Arch, - input: LeafTopK<'_>, - actual_k: usize, - output: &mut [LeafNeighbor], - norms: &[f32], - worst: &mut [f32], - ) where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, - M: KernelMetric, - u64: From<<::BitMask as SIMDMask>::Underlying>; -} - -struct ZeroSelection; -struct FixedSelection; -struct DynamicSelection; - -impl SlotSelection for ZeroSelection { - fn process( - _arch: F::Arch, - _input: LeafTopK<'_>, - actual_k: usize, - _output: &mut [LeafNeighbor], - _norms: &[f32], - _worst: &mut [f32], - ) where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, - M: KernelMetric, - u64: From<<::BitMask as SIMDMask>::Underlying>, - { - debug_assert_eq!(actual_k, 0); - } -} - -impl SlotSelection for FixedSelection { - fn process( - arch: F::Arch, - input: LeafTopK<'_>, - actual_k: usize, - output: &mut [LeafNeighbor], - norms: &[f32], - worst: &mut [f32], - ) where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, - M: KernelMetric, - u64: From<<::BitMask as SIMDMask>::Underlying>, - { - debug_assert!(actual_k <= N); - process_selected::(arch, input, actual_k, output, norms, worst); - } -} - -impl SlotSelection for DynamicSelection { - fn process( - arch: F::Arch, - input: LeafTopK<'_>, - actual_k: usize, - output: &mut [LeafNeighbor], - norms: &[f32], - worst: &mut [f32], - ) where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, - M: KernelMetric, - u64: From<<::BitMask as SIMDMask>::Underlying>, - { - process_selected::(arch, input, actual_k, output, norms, worst); - } -} - -/// Convert effective k into one fixed row representation or the dynamic fallback. +/// Convert neighbor count into fixed source storage or the dynamic fallback. /// /// This branch runs once per leaf. Fixed conversion uses `as_chunks_mut` once, /// avoiding per-candidate slice-to-array checks while retaining safe insertion. -fn process_selected( +fn process_neighbor_width( arch: F::Arch, - input: LeafTopK<'_>, - actual_k: usize, + input: LeafInput<'_>, + neighbor_count: usize, output: &mut [LeafNeighbor], norms: &[f32], worst: &mut [f32], @@ -571,16 +446,16 @@ fn process_selected( M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { - match actual_k { - 1 => process_fixed::(arch, input, output, norms, worst), - 2 => process_fixed::(arch, input, output, norms, worst), - 3 => process_fixed::(arch, input, output, norms, worst), - width => process_pairs::( + match neighbor_count { + 1 => process_fixed_width::(arch, input, output, norms, worst), + 2 => process_fixed_width::(arch, input, output, norms, worst), + 3 => process_fixed_width::(arch, input, output, norms, worst), + dynamic_count => process_pairs::( arch, input, - DynamicRows { + DynamicNeighborStorage { values: output, - width, + neighbor_count: dynamic_count, }, norms, worst, @@ -588,9 +463,9 @@ fn process_selected( } } -fn process_fixed( +fn process_fixed_width( arch: F::Arch, - input: LeafTopK<'_>, + input: LeafInput<'_>, output: &mut [LeafNeighbor], norms: &[f32], worst: &mut [f32], @@ -600,77 +475,83 @@ fn process_fixed( M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let (rows, remainder) = output.as_chunks_mut::(); + let (neighbor_lists, remainder) = output.as_chunks_mut::(); debug_assert!(remainder.is_empty()); - process_pairs::(arch, input, FixedRows(rows), norms, worst); + process_pairs::( + arch, + input, + FixedNeighborStorage(neighbor_lists), + norms, + worst, + ); } -/// Mutable row adapter used by the shared pair traversal. +/// Mutable neighbor-list adapter used by the shared pair traversal. /// /// Implementations own the exclusive output borrow for the whole scan. Each -/// insertion borrows one row briefly, so updates to the current row and earlier -/// endpoint rows cannot alias simultaneously. -trait NeighborRows { - fn len(&self) -> usize; - fn insert(&mut self, row: usize, position: u32, distance: f32) -> f32; +/// insertion borrows one source list briefly, so updates to the current source +/// and earlier targets cannot alias simultaneously. +trait NeighborStorage { + fn source_count(&self) -> usize; + fn insert(&mut self, source: usize, target: u32, distance: f32) -> f32; } -struct FixedRows<'a, const N: usize>(&'a mut [[LeafNeighbor; N]]); +struct FixedNeighborStorage<'a, const N: usize>(&'a mut [[LeafNeighbor; N]]); -impl NeighborRows for FixedRows<'_, N> { +impl NeighborStorage for FixedNeighborStorage<'_, N> { #[inline(always)] - fn len(&self) -> usize { + fn source_count(&self) -> usize { self.0.len() } #[inline(always)] - fn insert(&mut self, row: usize, position: u32, distance: f32) -> f32 { - insert_fixed(&mut self.0[row], position, distance) + fn insert(&mut self, source: usize, target: u32, distance: f32) -> f32 { + insert_fixed_neighbor(&mut self.0[source], target, distance) } } -struct DynamicRows<'a> { +struct DynamicNeighborStorage<'a> { values: &'a mut [LeafNeighbor], - width: usize, + neighbor_count: usize, } -impl NeighborRows for DynamicRows<'_> { +impl NeighborStorage for DynamicNeighborStorage<'_> { #[inline(always)] - fn len(&self) -> usize { - self.values.len() / self.width + fn source_count(&self) -> usize { + self.values.len() / self.neighbor_count } #[inline(always)] - fn insert(&mut self, row: usize, position: u32, distance: f32) -> f32 { - insert_dynamic( - &mut self.values[row * self.width..(row + 1) * self.width], - position, + fn insert(&mut self, source: usize, target: u32, distance: f32) -> f32 { + insert_dynamic_neighbor( + &mut self.values[source * self.neighbor_count..(source + 1) * self.neighbor_count], + target, distance, ) } } -/// Scan the strict lower triangle once and update both endpoint rows. +/// Scan the strict lower triangle once and update both endpoint sources. /// /// Invariants on entry: /// /// - `dots` is a validated square row-major matrix; -/// - `output` has one sorted ascending-distance row per point; -/// - `worst[row]` equals that row's last slot; +/// - `output` has one sorted neighbor list per source point; +/// - `worst[source]` equals that source's last slot; /// - `norms` has one value per point exactly when `M` requires scales. /// /// Each SIMD chunk computes both endpoint eligibility masks before mutation. -/// Multiple lanes compete for the current row, so row candidates recheck its -/// live cached threshold. Every column lane targets a distinct earlier row and -/// can use the precomputed mask directly. Scalar tails call the matching scalar -/// metric operation to preserve established rounding semantics. +/// Multiple lanes compete for the current source, so source candidates recheck +/// its live cached threshold. Every target lane belongs to a distinct earlier +/// source and can use the precomputed mask directly. Scalar tails call the +/// matching scalar metric operation to preserve established rounding semantics. /// -/// `M` is concrete before type erasure. `R` presents fixed array rows for common -/// widths or safe dynamic slices for the uncommon fallback. +/// `M` is concrete before type erasure. `R` presents fixed neighbor arrays for +/// common counts or safe dynamic slices for the uncommon fallback. #[inline(never)] fn process_pairs( arch: F::Arch, - input: LeafTopK<'_>, + input: LeafInput<'_>, mut output: R, norms: &[f32], worst: &mut [f32], @@ -678,135 +559,138 @@ fn process_pairs( F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: KernelMetric, - R: NeighborRows, + R: NeighborStorage, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let points = input.dots.nrows(); + let point_count = input.dots.nrows(); let dots = input.dots.as_slice(); let uses_norms = M::LEAF_SCALE.is_some(); let worst_ptr = worst.as_mut_ptr(); - for row in 1..points { - let row_start = row * points; - let row_norm = if uses_norms { - F::splat(arch, norms[row]) + for source in 1..point_count { + let source_start = source * point_count; + let source_scale = if uses_norms { + F::splat(arch, norms[source]) } else { F::default(arch) }; - // SAFETY: `row < points == worst.len()` after validation. - 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 pair_dots = unsafe { F::load_simd(arch, dots.as_ptr().add(row_start + column)) }; - let column_norms = if uses_norms { - // SAFETY: the full chunk lies below `row <= norms.len()`. - unsafe { F::load_simd(arch, norms.as_ptr().add(column)) } + // SAFETY: `source < point_count == worst.len()` after validation. + let mut source_worst = unsafe { *worst_ptr.add(source) }; + let mut target = 0; + + while target + F::LANES <= source { + // SAFETY: the full chunk is contained in this source's strict-lower prefix. + let pair_dots = unsafe { F::load_simd(arch, dots.as_ptr().add(source_start + target)) }; + let target_scales = if uses_norms { + // SAFETY: the full target chunk lies below `source <= norms.len()`. + unsafe { F::load_simd(arch, norms.as_ptr().add(target)) } } else { F::default(arch) }; - let distances = M::leaf_distance(arch, pair_dots, row_norm, column_norms); - // Every pair may improve the current row and its earlier endpoint. - // Derive both masks from the same distance vector before either side - // mutates its threshold. - let row_eligible = distances.lt_simd(F::splat(arch, row_worst)); - // SAFETY: the full chunk lies below `row`, so it is inside `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 distances = M::leaf_distance(arch, pair_dots, source_scale, target_scales); + // Every pair may improve the current source and its earlier target. + // Derive both masks before either endpoint mutates its threshold. + let source_eligible = distances.lt_simd(F::splat(arch, source_worst)); + // SAFETY: the full target chunk lies below `source`, so it is inside `worst`. + let target_worst = unsafe { F::load_simd(arch, worst_ptr.add(target)) }; + let target_eligible = distances.lt_simd(target_worst); + let source_bits = u64::from(source_eligible.bitmask().to_underlying()); + let target_bits = u64::from(target_eligible.bitmask().to_underlying()); + + if source_bits | target_bits != 0 { let values = distances.to_array(); let values = values.as_ref(); - let mut row_bits = row_bits; - while row_bits != 0 { - let lane = row_bits.trailing_zeros() as usize; - row_bits &= row_bits - 1; + let mut source_bits = source_bits; + while source_bits != 0 { + let lane = source_bits.trailing_zeros() as usize; + source_bits &= source_bits - 1; let distance = values[lane]; - if distance < row_worst { - row_worst = output.insert(row, (column + lane) as u32, distance); + if distance < source_worst { + source_worst = output.insert(source, (target + 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; - let new_worst = output.insert(target, row as u32, values[lane]); - // SAFETY: `target < row < worst.len()`. - unsafe { *worst_ptr.add(target) = new_worst }; + let mut target_bits = target_bits; + while target_bits != 0 { + let lane = target_bits.trailing_zeros() as usize; + target_bits &= target_bits - 1; + let target_source = target + lane; + let new_worst = output.insert(target_source, source as u32, values[lane]); + // SAFETY: `target_source < source < worst.len()`. + unsafe { *worst_ptr.add(target_source) = new_worst }; } } - column += F::LANES; + target += F::LANES; } - while column < row { - // SAFETY: the scalar tail remains in the strict lower triangle. - let dot = unsafe { *dots.get_unchecked(row_start + column) }; - let (row_norm, column_norm) = if uses_norms { - // SAFETY: `column < row < points == norms.len()`. - (norms[row], unsafe { *norms.get_unchecked(column) }) + while target < source { + // SAFETY: the scalar target remains in this source's strict-lower prefix. + let dot = unsafe { *dots.get_unchecked(source_start + target) }; + let (source_scale, target_scale) = if uses_norms { + // SAFETY: `target < source < point_count == norms.len()`. + (norms[source], unsafe { *norms.get_unchecked(target) }) } else { (0.0, 0.0) }; - let distance = M::leaf_distance_scalar(dot, row_norm, column_norm); - if distance < row_worst { - row_worst = output.insert(row, column as u32, distance); + let distance = M::leaf_distance_scalar(dot, source_scale, target_scale); + if distance < source_worst { + source_worst = output.insert(source, target as u32, distance); } - // SAFETY: `column < row < worst.len()`. - let column_worst = unsafe { *worst_ptr.add(column) }; - if distance < column_worst { - let new_worst = output.insert(column, row as u32, distance); - // SAFETY: `column < row < worst.len()`. - unsafe { *worst_ptr.add(column) = new_worst }; + // SAFETY: `target < source < worst.len()`. + let target_worst = unsafe { *worst_ptr.add(target) }; + if distance < target_worst { + let new_worst = output.insert(target, source as u32, distance); + // SAFETY: `target < source < worst.len()`. + unsafe { *worst_ptr.add(target) = new_worst }; } - column += 1; + target += 1; } - // SAFETY: `row < worst.len()`. - unsafe { *worst_ptr.add(row) = row_worst }; + // SAFETY: `source < worst.len()`. + unsafe { *worst_ptr.add(source) = source_worst }; } - debug_assert_eq!(output.len(), points); + debug_assert_eq!(output.source_count(), point_count); } -/// Insert into a fixed-width row and return its new worst distance. +/// Insert into a fixed-width neighbor list and return its new worst distance. /// /// Production widths one through three use straight-line shifts. Strict `<` /// comparisons preserve scan order for ties; callers already rejected NaN via /// the eligibility comparison. #[inline(always)] -fn insert_fixed(row: &mut [LeafNeighbor; N], position: u32, distance: f32) -> f32 { - let entry = LeafNeighbor::new(position, distance); +fn insert_fixed_neighbor( + neighbors: &mut [LeafNeighbor; N], + target: u32, + distance: f32, +) -> f32 { + let entry = LeafNeighbor::new(target, distance); match N { 1 => { - row[0] = entry; + neighbors[0] = entry; distance } 2 => { - let first = row[0]; + let first = neighbors[0]; if distance < first.distance { - row[0] = entry; - row[1] = first; + neighbors[0] = entry; + neighbors[1] = first; first.distance } else { - row[1] = entry; + neighbors[1] = entry; distance } } 3 => { - let (first, second) = (row[0], row[1]); + let (first, second) = (neighbors[0], neighbors[1]); if distance < first.distance { - row[0] = entry; - row[1] = first; - row[2] = second; + neighbors[0] = entry; + neighbors[1] = first; + neighbors[2] = second; } else if distance < second.distance { - row[1] = entry; - row[2] = second; + neighbors[1] = entry; + neighbors[2] = second; } else { - row[2] = entry; + neighbors[2] = entry; return distance; } second.distance @@ -815,20 +699,20 @@ fn insert_fixed(row: &mut [LeafNeighbor; N], position: u32, dist } } -/// Insert into a run-time-width row using the same stable ordering contract. +/// Insert into a run-time-width neighbor list using the same stable ordering contract. /// /// The candidate replaces the last slot, then bubbles toward the front. This -/// path is used only for k greater than three. +/// path is used only for neighbor counts greater than three. #[inline(always)] -fn insert_dynamic(row: &mut [LeafNeighbor], position: u32, distance: f32) -> f32 { - let last = row.len() - 1; - row[last] = LeafNeighbor::new(position, distance); +fn insert_dynamic_neighbor(neighbors: &mut [LeafNeighbor], target: u32, distance: f32) -> f32 { + let last = neighbors.len() - 1; + neighbors[last] = LeafNeighbor::new(target, distance); let mut index = last; - while index > 0 && row[index].distance < row[index - 1].distance { - row.swap(index, index - 1); + while index > 0 && neighbors[index].distance < neighbors[index - 1].distance { + neighbors.swap(index, index - 1); index -= 1; } - row[last].distance + neighbors[last].distance } #[cfg(test)] @@ -837,68 +721,70 @@ mod tests { use super::*; - fn dots(metric: Metric, points: usize) -> Vec { + fn test_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 { + for source in 0..points { + dots[source * points + source] = if metric == Metric::Cosine && source == 0 { 0.0 } else { - 1.0 + (row % 5) as f32 + 1.0 + (source % 5) as f32 }; - for column in 0..row { - dots[row * points + column] = - (((row * 17 + column * 11) % 23) as f32 - 11.0) * 0.03125; + for target in 0..source { + dots[source * points + target] = + (((source * 17 + target * 11) % 23) as f32 - 11.0) * 0.03125; } } dots } - fn input(dots: &[f32], points: usize) -> LeafTopK<'_> { - LeafTopK { + fn test_input(dots: &[f32], points: usize) -> LeafInput<'_> { + LeafInput { dots: MatrixView::try_from(dots, points, points).unwrap(), } } - #[test] - fn k_value_preserves_requested_width() { - for (requested, value) in [ - (0, KValue::Zero), - (1, KValue::One), - (2, KValue::Two), - (3, KValue::Three), - (4, KValue::Dynamic(4)), - (17, KValue::Dynamic(17)), - ] { - assert_eq!(KValue::from_requested(requested), value); - assert_eq!(value.requested(), requested); - } - } - // Differential oracle for traversal and dispatch only. It intentionally // shares `M::leaf_distance_scalar`; public API tests independently spell // out metric formulas and full sorting behavior. fn scalar_traversal_reference( - input: LeafTopK<'_>, - k: usize, + input: LeafInput<'_>, + neighbor_count: usize, output: &mut [LeafNeighbor], ) { - let points = input.dots.nrows(); - let norms: Vec<_> = (0..points) - .map(|row| M::LEAF_SCALE.transform(input.dots[(row, row)])) + let point_count = input.dots.nrows(); + let norms: Vec<_> = (0..point_count) + .map(|source| M::LEAF_SCALE.transform(input.dots[(source, source)])) .collect(); - let mut worst = vec![f32::INFINITY; points]; + let mut worst = vec![f32::INFINITY; point_count]; let uses_norms = M::LEAF_SCALE.is_some(); - for row in 1..points { - for column in 0..row { - let (row_norm, column_norm) = if uses_norms { - (norms[row], norms[column]) + for source in 1..point_count { + for target in 0..source { + let (source_scale, target_scale) = if uses_norms { + (norms[source], norms[target]) } else { (0.0, 0.0) }; - let distance = - M::leaf_distance_scalar(input.dots[(row, column)], row_norm, column_norm); - insert_reference(output, &mut worst, k, row, column as u32, distance); - insert_reference(output, &mut worst, k, column, row as u32, distance); + let distance = M::leaf_distance_scalar( + input.dots[(source, target)], + source_scale, + target_scale, + ); + insert_reference( + output, + &mut worst, + neighbor_count, + source, + target as u32, + distance, + ); + insert_reference( + output, + &mut worst, + neighbor_count, + target, + source as u32, + distance, + ); } } } @@ -906,30 +792,36 @@ mod tests { fn insert_reference( output: &mut [LeafNeighbor], worst: &mut [f32], - k: usize, - row: usize, - position: u32, + neighbor_count: usize, + source: usize, + target: u32, distance: f32, ) { - if distance.partial_cmp(&worst[row]) != Some(std::cmp::Ordering::Less) { + if distance.partial_cmp(&worst[source]) != Some(std::cmp::Ordering::Less) { return; } - worst[row] = insert_dynamic(&mut output[row * k..(row + 1) * k], position, distance); + worst[source] = insert_dynamic_neighbor( + &mut output[source * neighbor_count..(source + 1) * neighbor_count], + target, + distance, + ); } - fn scalar_for_metric( + fn run_scalar_traversal( metric: Metric, - input: LeafTopK<'_>, - k: usize, + input: LeafInput<'_>, + neighbor_count: usize, output: &mut [LeafNeighbor], ) { match metric { - Metric::L2 => scalar_traversal_reference::(input, k, output), - Metric::Cosine => scalar_traversal_reference::(input, k, output), + Metric::L2 => scalar_traversal_reference::(input, neighbor_count, output), + Metric::Cosine => scalar_traversal_reference::(input, neighbor_count, output), Metric::CosineNormalized => { - scalar_traversal_reference::(input, k, output) + scalar_traversal_reference::(input, neighbor_count, output) + } + Metric::InnerProduct => { + scalar_traversal_reference::(input, neighbor_count, output) } - Metric::InnerProduct => scalar_traversal_reference::(input, k, output), } } @@ -937,22 +829,22 @@ mod tests { // Point count controls SIMD chunking. Cover both sides of 4-, 8-, and // 16-lane boundaries, then the boundary around a second 16-lane chunk. for points in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let dots = dots(metric, points); - let input = input(&dots, points); + let dots = test_dots(metric, points); + let input = test_input(&dots, points); for requested_k in [1, 2, 3, 4] { - let k = requested_k.min(points - 1); - let kernel = LeafKernel::new(metric, requested_k); - let mut expected = vec![LeafNeighbor::default(); points * k]; + let leaf_k = requested_k.min(points - 1); + let kernel = LeafKernel::new(metric); + let mut expected = vec![LeafNeighbor::default(); points * leaf_k]; kernel .nearest_neighbors( input, - MutMatrixView::try_from(expected.as_mut_slice(), points, k).unwrap(), - &mut LeafTopKWorkspace::new(), + MutMatrixView::try_from(expected.as_mut_slice(), points, leaf_k).unwrap(), + &mut LeafKernelWorkspace::new(), ) .unwrap(); - let mut actual = vec![LeafNeighbor::default(); points * k]; - scalar_for_metric(metric, input, k, &mut actual); + let mut actual = vec![LeafNeighbor::default(); points * leaf_k]; + run_scalar_traversal(metric, input, leaf_k, &mut actual); assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); } @@ -984,8 +876,8 @@ mod tests { 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_reference(&mut output, &mut worst, 4, 0, position, distance); + for (target, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 0.5)] { + insert_reference(&mut output, &mut worst, 4, 0, target, distance); } insert_reference(&mut output, &mut worst, 4, 0, 5, f32::NAN); @@ -1025,21 +917,42 @@ mod tests { ); } + #[test] + fn prepared_kernel_accepts_different_neighbor_counts() { + let points = 7; + let dots = test_dots(Metric::L2, points); + let input = test_input(&dots, points); + let kernel = LeafKernel::new(Metric::L2); + let mut workspace = LeafKernelWorkspace::new(); + + for neighbor_count in [1, 3, 2] { + let mut output = vec![LeafNeighbor::default(); points * neighbor_count]; + kernel + .nearest_neighbors( + input, + MutMatrixView::try_from(output.as_mut_slice(), points, neighbor_count).unwrap(), + &mut workspace, + ) + .unwrap(); + assert!(output.iter().all(|neighbor| neighbor.target != u32::MAX)); + } + } + #[test] fn workspace_can_shrink_and_grow_between_calls() { - let kernel = LeafKernel::new(Metric::L2, 2); - let mut workspace = LeafTopKWorkspace::new(); + let kernel = LeafKernel::new(Metric::L2); + let mut workspace = LeafKernelWorkspace::new(); for points in [17, 7, 17] { - let dots = dots(Metric::L2, points); + let dots = test_dots(Metric::L2, points); let mut output = vec![LeafNeighbor::default(); points * 2]; kernel .nearest_neighbors( - input(&dots, points), + test_input(&dots, points), MutMatrixView::try_from(output.as_mut_slice(), points, 2).unwrap(), &mut workspace, ) .unwrap(); - assert!(output.iter().all(|neighbor| neighbor.position != u32::MAX)); + assert!(output.iter().all(|neighbor| neighbor.target != u32::MAX)); } } } diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index 992d9fa28..c94193d5d 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -9,10 +9,10 @@ //! passes it to a [`PartitionKernel`] prepared once for the build metric. Kernel //! preparation selects the runtime architecture and concrete metric type once; //! repeated stripes call a direct `diskann-wide` function pointer with no ISA or -//! metric branch in the row loop. +//! metric branch in the point loop. //! //! L2 deliberately omits the point norm because it is constant across every -//! leader in one row. Cosine consumes squared point norms and leader norms. NaN +//! leader for that point. Cosine consumes squared point norms and leader norms. NaN //! distances are not rankable, and equal distances retain leader scan order. //! //! ```text @@ -25,8 +25,8 @@ //! shape/scale validation -> SIMD chunks + scalar tail -> sorted leader IDs //! ``` //! -//! Each row owns a fixed-capacity sorted tracker. Its last retained distance is -//! the rejection threshold, so noncompetitive SIMD chunks avoid lane extraction. +//! Each point owns a fixed-capacity sorted tracker. Its last retained distance +//! is the rejection threshold, so noncompetitive SIMD chunks avoid lane extraction. use std::marker::PhantomData; @@ -43,11 +43,11 @@ use crate::kernel_metric::{visit_metric, KernelMetric, MetricVisitor, ScaleKind} /// Maximum number of leaders retained for one point. /// /// Supported PiPNN partition fanouts fit within 16. Keeping this as a fixed -/// stack tracker bounds per-row stack use and code size; larger requests are +/// stack tracker bounds per-point stack use and code size; larger requests are /// rejected rather than silently truncated. pub const MAX_PARTITION_FANOUT: usize = 16; -type TopK = [(u32, f32); MAX_PARTITION_FANOUT]; +type LeaderTracker = [(u32, f32); MAX_PARTITION_FANOUT]; /// Metric-specific normalization inputs for one partition tile. #[derive(Clone, Copy, Debug)] @@ -59,8 +59,8 @@ pub enum PartitionScales<'a> { }, /// Unnormalized cosine needs squared point norms and leader norms. Cosine { - /// Squared norm for every point row. - row_squared_norms: &'a [f32], + /// Squared norm for every point. + point_squared_norms: &'a [f32], /// Norm for every leader column. leader_norms: &'a [f32], }, @@ -70,8 +70,8 @@ pub enum PartitionScales<'a> { /// One row-major point-by-leader dot-product tile. #[derive(Clone, Copy, Debug)] -pub struct PartitionTopK<'a> { - /// Point rows by leader columns. +pub struct PartitionInput<'a> { + /// One point per matrix row and one leader per column. pub dots: MatrixView<'a, f32>, /// Normalization inputs matching the prepared metric. pub scales: PartitionScales<'a>, @@ -120,24 +120,24 @@ pub enum PartitionKernelError { }, /// The requested fanout cannot be represented by the fixed top-k tracker. #[error( - "invalid fanout {fanout}: must not exceed {leaders} leaders or kernel maximum {maximum}" + "invalid fanout {fanout}: must not exceed {leader_count} leaders or kernel maximum {maximum}" )] InvalidFanout { - /// Requested number of leaders per row. + /// Requested number of leaders per point. fanout: usize, /// Available leader count. - leaders: usize, + leader_count: 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, + /// A point did not contain enough rankable leaders to fill its output. + #[error("point {point} has fewer than {fanout} rankable leaders")] + InsufficientRankableLeaders { + /// Zero-based point position in the input tile. + point: usize, /// Requested number of leader positions. fanout: usize, }, @@ -148,10 +148,10 @@ pub enum PartitionKernelError { /// Input and output receive independent call lifetimes. The prepared handle /// stores neither view, so it remains `Copy + Send + Sync` across worker threads. #[derive(Debug)] -struct PartitionInput; +struct PartitionInputArg; -impl AddLifetime for PartitionInput { - type Of<'a> = PartitionTopK<'a>; +impl AddLifetime for PartitionInputArg { + type Of<'a> = PartitionInput<'a>; } #[derive(Debug)] @@ -161,7 +161,8 @@ impl AddLifetime for PartitionOutput { type Of<'a> = MutMatrixView<'a, u32>; } -type PartitionFn = Dispatched2, PartitionInput, PartitionOutput>; +type PartitionFn = + Dispatched2, PartitionInputArg, PartitionOutput>; /// A partition kernel prepared for one metric and the current CPU. /// @@ -179,14 +180,14 @@ impl PartitionKernel { diskann_wide::arch::dispatch1_no_features(PreparePartition, metric) } - /// Select the nearest leader positions for every input row. + /// Select the nearest leader positions for every input point. /// /// `output.nrows()` must equal `input.dots.nrows()`; its column count is the /// requested fanout. Results are ordered by ascending distance. For L2, the - /// score omits the point norm because it cannot affect within-row ranking. + /// score omits the point norm because it cannot affect that point's ranking. pub fn nearest_leaders( &self, - input: PartitionTopK<'_>, + input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, ) -> Result<(), PartitionKernelError> { self.run.call(input, output) @@ -231,7 +232,7 @@ where run: self.0.dispatch2::< PartitionEntry, Result<(), PartitionKernelError>, - PartitionInput, + PartitionInputArg, PartitionOutput, >(), } @@ -241,10 +242,10 @@ where /// Architecture/metric-specialized function-pointer destination. /// /// The zero-sized entry receives all stripe state as arguments. Validation must -/// complete before `process_rows` reaches unchecked contiguous SIMD loads. +/// complete before `process_points` reaches unchecked contiguous SIMD loads. struct PartitionEntry(PhantomData); -impl FTarget2, PartitionTopK<'_>, MutMatrixView<'_, u32>> +impl FTarget2, PartitionInput<'_>, MutMatrixView<'_, u32>> for PartitionEntry where A: Architecture, @@ -255,7 +256,7 @@ where { fn run( arch: A, - input: PartitionTopK<'_>, + input: PartitionInput<'_>, mut output: MutMatrixView<'_, u32>, ) -> Result<(), PartitionKernelError> { // Validation establishes matrix areas, backing lengths, scale units, @@ -266,15 +267,15 @@ where return Ok(()); } - process_rows::(arch, input.dots, scales, fanout, output.as_mut_slice()); + process_points::(arch, input.dots, scales, fanout, output.as_mut_slice()); // A sorted tracker can be underfilled only at its last slot. This keeps - // post-validation linear in rows rather than scanning every output ID. - if let Some(row) = output + // post-validation linear in points rather than scanning every output ID. + if let Some(point) = output .as_slice() .chunks_exact(fanout) - .position(|leaders| leaders[fanout - 1] == u32::MAX) + .position(|assignments| assignments[fanout - 1] == u32::MAX) { - return Err(PartitionKernelError::InsufficientRankableDistances { row, fanout }); + return Err(PartitionKernelError::InsufficientRankableLeaders { point, fanout }); } Ok(()) } @@ -286,8 +287,8 @@ where /// on associated `ScaleKind` constants that monomorphize out of hot loops. #[derive(Clone, Copy)] struct ScaleSlices<'a> { - rows: &'a [f32], - leaders: &'a [f32], + point_scales: &'a [f32], + leader_scales: &'a [f32], } /// Validate the complete partition-kernel safety and metric contract. @@ -296,32 +297,32 @@ struct ScaleSlices<'a> { /// `PartitionScales` variant must match concrete metric `M`, preventing plausible /// but incorrect norm units from crossing the interface. fn validate<'a, M: KernelMetric>( - input: PartitionTopK<'a>, + input: PartitionInput<'a>, output: &MutMatrixView<'_, u32>, ) -> Result, PartitionKernelError> { - let rows = input.dots.nrows(); - let leaders = input.dots.ncols(); + let point_count = input.dots.nrows(); + let leader_count = input.dots.ncols(); let fanout = output.ncols(); - let dots_len = checked_area("dot-product tile", rows, leaders)?; + let dots_len = checked_area("dot-product tile", point_count, leader_count)?; check_length("dot-product tile", input.dots.as_slice().len(), dots_len)?; let output_len = checked_area("output", output.nrows(), fanout)?; check_length("output", output.as_slice().len(), output_len)?; - if output.nrows() != rows { + if output.nrows() != point_count { return Err(PartitionKernelError::InvalidOutputShape { - expected_rows: rows, + expected_rows: point_count, actual_rows: output.nrows(), actual_cols: output.ncols(), }); } - if leaders > u32::MAX as usize { - return Err(PartitionKernelError::TooManyLeaders(leaders)); + if leader_count > u32::MAX as usize { + return Err(PartitionKernelError::TooManyLeaders(leader_count)); } - if fanout > MAX_PARTITION_FANOUT || fanout > leaders { + if fanout > MAX_PARTITION_FANOUT || fanout > leader_count { return Err(PartitionKernelError::InvalidFanout { fanout, - leaders, + leader_count, maximum: MAX_PARTITION_FANOUT, }); } @@ -333,22 +334,22 @@ fn validate<'a, M: KernelMetric>( leader_squared_norms, }, ) => ScaleSlices { - rows: &[], - leaders: leader_squared_norms, + point_scales: &[], + leader_scales: leader_squared_norms, }, ( Metric::Cosine, PartitionScales::Cosine { - row_squared_norms, + point_squared_norms, leader_norms, }, ) => ScaleSlices { - rows: row_squared_norms, - leaders: leader_norms, + point_scales: point_squared_norms, + leader_scales: leader_norms, }, (Metric::CosineNormalized | Metric::InnerProduct, PartitionScales::None) => ScaleSlices { - rows: &[], - leaders: &[], + point_scales: &[], + leader_scales: &[], }, (Metric::L2, _) => return Err(PartitionKernelError::InvalidScales { expected: "L2" }), (Metric::Cosine, _) => { @@ -367,14 +368,14 @@ fn validate<'a, M: KernelMetric>( }; check_length( - "row scales", - scales.rows.len(), - expected_scale_len(M::PARTITION_ROW_SCALE, rows), + "point scales", + scales.point_scales.len(), + expected_scale_len(M::PARTITION_POINT_SCALE, point_count), )?; check_length( "leader scales", - scales.leaders.len(), - expected_scale_len(M::PARTITION_LEADER_SCALE, leaders), + scales.leader_scales.len(), + expected_scale_len(M::PARTITION_LEADER_SCALE, leader_count), )?; Ok(scales) } @@ -412,20 +413,20 @@ fn check_length( } } -/// Convert each point-to-leader dot-product row into sorted top-fanout IDs. +/// Convert each point's leader scores into sorted top-fanout IDs. /// -/// Per-row flow: +/// Per-point flow: /// -/// 1. transform the row scale once according to concrete metric `M`; -/// 2. process full SIMD chunks, rejecting lanes against the tracker's last slot; -/// 3. process the tail with the scalar metric operation; -/// 4. copy the sorted tracker prefix to that row's output. +/// 1. transform the point scale once according to concrete metric `M`; +/// 2. process full SIMD leader groups, rejecting lanes against the last slot; +/// 3. process the remaining leaders with the scalar metric operation; +/// 4. copy the sorted tracker prefix to that point's output. /// -/// `top[..fanout]` remains sorted after every accepted candidate. Strict `<` +/// `tracker[..fanout]` remains sorted after every accepted candidate. Strict `<` /// preserves leader scan order for ties and makes NaNs non-rankable. L2 keeps /// historical bulk-FMA/scalar-tail rounding because changing it can alter graph /// assignment at near ties. -fn process_rows( +fn process_points( arch: F::Arch, dots: MatrixView<'_, f32>, scales: ScaleSlices<'_>, @@ -437,67 +438,71 @@ fn process_rows( M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let leaders = dots.ncols(); - for (row, (dot_row, output_row)) in dots + let leader_count = dots.ncols(); + for (point, (point_dots, point_output)) in dots .as_slice() - .chunks_exact(leaders) + .chunks_exact(leader_count) .zip(output.chunks_exact_mut(fanout)) .enumerate() { - let row_scale = if M::PARTITION_ROW_SCALE.is_some() { - M::PARTITION_ROW_SCALE.transform(scales.rows[row]) + let point_scale = if M::PARTITION_POINT_SCALE.is_some() { + M::PARTITION_POINT_SCALE.transform(scales.point_scales[point]) } else { 0.0 }; - let row_scale_vector = F::splat(arch, row_scale); - let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; - let full = leaders / F::LANES * F::LANES; + let point_scale_vector = F::splat(arch, point_scale); + let mut tracker = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + let full = leader_count / F::LANES * F::LANES; for base in (0..full).step_by(F::LANES) { - // SAFETY: `base + F::LANES <= full <= dot_row.len()`. - let dots = unsafe { F::load_simd(arch, dot_row.as_ptr().add(base)) }; + // SAFETY: `base + F::LANES <= full <= point_dots.len()`. + let point_dots = unsafe { F::load_simd(arch, point_dots.as_ptr().add(base)) }; let leader_scales = if M::PARTITION_LEADER_SCALE.is_some() { - // SAFETY: validation requires one leader scale per dot-product column. - unsafe { F::load_simd(arch, scales.leaders.as_ptr().add(base)) } + // SAFETY: validation requires one scale per leader. + unsafe { F::load_simd(arch, scales.leader_scales.as_ptr().add(base)) } } else { F::default(arch) }; - insert_lanes( - M::partition_distance(arch, dots, row_scale_vector, leader_scales), + insert_leader_lanes( + M::partition_distance(arch, point_dots, point_scale_vector, leader_scales), base, - &mut top, + &mut tracker, fanout, ); } - for (leader, &dot) in dot_row.iter().enumerate().skip(full) { + for (leader, &dot) in point_dots.iter().enumerate().skip(full) { let leader_scale = if M::PARTITION_LEADER_SCALE.is_some() { - M::PARTITION_LEADER_SCALE.transform(scales.leaders[leader]) + M::PARTITION_LEADER_SCALE.transform(scales.leader_scales[leader]) } else { 0.0 }; - insert_topk( - &mut top, + insert_leader( + &mut tracker, fanout, leader as u32, - M::partition_distance_scalar(dot, row_scale, leader_scale), + M::partition_distance_scalar(dot, point_scale, leader_scale), ); } - copy_ids(&top, output_row); + copy_leader_ids(&tracker, point_output); } } -/// Offer competitive SIMD lanes to a row tracker in increasing leader order. +/// Offer competitive SIMD lanes to a point tracker in increasing leader order. /// /// The broadcast threshold avoids materializing lanes when none can improve the /// last slot. Bit iteration follows low-to-high lane order, preserving scalar tie /// behavior across SIMD widths. -fn insert_lanes(distances: F, base: usize, top: &mut TopK, fanout: usize) -where +fn insert_leader_lanes( + distances: F, + first_leader: usize, + tracker: &mut LeaderTracker, + fanout: usize, +) where F: SIMDVector + SIMDPartialOrd, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let threshold = F::splat(distances.arch(), top[fanout - 1].1); + let threshold = F::splat(distances.arch(), tracker[fanout - 1].1); let eligible = distances.lt_simd(threshold); if eligible.none() { return; @@ -509,7 +514,7 @@ where while lanes != 0 { let lane = lanes.trailing_zeros() as usize; lanes &= lanes - 1; - insert_topk(top, fanout, (base + lane) as u32, values[lane]); + insert_leader(tracker, fanout, (first_leader + lane) as u32, values[lane]); } } @@ -519,23 +524,23 @@ where /// not enter, so scan order is the deterministic tie breaker and the last slot /// remains both rejection threshold and underfill sentinel. #[inline(always)] -fn insert_topk(top: &mut TopK, fanout: usize, leader: u32, distance: f32) { +fn insert_leader(tracker: &mut LeaderTracker, fanout: usize, leader: u32, distance: f32) { let threshold = fanout - 1; - if distance.partial_cmp(&top[threshold].1) != Some(std::cmp::Ordering::Less) { + if distance.partial_cmp(&tracker[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; + tracker[threshold] = (leader, distance); + let mut slot = threshold; + while slot > 0 && tracker[slot].1 < tracker[slot - 1].1 { + tracker.swap(slot, slot - 1); + slot -= 1; } } /// Publish only leader IDs; distances stay private tracker state. -fn copy_ids(top: &TopK, output: &mut [u32]) { - for (destination, &(leader, _)) in output.iter_mut().zip(top) { +fn copy_leader_ids(tracker: &LeaderTracker, assignments: &mut [u32]) { + for (destination, &(leader, _)) in assignments.iter_mut().zip(tracker) { *destination = leader; } } @@ -546,20 +551,20 @@ mod tests { use super::*; - fn data(metric: Metric, leaders: usize) -> (Vec, Vec, Vec) { - let dots = (0..2 * leaders) + fn test_data(metric: Metric, leader_count: usize) -> (Vec, Vec, Vec) { + let dots = (0..2 * leader_count) .map(|index| (((index * 13 + 7) % 29) as f32 - 14.0) * 0.125) .collect(); - let row_scales = if metric == Metric::Cosine { + let point_scales = if metric == Metric::Cosine { vec![0.0, 16.0] } else { Vec::new() }; let leader_scales = match metric { - Metric::L2 => (0..leaders) + Metric::L2 => (0..leader_count) .map(|leader| ((leader + 1) as f32).powi(2)) .collect(), - Metric::Cosine => (0..leaders) + Metric::Cosine => (0..leader_count) .map(|leader| { if leader == 0 { 0.0 @@ -570,29 +575,29 @@ mod tests { .collect(), Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), }; - (dots, row_scales, leader_scales) + (dots, point_scales, leader_scales) } - fn input<'a>( + fn test_input<'a>( metric: Metric, dots: &'a [f32], - rows: usize, - leaders: usize, - row_scales: &'a [f32], + point_count: usize, + leader_count: usize, + point_scales: &'a [f32], leader_scales: &'a [f32], - ) -> PartitionTopK<'a> { + ) -> PartitionInput<'a> { let scales = match metric { Metric::L2 => PartitionScales::L2 { leader_squared_norms: leader_scales, }, Metric::Cosine => PartitionScales::Cosine { - row_squared_norms: row_scales, + point_squared_norms: point_scales, leader_norms: leader_scales, }, Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, }; - PartitionTopK { - dots: MatrixView::try_from(dots, rows, leaders).unwrap(), + PartitionInput { + dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), scales, } } @@ -601,7 +606,7 @@ mod tests { // It intentionally shares `M::partition_distance_scalar`; public API tests // independently spell out ranking formulas and full sorting behavior. fn scalar_traversal_reference( - input: PartitionTopK<'_>, + input: PartitionInput<'_>, fanout: usize, output: &mut [u32], ) { @@ -609,55 +614,55 @@ mod tests { PartitionScales::L2 { leader_squared_norms, } => ScaleSlices { - rows: &[], - leaders: leader_squared_norms, + point_scales: &[], + leader_scales: leader_squared_norms, }, PartitionScales::Cosine { - row_squared_norms, + point_squared_norms, leader_norms, } => ScaleSlices { - rows: row_squared_norms, - leaders: leader_norms, + point_scales: point_squared_norms, + leader_scales: leader_norms, }, PartitionScales::None => ScaleSlices { - rows: &[], - leaders: &[], + point_scales: &[], + leader_scales: &[], }, }; - let leaders = input.dots.ncols(); - for (row, (dot_row, output_row)) in input + let leader_count = input.dots.ncols(); + for (point, (point_dots, point_output)) in input .dots .as_slice() - .chunks_exact(leaders) + .chunks_exact(leader_count) .zip(output.chunks_exact_mut(fanout)) .enumerate() { - let row_scale = if M::PARTITION_ROW_SCALE.is_some() { - M::PARTITION_ROW_SCALE.transform(scales.rows[row]) + let point_scale = if M::PARTITION_POINT_SCALE.is_some() { + M::PARTITION_POINT_SCALE.transform(scales.point_scales[point]) } else { 0.0 }; - let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; - for (leader, &dot) in dot_row.iter().enumerate() { + let mut tracker = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + for (leader, &dot) in point_dots.iter().enumerate() { let leader_scale = if M::PARTITION_LEADER_SCALE.is_some() { - M::PARTITION_LEADER_SCALE.transform(scales.leaders[leader]) + M::PARTITION_LEADER_SCALE.transform(scales.leader_scales[leader]) } else { 0.0 }; - insert_topk( - &mut top, + insert_leader( + &mut tracker, fanout, leader as u32, - M::partition_distance_scalar(dot, row_scale, leader_scale), + M::partition_distance_scalar(dot, point_scale, leader_scale), ); } - copy_ids(&top, output_row); + copy_leader_ids(&tracker, point_output); } } - fn scalar_for_metric( + fn run_scalar_traversal( metric: Metric, - input: PartitionTopK<'_>, + input: PartitionInput<'_>, fanout: usize, output: &mut [u32], ) { @@ -676,12 +681,19 @@ mod tests { fn assert_scalar_reference_matches_prepared_dispatch(metric: Metric) { // Leader count controls SIMD chunking. Exercise both sides of 4-, 8-, and // 16-lane boundaries, then a second 16-lane chunk. - for leaders in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let (dots, row_scales, leader_scales) = data(metric, leaders); - let input = input(metric, &dots, 2, leaders, &row_scales, &leader_scales); + for leader_count in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let (dots, point_scales, leader_scales) = test_data(metric, leader_count); + let input = test_input( + metric, + &dots, + 2, + leader_count, + &point_scales, + &leader_scales, + ); let kernel = PartitionKernel::new(metric); for fanout in [1, 2, 6, MAX_PARTITION_FANOUT] { - if fanout > leaders { + if fanout > leader_count { continue; } let mut expected = vec![u32::MAX; 2 * fanout]; @@ -693,10 +705,10 @@ mod tests { .unwrap(); let mut actual = vec![u32::MAX; 2 * fanout]; - scalar_for_metric(metric, input, fanout, &mut actual); + run_scalar_traversal(metric, input, fanout, &mut actual); assert_eq!( actual, expected, - "{metric:?}, leaders={leaders}, k={fanout}" + "{metric:?}, leaders={leader_count}, k={fanout}" ); } } @@ -737,31 +749,31 @@ mod tests { #[test] fn cosine_special_norms_match_scalar_and_prepared_dispatch() { - let leaders = 17; - let dots = vec![1.0; 4 * leaders]; - let row_scales = [0.0, f32::MIN_POSITIVE / 2.0, f32::MIN_POSITIVE, f32::NAN]; - let mut leader_scales = vec![1.0; leaders]; + let leader_count = 17; + let point_scales = [0.0, f32::MIN_POSITIVE / 2.0, f32::MIN_POSITIVE, f32::NAN]; + let dots = vec![1.0; point_scales.len() * leader_count]; + let mut leader_scales = vec![1.0; leader_count]; leader_scales[..4].copy_from_slice(&[ 0.0, f32::MIN_POSITIVE.sqrt() / 2.0, f32::MIN_POSITIVE.sqrt(), f32::NAN, ]); - let input = input( + let input = test_input( Metric::Cosine, &dots, - row_scales.len(), - leaders, - &row_scales, + point_scales.len(), + leader_count, + &point_scales, &leader_scales, ); - let mut expected = vec![u32::MAX; row_scales.len() * 2]; + let mut expected = vec![u32::MAX; point_scales.len() * 2]; scalar_traversal_reference::(input, 2, &mut expected); - let mut actual = vec![u32::MAX; row_scales.len() * 2]; + let mut actual = vec![u32::MAX; point_scales.len() * 2]; PartitionKernel::new(Metric::Cosine) .nearest_leaders( input, - MutMatrixView::try_from(actual.as_mut_slice(), row_scales.len(), 2).unwrap(), + MutMatrixView::try_from(actual.as_mut_slice(), point_scales.len(), 2).unwrap(), ) .unwrap(); @@ -784,12 +796,12 @@ mod tests { #[test] fn scalar_topk_orders_candidates_and_preserves_ties() { - let mut top = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + let mut tracker = [(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_leader(&mut tracker, 4, leader, distance); } - insert_topk(&mut top, 4, 5, f32::NAN); + insert_leader(&mut tracker, 4, 5, f32::NAN); - assert_eq!(top[..4], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); + assert_eq!(tracker[..4], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); } } diff --git a/diskann-pipnn/tests/leaf_kernel_api.rs b/diskann-pipnn/tests/leaf_kernel_api.rs index f82bfe11d..07b6f85dc 100644 --- a/diskann-pipnn/tests/leaf_kernel_api.rs +++ b/diskann-pipnn/tests/leaf_kernel_api.rs @@ -6,7 +6,8 @@ use std::cmp::Ordering; use diskann_pipnn::leaf_kernel::{ - leaf_output_len, LeafKernel, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, + leaf_neighbor_count, leaf_output_len, LeafInput, LeafKernel, LeafKernelError, + LeafKernelWorkspace, LeafNeighbor, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -15,29 +16,30 @@ const SIMD_BOUNDARY_POINTS: [usize; 9] = [7, 8, 9, 15, 16, 17, 64, 256, 512]; const ZERO_NORM_POSITION: usize = 0; const DISTINCT_NORM_POSITION: usize = 2; const NORM_PERIOD: usize = 5; -const ROW_MIXER: usize = 17; -const COLUMN_MIXER: usize = 11; +const SOURCE_MIXER: usize = 17; +const TARGET_MIXER: usize = 11; const MIX_MODULUS: usize = 23; const MIX_CENTER: f32 = 11.0; const DOT_SCALE: f32 = 1.0 / 32.0; -const TIED_COLUMNS: [usize; 2] = [1, 2]; +const TIED_TARGETS: [usize; 2] = [1, 2]; -fn differential_input(metric: Metric, points: usize) -> Vec { +fn differential_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 == ZERO_NORM_POSITION { + for source in 0..points { + dots[source * points + source] = if metric == Metric::Cosine && source == ZERO_NORM_POSITION + { 0.0 - } else if row == DISTINCT_NORM_POSITION { + } else if source == DISTINCT_NORM_POSITION { 2.0 } else { - 1.0 + (row % NORM_PERIOD) as f32 + 1.0 + (source % NORM_PERIOD) as f32 }; - for column in 0..row { + for target in 0..source { let pair = - ((row * ROW_MIXER + column * COLUMN_MIXER) % MIX_MODULUS) as f32 - MIX_CENTER; - dots[row * points + column] = if row == points - 1 && column == 0 { + ((source * SOURCE_MIXER + target * TARGET_MIXER) % MIX_MODULUS) as f32 - MIX_CENTER; + dots[source * points + target] = if source == points - 1 && target == 0 { f32::NAN - } else if TIED_COLUMNS.contains(&column) { + } else if TIED_TARGETS.contains(&target) { 0.5 } else { pair * DOT_SCALE @@ -47,22 +49,27 @@ fn differential_input(metric: Metric, points: usize) -> Vec { dots } -fn input(dots: &[f32], points: usize) -> LeafTopK<'_> { - LeafTopK { +fn test_input(dots: &[f32], points: usize) -> LeafInput<'_> { + LeafInput { dots: MatrixView::try_from(dots, points, points).unwrap(), } } -fn reference(dots: &[f32], points: usize, requested_k: usize, metric: Metric) -> Vec { - let k = requested_k.min(points.saturating_sub(1)); - let mut output = vec![LeafNeighbor::default(); points * k]; - if k == 0 { +fn brute_force_reference( + dots: &[f32], + points: usize, + requested_k: usize, + metric: Metric, +) -> Vec { + let leaf_k = requested_k.min(points.saturating_sub(1)); + let mut output = vec![LeafNeighbor::default(); points * leaf_k]; + if leaf_k == 0 { return output; } let norms: Vec<_> = (0..points) - .map(|row| { - let diagonal = dots[row * points + row]; + .map(|source| { + let diagonal = dots[source * points + source]; if metric == Metric::Cosine { if diagonal < f32::MIN_POSITIVE { 0.0 @@ -75,25 +82,25 @@ fn reference(dots: &[f32], points: usize, requested_k: usize, metric: Metric) -> }) .collect(); - for row in 0..points { + for source in 0..points { let mut candidates = Vec::with_capacity(points - 1); - for position in 0..points { - if position == row { + for target in 0..points { + if target == source { continue; } - let (lower_row, lower_column) = if row > position { - (row, position) + let (lower_source, lower_target) = if source > target { + (source, target) } else { - (position, row) + (target, source) }; - let dot = dots[lower_row * points + lower_column]; + let dot = dots[lower_source * points + lower_target]; let clamp = |distance: f32| if distance < 0.0 { 0.0 } else { distance }; let distance = match metric { - Metric::L2 => clamp(norms[row] + norms[position] - 2.0 * dot), + Metric::L2 => clamp(norms[source] + norms[target] - 2.0 * dot), Metric::CosineNormalized => clamp(1.0 - dot), Metric::InnerProduct => -dot, Metric::Cosine => { - let denominator = norms[row] * norms[position]; + let denominator = norms[source] * norms[target]; let similarity = if denominator == 0.0 { 0.0 } else { @@ -103,7 +110,7 @@ fn reference(dots: &[f32], points: usize, requested_k: usize, metric: Metric) -> } }; if distance.partial_cmp(&f32::INFINITY) == Some(Ordering::Less) { - candidates.push(LeafNeighbor::new(position as u32, distance)); + candidates.push(LeafNeighbor::new(target as u32, distance)); } } candidates.sort_by(|left, right| { @@ -111,24 +118,28 @@ fn reference(dots: &[f32], points: usize, requested_k: usize, metric: Metric) -> .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]); + let count = candidates.len().min(leaf_k); + output[source * leaf_k..source * leaf_k + count].copy_from_slice(&candidates[..count]); } output } -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 returned_k = LeafKernel::new(metric, k) +fn run_kernel( + dots: &[f32], + points: usize, + requested_k: usize, + metric: Metric, +) -> (usize, Vec) { + let leaf_k = leaf_neighbor_count(points, requested_k).unwrap(); + let mut output = vec![LeafNeighbor::default(); points * leaf_k]; + LeafKernel::new(metric) .nearest_neighbors( - input(dots, points), - MutMatrixView::try_from(output.as_mut_slice(), points, actual_k).unwrap(), - &mut LeafTopKWorkspace::new(), + test_input(dots, points), + MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), + &mut LeafKernelWorkspace::new(), ) .unwrap(); - assert_eq!(returned_k, actual_k); - (returned_k, output) + (leaf_k, output) } #[test] @@ -140,10 +151,10 @@ fn prepared_dispatch_matches_reference_across_simd_width_boundaries() { Metric::InnerProduct, ] { for points in SIMD_BOUNDARY_POINTS { - let dots = differential_input(metric, points); + let dots = differential_dots(metric, points); for requested_k in [1, 2, 3, 4, 5] { - let expected = reference(&dots, points, requested_k, metric); - let actual = run(&dots, points, requested_k, metric).1; + let expected = brute_force_reference(&dots, points, requested_k, metric); + let actual = run_kernel(&dots, points, requested_k, metric).1; assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); } } @@ -161,7 +172,7 @@ fn l2_scans_only_the_lower_triangle_and_breaks_ties_by_position() { ]; assert_eq!( - run(&dots, 4, 2, Metric::L2).1, + run_kernel(&dots, 4, 2, Metric::L2).1, [ LeafNeighbor::new(1, 1.0), LeafNeighbor::new(2, 1.0), @@ -189,10 +200,10 @@ fn supports_every_leaf_metric() { (Metric::CosineNormalized, [1, 2, 1]), (Metric::InnerProduct, [1, 2, 1]), ] { - let positions: Vec<_> = run(&dots, 3, 1, metric) + let positions: Vec<_> = run_kernel(&dots, 3, 1, metric) .1 .iter() - .map(|neighbor| neighbor.position) + .map(|neighbor| neighbor.target) .collect(); assert_eq!(positions, expected, "metric {metric:?}"); } @@ -207,7 +218,7 @@ fn cosine_treats_zero_norm_as_zero_similarity() { 0.0, 0.0, 1.0, ]; - let output = run(&dots, 3, 2, Metric::Cosine).1; + let output = run_kernel(&dots, 3, 2, Metric::Cosine).1; assert_eq!(output[0], LeafNeighbor::new(1, 1.0)); assert_eq!(output[1], LeafNeighbor::new(2, 1.0)); } @@ -216,23 +227,35 @@ fn cosine_treats_zero_norm_as_zero_similarity() { 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, + run_kernel(&out_of_range, 2, 1, Metric::L2).1[0].distance, + 0.0 + ); + assert_eq!( + run_kernel(&out_of_range, 2, 1, Metric::CosineNormalized).1[0].distance, + 0.0 + ); + assert_eq!( + run_kernel(&out_of_range, 2, 1, Metric::Cosine).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); + assert_eq!( + run_kernel(&opposite, 2, 1, Metric::Cosine).1[0].distance, + 3.0 + ); let subnormal = [f32::MIN_POSITIVE / 2.0, 0.0, 1.0, 1.0]; - assert_eq!(run(&subnormal, 2, 1, Metric::Cosine).1[0].distance, 1.0); + assert_eq!( + run_kernel(&subnormal, 2, 1, Metric::Cosine).1[0].distance, + 1.0 + ); let minimum_normal = [f32::MIN_POSITIVE, 0.0, f32::MIN_POSITIVE.sqrt(), 1.0]; assert_eq!( - run(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, + run_kernel(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, 0.0 ); } @@ -243,10 +266,10 @@ fn finite_max_distance_fills_the_final_simd_slot() { 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); + let (leaf_k, output) = run_kernel(&dots, points, points - 1, Metric::InnerProduct); + assert_eq!(leaf_k, 8); assert_eq!( - output[8 * actual_k + actual_k - 1], + output[8 * leaf_k + leaf_k - 1], LeafNeighbor::new(0, f32::MAX) ); } @@ -266,28 +289,28 @@ fn every_metric_ignores_nan_pairs() { Metric::CosineNormalized, Metric::InnerProduct, ] { - let output = run(&dots, 3, 1, metric).1; - assert_eq!(output[0].position, 2, "metric {metric:?}"); - assert_eq!(output[1].position, 2, "metric {metric:?}"); + let output = run_kernel(&dots, 3, 1, metric).1; + assert_eq!(output[0].target, 2, "metric {metric:?}"); + assert_eq!(output[1].target, 2, "metric {metric:?}"); } } #[test] -fn rejects_incomplete_neighbor_rows() { +fn rejects_sources_with_too_few_rankable_neighbors() { let dots = [1.0, 0.0, f32::NAN, 1.0]; let mut output = [LeafNeighbor::default(); 2]; - let error = LeafKernel::new(Metric::L2, 1) + let error = LeafKernel::new(Metric::L2) .nearest_neighbors( - input(&dots, 2), + test_input(&dots, 2), MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafTopKWorkspace::new(), + &mut LeafKernelWorkspace::new(), ) .unwrap_err(); assert_eq!( error, LeafKernelError::InsufficientRankableNeighbors { - row: 0, + source_index: 0, neighbors: 1 } ); @@ -301,57 +324,70 @@ fn clamps_k_to_available_non_self_neighbors() { 0.0, 1.0, 3.0, 0.0, 0.0, 1.0, ]; - let (actual_k, output) = run(&dots, 3, 99, Metric::L2); + let (leaf_k, output) = run_kernel(&dots, 3, 99, Metric::L2); - assert_eq!(actual_k, 2); - for (row, neighbors) in output.chunks_exact(actual_k).enumerate() { + assert_eq!(leaf_k, 2); + for (source, neighbors) in output.chunks_exact(leaf_k).enumerate() { assert!(neighbors .iter() - .all(|neighbor| neighbor.position as usize != row)); + .all(|neighbor| neighbor.target as usize != source)); } } #[test] fn accepts_empty_singleton_and_zero_k_inputs() { - for (dots, points, k, metric) in [ + for (dots, points, requested_k, metric) in [ (&[][..], 0, 2, Metric::L2), (&[4.0][..], 1, 2, Metric::Cosine), (&[1.0, 0.0, 0.0, 1.0][..], 2, 0, Metric::InnerProduct), ] { - assert_eq!(run(dots, points, k, metric).0, 0); + assert_eq!(run_kernel(dots, points, requested_k, metric).0, 0); } } #[test] -fn rejects_non_square_input_and_wrong_output_shape() { +fn rejects_non_square_input_and_invalid_output_dimensions() { let dots = [0.0; 6]; - let non_square = LeafTopK { + let non_square = LeafInput { dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), }; let mut output = [LeafNeighbor::default(); 2]; - let kernel = LeafKernel::new(Metric::L2, 1); + let kernel = LeafKernel::new(Metric::L2); assert_eq!( kernel.nearest_neighbors( non_square, MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafTopKWorkspace::new(), + &mut LeafKernelWorkspace::new(), ), Err(LeafKernelError::NonSquareDots { rows: 2, cols: 3 }) ); let square = [0.0; 9]; - let mut wrong = [LeafNeighbor::default(); 3]; + let mut wrong_rows = [LeafNeighbor::default(); 2]; + assert_eq!( + kernel.nearest_neighbors( + test_input(&square, 3), + MutMatrixView::try_from(&mut wrong_rows[..], 2, 1).unwrap(), + &mut LeafKernelWorkspace::new(), + ), + Err(LeafKernelError::InvalidOutputRows { + expected: 3, + actual: 2, + columns: 1, + }) + ); + + let mut too_many = [LeafNeighbor::default(); 9]; assert_eq!( - LeafKernel::new(Metric::L2, 2).nearest_neighbors( - input(&square, 3), - MutMatrixView::try_from(&mut wrong[..], 3, 1).unwrap(), - &mut LeafTopKWorkspace::new(), + kernel.nearest_neighbors( + test_input(&square, 3), + MutMatrixView::try_from(&mut too_many[..], 3, 3).unwrap(), + &mut LeafKernelWorkspace::new(), ), - Err(LeafKernelError::InvalidOutputShape { - expected_rows: 3, - expected_cols: 2, - actual_rows: 3, - actual_cols: 1, + Err(LeafKernelError::InvalidNeighborCount { + points: 3, + neighbors: 3, + maximum: 2, }) ); } @@ -360,16 +396,16 @@ fn rejects_non_square_input_and_wrong_output_shape() { fn cosine_zero_norm_masks_nan_norm_at_simd_boundaries() { for points in [9, 17] { let mut dots = vec![0.0; points * points]; - for row in 1..points { - dots[row * points + row] = f32::NAN; + for source in 1..points { + dots[source * points + source] = f32::NAN; } - let output = run(&dots, points, 1, Metric::Cosine).1; - for (row, neighbor) in output.iter().enumerate().skip(1) { + let output = run_kernel(&dots, points, 1, Metric::Cosine).1; + for (source, neighbor) in output.iter().enumerate().skip(1) { assert_eq!( *neighbor, LeafNeighbor::new(0, 1.0), - "n={points}, row={row}" + "n={points}, source={source}" ); } } diff --git a/diskann-pipnn/tests/partition_kernel_api.rs b/diskann-pipnn/tests/partition_kernel_api.rs index e5318b8d1..59c69549b 100644 --- a/diskann-pipnn/tests/partition_kernel_api.rs +++ b/diskann-pipnn/tests/partition_kernel_api.rs @@ -4,58 +4,58 @@ */ use diskann_pipnn::partition_kernel::{ - PartitionKernel, PartitionKernelError, PartitionScales, PartitionTopK, MAX_PARTITION_FANOUT, + PartitionInput, PartitionKernel, PartitionKernelError, PartitionScales, MAX_PARTITION_FANOUT, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; -fn input<'a>( +fn test_input<'a>( metric: Metric, dots: &'a [f32], - rows: usize, - leaders: usize, - row_scales: &'a [f32], + point_count: usize, + leader_count: usize, + point_scales: &'a [f32], leader_scales: &'a [f32], -) -> PartitionTopK<'a> { +) -> PartitionInput<'a> { let scales = match metric { Metric::L2 => PartitionScales::L2 { leader_squared_norms: leader_scales, }, Metric::Cosine => PartitionScales::Cosine { - row_squared_norms: row_scales, + point_squared_norms: point_scales, leader_norms: leader_scales, }, Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, }; - PartitionTopK { - dots: MatrixView::try_from(dots, rows, leaders).unwrap(), + PartitionInput { + dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), scales, } } -fn reference(input: PartitionTopK<'_>, fanout: usize, metric: Metric) -> Vec { - let rows = input.dots.nrows(); - let leaders = input.dots.ncols(); - let (row_scales, leader_scales) = match input.scales { +fn brute_force_reference(input: PartitionInput<'_>, fanout: usize, metric: Metric) -> Vec { + let point_count = input.dots.nrows(); + let leader_count = input.dots.ncols(); + let (point_scales, leader_scales) = match input.scales { PartitionScales::L2 { leader_squared_norms, } => (&[][..], leader_squared_norms), PartitionScales::Cosine { - row_squared_norms, + point_squared_norms, leader_norms, - } => (row_squared_norms, leader_norms), + } => (point_squared_norms, leader_norms), PartitionScales::None => (&[][..], &[][..]), }; - let mut output = vec![u32::MAX; rows * fanout]; - for (row, (dots, output)) in input + let mut assignments = vec![u32::MAX; point_count * fanout]; + for (point, (point_dots, point_assignments)) in input .dots .as_slice() - .chunks_exact(leaders) - .zip(output.chunks_exact_mut(fanout)) + .chunks_exact(leader_count) + .zip(assignments.chunks_exact_mut(fanout)) .enumerate() { - let row_scale = row_scales.get(row).copied().unwrap_or(0.0); - let mut candidates: Vec<_> = dots + let point_scale = point_scales.get(point).copied().unwrap_or(0.0); + let mut candidates: Vec<_> = point_dots .iter() .enumerate() .filter_map(|(leader, &dot)| { @@ -65,15 +65,15 @@ fn reference(input: PartitionTopK<'_>, fanout: usize, metric: Metric) -> Vec 1.0 - dot, Metric::InnerProduct => -dot, Metric::Cosine => { - let row_norm = if row_scale < f32::MIN_POSITIVE { + let point_norm = if point_scale < f32::MIN_POSITIVE { 0.0 } else { - row_scale.sqrt() + point_scale.sqrt() }; - 1.0 - if row_norm == 0.0 || leader_scale == 0.0 { + 1.0 - if point_norm == 0.0 || leader_scale == 0.0 { 0.0 } else { - dot / (row_norm * leader_scale) + dot / (point_norm * leader_scale) } } }; @@ -82,35 +82,35 @@ fn reference(input: PartitionTopK<'_>, fanout: usize, metric: Metric) -> Vec (Vec, Vec, Vec) { - let dots = (0..2 * leaders) +fn differential_data(metric: Metric, leader_count: usize) -> (Vec, Vec, Vec) { + let dots = (0..2 * leader_count) .map(|index| { - let leader = index % leaders; - let row = index / leaders; - let base = ((leader * 13 + row * 7) % 19) as f32 - 9.0; + let leader = index % leader_count; + let point = index / leader_count; + let base = ((leader * 13 + point * 7) % 19) as f32 - 9.0; if leader == 2 || leader == 3 { 1.0 - } else if leader + 1 == leaders { + } else if leader + 1 == leader_count { f32::NAN } else { base * 0.25 } }) .collect(); - let row_scales = if metric == Metric::Cosine { + let point_scales = if metric == Metric::Cosine { vec![0.0, 16.0] } else { Vec::new() }; let leader_scales = match metric { - Metric::Cosine => (0..leaders) + Metric::Cosine => (0..leader_count) .map(|leader| { if leader == 1 { 0.0 @@ -121,7 +121,7 @@ fn differential_input(metric: Metric, leaders: usize) -> (Vec, Vec, Ve } }) .collect(), - Metric::L2 => (0..leaders) + Metric::L2 => (0..leader_count) .map(|leader| { let norm = if leader == 2 || leader == 3 { 3.0 @@ -133,12 +133,12 @@ fn differential_input(metric: Metric, leaders: usize) -> (Vec, Vec, Ve .collect(), Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), }; - (dots, row_scales, leader_scales) + (dots, point_scales, leader_scales) } fn run( metric: Metric, - input: PartitionTopK<'_>, + input: PartitionInput<'_>, fanout: usize, ) -> Result, PartitionKernelError> { let mut output = vec![u32::MAX; input.dots.nrows() * fanout]; @@ -157,17 +157,24 @@ fn prepared_dispatch_matches_reference_across_simd_width_boundaries() { Metric::CosineNormalized, Metric::InnerProduct, ] { - for leaders in [7, 8, 9, 15, 16, 17] { - let (dots, row_scales, leader_scales) = differential_input(metric, leaders); - let input = input(metric, &dots, 2, leaders, &row_scales, &leader_scales); + for leader_count in [7, 8, 9, 15, 16, 17] { + let (dots, point_scales, leader_scales) = differential_data(metric, leader_count); + let input = test_input( + metric, + &dots, + 2, + leader_count, + &point_scales, + &leader_scales, + ); for fanout in [1, 2, 16] { - if fanout >= leaders { + if fanout >= leader_count { continue; } assert_eq!( run(metric, input, fanout).unwrap(), - reference(input, fanout, metric), - "{metric:?}, leaders={leaders}, k={fanout}" + brute_force_reference(input, fanout, metric), + "{metric:?}, leaders={leader_count}, k={fanout}" ); } } @@ -184,7 +191,12 @@ fn l2_keeps_the_first_leader_when_boundary_distances_tie() { let norms = [0.0, 1.0, 4.0, 9.0]; assert_eq!( - run(Metric::L2, input(Metric::L2, &dots, 2, 4, &[], &norms), 2).unwrap(), + run( + Metric::L2, + test_input(Metric::L2, &dots, 2, 4, &[], &norms), + 2 + ) + .unwrap(), [0, 1, 2, 1] ); } @@ -196,7 +208,7 @@ fn supports_every_partition_metric() { 1.0, 0.0, -1.0, 2.0, 6.0, 0.0, ]; - for (metric, rows, leaders, expected) in [ + for (metric, point_scales, leader_scales, expected) in [ (Metric::L2, &[][..], &[1.0, 4.0, 9.0][..], [0, 1, 1, 0]), ( Metric::Cosine, @@ -208,7 +220,12 @@ fn supports_every_partition_metric() { (Metric::InnerProduct, &[][..], &[][..], [0, 1, 1, 0]), ] { assert_eq!( - run(metric, input(metric, &dots, 2, 3, rows, leaders), 2).unwrap(), + run( + metric, + test_input(metric, &dots, 2, 3, point_scales, leader_scales), + 2, + ) + .unwrap(), expected, "metric {metric:?}" ); @@ -220,7 +237,7 @@ fn cosine_treats_a_zero_norm_as_zero_similarity() { assert_eq!( run( Metric::Cosine, - input(Metric::Cosine, &[100.0, -100.0], 1, 2, &[0.0], &[1.0, 1.0]), + test_input(Metric::Cosine, &[100.0, -100.0], 1, 2, &[0.0], &[1.0, 1.0]), 2, ) .unwrap(), @@ -235,7 +252,7 @@ fn finite_max_distance_fills_the_final_simd_slot() { assert_eq!( run( Metric::InnerProduct, - input(Metric::InnerProduct, &dots, 1, 8, &[], &[]), + test_input(Metric::InnerProduct, &dots, 1, 8, &[], &[]), 8 ) .unwrap(), @@ -248,7 +265,7 @@ fn ignores_nan_distances_without_displacing_finite_leaders() { assert_eq!( run( Metric::InnerProduct, - input(Metric::InnerProduct, &[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), + test_input(Metric::InnerProduct, &[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), 2, ) .unwrap(), @@ -257,34 +274,37 @@ fn ignores_nan_distances_without_displacing_finite_leaders() { } #[test] -fn rejects_rows_with_too_few_rankable_distances() { +fn rejects_points_with_too_few_rankable_leaders() { assert_eq!( run( Metric::InnerProduct, - input(Metric::InnerProduct, &[f32::NAN, 3.0], 1, 2, &[], &[]), + test_input(Metric::InnerProduct, &[f32::NAN, 3.0], 1, 2, &[], &[]), 2, ), - Err(PartitionKernelError::InsufficientRankableDistances { row: 0, fanout: 2 }) + Err(PartitionKernelError::InsufficientRankableLeaders { + point: 0, + fanout: 2, + }) ); } #[test] -fn accepts_empty_rows_zero_fanout_and_largest_leader_id() { +fn accepts_empty_points_zero_fanout_and_largest_leader_id() { run( Metric::InnerProduct, - input(Metric::InnerProduct, &[], 0, 3, &[], &[]), + test_input(Metric::InnerProduct, &[], 0, 3, &[], &[]), 2, ) .unwrap(); run( Metric::InnerProduct, - input(Metric::InnerProduct, &[1.0, 2.0, 3.0], 1, 3, &[], &[]), + test_input(Metric::InnerProduct, &[1.0, 2.0, 3.0], 1, 3, &[], &[]), 0, ) .unwrap(); run( Metric::InnerProduct, - input(Metric::InnerProduct, &[], 0, u32::MAX as usize, &[], &[]), + test_input(Metric::InnerProduct, &[], 0, u32::MAX as usize, &[], &[]), 0, ) .unwrap(); @@ -293,7 +313,7 @@ fn accepts_empty_rows_zero_fanout_and_largest_leader_id() { assert_eq!( run( Metric::InnerProduct, - input( + test_input( Metric::InnerProduct, &[], 0, @@ -310,7 +330,7 @@ fn accepts_empty_rows_zero_fanout_and_largest_leader_id() { #[test] fn rejects_wrong_output_scales_and_fanout() { let dots = [0.0; 6]; - let valid_input = input(Metric::InnerProduct, &dots, 2, 3, &[], &[]); + let valid_input = test_input(Metric::InnerProduct, &dots, 2, 3, &[], &[]); let mut wrong_output = [u32::MAX; 3]; assert_eq!( PartitionKernel::new(Metric::InnerProduct).nearest_leaders( @@ -324,7 +344,7 @@ fn rejects_wrong_output_scales_and_fanout() { }) ); - let wrong_scales = PartitionTopK { + let wrong_scales = PartitionInput { dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), scales: PartitionScales::None, }; @@ -337,7 +357,7 @@ fn rejects_wrong_output_scales_and_fanout() { run(Metric::InnerProduct, valid_input, MAX_PARTITION_FANOUT + 1,), Err(PartitionKernelError::InvalidFanout { fanout: MAX_PARTITION_FANOUT + 1, - leaders: 3, + leader_count: 3, maximum: MAX_PARTITION_FANOUT, }) ); @@ -346,12 +366,12 @@ fn rejects_wrong_output_scales_and_fanout() { assert_eq!( run( Metric::InnerProduct, - input(Metric::InnerProduct, &one, 1, 1, &[], &[]), + test_input(Metric::InnerProduct, &one, 1, 1, &[], &[]), 2, ), Err(PartitionKernelError::InvalidFanout { fanout: 2, - leaders: 1, + leader_count: 1, maximum: MAX_PARTITION_FANOUT, }) ); From c76710d273b84ec0f06d6003cd66ca5fa8dffee6 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:32:21 +0000 Subject: [PATCH 20/26] docs(pipnn): explain kernel pipeline --- diskann-pipnn/src/kernel_metric.rs | 127 +++++++++++++- diskann-pipnn/src/leaf_kernel.rs | 236 +++++++++++++++++++++++++- diskann-pipnn/src/lib.rs | 123 ++++++++++++-- diskann-pipnn/src/partition_kernel.rs | 192 +++++++++++++++++++++ 4 files changed, 662 insertions(+), 16 deletions(-) diff --git a/diskann-pipnn/src/kernel_metric.rs b/diskann-pipnn/src/kernel_metric.rs index 01aaca0d3..6d13bfecd 100644 --- a/diskann-pipnn/src/kernel_metric.rs +++ b/diskann-pipnn/src/kernel_metric.rs @@ -5,9 +5,60 @@ //! Metric marker types shared by the partition and leaf kernels. //! +//! PiPNN uses dense matrix multiplication to produce dot products in two places: +//! partitioning compares dataset points with sampled leaders, while leaf building +//! compares every pair of points inside one small group. A dot product alone is +//! not always the requested distance. Squared L2 also needs squared norms; +//! unnormalized cosine needs norms; normalized cosine and inner product do not. +//! This module defines that conversion once so both kernels rank candidates with +//! identical scale units, zero handling, NaN handling, and scalar/SIMD formulas. +//! +//! [`KernelMetric`] is private because callers choose public [`Metric`] values, +//! not formula implementations. [`ScaleKind`] records what auxiliary value a +//! formula consumes. Zero-sized markers ([`L2`], [`Cosine`], +//! [`CosineNormalized`], [`InnerProduct`]) let dispatch compile one concrete +//! formula into each prepared kernel. [`MetricVisitor`] and [`visit_metric`] +//! perform the one-time runtime-to-concrete conversion. +//! //! Runtime metric selection happens only while preparing a dispatched kernel. //! The hot loops receive a concrete marker type, allowing metric arithmetic and //! scale handling to inline without a per-point or per-chunk enum match. +//! +//! Every helper converts an already-computed dot product into an +//! ascending-order score: +//! +//! | Metric | Leaf distance for source `s`, target `t` | Partition score for point `p`, leader `l` | Scale storage | +//! | --- | --- | --- | --- | +//! | squared L2 | `max(0, ‖s‖² + ‖t‖² - 2(s·t))` | `‖l‖² - 2(p·l)` | squared norms | +//! | cosine | `max(0, 1 - (s·t)/(‖s‖‖t‖))` | `1 - (p·l)/(‖p‖‖l‖)` | squared source/point norms; leader norms | +//! | normalized cosine | `max(0, 1 - s·t)` | `1 - p·l` | none | +//! | inner product | `-(s·t)` | `-(p·l)` | none | +//! +//! L2 partition ranking omits `‖p‖²`: that term is constant across all leaders +//! considered for one point and cannot change their order. +//! +//! # Core flow +//! +//! 1. [`visit_metric`] maps runtime [`Metric`] to a zero-sized marker. +//! 2. Leaf or partition preparation combines that marker with selected CPU +//! architecture. +//! 3. Final function pointer is monomorphized over both choices. +//! 4. SIMD bulk and scalar-tail calls share this module's metric contract. +//! +//! # Numerical behavior +//! +//! Subnormal norms are treated as zero before cosine division. A zero-norm +//! cosine endpoint forces zero similarity and distance `1.0`, even when the +//! other endpoint or dot is NaN. Otherwise NaN remains NaN, allowing strict +//! top-k comparisons to reject it. L2 scalar partition tails retain historical +//! non-fused operation order because rounding can change leader assignment at +//! near ties. +//! +//! # Performance +//! +//! Metric selection costs one match per prepared kernel, not per point or SIMD +//! chunk. Associated [`ScaleKind`] constants remove unused scale loads after +//! monomorphization. Distance helpers are constant-time and allocation-free. use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; @@ -34,6 +85,11 @@ impl ScaleKind { /// DiskANN treats subnormal squared norms, and corresponding subnormal /// norms, as zero before division. Ordered comparisons intentionally leave /// NaN unchanged so later distance comparisons keep it non-rankable. + /// + /// `stored` is interpreted according to `self`. The return value is zero, + /// the original norm, the original squared norm, or its square root. This + /// operation is constant-time and normally specializes to one match arm + /// because `ScaleKind` comes from a [`KernelMetric`] associated constant. #[inline(always)] pub(crate) fn transform(self, stored: f32) -> f32 { match self { @@ -56,6 +112,9 @@ impl ScaleKind { } } + /// Return whether callers must supply this scale position. + /// + /// Calls use an associated constant, so this test compiles out of hot loops. pub(crate) const fn is_some(self) -> bool { !matches!(self, Self::None) } @@ -67,6 +126,12 @@ impl ScaleKind { /// Generic methods then inline metric arithmetic into the architecture-specific /// function pointer. Leaf and partition operations remain separate because L2 /// partition ranking deliberately omits the point norm. +/// +/// All methods return scores ordered from nearest to farthest. Implementations +/// follow the module-level zero/NaN contract; caller-side strict comparisons +/// leave scores that remain NaN non-rankable. Marker types carry no data; +/// associated scale constants and forced inlining +/// remove metric branches from dispatched loops. pub(crate) trait KernelMetric: Send + Sync + 'static { /// Runtime tag represented by this marker. const METRIC: Metric; @@ -78,30 +143,54 @@ pub(crate) trait KernelMetric: Send + Sync + 'static { const PARTITION_LEADER_SCALE: ScaleKind; /// SIMD distance for one leaf source against a lane group of earlier targets. + /// + /// `arch` is the selected architecture token. `dot` and `target_scale` hold + /// one target per lane; `source_scale` broadcasts the source scale. Scale + /// arguments are zero when [`Self::LEAF_SCALE`] is [`ScaleKind::None`]. The + /// return value contains one ascending-order distance per lane. fn leaf_distance(arch: F::Arch, dot: F, source_scale: F, target_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; /// Scalar-tail equivalent of `leaf_distance`. + /// + /// Inputs and return value represent one SIMD lane. Operation order is part + /// of graph determinism where an implementation documents it. fn leaf_distance_scalar(dot: f32, source_scale: f32, target_scale: f32) -> f32; /// SIMD ranking score for one point against a lane group of leaders. + /// + /// `arch` is the selected architecture token. `dot` and `leader_scale` hold + /// one leader per lane; `point_scale` broadcasts one point scale. Scale + /// arguments are zero when the corresponding associated kind is + /// [`ScaleKind::None`]. The return value contains one ascending-order score + /// per lane; point-constant terms may be omitted. fn partition_distance(arch: F::Arch, dot: F, point_scale: F, leader_scale: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; /// Scalar-tail equivalent of `partition_distance`. + /// + /// Inputs and return value represent one SIMD lane. Implementations preserve + /// any documented non-fused order used by existing graph builds. fn partition_distance_scalar(dot: f32, point_scale: f32, leader_scale: f32) -> f32; } /// Zero-sized metric markers used only for monomorphization. pub(crate) struct L2; +/// Unnormalized-cosine marker. pub(crate) struct Cosine; +/// Unit-normalized-cosine marker. pub(crate) struct CosineNormalized; +/// Negative-inner-product marker. pub(crate) struct InnerProduct; +/// Clamp negative SIMD roundoff to zero while preserving NaN lanes. +/// +/// One ordered self-comparison normalizes backend-specific SIMD `max` NaN +/// behavior; no lane branches or allocations are introduced. #[inline(always)] fn clamp_nonnegative(arch: F::Arch, distance: F) -> F where @@ -116,6 +205,7 @@ where .select(zero.max_simd(distance), distance) } +/// Scalar equivalent of [`clamp_nonnegative`]. #[inline(always)] fn clamp_nonnegative_scalar(distance: f32) -> f32 { if distance < 0.0 { @@ -128,8 +218,12 @@ fn clamp_nonnegative_scalar(distance: f32) -> f32 { /// Compute cosine distance while preserving DiskANN zero/NaN semantics. /// /// Zero lanes divide by one only to keep the operation defined, then explicitly -/// select zero similarity. NaN norms fail the zero comparison and propagate -/// through division, leaving the final distance non-rankable. +/// select zero similarity. A NaN norm fails its own zero comparison and +/// propagates through division unless the other endpoint takes the zero-norm +/// path; in that case zero similarity takes precedence. +/// +/// `dot`, `source_norm`, and `target_norm` each contain one pair per lane. The +/// return value is `1 - cosine_similarity`. All lane handling is branchless. #[inline(always)] fn cosine_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -168,6 +262,8 @@ impl KernelMetric for L2 { F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + // Reconstruct squared L2 from Gram-matrix entries. Negative values can + // arise only from floating-point roundoff, so clamp without hiding NaN. clamp_nonnegative( arch, source_scale + target_scale - F::splat(arch, 2.0) * dot, @@ -176,6 +272,7 @@ impl KernelMetric for L2 { #[inline(always)] fn leaf_distance_scalar(dot: f32, source_scale: f32, target_scale: f32) -> f32 { + // Keep scalar tail arithmetic in the same left-to-right shape. clamp_nonnegative_scalar(source_scale + target_scale - 2.0 * dot) } @@ -185,6 +282,8 @@ impl KernelMetric for L2 { F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + // Point norm is constant for this ranking. Bulk lanes retain the + // historical fused multiply-add used by partition assignment. F::splat(arch, -2.0).mul_add_simd(dot, leader_scale) } @@ -208,11 +307,14 @@ impl KernelMetric for Cosine { F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + // Leaf output stores metric distances, so clamp negative roundoff after + // applying zero-norm and NaN handling in `cosine_distance`. clamp_nonnegative(arch, cosine_distance(arch, dot, source_scale, target_scale)) } #[inline(always)] fn leaf_distance_scalar(dot: f32, source_scale: f32, target_scale: f32) -> f32 { + // Match the bulk path's distance clamp for the scalar tail. clamp_nonnegative_scalar(cosine_distance_scalar(dot, source_scale, target_scale)) } @@ -222,11 +324,14 @@ impl KernelMetric for Cosine { F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + // Partitioning consumes only score order, so no post-formula clamp is + // needed; omitting it preserves existing near-tie behavior. cosine_distance(arch, dot, point_scale, leader_scale) } #[inline(always)] fn partition_distance_scalar(dot: f32, point_scale: f32, leader_scale: f32) -> f32 { + // Preserve the same unclamped ranking score in the scalar tail. cosine_distance_scalar(dot, point_scale, leader_scale) } } @@ -243,11 +348,14 @@ impl KernelMetric for CosineNormalized { F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + // Unit-normalized inputs need no scale loads; only roundoff below zero + // is clamped in stored leaf distances. clamp_nonnegative(arch, F::splat(arch, 1.0) - dot) } #[inline(always)] fn leaf_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { + // Scalar tail mirrors the normalized-cosine bulk formula. clamp_nonnegative_scalar(1.0 - dot) } @@ -257,11 +365,13 @@ impl KernelMetric for CosineNormalized { F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + // Ranking needs only `1 - dot`; no norm memory is touched. F::splat(arch, 1.0) - dot } #[inline(always)] fn partition_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { + // Preserve the unclamped ranking score used by full SIMD groups. 1.0 - dot } } @@ -278,11 +388,14 @@ impl KernelMetric for InnerProduct { F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + // Negation converts maximum inner product into the common ascending + // distance order without scale loads. F::default(arch) - dot } #[inline(always)] fn leaf_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { + // Scalar tail uses the same ascending score. -dot } @@ -292,11 +405,13 @@ impl KernelMetric for InnerProduct { F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + // Partition leader ranking shares the negative-inner-product score. F::default(arch) - dot } #[inline(always)] fn partition_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { + // Scalar tail uses the same ascending score. -dot } } @@ -306,15 +421,23 @@ impl KernelMetric for InnerProduct { /// The visitor receives concrete `M`, allowing architecture and width wrappers /// to compose with metric arithmetic before producing the final function pointer. /// This avoids a nested metric trait object inside architecture dispatch. +/// Visitor execution occurs once during preparation and allocates nothing. pub(crate) trait MetricVisitor { /// Final caller-selected erased representation. type Output; /// Consume the visitor with one concrete metric marker. + /// + /// Returns the caller-defined erased representation, normally one prepared + /// architecture/metric-specific function pointer. fn visit(self) -> Self::Output; } /// Visit the concrete marker represented by a runtime metric tag. +/// +/// `metric` selects exactly one concrete marker; `visitor` constructs and +/// returns its erased output. This performs one four-way match and no allocation +/// during kernel preparation. pub(crate) fn visit_metric(metric: Metric, visitor: V) -> V::Output { match metric { Metric::L2 => visitor.visit::(), diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 25f5e841b..0cb7e6843 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -5,6 +5,25 @@ //! Prepared nearest-neighbor kernels over a leaf's lower dot-product matrix. //! +//! PiPNN partitioning produces small, overlapping groups of dataset points +//! called *leaves*. Points sharing a leaf are treated as likely neighbors. For +//! each leaf, the builder gathers its vectors into a matrix `A` (one vector per +//! row). `sgemm_aat_lower` computes the lower triangle of the Gram matrix +//! `A · Aᵀ`, so entry `(i, j)` is the dot product of leaf points `i` and `j`. +//! This module consumes that result and picks each point's `k` nearest non-self +//! points in the leaf. +//! +//! This module does not gather vectors, run GEMM, translate dataset IDs, merge +//! candidates from overlapping leaves, or prune final graph degree. Its output +//! uses leaf-local positions. The caller maps those positions back through the +//! leaf's dataset-ID array and normally offers each selected pair in both graph +//! directions before cross-leaf merge/pruning. +//! +//! Here *source* means the point whose `k`-neighbor output list is being built; +//! *target* means another point in the same leaf. Pair distance is symmetric, so +//! one strict-lower matrix entry is evaluated once and offered independently to +//! both endpoint source lists. +//! //! `sgemm_aat_lower` writes pair `(source, target)` only when `target <= source`. //! The kernel scans that strict lower triangle once and offers each distance to //! both endpoint points. A [`LeafKernel`] is prepared once for the build metric @@ -23,9 +42,97 @@ //! shape validation -> scale scratch -> strict-lower scan -> sorted neighbor slots //! ``` //! -//! `workspace.worst[source]` always mirrors the last retained slot for that -//! source point. The SIMD loop may update both endpoints of a pair, so this -//! mirror is the threshold shared by source and target candidate masks. +//! Between source iterations, `workspace.worst[point]` mirrors that point's last +//! retained slot. While one source is scanned, its threshold lives in local +//! `source_worst`; thresholds for earlier target points are updated in the +//! workspace immediately. The SIMD loop snapshots both endpoint thresholds +//! before either list changes, then writes the current source threshold back +//! after its strict-lower prefix is complete. +//! +//! # Main structures +//! +//! - [`LeafKernel`] is the reusable public handle containing one prepared direct +//! function pointer. +//! - [`LeafInput`] identifies the borrowed lower-triangular dot matrix. +//! - [`LeafKernelWorkspace`] owns norm and rejection-threshold scratch and is +//! reused by one worker across leaves. +//! - [`LeafNeighbor`] is one output slot containing leaf-local target position +//! plus distance. +//! - `process_neighbor_width` chooses fixed storage for widths one through three +//! or dynamic storage for larger widths. +//! - `process_pairs` is the shared SIMD/scalar strict-lower traversal; +//! `insert_fixed_neighbor` and `insert_dynamic_neighbor` maintain stable sorted +//! output for both endpoints. +//! +//! # Inputs and output +//! +//! For `n` leaf points, [`LeafInput::dots`] is an `n × n` row-major matrix from +//! `sgemm_aat_lower`. Diagonal entries provide norms when the metric needs them; +//! only strict-lower entries `(source, target)` with `target < source` provide +//! pair dots. Output is an `n × k` [`LeafNeighbor`] matrix. Every output row is +//! sorted by ascending distance and stores leaf-local target positions, not +//! dataset IDs. +//! +//! Distances are reconstructed from one pair dot and, when required, diagonal +//! entries of the Gram matrix. Smaller is better: +//! +//! | Prepared metric | Leaf distance | +//! | --- | --- | +//! | squared L2 | `max(0, ‖source‖² + ‖target‖² - 2(source·target))` | +//! | cosine | `max(0, 1 - (source·target)/(‖source‖‖target‖))` | +//! | normalized cosine | `max(0, 1 - source·target)` | +//! | inner product | `-(source·target)` | +//! +//! `CosineNormalized` assumes leaf vectors were normalized before GEMM. For +//! unnormalized cosine, a zero/subnormal norm gives zero similarity. NaN scores +//! never enter output because selection uses strict ordered comparisons. +//! +//! # Core flow +//! +//! 1. Validate matrix areas, backing lengths, point IDs, and output width. +//! 2. Build metric scales from diagonal dots and reset per-source thresholds. +//! 3. Scan every strict-lower pair once in SIMD groups plus scalar tails. +//! 4. Offer that distance to both pair endpoints using stable top-k insertion. +//! 5. Reject any source whose final slot remains unfilled. +//! +//! # Performance +//! +//! With `k > 0`, the kernel evaluates exactly `n(n - 1) / 2` pair distances; +//! `k = 0` returns before traversal. Widths `k = 1, 2, 3` use fixed arrays and +//! straight-line insertion, giving `O(n²)` work. +//! Larger widths use `O(k)` insertion, giving `O(n²k)` worst-case work. Scratch +//! is `O(n)` (`worst`, plus norms only when required); output is `O(nk)`. No +//! allocation occurs after a worker workspace has sufficient capacity. Runtime +//! architecture and metric selection happen once in [`LeafKernel::new`]. +//! +//! # Example +//! +//! ``` +//! use diskann_pipnn::leaf_kernel::{ +//! leaf_output_len, LeafInput, LeafKernel, LeafKernelWorkspace, LeafNeighbor, +//! }; +//! use diskann_utils::views::{MatrixView, MutMatrixView}; +//! use diskann_vector::distance::Metric; +//! +//! // Only the diagonal and strict lower triangle are consumed. +//! let dots = [ +//! 1.0, f32::NAN, f32::NAN, +//! 0.9, 1.0, f32::NAN, +//! 0.1, 0.2, 1.0, +//! ]; +//! let input = LeafInput { +//! dots: MatrixView::try_from(&dots[..], 3, 3).unwrap(), +//! }; +//! let mut neighbors = vec![LeafNeighbor::default(); leaf_output_len(3, 1).unwrap()]; +//! let output = MutMatrixView::try_from(&mut neighbors[..], 3, 1).unwrap(); +//! let mut workspace = LeafKernelWorkspace::new(); +//! +//! LeafKernel::new(Metric::CosineNormalized) +//! .nearest_neighbors(input, output, &mut workspace) +//! .unwrap(); +//! +//! assert_eq!(neighbors.iter().map(|neighbor| neighbor.target).collect::>(), [1, 0, 1]); +//! ``` use std::marker::PhantomData; @@ -50,6 +157,9 @@ pub struct LeafNeighbor { impl LeafNeighbor { /// Construct a leaf-local neighbor. + /// + /// `target` is a position in the current leaf and `distance` is its score + /// from the source represented by the containing output row. pub const fn new(target: u32, distance: f32) -> Self { Self { target, distance } } @@ -77,6 +187,9 @@ pub struct LeafKernelWorkspace { impl LeafKernelWorkspace { /// Construct an empty workspace. + /// + /// This does not allocate. First use grows buffers to the leaf point count; + /// later calls reuse capacity owned by the same worker. pub const fn new() -> Self { Self { norms: Vec::new(), @@ -158,6 +271,19 @@ pub enum LeafKernelError { } /// Return the usable non-self neighbor count for one leaf. +/// +/// `points` is the leaf point count and `requested_k` is the build-wide target. +/// The returned width is `min(requested_k, points - 1)`, allowing empty, +/// singleton, and small leaves without a second effective-k state. +/// +/// # Errors +/// +/// Returns [`LeafKernelError::TooManyPoints`] when leaf-local positions cannot +/// fit in `u32`. +/// +/// # Performance +/// +/// Constant-time and allocation-free. pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result { if points > u32::MAX as usize { return Err(LeafKernelError::TooManyPoints(points)); @@ -166,6 +292,18 @@ pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result Result { checked_area("output", points, leaf_neighbor_count(points, requested_k)?) } @@ -195,6 +333,9 @@ type LeafFn = Dispatched1, LeafCallArg>; /// /// Construct this once with [`LeafKernel::new`] and share it across leaf workers. /// Each output view carries its leaf-specific neighbor width. +/// +/// The handle stores only one direct function pointer. It borrows no leaf data +/// or workspace and is therefore `Copy`, `Send`, and `Sync`. #[derive(Clone, Copy, Debug)] pub struct LeafKernel { run: LeafFn, @@ -202,6 +343,14 @@ pub struct LeafKernel { impl LeafKernel { /// Prepare a leaf kernel for `metric` and the current CPU. + /// + /// The returned handle contains one architecture/metric-specialized function + /// pointer and can process any valid leaf size or neighbor width. + /// + /// # Performance + /// + /// Performs runtime architecture detection and one metric match once. + /// Reusing the handle keeps both decisions out of per-leaf hot loops. pub fn new(metric: Metric) -> Self { diskann_wide::arch::dispatch1_no_features(PrepareLeaf, metric) } @@ -211,6 +360,30 @@ impl LeafKernel { /// `output` must have one row per input point. Its column count is the /// neighbor count for this leaf and must not exceed `point_count - 1`. /// Equal distances retain pair scan order. + /// + /// `input` supplies the square lower-triangular dot matrix. `output` is + /// overwritten with sorted leaf-local neighbors. `workspace` is an exclusive + /// worker-owned scratch lease whose capacity is retained after return. + /// Successful return guarantees every source has exactly `output.ncols()` + /// rankable, non-self neighbors. + /// + /// # Core flow + /// + /// The prepared entry validates every view before mutation, prepares scales, + /// clears output and thresholds, scans the strict lower triangle once, then + /// verifies the final slot of every source. Each pair updates both endpoints. + /// + /// # Errors + /// + /// Returns [`LeafKernelError`] for invalid or overflowing shapes, excessive + /// point/neighbor counts, scratch allocation failure, or an underfilled + /// source caused by non-rankable distances. Validation errors leave output + /// and workspace contents unchanged. + /// + /// # Performance + /// + /// See module-level complexity. This call uses the prepared direct function + /// pointer; it performs no runtime ISA or metric dispatch. pub fn nearest_neighbors( &self, input: LeafInput<'_>, @@ -272,6 +445,11 @@ where /// /// This type is zero-sized. All per-leaf state, including output width, arrives /// through `LeafCall`; validation completes before pointer-based SIMD executes. +/// +/// Call order is fixed: validate without mutation, allocate/reset scratch, +/// initialize output, execute one specialized traversal, then verify fill state. +/// Keeping those phases in the dispatched destination makes every unchecked +/// load depend on one visible validation gate. struct LeafEntry(PhantomData); impl FTarget1, LeafCall<'_>> for LeafEntry @@ -287,6 +465,8 @@ where // unchecked loads below. No output or scratch mutation occurs on error. validate(call.input, &call.output)?; let neighbor_count = call.output.ncols(); + // Empty or singleton leaves request zero columns. Avoid touching scratch + // or output so this path remains allocation-free. if neighbor_count == 0 { return Ok(()); } @@ -297,6 +477,8 @@ where call.output.as_mut_slice().fill(LeafNeighbor::default()); call.workspace.worst.fill(f32::INFINITY); + // Width dispatch happens once per leaf. Common production widths become + // fixed arrays; uncommon widths retain the same traversal through slices. process_neighbor_width::( arch, call.input, @@ -305,6 +487,8 @@ where &call.workspace.norms, &mut call.workspace.worst, ); + // Sorted lists use the last slot as both worst-distance threshold and + // underfill sentinel, so one slot check per source proves full output. if let Some(source) = call .output .as_slice() @@ -325,6 +509,12 @@ where /// Matrix views are rechecked with `checked_mul` because the hot loop performs /// unchecked contiguous loads. Output columns are the leaf-specific neighbor /// width and cannot exceed the number of non-self points. +/// +/// `input` and `output` are borrowed only for inspection. Success returns no +/// value; it establishes square dots, exact backing lengths, representable local +/// IDs, and valid output width. Failure returns [`LeafKernelError`] before any +/// output or workspace mutation. Runtime is constant apart from view metadata +/// checks; matrix contents are not scanned. fn validate( input: LeafInput<'_>, output: &MutMatrixView<'_, LeafNeighbor>, @@ -373,6 +563,11 @@ fn validate( /// L2 stores diagonal squared norms; cosine converts diagonals to norms using /// DiskANN's zero threshold. Normalized cosine and inner product skip the norm /// allocation entirely. `worst` is reset separately after allocation succeeds. +/// +/// `input` supplies diagonal dots and `workspace` owns reusable vectors. Success +/// prepares one scale and one threshold per point when needed; allocation failure +/// is returned without entering SIMD traversal. Work is `O(n)`, with at most +/// `O(n)` retained capacity per buffer. fn prepare_workspace( input: LeafInput<'_>, workspace: &mut LeafKernelWorkspace, @@ -433,6 +628,12 @@ fn check_length( /// /// This branch runs once per leaf. Fixed conversion uses `as_chunks_mut` once, /// avoiding per-candidate slice-to-array checks while retaining safe insertion. +/// +/// `output` contains `point_count * neighbor_count` initialized slots; `norms` +/// and `worst` satisfy the invariants established by `prepare_workspace`. The +/// function writes output and thresholds in place and returns no value. Widths +/// one through three take the fixed path; all others pay one division per source +/// insertion to locate its dynamic slice. fn process_neighbor_width( arch: F::Arch, input: LeafInput<'_>, @@ -463,6 +664,11 @@ fn process_neighbor_width( } } +/// Reinterpret validated output as one fixed array per source, then run shared +/// pair traversal. +/// +/// `N` is one, two, or three. `as_chunks_mut` performs one safe shape split per +/// leaf, keeping array conversion out of candidate insertion. fn process_fixed_width( arch: F::Arch, input: LeafInput<'_>, @@ -492,7 +698,10 @@ fn process_fixed_width( /// insertion borrows one source list briefly, so updates to the current source /// and earlier targets cannot alias simultaneously. trait NeighborStorage { + /// Number of source neighbor lists owned by this adapter. fn source_count(&self) -> usize; + + /// Insert one source-target candidate and return that source's new threshold. fn insert(&mut self, source: usize, target: u32, distance: f32) -> f32; } @@ -548,6 +757,14 @@ impl NeighborStorage for DynamicNeighborStorage<'_> { /// /// `M` is concrete before type erasure. `R` presents fixed neighbor arrays for /// common counts or safe dynamic slices for the uncommon fallback. +/// +/// `input` supplies `n × n` dots, `output` owns `n` sorted lists, `norms` holds +/// metric scales when required, and `worst` mirrors every list's final distance. +/// The function mutates output and thresholds in place and returns no value. +/// It evaluates exactly `n(n - 1) / 2` pairs. SIMD computes up to `F::LANES` +/// distances together; accepted candidates still insert in scan order to keep +/// deterministic ties. Fixed widths cost constant work per accepted endpoint; +/// dynamic widths cost `O(k)` per insertion. #[inline(never)] fn process_pairs( arch: F::Arch, @@ -567,8 +784,12 @@ fn process_pairs( let uses_norms = M::LEAF_SCALE.is_some(); let worst_ptr = worst.as_mut_ptr(); + // `source` starts at one because source zero has no strict-lower targets; + // later sources still offer their pair back to source zero. for source in 1..point_count { let source_start = source * point_count; + // `uses_norms` comes from a metric associated constant. Specialization + // removes both branch and scale memory traffic for scale-free metrics. let source_scale = if uses_norms { F::splat(arch, norms[source]) } else { @@ -657,6 +878,11 @@ fn process_pairs( /// Production widths one through three use straight-line shifts. Strict `<` /// comparisons preserve scan order for ties; callers already rejected NaN via /// the eligibility comparison. +/// +/// `neighbors` is the sorted list for one source. `target` and `distance` are a +/// candidate already known to beat its final slot. The return value is the new +/// final-slot distance. Insertion is allocation-free and constant-time because +/// `N <= 3`. #[inline(always)] fn insert_fixed_neighbor( neighbors: &mut [LeafNeighbor; N], @@ -703,6 +929,10 @@ fn insert_fixed_neighbor( /// /// The candidate replaces the last slot, then bubbles toward the front. This /// path is used only for neighbor counts greater than three. +/// +/// `neighbors` is one non-empty sorted source list. `target` and `distance` are +/// already known to beat its final slot. The return value is the new final-slot +/// distance. Work is `O(k)` worst-case and allocation-free. #[inline(always)] fn insert_dynamic_neighbor(neighbors: &mut [LeafNeighbor], target: u32, distance: f32) -> f32 { let last = neighbors.len() - 1; diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index 9cfc3e25e..d02ff0ca3 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -3,13 +3,68 @@ * Licensed under the MIT license. */ -//! Provider-independent PiPNN graph construction. +//! Numerical kernels for provider-independent PiPNN graph construction. //! -//! The crate owns overlapping partition generation, leaf-local nearest-neighbor -//! construction, candidate merging, and optional graph-degree finalization. The -//! caller supplies contiguous data, DiskANN graph policy, and the Rayon pool. -//! Providers, start/frozen points, quantization, persistence, and search remain -//! outside this algorithm seam. +//! PiPNN means **Pick-in-Partitions Nearest Neighbors**. The wider algorithm +//! builds a graph for approximate nearest-neighbor search: every input vector +//! becomes one graph vertex, and its adjacency list stores other vectors worth +//! visiting during a later query. The APIs exposed in this layer provide PiPNN's +//! numerical selection kernels; they do not yet expose the full graph builder or +//! execute queries. +//! +//! Incremental builders such as Vamana find construction candidates by running +//! beam search against a partially built graph: they repeatedly follow graph +//! edges to discover nearby vertices, causing random memory access. PiPNN removes +//! that search from construction and uses three bulk stages instead: +//! +//! 1. **Partition.** Randomized Ball Carving samples points called *leaders*. +//! Every point is assigned to its nearest `fanout` leaders. Assigning to more +//! than one leader makes child groups overlap. Oversized groups are processed +//! recursively until bounded groups called *leaves* remain. +//! 2. **Pick within leaves.** Vectors in one leaf are contiguous enough for a +//! dense matrix multiplication to compute all pair dot products. Each point +//! picks its nearest leaf companions; selected pairs become candidate graph +//! edges. +//! 3. **Merge and prune.** Candidates from overlapping leaves are combined. +//! HashPrune can keep a bounded reservoir per source while edges stream in, +//! retaining the closest candidate for each residual-direction hash. The +//! alternative collects unique candidates directly. An optional final Vamana +//! RobustPrune selects a bounded, directionally diverse adjacency list. +//! +//! ```text +//! dataset points +//! │ +//! v +//! sample leaders + point/leader GEMM +//! │ +//! v +//! choose nearest leaders ──> overlapping child groups ──> recurse ──> leaves +//! │ +//! leaf all-pairs GEMM +//! │ +//! v +//! pick local neighbors +//! │ +//! v +//! merge/prune edges +//! │ +//! v +//! search graph +//! ``` +//! +//! This crate keeps GEMM separate from score selection: callers compute dense +//! dot-product matrices, then the kernels documented below convert those dots to +//! metric scores and retain top candidates. A *point* is a vector being assigned +//! during partitioning; a *leader* names a child group. In leaf selection, +//! *source* names the point whose output list is being built and *target* names +//! another point in that same leaf. +//! +//! The wider PiPNN pipeline owns overlapping partition generation, leaf-local +//! nearest-neighbor construction, candidate merging, and optional graph-degree +//! finalization. This layer exports the partition-assignment and leaf-selection +//! kernels used inside that pipeline. Callers supply their dot-product matrices, +//! output storage, and reusable scratch; providers, graph IDs, recursion, edge +//! merging, persistence, and search remain outside these kernel APIs. //! //! Numerical kernels include: //! @@ -18,11 +73,57 @@ //! - [`leaf_kernel::LeafKernel`] scans each leaf's lower-triangular dot-product //! matrix once and retains nearest non-self neighbors for both endpoints. //! -//! Callers prepare these small handles once per build metric (and leaf `k`) and -//! reuse them across stripes or leaves. Preparation uses `diskann-wide` to select -//! the runtime architecture and returns a direct function pointer; repeated calls -//! do not repeat ISA or metric dispatch. PiPNN itself never names instruction -//! sets. +//! # Main modules and structures +//! +//! ## [`partition_kernel`] +//! +//! Partition callers first compute a point-by-leader dot-product tile with GEMM. +//! [`partition_kernel::PartitionInput`] bundles that tile with typed +//! [`partition_kernel::PartitionScales`]. A prepared +//! [`partition_kernel::PartitionKernel`] writes sorted leader-local positions to +//! a caller-owned output matrix. Fanout is the output column count and is bounded +//! by [`partition_kernel::MAX_PARTITION_FANOUT`]. Module documentation describes +//! scale units, validation, `process_points`, and tracker insertion. +//! +//! ## [`leaf_kernel`] +//! +//! Leaf callers compute a lower-triangular point-by-point dot matrix with +//! `sgemm_aat_lower`. [`leaf_kernel::LeafInput`] borrows that matrix; +//! [`leaf_kernel::LeafKernelWorkspace`] owns reusable per-worker scratch; and +//! [`leaf_kernel::LeafKernel`] writes sorted [`leaf_kernel::LeafNeighbor`] values +//! to a caller-owned matrix. [`leaf_kernel::leaf_neighbor_count`] derives each +//! leaf's width from its point count and requested `k`. Module documentation +//! describes width selection, `process_pairs`, fixed/dynamic storage, and stable +//! endpoint insertion. +//! +//! ## `kernel_metric` +//! +//! This private module owns metric formulas, scale units, zero/NaN behavior, and +//! one-time runtime-to-concrete metric selection shared by both public kernels. +//! Keeping it private prevents callers from constructing a formula/scale mismatch. +//! +//! # Typical use +//! +//! 1. Prepare one partition and one leaf kernel for the build metric. +//! 2. Reuse the partition handle for every GEMM stripe, changing only borrowed +//! input/output views. +//! 3. Reuse the leaf handle for every leaf. Derive output width with +//! [`leaf_kernel::leaf_neighbor_count`] and lease one workspace per worker. +//! 4. Translate leaf-local positions to dataset IDs outside these kernels. +//! +//! Callers prepare these small handles once per build metric and reuse them +//! across stripes or leaves. Each output view supplies its call-specific fanout +//! or neighbor width. Preparation uses `diskann-wide` to select the runtime +//! architecture and returns a direct function pointer; repeated calls do not +//! repeat ISA or metric dispatch. PiPNN itself never names instruction sets. +//! +//! # Ownership and performance boundary +//! +//! Kernels borrow all matrices and mutate only caller-owned output/scratch. They +//! do not own providers, thread pools, GEMM buffers, graph IDs, or persistence. +//! Partition traversal performs one score per point-leader pair; leaf traversal +//! performs one score per unordered point pair. Detailed complexity and scratch +//! costs are documented in each module. mod kernel_metric; diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann-pipnn/src/partition_kernel.rs index c94193d5d..fa64152d1 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann-pipnn/src/partition_kernel.rs @@ -5,6 +5,25 @@ //! Prepared distance and top-k kernels for partition assignment. //! +//! PiPNN recursively turns a dataset into small, overlapping groups called +//! *leaves*. At one recursion node it samples several existing points as +//! *leaders*. Each leader represents one child group. Every point is assigned to +//! its nearest `fanout` leaders, so `fanout > 1` copies that point into multiple +//! children and creates overlap. Children larger than the configured leaf limit +//! are partitioned again. +//! +//! This module performs only the nearest-leader selection inside that stage. It +//! does not sample leaders, gather vectors, run GEMM, group point IDs, or recurse. +//! The caller gathers a stripe of points and all leaders, computes their dot +//! products as one general matrix multiplication (GEMM), and passes that matrix +//! here. +//! [`PartitionKernel::nearest_leaders`] converts dots to metric scores and writes +//! leader column positions; the caller uses those positions to form child groups. +//! +//! For example, output `[2, 5, 7]` for one point at fanout three means: add that +//! point to children represented by leader columns 2, 5, and 7. It does not mean +//! those leaders are final graph neighbors. +//! //! The caller computes a row-major `points · leadersᵀ` tile with GEMM, then //! passes it to a [`PartitionKernel`] prepared once for the build metric. Kernel //! preparation selects the runtime architecture and concrete metric type once; @@ -27,6 +46,87 @@ //! //! Each point owns a fixed-capacity sorted tracker. Its last retained distance //! is the rejection threshold, so noncompetitive SIMD chunks avoid lane extraction. +//! +//! # Main structures +//! +//! - [`PartitionKernel`] is the reusable public handle containing one prepared +//! direct function pointer. +//! - [`PartitionInput`] bundles borrowed point-leader dots with +//! [`PartitionScales`], whose variants make scale units explicit. +//! - `PartitionEntry` is the architecture/metric-specialized destination that +//! validates a call before entering pointer-based SIMD. +//! - `process_points` is the shared point traversal. Concrete metric scale kinds +//! specialize unary/no-scale and binary-scale formulas without separate +//! runtime row processors. +//! - `LeaderTracker`, `insert_leader_lanes`, and `insert_leader` maintain one +//! fixed-capacity, stable sorted prefix per point. +//! +//! # Inputs and output +//! +//! For `p` points and `l` leaders, [`PartitionInput::dots`] is the row-major +//! `p × l` GEMM result. [`PartitionScales`] supplies exactly the scale units +//! required by the prepared metric. Output is a `p × f` matrix of leader-local +//! positions, where `f = output.ncols()` is requested fanout. Every point's +//! output is sorted by ascending score. +//! +//! Scores are derived from one point-leader dot product. Smaller is better: +//! +//! | Prepared metric | Score | Required [`PartitionScales`] | +//! | --- | --- | --- | +//! | squared L2 | `‖leader‖² - 2(point·leader)` | [`PartitionScales::L2`] | +//! | cosine | `1 - (point·leader)/(‖point‖‖leader‖)` | [`PartitionScales::Cosine`] | +//! | normalized cosine | `1 - point·leader` | [`PartitionScales::None`] | +//! | inner product | `-(point·leader)` | [`PartitionScales::None`] | +//! +//! Squared L2 omits `‖point‖²` because adding the same value to every leader +//! cannot change their order. `CosineNormalized` assumes vectors were normalized +//! before GEMM; this kernel does not verify vector norms. +//! +//! # Core flow +//! +//! 1. Validate matrix areas, backing lengths, fanout, and metric scale variant. +//! 2. Transform one point scale outside its leader loop when required. +//! 3. Score full SIMD leader groups and reject noncompetitive groups by mask. +//! 4. Score scalar-tail leaders with the metric's scalar operation order. +//! 5. Copy sorted leader IDs and reject underfilled points. +//! +//! # Performance +//! +//! With `p > 0` and `f > 0`, the kernel evaluates exactly `p * l` scores; +//! empty stripes or zero fanout return before traversal. Competitive leaders +//! bubble through at most `f <= MAX_PARTITION_FANOUT` tracker slots, giving +//! `O(plf)` worst-case work and `O(pl)` score computation. Tracker storage is a +//! fixed `O(MAX_PARTITION_FANOUT)` stack array per point; output is `O(pf)` and +//! no heap allocation occurs. Whole SIMD groups with no score below the current +//! threshold avoid lane materialization. Runtime architecture and metric selection happen +//! once in [`PartitionKernel::new`], outside stripe processing. +//! +//! # Example +//! +//! ``` +//! use diskann_pipnn::partition_kernel::{ +//! PartitionInput, PartitionKernel, PartitionScales, +//! }; +//! use diskann_utils::views::{MatrixView, MutMatrixView}; +//! use diskann_vector::distance::Metric; +//! +//! let dots = [ +//! 0.8, 0.2, 0.5, +//! 0.1, 0.9, 0.3, +//! ]; +//! let input = PartitionInput { +//! dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), +//! scales: PartitionScales::None, +//! }; +//! let mut assignments = vec![u32::MAX; 2 * 2]; +//! let output = MutMatrixView::try_from(&mut assignments[..], 2, 2).unwrap(); +//! +//! PartitionKernel::new(Metric::CosineNormalized) +//! .nearest_leaders(input, output) +//! .unwrap(); +//! +//! assert_eq!(assignments, [0, 2, 1, 2]); +//! ``` use std::marker::PhantomData; @@ -50,6 +150,11 @@ pub const MAX_PARTITION_FANOUT: usize = 16; type LeaderTracker = [(u32, f32); MAX_PARTITION_FANOUT]; /// Metric-specific normalization inputs for one partition tile. +/// +/// Slice lengths are checked against dot-matrix dimensions before output +/// mutation. Names encode units: cosine points arrive as squared norms because +/// they come from the point matrix diagonal, while leaders are normalized once +/// by the partition caller and arrive as norms. #[derive(Clone, Copy, Debug)] pub enum PartitionScales<'a> { /// L2 needs only squared leader norms; the point norm cannot affect ranking. @@ -69,6 +174,10 @@ pub enum PartitionScales<'a> { } /// One row-major point-by-leader dot-product tile. +/// +/// Matrix rows are points, columns are leaders, and [`Self::scales`] must match +/// the metric used to prepare [`PartitionKernel`]. This value only borrows input; +/// the prepared kernel stores no tile state. #[derive(Clone, Copy, Debug)] pub struct PartitionInput<'a> { /// One point per matrix row and one leader per column. @@ -169,6 +278,9 @@ type PartitionFn = /// Construct this once with [`PartitionKernel::new`] and reuse it for every /// point stripe. The handle is a direct function pointer and is `Copy`, `Send`, /// and `Sync`. +/// +/// It stores no matrix or output borrow, so callers may share one handle across +/// Rayon workers while each call owns independent views. #[derive(Clone, Copy, Debug)] pub struct PartitionKernel { run: PartitionFn, @@ -176,6 +288,14 @@ pub struct PartitionKernel { impl PartitionKernel { /// Prepare a partition kernel for `metric` and the current CPU. + /// + /// The return value contains one architecture/metric-specialized function + /// pointer and can process any valid stripe shape and fanout. + /// + /// # Performance + /// + /// Performs runtime architecture detection and one metric match once. + /// Reusing the handle removes both decisions from point and leader loops. pub fn new(metric: Metric) -> Self { diskann_wide::arch::dispatch1_no_features(PreparePartition, metric) } @@ -185,6 +305,28 @@ impl PartitionKernel { /// `output.nrows()` must equal `input.dots.nrows()`; its column count is the /// requested fanout. Results are ordered by ascending distance. For L2, the /// score omits the point norm because it cannot affect that point's ranking. + /// + /// `input` supplies point-leader dots and typed metric scales. `output` is + /// overwritten with leader-local positions. Successful return guarantees + /// exactly `output.ncols()` rankable leaders for every point. + /// + /// # Core flow + /// + /// The prepared entry validates every view and scale slice before mutation, + /// runs one architecture/metric-specialized point traversal, then checks the + /// final tracker slot for underfill. + /// + /// # Errors + /// + /// Returns [`PartitionKernelError`] for overflowing or mismatched shapes, + /// wrong scale variants or lengths, excessive fanout/leader counts, or a + /// point with too few rankable scores. Validation errors leave output + /// unchanged. + /// + /// # Performance + /// + /// See module-level complexity. This call follows one prepared direct + /// function pointer and performs no runtime ISA or metric dispatch. pub fn nearest_leaders( &self, input: PartitionInput<'_>, @@ -243,6 +385,10 @@ where /// /// The zero-sized entry receives all stripe state as arguments. Validation must /// complete before `process_points` reaches unchecked contiguous SIMD loads. +/// +/// Call order is fixed: validate without mutation, handle empty work, execute one +/// specialized traversal, then verify each point's last assignment. Keeping the +/// phases together makes every unchecked load depend on one visible gate. struct PartitionEntry(PhantomData); impl FTarget2, PartitionInput<'_>, MutMatrixView<'_, u32>> @@ -263,10 +409,14 @@ where // and fanout bounds before any output mutation or unchecked load. let scales = validate::(input, &output)?; let fanout = output.ncols(); + // Zero fanout and empty stripes require no assignments. Return before + // constructing trackers or touching output. if fanout == 0 || input.dots.nrows() == 0 { return Ok(()); } + // Architecture and metric are concrete here; only stripe dimensions and + // fanout remain runtime values. process_points::(arch, input.dots, scales, fanout, output.as_mut_slice()); // A sorted tracker can be underfilled only at its last slot. This keeps // post-validation linear in points rather than scanning every output ID. @@ -296,6 +446,12 @@ struct ScaleSlices<'a> { /// Matrix areas are recomputed with `checked_mul` before pointer loads. The /// `PartitionScales` variant must match concrete metric `M`, preventing plausible /// but incorrect norm units from crossing the interface. +/// +/// `input` and `output` are inspected only. Success returns borrowed scale slices +/// normalized to the storage layout expected by `M`; it establishes exact +/// backing lengths, representable leader IDs, and bounded fanout. Failure returns +/// [`PartitionKernelError`] before output mutation. Runtime is constant apart +/// from view metadata checks; matrix and scale contents are not scanned. fn validate<'a, M: KernelMetric>( input: PartitionInput<'a>, output: &MutMatrixView<'_, u32>, @@ -327,6 +483,9 @@ fn validate<'a, M: KernelMetric>( }); } + // Match the public enum against the concrete marker before erasing it to + // slices. This prevents squared point norms from being mistaken for leader + // norms even though both representations are `&[f32]`. let scales = match (M::METRIC, input.scales) { ( Metric::L2, @@ -367,6 +526,8 @@ fn validate<'a, M: KernelMetric>( } }; + // After variant validation, associated scale kinds define exact lengths. + // Scale-free metrics must provide empty slices so stale data cannot be used. check_length( "point scales", scales.point_scales.len(), @@ -380,6 +541,9 @@ fn validate<'a, M: KernelMetric>( Ok(scales) } +/// Return required scale length after metric specialization. +/// +/// Associated `ScaleKind` constants make this choice compile away. const fn expected_scale_len(kind: ScaleKind, count: usize) -> usize { if kind.is_some() { count @@ -426,6 +590,12 @@ fn check_length( /// preserves leader scan order for ties and makes NaNs non-rankable. L2 keeps /// historical bulk-FMA/scalar-tail rounding because changing it can alter graph /// assignment at near ties. +/// +/// `dots` supplies `p × l` scores, `scales` contains validated metric inputs, +/// `fanout` is both tracker prefix length and output width, and `output` contains +/// `p * fanout` slots. The function writes leader IDs in place and returns no +/// value. It computes `p * l` scores; each competitive score may shift `O(fanout)` +/// tracker entries. Tracker memory is fixed on the stack and no allocation occurs. fn process_points( arch: F::Arch, dots: MatrixView<'_, f32>, @@ -439,12 +609,16 @@ fn process_points( u64: From<<::BitMask as SIMDMask>::Underlying>, { let leader_count = dots.ncols(); + // Each point is independent. Reinitialize the fixed tracker here so no + // assignment state or tie order leaks across points. for (point, (point_dots, point_output)) in dots .as_slice() .chunks_exact(leader_count) .zip(output.chunks_exact_mut(fanout)) .enumerate() { + // Transform once per point rather than once per leader. For metrics + // without a point scale, specialization removes this branch and load. let point_scale = if M::PARTITION_POINT_SCALE.is_some() { M::PARTITION_POINT_SCALE.transform(scales.point_scales[point]) } else { @@ -452,6 +626,8 @@ fn process_points( }; let point_scale_vector = F::splat(arch, point_scale); let mut tracker = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + // Split at the largest complete vector boundary. Scalar tail uses the + // metric's explicit scalar operation order, not a padded SIMD load. let full = leader_count / F::LANES * F::LANES; for base in (0..full).step_by(F::LANES) { @@ -471,6 +647,8 @@ fn process_points( ); } + // Tail values use scalar metric functions intentionally. Padding a SIMD + // group would risk out-of-bounds scale loads and different L2 rounding. for (leader, &dot) in point_dots.iter().enumerate().skip(full) { let leader_scale = if M::PARTITION_LEADER_SCALE.is_some() { M::PARTITION_LEADER_SCALE.transform(scales.leader_scales[leader]) @@ -484,6 +662,8 @@ fn process_points( M::partition_distance_scalar(dot, point_scale, leader_scale), ); } + // Distances are only tracker state; child-group construction needs leader + // column positions in deterministic nearest-first order. copy_leader_ids(&tracker, point_output); } } @@ -493,6 +673,11 @@ fn process_points( /// The broadcast threshold avoids materializing lanes when none can improve the /// last slot. Bit iteration follows low-to-high lane order, preserving scalar tie /// behavior across SIMD widths. +/// +/// `distances` contains consecutive leaders beginning at `first_leader`; +/// `tracker[..fanout]` is the point's sorted retained prefix. The function +/// mutates that tracker and returns no value. Rejected groups cost one comparison +/// and mask test; accepted lanes each pay `O(fanout)` worst-case insertion. fn insert_leader_lanes( distances: F, first_leader: usize, @@ -523,6 +708,10 @@ fn insert_leader_lanes( /// The last slot is overwritten, then bubbled left. Equal and NaN distances do /// not enter, so scan order is the deterministic tie breaker and the last slot /// remains both rejection threshold and underfill sentinel. +/// +/// `tracker[..fanout]` must already be sorted and `fanout` must be non-zero. +/// `leader` is a local column position. The function returns no value and shifts +/// at most `fanout - 1` entries without allocation. #[inline(always)] fn insert_leader(tracker: &mut LeaderTracker, fanout: usize, leader: u32, distance: f32) { let threshold = fanout - 1; @@ -539,6 +728,9 @@ fn insert_leader(tracker: &mut LeaderTracker, fanout: usize, leader: u32, distan } /// Publish only leader IDs; distances stay private tracker state. +/// +/// `assignments.len()` is validated fanout. Copying costs `O(fanout)` and leaves +/// tracker state available for the underfill sentinel check encoded in IDs. fn copy_leader_ids(tracker: &LeaderTracker, assignments: &mut [u32]) { for (destination, &(leader, _)) in assignments.iter_mut().zip(tracker) { *destination = leader; From 3bc430d2eec95dfa1983b002410d5fab5321a1df Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:53:35 +0000 Subject: [PATCH 21/26] docs(pipnn): state cosine threshold --- diskann-pipnn/src/kernel_metric.rs | 14 ++++++++------ diskann-pipnn/src/leaf_kernel.rs | 5 +++-- diskann-pipnn/src/lib.rs | 8 ++++---- 3 files changed, 15 insertions(+), 12 deletions(-) diff --git a/diskann-pipnn/src/kernel_metric.rs b/diskann-pipnn/src/kernel_metric.rs index 6d13bfecd..753692bd6 100644 --- a/diskann-pipnn/src/kernel_metric.rs +++ b/diskann-pipnn/src/kernel_metric.rs @@ -47,9 +47,10 @@ //! //! # Numerical behavior //! -//! Subnormal norms are treated as zero before cosine division. A zero-norm -//! cosine endpoint forces zero similarity and distance `1.0`, even when the -//! other endpoint or dot is NaN. Otherwise NaN remains NaN, allowing strict +//! Squared norms below [`f32::MIN_POSITIVE`], and norms below +//! `sqrt(f32::MIN_POSITIVE)`, are treated as zero before cosine division. A +//! zero-threshold endpoint forces zero similarity and distance `1.0`, even when +//! the other endpoint or dot is NaN. Otherwise NaN remains NaN, allowing strict //! top-k comparisons to reject it. L2 scalar partition tails retain historical //! non-fused operation order because rounding can change leader assignment at //! near ties. @@ -82,9 +83,10 @@ pub(crate) enum ScaleKind { impl ScaleKind { /// Convert stored scale to the arithmetic form required by a kernel. /// - /// DiskANN treats subnormal squared norms, and corresponding subnormal - /// norms, as zero before division. Ordered comparisons intentionally leave - /// NaN unchanged so later distance comparisons keep it non-rankable. + /// DiskANN treats squared norms below `f32::MIN_POSITIVE`, and norms below + /// `sqrt(f32::MIN_POSITIVE)`, as zero before division. Ordered comparisons + /// intentionally leave NaN unchanged so later distance comparisons keep it + /// non-rankable. /// /// `stored` is interpreted according to `self`. The return value is zero, /// the original norm, the original squared norm, or its square root. This diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann-pipnn/src/leaf_kernel.rs index 0cb7e6843..69b587d4e 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann-pipnn/src/leaf_kernel.rs @@ -84,8 +84,9 @@ //! | inner product | `-(source·target)` | //! //! `CosineNormalized` assumes leaf vectors were normalized before GEMM. For -//! unnormalized cosine, a zero/subnormal norm gives zero similarity. NaN scores -//! never enter output because selection uses strict ordered comparisons. +//! unnormalized cosine, a norm below `sqrt(f32::MIN_POSITIVE)` gives zero +//! similarity. NaN scores never enter output because selection uses strict +//! ordered comparisons. //! //! # Core flow //! diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index d02ff0ca3..7dfc101e2 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -21,10 +21,10 @@ //! Every point is assigned to its nearest `fanout` leaders. Assigning to more //! than one leader makes child groups overlap. Oversized groups are processed //! recursively until bounded groups called *leaves* remain. -//! 2. **Pick within leaves.** Vectors in one leaf are contiguous enough for a -//! dense matrix multiplication to compute all pair dot products. Each point -//! picks its nearest leaf companions; selected pairs become candidate graph -//! edges. +//! 2. **Pick within leaves.** Vectors in one leaf are contiguous enough for one +//! dense general matrix multiplication (GEMM) to compute all pair dot +//! products. Each point picks its nearest leaf companions; selected pairs +//! become candidate graph edges. //! 3. **Merge and prune.** Candidates from overlapping leaves are combined. //! HashPrune can keep a bounded reservoir per source while edges stream in, //! retaining the closest candidate for each residual-direction hash. The From 400fc133c98d0d251cb150624a21e1d4c9cf1467 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 07:08:43 +0000 Subject: [PATCH 22/26] refactor(pipnn): move kernels into diskann graph Keep PiPNN beside graph policy so later layers can reuse private RobustPrune state without publishing it across a crate boundary. Preserve independent kernel oracles while removing duplicate formula-sharing differential wrappers. --- .github/workflows/ci.yml | 6 +- .github/workflows/nightly.yml | 4 +- Cargo.lock | 10 -- Cargo.toml | 2 - diskann-pipnn/Cargo.toml | 20 --- diskann/Cargo.toml | 3 + diskann/src/graph/mod.rs | 3 + .../src/graph/pipnn}/kernel_metric.rs | 6 +- .../src/graph/pipnn}/leaf_kernel.rs | 118 +----------------- .../lib.rs => diskann/src/graph/pipnn/mod.rs | 2 +- .../src/graph/pipnn}/partition_kernel.rs | 115 +---------------- .../tests/pipnn_leaf_kernel.rs | 24 ++-- .../tests/pipnn_partition_kernel.rs | 8 +- 13 files changed, 40 insertions(+), 281 deletions(-) delete mode 100644 diskann-pipnn/Cargo.toml rename {diskann-pipnn/src => diskann/src/graph/pipnn}/kernel_metric.rs (99%) rename {diskann-pipnn/src => diskann/src/graph/pipnn}/leaf_kernel.rs (90%) rename diskann-pipnn/src/lib.rs => diskann/src/graph/pipnn/mod.rs (98%) rename {diskann-pipnn/src => diskann/src/graph/pipnn}/partition_kernel.rs (89%) rename diskann-pipnn/tests/leaf_kernel_api.rs => diskann/tests/pipnn_leaf_kernel.rs (95%) rename diskann-pipnn/tests/partition_kernel_api.rs => diskann/tests/pipnn_partition_kernel.rs (97%) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f435e3752..02bcc51bc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -355,7 +355,8 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ - --package diskann-pipnn \ + --package diskann \ + --features diskann/pipnn \ -- --skip compile_tests \ --skip pivots::tests::run_test_happy_path @@ -416,7 +417,8 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ - --package diskann-pipnn \ + --package diskann \ + --features diskann/pipnn \ -- --skip compile_tests test-workspace: diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index ccf910024..def3cbc98 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -20,7 +20,7 @@ env: DISKANN_FEATURES: >- virtual_storage,spherical-quantization,product-quantization,tracing, experimental_diversity_search,disk-index,flatbuffers,linalg,codegen, - multi-vector,bftree,inmem2,integration-test + multi-vector,bftree,inmem2,integration-test,pipnn defaults: run: @@ -142,6 +142,6 @@ jobs: cargo +nightly miri nextest run --locked \ --package diskann-quantization cargo +nightly miri test --locked \ - --package diskann-pipnn --lib + --package diskann --features pipnn --lib env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance diff --git a/Cargo.lock b/Cargo.lock index 059a0b2c6..50ed17b11 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -680,16 +680,6 @@ dependencies = [ "thiserror 2.0.17", ] -[[package]] -name = "diskann-pipnn" -version = "0.55.0" -dependencies = [ - "diskann-utils", - "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 ee1ece101..6394a7d31 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,6 @@ members = [ "diskann-quantization", # Algorithm "diskann", - "diskann-pipnn", # Providers "diskann-providers", "diskann-disk", @@ -60,7 +59,6 @@ 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-pipnn/Cargo.toml b/diskann-pipnn/Cargo.toml deleted file mode 100644 index 848fbce2a..000000000 --- a/diskann-pipnn/Cargo.toml +++ /dev/null @@ -1,20 +0,0 @@ -# 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-utils.workspace = true -diskann-vector.workspace = true -diskann-wide.workspace = true -thiserror.workspace = true - -[lints] -workspace = true diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 14f6ba074..72911c910 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -56,6 +56,9 @@ panic = "warn" [features] default = ["tracing"] +# Enable PiPNN batch graph construction. +pipnn = [] + # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/mod.rs b/diskann/src/graph/mod.rs index 6bf4be7dd..041186969 100644 --- a/diskann/src/graph/mod.rs +++ b/diskann/src/graph/mod.rs @@ -8,6 +8,9 @@ pub use search_output_buffer::{ BufferState, IdDistance, IdDistanceAssociatedData, SearchOutputBuffer, }; +#[cfg(feature = "pipnn")] +pub mod pipnn; + pub mod adjacencylist; pub use adjacencylist::AdjacencyList; diff --git a/diskann-pipnn/src/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs similarity index 99% rename from diskann-pipnn/src/kernel_metric.rs rename to diskann/src/graph/pipnn/kernel_metric.rs index 753692bd6..edb3e6a84 100644 --- a/diskann-pipnn/src/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -210,11 +210,7 @@ where /// Scalar equivalent of [`clamp_nonnegative`]. #[inline(always)] fn clamp_nonnegative_scalar(distance: f32) -> f32 { - if distance < 0.0 { - 0.0 - } else { - distance - } + if distance < 0.0 { 0.0 } else { distance } } /// Compute cosine distance while preserving DiskANN zero/NaN semantics. diff --git a/diskann-pipnn/src/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs similarity index 90% rename from diskann-pipnn/src/leaf_kernel.rs rename to diskann/src/graph/pipnn/leaf_kernel.rs index 69b587d4e..4a289f3b0 100644 --- a/diskann-pipnn/src/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -109,7 +109,7 @@ //! # Example //! //! ``` -//! use diskann_pipnn::leaf_kernel::{ +//! use diskann::graph::pipnn::leaf_kernel::{ //! leaf_output_len, LeafInput, LeafKernel, LeafKernelWorkspace, LeafNeighbor, //! }; //! use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -140,12 +140,12 @@ use std::marker::PhantomData; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use diskann_wide::{ + Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, arch::{self, Dispatched1, FTarget1}, lifetime::AddLifetime, - Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, }; -use crate::kernel_metric::{visit_metric, KernelMetric, MetricVisitor}; +use super::kernel_metric::{KernelMetric, MetricVisitor, visit_metric}; /// One leaf-local neighbor and its metric distance. #[derive(Clone, Copy, Debug, PartialEq)] @@ -948,8 +948,6 @@ fn insert_dynamic_neighbor(neighbors: &mut [LeafNeighbor], target: u32, distance #[cfg(test)] mod tests { - use crate::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; - use super::*; fn test_dots(metric: Metric, points: usize) -> Vec { @@ -974,52 +972,6 @@ mod tests { } } - // Differential oracle for traversal and dispatch only. It intentionally - // shares `M::leaf_distance_scalar`; public API tests independently spell - // out metric formulas and full sorting behavior. - fn scalar_traversal_reference( - input: LeafInput<'_>, - neighbor_count: usize, - output: &mut [LeafNeighbor], - ) { - let point_count = input.dots.nrows(); - let norms: Vec<_> = (0..point_count) - .map(|source| M::LEAF_SCALE.transform(input.dots[(source, source)])) - .collect(); - let mut worst = vec![f32::INFINITY; point_count]; - let uses_norms = M::LEAF_SCALE.is_some(); - for source in 1..point_count { - for target in 0..source { - let (source_scale, target_scale) = if uses_norms { - (norms[source], norms[target]) - } else { - (0.0, 0.0) - }; - let distance = M::leaf_distance_scalar( - input.dots[(source, target)], - source_scale, - target_scale, - ); - insert_reference( - output, - &mut worst, - neighbor_count, - source, - target as u32, - distance, - ); - insert_reference( - output, - &mut worst, - neighbor_count, - target, - source as u32, - distance, - ); - } - } - } - fn insert_reference( output: &mut [LeafNeighbor], worst: &mut [f32], @@ -1038,70 +990,6 @@ mod tests { ); } - fn run_scalar_traversal( - metric: Metric, - input: LeafInput<'_>, - neighbor_count: usize, - output: &mut [LeafNeighbor], - ) { - match metric { - Metric::L2 => scalar_traversal_reference::(input, neighbor_count, output), - Metric::Cosine => scalar_traversal_reference::(input, neighbor_count, output), - Metric::CosineNormalized => { - scalar_traversal_reference::(input, neighbor_count, output) - } - Metric::InnerProduct => { - scalar_traversal_reference::(input, neighbor_count, output) - } - } - } - - fn assert_scalar_reference_matches_prepared_dispatch(metric: Metric) { - // Point count controls SIMD chunking. Cover both sides of 4-, 8-, and - // 16-lane boundaries, then the boundary around a second 16-lane chunk. - for points in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let dots = test_dots(metric, points); - let input = test_input(&dots, points); - for requested_k in [1, 2, 3, 4] { - let leaf_k = requested_k.min(points - 1); - let kernel = LeafKernel::new(metric); - let mut expected = vec![LeafNeighbor::default(); points * leaf_k]; - kernel - .nearest_neighbors( - input, - MutMatrixView::try_from(expected.as_mut_slice(), points, leaf_k).unwrap(), - &mut LeafKernelWorkspace::new(), - ) - .unwrap(); - - let mut actual = vec![LeafNeighbor::default(); points * leaf_k]; - run_scalar_traversal(metric, input, leaf_k, &mut actual); - - assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); - } - } - } - - #[test] - fn l2_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_prepared_dispatch(Metric::L2); - } - - #[test] - fn cosine_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_prepared_dispatch(Metric::Cosine); - } - - #[test] - fn normalized_cosine_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_prepared_dispatch(Metric::CosineNormalized); - } - - #[test] - fn inner_product_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_prepared_dispatch(Metric::InnerProduct); - } - #[test] fn scalar_insertion_orders_candidates_and_rejects_nan() { let mut output = [LeafNeighbor::default(); 4]; diff --git a/diskann-pipnn/src/lib.rs b/diskann/src/graph/pipnn/mod.rs similarity index 98% rename from diskann-pipnn/src/lib.rs rename to diskann/src/graph/pipnn/mod.rs index 7dfc101e2..3e714e201 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -52,7 +52,7 @@ //! search graph //! ``` //! -//! This crate keeps GEMM separate from score selection: callers compute dense +//! This module keeps GEMM separate from score selection: callers compute dense //! dot-product matrices, then the kernels documented below convert those dots to //! metric scores and retain top candidates. A *point* is a vector being assigned //! during partitioning; a *leader* names a child group. In leaf selection, diff --git a/diskann-pipnn/src/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs similarity index 89% rename from diskann-pipnn/src/partition_kernel.rs rename to diskann/src/graph/pipnn/partition_kernel.rs index fa64152d1..3c88e5c9f 100644 --- a/diskann-pipnn/src/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -104,7 +104,7 @@ //! # Example //! //! ``` -//! use diskann_pipnn::partition_kernel::{ +//! use diskann::graph::pipnn::partition_kernel::{ //! PartitionInput, PartitionKernel, PartitionScales, //! }; //! use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -133,12 +133,12 @@ use std::marker::PhantomData; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use diskann_wide::{ + Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, arch::{self, Dispatched2, FTarget2}, lifetime::AddLifetime, - Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, }; -use crate::kernel_metric::{visit_metric, KernelMetric, MetricVisitor, ScaleKind}; +use super::kernel_metric::{KernelMetric, MetricVisitor, ScaleKind, visit_metric}; /// Maximum number of leaders retained for one point. /// @@ -545,11 +545,7 @@ fn validate<'a, M: KernelMetric>( /// /// Associated `ScaleKind` constants make this choice compile away. const fn expected_scale_len(kind: ScaleKind, count: usize) -> usize { - if kind.is_some() { - count - } else { - 0 - } + if kind.is_some() { count } else { 0 } } fn checked_area( @@ -739,37 +735,10 @@ fn copy_leader_ids(tracker: &LeaderTracker, assignments: &mut [u32]) { #[cfg(test)] mod tests { - use crate::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; + use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; use super::*; - fn test_data(metric: Metric, leader_count: usize) -> (Vec, Vec, Vec) { - let dots = (0..2 * leader_count) - .map(|index| (((index * 13 + 7) % 29) as f32 - 14.0) * 0.125) - .collect(); - let point_scales = if metric == Metric::Cosine { - vec![0.0, 16.0] - } else { - Vec::new() - }; - let leader_scales = match metric { - Metric::L2 => (0..leader_count) - .map(|leader| ((leader + 1) as f32).powi(2)) - .collect(), - Metric::Cosine => (0..leader_count) - .map(|leader| { - if leader == 0 { - 0.0 - } else { - (leader + 1) as f32 - } - }) - .collect(), - Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), - }; - (dots, point_scales, leader_scales) - } - fn test_input<'a>( metric: Metric, dots: &'a [f32], @@ -852,80 +821,6 @@ mod tests { } } - fn run_scalar_traversal( - metric: Metric, - input: PartitionInput<'_>, - fanout: usize, - output: &mut [u32], - ) { - match metric { - Metric::L2 => scalar_traversal_reference::(input, fanout, output), - Metric::Cosine => scalar_traversal_reference::(input, fanout, output), - Metric::CosineNormalized => { - scalar_traversal_reference::(input, fanout, output) - } - Metric::InnerProduct => { - scalar_traversal_reference::(input, fanout, output) - } - } - } - - fn assert_scalar_reference_matches_prepared_dispatch(metric: Metric) { - // Leader count controls SIMD chunking. Exercise both sides of 4-, 8-, and - // 16-lane boundaries, then a second 16-lane chunk. - for leader_count in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let (dots, point_scales, leader_scales) = test_data(metric, leader_count); - let input = test_input( - metric, - &dots, - 2, - leader_count, - &point_scales, - &leader_scales, - ); - let kernel = PartitionKernel::new(metric); - for fanout in [1, 2, 6, MAX_PARTITION_FANOUT] { - if fanout > leader_count { - continue; - } - let mut expected = vec![u32::MAX; 2 * fanout]; - kernel - .nearest_leaders( - input, - MutMatrixView::try_from(expected.as_mut_slice(), 2, fanout).unwrap(), - ) - .unwrap(); - - let mut actual = vec![u32::MAX; 2 * fanout]; - run_scalar_traversal(metric, input, fanout, &mut actual); - assert_eq!( - actual, expected, - "{metric:?}, leaders={leader_count}, k={fanout}" - ); - } - } - } - - #[test] - fn l2_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_prepared_dispatch(Metric::L2); - } - - #[test] - fn cosine_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_prepared_dispatch(Metric::Cosine); - } - - #[test] - fn normalized_cosine_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_prepared_dispatch(Metric::CosineNormalized); - } - - #[test] - fn inner_product_scalar_reference_matches_prepared_dispatch_at_lane_boundaries() { - assert_scalar_reference_matches_prepared_dispatch(Metric::InnerProduct); - } - #[test] fn scalar_distance_matches_metric_contract() { assert_eq!(L2::partition_distance_scalar(2.0, 0.0, 9.0), 5.0); diff --git a/diskann-pipnn/tests/leaf_kernel_api.rs b/diskann/tests/pipnn_leaf_kernel.rs similarity index 95% rename from diskann-pipnn/tests/leaf_kernel_api.rs rename to diskann/tests/pipnn_leaf_kernel.rs index 07b6f85dc..77837e639 100644 --- a/diskann-pipnn/tests/leaf_kernel_api.rs +++ b/diskann/tests/pipnn_leaf_kernel.rs @@ -3,16 +3,18 @@ * Licensed under the MIT license. */ +#![cfg(feature = "pipnn")] + use std::cmp::Ordering; -use diskann_pipnn::leaf_kernel::{ - leaf_neighbor_count, leaf_output_len, LeafInput, LeafKernel, LeafKernelError, - LeafKernelWorkspace, LeafNeighbor, +use diskann::graph::pipnn::leaf_kernel::{ + LeafInput, LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, + leaf_output_len, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; -const SIMD_BOUNDARY_POINTS: [usize; 9] = [7, 8, 9, 15, 16, 17, 64, 256, 512]; +const SIMD_BOUNDARY_POINTS: [usize; 15] = [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33, 64, 256, 512]; const ZERO_NORM_POSITION: usize = 0; const DISTINCT_NORM_POSITION: usize = 2; const NORM_PERIOD: usize = 5; @@ -37,9 +39,7 @@ fn differential_dots(metric: Metric, points: usize) -> Vec { for target in 0..source { let pair = ((source * SOURCE_MIXER + target * TARGET_MIXER) % MIX_MODULUS) as f32 - MIX_CENTER; - dots[source * points + target] = if source == points - 1 && target == 0 { - f32::NAN - } else if TIED_TARGETS.contains(&target) { + dots[source * points + target] = if TIED_TARGETS.contains(&target) { 0.5 } else { pair * DOT_SCALE @@ -224,7 +224,7 @@ fn cosine_treats_zero_norm_as_zero_similarity() { } #[test] -fn preserves_pipnn_metric_edge_semantics() { +fn clamps_negative_distances_and_preserves_cosine_extremes() { #[rustfmt::skip] let out_of_range = [1.0, 0.0, 2.0, 1.0]; assert_eq!( @@ -328,9 +328,11 @@ fn clamps_k_to_available_non_self_neighbors() { assert_eq!(leaf_k, 2); for (source, neighbors) in output.chunks_exact(leaf_k).enumerate() { - assert!(neighbors - .iter() - .all(|neighbor| neighbor.target as usize != source)); + assert!( + neighbors + .iter() + .all(|neighbor| neighbor.target as usize != source) + ); } } diff --git a/diskann-pipnn/tests/partition_kernel_api.rs b/diskann/tests/pipnn_partition_kernel.rs similarity index 97% rename from diskann-pipnn/tests/partition_kernel_api.rs rename to diskann/tests/pipnn_partition_kernel.rs index 59c69549b..e2246b97a 100644 --- a/diskann-pipnn/tests/partition_kernel_api.rs +++ b/diskann/tests/pipnn_partition_kernel.rs @@ -3,8 +3,10 @@ * Licensed under the MIT license. */ -use diskann_pipnn::partition_kernel::{ - PartitionInput, PartitionKernel, PartitionKernelError, PartitionScales, MAX_PARTITION_FANOUT, +#![cfg(feature = "pipnn")] + +use diskann::graph::pipnn::partition_kernel::{ + MAX_PARTITION_FANOUT, PartitionInput, PartitionKernel, PartitionKernelError, PartitionScales, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -157,7 +159,7 @@ fn prepared_dispatch_matches_reference_across_simd_width_boundaries() { Metric::CosineNormalized, Metric::InnerProduct, ] { - for leader_count in [7, 8, 9, 15, 16, 17] { + for leader_count in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let (dots, point_scales, leader_scales) = differential_data(metric, leader_count); let input = test_input( metric, From e5fa1369209d50152701c4c0c24654d5109fcc23 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 07:14:52 +0000 Subject: [PATCH 23/26] test(pipnn): document fixture panic policy --- diskann/tests/pipnn_leaf_kernel.rs | 5 +++++ diskann/tests/pipnn_partition_kernel.rs | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/diskann/tests/pipnn_leaf_kernel.rs b/diskann/tests/pipnn_leaf_kernel.rs index 77837e639..c6dc621d8 100644 --- a/diskann/tests/pipnn_leaf_kernel.rs +++ b/diskann/tests/pipnn_leaf_kernel.rs @@ -4,6 +4,11 @@ */ #![cfg(feature = "pipnn")] +#![allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] use std::cmp::Ordering; diff --git a/diskann/tests/pipnn_partition_kernel.rs b/diskann/tests/pipnn_partition_kernel.rs index e2246b97a..597c472d3 100644 --- a/diskann/tests/pipnn_partition_kernel.rs +++ b/diskann/tests/pipnn_partition_kernel.rs @@ -4,6 +4,10 @@ */ #![cfg(feature = "pipnn")] +#![allow( + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] use diskann::graph::pipnn::partition_kernel::{ MAX_PARTITION_FANOUT, PartitionInput, PartitionKernel, PartitionKernelError, PartitionScales, From 537b5ebc1d3984a73af6dce39e5f77ae6b5952d5 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 08:22:58 +0000 Subject: [PATCH 24/26] ci(pipnn): cover the in-crate feature --- .github/workflows/ci.yml | 3 ++- diskann/src/graph/pipnn/mod.rs | 8 +++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 02bcc51bc..3894a758e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ env: CARGO_TERM_COLOR: always # The features we want to explicitly test. For example, the `flatbuffers-build` feature # of `diskann-quantization` requires additional setup and so must not be included by default. - DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test" + DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test,pipnn" # Intel SDE version used for baseline and AVX-512 emulation jobs. SDE_VERSION: "sde-external-10.8.0-2026-03-15-lin" @@ -518,6 +518,7 @@ jobs: cargo llvm-cov nextest --locked \ --config "$RUST_CONFIG" \ --workspace \ + --features "${{ env.DISKANN_FEATURES }}" \ --lcov --output-path lcov.info - name: Generate miri code coverage diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 3e714e201..1fb729114 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -25,11 +25,9 @@ //! dense general matrix multiplication (GEMM) to compute all pair dot //! products. Each point picks its nearest leaf companions; selected pairs //! become candidate graph edges. -//! 3. **Merge and prune.** Candidates from overlapping leaves are combined. -//! HashPrune can keep a bounded reservoir per source while edges stream in, -//! retaining the closest candidate for each residual-direction hash. The -//! alternative collects unique candidates directly. An optional final Vamana -//! RobustPrune selects a bounded, directionally diverse adjacency list. +//! 3. **Merge and finalize.** Candidates from overlapping leaves are combined +//! into one bounded adjacency list per source. Later stack layers own the +//! candidate-merging and graph-policy details; these numerical kernels do not. //! //! ```text //! dataset points From 37349ee84da8197e44c5fe877878847577db2a9e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 09:13:19 +0000 Subject: [PATCH 25/26] bench(pipnn): add IAI kernel microbenchmarks --- Cargo.lock | 1 + diskann/Cargo.toml | 6 ++ diskann/benches/bench_main_iai.rs | 22 +++++ diskann/benches/benchmarks_iai/mod.rs | 6 ++ .../benches/benchmarks_iai/pipnn_kernels.rs | 96 +++++++++++++++++++ 5 files changed, 131 insertions(+) create mode 100644 diskann/benches/bench_main_iai.rs create mode 100644 diskann/benches/benchmarks_iai/mod.rs create mode 100644 diskann/benches/benchmarks_iai/pipnn_kernels.rs diff --git a/Cargo.lock b/Cargo.lock index 50ed17b11..0c38cc590 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -444,6 +444,7 @@ dependencies = [ "futures-util", "half", "hashbrown 0.16.1", + "iai-callgrind", "num-traits", "pin-project", "rand", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 72911c910..a28b1b815 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -32,6 +32,7 @@ diskann-wide = { workspace = true } dashmap = { workspace = true, optional = true } [dev-dependencies] +iai-callgrind.workspace = true futures-util = { workspace = true, default-features = false } pin-project.workspace = true rand.workspace = true @@ -41,6 +42,11 @@ serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "sync"] } dashmap = { workspace = true } +[[bench]] +name = "bench_main_iai" +harness = false +required-features = ["pipnn", "testing"] + # Some 'cfg's in the source tree will be flagged by `cargo clippy -j 2 --workspace --no-deps --all-targets -- -D warnings` [lints.rust] unexpected_cfgs = { level = "warn", check-cfg = ['cfg(coverage)'] } diff --git a/diskann/benches/bench_main_iai.rs b/diskann/benches/bench_main_iai.rs new file mode 100644 index 000000000..f91432787 --- /dev/null +++ b/diskann/benches/bench_main_iai.rs @@ -0,0 +1,22 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use benchmarks_iai::pipnn_kernels::pipnn_kernels; +use iai_callgrind::{EventKind, LibraryBenchmarkConfig, RegressionConfig, main}; + +mod benchmarks_iai; + +main!( + config = LibraryBenchmarkConfig::default() + .regression( + RegressionConfig::default().limits([ + (EventKind::Ir, 5.0), + (EventKind::EstimatedCycles, 5.0), + (EventKind::TotalRW, 5.0), + (EventKind::L1hits, 5.0), + ]) + ); + library_benchmark_groups = pipnn_kernels, +); diff --git a/diskann/benches/benchmarks_iai/mod.rs b/diskann/benches/benchmarks_iai/mod.rs new file mode 100644 index 000000000..82760c5e5 --- /dev/null +++ b/diskann/benches/benchmarks_iai/mod.rs @@ -0,0 +1,6 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +pub(crate) mod pipnn_kernels; diff --git a/diskann/benches/benchmarks_iai/pipnn_kernels.rs b/diskann/benches/benchmarks_iai/pipnn_kernels.rs new file mode 100644 index 000000000..0e6de738d --- /dev/null +++ b/diskann/benches/benchmarks_iai/pipnn_kernels.rs @@ -0,0 +1,96 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +#![allow( + clippy::unwrap_used, + reason = "deterministic benchmark fixture construction must abort on invalid setup" +)] + +use diskann::graph::pipnn::{ + leaf_kernel::{LeafInput, LeafKernel, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count}, + partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, +}; +use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_vector::distance::Metric; +use iai_callgrind::black_box; + +const PARTITION_POINTS: usize = 256; +const LEADERS: usize = 32; +const FANOUT: usize = 4; +const LEAF_POINTS: usize = 128; +const LEAF_K: usize = 3; + +type PartitionFixture = (PartitionKernel, Vec, Vec, Vec); +type LeafFixture = (LeafKernel, LeafKernelWorkspace, Vec, Vec); + +fn setup_partition() -> PartitionFixture { + let dots = (0..PARTITION_POINTS * LEADERS) + .map(|index| ((index * 17 + 11) % 257) as f32 / 257.0) + .collect(); + let leader_squared_norms = (0..LEADERS) + .map(|leader| 1.0 + leader as f32 / LEADERS as f32) + .collect(); + ( + PartitionKernel::new(Metric::L2), + dots, + leader_squared_norms, + vec![u32::MAX; PARTITION_POINTS * FANOUT], + ) +} + +#[iai_callgrind::library_benchmark(setup = setup_partition)] +fn assign_points_to_leaders(fixture: PartitionFixture) { + let (kernel, dots, leader_squared_norms, mut output) = fixture; + kernel + .nearest_leaders( + PartitionInput { + dots: MatrixView::try_from(dots.as_slice(), PARTITION_POINTS, LEADERS).unwrap(), + scales: PartitionScales::L2 { + leader_squared_norms: &leader_squared_norms, + }, + }, + MutMatrixView::try_from(output.as_mut_slice(), PARTITION_POINTS, FANOUT).unwrap(), + ) + .unwrap(); + black_box(output); +} + +fn setup_leaf() -> LeafFixture { + let mut dots = vec![f32::NAN; LEAF_POINTS * LEAF_POINTS]; + for source in 0..LEAF_POINTS { + dots[source * LEAF_POINTS + source] = 1.0 + (source % 7) as f32; + for target in 0..source { + dots[source * LEAF_POINTS + target] = + ((source * 17 + target * 11) % 257) as f32 / 257.0; + } + } + let neighbors = leaf_neighbor_count(LEAF_POINTS, LEAF_K).unwrap(); + ( + LeafKernel::new(Metric::L2), + LeafKernelWorkspace::new(), + dots, + vec![LeafNeighbor::default(); LEAF_POINTS * neighbors], + ) +} + +#[iai_callgrind::library_benchmark(setup = setup_leaf)] +fn select_leaf_neighbors(fixture: LeafFixture) { + let (kernel, mut workspace, dots, mut output) = fixture; + kernel + .nearest_neighbors( + LeafInput { + dots: MatrixView::try_from(dots.as_slice(), LEAF_POINTS, LEAF_POINTS).unwrap(), + }, + MutMatrixView::try_from(output.as_mut_slice(), LEAF_POINTS, LEAF_K).unwrap(), + &mut workspace, + ) + .unwrap(); + black_box(output); +} + +iai_callgrind::library_benchmark_group!( + name = pipnn_kernels; + benchmarks = assign_points_to_leaders, select_leaf_neighbors, +); From e265ecb78ad7b89384fb4cc104feacc15f5941b1 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 09:53:41 +0000 Subject: [PATCH 26/26] ci(pipnn): compile the shared IAI target --- .github/workflows/ci.yml | 2 +- .github/workflows/nightly.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3894a758e..6a2d51cdf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ env: CARGO_TERM_COLOR: always # The features we want to explicitly test. For example, the `flatbuffers-build` feature # of `diskann-quantization` requires additional setup and so must not be included by default. - DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test,pipnn" + DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test,pipnn,testing" # Intel SDE version used for baseline and AVX-512 emulation jobs. SDE_VERSION: "sde-external-10.8.0-2026-03-15-lin" diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index def3cbc98..c2ed9e443 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -20,7 +20,7 @@ env: DISKANN_FEATURES: >- virtual_storage,spherical-quantization,product-quantization,tracing, experimental_diversity_search,disk-index,flatbuffers,linalg,codegen, - multi-vector,bftree,inmem2,integration-test,pipnn + multi-vector,bftree,inmem2,integration-test,pipnn,testing defaults: run: