From 016c05b208b9d79dc74bfd99f7f89911429990d8 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/80] 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 700396de4e..1e7feeb9d0 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 7ee59d6b60..d0b6c38e70 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -207,6 +207,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 0000000000..e84d600d3e --- /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 3a177221b66c601551f8dd40cd8f9b7c3790853c 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/80] 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 2588eee92a..0b21b0c2cd 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 73d7d2d611..4721ee7774 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 0000000000..1ff7c3bfd5 --- /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 0000000000..a913ba138e --- /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 0000000000..198434b72b --- /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 0000000000..b39c05aa82 --- /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 0000000000..4028304b4d --- /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 0000000000..f3082f3d3e --- /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 318227ca5d..f0b7431a87 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 82cf391076..83779dacf0 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 836193a452..b93c861e7a 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 60ffa4477c..45cc64a4a2 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 054b249e8f..48ecea7aee 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 d38465f906..ff59119992 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 328dba4d26..7028b2fdcf 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 3158ffc1dd..d38de49de8 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 30d08e6cb1..d6adcb7b13 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 507dc1288f..1e8d707769 100644 --- a/diskann-wide/src/emulated.rs +++ b/diskann-wide/src/emulated.rs @@ -199,6 +199,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 @@ -902,6 +911,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 fc15661f1e..69bba4bc41 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! { @@ -1141,6 +1173,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 b5ae8cf646838155cd9cbd93b0d69a6a3ef8f732 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/80] 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 c874e1b361..29291313d6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -354,6 +354,7 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ + --package diskann-pipnn \ -- --skip compile_tests \ --skip pivots::tests::run_test_happy_path @@ -416,6 +417,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 0000000000..333233922f --- /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 798996460ded42b3ba228d3b20f4ff15bca44a67 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/80] 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 333233922f..22790c554f 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 a913ba138e..0953bee21c 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 0000000000..661f0f5582 --- /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 b39c05aa82..3525ae1574 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 0000000000..df51bfc66b --- /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 4028304b4d..2110ddffa4 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 07cff8c556b9a5e2c9258e637509d52ef3b7f51d 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/80] 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 0953bee21c..b1eeb61705 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 661f0f5582..becef50398 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 3525ae1574..4865005c47 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 df51bfc66b..feebc53487 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 2110ddffa4..5cead7e96d 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 f3082f3d3e..bb90a8b9f2 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 2efc9f6471fcb3c83ac5449fe19cd9835885a71a 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/80] 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 b1eeb61705..34e78a02b4 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 becef50398..90381bc1da 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 4865005c47..b2942fa650 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 feebc53487..aaa03917b7 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 c3c1f72fc8d003d53b182b425d628dfb252bfb74 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/80] 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 34e78a02b4..f8531d4295 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 6fef4174d3d21cb6a559decff04c59016e9b6f0a 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/80] 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 1e7feeb9d0..55ca5b2dfb 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 f8531d4295..7f75f228ba 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 90381bc1da..67878424f8 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 b2942fa650..bb193eb9b7 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 aaa03917b7..e7b33ead19 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 ce0293fc32c51c7fdf6d7cc49845a055a85118b4 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/80] 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 67878424f8..1dba1805c9 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 e7b33ead19..ad3e048a04 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 c73fe3b52d4b3589d0761d44404d7c6e91f56431 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/80] 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 7f75f228ba..2aee917f3d 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 bb193eb9b7..ab4d022555 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 09150f0c7d..b84b793b93 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. /// @@ -906,16 +906,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 27ff45625726c761586805b3ca449fb0784cce2a 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/80] 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 b84b793b93..a233622615 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. /// @@ -906,16 +906,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 16f733db3553381af5254175ef2f37c929b5cdfd 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/80] fix(pipnn): address kernel review feedback --- 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 +- 9 files changed, 363 insertions(+), 283 deletions(-) diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index d0b6c38e70..f434cfb7e2 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -82,6 +82,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; @@ -156,51 +180,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); @@ -217,34 +199,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 e84d600d3e..d19c92c761 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 22790c554f..52661f5195 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 2aee917f3d..c514de43e2 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 1dba1805c9..1bef128ab2 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 198434b72b..17a5dd708f 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 ab4d022555..8122244f9e 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 ad3e048a04..7da20fd580 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 5cead7e96d..5971f4cb68 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 ac29cb818b38049753ab87925dc63e0a3d80760a 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/80] 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 0b21b0c2cd..3198b891bd 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 1ff7c3bfd5..848fbce2a1 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 52661f5195..0000000000 --- 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 0000000000..f7e61d43b8 --- /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 c514de43e2..d0869bf61a 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 1bef128ab2..0000000000 --- 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 17a5dd708f..ad08807e01 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 8122244f9e..80d2a02e8f 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 7da20fd580..0000000000 --- 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 5971f4cb68..f82bfe11d2 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 bb90a8b9f2..0000000000 --- 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 0000000000..e5318b8d1c --- /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 8e63abb5656556d17923a21ccf12966c697401af 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/80] 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 ad08807e01..9cfc3e25e0 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 045f0126a6c14abbb2bab2808c54fb0cbfbf3b62 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/80] 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 f7e61d43b8..08bc2987d4 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 d0869bf61a..aac60c93f5 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 80d2a02e8f..e19560c6b0 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 0d2a8d95c199715990f0d093d0a35d290ade8a56 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/80] 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 aac60c93f5..9697b13c39 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 e19560c6b0..f01a3fecc5 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 6fef94db023bcd78142058de90355886fdf69a96 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/80] 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 08bc2987d4..181b642f78 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 9697b13c39..71276a9705 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 f01a3fecc5..992d9fa284 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 f97ca93f894e4cde12f479800df48f78b41cee42 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/80] 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 71276a9705..8efc76718b 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 47dcb740f39ebd3e03a620a3770c433dcb01c8ff 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/80] 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 181b642f78..01aaca0d3f 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 8efc76718b..25f5e841b8 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 992d9fa284..c94193d5d7 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 f82bfe11d2..07b6f85dcd 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 e5318b8d1c..59c69549b9 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 844a50c592e766c97a2875e3622b9e0e4f179361 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/80] 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 01aaca0d3f..6d13bfecd2 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 25f5e841b8..0cb7e6843c 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 9cfc3e25e0..d02ff0ca3b 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 c94193d5d7..fa64152d15 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 36c17636a0a871ab3d0730db2a55dbe00e82f024 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/80] 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 6d13bfecd2..753692bd61 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 0cb7e6843c..69b587d4ee 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 d02ff0ca3b..7dfc101e2a 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 161965f39188162975d378caac14a5d3a10083b2 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/80] 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 +- 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 +- 12 files changed, 38 insertions(+), 279 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 29291313d6..a70a00d6b1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -354,7 +354,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 @@ -417,7 +418,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/Cargo.lock b/Cargo.lock index 3198b891bd..2588eee92a 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 4721ee7774..73d7d2d611 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 848fbce2a1..0000000000 --- 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 14f6ba0746..72911c910f 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 374efc6443..b9a9fc2345 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 753692bd61..edb3e6a842 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 69b587d4ee..4a289f3b01 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 7dfc101e2a..3e714e201c 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 fa64152d15..3c88e5c9f2 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 07b6f85dcd..77837e6399 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 59c69549b9..e2246b97a1 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 e8cb9c7c2598a8a187ef8bdda075a331b887c92e 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/80] 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 77837e6399..c6dc621d8b 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 e2246b97a1..597c472d3d 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 c0d15a266c23b6700e4b86c3fbf64b44b947fb4f 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/80] ci(pipnn): cover the in-crate feature --- diskann/src/graph/pipnn/mod.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 3e714e201c..1fb729114c 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 f671b788ec63bc35ddd9ac029190bac1b0c89ee7 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/80] 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 2588eee92a..f9d909d4a2 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 72911c910f..a28b1b815d 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 0000000000..f914327879 --- /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 0000000000..82760c5e5e --- /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 0000000000..0e6de738da --- /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 32f323cd1f59dc99c5d923854ae5df52c5d45777 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:50:53 +0000 Subject: [PATCH 26/80] fix(ci): use workspace feature name for pipnn --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a70a00d6b1..73b45d0922 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -355,7 +355,7 @@ jobs: --package diskann-vector \ --package diskann-quantization \ --package diskann \ - --features diskann/pipnn \ + --features pipnn \ -- --skip compile_tests \ --skip pivots::tests::run_test_happy_path @@ -419,7 +419,7 @@ jobs: --package diskann-vector \ --package diskann-quantization \ --package diskann \ - --features diskann/pipnn \ + --features pipnn \ -- --skip compile_tests test-workspace: From 0f31011e2bbfdde63d810cfcca7b688b86aff1ac Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:17:14 +0000 Subject: [PATCH 27/80] fix(ci): validate every stacked layer --- .github/workflows/disk-benchmarks.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/disk-benchmarks.yml b/.github/workflows/disk-benchmarks.yml index 36753c846e..341726c119 100644 --- a/.github/workflows/disk-benchmarks.yml +++ b/.github/workflows/disk-benchmarks.yml @@ -19,6 +19,7 @@ on: pull_request: branches: - main + - "pipnn-stack/**" paths: - 'diskann/**' - 'diskann-disk/**' From 80c9f33c457d51ec6fedce2ac46ffb1e5136e34a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:37:47 +0000 Subject: [PATCH 28/80] ci(pipnn): run IAI and ARM feature gates --- .github/workflows/ci.yml | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 73b45d0922..e2e2d3fe3c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -70,6 +70,7 @@ jobs: - test-workspace - test-workspace-features - coverage + - iai-callgrind - vectorset-clippy - vectorset-fmt - vectorset-build @@ -93,6 +94,32 @@ jobs: steps: - run: exit 0 + iai-callgrind: + needs: basics + name: IAI-Callgrind microbenchmarks + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install Rust + run: rustup show + + - name: Install Valgrind + run: | + sudo apt-get update -qq + sudo apt-get install -y valgrind + + - name: Install iai-callgrind-runner + run: cargo install iai-callgrind-runner --version 0.14.2 --locked + + - uses: Swatinem/rust-cache@v2 + + - name: Run shared DiskANN IAI target + run: | + cargo bench --locked -p diskann \ + --bench bench_main_iai \ + --features pipnn,testing + fmt: name: format check runs-on: ubuntu-latest @@ -463,6 +490,7 @@ jobs: os: - windows-latest - ubuntu-latest + - ubuntu-24.04-arm steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 From 0ce4ef3d2c155d3ad2ff93414cbed53b9dc83faf Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:41:41 +0000 Subject: [PATCH 29/80] test(pipnn): keep regression tooling local Remove the submitted DiskANN microbenchmark target and co-locate numerical tests with their implementation files. --- .github/workflows/ci.yml | 27 -- Cargo.lock | 1 - diskann-linalg/src/lib.rs | 112 +++++ diskann-linalg/tests/sgemm_aat_lower.rs | 109 ----- 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 ---- diskann/src/graph/pipnn/leaf_kernel.rs | 424 +++++++++++++++++ diskann/src/graph/pipnn/partition_kernel.rs | 382 ++++++++++++++++ diskann/tests/pipnn_leaf_kernel.rs | 427 ------------------ diskann/tests/pipnn_partition_kernel.rs | 384 ---------------- 12 files changed, 918 insertions(+), 1078 deletions(-) delete mode 100644 diskann-linalg/tests/sgemm_aat_lower.rs delete mode 100644 diskann/benches/bench_main_iai.rs delete mode 100644 diskann/benches/benchmarks_iai/mod.rs delete mode 100644 diskann/benches/benchmarks_iai/pipnn_kernels.rs delete mode 100644 diskann/tests/pipnn_leaf_kernel.rs delete mode 100644 diskann/tests/pipnn_partition_kernel.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e2e2d3fe3c..30ee4ae606 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -70,7 +70,6 @@ jobs: - test-workspace - test-workspace-features - coverage - - iai-callgrind - vectorset-clippy - vectorset-fmt - vectorset-build @@ -94,32 +93,6 @@ jobs: steps: - run: exit 0 - iai-callgrind: - needs: basics - name: IAI-Callgrind microbenchmarks - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Install Rust - run: rustup show - - - name: Install Valgrind - run: | - sudo apt-get update -qq - sudo apt-get install -y valgrind - - - name: Install iai-callgrind-runner - run: cargo install iai-callgrind-runner --version 0.14.2 --locked - - - uses: Swatinem/rust-cache@v2 - - - name: Run shared DiskANN IAI target - run: | - cargo bench --locked -p diskann \ - --bench bench_main_iai \ - --features pipnn,testing - fmt: name: format check runs-on: ubuntu-latest diff --git a/Cargo.lock b/Cargo.lock index f9d909d4a2..2588eee92a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -444,7 +444,6 @@ dependencies = [ "futures-util", "half", "hashbrown 0.16.1", - "iai-callgrind", "num-traits", "pin-project", "rand", diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index f434cfb7e2..db9a7a6ff4 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -690,3 +690,115 @@ mod tests { } } } +#[cfg(test)] +#[allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] +mod sgemm_aat_lower_tests { + use super::{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(3, 2, &a, &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(2, 2, &[0.0; 3], &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(2, 2, &[0.0; 4], &mut c).unwrap_err(); + + assert_eq!( + error, + SgemmError::InvalidMatrixDimensions { + matrix_name: MatrixName::C, + expected_rows: 2, + expected_cols: 2, + actual_len: 3, + } + ); + } + + #[test] + fn rejects_input_size_overflow() { + let error = sgemm_aat_lower(usize::MAX, 2, &[], &mut []).unwrap_err(); + + assert_eq!( + error, + SgemmError::DimensionOverflow { + matrix_name: MatrixName::A, + rows: usize::MAX, + cols: 2, + } + ); + } + + #[test] + fn rejects_output_size_overflow() { + let error = sgemm_aat_lower(usize::MAX, 0, &[], &mut []).unwrap_err(); + + assert_eq!( + error, + SgemmError::DimensionOverflow { + matrix_name: MatrixName::C, + rows: usize::MAX, + cols: usize::MAX, + } + ); + } +} diff --git a/diskann-linalg/tests/sgemm_aat_lower.rs b/diskann-linalg/tests/sgemm_aat_lower.rs deleted file mode 100644 index d19c92c761..0000000000 --- a/diskann-linalg/tests/sgemm_aat_lower.rs +++ /dev/null @@ -1,109 +0,0 @@ -/* - * 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(3, 2, &a, &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(2, 2, &[0.0; 3], &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(2, 2, &[0.0; 4], &mut c).unwrap_err(); - - assert_eq!( - error, - SgemmError::InvalidMatrixDimensions { - matrix_name: MatrixName::C, - expected_rows: 2, - expected_cols: 2, - actual_len: 3, - } - ); -} - -#[test] -fn rejects_input_size_overflow() { - let error = sgemm_aat_lower(usize::MAX, 2, &[], &mut []).unwrap_err(); - - assert_eq!( - error, - SgemmError::DimensionOverflow { - matrix_name: MatrixName::A, - rows: usize::MAX, - cols: 2, - } - ); -} - -#[test] -fn rejects_output_size_overflow() { - let error = sgemm_aat_lower(usize::MAX, 0, &[], &mut []).unwrap_err(); - - assert_eq!( - error, - SgemmError::DimensionOverflow { - matrix_name: MatrixName::C, - rows: usize::MAX, - cols: usize::MAX, - } - ); -} diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index a28b1b815d..72911c910f 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -32,7 +32,6 @@ 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 @@ -42,11 +41,6 @@ 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 deleted file mode 100644 index f914327879..0000000000 --- a/diskann/benches/bench_main_iai.rs +++ /dev/null @@ -1,22 +0,0 @@ -/* - * 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 deleted file mode 100644 index 82760c5e5e..0000000000 --- a/diskann/benches/benchmarks_iai/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -/* - * 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 deleted file mode 100644 index 0e6de738da..0000000000 --- a/diskann/benches/benchmarks_iai/pipnn_kernels.rs +++ /dev/null @@ -1,96 +0,0 @@ -/* - * 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, -); diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 4a289f3b01..92107b1f38 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -1075,3 +1075,427 @@ mod tests { } } } +#[cfg(test)] +#[allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] +mod integration_tests { + use std::cmp::Ordering; + + use super::{ + 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; 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; + 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_TARGETS: [usize; 2] = [1, 2]; + + fn differential_dots(metric: Metric, points: usize) -> Vec { + let mut dots = vec![f32::NAN; points * points]; + for source in 0..points { + dots[source * points + source] = + if metric == Metric::Cosine && source == ZERO_NORM_POSITION { + 0.0 + } else if source == DISTINCT_NORM_POSITION { + 2.0 + } else { + 1.0 + (source % NORM_PERIOD) as f32 + }; + for target in 0..source { + let pair = ((source * SOURCE_MIXER + target * TARGET_MIXER) % MIX_MODULUS) as f32 + - MIX_CENTER; + dots[source * points + target] = if TIED_TARGETS.contains(&target) { + 0.5 + } else { + pair * DOT_SCALE + }; + } + } + dots + } + + fn test_input(dots: &[f32], points: usize) -> LeafInput<'_> { + LeafInput { + dots: MatrixView::try_from(dots, points, points).unwrap(), + } + } + + 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(|source| { + let diagonal = dots[source * points + source]; + if metric == Metric::Cosine { + if diagonal < f32::MIN_POSITIVE { + 0.0 + } else { + diagonal.sqrt() + } + } else { + diagonal + } + }) + .collect(); + + for source in 0..points { + let mut candidates = Vec::with_capacity(points - 1); + for target in 0..points { + if target == source { + continue; + } + let (lower_source, lower_target) = if source > target { + (source, target) + } else { + (target, source) + }; + 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[source] + norms[target] - 2.0 * dot), + Metric::CosineNormalized => clamp(1.0 - dot), + Metric::InnerProduct => -dot, + Metric::Cosine => { + let denominator = norms[source] * norms[target]; + let similarity = if denominator == 0.0 { + 0.0 + } else { + dot / denominator + }; + clamp(1.0 - similarity) + } + }; + if distance.partial_cmp(&f32::INFINITY) == Some(Ordering::Less) { + candidates.push(LeafNeighbor::new(target as u32, distance)); + } + } + candidates.sort_by(|left, right| { + left.distance + .partial_cmp(&right.distance) + .expect("NaN distances were filtered") + }); + let count = candidates.len().min(leaf_k); + output[source * leaf_k..source * leaf_k + count].copy_from_slice(&candidates[..count]); + } + output + } + + 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( + test_input(dots, points), + MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), + &mut LeafKernelWorkspace::new(), + ) + .unwrap(); + (leaf_k, output) + } + + #[test] + fn prepared_dispatch_matches_reference_across_simd_width_boundaries() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + for points in SIMD_BOUNDARY_POINTS { + let dots = differential_dots(metric, points); + for requested_k in [1, 2, 3, 4, 5] { + 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}"); + } + } + } + } + + #[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, + ]; + + assert_eq!( + run_kernel(&dots, 4, 2, Metric::L2).1, + [ + 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, + ]; + for (metric, expected) in [ + (Metric::L2, [1, 2, 1]), + (Metric::Cosine, [1, 2, 1]), + (Metric::CosineNormalized, [1, 2, 1]), + (Metric::InnerProduct, [1, 2, 1]), + ] { + let positions: Vec<_> = run_kernel(&dots, 3, 1, metric) + .1 + .iter() + .map(|neighbor| neighbor.target) + .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_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)); + } + + #[test] + fn clamps_negative_distances_and_preserves_cosine_extremes() { + #[rustfmt::skip] + let out_of_range = [1.0, 0.0, 2.0, 1.0]; + assert_eq!( + 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 + ); + + #[rustfmt::skip] + let opposite = [1.0, 0.0, -2.0, 1.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_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_kernel(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, + 0.0 + ); + } + + #[test] + fn finite_max_distance_fills_the_final_simd_slot() { + let points = 9; + let mut dots = vec![0.0; points * points]; + dots[8 * points] = -f32::MAX; + + let (leaf_k, output) = run_kernel(&dots, points, points - 1, Metric::InnerProduct); + assert_eq!(leaf_k, 8); + assert_eq!( + output[8 * leaf_k + leaf_k - 1], + LeafNeighbor::new(0, f32::MAX) + ); + } + + #[test] + fn every_metric_ignores_nan_pairs() { + #[rustfmt::skip] + let dots = [ + 1.0, 0.0, 0.0, + f32::NAN, 1.0, 0.0, + 0.5, 0.25, 1.0, + ]; + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let output = run_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_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) + .nearest_neighbors( + test_input(&dots, 2), + MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), + &mut LeafKernelWorkspace::new(), + ) + .unwrap_err(); + + assert_eq!( + error, + LeafKernelError::InsufficientRankableNeighbors { + source_index: 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 (leaf_k, output) = run_kernel(&dots, 3, 99, Metric::L2); + + 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) + ); + } + } + + #[test] + fn accepts_empty_singleton_and_zero_k_inputs() { + 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_kernel(dots, points, requested_k, metric).0, 0); + } + } + + #[test] + fn rejects_non_square_input_and_invalid_output_dimensions() { + let dots = [0.0; 6]; + let non_square = LeafInput { + dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), + }; + let mut output = [LeafNeighbor::default(); 2]; + let kernel = LeafKernel::new(Metric::L2); + assert_eq!( + kernel.nearest_neighbors( + non_square, + MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), + &mut LeafKernelWorkspace::new(), + ), + Err(LeafKernelError::NonSquareDots { rows: 2, cols: 3 }) + ); + + let square = [0.0; 9]; + 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!( + kernel.nearest_neighbors( + test_input(&square, 3), + MutMatrixView::try_from(&mut too_many[..], 3, 3).unwrap(), + &mut LeafKernelWorkspace::new(), + ), + Err(LeafKernelError::InvalidNeighborCount { + points: 3, + neighbors: 3, + maximum: 2, + }) + ); + } + + #[test] + fn cosine_zero_norm_masks_nan_norm_at_simd_boundaries() { + for points in [9, 17] { + let mut dots = vec![0.0; points * points]; + for source in 1..points { + dots[source * points + source] = f32::NAN; + } + + 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}, source={source}" + ); + } + } + } + + #[test] + fn output_length_rejects_unrepresentable_point_count() { + assert_eq!( + leaf_output_len(usize::MAX, 1), + Err(LeafKernelError::TooManyPoints(usize::MAX)) + ); + } +} diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 3c88e5c9f2..95b650fa01 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -892,3 +892,385 @@ mod tests { assert_eq!(tracker[..4], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); } } +#[cfg(test)] +#[allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] +mod integration_tests { + use super::{ + MAX_PARTITION_FANOUT, PartitionInput, PartitionKernel, PartitionKernelError, + PartitionScales, + }; + use diskann_utils::views::{MatrixView, MutMatrixView}; + use diskann_vector::distance::Metric; + + fn test_input<'a>( + metric: Metric, + dots: &'a [f32], + point_count: usize, + leader_count: usize, + point_scales: &'a [f32], + leader_scales: &'a [f32], + ) -> PartitionInput<'a> { + let scales = match metric { + Metric::L2 => PartitionScales::L2 { + leader_squared_norms: leader_scales, + }, + Metric::Cosine => PartitionScales::Cosine { + point_squared_norms: point_scales, + leader_norms: leader_scales, + }, + Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, + }; + PartitionInput { + dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), + 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 { + point_squared_norms, + leader_norms, + } => (point_squared_norms, leader_norms), + PartitionScales::None => (&[][..], &[][..]), + }; + let mut assignments = vec![u32::MAX; point_count * fanout]; + for (point, (point_dots, point_assignments)) in input + .dots + .as_slice() + .chunks_exact(leader_count) + .zip(assignments.chunks_exact_mut(fanout)) + .enumerate() + { + let point_scale = point_scales.get(point).copied().unwrap_or(0.0); + let mut candidates: Vec<_> = point_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 point_norm = if point_scale < f32::MIN_POSITIVE { + 0.0 + } else { + point_scale.sqrt() + }; + 1.0 - if point_norm == 0.0 || leader_scale == 0.0 { + 0.0 + } else { + dot / (point_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 point_assignments.iter_mut().zip(candidates) { + *destination = leader; + } + } + assignments + } + + fn differential_data(metric: Metric, leader_count: usize) -> (Vec, Vec, Vec) { + let dots = (0..2 * leader_count) + .map(|index| { + 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 == leader_count { + f32::NAN + } else { + base * 0.25 + } + }) + .collect(); + let point_scales = if metric == Metric::Cosine { + vec![0.0, 16.0] + } else { + Vec::new() + }; + let leader_scales = match metric { + Metric::Cosine => (0..leader_count) + .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..leader_count) + .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, point_scales, leader_scales) + } + + fn run( + metric: Metric, + input: PartitionInput<'_>, + 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 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, + &dots, + 2, + leader_count, + &point_scales, + &leader_scales, + ); + for fanout in [1, 2, 16] { + if fanout >= leader_count { + continue; + } + assert_eq!( + run(metric, input, fanout).unwrap(), + brute_force_reference(input, fanout, metric), + "{metric:?}, leaders={leader_count}, 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, + test_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, point_scales, leader_scales, 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, + test_input(metric, &dots, 2, 3, point_scales, leader_scales), + 2, + ) + .unwrap(), + expected, + "metric {metric:?}" + ); + } + } + + #[test] + fn cosine_treats_a_zero_norm_as_zero_similarity() { + assert_eq!( + run( + Metric::Cosine, + test_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, + test_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, + test_input(Metric::InnerProduct, &[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), + 2, + ) + .unwrap(), + [1, 2] + ); + } + + #[test] + fn rejects_points_with_too_few_rankable_leaders() { + assert_eq!( + run( + Metric::InnerProduct, + test_input(Metric::InnerProduct, &[f32::NAN, 3.0], 1, 2, &[], &[]), + 2, + ), + Err(PartitionKernelError::InsufficientRankableLeaders { + point: 0, + fanout: 2, + }) + ); + } + + #[test] + fn accepts_empty_points_zero_fanout_and_largest_leader_id() { + run( + Metric::InnerProduct, + test_input(Metric::InnerProduct, &[], 0, 3, &[], &[]), + 2, + ) + .unwrap(); + run( + Metric::InnerProduct, + test_input(Metric::InnerProduct, &[1.0, 2.0, 3.0], 1, 3, &[], &[]), + 0, + ) + .unwrap(); + run( + Metric::InnerProduct, + test_input(Metric::InnerProduct, &[], 0, u32::MAX as usize, &[], &[]), + 0, + ) + .unwrap(); + + #[cfg(target_pointer_width = "64")] + assert_eq!( + run( + Metric::InnerProduct, + test_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 = test_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 = PartitionInput { + 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, + leader_count: 3, + maximum: MAX_PARTITION_FANOUT, + }) + ); + + let one = [0.0]; + assert_eq!( + run( + Metric::InnerProduct, + test_input(Metric::InnerProduct, &one, 1, 1, &[], &[]), + 2, + ), + Err(PartitionKernelError::InvalidFanout { + fanout: 2, + leader_count: 1, + maximum: MAX_PARTITION_FANOUT, + }) + ); + } +} diff --git a/diskann/tests/pipnn_leaf_kernel.rs b/diskann/tests/pipnn_leaf_kernel.rs deleted file mode 100644 index c6dc621d8b..0000000000 --- a/diskann/tests/pipnn_leaf_kernel.rs +++ /dev/null @@ -1,427 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -#![cfg(feature = "pipnn")] -#![allow( - clippy::expect_used, - clippy::unwrap_used, - reason = "deterministic test fixture construction must abort on invalid setup" -)] - -use std::cmp::Ordering; - -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; 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; -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_TARGETS: [usize; 2] = [1, 2]; - -fn differential_dots(metric: Metric, points: usize) -> Vec { - let mut dots = vec![f32::NAN; points * points]; - for source in 0..points { - dots[source * points + source] = if metric == Metric::Cosine && source == ZERO_NORM_POSITION - { - 0.0 - } else if source == DISTINCT_NORM_POSITION { - 2.0 - } else { - 1.0 + (source % NORM_PERIOD) as f32 - }; - for target in 0..source { - let pair = - ((source * SOURCE_MIXER + target * TARGET_MIXER) % MIX_MODULUS) as f32 - MIX_CENTER; - dots[source * points + target] = if TIED_TARGETS.contains(&target) { - 0.5 - } else { - pair * DOT_SCALE - }; - } - } - dots -} - -fn test_input(dots: &[f32], points: usize) -> LeafInput<'_> { - LeafInput { - dots: MatrixView::try_from(dots, points, points).unwrap(), - } -} - -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(|source| { - let diagonal = dots[source * points + source]; - if metric == Metric::Cosine { - if diagonal < f32::MIN_POSITIVE { - 0.0 - } else { - diagonal.sqrt() - } - } else { - diagonal - } - }) - .collect(); - - for source in 0..points { - let mut candidates = Vec::with_capacity(points - 1); - for target in 0..points { - if target == source { - continue; - } - let (lower_source, lower_target) = if source > target { - (source, target) - } else { - (target, source) - }; - 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[source] + norms[target] - 2.0 * dot), - Metric::CosineNormalized => clamp(1.0 - dot), - Metric::InnerProduct => -dot, - Metric::Cosine => { - let denominator = norms[source] * norms[target]; - let similarity = if denominator == 0.0 { - 0.0 - } else { - dot / denominator - }; - clamp(1.0 - similarity) - } - }; - if distance.partial_cmp(&f32::INFINITY) == Some(Ordering::Less) { - candidates.push(LeafNeighbor::new(target as u32, distance)); - } - } - candidates.sort_by(|left, right| { - left.distance - .partial_cmp(&right.distance) - .expect("NaN distances were filtered") - }); - let count = candidates.len().min(leaf_k); - output[source * leaf_k..source * leaf_k + count].copy_from_slice(&candidates[..count]); - } - output -} - -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( - test_input(dots, points), - MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), - &mut LeafKernelWorkspace::new(), - ) - .unwrap(); - (leaf_k, output) -} - -#[test] -fn prepared_dispatch_matches_reference_across_simd_width_boundaries() { - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - for points in SIMD_BOUNDARY_POINTS { - let dots = differential_dots(metric, points); - for requested_k in [1, 2, 3, 4, 5] { - 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}"); - } - } - } -} - -#[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, - ]; - - assert_eq!( - run_kernel(&dots, 4, 2, Metric::L2).1, - [ - 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, - ]; - for (metric, expected) in [ - (Metric::L2, [1, 2, 1]), - (Metric::Cosine, [1, 2, 1]), - (Metric::CosineNormalized, [1, 2, 1]), - (Metric::InnerProduct, [1, 2, 1]), - ] { - let positions: Vec<_> = run_kernel(&dots, 3, 1, metric) - .1 - .iter() - .map(|neighbor| neighbor.target) - .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_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)); -} - -#[test] -fn clamps_negative_distances_and_preserves_cosine_extremes() { - #[rustfmt::skip] - let out_of_range = [1.0, 0.0, 2.0, 1.0]; - assert_eq!( - 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 - ); - - #[rustfmt::skip] - let opposite = [1.0, 0.0, -2.0, 1.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_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_kernel(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, - 0.0 - ); -} - -#[test] -fn finite_max_distance_fills_the_final_simd_slot() { - let points = 9; - let mut dots = vec![0.0; points * points]; - dots[8 * points] = -f32::MAX; - - let (leaf_k, output) = run_kernel(&dots, points, points - 1, Metric::InnerProduct); - assert_eq!(leaf_k, 8); - assert_eq!( - output[8 * leaf_k + leaf_k - 1], - LeafNeighbor::new(0, f32::MAX) - ); -} - -#[test] -fn every_metric_ignores_nan_pairs() { - #[rustfmt::skip] - let dots = [ - 1.0, 0.0, 0.0, - f32::NAN, 1.0, 0.0, - 0.5, 0.25, 1.0, - ]; - - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let output = run_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_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) - .nearest_neighbors( - test_input(&dots, 2), - MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::new(), - ) - .unwrap_err(); - - assert_eq!( - error, - LeafKernelError::InsufficientRankableNeighbors { - source_index: 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 (leaf_k, output) = run_kernel(&dots, 3, 99, Metric::L2); - - 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) - ); - } -} - -#[test] -fn accepts_empty_singleton_and_zero_k_inputs() { - 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_kernel(dots, points, requested_k, metric).0, 0); - } -} - -#[test] -fn rejects_non_square_input_and_invalid_output_dimensions() { - let dots = [0.0; 6]; - let non_square = LeafInput { - dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), - }; - let mut output = [LeafNeighbor::default(); 2]; - let kernel = LeafKernel::new(Metric::L2); - assert_eq!( - kernel.nearest_neighbors( - non_square, - MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::new(), - ), - Err(LeafKernelError::NonSquareDots { rows: 2, cols: 3 }) - ); - - let square = [0.0; 9]; - 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!( - kernel.nearest_neighbors( - test_input(&square, 3), - MutMatrixView::try_from(&mut too_many[..], 3, 3).unwrap(), - &mut LeafKernelWorkspace::new(), - ), - Err(LeafKernelError::InvalidNeighborCount { - points: 3, - neighbors: 3, - maximum: 2, - }) - ); -} - -#[test] -fn cosine_zero_norm_masks_nan_norm_at_simd_boundaries() { - for points in [9, 17] { - let mut dots = vec![0.0; points * points]; - for source in 1..points { - dots[source * points + source] = f32::NAN; - } - - 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}, source={source}" - ); - } - } -} - -#[test] -fn output_length_rejects_unrepresentable_point_count() { - assert_eq!( - leaf_output_len(usize::MAX, 1), - Err(LeafKernelError::TooManyPoints(usize::MAX)) - ); -} diff --git a/diskann/tests/pipnn_partition_kernel.rs b/diskann/tests/pipnn_partition_kernel.rs deleted file mode 100644 index 597c472d3d..0000000000 --- a/diskann/tests/pipnn_partition_kernel.rs +++ /dev/null @@ -1,384 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -#![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, -}; -use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_vector::distance::Metric; - -fn test_input<'a>( - metric: Metric, - dots: &'a [f32], - point_count: usize, - leader_count: usize, - point_scales: &'a [f32], - leader_scales: &'a [f32], -) -> PartitionInput<'a> { - let scales = match metric { - Metric::L2 => PartitionScales::L2 { - leader_squared_norms: leader_scales, - }, - Metric::Cosine => PartitionScales::Cosine { - point_squared_norms: point_scales, - leader_norms: leader_scales, - }, - Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, - }; - PartitionInput { - dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), - 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 { - point_squared_norms, - leader_norms, - } => (point_squared_norms, leader_norms), - PartitionScales::None => (&[][..], &[][..]), - }; - let mut assignments = vec![u32::MAX; point_count * fanout]; - for (point, (point_dots, point_assignments)) in input - .dots - .as_slice() - .chunks_exact(leader_count) - .zip(assignments.chunks_exact_mut(fanout)) - .enumerate() - { - let point_scale = point_scales.get(point).copied().unwrap_or(0.0); - let mut candidates: Vec<_> = point_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 point_norm = if point_scale < f32::MIN_POSITIVE { - 0.0 - } else { - point_scale.sqrt() - }; - 1.0 - if point_norm == 0.0 || leader_scale == 0.0 { - 0.0 - } else { - dot / (point_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 point_assignments.iter_mut().zip(candidates) { - *destination = leader; - } - } - assignments -} - -fn differential_data(metric: Metric, leader_count: usize) -> (Vec, Vec, Vec) { - let dots = (0..2 * leader_count) - .map(|index| { - 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 == leader_count { - f32::NAN - } else { - base * 0.25 - } - }) - .collect(); - let point_scales = if metric == Metric::Cosine { - vec![0.0, 16.0] - } else { - Vec::new() - }; - let leader_scales = match metric { - Metric::Cosine => (0..leader_count) - .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..leader_count) - .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, point_scales, leader_scales) -} - -fn run( - metric: Metric, - input: PartitionInput<'_>, - 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 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, - &dots, - 2, - leader_count, - &point_scales, - &leader_scales, - ); - for fanout in [1, 2, 16] { - if fanout >= leader_count { - continue; - } - assert_eq!( - run(metric, input, fanout).unwrap(), - brute_force_reference(input, fanout, metric), - "{metric:?}, leaders={leader_count}, 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, - test_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, point_scales, leader_scales, 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, - test_input(metric, &dots, 2, 3, point_scales, leader_scales), - 2, - ) - .unwrap(), - expected, - "metric {metric:?}" - ); - } -} - -#[test] -fn cosine_treats_a_zero_norm_as_zero_similarity() { - assert_eq!( - run( - Metric::Cosine, - test_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, - test_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, - test_input(Metric::InnerProduct, &[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), - 2, - ) - .unwrap(), - [1, 2] - ); -} - -#[test] -fn rejects_points_with_too_few_rankable_leaders() { - assert_eq!( - run( - Metric::InnerProduct, - test_input(Metric::InnerProduct, &[f32::NAN, 3.0], 1, 2, &[], &[]), - 2, - ), - Err(PartitionKernelError::InsufficientRankableLeaders { - point: 0, - fanout: 2, - }) - ); -} - -#[test] -fn accepts_empty_points_zero_fanout_and_largest_leader_id() { - run( - Metric::InnerProduct, - test_input(Metric::InnerProduct, &[], 0, 3, &[], &[]), - 2, - ) - .unwrap(); - run( - Metric::InnerProduct, - test_input(Metric::InnerProduct, &[1.0, 2.0, 3.0], 1, 3, &[], &[]), - 0, - ) - .unwrap(); - run( - Metric::InnerProduct, - test_input(Metric::InnerProduct, &[], 0, u32::MAX as usize, &[], &[]), - 0, - ) - .unwrap(); - - #[cfg(target_pointer_width = "64")] - assert_eq!( - run( - Metric::InnerProduct, - test_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 = test_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 = PartitionInput { - 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, - leader_count: 3, - maximum: MAX_PARTITION_FANOUT, - }) - ); - - let one = [0.0]; - assert_eq!( - run( - Metric::InnerProduct, - test_input(Metric::InnerProduct, &one, 1, 1, &[], &[]), - 2, - ), - Err(PartitionKernelError::InvalidFanout { - fanout: 2, - leader_count: 1, - maximum: MAX_PARTITION_FANOUT, - }) - ); -} From ace78abbc308f808f71045319d92664f7f303d88 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:11:21 +0000 Subject: [PATCH 30/80] ci(pipnn): validate numerical feature Add only the PiPNN feature and SDE package coverage; leave workflow triggers and existing formatting unchanged. --- .github/workflows/disk-benchmarks.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/disk-benchmarks.yml b/.github/workflows/disk-benchmarks.yml index 341726c119..36753c846e 100644 --- a/.github/workflows/disk-benchmarks.yml +++ b/.github/workflows/disk-benchmarks.yml @@ -19,7 +19,6 @@ on: pull_request: branches: - main - - "pipnn-stack/**" paths: - 'diskann/**' - 'diskann-disk/**' From 699c601e2a3be334b397f45166ab199875e7bc34 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 03:12:52 +0000 Subject: [PATCH 31/80] fix(pipnn): enforce SIMD safety invariants Establish matrix, scale, output, and scratch lengths in the safe functions that contain unchecked loads instead of relying only on upstream comments. --- diskann/src/graph/pipnn/leaf_kernel.rs | 30 ++++++++++++++++++--- diskann/src/graph/pipnn/partition_kernel.rs | 22 ++++++++++++++- 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 92107b1f38..3fdf728a1a 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -781,8 +781,30 @@ fn process_pairs( u64: From<<::BitMask as SIMDMask>::Underlying>, { let point_count = input.dots.nrows(); + assert_eq!( + input.dots.ncols(), + point_count, + "validated leaf dot matrix must be square" + ); + assert_eq!( + output.source_count(), + point_count, + "validated leaf output must have one list per point" + ); + assert_eq!( + worst.len(), + point_count, + "validated leaf thresholds must have one value per point" + ); let dots = input.dots.as_slice(); let uses_norms = M::LEAF_SCALE.is_some(); + if uses_norms { + assert_eq!( + norms.len(), + point_count, + "validated leaf norms must have one value per point" + ); + } let worst_ptr = worst.as_mut_ptr(); // `source` starts at one because source zero has no strict-lower targets; @@ -796,7 +818,7 @@ fn process_pairs( } else { F::default(arch) }; - // SAFETY: `source < point_count == worst.len()` after validation. + // SAFETY: `source < point_count == worst.len()` by the assertions above. let mut source_worst = unsafe { *worst_ptr.add(source) }; let mut target = 0; @@ -804,7 +826,8 @@ fn process_pairs( // 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()`. + // SAFETY: the full target chunk lies below `source < point_count`, and + // the assertion above established `norms.len() == point_count`. unsafe { F::load_simd(arch, norms.as_ptr().add(target)) } } else { F::default(arch) @@ -813,7 +836,8 @@ fn process_pairs( // 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`. + // SAFETY: the full target chunk lies below `source < point_count`, and + // the assertion above established `worst.len() == point_count`. 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()); diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 95b650fa01..9f2e420730 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -604,7 +604,26 @@ fn process_points( M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { + let point_count = dots.nrows(); let leader_count = dots.ncols(); + assert!( + leader_count > 0, + "validated partition input must contain leaders" + ); + if M::PARTITION_POINT_SCALE.is_some() { + assert_eq!( + scales.point_scales.len(), + point_count, + "validated point scales must match point count" + ); + } + if M::PARTITION_LEADER_SCALE.is_some() { + assert_eq!( + scales.leader_scales.len(), + leader_count, + "validated leader scales must match leader count" + ); + } // 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 @@ -630,7 +649,8 @@ fn process_points( // 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 scale per leader. + // SAFETY: the assertion above established one scale per leader, and + // `base + F::LANES <= full <= leader_count`. unsafe { F::load_simd(arch, scales.leader_scales.as_ptr().add(base)) } } else { F::default(arch) From b01f0b9899fb62c3b2b9a6d85607e1cfda470c2a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:46:15 +0000 Subject: [PATCH 32/80] refactor(pipnn): pass leaf matrices directly Remove the one-field LeafInput wrapper while keeping prepared metric and ISA dispatch unchanged. --- diskann/src/graph/pipnn/leaf_kernel.rs | 72 ++++++++++---------------- diskann/src/graph/pipnn/mod.rs | 4 +- 2 files changed, 28 insertions(+), 48 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 3fdf728a1a..00ba595b75 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -51,9 +51,8 @@ //! //! # Main structures //! -//! - [`LeafKernel`] is the reusable public handle containing one prepared direct +//! - [`LeafKernel`] is the reusable 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 @@ -66,7 +65,7 @@ //! //! # Inputs and output //! -//! For `n` leaf points, [`LeafInput::dots`] is an `n × n` row-major matrix from +//! For `n` leaf points, the input 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 @@ -110,7 +109,7 @@ //! //! ``` //! use diskann::graph::pipnn::leaf_kernel::{ -//! leaf_output_len, LeafInput, LeafKernel, LeafKernelWorkspace, LeafNeighbor, +//! leaf_output_len, LeafKernel, LeafKernelWorkspace, LeafNeighbor, //! }; //! use diskann_utils::views::{MatrixView, MutMatrixView}; //! use diskann_vector::distance::Metric; @@ -121,9 +120,7 @@ //! 0.9, 1.0, f32::NAN, //! 0.1, 0.2, 1.0, //! ]; -//! let input = LeafInput { -//! dots: MatrixView::try_from(&dots[..], 3, 3).unwrap(), -//! }; +//! let input = 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(); @@ -172,13 +169,6 @@ impl Default for LeafNeighbor { } } -/// Square lower-triangular dot-product matrix for one leaf. -#[derive(Clone, Copy, Debug)] -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 LeafKernelWorkspace { @@ -316,7 +306,7 @@ pub fn leaf_output_len(points: usize, requested_k: usize) -> Result { - input: LeafInput<'a>, + input: MatrixView<'a, f32>, output: MutMatrixView<'a, LeafNeighbor>, workspace: &'a mut LeafKernelWorkspace, } @@ -387,7 +377,7 @@ impl LeafKernel { /// pointer; it performs no runtime ISA or metric dispatch. pub fn nearest_neighbors( &self, - input: LeafInput<'_>, + input: MatrixView<'_, f32>, output: MutMatrixView<'_, LeafNeighbor>, workspace: &mut LeafKernelWorkspace, ) -> Result<(), LeafKernelError> { @@ -517,11 +507,11 @@ where /// output or workspace mutation. Runtime is constant apart from view metadata /// checks; matrix contents are not scanned. fn validate( - input: LeafInput<'_>, + input: MatrixView<'_, f32>, output: &MutMatrixView<'_, LeafNeighbor>, ) -> Result<(), LeafKernelError> { - let point_count = input.dots.nrows(); - let dot_columns = input.dots.ncols(); + let point_count = input.nrows(); + let dot_columns = input.ncols(); if point_count > u32::MAX as usize { return Err(LeafKernelError::TooManyPoints(point_count)); } @@ -532,11 +522,7 @@ fn validate( }); } 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, - )?; + check_length("leaf dot-product matrix", input.as_slice().len(), dots_len)?; let output_len = checked_area("output", output.nrows(), output.ncols())?; check_length("output", output.as_slice().len(), output_len)?; @@ -570,14 +556,14 @@ fn validate( /// is returned without entering SIMD traversal. Work is `O(n)`, with at most /// `O(n)` retained capacity per buffer. fn prepare_workspace( - input: LeafInput<'_>, + input: MatrixView<'_, f32>, workspace: &mut LeafKernelWorkspace, ) -> Result<(), LeafKernelError> { - let points = input.dots.nrows(); + let points = input.nrows(); if M::LEAF_SCALE.is_some() { resize("norms", &mut workspace.norms, points, 0.0)?; for (source, norm) in workspace.norms.iter_mut().enumerate() { - *norm = M::LEAF_SCALE.transform(input.dots[(source, source)]); + *norm = M::LEAF_SCALE.transform(input[(source, source)]); } } else { workspace.norms.clear(); @@ -637,7 +623,7 @@ fn check_length( /// insertion to locate its dynamic slice. fn process_neighbor_width( arch: F::Arch, - input: LeafInput<'_>, + input: MatrixView<'_, f32>, neighbor_count: usize, output: &mut [LeafNeighbor], norms: &[f32], @@ -672,7 +658,7 @@ fn process_neighbor_width( /// leaf, keeping array conversion out of candidate insertion. fn process_fixed_width( arch: F::Arch, - input: LeafInput<'_>, + input: MatrixView<'_, f32>, output: &mut [LeafNeighbor], norms: &[f32], worst: &mut [f32], @@ -769,7 +755,7 @@ impl NeighborStorage for DynamicNeighborStorage<'_> { #[inline(never)] fn process_pairs( arch: F::Arch, - input: LeafInput<'_>, + input: MatrixView<'_, f32>, mut output: R, norms: &[f32], worst: &mut [f32], @@ -780,9 +766,9 @@ fn process_pairs( R: NeighborStorage, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let point_count = input.dots.nrows(); + let point_count = input.nrows(); assert_eq!( - input.dots.ncols(), + input.ncols(), point_count, "validated leaf dot matrix must be square" ); @@ -796,7 +782,7 @@ fn process_pairs( point_count, "validated leaf thresholds must have one value per point" ); - let dots = input.dots.as_slice(); + let dots = input.as_slice(); let uses_norms = M::LEAF_SCALE.is_some(); if uses_norms { assert_eq!( @@ -990,10 +976,8 @@ mod tests { dots } - fn test_input(dots: &[f32], points: usize) -> LeafInput<'_> { - LeafInput { - dots: MatrixView::try_from(dots, points, points).unwrap(), - } + fn test_input(dots: &[f32], points: usize) -> MatrixView<'_, f32> { + MatrixView::try_from(dots, points, points).unwrap() } fn insert_reference( @@ -1109,8 +1093,8 @@ mod integration_tests { use std::cmp::Ordering; use super::{ - LeafInput, LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, - leaf_neighbor_count, leaf_output_len, + LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, + leaf_output_len, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -1151,10 +1135,8 @@ mod integration_tests { dots } - fn test_input(dots: &[f32], points: usize) -> LeafInput<'_> { - LeafInput { - dots: MatrixView::try_from(dots, points, points).unwrap(), - } + fn test_input(dots: &[f32], points: usize) -> MatrixView<'_, f32> { + MatrixView::try_from(dots, points, points).unwrap() } fn brute_force_reference( @@ -1452,9 +1434,7 @@ mod integration_tests { #[test] fn rejects_non_square_input_and_invalid_output_dimensions() { let dots = [0.0; 6]; - let non_square = LeafInput { - dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), - }; + let non_square = MatrixView::try_from(&dots[..], 2, 3).unwrap(); let mut output = [LeafNeighbor::default(); 2]; let kernel = LeafKernel::new(Metric::L2); assert_eq!( diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 1fb729114c..1d2b0b87b7 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -86,8 +86,8 @@ //! ## [`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 +//! `sgemm_aat_lower`. [`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 From a21ee9c98291bab56a8245f09b2377a415b01550 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 05:52:15 +0000 Subject: [PATCH 33/80] refactor(pipnn): cap leaf kernel at k three Production uses fixed k=2/3. Reject larger widths and remove the duplicate dynamic storage and insertion path. --- diskann/src/graph/pipnn/leaf_kernel.rs | 269 +++++++++---------------- 1 file changed, 100 insertions(+), 169 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 00ba595b75..30f18f584c 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -57,20 +57,18 @@ //! 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_neighbor_width` chooses fixed storage for widths one through three. //! - `process_pairs` is the shared SIMD/scalar strict-lower traversal; -//! `insert_fixed_neighbor` and `insert_dynamic_neighbor` maintain stable sorted -//! output for both endpoints. +//! `insert_fixed_neighbor` maintains stable sorted output for both endpoints. //! //! # Inputs and output //! //! For `n` leaf points, the input 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. +//! pair dots. Output is an `n × k` [`LeafNeighbor`] matrix with `k <= 3`. 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: @@ -98,12 +96,11 @@ //! # 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`]. +//! `k = 0` returns before traversal. Supported widths `k = 1, 2, 3` use fixed +//! arrays and straight-line insertion, giving `O(n²)` 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 //! @@ -144,6 +141,9 @@ use diskann_wide::{ use super::kernel_metric::{KernelMetric, MetricVisitor, visit_metric}; +/// Largest leaf-local neighbor count supported by the fixed insertion kernel. +pub const MAX_LEAF_NEIGHBORS: usize = 3; + /// One leaf-local neighbor and its metric distance. #[derive(Clone, Copy, Debug, PartialEq)] pub struct LeafNeighbor { @@ -233,7 +233,7 @@ pub enum LeafKernelError { /// Supplied neighbor columns. columns: usize, }, - /// A source requests more non-self neighbors than the leaf contains. + /// A source requests more neighbors than the leaf or fixed kernel supports. #[error("invalid leaf neighbor count {neighbors} for {points} points; maximum is {maximum}")] InvalidNeighborCount { /// Point count in the leaf. @@ -264,13 +264,15 @@ 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. +/// Values above [`MAX_LEAF_NEIGHBORS`] are rejected. Otherwise 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`. +/// fit in `u32`, or [`LeafKernelError::InvalidNeighborCount`] when `requested_k` +/// exceeds [`MAX_LEAF_NEIGHBORS`]. /// /// # Performance /// @@ -279,6 +281,13 @@ pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result u32::MAX as usize { return Err(LeafKernelError::TooManyPoints(points)); } + if requested_k > MAX_LEAF_NEIGHBORS { + return Err(LeafKernelError::InvalidNeighborCount { + points, + neighbors: requested_k, + maximum: MAX_LEAF_NEIGHBORS, + }); + } Ok(requested_k.min(points.saturating_sub(1))) } @@ -349,7 +358,8 @@ impl LeafKernel { /// Select the nearest non-self leaf positions for every source point. /// /// `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`. + /// neighbor count for this leaf and must not exceed either `point_count - 1` + /// or [`MAX_LEAF_NEIGHBORS`]. /// Equal distances retain pair scan order. /// /// `input` supplies the square lower-triangular dot matrix. `output` is @@ -533,7 +543,7 @@ fn validate( columns: output.ncols(), }); } - let maximum_neighbors = point_count.saturating_sub(1); + let maximum_neighbors = point_count.saturating_sub(1).min(MAX_LEAF_NEIGHBORS); let neighbor_count = output.ncols(); if neighbor_count > maximum_neighbors { return Err(LeafKernelError::InvalidNeighborCount { @@ -611,16 +621,12 @@ fn check_length( } } -/// Convert neighbor count into fixed source storage or the dynamic fallback. +/// Convert the validated neighbor count into fixed source storage. /// -/// 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. +/// This branch runs once per leaf. `as_chunks_mut` performs one safe conversion, +/// avoiding per-candidate slice-to-array checks. `output` contains +/// `point_count * neighbor_count` initialized slots; `norms` and `worst` satisfy +/// the invariants established by `prepare_workspace`. fn process_neighbor_width( arch: F::Arch, input: MatrixView<'_, f32>, @@ -638,16 +644,7 @@ fn process_neighbor_width( 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, - DynamicNeighborStorage { - values: output, - neighbor_count: dynamic_count, - }, - norms, - worst, - ), + _ => unreachable!("validated leaf neighbor count must be in 1..=3"), } } @@ -670,61 +667,7 @@ fn process_fixed_width( { let (neighbor_lists, remainder) = output.as_chunks_mut::(); debug_assert!(remainder.is_empty()); - process_pairs::( - arch, - input, - FixedNeighborStorage(neighbor_lists), - norms, - worst, - ); -} - -/// Mutable neighbor-list adapter used by the shared pair traversal. -/// -/// Implementations own the exclusive output borrow for the whole scan. Each -/// 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; -} - -struct FixedNeighborStorage<'a, const N: usize>(&'a mut [[LeafNeighbor; N]]); - -impl NeighborStorage for FixedNeighborStorage<'_, N> { - #[inline(always)] - fn source_count(&self) -> usize { - self.0.len() - } - - #[inline(always)] - fn insert(&mut self, source: usize, target: u32, distance: f32) -> f32 { - insert_fixed_neighbor(&mut self.0[source], target, distance) - } -} - -struct DynamicNeighborStorage<'a> { - values: &'a mut [LeafNeighbor], - neighbor_count: usize, -} - -impl NeighborStorage for DynamicNeighborStorage<'_> { - #[inline(always)] - fn source_count(&self) -> usize { - self.values.len() / self.neighbor_count - } - - #[inline(always)] - 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, - ) - } + process_pairs::(arch, input, neighbor_lists, norms, worst); } /// Scan the strict lower triangle once and update both endpoint sources. @@ -742,28 +685,23 @@ impl NeighborStorage for DynamicNeighborStorage<'_> { /// 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 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. +/// `M` is concrete before type erasure and `N` is one through three. `input` +/// supplies `n × n` dots, `output` owns `n` sorted fixed-size lists, `norms` +/// holds metric scales when required, and `worst` mirrors every list's final +/// distance. The function 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. #[inline(never)] -fn process_pairs( +fn process_pairs( arch: F::Arch, input: MatrixView<'_, f32>, - mut output: R, + output: &mut [[LeafNeighbor; N]], norms: &[f32], worst: &mut [f32], ) where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: KernelMetric, - R: NeighborStorage, u64: From<<::BitMask as SIMDMask>::Underlying>, { let point_count = input.nrows(); @@ -773,7 +711,7 @@ fn process_pairs( "validated leaf dot matrix must be square" ); assert_eq!( - output.source_count(), + output.len(), point_count, "validated leaf output must have one list per point" ); @@ -838,7 +776,11 @@ fn process_pairs( source_bits &= source_bits - 1; let distance = values[lane]; if distance < source_worst { - source_worst = output.insert(source, (target + lane) as u32, distance); + source_worst = insert_fixed_neighbor( + &mut output[source], + (target + lane) as u32, + distance, + ); } } @@ -847,7 +789,11 @@ fn process_pairs( 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]); + let new_worst = insert_fixed_neighbor( + &mut output[target_source], + source as u32, + values[lane], + ); // SAFETY: `target_source < source < worst.len()`. unsafe { *worst_ptr.add(target_source) = new_worst }; } @@ -866,12 +812,12 @@ fn process_pairs( }; let distance = M::leaf_distance_scalar(dot, source_scale, target_scale); if distance < source_worst { - source_worst = output.insert(source, target as u32, distance); + source_worst = insert_fixed_neighbor(&mut output[source], target as u32, distance); } // 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); + let new_worst = insert_fixed_neighbor(&mut output[target], source as u32, distance); // SAFETY: `target < source < worst.len()`. unsafe { *worst_ptr.add(target) = new_worst }; } @@ -881,7 +827,7 @@ fn process_pairs( unsafe { *worst_ptr.add(source) = source_worst }; } - debug_assert_eq!(output.source_count(), point_count); + debug_assert_eq!(output.len(), point_count); } /// Insert into a fixed-width neighbor list and return its new worst distance. @@ -936,26 +882,6 @@ fn insert_fixed_neighbor( } } -/// 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 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; - neighbors[last] = LeafNeighbor::new(target, distance); - let mut index = last; - while index > 0 && neighbors[index].distance < neighbors[index - 1].distance { - neighbors.swap(index, index - 1); - index -= 1; - } - neighbors[last].distance -} - #[cfg(test)] mod tests { use super::*; @@ -980,33 +906,16 @@ mod tests { MatrixView::try_from(dots, points, points).unwrap() } - fn insert_reference( - output: &mut [LeafNeighbor], - worst: &mut [f32], - neighbor_count: usize, - source: usize, - target: u32, - distance: f32, - ) { - if distance.partial_cmp(&worst[source]) != Some(std::cmp::Ordering::Less) { - return; - } - worst[source] = insert_dynamic_neighbor( - &mut output[source * neighbor_count..(source + 1) * neighbor_count], - target, - distance, - ); - } - #[test] - fn scalar_insertion_orders_candidates_and_rejects_nan() { - let mut output = [LeafNeighbor::default(); 4]; - let mut worst = [f32::INFINITY]; + fn fixed_insertion_orders_candidates() { + let mut output = [LeafNeighbor::default(); 3]; + let mut worst = f32::INFINITY; 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); + if distance < worst { + worst = insert_fixed_neighbor(&mut output, target, distance); + } } - insert_reference(&mut output, &mut worst, 4, 0, 5, f32::NAN); assert_eq!( output, @@ -1014,17 +923,24 @@ mod tests { 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]); + assert_eq!(worst, 2.0); } #[test] - fn output_length_clamps_to_non_self_neighbors() { + fn output_length_clamps_to_non_self_neighbors_and_rejects_large_k() { 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(4, 3).unwrap(), 12); + assert_eq!( + leaf_output_len(4, 4), + Err(LeafKernelError::InvalidNeighborCount { + points: 4, + neighbors: 4, + maximum: MAX_LEAF_NEIGHBORS, + }) + ); #[cfg(target_pointer_width = "64")] assert_eq!( leaf_output_len(u32::MAX as usize + 1, 1), @@ -1093,8 +1009,8 @@ mod integration_tests { use std::cmp::Ordering; use super::{ - LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, - leaf_output_len, + LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, MAX_LEAF_NEIGHBORS, + leaf_neighbor_count, leaf_output_len, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -1236,7 +1152,7 @@ mod integration_tests { ] { for points in SIMD_BOUNDARY_POINTS { let dots = differential_dots(metric, points); - for requested_k in [1, 2, 3, 4, 5] { + for requested_k in [1, 2, 3] { 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}"); @@ -1345,15 +1261,15 @@ mod integration_tests { } #[test] - fn finite_max_distance_fills_the_final_simd_slot() { - let points = 9; + fn finite_max_distance_fills_the_final_fixed_slot() { + let points = 4; let mut dots = vec![0.0; points * points]; - dots[8 * points] = -f32::MAX; + dots[3 * points] = -f32::MAX; - let (leaf_k, output) = run_kernel(&dots, points, points - 1, Metric::InnerProduct); - assert_eq!(leaf_k, 8); + let (leaf_k, output) = run_kernel(&dots, points, MAX_LEAF_NEIGHBORS, Metric::InnerProduct); + assert_eq!(leaf_k, MAX_LEAF_NEIGHBORS); assert_eq!( - output[8 * leaf_k + leaf_k - 1], + output[3 * leaf_k + leaf_k - 1], LeafNeighbor::new(0, f32::MAX) ); } @@ -1408,7 +1324,7 @@ mod integration_tests { 0.0, 1.0, 3.0, 0.0, 0.0, 1.0, ]; - let (leaf_k, output) = run_kernel(&dots, 3, 99, Metric::L2); + let (leaf_k, output) = run_kernel(&dots, 3, MAX_LEAF_NEIGHBORS, Metric::L2); assert_eq!(leaf_k, 2); for (source, neighbors) in output.chunks_exact(leaf_k).enumerate() { @@ -1474,6 +1390,21 @@ mod integration_tests { maximum: 2, }) ); + + let square = [0.0; 25]; + let mut too_wide = [LeafNeighbor::default(); 20]; + assert_eq!( + kernel.nearest_neighbors( + test_input(&square, 5), + MutMatrixView::try_from(&mut too_wide[..], 5, 4).unwrap(), + &mut LeafKernelWorkspace::new(), + ), + Err(LeafKernelError::InvalidNeighborCount { + points: 5, + neighbors: 4, + maximum: MAX_LEAF_NEIGHBORS, + }) + ); } #[test] From 41826253700e9a225d583b811883cdc3cbe6b18c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:23:41 +0000 Subject: [PATCH 34/80] fix(pipnn): return errors for invalid leaf widths Keep supported fixed widths panic-free and remove internal unreachable branches. --- diskann/src/graph/pipnn/leaf_kernel.rs | 64 ++++++++++---------------- 1 file changed, 24 insertions(+), 40 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 30f18f584c..ad2c88f5c0 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -97,7 +97,7 @@ //! //! With `k > 0`, the kernel evaluates exactly `n(n - 1) / 2` pair distances; //! `k = 0` returns before traversal. Supported widths `k = 1, 2, 3` use fixed -//! arrays and straight-line insertion, giving `O(n²)` work. Scratch is `O(n)` +//! arrays and bounded insertion, giving `O(n²)` 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`]. @@ -487,7 +487,7 @@ where call.output.as_mut_slice(), &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 @@ -634,7 +634,8 @@ fn process_neighbor_width( output: &mut [LeafNeighbor], norms: &[f32], worst: &mut [f32], -) where +) -> Result<(), LeafKernelError> +where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: KernelMetric, @@ -644,8 +645,15 @@ fn process_neighbor_width( 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), - _ => unreachable!("validated leaf neighbor count must be in 1..=3"), + _ => { + return Err(LeafKernelError::InvalidNeighborCount { + points: input.nrows(), + neighbors: neighbor_count, + maximum: MAX_LEAF_NEIGHBORS, + }); + } } + Ok(()) } /// Reinterpret validated output as one fixed array per source, then run shared @@ -832,54 +840,30 @@ fn process_pairs( /// Insert into a fixed-width neighbor list and return its new worst distance. /// -/// Production widths one through three use straight-line shifts. Strict `<` +/// Width is a compile-time constant from one through three. 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`. +/// final-slot distance. Insertion is allocation-free and bounded by three swaps. #[inline(always)] fn insert_fixed_neighbor( neighbors: &mut [LeafNeighbor; N], target: u32, distance: f32, ) -> f32 { - let entry = LeafNeighbor::new(target, distance); - match N { - 1 => { - neighbors[0] = entry; - distance - } - 2 => { - let first = neighbors[0]; - if distance < first.distance { - neighbors[0] = entry; - neighbors[1] = first; - first.distance - } else { - neighbors[1] = entry; - distance - } - } - 3 => { - let (first, second) = (neighbors[0], neighbors[1]); - if distance < first.distance { - neighbors[0] = entry; - neighbors[1] = first; - neighbors[2] = second; - } else if distance < second.distance { - neighbors[1] = entry; - neighbors[2] = second; - } else { - neighbors[2] = entry; - return distance; - } - second.distance - } - _ => unreachable!("fixed leaf widths are one through three"), + if N == 0 { + return f32::INFINITY; + } + let last = N - 1; + neighbors[last] = LeafNeighbor::new(target, distance); + let mut index = last; + while index > 0 && neighbors[index].distance < neighbors[index - 1].distance { + neighbors.swap(index, index - 1); + index -= 1; } + neighbors[last].distance } #[cfg(test)] From 66f51bf8e719fd88a8f0efd6b1d2562bfce2da00 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:34:27 +0000 Subject: [PATCH 35/80] fix(utils): reject overflowing matrix shapes Make every MatrixBase constructor preserve the backing-length invariant before PiPNN relies on MatrixView shape. --- diskann-utils/src/views.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs index a9352918c9..a20b8b5fcf 100644 --- a/diskann-utils/src/views.rs +++ b/diskann-utils/src/views.rs @@ -195,8 +195,14 @@ impl MatrixBase> { where U: Generator, { - let data: Box<[T]> = (0..nrows * ncols).map(|_| generator.generate()).collect(); - debug_assert_eq!(data.len(), nrows * ncols); + let len = nrows.checked_mul(ncols); + assert!( + len.is_some(), + "matrix shape {nrows} x {ncols} overflows usize" + ); + let len = len.unwrap_or(0); + let data: Box<[T]> = (0..len).map(|_| generator.generate()).collect(); + debug_assert_eq!(data.len(), len); Self { data, nrows, ncols } } } @@ -211,7 +217,7 @@ where /// The length of the base must be equal to `nrows * ncols`. pub fn try_from(data: T, nrows: usize, ncols: usize) -> Result> { let len = data.as_slice().len(); - if len != nrows * ncols { + if nrows.checked_mul(ncols) != Some(len) { Err(TryFromError { data, nrows, ncols }) } else { Ok(Self { data, nrows, ncols }) @@ -1053,6 +1059,8 @@ mod tests { m.unwrap_err().to_string(), "tried to construct a matrix view with 5 rows and 4 columns over a slice of length 12" ); + + assert!(MatrixView::try_from(&[] as &[usize], usize::MAX, 2).is_err()); } #[test] @@ -1428,6 +1436,8 @@ mod tests { assert_eq!(m.nrows(), 5); assert_eq!(m.ncols(), 1); assert!(m.as_slice().iter().all(|&x| x == 9)); + + assert!(std::panic::catch_unwind(|| Matrix::new(0, usize::MAX, 2)).is_err()); } #[test] From ae89cf75901fc672168de70f1deec898065f01d0 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:34:27 +0000 Subject: [PATCH 36/80] refactor(pipnn): trust validated matrix views Remove duplicate backing-area checks and trim kernel module docs to contracts, invariants, and complexity. --- diskann/src/graph/pipnn/kernel_metric.rs | 66 ++----- diskann/src/graph/pipnn/leaf_kernel.rs | 181 +++---------------- diskann/src/graph/pipnn/partition_kernel.rs | 184 ++------------------ 3 files changed, 50 insertions(+), 381 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index edb3e6a842..a6b0c3a869 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -3,63 +3,21 @@ * Licensed under the MIT license. */ -//! Metric marker types shared by the partition and leaf kernels. +//! Metric formulas shared by PiPNN 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. +//! Runtime [`Metric`] selection happens once while preparing a dispatched +//! kernel. Zero-sized marker types then monomorphize scalar and SIMD formulas, +//! so hot loops contain neither metric matches nor trait objects. //! -//! [`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. +//! All formulas produce ascending scores. L2 uses squared norms; unnormalized +//! cosine uses norms with zero/subnormal inputs mapped to zero similarity; +//! normalized cosine and inner product need no scales. Ordered comparisons leave +//! NaN non-rankable. The L2 partition scalar tail deliberately keeps its +//! non-fused operation order because rounding can change leader ties. //! -//! 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 -//! -//! 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. -//! -//! # 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. +//! [`ScaleKind`] records required scale representation. [`KernelMetric`] owns +//! leaf and partition formulas. [`MetricVisitor`] performs runtime-to-marker +//! conversion before the final architecture-specific function pointer is stored. use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index ad2c88f5c0..d60c821b89 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -3,131 +3,26 @@ * Licensed under the MIT license. */ -//! Prepared nearest-neighbor kernels over a leaf's lower dot-product matrix. +//! Prepared leaf-local top-k selection over a lower-triangular Gram 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. +//! Caller supplies an `n × n` [`MatrixView`] produced by `sgemm_aat_lower`. +//! Diagonal entries provide metric scales; only strict-lower pair dots are read. +//! Each pair is evaluated once and offered to both endpoint rows. //! -//! 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. +//! Output is an `n × k` matrix of sorted [`LeafNeighbor`] values with leaf-local +//! targets. Supported `k` is zero through [`MAX_LEAF_NEIGHBORS`]; positive widths +//! use fixed arrays. Strict comparisons preserve encounter order for ties and +//! reject NaN. L2, cosine, normalized cosine, and inner product share the same +//! scalar/SIMD traversal. //! -//! 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. +//! [`LeafKernel::new`] selects metric and runtime architecture once, storing one +//! direct function pointer reused across leaves. Every call validates square +//! shape, row count, local-ID bounds, and output width before scratch mutation or +//! unchecked SIMD loads. [`LeafKernelWorkspace`] retains per-worker norm and +//! threshold buffers. //! -//! `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 -//! 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 + runtime architecture -//! │ -//! v -//! prepared Dispatched1 handle -//! │ reused with input + output.ncols() -//! v -//! shape validation -> scale scratch -> strict-lower scan -> sorted neighbor slots -//! ``` -//! -//! 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 handle containing one prepared direct -//! function pointer. -//! - [`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. -//! - `process_pairs` is the shared SIMD/scalar strict-lower traversal; -//! `insert_fixed_neighbor` maintains stable sorted output for both endpoints. -//! -//! # Inputs and output -//! -//! For `n` leaf points, the input 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 with `k <= 3`. 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 norm below `sqrt(f32::MIN_POSITIVE)` 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. Supported widths `k = 1, 2, 3` use fixed -//! arrays and bounded insertion, giving `O(n²)` 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::graph::pipnn::leaf_kernel::{ -//! leaf_output_len, 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 = 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]); -//! ``` +//! Work is `n(n - 1) / 2` distance evaluations with constant bounded insertion; +//! scratch is `O(n)` and output is `O(nk)`. use std::marker::PhantomData; @@ -213,16 +108,6 @@ pub enum LeafKernelError { /// Declared column count. cols: usize, }, - /// 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. - buffer: &'static str, - /// Required length. - expected: usize, - /// Supplied length. - actual: usize, - }, /// The output matrix does not have one row per input point. #[error("invalid output row count: expected {expected}, got {actual} with {columns} columns")] InvalidOutputRows { @@ -507,15 +392,10 @@ 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 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. +/// `MatrixView` and `MutMatrixView` construction guarantee exact, non-overflowing +/// backing lengths. This check establishes square dots, representable local IDs, +/// and an output width bounded by the point count and fixed kernel capacity. +/// Failure returns [`LeafKernelError`] before output or workspace mutation. fn validate( input: MatrixView<'_, f32>, output: &MutMatrixView<'_, LeafNeighbor>, @@ -531,11 +411,6 @@ fn validate( cols: dot_columns, }); } - let dots_len = checked_area("leaf dot-product matrix", point_count, dot_columns)?; - check_length("leaf dot-product matrix", input.as_slice().len(), dots_len)?; - let output_len = checked_area("output", output.nrows(), output.ncols())?; - check_length("output", output.as_slice().len(), output_len)?; - if output.nrows() != point_count { return Err(LeafKernelError::InvalidOutputRows { expected: point_count, @@ -605,22 +480,6 @@ fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> Result Result<(), LeafKernelError> { - if actual == expected { - Ok(()) - } else { - Err(LeafKernelError::InvalidBufferLength { - buffer, - expected, - actual, - }) - } -} - /// Convert the validated neighbor count into fixed source storage. /// /// This branch runs once per leaf. `as_chunks_mut` performs one safe conversion, diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 9f2e420730..c1cb81e4df 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -3,130 +3,22 @@ * Licensed under the MIT license. */ -//! Prepared distance and top-k kernels for partition assignment. +//! Prepared nearest-leader selection for PiPNN 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. +//! Caller supplies a row-major point-by-leader dot matrix plus metric-specific +//! [`PartitionScales`]. Output contains sorted leader-column positions for each +//! point; fanout is the output width and cannot exceed +//! [`MAX_PARTITION_FANOUT`] or the leader count. //! -//! 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. +//! [`PartitionKernel::new`] selects metric and runtime architecture once and +//! stores a direct function pointer reused by every point stripe. Calls validate +//! row counts, scale variants and lengths, fanout, and leader-ID representation +//! before output mutation or unchecked SIMD loads. //! -//! 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; -//! repeated stripes call a direct `diskann-wide` function pointer with no ISA or -//! metric branch in the point loop. -//! -//! L2 deliberately omits the point norm because it is constant across every -//! 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 -//! metric + runtime architecture -//! │ -//! v -//! prepared Dispatched2 handle -//! │ reused for every point stripe -//! v -//! shape/scale validation -> SIMD chunks + scalar tail -> sorted leader IDs -//! ``` -//! -//! 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::graph::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]); -//! ``` +//! L2 omits the point norm because it cannot change one point's leader order. +//! Strict comparisons preserve scan order for ties and leave NaN non-rankable. +//! Each point evaluates every leader; competitive scores move through a fixed +//! stack tracker of at most [`MAX_PARTITION_FANOUT`] entries. use std::marker::PhantomData; @@ -189,16 +81,6 @@ pub struct PartitionInput<'a> { /// Validation error returned by [`PartitionKernel::nearest_leaders`]. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub enum PartitionKernelError { - /// A declared matrix 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, - }, /// 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" @@ -443,15 +325,11 @@ struct ScaleSlices<'a> { /// 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. -/// -/// `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. +/// `MatrixView` and `MutMatrixView` construction guarantee exact, +/// non-overflowing backing lengths. The `PartitionScales` variant must match +/// concrete metric `M`, preventing plausible but incorrect norm units from +/// crossing the interface. Success returns normalized scale slices and +/// establishes representable leader IDs plus bounded fanout. fn validate<'a, M: KernelMetric>( input: PartitionInput<'a>, output: &MutMatrixView<'_, u32>, @@ -460,11 +338,6 @@ fn validate<'a, M: KernelMetric>( let leader_count = input.dots.ncols(); let fanout = output.ncols(); - 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() != point_count { return Err(PartitionKernelError::InvalidOutputShape { expected_rows: point_count, @@ -548,15 +421,6 @@ const fn expected_scale_len(kind: ScaleKind, count: usize) -> usize { if kind.is_some() { count } else { 0 } } -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, @@ -889,18 +753,6 @@ mod tests { 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 tracker = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; From af4470bac8ca7a6113bd94f031812de269fc3189 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:53:54 +0000 Subject: [PATCH 37/80] perf(pipnn): retain measured fixed-k insertion Keep panic-free invalid-width handling while restoring the k=1/2/3 shifts that beat the generic bubble loop in Callgrind. --- diskann/src/graph/pipnn/leaf_kernel.rs | 49 +++++++++++++++++++------- 1 file changed, 37 insertions(+), 12 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index d60c821b89..f9426bf9db 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -701,28 +701,53 @@ fn process_pairs( /// /// Width is a compile-time constant from one through three. Strict `<` /// comparisons preserve scan order for ties; callers already rejected NaN via -/// the eligibility comparison. +/// the eligibility comparison. Explicit shifts save about 0.5% estimated cycles +/// versus the generic bubble loop in the local Callgrind `k=3` fixture. /// /// `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 bounded by three swaps. +/// final-slot distance. Unsupported instantiations return an underfill sentinel; +/// `process_neighbor_width` never constructs them. #[inline(always)] fn insert_fixed_neighbor( neighbors: &mut [LeafNeighbor; N], target: u32, distance: f32, ) -> f32 { - if N == 0 { - return f32::INFINITY; - } - let last = N - 1; - neighbors[last] = LeafNeighbor::new(target, distance); - let mut index = last; - while index > 0 && neighbors[index].distance < neighbors[index - 1].distance { - neighbors.swap(index, index - 1); - index -= 1; + let entry = LeafNeighbor::new(target, distance); + match N { + 1 => { + neighbors[0] = entry; + distance + } + 2 => { + let first = neighbors[0]; + if distance < first.distance { + neighbors[0] = entry; + neighbors[1] = first; + first.distance + } else { + neighbors[1] = entry; + distance + } + } + 3 => { + let (first, second) = (neighbors[0], neighbors[1]); + if distance < first.distance { + neighbors[0] = entry; + neighbors[1] = first; + neighbors[2] = second; + } else if distance < second.distance { + neighbors[1] = entry; + neighbors[2] = second; + } else { + neighbors[2] = entry; + return distance; + } + second.distance + } + _ => f32::INFINITY, } - neighbors[last].distance } #[cfg(test)] From f20581aa6681c3a1085e167f4959a14a51c2a227 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:58:29 +0000 Subject: [PATCH 38/80] docs(pipnn): remove stale dynamic-k reference --- diskann/src/graph/pipnn/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 1d2b0b87b7..b1aa678676 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -91,8 +91,8 @@ //! [`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. +//! describes fixed-width selection, `process_pairs`, and stable endpoint +//! insertion. //! //! ## `kernel_metric` //! From 0826436b3c9670626859a7b5abb8a7fb73cb33cb Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 07:12:59 +0000 Subject: [PATCH 39/80] refactor(pipnn): constrain SIMD lanes locally Use the existing f32x16 ConstLanes contract so PiPNN can read to_array results without changing diskann-wide traits. --- diskann-wide/src/traits.rs | 2 +- diskann/src/graph/pipnn/leaf_kernel.rs | 11 +++++------ diskann/src/graph/pipnn/partition_kernel.rs | 9 ++++----- 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/diskann-wide/src/traits.rs b/diskann-wide/src/traits.rs index a233622615..09150f0c7d 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: AsRef<[T]> + AsMut<[T]>; + type Type; } /// Map scalar + lengths to arrays. diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index f9426bf9db..807ac17ac9 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -29,7 +29,7 @@ use std::marker::PhantomData; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use diskann_wide::{ - Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, + Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, arch::{self, Dispatched1, FTarget1}, lifetime::AddLifetime, }; @@ -495,7 +495,7 @@ fn process_neighbor_width( worst: &mut [f32], ) -> Result<(), LeafKernelError> where - F: SIMDVector + SIMDFloat + std::ops::Div, + F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, @@ -527,7 +527,7 @@ fn process_fixed_width( norms: &[f32], worst: &mut [f32], ) where - F: SIMDVector + SIMDFloat + std::ops::Div, + F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, @@ -566,7 +566,7 @@ fn process_pairs( norms: &[f32], worst: &mut [f32], ) where - F: SIMDVector + SIMDFloat + std::ops::Div, + F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, @@ -635,8 +635,7 @@ fn process_pairs( 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 values: [f32; 16] = distances.to_array(); let mut source_bits = source_bits; while source_bits != 0 { let lane = source_bits.trailing_zeros() as usize; diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index c1cb81e4df..ffe4182353 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -25,7 +25,7 @@ use std::marker::PhantomData; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use diskann_wide::{ - Architecture, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, + Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, arch::{self, Dispatched2, FTarget2}, lifetime::AddLifetime, }; @@ -463,7 +463,7 @@ fn process_points( fanout: usize, output: &mut [u32], ) where - F: SIMDVector + SIMDFloat + std::ops::Div, + F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, @@ -564,7 +564,7 @@ fn insert_leader_lanes( tracker: &mut LeaderTracker, fanout: usize, ) where - F: SIMDVector + SIMDPartialOrd, + F: SIMDVector> + SIMDPartialOrd, u64: From<<::BitMask as SIMDMask>::Underlying>, { let threshold = F::splat(distances.arch(), tracker[fanout - 1].1); @@ -573,8 +573,7 @@ fn insert_leader_lanes( return; } - let values = distances.to_array(); - let values = values.as_ref(); + let values: [f32; 16] = distances.to_array(); let mut lanes = u64::from(eligible.bitmask().to_underlying()); while lanes != 0 { let lane = lanes.trailing_zeros() as usize; From d2beeee97e842168191c2d0e3ba4c41d084163bd Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:38:56 +0000 Subject: [PATCH 40/80] refactor(pipnn): remove stored kernel pointers Dispatch runtime convenience calls immediately, expose direct generic stage kernels, and replace the arbitrary fanout cap with caller-owned reusable tracker storage. --- diskann/src/graph/pipnn/leaf_kernel.rs | 220 +++++-------- diskann/src/graph/pipnn/partition_kernel.rs | 326 ++++++++------------ 2 files changed, 215 insertions(+), 331 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 807ac17ac9..209bf4537e 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -15,23 +15,20 @@ //! reject NaN. L2, cosine, normalized cosine, and inner product share the same //! scalar/SIMD traversal. //! -//! [`LeafKernel::new`] selects metric and runtime architecture once, storing one -//! direct function pointer reused across leaves. Every call validates square -//! shape, row count, local-ID bounds, and output width before scratch mutation or -//! unchecked SIMD loads. [`LeafKernelWorkspace`] retains per-worker norm and -//! threshold buffers. +//! Runtime callers may use [`LeafKernel`]; production dispatches once at the leaf +//! stage boundary and calls the generic kernel directly. Every call validates +//! square shape, row count, local-ID bounds, and output width before scratch +//! mutation or unchecked SIMD loads. [`LeafKernelWorkspace`] retains per-worker +//! norm and threshold buffers. //! //! Work is `n(n - 1) / 2` distance evaluations with constant bounded insertion; //! scratch is `O(n)` and output is `O(nk)`. -use std::marker::PhantomData; - use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Dispatched1, FTarget1}, - lifetime::AddLifetime, + arch::{self, Target1}, }; use super::kernel_metric::{KernelMetric, MetricVisitor, visit_metric}; @@ -193,11 +190,7 @@ pub fn leaf_output_len(points: usize, requested_k: usize) -> Result { input: MatrixView<'a, f32>, @@ -205,140 +198,98 @@ struct LeafCall<'a> { workspace: &'a mut LeafKernelWorkspace, } -#[derive(Debug)] -struct LeafCallArg; - -impl AddLifetime for LeafCallArg { - type Of<'a> = LeafCall<'a>; -} - -type LeafFn = Dispatched1, LeafCallArg>; - -/// A leaf kernel prepared for one metric and the current CPU. -/// -/// Construct this once with [`LeafKernel::new`] and share it across leaf workers. -/// Each output view carries its leaf-specific neighbor width. +/// Leaf-kernel convenience API for callers with a runtime [`Metric`]. /// -/// The handle stores only one direct function pointer. It borrows no leaf data -/// or workspace and is therefore `Copy`, `Send`, and `Sync`. +/// Each call performs runtime architecture and metric selection, then invokes +/// the same generic kernel used by the production stage-level dispatch. #[derive(Clone, Copy, Debug)] pub struct LeafKernel { - run: LeafFn, + metric: Metric, } 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) + /// Construct a kernel selector for `metric`. + pub const fn new(metric: Metric) -> Self { + Self { metric } } /// Select the nearest non-self leaf positions for every source point. /// /// `output` must have one row per input point. Its column count is the /// neighbor count for this leaf and must not exceed either `point_count - 1` - /// or [`MAX_LEAF_NEIGHBORS`]. - /// 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. + /// or [`MAX_LEAF_NEIGHBORS`]. Equal distances retain pair scan order. /// /// # Errors /// - /// Returns [`LeafKernelError`] for invalid or overflowing shapes, excessive + /// Returns [`LeafKernelError`] for incompatible 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. + /// source caused by non-rankable distances. pub fn nearest_neighbors( &self, input: MatrixView<'_, f32>, output: MutMatrixView<'_, LeafNeighbor>, workspace: &mut LeafKernelWorkspace, ) -> Result<(), LeafKernelError> { - self.run.call(LeafCall { - input, - output, - workspace, - }) + arch::dispatch1_no_features( + RunLeaf { + metric: self.metric, + }, + LeafCall { + input, + output, + workspace, + }, + ) } } -/// 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; +struct RunLeaf { + metric: Metric, +} -impl arch::Target1 for PrepareLeaf +impl Target1, LeafCall<'_>> for RunLeaf 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 { - visit_metric(metric, BuildLeaf(arch)) + fn run(self, arch: A, call: LeafCall<'_>) -> Result<(), LeafKernelError> { + visit_metric(self.metric, ExecuteLeaf { arch, call }) } } -/// Metric visitor holding a concrete architecture. -/// -/// `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); +struct ExecuteLeaf<'a, A> { + arch: A, + call: LeafCall<'a>, +} -impl MetricVisitor for BuildLeaf +impl MetricVisitor for ExecuteLeaf<'_, A> where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { - type Output = LeafKernel; + type Output = Result<(), LeafKernelError>; fn visit(self) -> Self::Output { - LeafKernel { - run: self - .0 - .dispatch1::, Result<(), LeafKernelError>, LeafCallArg>(), - } + nearest_neighbors_for::( + self.arch, + self.call.input, + self.call.output, + self.call.workspace, + ) } } -/// Architecture/metric-specialized function-pointer destination. -/// -/// 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 +/// Architecture/metric-specialized leaf kernel used by stage-level dispatch. +pub(crate) fn nearest_neighbors_for( + arch: A, + input: MatrixView<'_, f32>, + mut output: MutMatrixView<'_, LeafNeighbor>, + workspace: &mut LeafKernelWorkspace, +) -> Result<(), LeafKernelError> where A: Architecture, A::f32x16: std::ops::Div, @@ -346,48 +297,35 @@ where M: KernelMetric, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { - 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. - 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(()); - } + validate(input, &output)?; + let neighbor_count = output.ncols(); + if neighbor_count == 0 { + return Ok(()); + } - // 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); - - // 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, - neighbor_count, - call.output.as_mut_slice(), - &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() - .chunks_exact(neighbor_count) - .position(|neighbors| neighbors[neighbor_count - 1].target == u32::MAX) - { - return Err(LeafKernelError::InsufficientRankableNeighbors { - source_index: source, - neighbors: neighbor_count, - }); - } - Ok(()) + prepare_workspace::(input, workspace)?; + output.as_mut_slice().fill(LeafNeighbor::default()); + workspace.worst.fill(f32::INFINITY); + + process_neighbor_width::( + arch, + input, + neighbor_count, + output.as_mut_slice(), + &workspace.norms, + &mut workspace.worst, + )?; + if let Some(source) = output + .as_slice() + .chunks_exact(neighbor_count) + .position(|neighbors| neighbors[neighbor_count - 1].target == u32::MAX) + { + return Err(LeafKernelError::InsufficientRankableNeighbors { + source_index: source, + neighbors: neighbor_count, + }); } + Ok(()) } /// Validate the complete safety contract before dispatched SIMD executes. diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index ffe4182353..ec738afd91 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -7,39 +7,50 @@ //! //! Caller supplies a row-major point-by-leader dot matrix plus metric-specific //! [`PartitionScales`]. Output contains sorted leader-column positions for each -//! point; fanout is the output width and cannot exceed -//! [`MAX_PARTITION_FANOUT`] or the leader count. +//! point; fanout is the output width and cannot exceed the leader count. //! -//! [`PartitionKernel::new`] selects metric and runtime architecture once and -//! stores a direct function pointer reused by every point stripe. Calls validate -//! row counts, scale variants and lengths, fanout, and leader-ID representation -//! before output mutation or unchecked SIMD loads. +//! Runtime callers may use [`PartitionKernel`]; production dispatches once at +//! the partition-stage boundary and calls the generic kernel directly. Calls +//! validate row counts, scale variants and lengths, fanout, and leader-ID +//! representation before output mutation or unchecked SIMD loads. //! //! L2 omits the point norm because it cannot change one point's leader order. //! Strict comparisons preserve scan order for ties and leave NaN non-rankable. -//! Each point evaluates every leader; competitive scores move through a fixed -//! stack tracker of at most [`MAX_PARTITION_FANOUT`] entries. - -use std::marker::PhantomData; +//! Each point evaluates every leader; competitive scores move through a +//! caller-owned tracker reused across points. use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, - arch::{self, Dispatched2, FTarget2}, - lifetime::AddLifetime, + arch::{self, Target1}, }; use super::kernel_metric::{KernelMetric, MetricVisitor, ScaleKind, visit_metric}; -/// Maximum number of leaders retained for one point. -/// -/// Supported PiPNN partition fanouts fit within 16. Keeping this as a fixed -/// 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; +/// Reusable nearest-leader tracker for one partition worker. +#[derive(Debug, Default)] +pub struct PartitionKernelWorkspace { + tracker: Vec<(u32, f32)>, +} + +impl PartitionKernelWorkspace { + /// Construct an empty allocation-free workspace. + pub const fn new() -> Self { + Self { + tracker: Vec::new(), + } + } -type LeaderTracker = [(u32, f32); MAX_PARTITION_FANOUT]; + fn prepare(&mut self, fanout: usize) -> Result<(), PartitionKernelError> { + let additional = fanout.saturating_sub(self.tracker.len()); + self.tracker + .try_reserve(additional) + .map_err(|_| PartitionKernelError::Allocation { additional })?; + self.tracker.resize(fanout, (u32::MAX, f32::INFINITY)); + Ok(()) + } +} /// Metric-specific normalization inputs for one partition tile. /// @@ -109,17 +120,19 @@ pub enum PartitionKernelError { /// 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 {leader_count} leaders or kernel maximum {maximum}" - )] + /// The requested fanout exceeds the available leader count. + #[error("invalid fanout {fanout}: must not exceed {leader_count} leaders")] InvalidFanout { /// Requested number of leaders per point. fanout: usize, /// Available leader count. leader_count: usize, - /// Kernel maximum. - maximum: usize, + }, + /// Reusable tracker storage could not be reserved. + #[error("failed to reserve {additional} partition tracker entries")] + Allocation { + /// Additional entries requested from the allocator. + additional: usize, }, /// Leader positions cannot be represented as `u32`. #[error("leader count {0} exceeds the u32 position limit")] @@ -134,147 +147,96 @@ 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 PartitionInputArg; - -impl AddLifetime for PartitionInputArg { - type Of<'a> = PartitionInput<'a>; -} - +/// Inputs for one immediate architecture/metric dispatch. #[derive(Debug)] -struct PartitionOutput; - -impl AddLifetime for PartitionOutput { - type Of<'a> = MutMatrixView<'a, u32>; +struct PartitionCall<'a> { + input: PartitionInput<'a>, + output: MutMatrixView<'a, u32>, + workspace: &'a mut PartitionKernelWorkspace, } -type PartitionFn = - Dispatched2, PartitionInputArg, PartitionOutput>; - -/// A partition kernel prepared for one metric and the current CPU. -/// -/// 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. +/// Partition-kernel convenience API for callers with a runtime [`Metric`]. #[derive(Clone, Copy, Debug)] pub struct PartitionKernel { - run: PartitionFn, + metric: Metric, } 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) + /// Construct a kernel selector for `metric`. + pub const fn new(metric: Metric) -> Self { + Self { metric } } - /// 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 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 + /// Select nearest leader positions for every input point. /// - /// See module-level complexity. This call follows one prepared direct - /// function pointer and performs no runtime ISA or metric dispatch. + /// `output.nrows()` must equal `input.dots.nrows()` and fanout, represented + /// by `output.ncols()`, must not exceed the leader count. pub fn nearest_leaders( &self, input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, + workspace: &mut PartitionKernelWorkspace, ) -> Result<(), PartitionKernelError> { - self.run.call(input, output) + arch::dispatch1_no_features( + RunPartition { + metric: self.metric, + }, + PartitionCall { + input, + output, + workspace, + }, + ) } } -/// 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; +struct RunPartition { + metric: Metric, +} -impl arch::Target1 for PreparePartition +impl Target1, PartitionCall<'_>> for RunPartition 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 { - visit_metric(metric, BuildPartition(arch)) + fn run(self, arch: A, call: PartitionCall<'_>) -> Result<(), PartitionKernelError> { + visit_metric(self.metric, ExecutePartition { arch, call }) } } -/// BYO-type-erasure visitor holding a concrete architecture. -/// -/// `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); +struct ExecutePartition<'a, A> { + arch: A, + call: PartitionCall<'a>, +} -impl MetricVisitor for BuildPartition +impl MetricVisitor for ExecutePartition<'_, A> where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { - type Output = PartitionKernel; + type Output = Result<(), PartitionKernelError>; fn visit(self) -> Self::Output { - PartitionKernel { - run: self.0.dispatch2::< - PartitionEntry, - Result<(), PartitionKernelError>, - PartitionInputArg, - PartitionOutput, - >(), - } + nearest_leaders_for::( + self.arch, + self.call.input, + self.call.output, + self.call.workspace, + ) } } -/// Architecture/metric-specialized function-pointer destination. -/// -/// 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>> - for PartitionEntry +/// Architecture/metric-specialized partition kernel used by stage dispatch. +pub(crate) fn nearest_leaders_for( + arch: A, + input: PartitionInput<'_>, + mut output: MutMatrixView<'_, u32>, + workspace: &mut PartitionKernelWorkspace, +) -> Result<(), PartitionKernelError> where A: Architecture, A::f32x16: std::ops::Div, @@ -282,35 +244,28 @@ where u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, M: KernelMetric, { - fn run( - arch: A, - input: PartitionInput<'_>, - 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(); - // Zero fanout and empty stripes require no assignments. Return before - // constructing trackers or touching output. - if fanout == 0 || input.dots.nrows() == 0 { - return Ok(()); - } + let scales = validate::(input, &output)?; + let fanout = output.ncols(); + 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. - if let Some(point) = output - .as_slice() - .chunks_exact(fanout) - .position(|assignments| assignments[fanout - 1] == u32::MAX) - { - return Err(PartitionKernelError::InsufficientRankableLeaders { point, fanout }); - } - Ok(()) + workspace.prepare(fanout)?; + process_points::( + arch, + input.dots, + scales, + output.as_mut_slice(), + &mut workspace.tracker, + ); + if let Some(point) = output + .as_slice() + .chunks_exact(fanout) + .position(|assignments| assignments[fanout - 1] == u32::MAX) + { + return Err(PartitionKernelError::InsufficientRankableLeaders { point, fanout }); } + Ok(()) } /// Validated scale slices in the storage form required by `M`. @@ -348,11 +303,10 @@ fn validate<'a, M: KernelMetric>( if leader_count > u32::MAX as usize { return Err(PartitionKernelError::TooManyLeaders(leader_count)); } - if fanout > MAX_PARTITION_FANOUT || fanout > leader_count { + if fanout > leader_count { return Err(PartitionKernelError::InvalidFanout { fanout, leader_count, - maximum: MAX_PARTITION_FANOUT, }); } @@ -460,8 +414,8 @@ fn process_points( arch: F::Arch, dots: MatrixView<'_, f32>, scales: ScaleSlices<'_>, - fanout: usize, output: &mut [u32], + tracker: &mut [(u32, f32)], ) where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, @@ -488,14 +442,15 @@ fn process_points( "validated leader scales must match leader count" ); } - // Each point is independent. Reinitialize the fixed tracker here so no - // assignment state or tie order leaks across points. + let fanout = tracker.len(); + // Each point is independent. Reset the caller-owned tracker so no assignment + // state or tie order leaks across rows. for (point, (point_dots, point_output)) in dots - .as_slice() - .chunks_exact(leader_count) + .row_iter() .zip(output.chunks_exact_mut(fanout)) .enumerate() { + tracker.fill((u32::MAX, f32::INFINITY)); // 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() { @@ -504,7 +459,6 @@ fn process_points( 0.0 }; 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; @@ -522,8 +476,7 @@ fn process_points( insert_leader_lanes( M::partition_distance(arch, point_dots, point_scale_vector, leader_scales), base, - &mut tracker, - fanout, + tracker, ); } @@ -536,15 +489,14 @@ fn process_points( 0.0 }; insert_leader( - &mut tracker, - fanout, + tracker, leader as u32, 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); + copy_leader_ids(tracker, point_output); } } @@ -558,16 +510,12 @@ fn process_points( /// `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, - tracker: &mut LeaderTracker, - fanout: usize, -) where +fn insert_leader_lanes(distances: F, first_leader: usize, tracker: &mut [(u32, f32)]) +where F: SIMDVector> + SIMDPartialOrd, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let threshold = F::splat(distances.arch(), tracker[fanout - 1].1); + let threshold = F::splat(distances.arch(), tracker[tracker.len() - 1].1); let eligible = distances.lt_simd(threshold); if eligible.none() { return; @@ -578,7 +526,7 @@ fn insert_leader_lanes( while lanes != 0 { let lane = lanes.trailing_zeros() as usize; lanes &= lanes - 1; - insert_leader(tracker, fanout, (first_leader + lane) as u32, values[lane]); + insert_leader(tracker, (first_leader + lane) as u32, values[lane]); } } @@ -592,8 +540,8 @@ fn insert_leader_lanes( /// `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; +fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, distance: f32) { + let threshold = tracker.len() - 1; if distance.partial_cmp(&tracker[threshold].1) != Some(std::cmp::Ordering::Less) { return; } @@ -610,7 +558,7 @@ fn insert_leader(tracker: &mut LeaderTracker, fanout: usize, leader: u32, distan /// /// `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]) { +fn copy_leader_ids(tracker: &[(u32, f32)], assignments: &mut [u32]) { for (destination, &(leader, _)) in assignments.iter_mut().zip(tracker) { *destination = leader; } @@ -673,11 +621,9 @@ mod tests { leader_scales: &[], }, }; - let leader_count = input.dots.ncols(); for (point, (point_dots, point_output)) in input .dots - .as_slice() - .chunks_exact(leader_count) + .row_iter() .zip(output.chunks_exact_mut(fanout)) .enumerate() { @@ -686,7 +632,7 @@ mod tests { } else { 0.0 }; - let mut tracker = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + let mut tracker = vec![(u32::MAX, f32::INFINITY); 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.leader_scales[leader]) @@ -695,7 +641,6 @@ mod tests { }; insert_leader( &mut tracker, - fanout, leader as u32, M::partition_distance_scalar(dot, point_scale, leader_scale), ); @@ -744,6 +689,7 @@ mod tests { .nearest_leaders( input, MutMatrixView::try_from(actual.as_mut_slice(), point_scales.len(), 2).unwrap(), + &mut PartitionKernelWorkspace::new(), ) .unwrap(); @@ -754,13 +700,13 @@ mod tests { #[test] fn scalar_topk_orders_candidates_and_preserves_ties() { - let mut tracker = [(u32::MAX, f32::INFINITY); MAX_PARTITION_FANOUT]; + let mut tracker = vec![(u32::MAX, f32::INFINITY); 4]; for (leader, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 1.0)] { - insert_leader(&mut tracker, 4, leader, distance); + insert_leader(&mut tracker, leader, distance); } - insert_leader(&mut tracker, 4, 5, f32::NAN); + insert_leader(&mut tracker, 5, f32::NAN); - assert_eq!(tracker[..4], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); + assert_eq!(tracker[..], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); } } #[cfg(test)] @@ -771,7 +717,7 @@ mod tests { )] mod integration_tests { use super::{ - MAX_PARTITION_FANOUT, PartitionInput, PartitionKernel, PartitionKernelError, + PartitionInput, PartitionKernel, PartitionKernelError, PartitionKernelWorkspace, PartitionScales, }; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -913,6 +859,7 @@ mod integration_tests { PartitionKernel::new(metric).nearest_leaders( input, MutMatrixView::try_from(output.as_mut_slice(), input.dots.nrows(), fanout).unwrap(), + &mut PartitionKernelWorkspace::new(), )?; Ok(output) } @@ -1104,6 +1051,7 @@ mod integration_tests { PartitionKernel::new(Metric::InnerProduct).nearest_leaders( valid_input, MutMatrixView::try_from(&mut wrong_output[..], 1, 3).unwrap(), + &mut PartitionKernelWorkspace::new(), ), Err(PartitionKernelError::InvalidOutputShape { expected_rows: 2, @@ -1122,11 +1070,10 @@ mod integration_tests { ); assert_eq!( - run(Metric::InnerProduct, valid_input, MAX_PARTITION_FANOUT + 1,), + run(Metric::InnerProduct, valid_input, 4), Err(PartitionKernelError::InvalidFanout { - fanout: MAX_PARTITION_FANOUT + 1, + fanout: 4, leader_count: 3, - maximum: MAX_PARTITION_FANOUT, }) ); @@ -1140,7 +1087,6 @@ mod integration_tests { Err(PartitionKernelError::InvalidFanout { fanout: 2, leader_count: 1, - maximum: MAX_PARTITION_FANOUT, }) ); } From c3bf708b3f7ec7226f943b7091ba21694bc45ef8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:50:27 +0000 Subject: [PATCH 41/80] test(pipnn): cover runtime partition fanout --- diskann/src/graph/pipnn/partition_kernel.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index ec738afd91..0a65e4a08a 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -708,6 +708,18 @@ mod tests { assert_eq!(tracker[..], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); } + + #[test] + fn workspace_reuses_runtime_fanout_capacity() { + let mut workspace = PartitionKernelWorkspace::new(); + workspace.prepare(32).unwrap(); + let allocation = workspace.tracker.as_ptr(); + + workspace.prepare(3).unwrap(); + + assert_eq!(workspace.tracker.as_ptr(), allocation); + assert_eq!(workspace.tracker.len(), 3); + } } #[cfg(test)] #[allow( @@ -882,7 +894,7 @@ mod integration_tests { &point_scales, &leader_scales, ); - for fanout in [1, 2, 16] { + for fanout in [1, 2, 16, 17, 32] { if fanout >= leader_count { continue; } From 0a16a6225cc728e23bf96642cea6784e7a4cdc89 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:00:34 +0000 Subject: [PATCH 42/80] refactor(pipnn): make kernels generic-only --- diskann/src/graph/pipnn/kernel_metric.rs | 51 +--- diskann/src/graph/pipnn/leaf_kernel.rs | 262 +++++++++----------- diskann/src/graph/pipnn/mod.rs | 145 +++-------- diskann/src/graph/pipnn/partition_kernel.rs | 202 +++++++-------- 4 files changed, 247 insertions(+), 413 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index a6b0c3a869..3859467ca5 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -5,9 +5,9 @@ //! Metric formulas shared by PiPNN partition and leaf kernels. //! -//! Runtime [`Metric`] selection happens once while preparing a dispatched -//! kernel. Zero-sized marker types then monomorphize scalar and SIMD formulas, -//! so hot loops contain neither metric matches nor trait objects. +//! The build boundary converts runtime [`Metric`] into one zero-sized marker +//! type. That concrete type is carried through partition and leaf construction, +//! so scalar and SIMD hot loops contain neither metric matches nor trait objects. //! //! All formulas produce ascending scores. L2 uses squared norms; unnormalized //! cosine uses norms with zero/subnormal inputs mapped to zero similarity; @@ -15,9 +15,9 @@ //! NaN non-rankable. The L2 partition scalar tail deliberately keeps its //! non-fused operation order because rounding can change leader ties. //! -//! [`ScaleKind`] records required scale representation. [`KernelMetric`] owns -//! leaf and partition formulas. [`MetricVisitor`] performs runtime-to-marker -//! conversion before the final architecture-specific function pointer is stored. +//! [`ScaleKind`] records required scale representation and [`KernelMetric`] +//! owns the leaf and partition formulas. Runtime selection belongs to the build +//! entry point; this module contains no dispatch or type erasure. use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; @@ -82,10 +82,10 @@ 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 point norm. +/// Runtime `Metric` is converted to one implementor at the build boundary. +/// Generic methods then inline metric arithmetic through the complete partition +/// and leaf stages. Those 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 @@ -372,37 +372,6 @@ 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. -/// 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::(), - Metric::Cosine => visitor.visit::(), - Metric::CosineNormalized => visitor.visit::(), - Metric::InnerProduct => visitor.visit::(), - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 209bf4537e..28fffabc31 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Prepared leaf-local top-k selection over a lower-triangular Gram matrix. +//! Leaf-local top-k selection over a lower-triangular Gram matrix. //! //! Caller supplies an `n × n` [`MatrixView`] produced by `sgemm_aat_lower`. //! Diagonal entries provide metric scales; only strict-lower pair dots are read. @@ -15,8 +15,8 @@ //! reject NaN. L2, cosine, normalized cosine, and inner product share the same //! scalar/SIMD traversal. //! -//! Runtime callers may use [`LeafKernel`]; production dispatches once at the leaf -//! stage boundary and calls the generic kernel directly. Every call validates +//! Callers select architecture and metric once outside the leaf loop, then call +//! [`nearest_neighbors`] with concrete `A` and `M` types. Every call validates //! square shape, row count, local-ID bounds, and output width before scratch //! mutation or unchecked SIMD loads. [`LeafKernelWorkspace`] retains per-worker //! norm and threshold buffers. @@ -25,13 +25,9 @@ //! scratch is `O(n)` and output is `O(nk)`. use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_vector::distance::Metric; -use diskann_wide::{ - Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, -}; +use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; -use super::kernel_metric::{KernelMetric, MetricVisitor, visit_metric}; +use super::kernel_metric::KernelMetric; /// Largest leaf-local neighbor count supported by the fixed insertion kernel. pub const MAX_LEAF_NEIGHBORS: usize = 3; @@ -81,7 +77,7 @@ impl LeafKernelWorkspace { } } -/// Validation or allocation error returned by [`LeafKernel::nearest_neighbors`]. +/// Validation or allocation error returned by [`nearest_neighbors`]. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub enum LeafKernelError { /// The point count cannot be represented in leaf-local `u32` positions. @@ -173,7 +169,7 @@ pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result Result { - input: MatrixView<'a, f32>, - output: MutMatrixView<'a, LeafNeighbor>, - workspace: &'a mut LeafKernelWorkspace, -} - -/// Leaf-kernel convenience API for callers with a runtime [`Metric`]. +/// Select the nearest non-self leaf positions for every source point. /// -/// Each call performs runtime architecture and metric selection, then invokes -/// the same generic kernel used by the production stage-level dispatch. -#[derive(Clone, Copy, Debug)] -pub struct LeafKernel { - metric: Metric, -} - -impl LeafKernel { - /// Construct a kernel selector for `metric`. - pub const fn new(metric: Metric) -> Self { - Self { metric } - } - - /// Select the nearest non-self leaf positions for every source point. - /// - /// `output` must have one row per input point. Its column count is the - /// neighbor count for this leaf and must not exceed either `point_count - 1` - /// or [`MAX_LEAF_NEIGHBORS`]. Equal distances retain pair scan order. - /// - /// # Errors - /// - /// Returns [`LeafKernelError`] for incompatible shapes, excessive - /// point/neighbor counts, scratch allocation failure, or an underfilled - /// source caused by non-rankable distances. - pub fn nearest_neighbors( - &self, - input: MatrixView<'_, f32>, - output: MutMatrixView<'_, LeafNeighbor>, - workspace: &mut LeafKernelWorkspace, - ) -> Result<(), LeafKernelError> { - arch::dispatch1_no_features( - RunLeaf { - metric: self.metric, - }, - LeafCall { - input, - output, - workspace, - }, - ) - } -} - -struct RunLeaf { - metric: Metric, -} - -impl Target1, LeafCall<'_>> for RunLeaf -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, -{ - fn run(self, arch: A, call: LeafCall<'_>) -> Result<(), LeafKernelError> { - visit_metric(self.metric, ExecuteLeaf { arch, call }) - } -} - -struct ExecuteLeaf<'a, A> { - arch: A, - call: LeafCall<'a>, -} - -impl MetricVisitor for ExecuteLeaf<'_, A> -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, -{ - type Output = Result<(), LeafKernelError>; - - fn visit(self) -> Self::Output { - nearest_neighbors_for::( - self.arch, - self.call.input, - self.call.output, - self.call.workspace, - ) - } -} - -/// Architecture/metric-specialized leaf kernel used by stage-level dispatch. -pub(crate) fn nearest_neighbors_for( +/// `output` must have one row per input point. Its column count is the neighbor +/// count for this leaf and must not exceed either `point_count - 1` or +/// [`MAX_LEAF_NEIGHBORS`]. Equal distances retain pair scan order. `A` and `M` +/// must already have been selected at the enclosing build boundary. +/// +/// # Errors +/// +/// Returns [`LeafKernelError`] for incompatible shapes, excessive +/// point/neighbor counts, scratch allocation failure, or an underfilled source +/// caused by non-rankable distances. +pub(crate) fn nearest_neighbors( arch: A, input: MatrixView<'_, f32>, mut output: MutMatrixView<'_, LeafNeighbor>, @@ -687,9 +601,68 @@ fn insert_fixed_neighbor( } } +#[cfg(test)] +struct DispatchedLeafCall<'a> { + input: MatrixView<'a, f32>, + output: MutMatrixView<'a, LeafNeighbor>, + workspace: &'a mut LeafKernelWorkspace, +} + +#[cfg(test)] +struct DispatchLeafForTest(diskann_vector::distance::Metric); + +#[cfg(test)] +impl diskann_wide::arch::Target1, DispatchedLeafCall<'_>> + for DispatchLeafForTest +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + fn run(self, arch: A, call: DispatchedLeafCall<'_>) -> Result<(), LeafKernelError> { + use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use diskann_vector::distance::Metric; + + match self.0 { + Metric::L2 => nearest_neighbors::(arch, call.input, call.output, call.workspace), + Metric::Cosine => { + nearest_neighbors::(arch, call.input, call.output, call.workspace) + } + Metric::CosineNormalized => nearest_neighbors::( + arch, + call.input, + call.output, + call.workspace, + ), + Metric::InnerProduct => { + nearest_neighbors::(arch, call.input, call.output, call.workspace) + } + } + } +} + +#[cfg(test)] +fn dispatch_nearest_neighbors( + metric: diskann_vector::distance::Metric, + input: MatrixView<'_, f32>, + output: MutMatrixView<'_, LeafNeighbor>, + workspace: &mut LeafKernelWorkspace, +) -> Result<(), LeafKernelError> { + diskann_wide::arch::dispatch1_no_features( + DispatchLeafForTest(metric), + DispatchedLeafCall { + input, + output, + workspace, + }, + ) +} + #[cfg(test)] mod tests { use super::*; + use diskann_vector::distance::Metric; fn test_dots(metric: Metric, points: usize) -> Vec { let mut dots = vec![f32::NAN; points * points]; @@ -766,40 +739,38 @@ mod tests { } #[test] - fn prepared_kernel_accepts_different_neighbor_counts() { + fn 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(); + dispatch_nearest_neighbors( + Metric::L2, + 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); let mut workspace = LeafKernelWorkspace::new(); for points in [17, 7, 17] { let dots = test_dots(Metric::L2, points); let mut output = vec![LeafNeighbor::default(); points * 2]; - kernel - .nearest_neighbors( - test_input(&dots, points), - MutMatrixView::try_from(output.as_mut_slice(), points, 2).unwrap(), - &mut workspace, - ) - .unwrap(); + dispatch_nearest_neighbors( + Metric::L2, + test_input(&dots, points), + MutMatrixView::try_from(output.as_mut_slice(), points, 2).unwrap(), + &mut workspace, + ) + .unwrap(); assert!(output.iter().all(|neighbor| neighbor.target != u32::MAX)); } } @@ -814,8 +785,8 @@ mod integration_tests { use std::cmp::Ordering; use super::{ - LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, MAX_LEAF_NEIGHBORS, - leaf_neighbor_count, leaf_output_len, + LeafKernelError, LeafKernelWorkspace, LeafNeighbor, MAX_LEAF_NEIGHBORS, + dispatch_nearest_neighbors, leaf_neighbor_count, leaf_output_len, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -937,18 +908,18 @@ mod integration_tests { ) -> (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( - test_input(dots, points), - MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), - &mut LeafKernelWorkspace::new(), - ) - .unwrap(); + dispatch_nearest_neighbors( + metric, + test_input(dots, points), + MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), + &mut LeafKernelWorkspace::new(), + ) + .unwrap(); (leaf_k, output) } #[test] - fn prepared_dispatch_matches_reference_across_simd_width_boundaries() { + fn dispatched_kernel_matches_reference_across_simd_width_boundaries() { for metric in [ Metric::L2, Metric::Cosine, @@ -1104,13 +1075,13 @@ mod integration_tests { 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) - .nearest_neighbors( - test_input(&dots, 2), - MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::new(), - ) - .unwrap_err(); + let error = dispatch_nearest_neighbors( + Metric::L2, + test_input(&dots, 2), + MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), + &mut LeafKernelWorkspace::new(), + ) + .unwrap_err(); assert_eq!( error, @@ -1157,9 +1128,9 @@ mod integration_tests { let dots = [0.0; 6]; let non_square = MatrixView::try_from(&dots[..], 2, 3).unwrap(); let mut output = [LeafNeighbor::default(); 2]; - let kernel = LeafKernel::new(Metric::L2); assert_eq!( - kernel.nearest_neighbors( + dispatch_nearest_neighbors( + Metric::L2, non_square, MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), &mut LeafKernelWorkspace::new(), @@ -1170,7 +1141,8 @@ mod integration_tests { let square = [0.0; 9]; let mut wrong_rows = [LeafNeighbor::default(); 2]; assert_eq!( - kernel.nearest_neighbors( + dispatch_nearest_neighbors( + Metric::L2, test_input(&square, 3), MutMatrixView::try_from(&mut wrong_rows[..], 2, 1).unwrap(), &mut LeafKernelWorkspace::new(), @@ -1184,7 +1156,8 @@ mod integration_tests { let mut too_many = [LeafNeighbor::default(); 9]; assert_eq!( - kernel.nearest_neighbors( + dispatch_nearest_neighbors( + Metric::L2, test_input(&square, 3), MutMatrixView::try_from(&mut too_many[..], 3, 3).unwrap(), &mut LeafKernelWorkspace::new(), @@ -1199,7 +1172,8 @@ mod integration_tests { let square = [0.0; 25]; let mut too_wide = [LeafNeighbor::default(); 20]; assert_eq!( - kernel.nearest_neighbors( + dispatch_nearest_neighbors( + Metric::L2, test_input(&square, 5), MutMatrixView::try_from(&mut too_wide[..], 5, 4).unwrap(), &mut LeafKernelWorkspace::new(), diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index b1aa678676..876077175e 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,125 +5,36 @@ //! Numerical kernels for provider-independent PiPNN graph construction. //! -//! 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 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 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 -//! │ -//! 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 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, -//! *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: -//! -//! - [`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. -//! -//! # 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::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 fixed-width selection, `process_pairs`, 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. +//! PiPNN assigns points to overlapping leader partitions, computes one +//! lower-triangular all-pairs matrix per bounded leaf, and merges the selected +//! leaf neighbors into graph candidates. This layer contains only the two score +//! selection kernels; later layers own partition recursion, GEMM, graph IDs, +//! candidate merging, final pruning, providers, and persistence. +//! +//! - [`partition_kernel`] converts a point-by-leader dot-product tile into sorted +//! leader-column positions. Output width is runtime fanout; reusable tracker +//! storage grows to that width and is reused across point rows. +//! - [`leaf_kernel`] scans each strict-lower-triangle pair once, updates both +//! endpoints, and retains up to three leaf-local neighbors per point. +//! - `kernel_metric` owns the scalar/SIMD formulas and exact scale units shared +//! by both kernels. +//! +//! Runtime architecture and metric selection intentionally do not live in the +//! kernels. The enclosing graph build selects concrete `A` and `M` types once, +//! then carries them through partition and leaf work. This keeps metric matches, +//! trait objects, and stored function pointers out of the hot loops. +//! +//! Both kernels validate view relationships and metric scale layouts before +//! unchecked SIMD access. They borrow all matrices and mutate only caller-owned +//! output and reusable workspace. Partition work is one score per point-leader +//! pair; leaf work is one score per unordered point pair. +// This stack layer introduces the numerical kernels before the following core +// layer wires them into graph construction. +#[allow(dead_code)] mod kernel_metric; +#[allow(dead_code)] pub mod leaf_kernel; +#[allow(dead_code)] pub mod partition_kernel; diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 0a65e4a08a..7beeb443b5 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -3,16 +3,16 @@ * Licensed under the MIT license. */ -//! Prepared nearest-leader selection for PiPNN partition assignment. +//! Nearest-leader selection for PiPNN partition assignment. //! //! Caller supplies a row-major point-by-leader dot matrix plus metric-specific //! [`PartitionScales`]. Output contains sorted leader-column positions for each //! point; fanout is the output width and cannot exceed the leader count. //! -//! Runtime callers may use [`PartitionKernel`]; production dispatches once at -//! the partition-stage boundary and calls the generic kernel directly. Calls -//! validate row counts, scale variants and lengths, fanout, and leader-ID -//! representation before output mutation or unchecked SIMD loads. +//! Callers select architecture and metric once outside partition recursion, +//! then call [`nearest_leaders`] with concrete `A` and `M` types. Calls validate +//! row counts, scale variants and lengths, fanout, and leader-ID representation +//! before output mutation or unchecked SIMD loads. //! //! L2 omits the point norm because it cannot change one point's leader order. //! Strict comparisons preserve scan order for ties and leave NaN non-rankable. @@ -23,10 +23,9 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, - arch::{self, Target1}, }; -use super::kernel_metric::{KernelMetric, MetricVisitor, ScaleKind, visit_metric}; +use super::kernel_metric::{KernelMetric, ScaleKind}; /// Reusable nearest-leader tracker for one partition worker. #[derive(Debug, Default)] @@ -79,17 +78,17 @@ 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. +/// concrete metric selected by the enclosing build. This value only borrows +/// input; the numerical kernel stores no tile state. #[derive(Clone, Copy, Debug)] 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. + /// Normalization inputs matching concrete metric `M`. pub scales: PartitionScales<'a>, } -/// Validation error returned by [`PartitionKernel::nearest_leaders`]. +/// Validation error returned by [`nearest_leaders`]. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub enum PartitionKernelError { /// The output matrix does not match the input row count. @@ -114,8 +113,8 @@ pub enum PartitionKernelError { /// 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")] + /// Scale inputs do not match concrete metric `M`. + #[error("partition scales do not match selected {expected} metric")] InvalidScales { /// Expected scale layout. expected: &'static str, @@ -147,91 +146,12 @@ pub enum PartitionKernelError { }, } -/// Inputs for one immediate architecture/metric dispatch. -#[derive(Debug)] -struct PartitionCall<'a> { - input: PartitionInput<'a>, - output: MutMatrixView<'a, u32>, - workspace: &'a mut PartitionKernelWorkspace, -} - -/// Partition-kernel convenience API for callers with a runtime [`Metric`]. -#[derive(Clone, Copy, Debug)] -pub struct PartitionKernel { - metric: Metric, -} - -impl PartitionKernel { - /// Construct a kernel selector for `metric`. - pub const fn new(metric: Metric) -> Self { - Self { metric } - } - - /// Select nearest leader positions for every input point. - /// - /// `output.nrows()` must equal `input.dots.nrows()` and fanout, represented - /// by `output.ncols()`, must not exceed the leader count. - pub fn nearest_leaders( - &self, - input: PartitionInput<'_>, - output: MutMatrixView<'_, u32>, - workspace: &mut PartitionKernelWorkspace, - ) -> Result<(), PartitionKernelError> { - arch::dispatch1_no_features( - RunPartition { - metric: self.metric, - }, - PartitionCall { - input, - output, - workspace, - }, - ) - } -} - -struct RunPartition { - metric: Metric, -} - -impl Target1, PartitionCall<'_>> for RunPartition -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, -{ - fn run(self, arch: A, call: PartitionCall<'_>) -> Result<(), PartitionKernelError> { - visit_metric(self.metric, ExecutePartition { arch, call }) - } -} - -struct ExecutePartition<'a, A> { - arch: A, - call: PartitionCall<'a>, -} - -impl MetricVisitor for ExecutePartition<'_, A> -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, -{ - type Output = Result<(), PartitionKernelError>; - - fn visit(self) -> Self::Output { - nearest_leaders_for::( - self.arch, - self.call.input, - self.call.output, - self.call.workspace, - ) - } -} - -/// Architecture/metric-specialized partition kernel used by stage dispatch. -pub(crate) fn nearest_leaders_for( +/// Select nearest leader positions for every input point. +/// +/// `output.nrows()` must equal `input.dots.nrows()` and fanout, represented by +/// `output.ncols()`, must not exceed the leader count. `A` and `M` must already +/// have been selected at the enclosing build boundary. +pub(crate) fn nearest_leaders( arch: A, input: PartitionInput<'_>, mut output: MutMatrixView<'_, u32>, @@ -564,6 +484,64 @@ fn copy_leader_ids(tracker: &[(u32, f32)], assignments: &mut [u32]) { } } +#[cfg(test)] +struct DispatchedPartitionCall<'a> { + input: PartitionInput<'a>, + output: MutMatrixView<'a, u32>, + workspace: &'a mut PartitionKernelWorkspace, +} + +#[cfg(test)] +struct DispatchPartitionForTest(Metric); + +#[cfg(test)] +impl + diskann_wide::arch::Target1, DispatchedPartitionCall<'_>> + for DispatchPartitionForTest +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + fn run(self, arch: A, call: DispatchedPartitionCall<'_>) -> Result<(), PartitionKernelError> { + use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + + match self.0 { + Metric::L2 => nearest_leaders::(arch, call.input, call.output, call.workspace), + Metric::Cosine => { + nearest_leaders::(arch, call.input, call.output, call.workspace) + } + Metric::CosineNormalized => nearest_leaders::( + arch, + call.input, + call.output, + call.workspace, + ), + Metric::InnerProduct => { + nearest_leaders::(arch, call.input, call.output, call.workspace) + } + } + } +} + +#[cfg(test)] +fn dispatch_nearest_leaders( + metric: Metric, + input: PartitionInput<'_>, + output: MutMatrixView<'_, u32>, + workspace: &mut PartitionKernelWorkspace, +) -> Result<(), PartitionKernelError> { + diskann_wide::arch::dispatch1_no_features( + DispatchPartitionForTest(metric), + DispatchedPartitionCall { + input, + output, + workspace, + }, + ) +} + #[cfg(test)] mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; @@ -663,7 +641,7 @@ mod tests { } #[test] - fn cosine_special_norms_match_scalar_and_prepared_dispatch() { + fn cosine_special_norms_match_scalar_and_dispatched_kernel() { 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]; @@ -685,13 +663,13 @@ mod tests { let mut expected = vec![u32::MAX; point_scales.len() * 2]; scalar_traversal_reference::(input, 2, &mut expected); let mut actual = vec![u32::MAX; point_scales.len() * 2]; - PartitionKernel::new(Metric::Cosine) - .nearest_leaders( - input, - MutMatrixView::try_from(actual.as_mut_slice(), point_scales.len(), 2).unwrap(), - &mut PartitionKernelWorkspace::new(), - ) - .unwrap(); + dispatch_nearest_leaders( + Metric::Cosine, + input, + MutMatrixView::try_from(actual.as_mut_slice(), point_scales.len(), 2).unwrap(), + &mut PartitionKernelWorkspace::new(), + ) + .unwrap(); assert_eq!(actual, expected); assert_eq!(&actual[..4], &[0, 1, 0, 1]); @@ -729,8 +707,8 @@ mod tests { )] mod integration_tests { use super::{ - PartitionInput, PartitionKernel, PartitionKernelError, PartitionKernelWorkspace, - PartitionScales, + PartitionInput, PartitionKernelError, PartitionKernelWorkspace, PartitionScales, + dispatch_nearest_leaders, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -868,7 +846,8 @@ mod integration_tests { fanout: usize, ) -> Result, PartitionKernelError> { let mut output = vec![u32::MAX; input.dots.nrows() * fanout]; - PartitionKernel::new(metric).nearest_leaders( + dispatch_nearest_leaders( + metric, input, MutMatrixView::try_from(output.as_mut_slice(), input.dots.nrows(), fanout).unwrap(), &mut PartitionKernelWorkspace::new(), @@ -877,7 +856,7 @@ mod integration_tests { } #[test] - fn prepared_dispatch_matches_reference_across_simd_width_boundaries() { + fn dispatched_kernel_matches_reference_across_simd_width_boundaries() { for metric in [ Metric::L2, Metric::Cosine, @@ -1060,7 +1039,8 @@ mod integration_tests { 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( + dispatch_nearest_leaders( + Metric::InnerProduct, valid_input, MutMatrixView::try_from(&mut wrong_output[..], 1, 3).unwrap(), &mut PartitionKernelWorkspace::new(), From cfdfc3c4ce592ba7f665fe0afdb564e54f6e0df4 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:23:58 +0000 Subject: [PATCH 43/80] refactor(pipnn): keep numerical kernels private --- diskann/src/graph/pipnn/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 876077175e..cb478ec113 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -35,6 +35,6 @@ mod kernel_metric; #[allow(dead_code)] -pub mod leaf_kernel; +mod leaf_kernel; #[allow(dead_code)] -pub mod partition_kernel; +mod partition_kernel; From 0e3e0555d24438bf4c33a46cefbb998f39149a7c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:30:30 +0000 Subject: [PATCH 44/80] refactor(pipnn): use default kernel workspaces --- diskann/src/graph/pipnn/leaf_kernel.rs | 29 ++++++--------------- diskann/src/graph/pipnn/partition_kernel.rs | 15 +++-------- 2 files changed, 12 insertions(+), 32 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 28fffabc31..229e918c3b 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -64,19 +64,6 @@ pub struct LeafKernelWorkspace { worst: Vec, } -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(), - worst: Vec::new(), - } - } -} - /// Validation or allocation error returned by [`nearest_neighbors`]. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub enum LeafKernelError { @@ -743,7 +730,7 @@ mod tests { let points = 7; let dots = test_dots(Metric::L2, points); let input = test_input(&dots, points); - let mut workspace = LeafKernelWorkspace::new(); + let mut workspace = LeafKernelWorkspace::default(); for neighbor_count in [1, 3, 2] { let mut output = vec![LeafNeighbor::default(); points * neighbor_count]; @@ -760,7 +747,7 @@ mod tests { #[test] fn workspace_can_shrink_and_grow_between_calls() { - let mut workspace = LeafKernelWorkspace::new(); + let mut workspace = LeafKernelWorkspace::default(); for points in [17, 7, 17] { let dots = test_dots(Metric::L2, points); let mut output = vec![LeafNeighbor::default(); points * 2]; @@ -912,7 +899,7 @@ mod integration_tests { metric, test_input(dots, points), MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), - &mut LeafKernelWorkspace::new(), + &mut LeafKernelWorkspace::default(), ) .unwrap(); (leaf_k, output) @@ -1079,7 +1066,7 @@ mod integration_tests { Metric::L2, test_input(&dots, 2), MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::new(), + &mut LeafKernelWorkspace::default(), ) .unwrap_err(); @@ -1133,7 +1120,7 @@ mod integration_tests { Metric::L2, non_square, MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::new(), + &mut LeafKernelWorkspace::default(), ), Err(LeafKernelError::NonSquareDots { rows: 2, cols: 3 }) ); @@ -1145,7 +1132,7 @@ mod integration_tests { Metric::L2, test_input(&square, 3), MutMatrixView::try_from(&mut wrong_rows[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::new(), + &mut LeafKernelWorkspace::default(), ), Err(LeafKernelError::InvalidOutputRows { expected: 3, @@ -1160,7 +1147,7 @@ mod integration_tests { Metric::L2, test_input(&square, 3), MutMatrixView::try_from(&mut too_many[..], 3, 3).unwrap(), - &mut LeafKernelWorkspace::new(), + &mut LeafKernelWorkspace::default(), ), Err(LeafKernelError::InvalidNeighborCount { points: 3, @@ -1176,7 +1163,7 @@ mod integration_tests { Metric::L2, test_input(&square, 5), MutMatrixView::try_from(&mut too_wide[..], 5, 4).unwrap(), - &mut LeafKernelWorkspace::new(), + &mut LeafKernelWorkspace::default(), ), Err(LeafKernelError::InvalidNeighborCount { points: 5, diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 7beeb443b5..564e6df005 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -34,13 +34,6 @@ pub struct PartitionKernelWorkspace { } impl PartitionKernelWorkspace { - /// Construct an empty allocation-free workspace. - pub const fn new() -> Self { - Self { - tracker: Vec::new(), - } - } - fn prepare(&mut self, fanout: usize) -> Result<(), PartitionKernelError> { let additional = fanout.saturating_sub(self.tracker.len()); self.tracker @@ -667,7 +660,7 @@ mod tests { Metric::Cosine, input, MutMatrixView::try_from(actual.as_mut_slice(), point_scales.len(), 2).unwrap(), - &mut PartitionKernelWorkspace::new(), + &mut PartitionKernelWorkspace::default(), ) .unwrap(); @@ -689,7 +682,7 @@ mod tests { #[test] fn workspace_reuses_runtime_fanout_capacity() { - let mut workspace = PartitionKernelWorkspace::new(); + let mut workspace = PartitionKernelWorkspace::default(); workspace.prepare(32).unwrap(); let allocation = workspace.tracker.as_ptr(); @@ -850,7 +843,7 @@ mod integration_tests { metric, input, MutMatrixView::try_from(output.as_mut_slice(), input.dots.nrows(), fanout).unwrap(), - &mut PartitionKernelWorkspace::new(), + &mut PartitionKernelWorkspace::default(), )?; Ok(output) } @@ -1043,7 +1036,7 @@ mod integration_tests { Metric::InnerProduct, valid_input, MutMatrixView::try_from(&mut wrong_output[..], 1, 3).unwrap(), - &mut PartitionKernelWorkspace::new(), + &mut PartitionKernelWorkspace::default(), ), Err(PartitionKernelError::InvalidOutputShape { expected_rows: 2, From 9a08dc2b5d754c3140d2f955242b4568415496d1 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:53:16 +0000 Subject: [PATCH 45/80] docs(pipnn): align kernel comments with implementation --- diskann-linalg/src/faer.rs | 13 +- diskann-linalg/src/lib.rs | 12 +- diskann/src/graph/pipnn/kernel_metric.rs | 136 ++++++++--------- diskann/src/graph/pipnn/leaf_kernel.rs | 152 +++++++++---------- diskann/src/graph/pipnn/mod.rs | 40 ++--- diskann/src/graph/pipnn/partition_kernel.rs | 157 +++++++++----------- 6 files changed, 229 insertions(+), 281 deletions(-) diff --git a/diskann-linalg/src/faer.rs b/diskann-linalg/src/faer.rs index 55ca5b2dfb..f4f498a13b 100644 --- a/diskann-linalg/src/faer.rs +++ b/diskann-linalg/src/faer.rs @@ -53,14 +53,13 @@ pub(super) fn sgemm_impl( faer::linalg::matmul::matmul(c, beta, a, b, alpha, Par::Seq) } -/// Implements the public lower-triangular AAT operation. +/// Compute the lower triangle of `A * Aᵀ` with Faer. /// -/// 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. +/// Leaf selection reads each symmetric pair once. It does not read the upper +/// triangle. `BlockStructure::TriangularLower` prevents writes to that triangle. +/// +/// `sgemm_aat_lower` checks both slice lengths and both size products. Therefore, +/// the Faer matrix views stay inside 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-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index db9a7a6ff4..778fb94cf8 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -189,16 +189,16 @@ pub fn sgemm( Ok(()) } -/// Computes the lower triangle of $C = A A^\mathsf{T}$ for a dense row-major -/// $m \times k$ matrix $A$. +/// Compute the lower triangle of $C = A A^\mathsf{T}$. /// -/// The lower triangle, including the diagonal, is overwritten. The upper -/// triangle of the $m \times m$ destination is left unchanged. +/// `A` is a dense row-major $m \times k$ matrix. The function overwrites the +/// lower triangle of `C`, including its diagonal. It does not change the upper +/// triangle. /// /// # Errors /// -/// Returns an error if a matrix-size calculation overflows or either slice does -/// not match its declared dimensions. +/// Returns an error if a size product overflows. It also returns an error if a +/// slice length does not match its declared matrix shape. 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)?; diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 3859467ca5..193dbc1f29 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -3,29 +3,29 @@ * Licensed under the MIT license. */ -//! Metric formulas shared by PiPNN partition and leaf kernels. +//! Metric formulas for the PiPNN partition and leaf kernels. //! -//! The build boundary converts runtime [`Metric`] into one zero-sized marker -//! type. That concrete type is carried through partition and leaf construction, -//! so scalar and SIMD hot loops contain neither metric matches nor trait objects. +//! `build_graph` maps each runtime [`Metric`] to one zero-sized marker type. The +//! partition and leaf functions receive that concrete type. Their hot loops use +//! no metric match or trait object. //! -//! All formulas produce ascending scores. L2 uses squared norms; unnormalized -//! cosine uses norms with zero/subnormal inputs mapped to zero similarity; -//! normalized cosine and inner product need no scales. Ordered comparisons leave -//! NaN non-rankable. The L2 partition scalar tail deliberately keeps its -//! non-fused operation order because rounding can change leader ties. +//! Each formula returns an ascending score. L2 uses squared norms. Cosine uses +//! norms and maps a zero norm to zero similarity. Normalized cosine and inner +//! product do not use norms. Ordered comparisons do not rank NaN. //! -//! [`ScaleKind`] records required scale representation and [`KernelMetric`] -//! owns the leaf and partition formulas. Runtime selection belongs to the build -//! entry point; this module contains no dispatch or type erasure. +//! The L2 partition SIMD path uses fused arithmetic. Its scalar tail uses +//! non-fused arithmetic. This operation order is part of the tie-order contract. +//! +//! [`ScaleKind`] defines the stored norm unit. [`KernelMetric`] defines the leaf +//! and partition formulas. use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; -/// Stored scale representation consumed by one kernel position. +/// Stored norm representation for one kernel input. /// -/// Associated constants on `KernelMetric` let the compiler remove unused scale -/// loads and allocations after metric selection. +/// `KernelMetric` uses associated constants for these values. The compiler +/// removes unused norm loads and workspace from each concrete metric. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum ScaleKind { /// Metric does not read this scale position. @@ -39,17 +39,14 @@ pub(crate) enum ScaleKind { } impl ScaleKind { - /// Convert stored scale to the arithmetic form required by a kernel. + /// Convert a stored norm to the unit that the kernel requires. /// - /// 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. + /// A squared norm below `f32::MIN_POSITIVE` becomes zero. A norm below + /// `sqrt(f32::MIN_POSITIVE)` also becomes zero. The function does not change + /// NaN, so the kernel does not rank it. /// - /// `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. + /// The concrete [`KernelMetric`] supplies `self` as an associated constant. + /// The compiler selects one match arm for each kernel instance. #[inline(always)] pub(crate) fn transform(self, stored: f32) -> f32 { match self { @@ -72,26 +69,22 @@ impl ScaleKind { } } - /// Return whether callers must supply this scale position. + /// Return `true` when the metric requires this norm input. /// - /// Calls use an associated constant, so this test compiles out of hot loops. + /// The compiler removes this test from each concrete metric loop. pub(crate) const fn is_some(self) -> bool { !matches!(self, Self::None) } } -/// Concrete metric contract shared by leaf and partition hot loops. +/// Metric contract for the leaf and partition hot loops. /// -/// Runtime `Metric` is converted to one implementor at the build boundary. -/// Generic methods then inline metric arithmetic through the complete partition -/// and leaf stages. Those operations remain separate because L2 partition -/// ranking deliberately omits the point norm. +/// `build_graph` selects one implementation. Generic calls inline its arithmetic +/// through the complete build. Leaf and partition formulas are separate because +/// L2 partition ranking does not need 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. +/// Each method returns an ascending score. Strict comparisons do not rank NaN. +/// Marker types contain no data. Associated constants remove unused norm work. pub(crate) trait KernelMetric: Send + Sync + 'static { /// Runtime tag represented by this marker. const METRIC: Metric; @@ -102,39 +95,37 @@ pub(crate) trait KernelMetric: Send + Sync + 'static { /// Leader-column scale representation used by partition assignment. const PARTITION_LEADER_SCALE: ScaleKind; - /// SIMD distance for one leaf source against a lane group of earlier targets. + /// Compute SIMD distances from one leaf source to 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. + /// `dot` and `target_scale` contain one target per lane. `source_scale` + /// contains the source norm in each lane. A metric without norms receives + /// zero for both norm arguments. The result contains one 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`. + /// Compute the 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. + /// The inputs and result represent one SIMD lane. Each implementation + /// documents any required operation order. 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. + /// Compute SIMD scores from one point to a 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. + /// `dot` and `leader_scale` contain one leader per lane. `point_scale` + /// contains the point norm in each lane. A metric without a norm receives + /// zero for that argument. The formula can omit terms that are constant for + /// all 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`. + /// Compute the 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. + /// The inputs and result represent one SIMD lane. Each implementation uses + /// its documented operation order. fn partition_distance_scalar(dot: f32, point_scale: f32, leader_scale: f32) -> f32; } @@ -147,10 +138,10 @@ pub(crate) struct CosineNormalized; /// Negative-inner-product marker. pub(crate) struct InnerProduct; -/// Clamp negative SIMD roundoff to zero while preserving NaN lanes. +/// Clamp negative SIMD roundoff to zero and keep NaN lanes unchanged. /// -/// One ordered self-comparison normalizes backend-specific SIMD `max` NaN -/// behavior; no lane branches or allocations are introduced. +/// SIMD `max` has architecture-specific NaN behavior. The ordered self-test +/// selects the original value for each NaN lane. #[inline(always)] fn clamp_nonnegative(arch: F::Arch, distance: F) -> F where @@ -158,8 +149,8 @@ where 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. + // Select the original value for NaN lanes. This gives all architectures the + // same non-rankable NaN result. distance .eq_simd(distance) .select(zero.max_simd(distance), distance) @@ -171,15 +162,14 @@ fn clamp_nonnegative_scalar(distance: f32) -> f32 { if distance < 0.0 { 0.0 } else { distance } } -/// Compute cosine distance while preserving DiskANN zero/NaN semantics. +/// Compute cosine distance with the DiskANN zero-norm and NaN rules. /// -/// Zero lanes divide by one only to keep the operation defined, then explicitly -/// 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. +/// A zero-norm lane divides by one and then selects zero similarity. A NaN norm +/// propagates through division. If the other norm is zero, 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. +/// Each input contains one point pair per lane. The result is +/// `1 - cosine_similarity`. The function uses no lane branch. #[inline(always)] fn cosine_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -238,15 +228,15 @@ 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. + // The point norm is constant for this ranking. The SIMD path uses the + // fused multiply-add operation that defines its tie order. 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. + // The scalar tail uses non-fused subtraction. A fused operation can + // change rounding and select a different leader at a tie. leader_scale - 2.0 * dot } } @@ -280,14 +270,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. + // Partitioning uses only score order. Do not clamp the score because a + // clamp can change the order of near ties. 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. + // Use the same unclamped score in the scalar tail. cosine_distance_scalar(dot, point_scale, leader_scale) } } @@ -327,7 +317,7 @@ impl KernelMetric for CosineNormalized { #[inline(always)] fn partition_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { - // Preserve the unclamped ranking score used by full SIMD groups. + // Use the same unclamped score as the SIMD path. 1.0 - dot } } diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 229e918c3b..70a91720b1 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -3,26 +3,26 @@ * Licensed under the MIT license. */ -//! Leaf-local top-k selection over a lower-triangular Gram matrix. +//! Leaf-local top-k selection from a lower-triangular Gram matrix. //! -//! Caller supplies an `n × n` [`MatrixView`] produced by `sgemm_aat_lower`. -//! Diagonal entries provide metric scales; only strict-lower pair dots are read. -//! Each pair is evaluated once and offered to both endpoint rows. +//! The input is an `n × n` [`MatrixView`] from `sgemm_aat_lower`. The diagonal +//! contains metric norms. The kernel reads only the strict lower triangle. It +//! evaluates each point pair once and updates both points. //! -//! Output is an `n × k` matrix of sorted [`LeafNeighbor`] values with leaf-local -//! targets. Supported `k` is zero through [`MAX_LEAF_NEIGHBORS`]; positive widths -//! use fixed arrays. Strict comparisons preserve encounter order for ties and -//! reject NaN. L2, cosine, normalized cosine, and inner product share the same -//! scalar/SIMD traversal. +//! The output is an `n × k` matrix of sorted [`LeafNeighbor`] values. Each target +//! is a position in the leaf. The kernel supports `k` from zero through +//! [`MAX_LEAF_NEIGHBORS`]. Positive widths use fixed arrays. //! -//! Callers select architecture and metric once outside the leaf loop, then call -//! [`nearest_neighbors`] with concrete `A` and `M` types. Every call validates -//! square shape, row count, local-ID bounds, and output width before scratch -//! mutation or unchecked SIMD loads. [`LeafKernelWorkspace`] retains per-worker -//! norm and threshold buffers. +//! Strict comparisons keep scan order for equal distances. They do not rank NaN. +//! All supported metrics use the same scalar and SIMD traversal. //! -//! Work is `n(n - 1) / 2` distance evaluations with constant bounded insertion; -//! scratch is `O(n)` and output is `O(nk)`. +//! The caller supplies concrete architecture `A` and metric `M`. The function +//! checks all shapes and local-ID bounds before it changes workspace or uses an +//! unchecked SIMD load. [`LeafKernelWorkspace`] stores reusable norms and +//! rejection thresholds. +//! +//! The kernel evaluates `n(n - 1) / 2` distances. Scratch size is `O(n)`. Output +//! size is `O(nk)`. use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; @@ -44,8 +44,8 @@ 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. + /// `target` is a position in the leaf. `distance` is its score relative to + /// the source of the output row. pub const fn new(target: u32, distance: f32) -> Self { Self { target, distance } } @@ -126,12 +126,11 @@ pub enum LeafKernelError { }, } -/// Return the usable non-self neighbor count for one leaf. +/// Return the non-self neighbor count for one leaf. /// -/// `points` is the leaf point count and `requested_k` is the build-wide target. -/// Values above [`MAX_LEAF_NEIGHBORS`] are rejected. Otherwise the returned -/// width is `min(requested_k, points - 1)`, allowing empty, singleton, and small -/// leaves without a second effective-k state. +/// `points` is the number of points in the leaf. `requested_k` is the configured +/// neighbor count. The result is `min(requested_k, points - 1)`. The function +/// rejects a value above [`MAX_LEAF_NEIGHBORS`]. /// /// # Errors /// @@ -158,8 +157,7 @@ pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result Result( arch: A, input: MatrixView<'_, f32>, @@ -229,12 +225,11 @@ where Ok(()) } -/// Validate the complete safety contract before dispatched SIMD executes. +/// Check the safety conditions for the SIMD kernel. /// -/// `MatrixView` and `MutMatrixView` construction guarantee exact, non-overflowing -/// backing lengths. This check establishes square dots, representable local IDs, -/// and an output width bounded by the point count and fixed kernel capacity. -/// Failure returns [`LeafKernelError`] before output or workspace mutation. +/// The matrix views already prove their backing lengths. This function checks +/// that the dot matrix is square. It also checks local-ID range and output width. +/// An error occurs before the kernel changes output or workspace. fn validate( input: MatrixView<'_, f32>, output: &MutMatrixView<'_, LeafNeighbor>, @@ -269,16 +264,13 @@ fn validate( Ok(()) } -/// Prepare metric-specific scale and threshold scratch. +/// Prepare metric norms and rejection thresholds. /// -/// 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. +/// L2 uses each diagonal value as a squared norm. Cosine converts each diagonal +/// value to a norm. Normalized cosine and inner product clear the norm buffer. /// -/// `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. +/// The workspace keeps at most `O(n)` capacity in each buffer. An allocation +/// error occurs before SIMD traversal. fn prepare_workspace( input: MatrixView<'_, f32>, workspace: &mut LeafKernelWorkspace, @@ -319,12 +311,10 @@ fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> Result( arch: F::Arch, input: MatrixView<'_, f32>, @@ -354,11 +344,10 @@ where Ok(()) } -/// Reinterpret validated output as one fixed array per source, then run shared -/// pair traversal. +/// Split the validated output into one fixed array for each source. /// -/// `N` is one, two, or three. `as_chunks_mut` performs one safe shape split per -/// leaf, keeping array conversion out of candidate insertion. +/// `N` is one, two, or three. The function performs one safe split for each +/// leaf. It then starts the shared pair traversal. fn process_fixed_width( arch: F::Arch, input: MatrixView<'_, f32>, @@ -376,27 +365,22 @@ fn process_fixed_width( process_pairs::(arch, input, neighbor_lists, norms, worst); } -/// Scan the strict lower triangle once and update both endpoint sources. +/// Scan the strict lower triangle and update both points of each pair. /// -/// Invariants on entry: +/// Entry conditions: /// -/// - `dots` is a validated square row-major matrix; -/// - `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. +/// - `input` is a square row-major matrix. +/// - `output` has one sorted list for each point. +/// - `worst[source]` equals the distance in the last output slot. +/// - `norms` has one value per point when metric `M` requires norms. /// -/// Each SIMD chunk computes both endpoint eligibility masks before mutation. -/// 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. +/// Each SIMD chunk computes both eligibility masks before it changes output. +/// Several lanes can update the current source. Each such lane checks the current +/// threshold again. Each target lane updates a different earlier source. /// -/// `M` is concrete before type erasure and `N` is one through three. `input` -/// supplies `n × n` dots, `output` owns `n` sorted fixed-size lists, `norms` -/// holds metric scales when required, and `worst` mirrors every list's final -/// distance. The function 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. +/// The scalar tail uses the scalar formula for `M`. This keeps its specified +/// rounding order. The function evaluates exactly `n(n - 1) / 2` pairs. It +/// inserts accepted candidates in scan order. #[inline(never)] fn process_pairs( arch: F::Arch, @@ -437,12 +421,12 @@ fn process_pairs( } 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. + // Source zero has no earlier target. Each source after zero can still add + // itself to the neighbor list of 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. + // `uses_norms` is an associated constant of `M`. The compiler removes + // this branch and all norm loads from metrics that do not use norms. let source_scale = if uses_norms { F::splat(arch, norms[source]) } else { @@ -535,17 +519,15 @@ fn process_pairs( debug_assert_eq!(output.len(), point_count); } -/// Insert into a fixed-width neighbor list and return its new worst distance. +/// Insert one candidate into a fixed-width sorted list. /// -/// Width is a compile-time constant from one through three. Strict `<` -/// comparisons preserve scan order for ties; callers already rejected NaN via -/// the eligibility comparison. Explicit shifts save about 0.5% estimated cycles -/// versus the generic bubble loop in the local Callgrind `k=3` fixture. +/// Width `N` is one, two, or three. The caller has already proved that the +/// candidate is better than the last slot. Strict comparisons keep scan order +/// for equal distances. The eligibility test has already rejected NaN. /// -/// `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. Unsupported instantiations return an underfill sentinel; -/// `process_neighbor_width` never constructs them. +/// The function returns the new last-slot distance. Explicit shifts used 0.5% +/// fewer estimated cycles than a generic bubble loop in the local `k=3` +/// Callgrind fixture. An unsupported `N` returns the underfill sentinel. #[inline(always)] fn insert_fixed_neighbor( neighbors: &mut [LeafNeighbor; N], diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index cb478ec113..1dfd90c15d 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,34 +3,26 @@ * Licensed under the MIT license. */ -//! Numerical kernels for provider-independent PiPNN graph construction. +//! Numerical kernels for PiPNN graph construction. //! -//! PiPNN assigns points to overlapping leader partitions, computes one -//! lower-triangular all-pairs matrix per bounded leaf, and merges the selected -//! leaf neighbors into graph candidates. This layer contains only the two score -//! selection kernels; later layers own partition recursion, GEMM, graph IDs, -//! candidate merging, final pruning, providers, and persistence. +//! [`partition_kernel`] converts point-to-leader dot products into sorted leader +//! positions. The output width sets the fanout. One workspace stores the +//! runtime-sized tracker and reuses it for each point. //! -//! - [`partition_kernel`] converts a point-by-leader dot-product tile into sorted -//! leader-column positions. Output width is runtime fanout; reusable tracker -//! storage grows to that width and is reused across point rows. -//! - [`leaf_kernel`] scans each strict-lower-triangle pair once, updates both -//! endpoints, and retains up to three leaf-local neighbors per point. -//! - `kernel_metric` owns the scalar/SIMD formulas and exact scale units shared -//! by both kernels. +//! [`leaf_kernel`] reads a lower-triangular Gram matrix. It evaluates each point +//! pair once and updates both points. Each point retains at most three local +//! neighbors. //! -//! Runtime architecture and metric selection intentionally do not live in the -//! kernels. The enclosing graph build selects concrete `A` and `M` types once, -//! then carries them through partition and leaf work. This keeps metric matches, -//! trait objects, and stored function pointers out of the hot loops. +//! `kernel_metric` defines the scalar and SIMD formulas. It also defines the +//! required norm units for each metric. //! -//! Both kernels validate view relationships and metric scale layouts before -//! unchecked SIMD access. They borrow all matrices and mutate only caller-owned -//! output and reusable workspace. Partition work is one score per point-leader -//! pair; leaf work is one score per unordered point pair. - -// This stack layer introduces the numerical kernels before the following core -// layer wires them into graph construction. +//! The graph builder selects architecture `A` and metric `M` once. It passes +//! these concrete types to both kernels. The hot loops use no metric match, +//! trait object, or stored function pointer. +//! +//! Each kernel checks all view and scale relationships before unchecked SIMD +//! access. The kernels borrow their matrices. They write only to caller-owned +//! output and workspace. #[allow(dead_code)] mod kernel_metric; diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 564e6df005..a8d4dee4a7 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -5,19 +5,18 @@ //! Nearest-leader selection for PiPNN partition assignment. //! -//! Caller supplies a row-major point-by-leader dot matrix plus metric-specific -//! [`PartitionScales`]. Output contains sorted leader-column positions for each -//! point; fanout is the output width and cannot exceed the leader count. +//! The input contains a row-major point-to-leader dot matrix and +//! metric-specific [`PartitionScales`]. The output contains sorted leader-column +//! positions for each point. Its width sets the fanout and cannot exceed the +//! leader count. //! -//! Callers select architecture and metric once outside partition recursion, -//! then call [`nearest_leaders`] with concrete `A` and `M` types. Calls validate -//! row counts, scale variants and lengths, fanout, and leader-ID representation -//! before output mutation or unchecked SIMD loads. +//! The caller supplies concrete architecture `A` and metric `M`. The function +//! checks row counts, scale units, scale lengths, fanout, and leader-ID range. +//! These checks occur before output changes or unchecked SIMD loads. //! -//! L2 omits the point norm because it cannot change one point's leader order. -//! Strict comparisons preserve scan order for ties and leave NaN non-rankable. -//! Each point evaluates every leader; competitive scores move through a -//! caller-owned tracker reused across points. +//! L2 omits the point norm because it is constant for one point. Strict +//! comparisons keep leader scan order for equal scores. They do not rank NaN. +//! One runtime-sized workspace tracks the nearest leaders for each point. use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -44,12 +43,11 @@ impl PartitionKernelWorkspace { } } -/// Metric-specific normalization inputs for one partition tile. +/// Metric-specific norm 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. +/// The kernel checks each slice length before it changes output. Cosine point +/// values are squared norms from a matrix diagonal. Cosine leader values are +/// norms that partition setup computes once. #[derive(Clone, Copy, Debug)] pub enum PartitionScales<'a> { /// L2 needs only squared leader norms; the point norm cannot affect ranking. @@ -68,11 +66,10 @@ pub enum PartitionScales<'a> { None, } -/// One row-major point-by-leader dot-product tile. +/// One row-major point-to-leader dot-product tile. /// -/// Matrix rows are points, columns are leaders, and [`Self::scales`] must match -/// concrete metric selected by the enclosing build. This value only borrows -/// input; the numerical kernel stores no tile state. +/// Rows are points and columns are leaders. [`Self::scales`] must match concrete +/// metric `M`. This value borrows all input and stores no kernel state. #[derive(Clone, Copy, Debug)] pub struct PartitionInput<'a> { /// One point per matrix row and one leader per column. @@ -139,11 +136,10 @@ pub enum PartitionKernelError { }, } -/// Select nearest leader positions for every input point. +/// Select the nearest leader positions for each input point. /// -/// `output.nrows()` must equal `input.dots.nrows()` and fanout, represented by -/// `output.ncols()`, must not exceed the leader count. `A` and `M` must already -/// have been selected at the enclosing build boundary. +/// `output.nrows()` must equal `input.dots.nrows()`. `output.ncols()` sets the +/// fanout and must not exceed the leader count. pub(crate) fn nearest_leaders( arch: A, input: PartitionInput<'_>, @@ -181,23 +177,21 @@ where Ok(()) } -/// Validated scale slices in the storage form required by `M`. +/// Checked norm slices in the storage form that `M` requires. /// -/// Empty slices are intentional for metrics that omit a scale; consumers branch -/// on associated `ScaleKind` constants that monomorphize out of hot loops. +/// A metric that does not use a norm receives an empty slice. Associated +/// `ScaleKind` constants remove these branches from concrete metric loops. #[derive(Clone, Copy)] struct ScaleSlices<'a> { point_scales: &'a [f32], leader_scales: &'a [f32], } -/// Validate the complete partition-kernel safety and metric contract. +/// Check the safety and metric conditions for partition selection. /// -/// `MatrixView` and `MutMatrixView` construction guarantee exact, -/// non-overflowing backing lengths. The `PartitionScales` variant must match -/// concrete metric `M`, preventing plausible but incorrect norm units from -/// crossing the interface. Success returns normalized scale slices and -/// establishes representable leader IDs plus bounded fanout. +/// The matrix views already prove their backing lengths. This function checks +/// row counts, leader-ID range, fanout, scale variant, and scale lengths. A +/// successful result contains norm slices in the units that `M` requires. fn validate<'a, M: KernelMetric>( input: PartitionInput<'a>, output: &MutMatrixView<'_, u32>, @@ -223,9 +217,8 @@ 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]`. + // Match the scale variant to concrete metric `M` before extracting slices. + // This prevents use of a squared point norm as a leader norm. let scales = match (M::METRIC, input.scales) { ( Metric::L2, @@ -266,8 +259,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. + // The associated scale kinds define the exact slice lengths. A metric that + // does not use a scale must receive an empty slice. check_length( "point scales", scales.point_scales.len(), @@ -281,9 +274,9 @@ fn validate<'a, M: KernelMetric>( Ok(scales) } -/// Return required scale length after metric specialization. +/// Return the required norm-slice length for one concrete metric. /// -/// Associated `ScaleKind` constants make this choice compile away. +/// The compiler removes this choice from the metric loop. const fn expected_scale_len(kind: ScaleKind, count: usize) -> usize { if kind.is_some() { count } else { 0 } } @@ -304,25 +297,21 @@ fn check_length( } } -/// Convert each point's leader scores into sorted top-fanout IDs. +/// Convert each point's leader scores into sorted leader IDs. /// -/// Per-point flow: +/// For each point, the function does these steps: /// -/// 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. +/// 1. Convert the point norm to the unit that `M` requires. +/// 2. Process complete SIMD groups. +/// 3. Process the scalar tail. +/// 4. Copy the sorted leader IDs to the output row. /// -/// `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. +/// `tracker` has `fanout` entries and stays sorted after each insertion. Strict +/// comparisons keep leader scan order for equal scores. They do not rank NaN. /// -/// `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. +/// The function computes `p * l` scores. One insertion can move `O(fanout)` +/// entries. The caller allocates the runtime-sized tracker once and reuses it for +/// all point rows. fn process_points( arch: F::Arch, dots: MatrixView<'_, f32>, @@ -356,24 +345,24 @@ fn process_points( ); } let fanout = tracker.len(); - // Each point is independent. Reset the caller-owned tracker so no assignment - // state or tie order leaks across rows. + // Reset the tracker for each point. No assignment state can pass from one + // output row to another. for (point, (point_dots, point_output)) in dots .row_iter() .zip(output.chunks_exact_mut(fanout)) .enumerate() { tracker.fill((u32::MAX, f32::INFINITY)); - // Transform once per point rather than once per leader. For metrics - // without a point scale, specialization removes this branch and load. + // Convert the point norm once for this row. The compiler removes this + // branch and load from metrics that do not use a point norm. let point_scale = if M::PARTITION_POINT_SCALE.is_some() { M::PARTITION_POINT_SCALE.transform(scales.point_scales[point]) } else { 0.0 }; let point_scale_vector = F::splat(arch, point_scale); - // Split at the largest complete vector boundary. Scalar tail uses the - // metric's explicit scalar operation order, not a padded SIMD load. + // Process all complete SIMD groups first. The scalar tail uses the + // metric's scalar operation order. let full = leader_count / F::LANES * F::LANES; for base in (0..full).step_by(F::LANES) { @@ -393,8 +382,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. + // Use scalar formulas for the tail. A padded SIMD load can read past the + // norm slice and can change 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]) @@ -407,22 +396,20 @@ 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. + // Distances stay in the workspace. Partition construction needs only the + // leader-column positions in nearest-first order. copy_leader_ids(tracker, point_output); } } -/// Offer competitive SIMD lanes to a point tracker in increasing leader order. +/// Insert competitive SIMD lanes 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. +/// One broadcast comparison rejects a group that cannot improve the last slot. +/// Bit iteration proceeds from low lane to high lane. This order matches scalar +/// tie behavior for all 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. +/// `distances` starts at `first_leader`. `tracker` is sorted. Each accepted lane +/// can move `O(fanout)` entries. fn insert_leader_lanes(distances: F, first_leader: usize, tracker: &mut [(u32, f32)]) where F: SIMDVector> + SIMDPartialOrd, @@ -443,15 +430,14 @@ where } } -/// Insert one strictly better candidate while preserving sorted-prefix state. +/// Insert one better candidate into a sorted tracker. /// -/// 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. +/// The function replaces the last slot and moves the new value to the left. +/// Equal scores and NaN do not enter. Thus, scan order resolves equal scores. +/// The last slot is also the 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. +/// `tracker` must be sorted and non-empty. `leader` is a local column position. +/// One insertion moves at most `fanout - 1` entries. #[inline(always)] fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, distance: f32) { let threshold = tracker.len() - 1; @@ -467,10 +453,10 @@ fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, distance: f32) { } } -/// Publish only leader IDs; distances stay private tracker state. +/// Copy the retained leader IDs to one output row. /// -/// `assignments.len()` is validated fanout. Copying costs `O(fanout)` and leaves -/// tracker state available for the underfill sentinel check encoded in IDs. +/// `assignments.len()` equals the checked fanout. The tracker keeps its distances +/// for the underfill check. fn copy_leader_ids(tracker: &[(u32, f32)], assignments: &mut [u32]) { for (destination, &(leader, _)) in assignments.iter_mut().zip(tracker) { *destination = leader; @@ -565,9 +551,8 @@ mod tests { } } - // 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. + // This oracle checks SIMD chunking, scalar tails, and tracker order. It uses + // `M::partition_distance_scalar`. Separate tests define each formula directly. fn scalar_traversal_reference( input: PartitionInput<'_>, fanout: usize, From 8499c7d912fab29a6b69d3cbee86da24943d4771 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:01:29 +0000 Subject: [PATCH 46/80] docs(pipnn): remove kernel performance notes --- diskann/src/graph/pipnn/kernel_metric.rs | 10 +++------- diskann/src/graph/pipnn/leaf_kernel.rs | 21 +++------------------ diskann/src/graph/pipnn/mod.rs | 3 +-- diskann/src/graph/pipnn/partition_kernel.rs | 14 +++++--------- 4 files changed, 12 insertions(+), 36 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 193dbc1f29..eae0a2806f 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -6,8 +6,7 @@ //! Metric formulas for the PiPNN partition and leaf kernels. //! //! `build_graph` maps each runtime [`Metric`] to one zero-sized marker type. The -//! partition and leaf functions receive that concrete type. Their hot loops use -//! no metric match or trait object. +//! partition and leaf functions receive that concrete type. //! //! Each formula returns an ascending score. L2 uses squared norms. Cosine uses //! norms and maps a zero norm to zero similarity. Normalized cosine and inner @@ -24,8 +23,7 @@ use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; /// Stored norm representation for one kernel input. /// -/// `KernelMetric` uses associated constants for these values. The compiler -/// removes unused norm loads and workspace from each concrete metric. +/// `KernelMetric` supplies this value as an associated constant. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum ScaleKind { /// Metric does not read this scale position. @@ -70,14 +68,12 @@ impl ScaleKind { } /// Return `true` when the metric requires this norm input. - /// - /// The compiler removes this test from each concrete metric loop. pub(crate) const fn is_some(self) -> bool { !matches!(self, Self::None) } } -/// Metric contract for the leaf and partition hot loops. +/// Metric contract for leaf and partition selection. /// /// `build_graph` selects one implementation. Generic calls inline its arithmetic /// through the complete build. Leaf and partition formulas are separate because diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 70a91720b1..88e94395c0 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -20,9 +20,6 @@ //! checks all shapes and local-ID bounds before it changes workspace or uses an //! unchecked SIMD load. [`LeafKernelWorkspace`] stores reusable norms and //! rejection thresholds. -//! -//! The kernel evaluates `n(n - 1) / 2` distances. Scratch size is `O(n)`. Output -//! size is `O(nk)`. use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; @@ -137,10 +134,6 @@ pub enum LeafKernelError { /// Returns [`LeafKernelError::TooManyPoints`] when leaf-local positions cannot /// fit in `u32`, or [`LeafKernelError::InvalidNeighborCount`] when `requested_k` /// exceeds [`MAX_LEAF_NEIGHBORS`]. -/// -/// # 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)); @@ -163,10 +156,6 @@ pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result Result { checked_area("output", points, leaf_neighbor_count(points, requested_k)?) } @@ -269,8 +258,7 @@ fn validate( /// L2 uses each diagonal value as a squared norm. Cosine converts each diagonal /// value to a norm. Normalized cosine and inner product clear the norm buffer. /// -/// The workspace keeps at most `O(n)` capacity in each buffer. An allocation -/// error occurs before SIMD traversal. +/// An allocation error occurs before SIMD traversal. fn prepare_workspace( input: MatrixView<'_, f32>, workspace: &mut LeafKernelWorkspace, @@ -425,8 +413,6 @@ fn process_pairs( // itself to the neighbor list of source zero. for source in 1..point_count { let source_start = source * point_count; - // `uses_norms` is an associated constant of `M`. The compiler removes - // this branch and all norm loads from metrics that do not use norms. let source_scale = if uses_norms { F::splat(arch, norms[source]) } else { @@ -525,9 +511,8 @@ fn process_pairs( /// candidate is better than the last slot. Strict comparisons keep scan order /// for equal distances. The eligibility test has already rejected NaN. /// -/// The function returns the new last-slot distance. Explicit shifts used 0.5% -/// fewer estimated cycles than a generic bubble loop in the local `k=3` -/// Callgrind fixture. An unsupported `N` returns the underfill sentinel. +/// The function returns the new last-slot distance. An unsupported `N` returns +/// the underfill sentinel. #[inline(always)] fn insert_fixed_neighbor( neighbors: &mut [LeafNeighbor; N], diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 1dfd90c15d..1c51887d81 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -17,8 +17,7 @@ //! required norm units for each metric. //! //! The graph builder selects architecture `A` and metric `M` once. It passes -//! these concrete types to both kernels. The hot loops use no metric match, -//! trait object, or stored function pointer. +//! these concrete types to both kernels. //! //! Each kernel checks all view and scale relationships before unchecked SIMD //! access. The kernels borrow their matrices. They write only to caller-owned diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index a8d4dee4a7..246b718cc7 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -275,8 +275,6 @@ fn validate<'a, M: KernelMetric>( } /// Return the required norm-slice length for one concrete metric. -/// -/// The compiler removes this choice from the metric loop. const fn expected_scale_len(kind: ScaleKind, count: usize) -> usize { if kind.is_some() { count } else { 0 } } @@ -309,9 +307,8 @@ fn check_length( /// `tracker` has `fanout` entries and stays sorted after each insertion. Strict /// comparisons keep leader scan order for equal scores. They do not rank NaN. /// -/// The function computes `p * l` scores. One insertion can move `O(fanout)` -/// entries. The caller allocates the runtime-sized tracker once and reuses it for -/// all point rows. +/// The caller allocates the runtime-sized tracker once and reuses it for all +/// point rows. fn process_points( arch: F::Arch, dots: MatrixView<'_, f32>, @@ -353,8 +350,8 @@ fn process_points( .enumerate() { tracker.fill((u32::MAX, f32::INFINITY)); - // Convert the point norm once for this row. The compiler removes this - // branch and load from metrics that do not use a point norm. + // Convert the point norm once for this row. Metrics without a point norm + // use zero. let point_scale = if M::PARTITION_POINT_SCALE.is_some() { M::PARTITION_POINT_SCALE.transform(scales.point_scales[point]) } else { @@ -408,8 +405,7 @@ fn process_points( /// Bit iteration proceeds from low lane to high lane. This order matches scalar /// tie behavior for all SIMD widths. /// -/// `distances` starts at `first_leader`. `tracker` is sorted. Each accepted lane -/// can move `O(fanout)` entries. +/// `distances` starts at `first_leader`. `tracker` is sorted. fn insert_leader_lanes(distances: F, first_leader: usize, tracker: &mut [(u32, f32)]) where F: SIMDVector> + SIMDPartialOrd, From dc06de7bd496138fc1f97f2c82beee8c9affc2a2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:40:42 +0000 Subject: [PATCH 47/80] refactor(pipnn): report partition underfill locally --- diskann/src/graph/pipnn/partition_kernel.rs | 36 +++++++++------------ 1 file changed, 15 insertions(+), 21 deletions(-) diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 246b718cc7..dc4b444036 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -143,7 +143,7 @@ pub enum PartitionKernelError { pub(crate) fn nearest_leaders( arch: A, input: PartitionInput<'_>, - mut output: MutMatrixView<'_, u32>, + output: MutMatrixView<'_, u32>, workspace: &mut PartitionKernelWorkspace, ) -> Result<(), PartitionKernelError> where @@ -160,21 +160,7 @@ where } workspace.prepare(fanout)?; - process_points::( - arch, - input.dots, - scales, - output.as_mut_slice(), - &mut workspace.tracker, - ); - if let Some(point) = output - .as_slice() - .chunks_exact(fanout) - .position(|assignments| assignments[fanout - 1] == u32::MAX) - { - return Err(PartitionKernelError::InsufficientRankableLeaders { point, fanout }); - } - Ok(()) + process_points::(arch, input.dots, scales, output, &mut workspace.tracker) } /// Checked norm slices in the storage form that `M` requires. @@ -302,7 +288,8 @@ fn check_length( /// 1. Convert the point norm to the unit that `M` requires. /// 2. Process complete SIMD groups. /// 3. Process the scalar tail. -/// 4. Copy the sorted leader IDs to the output row. +/// 4. Check that the tracker is full. +/// 5. Copy the sorted leader IDs to the output row. /// /// `tracker` has `fanout` entries and stays sorted after each insertion. Strict /// comparisons keep leader scan order for equal scores. They do not rank NaN. @@ -313,9 +300,10 @@ fn process_points( arch: F::Arch, dots: MatrixView<'_, f32>, scales: ScaleSlices<'_>, - output: &mut [u32], + mut output: MutMatrixView<'_, u32>, tracker: &mut [(u32, f32)], -) where +) -> Result<(), PartitionKernelError> +where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: KernelMetric, @@ -341,12 +329,14 @@ fn process_points( "validated leader scales must match leader count" ); } - let fanout = tracker.len(); + let fanout = output.ncols(); + debug_assert!(fanout > 0); + debug_assert_eq!(tracker.len(), fanout); // Reset the tracker for each point. No assignment state can pass from one // output row to another. for (point, (point_dots, point_output)) in dots .row_iter() - .zip(output.chunks_exact_mut(fanout)) + .zip(output.as_mut_slice().chunks_exact_mut(fanout)) .enumerate() { tracker.fill((u32::MAX, f32::INFINITY)); @@ -393,10 +383,14 @@ fn process_points( M::partition_distance_scalar(dot, point_scale, leader_scale), ); } + if tracker[fanout - 1].0 == u32::MAX { + return Err(PartitionKernelError::InsufficientRankableLeaders { point, fanout }); + } // Distances stay in the workspace. Partition construction needs only the // leader-column positions in nearest-first order. copy_leader_ids(tracker, point_output); } + Ok(()) } /// Insert competitive SIMD lanes in increasing leader order. From 0db935049c1ede4c40ef482cbbd0b060149f13ed Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:46:14 +0000 Subject: [PATCH 48/80] refactor(pipnn): remove unused leaf output helper --- diskann/src/graph/pipnn/leaf_kernel.rs | 61 +++----------------------- 1 file changed, 7 insertions(+), 54 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 88e94395c0..9e5f7c1e30 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -75,16 +75,6 @@ pub enum LeafKernelError { /// 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. - buffer: &'static str, - /// Declared row count. - rows: usize, - /// Declared column count. - cols: usize, - }, /// The output matrix does not have one row per input point. #[error("invalid output row count: expected {expected}, got {actual} with {columns} columns")] InvalidOutputRows { @@ -148,18 +138,6 @@ pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result Result { - checked_area("output", points, leaf_neighbor_count(points, requested_k)?) -} - /// Select the nearest non-self positions for each point in a leaf. /// /// `output` must have one row for each input point. Its column count is the @@ -294,11 +272,6 @@ fn resize( Ok(()) } -fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> Result { - rows.checked_mul(cols) - .ok_or(LeafKernelError::ShapeOverflow { buffer, rows, cols }) -} - /// Select fixed-width storage for the validated neighbor count. /// /// This match runs once for each leaf. `as_chunks_mut` converts the output once. @@ -661,12 +634,12 @@ mod tests { } #[test] - fn output_length_clamps_to_non_self_neighbors_and_rejects_large_k() { - assert_eq!(leaf_output_len(0, 3).unwrap(), 0); - assert_eq!(leaf_output_len(1, 3).unwrap(), 0); - assert_eq!(leaf_output_len(4, 3).unwrap(), 12); + fn neighbor_count_clamps_to_non_self_neighbors_and_rejects_large_k() { + assert_eq!(leaf_neighbor_count(0, 3).unwrap(), 0); + assert_eq!(leaf_neighbor_count(1, 3).unwrap(), 0); + assert_eq!(leaf_neighbor_count(4, 3).unwrap(), 3); assert_eq!( - leaf_output_len(4, 4), + leaf_neighbor_count(4, 4), Err(LeafKernelError::InvalidNeighborCount { points: 4, neighbors: 4, @@ -675,23 +648,11 @@ mod tests { ); #[cfg(target_pointer_width = "64")] assert_eq!( - leaf_output_len(u32::MAX as usize + 1, 1), + leaf_neighbor_count(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 kernel_accepts_different_neighbor_counts() { let points = 7; @@ -740,7 +701,7 @@ mod integration_tests { use super::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, MAX_LEAF_NEIGHBORS, - dispatch_nearest_neighbors, leaf_neighbor_count, leaf_output_len, + dispatch_nearest_neighbors, leaf_neighbor_count, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -1158,12 +1119,4 @@ mod integration_tests { } } } - - #[test] - fn output_length_rejects_unrepresentable_point_count() { - assert_eq!( - leaf_output_len(usize::MAX, 1), - Err(LeafKernelError::TooManyPoints(usize::MAX)) - ); - } } From dfc7914eebc10a29370432c76cd3c26d7767c0a0 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:58:02 +0000 Subject: [PATCH 49/80] refactor(pipnn): remove kernel assertions --- diskann/src/graph/pipnn/leaf_kernel.rs | 38 ++++----------------- diskann/src/graph/pipnn/partition_kernel.rs | 23 +------------ 2 files changed, 8 insertions(+), 53 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 9e5f7c1e30..f536cf7660 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -321,8 +321,7 @@ fn process_fixed_width( M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let (neighbor_lists, remainder) = output.as_chunks_mut::(); - debug_assert!(remainder.is_empty()); + let (neighbor_lists, _) = output.as_chunks_mut::(); process_pairs::(arch, input, neighbor_lists, norms, worst); } @@ -356,30 +355,8 @@ fn process_pairs( u64: From<<::BitMask as SIMDMask>::Underlying>, { let point_count = input.nrows(); - assert_eq!( - input.ncols(), - point_count, - "validated leaf dot matrix must be square" - ); - assert_eq!( - output.len(), - point_count, - "validated leaf output must have one list per point" - ); - assert_eq!( - worst.len(), - point_count, - "validated leaf thresholds must have one value per point" - ); let dots = input.as_slice(); let uses_norms = M::LEAF_SCALE.is_some(); - if uses_norms { - assert_eq!( - norms.len(), - point_count, - "validated leaf norms must have one value per point" - ); - } let worst_ptr = worst.as_mut_ptr(); // Source zero has no earlier target. Each source after zero can still add @@ -391,7 +368,8 @@ fn process_pairs( } else { F::default(arch) }; - // SAFETY: `source < point_count == worst.len()` by the assertions above. + // SAFETY: `validate` and `prepare_workspace` established + // `source < point_count == worst.len()`. let mut source_worst = unsafe { *worst_ptr.add(source) }; let mut target = 0; @@ -399,8 +377,8 @@ fn process_pairs( // 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 < point_count`, and - // the assertion above established `norms.len() == point_count`. + // SAFETY: the full target chunk is below `source < point_count`. + // `prepare_workspace` created one norm for each point. unsafe { F::load_simd(arch, norms.as_ptr().add(target)) } } else { F::default(arch) @@ -409,8 +387,8 @@ fn process_pairs( // 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 < point_count`, and - // the assertion above established `worst.len() == point_count`. + // SAFETY: the full target chunk is below `source < point_count`. + // `prepare_workspace` created one threshold for each point. 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()); @@ -474,8 +452,6 @@ fn process_pairs( // SAFETY: `source < worst.len()`. unsafe { *worst_ptr.add(source) = source_worst }; } - - debug_assert_eq!(output.len(), point_count); } /// Insert one candidate into a fixed-width sorted list. diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index dc4b444036..71c9a4f8c3 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -309,29 +309,8 @@ where M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let point_count = dots.nrows(); let leader_count = dots.ncols(); - assert!( - leader_count > 0, - "validated partition input must contain leaders" - ); - if M::PARTITION_POINT_SCALE.is_some() { - assert_eq!( - scales.point_scales.len(), - point_count, - "validated point scales must match point count" - ); - } - if M::PARTITION_LEADER_SCALE.is_some() { - assert_eq!( - scales.leader_scales.len(), - leader_count, - "validated leader scales must match leader count" - ); - } let fanout = output.ncols(); - debug_assert!(fanout > 0); - debug_assert_eq!(tracker.len(), fanout); // Reset the tracker for each point. No assignment state can pass from one // output row to another. for (point, (point_dots, point_output)) in dots @@ -356,7 +335,7 @@ where // 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: the assertion above established one scale per leader, and + // SAFETY: `validate` established one scale per leader, and // `base + F::LANES <= full <= leader_count`. unsafe { F::load_simd(arch, scales.leader_scales.as_ptr().add(base)) } } else { From efa4ea60f02bc835630ff1e85457c4c06e7218d8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:41:36 +0000 Subject: [PATCH 50/80] docs(pipnn): state partition row contract --- diskann/src/graph/pipnn/partition_kernel.rs | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 71c9a4f8c3..1054631aa4 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -281,21 +281,14 @@ fn check_length( } } -/// Convert each point's leader scores into sorted leader IDs. +/// Compute the nearest leader IDs for each point. /// -/// For each point, the function does these steps: +/// The function converts each dot-product row to metric `M` scores. It writes +/// `output.ncols()` leader-column IDs in nearest-first order. Equal scores keep +/// leader-column order. NaN and positive infinity do not enter the tracker. /// -/// 1. Convert the point norm to the unit that `M` requires. -/// 2. Process complete SIMD groups. -/// 3. Process the scalar tail. -/// 4. Check that the tracker is full. -/// 5. Copy the sorted leader IDs to the output row. -/// -/// `tracker` has `fanout` entries and stays sorted after each insertion. Strict -/// comparisons keep leader scan order for equal scores. They do not rank NaN. -/// -/// The caller allocates the runtime-sized tracker once and reuses it for all -/// point rows. +/// `tracker.len()` must equal `output.ncols()`. The function resets the tracker +/// for each point and returns an error if rankable scores do not fill it. fn process_points( arch: F::Arch, dots: MatrixView<'_, f32>, From 89a750c623585a36ab71178dc498348c9de346b7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:52:14 +0000 Subject: [PATCH 51/80] refactor(pipnn): name kernel responsibilities --- diskann/src/graph/pipnn/leaf_kernel.rs | 47 +++++++++------------ diskann/src/graph/pipnn/partition_kernel.rs | 9 +++- 2 files changed, 26 insertions(+), 30 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index f536cf7660..4b07b73d14 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -171,7 +171,7 @@ where output.as_mut_slice().fill(LeafNeighbor::default()); workspace.worst.fill(f32::INFINITY); - process_neighbor_width::( + dispatch_neighbor_count::( arch, input, neighbor_count, @@ -272,11 +272,11 @@ fn resize( Ok(()) } -/// Select fixed-width storage for the validated neighbor count. +/// Dispatch a runtime neighbor count to a fixed-width leaf scan. /// -/// This match runs once for each leaf. `as_chunks_mut` converts the output once. -/// Candidate insertion does not convert slices to arrays. -fn process_neighbor_width( +/// Valid counts are one, two, and three. Any other count returns +/// [`LeafKernelError::InvalidNeighborCount`]. +fn dispatch_neighbor_count( arch: F::Arch, input: MatrixView<'_, f32>, neighbor_count: usize, @@ -291,9 +291,9 @@ where u64: From<<::BitMask as SIMDMask>::Underlying>, { 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), + 1 => run_neighbor_count::(arch, input, output, norms, worst), + 2 => run_neighbor_count::(arch, input, output, norms, worst), + 3 => run_neighbor_count::(arch, input, output, norms, worst), _ => { return Err(LeafKernelError::InvalidNeighborCount { points: input.nrows(), @@ -305,11 +305,11 @@ where Ok(()) } -/// Split the validated output into one fixed array for each source. +/// Run leaf selection with compile-time neighbor count `N`. /// -/// `N` is one, two, or three. The function performs one safe split for each -/// leaf. It then starts the shared pair traversal. -fn process_fixed_width( +/// The function views the flat output as one `[LeafNeighbor; N]` row per point. +/// It then scans all point pairs. +fn run_neighbor_count( arch: F::Arch, input: MatrixView<'_, f32>, output: &mut [LeafNeighbor], @@ -322,27 +322,18 @@ fn process_fixed_width( u64: From<<::BitMask as SIMDMask>::Underlying>, { let (neighbor_lists, _) = output.as_chunks_mut::(); - process_pairs::(arch, input, neighbor_lists, norms, worst); + scan_point_pairs::(arch, input, neighbor_lists, norms, worst); } -/// Scan the strict lower triangle and update both points of each pair. +/// Select neighbors from all unordered point pairs in one leaf. /// -/// Entry conditions: +/// The function reads the strict lower triangle once. It offers each distance to +/// both endpoint lists. SIMD groups and the scalar tail preserve pair scan order. /// -/// - `input` is a square row-major matrix. -/// - `output` has one sorted list for each point. -/// - `worst[source]` equals the distance in the last output slot. -/// - `norms` has one value per point when metric `M` requires norms. -/// -/// Each SIMD chunk computes both eligibility masks before it changes output. -/// Several lanes can update the current source. Each such lane checks the current -/// threshold again. Each target lane updates a different earlier source. -/// -/// The scalar tail uses the scalar formula for `M`. This keeps its specified -/// rounding order. The function evaluates exactly `n(n - 1) / 2` pairs. It -/// inserts accepted candidates in scan order. +/// `input` must be square. `output` and `worst` must have one row per point. +/// `norms` must have one value per point when metric `M` requires norms. #[inline(never)] -fn process_pairs( +fn scan_point_pairs( arch: F::Arch, input: MatrixView<'_, f32>, output: &mut [[LeafNeighbor; N]], diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 1054631aa4..efe5fdd610 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -140,6 +140,11 @@ pub enum PartitionKernelError { /// /// `output.nrows()` must equal `input.dots.nrows()`. `output.ncols()` sets the /// fanout and must not exceed the leader count. +/// +/// # Errors +/// +/// Returns an error for an invalid shape, scale input, fanout, or allocation. +/// It also returns an error when fewer than `fanout` scores are rankable. pub(crate) fn nearest_leaders( arch: A, input: PartitionInput<'_>, @@ -160,7 +165,7 @@ where } workspace.prepare(fanout)?; - process_points::(arch, input.dots, scales, output, &mut workspace.tracker) + select_point_leaders::(arch, input.dots, scales, output, &mut workspace.tracker) } /// Checked norm slices in the storage form that `M` requires. @@ -289,7 +294,7 @@ fn check_length( /// /// `tracker.len()` must equal `output.ncols()`. The function resets the tracker /// for each point and returns an error if rankable scores do not fill it. -fn process_points( +fn select_point_leaders( arch: F::Arch, dots: MatrixView<'_, f32>, scales: ScaleSlices<'_>, From a6729df13dcc1cbb0f7b75c4243edfbdf92dd628 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:58:05 +0000 Subject: [PATCH 52/80] docs(pipnn): define partition leader semantics --- diskann/src/graph/pipnn/leaf_kernel.rs | 10 +-- diskann/src/graph/pipnn/partition_kernel.rs | 89 ++++++++++----------- 2 files changed, 46 insertions(+), 53 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 4b07b73d14..aabe7666e3 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -140,9 +140,9 @@ pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result( /// The function reads the strict lower triangle once. It offers each distance to /// both endpoint lists. SIMD groups and the scalar tail preserve pair scan order. /// -/// `input` must be square. `output` and `worst` must have one row per point. -/// `norms` must have one value per point when metric `M` requires norms. +/// `input` supplies square dot products. `output` and `worst` contain one row +/// for each point. `norms` contains one value per point when `M` uses norms. #[inline(never)] fn scan_point_pairs( arch: F::Arch, diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index efe5fdd610..7c29783b06 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -3,20 +3,19 @@ * Licensed under the MIT license. */ -//! Nearest-leader selection for PiPNN partition assignment. +//! Select partition centers for PiPNN point assignment. //! -//! The input contains a row-major point-to-leader dot matrix and -//! metric-specific [`PartitionScales`]. The output contains sorted leader-column -//! positions for each point. Its width sets the fanout and cannot exceed the -//! leader count. +//! A leader is a sampled dataset point that represents one child partition. +//! Each input row contains dot products from one assigned point to all sampled +//! leaders. Each output row contains the nearest leader-column IDs. The scatter +//! step uses each column ID as a child-partition ID. //! //! The caller supplies concrete architecture `A` and metric `M`. The function //! checks row counts, scale units, scale lengths, fanout, and leader-ID range. //! These checks occur before output changes or unchecked SIMD loads. //! -//! L2 omits the point norm because it is constant for one point. Strict -//! comparisons keep leader scan order for equal scores. They do not rank NaN. -//! One runtime-sized workspace tracks the nearest leaders for each point. +//! L2 omits the assigned point's norm because it is constant across all sampled +//! leaders. Equal scores keep sampled-leader order. NaN is not rankable. use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -26,7 +25,9 @@ use diskann_wide::{ use super::kernel_metric::{KernelMetric, ScaleKind}; -/// Reusable nearest-leader tracker for one partition worker. +/// Reusable nearest-center state for one partition worker. +/// +/// Each entry contains a sampled leader's matrix-column ID and its metric score. #[derive(Debug, Default)] pub struct PartitionKernelWorkspace { tracker: Vec<(u32, f32)>, @@ -43,36 +44,35 @@ impl PartitionKernelWorkspace { } } -/// Metric-specific norm inputs for one partition tile. +/// Metric norms for one point-to-leader tile. /// -/// The kernel checks each slice length before it changes output. Cosine point -/// values are squared norms from a matrix diagonal. Cosine leader values are -/// norms that partition setup computes once. +/// Cosine point values are squared norms for the points being assigned. Cosine +/// leader values are norms for the sampled partition centers. #[derive(Clone, Copy, Debug)] pub enum PartitionScales<'a> { - /// L2 needs only squared leader norms; the point norm cannot affect ranking. + /// L2 uses the squared norm of each sampled partition center. L2 { - /// Squared norm for every leader column. + /// Squared norm for each sampled leader. leader_squared_norms: &'a [f32], }, - /// Unnormalized cosine needs squared point norms and leader norms. + /// Unnormalized cosine uses norms for assigned points and sampled leaders. Cosine { /// Squared norm for every point. point_squared_norms: &'a [f32], - /// Norm for every leader column. + /// Norm for each sampled leader. leader_norms: &'a [f32], }, /// Normalized cosine and inner product need no normalization inputs. None, } -/// One row-major point-to-leader dot-product tile. +/// Dot products between assigned points and sampled partition centers. /// -/// Rows are points and columns are leaders. [`Self::scales`] must match concrete -/// metric `M`. This value borrows all input and stores no kernel state. +/// Each row is one point being assigned. Each column is one sampled leader. +/// [`Self::scales`] supplies the norm layout for metric `M`. #[derive(Clone, Copy, Debug)] pub struct PartitionInput<'a> { - /// One point per matrix row and one leader per column. + /// One assigned point per row and one sampled leader per column. pub dots: MatrixView<'a, f32>, /// Normalization inputs matching concrete metric `M`. pub scales: PartitionScales<'a>, @@ -136,10 +136,10 @@ pub enum PartitionKernelError { }, } -/// Select the nearest leader positions for each input point. +/// Select the nearest sampled partition centers for each input point. /// -/// `output.nrows()` must equal `input.dots.nrows()`. `output.ncols()` sets the -/// fanout and must not exceed the leader count. +/// The output width is the fanout. Each output value is a leader's column ID in +/// `input.dots`. Partition scatter uses that ID to select a child partition. /// /// # Errors /// @@ -286,14 +286,14 @@ fn check_length( } } -/// Compute the nearest leader IDs for each point. +/// Rank sampled partition centers for each assigned point. /// -/// The function converts each dot-product row to metric `M` scores. It writes -/// `output.ncols()` leader-column IDs in nearest-first order. Equal scores keep -/// leader-column order. NaN and positive infinity do not enter the tracker. +/// The function converts point-to-leader dot products to metric `M` scores. It +/// keeps the nearest `output.ncols()` centers in sampled-leader order for ties. +/// NaN and positive infinity are not rankable. /// -/// `tracker.len()` must equal `output.ncols()`. The function resets the tracker -/// for each point and returns an error if rankable scores do not fill it. +/// `tracker` stores the retained center-column IDs and scores for the current +/// point. The function resets this state before it processes another point. fn select_point_leaders( arch: F::Arch, dots: MatrixView<'_, f32>, @@ -363,20 +363,17 @@ where if tracker[fanout - 1].0 == u32::MAX { return Err(PartitionKernelError::InsufficientRankableLeaders { point, fanout }); } - // Distances stay in the workspace. Partition construction needs only the - // leader-column positions in nearest-first order. + // Scatter needs the sampled-center column IDs. Metric scores remain in + // the worker workspace. copy_leader_ids(tracker, point_output); } Ok(()) } -/// Insert competitive SIMD lanes in increasing leader order. -/// -/// One broadcast comparison rejects a group that cannot improve the last slot. -/// Bit iteration proceeds from low lane to high lane. This order matches scalar -/// tie behavior for all SIMD widths. +/// Offer one SIMD group of sampled centers to the current point's tracker. /// -/// `distances` starts at `first_leader`. `tracker` is sorted. +/// `first_leader` is the matrix-column ID of the first lane. Lanes enter in +/// sampled-leader order, which preserves tie order. fn insert_leader_lanes(distances: F, first_leader: usize, tracker: &mut [(u32, f32)]) where F: SIMDVector> + SIMDPartialOrd, @@ -397,14 +394,11 @@ where } } -/// Insert one better candidate into a sorted tracker. -/// -/// The function replaces the last slot and moves the new value to the left. -/// Equal scores and NaN do not enter. Thus, scan order resolves equal scores. -/// The last slot is also the rejection threshold and underfill sentinel. +/// Insert one sampled partition center into the current point's retained set. /// -/// `tracker` must be sorted and non-empty. `leader` is a local column position. -/// One insertion moves at most `fanout - 1` entries. +/// `leader` is the center's column ID in the point-to-leader matrix. `tracker` +/// stores retained centers in nearest-first order. Equal scores and NaN do not +/// enter, so sampled-leader order resolves ties. #[inline(always)] fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, distance: f32) { let threshold = tracker.len() - 1; @@ -420,10 +414,9 @@ fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, distance: f32) { } } -/// Copy the retained leader IDs to one output row. +/// Write retained center-column IDs for partition scatter. /// -/// `assignments.len()` equals the checked fanout. The tracker keeps its distances -/// for the underfill check. +/// `assignments` is the current point's fanout-sized output. fn copy_leader_ids(tracker: &[(u32, f32)], assignments: &mut [u32]) { for (destination, &(leader, _)) in assignments.iter_mut().zip(tracker) { *destination = leader; From 2c15a6544f27c0df2b0947d20d5ddcd8f00f82f1 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:28:20 +0000 Subject: [PATCH 53/80] docs(pipnn): state norm conversion behavior --- diskann/src/graph/pipnn/kernel_metric.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index eae0a2806f..eb82fa83a1 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -30,7 +30,7 @@ pub(crate) enum ScaleKind { None, /// Stored value is already a squared norm. SquaredNorm, - /// Stored value is a squared norm that must become a norm. + /// Stored value is a squared norm; the kernel takes its square root. NormFromSquared, /// Stored value is already a norm. Norm, From 22787fd2e5416a32b3a36efce8964ed853dc60c1 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:31:18 +0000 Subject: [PATCH 54/80] refactor(pipnn): name fixed-width leaf scan --- diskann/src/graph/pipnn/leaf_kernel.rs | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index aabe7666e3..a5d53eff3e 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -291,9 +291,9 @@ where u64: From<<::BitMask as SIMDMask>::Underlying>, { match neighbor_count { - 1 => run_neighbor_count::(arch, input, output, norms, worst), - 2 => run_neighbor_count::(arch, input, output, norms, worst), - 3 => run_neighbor_count::(arch, input, output, norms, worst), + 1 => scan_pairs_for_neighbor_count::(arch, input, output, norms, worst), + 2 => scan_pairs_for_neighbor_count::(arch, input, output, norms, worst), + 3 => scan_pairs_for_neighbor_count::(arch, input, output, norms, worst), _ => { return Err(LeafKernelError::InvalidNeighborCount { points: input.nrows(), @@ -305,11 +305,8 @@ where Ok(()) } -/// Run leaf selection with compile-time neighbor count `N`. -/// -/// The function views the flat output as one `[LeafNeighbor; N]` row per point. -/// It then scans all point pairs. -fn run_neighbor_count( +/// Scan all leaf point pairs while retaining `N` neighbors for each point. +fn scan_pairs_for_neighbor_count( arch: F::Arch, input: MatrixView<'_, f32>, output: &mut [LeafNeighbor], From 64963c518afe392db81b0484ce3b96f1471488d3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:34:39 +0000 Subject: [PATCH 55/80] docs(pipnn): remove type-restatement comments --- diskann/src/graph/pipnn/leaf_kernel.rs | 17 +--------- diskann/src/graph/pipnn/partition_kernel.rs | 35 +++------------------ 2 files changed, 5 insertions(+), 47 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index a5d53eff3e..4a158c0092 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -69,46 +69,31 @@ pub enum LeafKernelError { TooManyPoints(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, - }, + NonSquareDots { rows: usize, cols: usize }, /// 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: usize, - /// Supplied row count. actual: usize, - /// Supplied neighbor columns. columns: usize, }, /// A source requests more neighbors than the leaf or fixed kernel supports. #[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}")] Allocation { - /// Name of the temporary buffer. buffer: &'static str, - /// Additional element capacity requested. additional: usize, }, /// 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 source position in the leaf. source_index: usize, - /// Required number of non-self neighbors. neighbors: usize, }, } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 7c29783b06..e1503d4649 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -72,9 +72,7 @@ pub enum PartitionScales<'a> { /// [`Self::scales`] supplies the norm layout for metric `M`. #[derive(Clone, Copy, Debug)] pub struct PartitionInput<'a> { - /// One assigned point per row and one sampled leader per column. pub dots: MatrixView<'a, f32>, - /// Normalization inputs matching concrete metric `M`. pub scales: PartitionScales<'a>, } @@ -86,54 +84,32 @@ pub enum PartitionKernelError { "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 scale buffer. buffer: &'static str, - /// Required length. expected: usize, - /// Supplied length. actual: usize, }, /// Scale inputs do not match concrete metric `M`. #[error("partition scales do not match selected {expected} metric")] - InvalidScales { - /// Expected scale layout. - expected: &'static str, - }, + InvalidScales { expected: &'static str }, /// The requested fanout exceeds the available leader count. #[error("invalid fanout {fanout}: must not exceed {leader_count} leaders")] - InvalidFanout { - /// Requested number of leaders per point. - fanout: usize, - /// Available leader count. - leader_count: usize, - }, + InvalidFanout { fanout: usize, leader_count: usize }, /// Reusable tracker storage could not be reserved. #[error("failed to reserve {additional} partition tracker entries")] - Allocation { - /// Additional entries requested from the allocator. - additional: usize, - }, + Allocation { additional: usize }, /// Leader positions cannot be represented as `u32`. #[error("leader count {0} exceeds the u32 position limit")] TooManyLeaders(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, - }, + InsufficientRankableLeaders { point: usize, fanout: usize }, } /// Select the nearest sampled partition centers for each input point. @@ -265,7 +241,6 @@ fn validate<'a, M: KernelMetric>( Ok(scales) } -/// Return the required norm-slice length for one concrete metric. const fn expected_scale_len(kind: ScaleKind, count: usize) -> usize { if kind.is_some() { count } else { 0 } } @@ -415,8 +390,6 @@ fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, distance: f32) { } /// Write retained center-column IDs for partition scatter. -/// -/// `assignments` is the current point's fanout-sized output. fn copy_leader_ids(tracker: &[(u32, f32)], assignments: &mut [u32]) { for (destination, &(leader, _)) in assignments.iter_mut().zip(tracker) { *destination = leader; From c069460c47363f189d9d5bc7b164fd4ba0df7ced Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:38:19 +0000 Subject: [PATCH 56/80] docs(pipnn): describe leaf insertion by domain role --- diskann/src/graph/pipnn/leaf_kernel.rs | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 4a158c0092..89d2d46ab1 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -427,14 +427,11 @@ fn scan_point_pairs( } } -/// Insert one candidate into a fixed-width sorted list. +/// Insert one target point into a source point's retained neighbor set. /// -/// Width `N` is one, two, or three. The caller has already proved that the -/// candidate is better than the last slot. Strict comparisons keep scan order -/// for equal distances. The eligibility test has already rejected NaN. -/// -/// The function returns the new last-slot distance. An unsupported `N` returns -/// the underfill sentinel. +/// `N` is the configured leaf neighbor count. The candidate is closer than the +/// current farthest neighbor. Equal distances keep pair scan order. The function +/// returns the new farthest retained distance. #[inline(always)] fn insert_fixed_neighbor( neighbors: &mut [LeafNeighbor; N], From 185db44c9b9fb8706e73a81a6aabf89aa313fbda Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:46:44 +0000 Subject: [PATCH 57/80] refactor(pipnn): name kernel test helpers --- diskann/src/graph/pipnn/leaf_kernel.rs | 33 +++++++++++---------- diskann/src/graph/pipnn/partition_kernel.rs | 30 +++++++++---------- 2 files changed, 32 insertions(+), 31 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 89d2d46ab1..4817183879 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -761,7 +761,7 @@ mod integration_tests { output } - fn run_kernel( + fn run_leaf_kernel( dots: &[f32], points: usize, requested_k: usize, @@ -791,7 +791,7 @@ mod integration_tests { let dots = differential_dots(metric, points); for requested_k in [1, 2, 3] { let expected = brute_force_reference(&dots, points, requested_k, metric); - let actual = run_kernel(&dots, points, requested_k, metric).1; + let actual = run_leaf_kernel(&dots, points, requested_k, metric).1; assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); } } @@ -809,7 +809,7 @@ mod integration_tests { ]; assert_eq!( - run_kernel(&dots, 4, 2, Metric::L2).1, + run_leaf_kernel(&dots, 4, 2, Metric::L2).1, [ LeafNeighbor::new(1, 1.0), LeafNeighbor::new(2, 1.0), @@ -837,7 +837,7 @@ mod integration_tests { (Metric::CosineNormalized, [1, 2, 1]), (Metric::InnerProduct, [1, 2, 1]), ] { - let positions: Vec<_> = run_kernel(&dots, 3, 1, metric) + let positions: Vec<_> = run_leaf_kernel(&dots, 3, 1, metric) .1 .iter() .map(|neighbor| neighbor.target) @@ -855,7 +855,7 @@ mod integration_tests { 0.0, 0.0, 1.0, ]; - let output = run_kernel(&dots, 3, 2, Metric::Cosine).1; + let output = run_leaf_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)); } @@ -865,34 +865,34 @@ mod integration_tests { #[rustfmt::skip] let out_of_range = [1.0, 0.0, 2.0, 1.0]; assert_eq!( - run_kernel(&out_of_range, 2, 1, Metric::L2).1[0].distance, + run_leaf_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, + run_leaf_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, + run_leaf_kernel(&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_kernel(&opposite, 2, 1, Metric::Cosine).1[0].distance, + run_leaf_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_kernel(&subnormal, 2, 1, Metric::Cosine).1[0].distance, + run_leaf_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_kernel(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, + run_leaf_kernel(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, 0.0 ); } @@ -903,7 +903,8 @@ mod integration_tests { let mut dots = vec![0.0; points * points]; dots[3 * points] = -f32::MAX; - let (leaf_k, output) = run_kernel(&dots, points, MAX_LEAF_NEIGHBORS, Metric::InnerProduct); + let (leaf_k, output) = + run_leaf_kernel(&dots, points, MAX_LEAF_NEIGHBORS, Metric::InnerProduct); assert_eq!(leaf_k, MAX_LEAF_NEIGHBORS); assert_eq!( output[3 * leaf_k + leaf_k - 1], @@ -926,7 +927,7 @@ mod integration_tests { Metric::CosineNormalized, Metric::InnerProduct, ] { - let output = run_kernel(&dots, 3, 1, metric).1; + let output = run_leaf_kernel(&dots, 3, 1, metric).1; assert_eq!(output[0].target, 2, "metric {metric:?}"); assert_eq!(output[1].target, 2, "metric {metric:?}"); } @@ -961,7 +962,7 @@ mod integration_tests { 0.0, 1.0, 3.0, 0.0, 0.0, 1.0, ]; - let (leaf_k, output) = run_kernel(&dots, 3, MAX_LEAF_NEIGHBORS, Metric::L2); + let (leaf_k, output) = run_leaf_kernel(&dots, 3, MAX_LEAF_NEIGHBORS, Metric::L2); assert_eq!(leaf_k, 2); for (source, neighbors) in output.chunks_exact(leaf_k).enumerate() { @@ -980,7 +981,7 @@ mod integration_tests { (&[4.0][..], 1, 2, Metric::Cosine), (&[1.0, 0.0, 0.0, 1.0][..], 2, 0, Metric::InnerProduct), ] { - assert_eq!(run_kernel(dots, points, requested_k, metric).0, 0); + assert_eq!(run_leaf_kernel(dots, points, requested_k, metric).0, 0); } } @@ -1055,7 +1056,7 @@ mod integration_tests { dots[source * points + source] = f32::NAN; } - let output = run_kernel(&dots, points, 1, Metric::Cosine).1; + let output = run_leaf_kernel(&dots, points, 1, Metric::Cosine).1; for (source, neighbor) in output.iter().enumerate().skip(1) { assert_eq!( *neighbor, diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index e1503d4649..4bbbe141b0 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -751,7 +751,7 @@ mod integration_tests { (dots, point_scales, leader_scales) } - fn run( + fn run_partition_kernel( metric: Metric, input: PartitionInput<'_>, fanout: usize, @@ -789,7 +789,7 @@ mod integration_tests { continue; } assert_eq!( - run(metric, input, fanout).unwrap(), + run_partition_kernel(metric, input, fanout).unwrap(), brute_force_reference(input, fanout, metric), "{metric:?}, leaders={leader_count}, k={fanout}" ); @@ -808,7 +808,7 @@ mod integration_tests { let norms = [0.0, 1.0, 4.0, 9.0]; assert_eq!( - run( + run_partition_kernel( Metric::L2, test_input(Metric::L2, &dots, 2, 4, &[], &norms), 2 @@ -837,7 +837,7 @@ mod integration_tests { (Metric::InnerProduct, &[][..], &[][..], [0, 1, 1, 0]), ] { assert_eq!( - run( + run_partition_kernel( metric, test_input(metric, &dots, 2, 3, point_scales, leader_scales), 2, @@ -852,7 +852,7 @@ mod integration_tests { #[test] fn cosine_treats_a_zero_norm_as_zero_similarity() { assert_eq!( - run( + run_partition_kernel( Metric::Cosine, test_input(Metric::Cosine, &[100.0, -100.0], 1, 2, &[0.0], &[1.0, 1.0]), 2, @@ -867,7 +867,7 @@ mod integration_tests { let mut dots = [0.0; 8]; dots[7] = -f32::MAX; assert_eq!( - run( + run_partition_kernel( Metric::InnerProduct, test_input(Metric::InnerProduct, &dots, 1, 8, &[], &[]), 8 @@ -880,7 +880,7 @@ mod integration_tests { #[test] fn ignores_nan_distances_without_displacing_finite_leaders() { assert_eq!( - run( + run_partition_kernel( Metric::InnerProduct, test_input(Metric::InnerProduct, &[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), 2, @@ -893,7 +893,7 @@ mod integration_tests { #[test] fn rejects_points_with_too_few_rankable_leaders() { assert_eq!( - run( + run_partition_kernel( Metric::InnerProduct, test_input(Metric::InnerProduct, &[f32::NAN, 3.0], 1, 2, &[], &[]), 2, @@ -907,19 +907,19 @@ mod integration_tests { #[test] fn accepts_empty_points_zero_fanout_and_largest_leader_id() { - run( + run_partition_kernel( Metric::InnerProduct, test_input(Metric::InnerProduct, &[], 0, 3, &[], &[]), 2, ) .unwrap(); - run( + run_partition_kernel( Metric::InnerProduct, test_input(Metric::InnerProduct, &[1.0, 2.0, 3.0], 1, 3, &[], &[]), 0, ) .unwrap(); - run( + run_partition_kernel( Metric::InnerProduct, test_input(Metric::InnerProduct, &[], 0, u32::MAX as usize, &[], &[]), 0, @@ -928,7 +928,7 @@ mod integration_tests { #[cfg(target_pointer_width = "64")] assert_eq!( - run( + run_partition_kernel( Metric::InnerProduct, test_input( Metric::InnerProduct, @@ -968,12 +968,12 @@ mod integration_tests { scales: PartitionScales::None, }; assert_eq!( - run(Metric::L2, wrong_scales, 2), + run_partition_kernel(Metric::L2, wrong_scales, 2), Err(PartitionKernelError::InvalidScales { expected: "L2" }) ); assert_eq!( - run(Metric::InnerProduct, valid_input, 4), + run_partition_kernel(Metric::InnerProduct, valid_input, 4), Err(PartitionKernelError::InvalidFanout { fanout: 4, leader_count: 3, @@ -982,7 +982,7 @@ mod integration_tests { let one = [0.0]; assert_eq!( - run( + run_partition_kernel( Metric::InnerProduct, test_input(Metric::InnerProduct, &one, 1, 1, &[], &[]), 2, From 578ab9f6022b1469a54b330bddd2748c3e18b8e6 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 05:48:55 +0000 Subject: [PATCH 58/80] refactor(pipnn): inline fixed-width leaf dispatch --- diskann/src/graph/pipnn/leaf_kernel.rs | 91 +++++++++----------------- 1 file changed, 32 insertions(+), 59 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 4817183879..9eeb1018c5 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -156,14 +156,36 @@ where output.as_mut_slice().fill(LeafNeighbor::default()); workspace.worst.fill(f32::INFINITY); - dispatch_neighbor_count::( - arch, - input, - neighbor_count, - output.as_mut_slice(), - &workspace.norms, - &mut workspace.worst, - )?; + match neighbor_count { + 1 => scan_point_pairs::( + arch, + input, + output.as_mut_slice(), + &workspace.norms, + &mut workspace.worst, + ), + 2 => scan_point_pairs::( + arch, + input, + output.as_mut_slice(), + &workspace.norms, + &mut workspace.worst, + ), + 3 => scan_point_pairs::( + arch, + input, + output.as_mut_slice(), + &workspace.norms, + &mut workspace.worst, + ), + _ => { + return Err(LeafKernelError::InvalidNeighborCount { + points: input.nrows(), + neighbors: neighbor_count, + maximum: MAX_LEAF_NEIGHBORS, + }); + } + } if let Some(source) = output .as_slice() .chunks_exact(neighbor_count) @@ -257,56 +279,6 @@ fn resize( Ok(()) } -/// Dispatch a runtime neighbor count to a fixed-width leaf scan. -/// -/// Valid counts are one, two, and three. Any other count returns -/// [`LeafKernelError::InvalidNeighborCount`]. -fn dispatch_neighbor_count( - arch: F::Arch, - input: MatrixView<'_, f32>, - neighbor_count: usize, - output: &mut [LeafNeighbor], - norms: &[f32], - worst: &mut [f32], -) -> Result<(), LeafKernelError> -where - F: SIMDVector> + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, - M: KernelMetric, - u64: From<<::BitMask as SIMDMask>::Underlying>, -{ - match neighbor_count { - 1 => scan_pairs_for_neighbor_count::(arch, input, output, norms, worst), - 2 => scan_pairs_for_neighbor_count::(arch, input, output, norms, worst), - 3 => scan_pairs_for_neighbor_count::(arch, input, output, norms, worst), - _ => { - return Err(LeafKernelError::InvalidNeighborCount { - points: input.nrows(), - neighbors: neighbor_count, - maximum: MAX_LEAF_NEIGHBORS, - }); - } - } - Ok(()) -} - -/// Scan all leaf point pairs while retaining `N` neighbors for each point. -fn scan_pairs_for_neighbor_count( - arch: F::Arch, - input: MatrixView<'_, f32>, - 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>, -{ - let (neighbor_lists, _) = output.as_chunks_mut::(); - scan_point_pairs::(arch, input, neighbor_lists, norms, worst); -} - /// Select neighbors from all unordered point pairs in one leaf. /// /// The function reads the strict lower triangle once. It offers each distance to @@ -318,7 +290,7 @@ fn scan_pairs_for_neighbor_count( fn scan_point_pairs( arch: F::Arch, input: MatrixView<'_, f32>, - output: &mut [[LeafNeighbor; N]], + output: &mut [LeafNeighbor], norms: &[f32], worst: &mut [f32], ) where @@ -327,6 +299,7 @@ fn scan_point_pairs( M: KernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { + let (output, _) = output.as_chunks_mut::(); let point_count = input.nrows(); let dots = input.as_slice(); let uses_norms = M::LEAF_SCALE.is_some(); From 3f54a317b76935e2345e525391ee74ac17f6d0be Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:03:58 +0000 Subject: [PATCH 59/80] refactor(pipnn): inline leader ID publication --- diskann/src/graph/pipnn/partition_kernel.rs | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 4bbbe141b0..142c14597d 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -340,7 +340,9 @@ where } // Scatter needs the sampled-center column IDs. Metric scores remain in // the worker workspace. - copy_leader_ids(tracker, point_output); + for (destination, &(leader, _)) in point_output.iter_mut().zip(tracker.iter()) { + *destination = leader; + } } Ok(()) } @@ -389,13 +391,6 @@ fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, distance: f32) { } } -/// Write retained center-column IDs for partition scatter. -fn copy_leader_ids(tracker: &[(u32, f32)], assignments: &mut [u32]) { - for (destination, &(leader, _)) in assignments.iter_mut().zip(tracker) { - *destination = leader; - } -} - #[cfg(test)] struct DispatchedPartitionCall<'a> { input: PartitionInput<'a>, @@ -534,7 +529,9 @@ mod tests { M::partition_distance_scalar(dot, point_scale, leader_scale), ); } - copy_leader_ids(&tracker, point_output); + for (destination, &(leader, _)) in point_output.iter_mut().zip(&tracker) { + *destination = leader; + } } } From 70ce8e2a82fcf505b4777314effdbd5f18f5fba2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:32:43 +0000 Subject: [PATCH 60/80] refactor(matrix): preserve constructor behavior --- diskann-utils/src/views.rs | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs index a20b8b5fcf..1f0cfdd4d3 100644 --- a/diskann-utils/src/views.rs +++ b/diskann-utils/src/views.rs @@ -195,14 +195,8 @@ impl MatrixBase> { where U: Generator, { - let len = nrows.checked_mul(ncols); - assert!( - len.is_some(), - "matrix shape {nrows} x {ncols} overflows usize" - ); - let len = len.unwrap_or(0); - let data: Box<[T]> = (0..len).map(|_| generator.generate()).collect(); - debug_assert_eq!(data.len(), len); + let data: Box<[T]> = (0..nrows * ncols).map(|_| generator.generate()).collect(); + debug_assert_eq!(data.len(), nrows * ncols); Self { data, nrows, ncols } } } @@ -1436,8 +1430,6 @@ mod tests { assert_eq!(m.nrows(), 5); assert_eq!(m.ncols(), 1); assert!(m.as_slice().iter().all(|&x| x == 9)); - - assert!(std::panic::catch_unwind(|| Matrix::new(0, usize::MAX, 2)).is_err()); } #[test] From 47c20890a4da9b5cbc1b9df526782e10739d0f25 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:44:13 +0000 Subject: [PATCH 61/80] refactor(pipnn): split stage metric contracts --- diskann/src/graph/pipnn/kernel_metric.rs | 384 +++--------------- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 110 +++++ .../graph/pipnn/kernel_metric/partition.rs | 108 +++++ diskann/src/graph/pipnn/leaf_kernel.rs | 60 +-- diskann/src/graph/pipnn/mod.rs | 6 +- diskann/src/graph/pipnn/partition_kernel.rs | 379 +++++++++-------- 6 files changed, 495 insertions(+), 552 deletions(-) create mode 100644 diskann/src/graph/pipnn/kernel_metric/leaf.rs create mode 100644 diskann/src/graph/pipnn/kernel_metric/partition.rs diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index eb82fa83a1..3464b71296 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -3,171 +3,71 @@ * Licensed under the MIT license. */ -//! Metric formulas for the PiPNN partition and leaf kernels. +//! Shared metric definitions for the PiPNN numerical kernels. //! -//! `build_graph` maps each runtime [`Metric`] to one zero-sized marker type. The -//! partition and leaf functions receive that concrete type. -//! -//! Each formula returns an ascending score. L2 uses squared norms. Cosine uses -//! norms and maps a zero norm to zero similarity. Normalized cosine and inner -//! product do not use norms. Ordered comparisons do not rank NaN. -//! -//! The L2 partition SIMD path uses fused arithmetic. Its scalar tail uses -//! non-fused arithmetic. This operation order is part of the tie-order contract. -//! -//! [`ScaleKind`] defines the stored norm unit. [`KernelMetric`] defines the leaf -//! and partition formulas. +//! `build_graph` maps each runtime [`Metric`] to one marker type. Leaf and +//! partition kernels use separate traits for that marker. Both traits use the +//! common cosine and norm functions in this module. + +mod leaf; +mod partition; + +pub(super) use leaf::LeafKernelMetric; +pub(super) use partition::PartitionKernelMetric; use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; -/// Stored norm representation for one kernel input. -/// -/// `KernelMetric` supplies this value as an associated constant. -#[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; the kernel takes its square root. - NormFromSquared, - /// Stored value is already a norm. - Norm, +/// Identify one metric across all PiPNN build stages. +pub(super) trait MetricTag: Send + Sync + 'static { + /// Runtime metric represented by this marker. + const METRIC: Metric; } -impl ScaleKind { - /// Convert a stored norm to the unit that the kernel requires. - /// - /// A squared norm below `f32::MIN_POSITIVE` becomes zero. A norm below - /// `sqrt(f32::MIN_POSITIVE)` also becomes zero. The function does not change - /// NaN, so the kernel does not rank it. - /// - /// The concrete [`KernelMetric`] supplies `self` as an associated constant. - /// The compiler selects one match arm for each kernel instance. - #[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() - } - } - } - } +/// Squared-L2 marker. +pub(super) struct L2; +/// Unnormalized-cosine marker. +pub(super) struct Cosine; +/// Unit-normalized-cosine marker. +pub(super) struct CosineNormalized; +/// Negative-inner-product marker. +pub(super) struct InnerProduct; - /// Return `true` when the metric requires this norm input. - pub(crate) const fn is_some(self) -> bool { - !matches!(self, Self::None) - } +impl MetricTag for L2 { + const METRIC: Metric = Metric::L2; } -/// Metric contract for leaf and partition selection. -/// -/// `build_graph` selects one implementation. Generic calls inline its arithmetic -/// through the complete build. Leaf and partition formulas are separate because -/// L2 partition ranking does not need the point norm. -/// -/// Each method returns an ascending score. Strict comparisons do not rank NaN. -/// Marker types contain no data. Associated constants remove unused norm work. -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 scale representation used by partition assignment. - const PARTITION_POINT_SCALE: ScaleKind; - /// Leader-column scale representation used by partition assignment. - const PARTITION_LEADER_SCALE: ScaleKind; - - /// Compute SIMD distances from one leaf source to earlier targets. - /// - /// `dot` and `target_scale` contain one target per lane. `source_scale` - /// contains the source norm in each lane. A metric without norms receives - /// zero for both norm arguments. The result contains one 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; - - /// Compute the scalar-tail equivalent of `leaf_distance`. - /// - /// The inputs and result represent one SIMD lane. Each implementation - /// documents any required operation order. - fn leaf_distance_scalar(dot: f32, source_scale: f32, target_scale: f32) -> f32; - - /// Compute SIMD scores from one point to a group of leaders. - /// - /// `dot` and `leader_scale` contain one leader per lane. `point_scale` - /// contains the point norm in each lane. A metric without a norm receives - /// zero for that argument. The formula can omit terms that are constant for - /// all 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; - - /// Compute the scalar-tail equivalent of `partition_distance`. - /// - /// The inputs and result represent one SIMD lane. Each implementation uses - /// its documented operation order. - fn partition_distance_scalar(dot: f32, point_scale: f32, leader_scale: f32) -> f32; +impl MetricTag for Cosine { + const METRIC: Metric = Metric::Cosine; } -/// 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; +impl MetricTag for CosineNormalized { + const METRIC: Metric = Metric::CosineNormalized; +} -/// Clamp negative SIMD roundoff to zero and keep NaN lanes unchanged. -/// -/// SIMD `max` has architecture-specific NaN behavior. The ordered self-test -/// selects the original value for each NaN lane. -#[inline(always)] -fn clamp_nonnegative(arch: F::Arch, distance: F) -> F -where - F: SIMDVector + SIMDFloat, - F::Mask: SIMDSelect, -{ - let zero = F::default(arch); - // Select the original value for NaN lanes. This gives all architectures the - // same non-rankable NaN result. - distance - .eq_simd(distance) - .select(zero.max_simd(distance), distance) +impl MetricTag for InnerProduct { + const METRIC: Metric = Metric::InnerProduct; } -/// Scalar equivalent of [`clamp_nonnegative`]. +/// Convert a squared norm to a norm. +/// +/// The function maps subnormal squared norms to zero. It preserves NaN so that +/// kernel comparisons do not rank an invalid value. #[inline(always)] -fn clamp_nonnegative_scalar(distance: f32) -> f32 { - if distance < 0.0 { 0.0 } else { distance } +pub(super) fn norm_from_squared(squared_norm: f32) -> f32 { + if squared_norm < f32::MIN_POSITIVE { + 0.0 + } else { + squared_norm.sqrt() + } } /// Compute cosine distance with the DiskANN zero-norm and NaN rules. /// -/// A zero-norm lane divides by one and then selects zero similarity. A NaN norm -/// propagates through division. If the other norm is zero, zero similarity takes -/// precedence. -/// -/// Each input contains one point pair per lane. The result is -/// `1 - cosine_similarity`. The function uses no lane branch. +/// Each lane contains one point pair. A zero norm produces zero similarity. A +/// NaN norm remains NaN unless the other norm is zero. #[inline(always)] -fn cosine_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F +pub(super) 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, @@ -183,8 +83,9 @@ where one - cosine } +/// Compute scalar cosine distance with the DiskANN zero-norm rules. #[inline(always)] -fn cosine_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32 { +pub(super) 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 { @@ -192,201 +93,18 @@ fn cosine_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32 { } } -impl KernelMetric for L2 { - const METRIC: Metric = Metric::L2; - const LEAF_SCALE: ScaleKind = ScaleKind::SquaredNorm; - const PARTITION_POINT_SCALE: ScaleKind = ScaleKind::None; - const PARTITION_LEADER_SCALE: ScaleKind = ScaleKind::SquaredNorm; - - #[inline(always)] - 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, - { - // 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, - ) - } - - #[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) - } - - #[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, - { - // The point norm is constant for this ranking. The SIMD path uses the - // fused multiply-add operation that defines its tie order. - F::splat(arch, -2.0).mul_add_simd(dot, leader_scale) - } - - #[inline(always)] - fn partition_distance_scalar(dot: f32, _: f32, leader_scale: f32) -> f32 { - // The scalar tail uses non-fused subtraction. A fused operation can - // change rounding and select a different leader at a tie. - leader_scale - 2.0 * dot - } -} - -impl KernelMetric for Cosine { - const METRIC: Metric = Metric::Cosine; - const LEAF_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, source_scale: F, target_scale: F) -> F - where - 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)) - } - - #[inline(always)] - 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, - { - // Partitioning uses only score order. Do not clamp the score because a - // clamp can change the order of near ties. - cosine_distance(arch, dot, point_scale, leader_scale) - } - - #[inline(always)] - fn partition_distance_scalar(dot: f32, point_scale: f32, leader_scale: f32) -> f32 { - // Use the same unclamped score in the scalar tail. - 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_POINT_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, - { - // 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) - } - - #[inline(always)] - fn partition_distance(arch: F::Arch, dot: F, _: F, _: F) -> F - where - 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 { - // Use the same unclamped score as the SIMD path. - 1.0 - dot - } -} - -impl KernelMetric for InnerProduct { - const METRIC: Metric = Metric::InnerProduct; - const LEAF_SCALE: ScaleKind = ScaleKind::None; - const PARTITION_POINT_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, - { - // 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 - } - - #[inline(always)] - fn partition_distance(arch: F::Arch, dot: F, _: F, _: F) -> F - where - 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 - } -} - #[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()); + fn norm_from_squared_applies_zero_threshold_without_erasing_nan() { + assert_eq!(norm_from_squared(-0.0).to_bits(), 0.0f32.to_bits()); + assert_eq!(norm_from_squared(f32::MIN_POSITIVE / 2.0), 0.0); 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), + norm_from_squared(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); + assert!(norm_from_squared(f32::NAN).is_nan()); } } diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs new file mode 100644 index 0000000000..81bf8cd134 --- /dev/null +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -0,0 +1,110 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Metric formulas for leaf-local neighbor selection. + +use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; + +use super::{ + Cosine, CosineNormalized, InnerProduct, L2, MetricTag, cosine_distance, cosine_distance_scalar, +}; + +/// Metric contract for leaf-local neighbor selection. +/// +/// Each function returns an ascending distance. The leaf kernel supplies squared +/// norms to L2 and norms to cosine. Dot-only metrics receive zero norm values. +pub(in super::super) trait LeafKernelMetric: MetricTag { + /// Compute SIMD distances from one source to earlier leaf targets. + fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect; + + /// Compute one scalar-tail leaf distance. + fn leaf_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32; +} + +/// Clamp negative SIMD roundoff to zero and preserve NaN lanes. +#[inline(always)] +fn clamp_nonnegative(arch: F::Arch, distance: F) -> F +where + F: SIMDVector + SIMDFloat, + F::Mask: SIMDSelect, +{ + let zero = F::default(arch); + distance + .eq_simd(distance) + .select(zero.max_simd(distance), distance) +} + +/// Clamp negative scalar roundoff to zero and preserve NaN. +#[inline(always)] +fn clamp_nonnegative_scalar(distance: f32) -> f32 { + if distance < 0.0 { 0.0 } else { distance } +} + +impl LeafKernelMetric for L2 { + #[inline(always)] + fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + clamp_nonnegative(arch, source_norm + target_norm - F::splat(arch, 2.0) * dot) + } + + #[inline(always)] + fn leaf_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32 { + clamp_nonnegative_scalar(source_norm + target_norm - 2.0 * dot) + } +} + +impl LeafKernelMetric for Cosine { + #[inline(always)] + fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + clamp_nonnegative(arch, cosine_distance(arch, dot, source_norm, target_norm)) + } + + #[inline(always)] + fn leaf_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32 { + clamp_nonnegative_scalar(cosine_distance_scalar(dot, source_norm, target_norm)) + } +} + +impl LeafKernelMetric for CosineNormalized { + #[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) + } +} + +impl LeafKernelMetric for InnerProduct { + #[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 + } +} diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs new file mode 100644 index 0000000000..f71c377462 --- /dev/null +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -0,0 +1,108 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Metric formulas for partition-leader ranking. + +use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; + +use super::{ + Cosine, CosineNormalized, InnerProduct, L2, MetricTag, cosine_distance, cosine_distance_scalar, +}; + +/// Metric contract for partition-leader ranking. +/// +/// Each function returns an ascending score. L2 receives a squared leader norm. +/// Cosine receives point and leader norms. Dot-only metrics receive zero norms. +pub(in super::super) trait PartitionKernelMetric: MetricTag { + /// Compute SIMD ranking scores for one point and multiple leaders. + fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect; + + /// Compute one scalar-tail partition ranking score. + fn partition_ranking_scalar(dot: f32, point_norm: f32, leader_norm: f32) -> f32; +} + +impl PartitionKernelMetric for L2 { + #[inline(always)] + fn partition_ranking(arch: F::Arch, dot: F, _: F, leader_norm: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + // The point norm is constant for all leaders. Fused arithmetic defines + // the ranking order for complete SIMD groups. + F::splat(arch, -2.0).mul_add_simd(dot, leader_norm) + } + + #[inline(always)] + fn partition_ranking_scalar(dot: f32, _: f32, leader_norm: f32) -> f32 { + // Non-fused arithmetic defines the ranking order for the scalar tail. + leader_norm - 2.0 * dot + } +} + +impl PartitionKernelMetric for Cosine { + #[inline(always)] + fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F + where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + { + cosine_distance(arch, dot, point_norm, leader_norm) + } + + #[inline(always)] + fn partition_ranking_scalar(dot: f32, point_norm: f32, leader_norm: f32) -> f32 { + cosine_distance_scalar(dot, point_norm, leader_norm) + } +} + +impl PartitionKernelMetric for CosineNormalized { + #[inline(always)] + fn partition_ranking(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_ranking_scalar(dot: f32, _: f32, _: f32) -> f32 { + 1.0 - dot + } +} + +impl PartitionKernelMetric for InnerProduct { + #[inline(always)] + fn partition_ranking(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_ranking_scalar(dot: f32, _: f32, _: f32) -> f32 { + -dot + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn l2_scalar_ranking_preserves_non_fused_rounding() { + let scalar = L2::partition_ranking_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/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 9eeb1018c5..76d127b51e 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -22,20 +22,21 @@ //! rejection thresholds. use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_vector::distance::Metric; use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; -use super::kernel_metric::KernelMetric; +use super::kernel_metric::{LeafKernelMetric, MetricTag, norm_from_squared}; /// Largest leaf-local neighbor count supported by the fixed insertion kernel. -pub const MAX_LEAF_NEIGHBORS: usize = 3; +pub(super) const MAX_LEAF_NEIGHBORS: usize = 3; /// One leaf-local neighbor and its metric distance. #[derive(Clone, Copy, Debug, PartialEq)] -pub struct LeafNeighbor { +pub(super) struct LeafNeighbor { /// Target position in the leaf, not a dataset ID. - pub target: u32, + pub(super) target: u32, /// Distance from the source point to `target`. - pub distance: f32, + pub(super) distance: f32, } impl LeafNeighbor { @@ -43,7 +44,7 @@ impl LeafNeighbor { /// /// `target` is a position in the leaf. `distance` is its score relative to /// the source of the output row. - pub const fn new(target: u32, distance: f32) -> Self { + pub(super) const fn new(target: u32, distance: f32) -> Self { Self { target, distance } } } @@ -56,14 +57,14 @@ impl Default for LeafNeighbor { /// Reusable temporary storage for leaf top-k selection. #[derive(Debug, Default)] -pub struct LeafKernelWorkspace { +pub(super) struct LeafKernelWorkspace { norms: Vec, worst: Vec, } /// Validation or allocation error returned by [`nearest_neighbors`]. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] -pub enum LeafKernelError { +pub(super) enum LeafKernelError { /// The point count cannot be represented in leaf-local `u32` positions. #[error("point count {0} exceeds the u32 position limit")] TooManyPoints(usize), @@ -109,7 +110,10 @@ pub enum LeafKernelError { /// Returns [`LeafKernelError::TooManyPoints`] when leaf-local positions cannot /// fit in `u32`, or [`LeafKernelError::InvalidNeighborCount`] when `requested_k` /// exceeds [`MAX_LEAF_NEIGHBORS`]. -pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result { +pub(super) fn leaf_neighbor_count( + points: usize, + requested_k: usize, +) -> Result { if points > u32::MAX as usize { return Err(LeafKernelError::TooManyPoints(points)); } @@ -133,7 +137,7 @@ pub fn leaf_neighbor_count(points: usize, requested_k: usize) -> Result( +pub(super) fn nearest_neighbors( arch: A, input: MatrixView<'_, f32>, mut output: MutMatrixView<'_, LeafNeighbor>, @@ -143,7 +147,7 @@ where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, - M: KernelMetric, + M: LeafKernelMetric, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { validate(input, &output)?; @@ -244,18 +248,24 @@ fn validate( /// value to a norm. Normalized cosine and inner product clear the norm buffer. /// /// An allocation error occurs before SIMD traversal. -fn prepare_workspace( +fn prepare_workspace( input: MatrixView<'_, f32>, workspace: &mut LeafKernelWorkspace, ) -> Result<(), LeafKernelError> { let points = input.nrows(); - if M::LEAF_SCALE.is_some() { - resize("norms", &mut workspace.norms, points, 0.0)?; - for (source, norm) in workspace.norms.iter_mut().enumerate() { - *norm = M::LEAF_SCALE.transform(input[(source, source)]); + match ::METRIC { + Metric::L2 | Metric::Cosine => { + resize("norms", &mut workspace.norms, points, 0.0)?; + for (source, norm) in workspace.norms.iter_mut().enumerate() { + let squared_norm = input[(source, source)]; + *norm = if ::METRIC == Metric::Cosine { + norm_from_squared(squared_norm) + } else { + squared_norm + }; + } } - } else { - workspace.norms.clear(); + Metric::CosineNormalized | Metric::InnerProduct => workspace.norms.clear(), } resize( "worst distances", @@ -296,20 +306,20 @@ fn scan_point_pairs( ) where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, - M: KernelMetric, + M: LeafKernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { let (output, _) = output.as_chunks_mut::(); let point_count = input.nrows(); let dots = input.as_slice(); - let uses_norms = M::LEAF_SCALE.is_some(); + let uses_norms = matches!(::METRIC, Metric::L2 | Metric::Cosine); let worst_ptr = worst.as_mut_ptr(); // Source zero has no earlier target. Each source after zero can still add // itself to the neighbor list of source zero. for source in 1..point_count { let source_start = source * point_count; - let source_scale = if uses_norms { + let source_norm = if uses_norms { F::splat(arch, norms[source]) } else { F::default(arch) @@ -322,14 +332,14 @@ fn scan_point_pairs( 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 { + let target_norms = if uses_norms { // SAFETY: the full target chunk is below `source < point_count`. // `prepare_workspace` created one norm for each point. unsafe { F::load_simd(arch, norms.as_ptr().add(target)) } } else { F::default(arch) }; - let distances = M::leaf_distance(arch, pair_dots, source_scale, target_scales); + let distances = M::leaf_distance(arch, pair_dots, source_norm, target_norms); // 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)); @@ -376,13 +386,13 @@ fn scan_point_pairs( 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 { + let (source_norm, target_norm) = 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, source_scale, target_scale); + let distance = M::leaf_distance_scalar(dot, source_norm, target_norm); if distance < source_worst { source_worst = insert_fixed_neighbor(&mut output[source], target as u32, distance); } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 1c51887d81..6da4dd1e14 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -13,13 +13,13 @@ //! pair once and updates both points. Each point retains at most three local //! neighbors. //! -//! `kernel_metric` defines the scalar and SIMD formulas. It also defines the -//! required norm units for each metric. +//! `kernel_metric` defines metric markers and shared math. Separate leaf and +//! partition traits define each kernel's scalar and SIMD formulas. //! //! The graph builder selects architecture `A` and metric `M` once. It passes //! these concrete types to both kernels. //! -//! Each kernel checks all view and scale relationships before unchecked SIMD +//! Each kernel checks all view and norm relationships before unchecked SIMD //! access. The kernels borrow their matrices. They write only to caller-owned //! output and workspace. #[allow(dead_code)] diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 142c14597d..77de376900 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -11,7 +11,7 @@ //! step uses each column ID as a child-partition ID. //! //! The caller supplies concrete architecture `A` and metric `M`. The function -//! checks row counts, scale units, scale lengths, fanout, and leader-ID range. +//! checks row counts, norm units, norm lengths, fanout, and leader-ID range. //! These checks occur before output changes or unchecked SIMD loads. //! //! L2 omits the assigned point's norm because it is constant across all sampled @@ -23,13 +23,13 @@ use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, }; -use super::kernel_metric::{KernelMetric, ScaleKind}; +use super::kernel_metric::{MetricTag, PartitionKernelMetric, norm_from_squared}; /// Reusable nearest-center state for one partition worker. /// /// Each entry contains a sampled leader's matrix-column ID and its metric score. #[derive(Debug, Default)] -pub struct PartitionKernelWorkspace { +pub(super) struct PartitionKernelWorkspace { tracker: Vec<(u32, f32)>, } @@ -44,12 +44,11 @@ impl PartitionKernelWorkspace { } } -/// Metric norms for one point-to-leader tile. +/// Norm values for one point-to-leader tile. /// -/// Cosine point values are squared norms for the points being assigned. Cosine -/// leader values are norms for the sampled partition centers. +/// Cosine point values are squared norms. Cosine leader values are norms. #[derive(Clone, Copy, Debug)] -pub enum PartitionScales<'a> { +pub(super) enum PartitionNorms<'a> { /// L2 uses the squared norm of each sampled partition center. L2 { /// Squared norm for each sampled leader. @@ -69,16 +68,16 @@ pub enum PartitionScales<'a> { /// Dot products between assigned points and sampled partition centers. /// /// Each row is one point being assigned. Each column is one sampled leader. -/// [`Self::scales`] supplies the norm layout for metric `M`. +/// [`Self::norms`] supplies the norm layout for metric `M`. #[derive(Clone, Copy, Debug)] -pub struct PartitionInput<'a> { - pub dots: MatrixView<'a, f32>, - pub scales: PartitionScales<'a>, +pub(super) struct PartitionInput<'a> { + pub(super) dots: MatrixView<'a, f32>, + pub(super) norms: PartitionNorms<'a>, } /// Validation error returned by [`nearest_leaders`]. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] -pub enum PartitionKernelError { +pub(super) enum PartitionKernelError { /// 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" @@ -88,16 +87,16 @@ pub enum PartitionKernelError { actual_rows: usize, actual_cols: usize, }, - /// A metric-specific scale slice has the wrong length. + /// A metric-specific norm slice has the wrong length. #[error("invalid {buffer} length: expected {expected}, got {actual}")] InvalidBufferLength { buffer: &'static str, expected: usize, actual: usize, }, - /// Scale inputs do not match concrete metric `M`. - #[error("partition scales do not match selected {expected} metric")] - InvalidScales { expected: &'static str }, + /// Norm inputs do not match concrete metric `M`. + #[error("partition norms do not match selected {expected} metric")] + InvalidNorms { expected: &'static str }, /// The requested fanout exceeds the available leader count. #[error("invalid fanout {fanout}: must not exceed {leader_count} leaders")] InvalidFanout { fanout: usize, leader_count: usize }, @@ -119,9 +118,9 @@ pub enum PartitionKernelError { /// /// # Errors /// -/// Returns an error for an invalid shape, scale input, fanout, or allocation. +/// Returns an error for an invalid shape, norm input, fanout, or allocation. /// It also returns an error when fewer than `fanout` scores are rankable. -pub(crate) fn nearest_leaders( +pub(super) fn nearest_leaders( arch: A, input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, @@ -132,37 +131,33 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, { - let scales = validate::(input, &output)?; + let norms = validate::(input, &output)?; let fanout = output.ncols(); if fanout == 0 || input.dots.nrows() == 0 { return Ok(()); } workspace.prepare(fanout)?; - select_point_leaders::(arch, input.dots, scales, output, &mut workspace.tracker) + select_point_leaders::(arch, input.dots, norms, output, &mut workspace.tracker) } -/// Checked norm slices in the storage form that `M` requires. -/// -/// A metric that does not use a norm receives an empty slice. Associated -/// `ScaleKind` constants remove these branches from concrete metric loops. +/// Checked norm slices for one concrete metric. #[derive(Clone, Copy)] -struct ScaleSlices<'a> { - point_scales: &'a [f32], - leader_scales: &'a [f32], +struct PartitionNormSlices<'a> { + point_squared_norms: &'a [f32], + leader_norm_values: &'a [f32], } /// Check the safety and metric conditions for partition selection. /// -/// The matrix views already prove their backing lengths. This function checks -/// row counts, leader-ID range, fanout, scale variant, and scale lengths. A -/// successful result contains norm slices in the units that `M` requires. -fn validate<'a, M: KernelMetric>( +/// Matrix views prove their backing lengths. This function checks row counts, +/// leader-ID range, fanout, norm variant, and norm lengths. +fn validate<'a, M: PartitionKernelMetric>( input: PartitionInput<'a>, output: &MutMatrixView<'_, u32>, -) -> Result, PartitionKernelError> { +) -> Result, PartitionKernelError> { let point_count = input.dots.nrows(); let leader_count = input.dots.ncols(); let fanout = output.ncols(); @@ -184,65 +179,48 @@ fn validate<'a, M: KernelMetric>( }); } - // Match the scale variant to concrete metric `M` before extracting slices. - // This prevents use of a squared point norm as a leader norm. - let scales = match (M::METRIC, input.scales) { + match (::METRIC, input.norms) { ( Metric::L2, - PartitionScales::L2 { + PartitionNorms::L2 { leader_squared_norms, }, - ) => ScaleSlices { - point_scales: &[], - leader_scales: leader_squared_norms, - }, + ) => { + check_length("leader norms", leader_squared_norms.len(), leader_count)?; + Ok(PartitionNormSlices { + point_squared_norms: &[], + leader_norm_values: leader_squared_norms, + }) + } ( Metric::Cosine, - PartitionScales::Cosine { + PartitionNorms::Cosine { point_squared_norms, leader_norms, }, - ) => ScaleSlices { - point_scales: point_squared_norms, - leader_scales: leader_norms, - }, - (Metric::CosineNormalized | Metric::InnerProduct, PartitionScales::None) => ScaleSlices { - point_scales: &[], - leader_scales: &[], - }, - (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", - }); + ) => { + check_length("point norms", point_squared_norms.len(), point_count)?; + check_length("leader norms", leader_norms.len(), leader_count)?; + Ok(PartitionNormSlices { + point_squared_norms, + leader_norm_values: leader_norms, + }) } - (Metric::InnerProduct, _) => { - return Err(PartitionKernelError::InvalidScales { - expected: "inner product", - }); + (Metric::CosineNormalized | Metric::InnerProduct, PartitionNorms::None) => { + Ok(PartitionNormSlices { + point_squared_norms: &[], + leader_norm_values: &[], + }) } - }; - - // The associated scale kinds define the exact slice lengths. A metric that - // does not use a scale must receive an empty slice. - check_length( - "point scales", - scales.point_scales.len(), - expected_scale_len(M::PARTITION_POINT_SCALE, point_count), - )?; - check_length( - "leader scales", - scales.leader_scales.len(), - expected_scale_len(M::PARTITION_LEADER_SCALE, leader_count), - )?; - Ok(scales) -} - -const fn expected_scale_len(kind: ScaleKind, count: usize) -> usize { - if kind.is_some() { count } else { 0 } + (Metric::L2, _) => Err(PartitionKernelError::InvalidNorms { expected: "L2" }), + (Metric::Cosine, _) => Err(PartitionKernelError::InvalidNorms { expected: "cosine" }), + (Metric::CosineNormalized, _) => Err(PartitionKernelError::InvalidNorms { + expected: "normalized cosine", + }), + (Metric::InnerProduct, _) => Err(PartitionKernelError::InvalidNorms { + expected: "inner product", + }), + } } fn check_length( @@ -272,18 +250,21 @@ fn check_length( fn select_point_leaders( arch: F::Arch, dots: MatrixView<'_, f32>, - scales: ScaleSlices<'_>, + norms: PartitionNormSlices<'_>, mut output: MutMatrixView<'_, u32>, tracker: &mut [(u32, f32)], ) -> Result<(), PartitionKernelError> where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, - M: KernelMetric, + M: PartitionKernelMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { let leader_count = dots.ncols(); let fanout = output.ncols(); + let metric = ::METRIC; + let uses_point_norm = metric == Metric::Cosine; + let uses_leader_norm = matches!(metric, Metric::L2 | Metric::Cosine); // Reset the tracker for each point. No assignment state can pass from one // output row to another. for (point, (point_dots, point_output)) in dots @@ -292,14 +273,12 @@ where .enumerate() { tracker.fill((u32::MAX, f32::INFINITY)); - // Convert the point norm once for this row. Metrics without a point norm - // use zero. - let point_scale = if M::PARTITION_POINT_SCALE.is_some() { - M::PARTITION_POINT_SCALE.transform(scales.point_scales[point]) + let point_norm = if uses_point_norm { + norm_from_squared(norms.point_squared_norms[point]) } else { 0.0 }; - let point_scale_vector = F::splat(arch, point_scale); + let point_norm_vector = F::splat(arch, point_norm); // Process all complete SIMD groups first. The scalar tail uses the // metric's scalar operation order. let full = leader_count / F::LANES * F::LANES; @@ -307,15 +286,15 @@ where for base in (0..full).step_by(F::LANES) { // 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: `validate` established one scale per leader, and + let leader_norms = if uses_leader_norm { + // SAFETY: `validate` established one norm value per leader. // `base + F::LANES <= full <= leader_count`. - unsafe { F::load_simd(arch, scales.leader_scales.as_ptr().add(base)) } + unsafe { F::load_simd(arch, norms.leader_norm_values.as_ptr().add(base)) } } else { F::default(arch) }; insert_leader_lanes( - M::partition_distance(arch, point_dots, point_scale_vector, leader_scales), + M::partition_ranking(arch, point_dots, point_norm_vector, leader_norms), base, tracker, ); @@ -324,15 +303,15 @@ where // Use scalar formulas for the tail. A padded SIMD load can read past the // norm slice and can change 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]) + let leader_norm = if uses_leader_norm { + norms.leader_norm_values[leader] } else { 0.0 }; insert_leader( tracker, leader as u32, - M::partition_distance_scalar(dot, point_scale, leader_scale), + M::partition_ranking_scalar(dot, point_norm, leader_norm), ); } if tracker[fanout - 1].0 == u32::MAX { @@ -351,18 +330,18 @@ where /// /// `first_leader` is the matrix-column ID of the first lane. Lanes enter in /// sampled-leader order, which preserves tie order. -fn insert_leader_lanes(distances: F, first_leader: usize, tracker: &mut [(u32, f32)]) +fn insert_leader_lanes(scores: F, first_leader: usize, tracker: &mut [(u32, f32)]) where F: SIMDVector> + SIMDPartialOrd, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let threshold = F::splat(distances.arch(), tracker[tracker.len() - 1].1); - let eligible = distances.lt_simd(threshold); + let threshold = F::splat(scores.arch(), tracker[tracker.len() - 1].1); + let eligible = scores.lt_simd(threshold); if eligible.none() { return; } - let values: [f32; 16] = distances.to_array(); + let values: [f32; 16] = scores.to_array(); let mut lanes = u64::from(eligible.bitmask().to_underlying()); while lanes != 0 { let lane = lanes.trailing_zeros() as usize; @@ -377,13 +356,13 @@ where /// stores retained centers in nearest-first order. Equal scores and NaN do not /// enter, so sampled-leader order resolves ties. #[inline(always)] -fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, distance: f32) { +fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, score: f32) { let threshold = tracker.len() - 1; - if distance.partial_cmp(&tracker[threshold].1) != Some(std::cmp::Ordering::Less) { + if score.partial_cmp(&tracker[threshold].1) != Some(std::cmp::Ordering::Less) { return; } - tracker[threshold] = (leader, distance); + tracker[threshold] = (leader, score); let mut slot = threshold; while slot > 0 && tracker[slot].1 < tracker[slot - 1].1 { tracker.swap(slot, slot - 1); @@ -451,7 +430,9 @@ fn dispatch_nearest_leaders( #[cfg(test)] mod tests { - use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; + use super::super::kernel_metric::{ + Cosine, CosineNormalized, InnerProduct, L2, PartitionKernelMetric, + }; use super::*; @@ -460,73 +441,74 @@ mod tests { dots: &'a [f32], point_count: usize, leader_count: usize, - point_scales: &'a [f32], - leader_scales: &'a [f32], + point_squared_norms: &'a [f32], + leader_norm_values: &'a [f32], ) -> PartitionInput<'a> { - let scales = match metric { - Metric::L2 => PartitionScales::L2 { - leader_squared_norms: leader_scales, + let norms = match metric { + Metric::L2 => PartitionNorms::L2 { + leader_squared_norms: leader_norm_values, }, - Metric::Cosine => PartitionScales::Cosine { - point_squared_norms: point_scales, - leader_norms: leader_scales, + Metric::Cosine => PartitionNorms::Cosine { + point_squared_norms, + leader_norms: leader_norm_values, }, - Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, + Metric::CosineNormalized | Metric::InnerProduct => PartitionNorms::None, }; PartitionInput { dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), - scales, + norms, } } // This oracle checks SIMD chunking, scalar tails, and tracker order. It uses - // `M::partition_distance_scalar`. Separate tests define each formula directly. - fn scalar_traversal_reference( + // the scalar ranking formula for metric `M`. + fn scalar_traversal_reference( input: PartitionInput<'_>, fanout: usize, output: &mut [u32], ) { - let scales = match input.scales { - PartitionScales::L2 { + let norms = match input.norms { + PartitionNorms::L2 { leader_squared_norms, - } => ScaleSlices { - point_scales: &[], - leader_scales: leader_squared_norms, + } => PartitionNormSlices { + point_squared_norms: &[], + leader_norm_values: leader_squared_norms, }, - PartitionScales::Cosine { + PartitionNorms::Cosine { point_squared_norms, leader_norms, - } => ScaleSlices { - point_scales: point_squared_norms, - leader_scales: leader_norms, + } => PartitionNormSlices { + point_squared_norms, + leader_norm_values: leader_norms, }, - PartitionScales::None => ScaleSlices { - point_scales: &[], - leader_scales: &[], + PartitionNorms::None => PartitionNormSlices { + point_squared_norms: &[], + leader_norm_values: &[], }, }; + let metric = ::METRIC; for (point, (point_dots, point_output)) in input .dots .row_iter() .zip(output.chunks_exact_mut(fanout)) .enumerate() { - let point_scale = if M::PARTITION_POINT_SCALE.is_some() { - M::PARTITION_POINT_SCALE.transform(scales.point_scales[point]) + let point_norm = if metric == Metric::Cosine { + norm_from_squared(norms.point_squared_norms[point]) } else { 0.0 }; let mut tracker = vec![(u32::MAX, f32::INFINITY); 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.leader_scales[leader]) + let leader_norm = if matches!(metric, Metric::L2 | Metric::Cosine) { + norms.leader_norm_values[leader] } else { 0.0 }; insert_leader( &mut tracker, leader as u32, - M::partition_distance_scalar(dot, point_scale, leader_scale), + M::partition_ranking_scalar(dot, point_norm, leader_norm), ); } for (destination, &(leader, _)) in point_output.iter_mut().zip(&tracker) { @@ -536,25 +518,25 @@ mod tests { } #[test] - fn scalar_distance_matches_metric_contract() { - assert_eq!(L2::partition_distance_scalar(2.0, 0.0, 9.0), 5.0); + fn scalar_ranking_matches_metric_contract() { + assert_eq!(L2::partition_ranking_scalar(2.0, 0.0, 9.0), 5.0); assert_eq!( - CosineNormalized::partition_distance_scalar(0.25, 0.0, 0.0), + CosineNormalized::partition_ranking_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()); + assert_eq!(InnerProduct::partition_ranking_scalar(3.0, 0.0, 0.0), -3.0); + assert_eq!(Cosine::partition_ranking_scalar(4.0, 2.0, 4.0), 0.5); + assert_eq!(Cosine::partition_ranking_scalar(4.0, 0.0, 4.0), 1.0); + assert!(Cosine::partition_ranking_scalar(1.0, f32::NAN, 1.0).is_nan()); } #[test] fn cosine_special_norms_match_scalar_and_dispatched_kernel() { 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(&[ + let point_squared_norms = [0.0, f32::MIN_POSITIVE / 2.0, f32::MIN_POSITIVE, f32::NAN]; + let dots = vec![1.0; point_squared_norms.len() * leader_count]; + let mut leader_norms = vec![1.0; leader_count]; + leader_norms[..4].copy_from_slice(&[ 0.0, f32::MIN_POSITIVE.sqrt() / 2.0, f32::MIN_POSITIVE.sqrt(), @@ -563,18 +545,18 @@ mod tests { let input = test_input( Metric::Cosine, &dots, - point_scales.len(), + point_squared_norms.len(), leader_count, - &point_scales, - &leader_scales, + &point_squared_norms, + &leader_norms, ); - let mut expected = vec![u32::MAX; point_scales.len() * 2]; + let mut expected = vec![u32::MAX; point_squared_norms.len() * 2]; scalar_traversal_reference::(input, 2, &mut expected); - let mut actual = vec![u32::MAX; point_scales.len() * 2]; + let mut actual = vec![u32::MAX; point_squared_norms.len() * 2]; dispatch_nearest_leaders( Metric::Cosine, input, - MutMatrixView::try_from(actual.as_mut_slice(), point_scales.len(), 2).unwrap(), + MutMatrixView::try_from(actual.as_mut_slice(), point_squared_norms.len(), 2).unwrap(), &mut PartitionKernelWorkspace::default(), ) .unwrap(); @@ -615,8 +597,8 @@ mod tests { )] mod integration_tests { use super::{ - PartitionInput, PartitionKernelError, PartitionKernelWorkspace, PartitionScales, - dispatch_nearest_leaders, + PartitionInput, PartitionKernelError, PartitionKernelWorkspace, PartitionNorms, + dispatch_nearest_leaders, norm_from_squared, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -626,37 +608,37 @@ mod integration_tests { dots: &'a [f32], point_count: usize, leader_count: usize, - point_scales: &'a [f32], - leader_scales: &'a [f32], + point_squared_norms: &'a [f32], + leader_norm_values: &'a [f32], ) -> PartitionInput<'a> { - let scales = match metric { - Metric::L2 => PartitionScales::L2 { - leader_squared_norms: leader_scales, + let norms = match metric { + Metric::L2 => PartitionNorms::L2 { + leader_squared_norms: leader_norm_values, }, - Metric::Cosine => PartitionScales::Cosine { - point_squared_norms: point_scales, - leader_norms: leader_scales, + Metric::Cosine => PartitionNorms::Cosine { + point_squared_norms, + leader_norms: leader_norm_values, }, - Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, + Metric::CosineNormalized | Metric::InnerProduct => PartitionNorms::None, }; PartitionInput { dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), - scales, + norms, } } 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 { + let (point_squared_norms, leader_norm_values) = match input.norms { + PartitionNorms::L2 { leader_squared_norms, } => (&[][..], leader_squared_norms), - PartitionScales::Cosine { + PartitionNorms::Cosine { point_squared_norms, leader_norms, } => (point_squared_norms, leader_norms), - PartitionScales::None => (&[][..], &[][..]), + PartitionNorms::None => (&[][..], &[][..]), }; let mut assignments = vec![u32::MAX; point_count * fanout]; for (point, (point_dots, point_assignments)) in input @@ -666,31 +648,27 @@ mod integration_tests { .zip(assignments.chunks_exact_mut(fanout)) .enumerate() { - let point_scale = point_scales.get(point).copied().unwrap_or(0.0); + let point_squared_norm = point_squared_norms.get(point).copied().unwrap_or(0.0); let mut candidates: Vec<_> = point_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, + let leader_norm = leader_norm_values.get(leader).copied().unwrap_or(0.0); + let score = match metric { + Metric::L2 => leader_norm - 2.0 * dot, Metric::CosineNormalized => 1.0 - dot, Metric::InnerProduct => -dot, Metric::Cosine => { - let point_norm = if point_scale < f32::MIN_POSITIVE { + let point_norm = norm_from_squared(point_squared_norm); + 1.0 - if point_norm == 0.0 || leader_norm == 0.0 { 0.0 } else { - point_scale.sqrt() - }; - 1.0 - if point_norm == 0.0 || leader_scale == 0.0 { - 0.0 - } else { - dot / (point_norm * leader_scale) + dot / (point_norm * leader_norm) } } }; - (distance.partial_cmp(&f32::INFINITY) == Some(std::cmp::Ordering::Less)) - .then_some((leader as u32, distance)) + (score.partial_cmp(&f32::INFINITY) == Some(std::cmp::Ordering::Less)) + .then_some((leader as u32, score)) }) .collect(); candidates.sort_by(|left, right| left.1.partial_cmp(&right.1).unwrap()); @@ -716,12 +694,12 @@ mod integration_tests { } }) .collect(); - let point_scales = if metric == Metric::Cosine { + let point_squared_norms = if metric == Metric::Cosine { vec![0.0, 16.0] } else { Vec::new() }; - let leader_scales = match metric { + let leader_norm_values = match metric { Metric::Cosine => (0..leader_count) .map(|leader| { if leader == 1 { @@ -745,7 +723,7 @@ mod integration_tests { .collect(), Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), }; - (dots, point_scales, leader_scales) + (dots, point_squared_norms, leader_norm_values) } fn run_partition_kernel( @@ -772,14 +750,15 @@ mod integration_tests { Metric::InnerProduct, ] { 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 (dots, point_squared_norms, leader_norm_values) = + differential_data(metric, leader_count); let input = test_input( metric, &dots, 2, leader_count, - &point_scales, - &leader_scales, + &point_squared_norms, + &leader_norm_values, ); for fanout in [1, 2, 16, 17, 32] { if fanout >= leader_count { @@ -815,6 +794,24 @@ mod integration_tests { ); } + #[test] + fn l2_scalar_tail_can_outrank_a_fused_simd_lane() { + let mut dots = [0.0; 17]; + dots[0] = f32::MAX; + dots[16] = f32::MAX; + let leader_squared_norms = [f32::MAX; 17]; + + assert_eq!( + run_partition_kernel( + Metric::L2, + test_input(Metric::L2, &dots, 1, 17, &[], &leader_squared_norms,), + 1, + ) + .unwrap(), + [16] + ); + } + #[test] fn supports_every_partition_metric() { #[rustfmt::skip] @@ -822,7 +819,7 @@ mod integration_tests { 1.0, 0.0, -1.0, 2.0, 6.0, 0.0, ]; - for (metric, point_scales, leader_scales, expected) in [ + for (metric, point_squared_norms, leader_norm_values, expected) in [ (Metric::L2, &[][..], &[1.0, 4.0, 9.0][..], [0, 1, 1, 0]), ( Metric::Cosine, @@ -836,7 +833,7 @@ mod integration_tests { assert_eq!( run_partition_kernel( metric, - test_input(metric, &dots, 2, 3, point_scales, leader_scales), + test_input(metric, &dots, 2, 3, point_squared_norms, leader_norm_values), 2, ) .unwrap(), @@ -942,7 +939,7 @@ mod integration_tests { } #[test] - fn rejects_wrong_output_scales_and_fanout() { + fn rejects_wrong_output_norms_and_fanout() { let dots = [0.0; 6]; let valid_input = test_input(Metric::InnerProduct, &dots, 2, 3, &[], &[]); let mut wrong_output = [u32::MAX; 3]; @@ -960,13 +957,13 @@ mod integration_tests { }) ); - let wrong_scales = PartitionInput { + let wrong_norms = PartitionInput { dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), - scales: PartitionScales::None, + norms: PartitionNorms::None, }; assert_eq!( - run_partition_kernel(Metric::L2, wrong_scales, 2), - Err(PartitionKernelError::InvalidScales { expected: "L2" }) + run_partition_kernel(Metric::L2, wrong_norms, 2), + Err(PartitionKernelError::InvalidNorms { expected: "L2" }) ); assert_eq!( From 44572dd05e24d7fd16b7616a619fae464a3d7dea Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:10:15 +0000 Subject: [PATCH 62/80] perf(pipnn): specialize stage norm handling --- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 34 +++++++++++++++ .../graph/pipnn/kernel_metric/partition.rs | 41 +++++++++++++++++++ diskann/src/graph/pipnn/leaf_kernel.rs | 23 ++++------- diskann/src/graph/pipnn/partition_kernel.rs | 14 +++---- 4 files changed, 90 insertions(+), 22 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 81bf8cd134..c3376b4e23 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -16,6 +16,12 @@ use super::{ /// Each function returns an ascending distance. The leaf kernel supplies squared /// norms to L2 and norms to cosine. Dot-only metrics receive zero norm values. pub(in super::super) trait LeafKernelMetric: MetricTag { + /// True when the metric reads leaf norms. + const USES_NORMS: bool; + + /// Convert one Gram diagonal value to the norm unit for this metric. + fn prepare_norm(squared_norm: f32) -> f32; + /// Compute SIMD distances from one source to earlier leaf targets. fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -46,6 +52,13 @@ fn clamp_nonnegative_scalar(distance: f32) -> f32 { } impl LeafKernelMetric for L2 { + const USES_NORMS: bool = true; + + #[inline(always)] + fn prepare_norm(squared_norm: f32) -> f32 { + squared_norm + } + #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -62,6 +75,13 @@ impl LeafKernelMetric for L2 { } impl LeafKernelMetric for Cosine { + const USES_NORMS: bool = true; + + #[inline(always)] + fn prepare_norm(squared_norm: f32) -> f32 { + super::norm_from_squared(squared_norm) + } + #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -78,6 +98,13 @@ impl LeafKernelMetric for Cosine { } impl LeafKernelMetric for CosineNormalized { + const USES_NORMS: bool = false; + + #[inline(always)] + fn prepare_norm(_: f32) -> f32 { + 0.0 + } + #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F where @@ -94,6 +121,13 @@ impl LeafKernelMetric for CosineNormalized { } impl LeafKernelMetric for InnerProduct { + const USES_NORMS: bool = false; + + #[inline(always)] + fn prepare_norm(_: f32) -> f32 { + 0.0 + } + #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F where diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index f71c377462..b6248a96c7 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -16,6 +16,15 @@ use super::{ /// Each function returns an ascending score. L2 receives a squared leader norm. /// Cosine receives point and leader norms. Dot-only metrics receive zero norms. pub(in super::super) trait PartitionKernelMetric: MetricTag { + /// True when the metric reads a point norm. + const USES_POINT_NORM: bool; + + /// True when the metric reads leader norm values. + const USES_LEADER_NORMS: bool; + + /// Convert one squared point norm to the unit for this metric. + fn prepare_point_norm(squared_norm: f32) -> f32; + /// Compute SIMD ranking scores for one point and multiple leaders. fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F where @@ -27,6 +36,14 @@ pub(in super::super) trait PartitionKernelMetric: MetricTag { } impl PartitionKernelMetric for L2 { + const USES_POINT_NORM: bool = false; + const USES_LEADER_NORMS: bool = true; + + #[inline(always)] + fn prepare_point_norm(_: f32) -> f32 { + 0.0 + } + #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, leader_norm: F) -> F where @@ -46,6 +63,14 @@ impl PartitionKernelMetric for L2 { } impl PartitionKernelMetric for Cosine { + const USES_POINT_NORM: bool = true; + const USES_LEADER_NORMS: bool = true; + + #[inline(always)] + fn prepare_point_norm(squared_norm: f32) -> f32 { + super::norm_from_squared(squared_norm) + } + #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F where @@ -62,6 +87,14 @@ impl PartitionKernelMetric for Cosine { } impl PartitionKernelMetric for CosineNormalized { + const USES_POINT_NORM: bool = false; + const USES_LEADER_NORMS: bool = false; + + #[inline(always)] + fn prepare_point_norm(_: f32) -> f32 { + 0.0 + } + #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, _: F) -> F where @@ -78,6 +111,14 @@ impl PartitionKernelMetric for CosineNormalized { } impl PartitionKernelMetric for InnerProduct { + const USES_POINT_NORM: bool = false; + const USES_LEADER_NORMS: bool = false; + + #[inline(always)] + fn prepare_point_norm(_: f32) -> f32 { + 0.0 + } + #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, _: F) -> F where diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 76d127b51e..646d75e1b4 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -22,10 +22,9 @@ //! rejection thresholds. use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_vector::distance::Metric; use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; -use super::kernel_metric::{LeafKernelMetric, MetricTag, norm_from_squared}; +use super::kernel_metric::LeafKernelMetric; /// Largest leaf-local neighbor count supported by the fixed insertion kernel. pub(super) const MAX_LEAF_NEIGHBORS: usize = 3; @@ -253,19 +252,13 @@ fn prepare_workspace( workspace: &mut LeafKernelWorkspace, ) -> Result<(), LeafKernelError> { let points = input.nrows(); - match ::METRIC { - Metric::L2 | Metric::Cosine => { - resize("norms", &mut workspace.norms, points, 0.0)?; - for (source, norm) in workspace.norms.iter_mut().enumerate() { - let squared_norm = input[(source, source)]; - *norm = if ::METRIC == Metric::Cosine { - norm_from_squared(squared_norm) - } else { - squared_norm - }; - } + if M::USES_NORMS { + resize("norms", &mut workspace.norms, points, 0.0)?; + for (source, norm) in workspace.norms.iter_mut().enumerate() { + *norm = M::prepare_norm(input[(source, source)]); } - Metric::CosineNormalized | Metric::InnerProduct => workspace.norms.clear(), + } else { + workspace.norms.clear(); } resize( "worst distances", @@ -312,7 +305,7 @@ fn scan_point_pairs( let (output, _) = output.as_chunks_mut::(); let point_count = input.nrows(); let dots = input.as_slice(); - let uses_norms = matches!(::METRIC, Metric::L2 | Metric::Cosine); + let uses_norms = M::USES_NORMS; let worst_ptr = worst.as_mut_ptr(); // Source zero has no earlier target. Each source after zero can still add diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 77de376900..4ff495d62c 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -23,7 +23,7 @@ use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, }; -use super::kernel_metric::{MetricTag, PartitionKernelMetric, norm_from_squared}; +use super::kernel_metric::{MetricTag, PartitionKernelMetric}; /// Reusable nearest-center state for one partition worker. /// @@ -262,9 +262,8 @@ where { let leader_count = dots.ncols(); let fanout = output.ncols(); - let metric = ::METRIC; - let uses_point_norm = metric == Metric::Cosine; - let uses_leader_norm = matches!(metric, Metric::L2 | Metric::Cosine); + let uses_point_norm = M::USES_POINT_NORM; + let uses_leader_norm = M::USES_LEADER_NORMS; // Reset the tracker for each point. No assignment state can pass from one // output row to another. for (point, (point_dots, point_output)) in dots @@ -274,7 +273,7 @@ where { tracker.fill((u32::MAX, f32::INFINITY)); let point_norm = if uses_point_norm { - norm_from_squared(norms.point_squared_norms[point]) + M::prepare_point_norm(norms.point_squared_norms[point]) } else { 0.0 }; @@ -431,7 +430,7 @@ fn dispatch_nearest_leaders( #[cfg(test)] mod tests { use super::super::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, PartitionKernelMetric, + Cosine, CosineNormalized, InnerProduct, L2, PartitionKernelMetric, norm_from_squared, }; use super::*; @@ -596,9 +595,10 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod integration_tests { + use super::super::kernel_metric::norm_from_squared; use super::{ PartitionInput, PartitionKernelError, PartitionKernelWorkspace, PartitionNorms, - dispatch_nearest_leaders, norm_from_squared, + dispatch_nearest_leaders, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; From 6bc1a96321507456939b2ecfd9653a4884dd4ce3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:56:51 +0000 Subject: [PATCH 63/80] docs(pipnn): use active metric comments --- diskann/src/graph/pipnn/kernel_metric.rs | 20 ++++++++--------- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 16 +++++++------- .../graph/pipnn/kernel_metric/partition.rs | 14 ++++++------ diskann/src/graph/pipnn/partition_kernel.rs | 22 +++++++++---------- 4 files changed, 36 insertions(+), 36 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 3464b71296..fabc727e7d 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Shared metric definitions for the PiPNN numerical kernels. +//! This module defines shared metrics for the PiPNN numerical kernels. //! //! `build_graph` maps each runtime [`Metric`] to one marker type. Leaf and //! partition kernels use separate traits for that marker. Both traits use the @@ -18,19 +18,19 @@ pub(super) use partition::PartitionKernelMetric; use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; -/// Identify one metric across all PiPNN build stages. +/// This trait identifies one metric across all PiPNN build stages. pub(super) trait MetricTag: Send + Sync + 'static { - /// Runtime metric represented by this marker. + /// This value identifies the runtime metric. const METRIC: Metric; } -/// Squared-L2 marker. +/// This marker selects squared L2. pub(super) struct L2; -/// Unnormalized-cosine marker. +/// This marker selects unnormalized cosine. pub(super) struct Cosine; -/// Unit-normalized-cosine marker. +/// This marker selects unit-normalized cosine. pub(super) struct CosineNormalized; -/// Negative-inner-product marker. +/// This marker selects negative inner product. pub(super) struct InnerProduct; impl MetricTag for L2 { @@ -49,7 +49,7 @@ impl MetricTag for InnerProduct { const METRIC: Metric = Metric::InnerProduct; } -/// Convert a squared norm to a norm. +/// This function converts a squared norm to a norm. /// /// The function maps subnormal squared norms to zero. It preserves NaN so that /// kernel comparisons do not rank an invalid value. @@ -62,7 +62,7 @@ pub(super) fn norm_from_squared(squared_norm: f32) -> f32 { } } -/// Compute cosine distance with the DiskANN zero-norm and NaN rules. +/// This function computes cosine distance with the DiskANN zero-norm and NaN rules. /// /// Each lane contains one point pair. A zero norm produces zero similarity. A /// NaN norm remains NaN unless the other norm is zero. @@ -83,7 +83,7 @@ where one - cosine } -/// Compute scalar cosine distance with the DiskANN zero-norm rules. +/// This function computes scalar cosine distance with the DiskANN zero-norm rules. #[inline(always)] pub(super) 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() { diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index c3376b4e23..58c45e95fb 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Metric formulas for leaf-local neighbor selection. +//! This module defines metric formulas for leaf-local neighbor selection. use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; @@ -11,28 +11,28 @@ use super::{ Cosine, CosineNormalized, InnerProduct, L2, MetricTag, cosine_distance, cosine_distance_scalar, }; -/// Metric contract for leaf-local neighbor selection. +/// This trait defines metric operations for leaf-local neighbor selection. /// /// Each function returns an ascending distance. The leaf kernel supplies squared /// norms to L2 and norms to cosine. Dot-only metrics receive zero norm values. pub(in super::super) trait LeafKernelMetric: MetricTag { - /// True when the metric reads leaf norms. + /// This value is true when the metric reads leaf norms. const USES_NORMS: bool; - /// Convert one Gram diagonal value to the norm unit for this metric. + /// This function converts one Gram diagonal value to the required norm unit. fn prepare_norm(squared_norm: f32) -> f32; - /// Compute SIMD distances from one source to earlier leaf targets. + /// This function computes SIMD distances from one source to earlier targets. fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; - /// Compute one scalar-tail leaf distance. + /// This function computes one scalar-tail leaf distance. fn leaf_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32; } -/// Clamp negative SIMD roundoff to zero and preserve NaN lanes. +/// This function clamps negative SIMD roundoff to zero and preserves NaN lanes. #[inline(always)] fn clamp_nonnegative(arch: F::Arch, distance: F) -> F where @@ -45,7 +45,7 @@ where .select(zero.max_simd(distance), distance) } -/// Clamp negative scalar roundoff to zero and preserve NaN. +/// This function clamps negative scalar roundoff to zero and preserves NaN. #[inline(always)] fn clamp_nonnegative_scalar(distance: f32) -> f32 { if distance < 0.0 { 0.0 } else { distance } diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index b6248a96c7..9ef6b36008 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! Metric formulas for partition-leader ranking. +//! This module defines metric formulas for partition-leader ranking. use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; @@ -11,27 +11,27 @@ use super::{ Cosine, CosineNormalized, InnerProduct, L2, MetricTag, cosine_distance, cosine_distance_scalar, }; -/// Metric contract for partition-leader ranking. +/// This trait defines metric operations for partition-leader ranking. /// /// Each function returns an ascending score. L2 receives a squared leader norm. /// Cosine receives point and leader norms. Dot-only metrics receive zero norms. pub(in super::super) trait PartitionKernelMetric: MetricTag { - /// True when the metric reads a point norm. + /// This value is true when the metric reads a point norm. const USES_POINT_NORM: bool; - /// True when the metric reads leader norm values. + /// This value is true when the metric reads leader norm values. const USES_LEADER_NORMS: bool; - /// Convert one squared point norm to the unit for this metric. + /// This function converts one squared point norm to the required unit. fn prepare_point_norm(squared_norm: f32) -> f32; - /// Compute SIMD ranking scores for one point and multiple leaders. + /// This function computes SIMD ranking scores for one point and multiple leaders. fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; - /// Compute one scalar-tail partition ranking score. + /// This function computes one scalar-tail partition ranking score. fn partition_ranking_scalar(dot: f32, point_norm: f32, leader_norm: f32) -> f32; } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 4ff495d62c..09572597d0 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -44,24 +44,24 @@ impl PartitionKernelWorkspace { } } -/// Norm values for one point-to-leader tile. +/// This enum stores norm values for one point-to-leader tile. /// -/// Cosine point values are squared norms. Cosine leader values are norms. +/// Cosine points use squared norms. Cosine leaders use norms. #[derive(Clone, Copy, Debug)] pub(super) enum PartitionNorms<'a> { - /// L2 uses the squared norm of each sampled partition center. + /// This variant stores squared norms for L2 leaders. L2 { - /// Squared norm for each sampled leader. + /// This slice contains one squared norm for each sampled leader. leader_squared_norms: &'a [f32], }, - /// Unnormalized cosine uses norms for assigned points and sampled leaders. + /// This variant stores norms for unnormalized cosine. Cosine { - /// Squared norm for every point. + /// This slice contains one squared norm for each point. point_squared_norms: &'a [f32], - /// Norm for each sampled leader. + /// This slice contains one norm for each sampled leader. leader_norms: &'a [f32], }, - /// Normalized cosine and inner product need no normalization inputs. + /// This variant provides no norms for normalized cosine or inner product. None, } @@ -118,7 +118,7 @@ pub(super) enum PartitionKernelError { /// /// # Errors /// -/// Returns an error for an invalid shape, norm input, fanout, or allocation. +/// The function returns an error for an invalid shape, norm input, fanout, or allocation. /// It also returns an error when fewer than `fanout` scores are rankable. pub(super) fn nearest_leaders( arch: A, @@ -143,14 +143,14 @@ where select_point_leaders::(arch, input.dots, norms, output, &mut workspace.tracker) } -/// Checked norm slices for one concrete metric. +/// This structure stores checked norm slices for one concrete metric. #[derive(Clone, Copy)] struct PartitionNormSlices<'a> { point_squared_norms: &'a [f32], leader_norm_values: &'a [f32], } -/// Check the safety and metric conditions for partition selection. +/// This function checks safety and metric conditions for partition selection. /// /// Matrix views prove their backing lengths. This function checks row counts, /// leader-ID range, fanout, norm variant, and norm lengths. From 640cd7d6e4b12bd0e51c06ede47c51e957180c7b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:01:31 +0000 Subject: [PATCH 64/80] test(pipnn): use norm-specific terminology --- diskann/src/graph/pipnn/leaf_kernel.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 646d75e1b4..44d7ebc8a6 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -637,7 +637,7 @@ mod integration_tests { 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 DOT_FACTOR: f32 = 1.0 / 32.0; const TIED_TARGETS: [usize; 2] = [1, 2]; fn differential_dots(metric: Metric, points: usize) -> Vec { @@ -657,7 +657,7 @@ mod integration_tests { dots[source * points + target] = if TIED_TARGETS.contains(&target) { 0.5 } else { - pair * DOT_SCALE + pair * DOT_FACTOR }; } } From 552c48178a8aee09c4f1a48f6112a7a1341ac76c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:33:47 +0000 Subject: [PATCH 65/80] perf(pipnn): bound leaf SIMD prefix once --- diskann/src/graph/pipnn/leaf_kernel.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 44d7ebc8a6..07429cfe19 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -321,8 +321,9 @@ fn scan_point_pairs( // `source < point_count == worst.len()`. let mut source_worst = unsafe { *worst_ptr.add(source) }; let mut target = 0; + let full = source / F::LANES * F::LANES; - while target + F::LANES <= source { + while target < full { // 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_norms = if uses_norms { From 9fa980e72859f832fcad4a2d27ebea8a85585eca Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:39:28 +0000 Subject: [PATCH 66/80] refactor(pipnn): remove metric tag indirection --- diskann/src/graph/pipnn/kernel_metric.rs | 33 +------------------ diskann/src/graph/pipnn/kernel_metric/leaf.rs | 15 +++------ .../graph/pipnn/kernel_metric/partition.rs | 23 ++++++------- diskann/src/graph/pipnn/partition_kernel.rs | 6 ++-- 4 files changed, 18 insertions(+), 59 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index fabc727e7d..1642b037cf 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -3,11 +3,7 @@ * Licensed under the MIT license. */ -//! This module defines shared metrics for the PiPNN numerical kernels. -//! -//! `build_graph` maps each runtime [`Metric`] to one marker type. Leaf and -//! partition kernels use separate traits for that marker. Both traits use the -//! common cosine and norm functions in this module. +//! This module defines metric markers and shared numerical functions. mod leaf; mod partition; @@ -15,40 +11,13 @@ mod partition; pub(super) use leaf::LeafKernelMetric; pub(super) use partition::PartitionKernelMetric; -use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; -/// This trait identifies one metric across all PiPNN build stages. -pub(super) trait MetricTag: Send + Sync + 'static { - /// This value identifies the runtime metric. - const METRIC: Metric; -} - -/// This marker selects squared L2. pub(super) struct L2; -/// This marker selects unnormalized cosine. pub(super) struct Cosine; -/// This marker selects unit-normalized cosine. pub(super) struct CosineNormalized; -/// This marker selects negative inner product. pub(super) struct InnerProduct; -impl MetricTag for L2 { - const METRIC: Metric = Metric::L2; -} - -impl MetricTag for Cosine { - const METRIC: Metric = Metric::Cosine; -} - -impl MetricTag for CosineNormalized { - const METRIC: Metric = Metric::CosineNormalized; -} - -impl MetricTag for InnerProduct { - const METRIC: Metric = Metric::InnerProduct; -} - /// This function converts a squared norm to a norm. /// /// The function maps subnormal squared norms to zero. It preserves NaN so that diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 58c45e95fb..6390aa1ea7 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -7,28 +7,21 @@ use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; -use super::{ - Cosine, CosineNormalized, InnerProduct, L2, MetricTag, cosine_distance, cosine_distance_scalar, -}; +use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar}; -/// This trait defines metric operations for leaf-local neighbor selection. +/// Leaf formulas return ascending distances. /// -/// Each function returns an ascending distance. The leaf kernel supplies squared -/// norms to L2 and norms to cosine. Dot-only metrics receive zero norm values. -pub(in super::super) trait LeafKernelMetric: MetricTag { - /// This value is true when the metric reads leaf norms. +/// L2 uses squared norms. Cosine uses norms. Dot-only metrics ignore norms. +pub(in super::super) trait LeafKernelMetric: Send + Sync + 'static { const USES_NORMS: bool; - /// This function converts one Gram diagonal value to the required norm unit. fn prepare_norm(squared_norm: f32) -> f32; - /// This function computes SIMD distances from one source to earlier targets. fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; - /// This function computes one scalar-tail leaf distance. fn leaf_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32; } diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 9ef6b36008..6d6ff17dbc 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -5,37 +5,31 @@ //! This module defines metric formulas for partition-leader ranking. +use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; -use super::{ - Cosine, CosineNormalized, InnerProduct, L2, MetricTag, cosine_distance, cosine_distance_scalar, -}; +use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar}; -/// This trait defines metric operations for partition-leader ranking. +/// Partition formulas return ascending ranking scores. /// -/// Each function returns an ascending score. L2 receives a squared leader norm. -/// Cosine receives point and leader norms. Dot-only metrics receive zero norms. -pub(in super::super) trait PartitionKernelMetric: MetricTag { - /// This value is true when the metric reads a point norm. +/// L2 uses squared leader norms. Cosine uses point and leader norms. +pub(in super::super) trait PartitionKernelMetric: Send + Sync + 'static { + const METRIC: Metric; const USES_POINT_NORM: bool; - - /// This value is true when the metric reads leader norm values. const USES_LEADER_NORMS: bool; - /// This function converts one squared point norm to the required unit. fn prepare_point_norm(squared_norm: f32) -> f32; - /// This function computes SIMD ranking scores for one point and multiple leaders. fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; - /// This function computes one scalar-tail partition ranking score. fn partition_ranking_scalar(dot: f32, point_norm: f32, leader_norm: f32) -> f32; } impl PartitionKernelMetric for L2 { + const METRIC: Metric = Metric::L2; const USES_POINT_NORM: bool = false; const USES_LEADER_NORMS: bool = true; @@ -63,6 +57,7 @@ impl PartitionKernelMetric for L2 { } impl PartitionKernelMetric for Cosine { + const METRIC: Metric = Metric::Cosine; const USES_POINT_NORM: bool = true; const USES_LEADER_NORMS: bool = true; @@ -87,6 +82,7 @@ impl PartitionKernelMetric for Cosine { } impl PartitionKernelMetric for CosineNormalized { + const METRIC: Metric = Metric::CosineNormalized; const USES_POINT_NORM: bool = false; const USES_LEADER_NORMS: bool = false; @@ -111,6 +107,7 @@ impl PartitionKernelMetric for CosineNormalized { } impl PartitionKernelMetric for InnerProduct { + const METRIC: Metric = Metric::InnerProduct; const USES_POINT_NORM: bool = false; const USES_LEADER_NORMS: bool = false; diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 09572597d0..906e05fae0 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -23,7 +23,7 @@ use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, }; -use super::kernel_metric::{MetricTag, PartitionKernelMetric}; +use super::kernel_metric::PartitionKernelMetric; /// Reusable nearest-center state for one partition worker. /// @@ -179,7 +179,7 @@ fn validate<'a, M: PartitionKernelMetric>( }); } - match (::METRIC, input.norms) { + match (M::METRIC, input.norms) { ( Metric::L2, PartitionNorms::L2 { @@ -485,7 +485,7 @@ mod tests { leader_norm_values: &[], }, }; - let metric = ::METRIC; + let metric = M::METRIC; for (point, (point_dots, point_output)) in input .dots .row_iter() From b42285b48db08bd2f2225126873a39cd4231471a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:02:21 +0000 Subject: [PATCH 67/80] docs(pipnn): keep metric comments factual --- diskann/src/graph/pipnn/kernel_metric.rs | 6 ++---- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 5 +---- .../src/graph/pipnn/kernel_metric/partition.rs | 5 +---- diskann/src/graph/pipnn/partition_kernel.rs | 16 ++-------------- 4 files changed, 6 insertions(+), 26 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 1642b037cf..c6dcf71fe8 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -3,7 +3,7 @@ * Licensed under the MIT license. */ -//! This module defines metric markers and shared numerical functions. +//! This module provides metric markers and shared numerical functions. mod leaf; mod partition; @@ -20,8 +20,7 @@ pub(super) struct InnerProduct; /// This function converts a squared norm to a norm. /// -/// The function maps subnormal squared norms to zero. It preserves NaN so that -/// kernel comparisons do not rank an invalid value. +/// It maps subnormal values to zero and preserves NaN. #[inline(always)] pub(super) fn norm_from_squared(squared_norm: f32) -> f32 { if squared_norm < f32::MIN_POSITIVE { @@ -52,7 +51,6 @@ where one - cosine } -/// This function computes scalar cosine distance with the DiskANN zero-norm rules. #[inline(always)] pub(super) 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() { diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 6390aa1ea7..ca4dce48be 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -3,15 +3,12 @@ * Licensed under the MIT license. */ -//! This module defines metric formulas for leaf-local neighbor selection. - use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar}; /// Leaf formulas return ascending distances. -/// -/// L2 uses squared norms. Cosine uses norms. Dot-only metrics ignore norms. +/// L2 uses squared norms. Cosine uses norms. Other metrics ignore norms. pub(in super::super) trait LeafKernelMetric: Send + Sync + 'static { const USES_NORMS: bool; diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 6d6ff17dbc..b8a308f029 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -3,15 +3,12 @@ * Licensed under the MIT license. */ -//! This module defines metric formulas for partition-leader ranking. - use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar}; -/// Partition formulas return ascending ranking scores. -/// +/// Partition formulas return ascending scores. /// L2 uses squared leader norms. Cosine uses point and leader norms. pub(in super::super) trait PartitionKernelMetric: Send + Sync + 'static { const METRIC: Metric; diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 906e05fae0..4b52fe32ad 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -44,24 +44,16 @@ impl PartitionKernelWorkspace { } } -/// This enum stores norm values for one point-to-leader tile. -/// -/// Cosine points use squared norms. Cosine leaders use norms. +/// L2 uses squared leader norms. Cosine uses squared point norms and leader norms. #[derive(Clone, Copy, Debug)] pub(super) enum PartitionNorms<'a> { - /// This variant stores squared norms for L2 leaders. L2 { - /// This slice contains one squared norm for each sampled leader. leader_squared_norms: &'a [f32], }, - /// This variant stores norms for unnormalized cosine. Cosine { - /// This slice contains one squared norm for each point. point_squared_norms: &'a [f32], - /// This slice contains one norm for each sampled leader. leader_norms: &'a [f32], }, - /// This variant provides no norms for normalized cosine or inner product. None, } @@ -143,17 +135,13 @@ where select_point_leaders::(arch, input.dots, norms, output, &mut workspace.tracker) } -/// This structure stores checked norm slices for one concrete metric. #[derive(Clone, Copy)] struct PartitionNormSlices<'a> { point_squared_norms: &'a [f32], leader_norm_values: &'a [f32], } -/// This function checks safety and metric conditions for partition selection. -/// -/// Matrix views prove their backing lengths. This function checks row counts, -/// leader-ID range, fanout, norm variant, and norm lengths. +/// This function checks row counts, leader IDs, fanout, and norm lengths. fn validate<'a, M: PartitionKernelMetric>( input: PartitionInput<'a>, output: &MutMatrixView<'_, u32>, From 2592d70e608dd2aa0a1af474dc1c2410ca759c11 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 07:49:32 +0000 Subject: [PATCH 68/80] refactor(pipnn): pass prepared norms to kernels --- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 32 -- .../graph/pipnn/kernel_metric/partition.rs | 43 --- diskann/src/graph/pipnn/leaf_kernel.rs | 160 ++++++---- diskann/src/graph/pipnn/partition_kernel.rs | 276 +++++++----------- 4 files changed, 214 insertions(+), 297 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index ca4dce48be..045d1178e7 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -10,10 +10,6 @@ use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_ /// Leaf formulas return ascending distances. /// L2 uses squared norms. Cosine uses norms. Other metrics ignore norms. pub(in super::super) trait LeafKernelMetric: Send + Sync + 'static { - const USES_NORMS: bool; - - fn prepare_norm(squared_norm: f32) -> f32; - fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, @@ -42,13 +38,6 @@ fn clamp_nonnegative_scalar(distance: f32) -> f32 { } impl LeafKernelMetric for L2 { - const USES_NORMS: bool = true; - - #[inline(always)] - fn prepare_norm(squared_norm: f32) -> f32 { - squared_norm - } - #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -65,13 +54,6 @@ impl LeafKernelMetric for L2 { } impl LeafKernelMetric for Cosine { - const USES_NORMS: bool = true; - - #[inline(always)] - fn prepare_norm(squared_norm: f32) -> f32 { - super::norm_from_squared(squared_norm) - } - #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -88,13 +70,6 @@ impl LeafKernelMetric for Cosine { } impl LeafKernelMetric for CosineNormalized { - const USES_NORMS: bool = false; - - #[inline(always)] - fn prepare_norm(_: f32) -> f32 { - 0.0 - } - #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F where @@ -111,13 +86,6 @@ impl LeafKernelMetric for CosineNormalized { } impl LeafKernelMetric for InnerProduct { - const USES_NORMS: bool = false; - - #[inline(always)] - fn prepare_norm(_: f32) -> f32 { - 0.0 - } - #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F where diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index b8a308f029..b672c2063d 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -3,7 +3,6 @@ * Licensed under the MIT license. */ -use diskann_vector::distance::Metric; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar}; @@ -11,12 +10,6 @@ use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_ /// Partition formulas return ascending scores. /// L2 uses squared leader norms. Cosine uses point and leader norms. pub(in super::super) trait PartitionKernelMetric: Send + Sync + 'static { - const METRIC: Metric; - const USES_POINT_NORM: bool; - const USES_LEADER_NORMS: bool; - - fn prepare_point_norm(squared_norm: f32) -> f32; - fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, @@ -26,15 +19,6 @@ pub(in super::super) trait PartitionKernelMetric: Send + Sync + 'static { } impl PartitionKernelMetric for L2 { - const METRIC: Metric = Metric::L2; - const USES_POINT_NORM: bool = false; - const USES_LEADER_NORMS: bool = true; - - #[inline(always)] - fn prepare_point_norm(_: f32) -> f32 { - 0.0 - } - #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, leader_norm: F) -> F where @@ -54,15 +38,6 @@ impl PartitionKernelMetric for L2 { } impl PartitionKernelMetric for Cosine { - const METRIC: Metric = Metric::Cosine; - const USES_POINT_NORM: bool = true; - const USES_LEADER_NORMS: bool = true; - - #[inline(always)] - fn prepare_point_norm(squared_norm: f32) -> f32 { - super::norm_from_squared(squared_norm) - } - #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F where @@ -79,15 +54,6 @@ impl PartitionKernelMetric for Cosine { } impl PartitionKernelMetric for CosineNormalized { - const METRIC: Metric = Metric::CosineNormalized; - const USES_POINT_NORM: bool = false; - const USES_LEADER_NORMS: bool = false; - - #[inline(always)] - fn prepare_point_norm(_: f32) -> f32 { - 0.0 - } - #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, _: F) -> F where @@ -104,15 +70,6 @@ impl PartitionKernelMetric for CosineNormalized { } impl PartitionKernelMetric for InnerProduct { - const METRIC: Metric = Metric::InnerProduct; - const USES_POINT_NORM: bool = false; - const USES_LEADER_NORMS: bool = false; - - #[inline(always)] - fn prepare_point_norm(_: f32) -> f32 { - 0.0 - } - #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, _: F) -> F where diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 07429cfe19..69bf5beb9a 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -57,7 +57,6 @@ impl Default for LeafNeighbor { /// Reusable temporary storage for leaf top-k selection. #[derive(Debug, Default)] pub(super) struct LeafKernelWorkspace { - norms: Vec, worst: Vec, } @@ -84,6 +83,9 @@ pub(super) enum LeafKernelError { neighbors: usize, maximum: usize, }, + /// The prepared norm count does not match the point count. + #[error("invalid leaf norm count: expected {expected}, got {actual}")] + InvalidNormCount { expected: usize, actual: usize }, /// Temporary storage could not be reserved. #[error("failed to reserve {additional} values for {buffer}")] Allocation { @@ -139,6 +141,7 @@ pub(super) fn leaf_neighbor_count( pub(super) fn nearest_neighbors( arch: A, input: MatrixView<'_, f32>, + norms: &[f32], mut output: MutMatrixView<'_, LeafNeighbor>, workspace: &mut LeafKernelWorkspace, ) -> Result<(), LeafKernelError> @@ -149,13 +152,18 @@ where M: LeafKernelMetric, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { - validate(input, &output)?; + validate(input, norms, &output)?; let neighbor_count = output.ncols(); if neighbor_count == 0 { return Ok(()); } - prepare_workspace::(input, workspace)?; + resize( + "worst distances", + &mut workspace.worst, + input.nrows(), + f32::INFINITY, + )?; output.as_mut_slice().fill(LeafNeighbor::default()); workspace.worst.fill(f32::INFINITY); @@ -164,21 +172,24 @@ where arch, input, output.as_mut_slice(), - &workspace.norms, + norms, + !norms.is_empty(), &mut workspace.worst, ), 2 => scan_point_pairs::( arch, input, output.as_mut_slice(), - &workspace.norms, + norms, + !norms.is_empty(), &mut workspace.worst, ), 3 => scan_point_pairs::( arch, input, output.as_mut_slice(), - &workspace.norms, + norms, + !norms.is_empty(), &mut workspace.worst, ), _ => { @@ -209,6 +220,7 @@ where /// An error occurs before the kernel changes output or workspace. fn validate( input: MatrixView<'_, f32>, + norms: &[f32], output: &MutMatrixView<'_, LeafNeighbor>, ) -> Result<(), LeafKernelError> { let point_count = input.nrows(); @@ -222,6 +234,12 @@ fn validate( cols: dot_columns, }); } + if !norms.is_empty() && norms.len() != point_count { + return Err(LeafKernelError::InvalidNormCount { + expected: point_count, + actual: norms.len(), + }); + } if output.nrows() != point_count { return Err(LeafKernelError::InvalidOutputRows { expected: point_count, @@ -241,33 +259,6 @@ fn validate( Ok(()) } -/// Prepare metric norms and rejection thresholds. -/// -/// L2 uses each diagonal value as a squared norm. Cosine converts each diagonal -/// value to a norm. Normalized cosine and inner product clear the norm buffer. -/// -/// An allocation error occurs before SIMD traversal. -fn prepare_workspace( - input: MatrixView<'_, f32>, - workspace: &mut LeafKernelWorkspace, -) -> Result<(), LeafKernelError> { - let points = input.nrows(); - if M::USES_NORMS { - resize("norms", &mut workspace.norms, points, 0.0)?; - for (source, norm) in workspace.norms.iter_mut().enumerate() { - *norm = M::prepare_norm(input[(source, source)]); - } - } else { - workspace.norms.clear(); - } - resize( - "worst distances", - &mut workspace.worst, - points, - f32::INFINITY, - ) -} - fn resize( buffer: &'static str, values: &mut Vec, @@ -287,14 +278,14 @@ fn resize( /// The function reads the strict lower triangle once. It offers each distance to /// both endpoint lists. SIMD groups and the scalar tail preserve pair scan order. /// -/// `input` supplies square dot products. `output` and `worst` contain one row -/// for each point. `norms` contains one value per point when `M` uses norms. +/// `input` supplies square dot products. `norms` contains prepared norm values. #[inline(never)] fn scan_point_pairs( arch: F::Arch, input: MatrixView<'_, f32>, output: &mut [LeafNeighbor], norms: &[f32], + uses_norms: bool, worst: &mut [f32], ) where F: SIMDVector> + SIMDFloat + std::ops::Div, @@ -305,7 +296,6 @@ fn scan_point_pairs( let (output, _) = output.as_chunks_mut::(); let point_count = input.nrows(); let dots = input.as_slice(); - let uses_norms = M::USES_NORMS; let worst_ptr = worst.as_mut_ptr(); // Source zero has no earlier target. Each source after zero can still add @@ -454,6 +444,7 @@ fn insert_fixed_neighbor( #[cfg(test)] struct DispatchedLeafCall<'a> { input: MatrixView<'a, f32>, + norms: &'a [f32], output: MutMatrixView<'a, LeafNeighbor>, workspace: &'a mut LeafKernelWorkspace, } @@ -475,19 +466,34 @@ where use diskann_vector::distance::Metric; match self.0 { - Metric::L2 => nearest_neighbors::(arch, call.input, call.output, call.workspace), - Metric::Cosine => { - nearest_neighbors::(arch, call.input, call.output, call.workspace) - } + Metric::L2 => nearest_neighbors::( + arch, + call.input, + call.norms, + call.output, + call.workspace, + ), + Metric::Cosine => nearest_neighbors::( + arch, + call.input, + call.norms, + call.output, + call.workspace, + ), Metric::CosineNormalized => nearest_neighbors::( arch, call.input, + call.norms, + call.output, + call.workspace, + ), + Metric::InnerProduct => nearest_neighbors::( + arch, + call.input, + call.norms, call.output, call.workspace, ), - Metric::InnerProduct => { - nearest_neighbors::(arch, call.input, call.output, call.workspace) - } } } } @@ -496,6 +502,7 @@ where fn dispatch_nearest_neighbors( metric: diskann_vector::distance::Metric, input: MatrixView<'_, f32>, + norms: &[f32], output: MutMatrixView<'_, LeafNeighbor>, workspace: &mut LeafKernelWorkspace, ) -> Result<(), LeafKernelError> { @@ -503,12 +510,31 @@ fn dispatch_nearest_neighbors( DispatchLeafForTest(metric), DispatchedLeafCall { input, + norms, output, workspace, }, ) } +#[cfg(test)] +fn prepared_test_norms( + metric: diskann_vector::distance::Metric, + input: MatrixView<'_, f32>, +) -> Vec { + use diskann_vector::distance::Metric; + + match metric { + Metric::L2 => (0..input.nrows()) + .map(|point| input[(point, point)]) + .collect(), + Metric::Cosine => (0..input.nrows()) + .map(|point| super::kernel_metric::norm_from_squared(input[(point, point)])) + .collect(), + Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), + } +} + #[cfg(test)] mod tests { use super::*; @@ -581,6 +607,7 @@ mod tests { let points = 7; let dots = test_dots(Metric::L2, points); let input = test_input(&dots, points); + let norms = prepared_test_norms(Metric::L2, input); let mut workspace = LeafKernelWorkspace::default(); for neighbor_count in [1, 3, 2] { @@ -588,6 +615,7 @@ mod tests { dispatch_nearest_neighbors( Metric::L2, input, + &norms, MutMatrixView::try_from(output.as_mut_slice(), points, neighbor_count).unwrap(), &mut workspace, ) @@ -602,9 +630,12 @@ mod tests { for points in [17, 7, 17] { let dots = test_dots(Metric::L2, points); let mut output = vec![LeafNeighbor::default(); points * 2]; + let input = test_input(&dots, points); + let norms = prepared_test_norms(Metric::L2, input); dispatch_nearest_neighbors( Metric::L2, - test_input(&dots, points), + input, + &norms, MutMatrixView::try_from(output.as_mut_slice(), points, 2).unwrap(), &mut workspace, ) @@ -624,7 +655,7 @@ mod integration_tests { use super::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, MAX_LEAF_NEIGHBORS, - dispatch_nearest_neighbors, leaf_neighbor_count, + dispatch_nearest_neighbors, leaf_neighbor_count, prepared_test_norms, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -745,10 +776,13 @@ mod integration_tests { metric: Metric, ) -> (usize, Vec) { let leaf_k = leaf_neighbor_count(points, requested_k).unwrap(); + let input = test_input(dots, points); + let norms = prepared_test_norms(metric, input); let mut output = vec![LeafNeighbor::default(); points * leaf_k]; dispatch_nearest_neighbors( metric, - test_input(dots, points), + input, + &norms, MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), &mut LeafKernelWorkspace::default(), ) @@ -914,9 +948,12 @@ mod integration_tests { 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 input = test_input(&dots, 2); + let norms = prepared_test_norms(Metric::L2, input); let error = dispatch_nearest_neighbors( Metric::L2, - test_input(&dots, 2), + input, + &norms, MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), &mut LeafKernelWorkspace::default(), ) @@ -971,6 +1008,7 @@ mod integration_tests { dispatch_nearest_neighbors( Metric::L2, non_square, + &[], MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), &mut LeafKernelWorkspace::default(), ), @@ -978,11 +1016,29 @@ mod integration_tests { ); let square = [0.0; 9]; + let square_input = test_input(&square, 3); + let square_norms = prepared_test_norms(Metric::L2, square_input); + let mut valid_output = [LeafNeighbor::default(); 3]; + assert_eq!( + dispatch_nearest_neighbors( + Metric::L2, + square_input, + &square_norms[..2], + MutMatrixView::try_from(&mut valid_output[..], 3, 1).unwrap(), + &mut LeafKernelWorkspace::default(), + ), + Err(LeafKernelError::InvalidNormCount { + expected: 3, + actual: 2, + }) + ); + let mut wrong_rows = [LeafNeighbor::default(); 2]; assert_eq!( dispatch_nearest_neighbors( Metric::L2, - test_input(&square, 3), + square_input, + &square_norms, MutMatrixView::try_from(&mut wrong_rows[..], 2, 1).unwrap(), &mut LeafKernelWorkspace::default(), ), @@ -997,7 +1053,8 @@ mod integration_tests { assert_eq!( dispatch_nearest_neighbors( Metric::L2, - test_input(&square, 3), + square_input, + &square_norms, MutMatrixView::try_from(&mut too_many[..], 3, 3).unwrap(), &mut LeafKernelWorkspace::default(), ), @@ -1009,11 +1066,14 @@ mod integration_tests { ); let square = [0.0; 25]; + let square_input = test_input(&square, 5); + let square_norms = prepared_test_norms(Metric::L2, square_input); let mut too_wide = [LeafNeighbor::default(); 20]; assert_eq!( dispatch_nearest_neighbors( Metric::L2, - test_input(&square, 5), + square_input, + &square_norms, MutMatrixView::try_from(&mut too_wide[..], 5, 4).unwrap(), &mut LeafKernelWorkspace::default(), ), diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 4b52fe32ad..0cb9c1b227 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -44,17 +44,11 @@ impl PartitionKernelWorkspace { } } -/// L2 uses squared leader norms. Cosine uses squared point norms and leader norms. +/// L2 uses squared leader norms. Cosine uses point and leader norms. #[derive(Clone, Copy, Debug)] -pub(super) enum PartitionNorms<'a> { - L2 { - leader_squared_norms: &'a [f32], - }, - Cosine { - point_squared_norms: &'a [f32], - leader_norms: &'a [f32], - }, - None, +pub(super) struct PartitionNorms<'a> { + pub(super) point_norms: &'a [f32], + pub(super) leader_norms: &'a [f32], } /// Dot products between assigned points and sampled partition centers. @@ -86,9 +80,6 @@ pub(super) enum PartitionKernelError { expected: usize, actual: usize, }, - /// Norm inputs do not match concrete metric `M`. - #[error("partition norms do not match selected {expected} metric")] - InvalidNorms { expected: &'static str }, /// The requested fanout exceeds the available leader count. #[error("invalid fanout {fanout}: must not exceed {leader_count} leaders")] InvalidFanout { fanout: usize, leader_count: usize }, @@ -114,6 +105,7 @@ pub(super) enum PartitionKernelError { /// It also returns an error when fewer than `fanout` scores are rankable. pub(super) fn nearest_leaders( arch: A, + metric: Metric, input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, workspace: &mut PartitionKernelWorkspace, @@ -125,27 +117,29 @@ where u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, M: PartitionKernelMetric, { - let norms = validate::(input, &output)?; + let norms = validate(metric, input, &output)?; let fanout = output.ncols(); if fanout == 0 || input.dots.nrows() == 0 { return Ok(()); } workspace.prepare(fanout)?; - select_point_leaders::(arch, input.dots, norms, output, &mut workspace.tracker) -} - -#[derive(Clone, Copy)] -struct PartitionNormSlices<'a> { - point_squared_norms: &'a [f32], - leader_norm_values: &'a [f32], + select_point_leaders::( + arch, + metric, + input.dots, + norms, + output, + &mut workspace.tracker, + ) } /// This function checks row counts, leader IDs, fanout, and norm lengths. -fn validate<'a, M: PartitionKernelMetric>( +fn validate<'a>( + metric: Metric, input: PartitionInput<'a>, output: &MutMatrixView<'_, u32>, -) -> Result, PartitionKernelError> { +) -> Result, PartitionKernelError> { let point_count = input.dots.nrows(); let leader_count = input.dots.ncols(); let fanout = output.ncols(); @@ -167,48 +161,22 @@ fn validate<'a, M: PartitionKernelMetric>( }); } - match (M::METRIC, input.norms) { - ( - Metric::L2, - PartitionNorms::L2 { - leader_squared_norms, - }, - ) => { - check_length("leader norms", leader_squared_norms.len(), leader_count)?; - Ok(PartitionNormSlices { - point_squared_norms: &[], - leader_norm_values: leader_squared_norms, - }) - } - ( - Metric::Cosine, - PartitionNorms::Cosine { - point_squared_norms, - leader_norms, - }, - ) => { - check_length("point norms", point_squared_norms.len(), point_count)?; - check_length("leader norms", leader_norms.len(), leader_count)?; - Ok(PartitionNormSlices { - point_squared_norms, - leader_norm_values: leader_norms, - }) - } - (Metric::CosineNormalized | Metric::InnerProduct, PartitionNorms::None) => { - Ok(PartitionNormSlices { - point_squared_norms: &[], - leader_norm_values: &[], - }) - } - (Metric::L2, _) => Err(PartitionKernelError::InvalidNorms { expected: "L2" }), - (Metric::Cosine, _) => Err(PartitionKernelError::InvalidNorms { expected: "cosine" }), - (Metric::CosineNormalized, _) => Err(PartitionKernelError::InvalidNorms { - expected: "normalized cosine", - }), - (Metric::InnerProduct, _) => Err(PartitionKernelError::InvalidNorms { - expected: "inner product", - }), - } + let (point_norm_count, leader_norm_count) = match metric { + Metric::L2 => (0, leader_count), + Metric::Cosine => (point_count, leader_count), + Metric::CosineNormalized | Metric::InnerProduct => (0, 0), + }; + check_length( + "point norms", + input.norms.point_norms.len(), + point_norm_count, + )?; + check_length( + "leader norms", + input.norms.leader_norms.len(), + leader_norm_count, + )?; + Ok(input.norms) } fn check_length( @@ -237,8 +205,9 @@ fn check_length( /// point. The function resets this state before it processes another point. fn select_point_leaders( arch: F::Arch, + metric: Metric, dots: MatrixView<'_, f32>, - norms: PartitionNormSlices<'_>, + norms: PartitionNorms<'_>, mut output: MutMatrixView<'_, u32>, tracker: &mut [(u32, f32)], ) -> Result<(), PartitionKernelError> @@ -250,8 +219,8 @@ where { let leader_count = dots.ncols(); let fanout = output.ncols(); - let uses_point_norm = M::USES_POINT_NORM; - let uses_leader_norm = M::USES_LEADER_NORMS; + let uses_point_norm = metric == Metric::Cosine; + let uses_leader_norm = matches!(metric, Metric::L2 | Metric::Cosine); // Reset the tracker for each point. No assignment state can pass from one // output row to another. for (point, (point_dots, point_output)) in dots @@ -261,7 +230,7 @@ where { tracker.fill((u32::MAX, f32::INFINITY)); let point_norm = if uses_point_norm { - M::prepare_point_norm(norms.point_squared_norms[point]) + norms.point_norms[point] } else { 0.0 }; @@ -276,7 +245,7 @@ where let leader_norms = if uses_leader_norm { // SAFETY: `validate` established one norm value per leader. // `base + F::LANES <= full <= leader_count`. - unsafe { F::load_simd(arch, norms.leader_norm_values.as_ptr().add(base)) } + unsafe { F::load_simd(arch, norms.leader_norms.as_ptr().add(base)) } } else { F::default(arch) }; @@ -291,7 +260,7 @@ where // norm slice and can change L2 rounding. for (leader, &dot) in point_dots.iter().enumerate().skip(full) { let leader_norm = if uses_leader_norm { - norms.leader_norm_values[leader] + norms.leader_norms[leader] } else { 0.0 }; @@ -381,19 +350,30 @@ where use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => nearest_leaders::(arch, call.input, call.output, call.workspace), - Metric::Cosine => { - nearest_leaders::(arch, call.input, call.output, call.workspace) + Metric::L2 => { + nearest_leaders::(arch, Metric::L2, call.input, call.output, call.workspace) } + Metric::Cosine => nearest_leaders::( + arch, + Metric::Cosine, + call.input, + call.output, + call.workspace, + ), Metric::CosineNormalized => nearest_leaders::( arch, + Metric::CosineNormalized, + call.input, + call.output, + call.workspace, + ), + Metric::InnerProduct => nearest_leaders::( + arch, + Metric::InnerProduct, call.input, call.output, call.workspace, ), - Metric::InnerProduct => { - nearest_leaders::(arch, call.input, call.output, call.workspace) - } } } } @@ -418,77 +398,47 @@ fn dispatch_nearest_leaders( #[cfg(test)] mod tests { use super::super::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, PartitionKernelMetric, norm_from_squared, + Cosine, CosineNormalized, InnerProduct, L2, PartitionKernelMetric, }; use super::*; fn test_input<'a>( - metric: Metric, + _metric: Metric, dots: &'a [f32], point_count: usize, leader_count: usize, - point_squared_norms: &'a [f32], - leader_norm_values: &'a [f32], + point_norms: &'a [f32], + leader_norms: &'a [f32], ) -> PartitionInput<'a> { - let norms = match metric { - Metric::L2 => PartitionNorms::L2 { - leader_squared_norms: leader_norm_values, - }, - Metric::Cosine => PartitionNorms::Cosine { - point_squared_norms, - leader_norms: leader_norm_values, - }, - Metric::CosineNormalized | Metric::InnerProduct => PartitionNorms::None, - }; PartitionInput { dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), - norms, + norms: PartitionNorms { + point_norms, + leader_norms, + }, } } // This oracle checks SIMD chunking, scalar tails, and tracker order. It uses // the scalar ranking formula for metric `M`. fn scalar_traversal_reference( + metric: Metric, input: PartitionInput<'_>, fanout: usize, output: &mut [u32], ) { - let norms = match input.norms { - PartitionNorms::L2 { - leader_squared_norms, - } => PartitionNormSlices { - point_squared_norms: &[], - leader_norm_values: leader_squared_norms, - }, - PartitionNorms::Cosine { - point_squared_norms, - leader_norms, - } => PartitionNormSlices { - point_squared_norms, - leader_norm_values: leader_norms, - }, - PartitionNorms::None => PartitionNormSlices { - point_squared_norms: &[], - leader_norm_values: &[], - }, - }; - let metric = M::METRIC; for (point, (point_dots, point_output)) in input .dots .row_iter() .zip(output.chunks_exact_mut(fanout)) .enumerate() { - let point_norm = if metric == Metric::Cosine { - norm_from_squared(norms.point_squared_norms[point]) - } else { - 0.0 - }; + let point_norm = input.norms.point_norms.get(point).copied().unwrap_or(0.0); let mut tracker = vec![(u32::MAX, f32::INFINITY); fanout]; for (leader, &dot) in point_dots.iter().enumerate() { let leader_norm = if matches!(metric, Metric::L2 | Metric::Cosine) { - norms.leader_norm_values[leader] + input.norms.leader_norms[leader] } else { 0.0 }; @@ -520,8 +470,8 @@ mod tests { #[test] fn cosine_special_norms_match_scalar_and_dispatched_kernel() { let leader_count = 17; - let point_squared_norms = [0.0, f32::MIN_POSITIVE / 2.0, f32::MIN_POSITIVE, f32::NAN]; - let dots = vec![1.0; point_squared_norms.len() * leader_count]; + let point_norms = [0.0, 0.0, f32::MIN_POSITIVE.sqrt(), f32::NAN]; + let dots = vec![1.0; point_norms.len() * leader_count]; let mut leader_norms = vec![1.0; leader_count]; leader_norms[..4].copy_from_slice(&[ 0.0, @@ -532,18 +482,18 @@ mod tests { let input = test_input( Metric::Cosine, &dots, - point_squared_norms.len(), + point_norms.len(), leader_count, - &point_squared_norms, + &point_norms, &leader_norms, ); - let mut expected = vec![u32::MAX; point_squared_norms.len() * 2]; - scalar_traversal_reference::(input, 2, &mut expected); - let mut actual = vec![u32::MAX; point_squared_norms.len() * 2]; + let mut expected = vec![u32::MAX; point_norms.len() * 2]; + scalar_traversal_reference::(Metric::Cosine, input, 2, &mut expected); + let mut actual = vec![u32::MAX; point_norms.len() * 2]; dispatch_nearest_leaders( Metric::Cosine, input, - MutMatrixView::try_from(actual.as_mut_slice(), point_squared_norms.len(), 2).unwrap(), + MutMatrixView::try_from(actual.as_mut_slice(), point_norms.len(), 2).unwrap(), &mut PartitionKernelWorkspace::default(), ) .unwrap(); @@ -583,7 +533,6 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod integration_tests { - use super::super::kernel_metric::norm_from_squared; use super::{ PartitionInput, PartitionKernelError, PartitionKernelWorkspace, PartitionNorms, dispatch_nearest_leaders, @@ -592,42 +541,27 @@ mod integration_tests { use diskann_vector::distance::Metric; fn test_input<'a>( - metric: Metric, + _metric: Metric, dots: &'a [f32], point_count: usize, leader_count: usize, - point_squared_norms: &'a [f32], - leader_norm_values: &'a [f32], + point_norms: &'a [f32], + leader_norms: &'a [f32], ) -> PartitionInput<'a> { - let norms = match metric { - Metric::L2 => PartitionNorms::L2 { - leader_squared_norms: leader_norm_values, - }, - Metric::Cosine => PartitionNorms::Cosine { - point_squared_norms, - leader_norms: leader_norm_values, - }, - Metric::CosineNormalized | Metric::InnerProduct => PartitionNorms::None, - }; PartitionInput { dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), - norms, + norms: PartitionNorms { + point_norms, + leader_norms, + }, } } 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_squared_norms, leader_norm_values) = match input.norms { - PartitionNorms::L2 { - leader_squared_norms, - } => (&[][..], leader_squared_norms), - PartitionNorms::Cosine { - point_squared_norms, - leader_norms, - } => (point_squared_norms, leader_norms), - PartitionNorms::None => (&[][..], &[][..]), - }; + let point_norms = input.norms.point_norms; + let leader_norms = input.norms.leader_norms; let mut assignments = vec![u32::MAX; point_count * fanout]; for (point, (point_dots, point_assignments)) in input .dots @@ -636,18 +570,17 @@ mod integration_tests { .zip(assignments.chunks_exact_mut(fanout)) .enumerate() { - let point_squared_norm = point_squared_norms.get(point).copied().unwrap_or(0.0); + let point_norm = point_norms.get(point).copied().unwrap_or(0.0); let mut candidates: Vec<_> = point_dots .iter() .enumerate() .filter_map(|(leader, &dot)| { - let leader_norm = leader_norm_values.get(leader).copied().unwrap_or(0.0); + let leader_norm = leader_norms.get(leader).copied().unwrap_or(0.0); let score = match metric { Metric::L2 => leader_norm - 2.0 * dot, Metric::CosineNormalized => 1.0 - dot, Metric::InnerProduct => -dot, Metric::Cosine => { - let point_norm = norm_from_squared(point_squared_norm); 1.0 - if point_norm == 0.0 || leader_norm == 0.0 { 0.0 } else { @@ -682,12 +615,12 @@ mod integration_tests { } }) .collect(); - let point_squared_norms = if metric == Metric::Cosine { - vec![0.0, 16.0] + let point_norms = if metric == Metric::Cosine { + vec![0.0, 4.0] } else { Vec::new() }; - let leader_norm_values = match metric { + let leader_norms = match metric { Metric::Cosine => (0..leader_count) .map(|leader| { if leader == 1 { @@ -711,7 +644,7 @@ mod integration_tests { .collect(), Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), }; - (dots, point_squared_norms, leader_norm_values) + (dots, point_norms, leader_norms) } fn run_partition_kernel( @@ -738,16 +671,8 @@ mod integration_tests { Metric::InnerProduct, ] { for leader_count in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let (dots, point_squared_norms, leader_norm_values) = - differential_data(metric, leader_count); - let input = test_input( - metric, - &dots, - 2, - leader_count, - &point_squared_norms, - &leader_norm_values, - ); + let (dots, point_norms, leader_norms) = differential_data(metric, leader_count); + let input = test_input(metric, &dots, 2, leader_count, &point_norms, &leader_norms); for fanout in [1, 2, 16, 17, 32] { if fanout >= leader_count { continue; @@ -807,7 +732,7 @@ mod integration_tests { 1.0, 0.0, -1.0, 2.0, 6.0, 0.0, ]; - for (metric, point_squared_norms, leader_norm_values, expected) in [ + for (metric, point_norms, leader_norms, expected) in [ (Metric::L2, &[][..], &[1.0, 4.0, 9.0][..], [0, 1, 1, 0]), ( Metric::Cosine, @@ -821,7 +746,7 @@ mod integration_tests { assert_eq!( run_partition_kernel( metric, - test_input(metric, &dots, 2, 3, point_squared_norms, leader_norm_values), + test_input(metric, &dots, 2, 3, point_norms, leader_norms), 2, ) .unwrap(), @@ -947,11 +872,18 @@ mod integration_tests { let wrong_norms = PartitionInput { dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), - norms: PartitionNorms::None, + norms: PartitionNorms { + point_norms: &[], + leader_norms: &[], + }, }; assert_eq!( run_partition_kernel(Metric::L2, wrong_norms, 2), - Err(PartitionKernelError::InvalidNorms { expected: "L2" }) + Err(PartitionKernelError::InvalidBufferLength { + buffer: "leader norms", + expected: 3, + actual: 0, + }) ); assert_eq!( From ac6e17407b10db9ad0f2c92bd2758e240297409f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:15:20 +0000 Subject: [PATCH 69/80] perf(pipnn): specialize prepared norm layouts --- diskann/src/graph/pipnn/leaf_kernel.rs | 41 +++++++++++++------ diskann/src/graph/pipnn/partition_kernel.rs | 44 ++++++++++++++------- 2 files changed, 58 insertions(+), 27 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 69bf5beb9a..b1141360ca 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -167,29 +167,47 @@ where output.as_mut_slice().fill(LeafNeighbor::default()); workspace.worst.fill(f32::INFINITY); - match neighbor_count { - 1 => scan_point_pairs::( + match (norms.is_empty(), neighbor_count) { + (false, 1) => scan_point_pairs::( arch, input, output.as_mut_slice(), norms, - !norms.is_empty(), &mut workspace.worst, ), - 2 => scan_point_pairs::( + (false, 2) => scan_point_pairs::( arch, input, output.as_mut_slice(), norms, - !norms.is_empty(), &mut workspace.worst, ), - 3 => scan_point_pairs::( + (false, 3) => scan_point_pairs::( + arch, + input, + output.as_mut_slice(), + norms, + &mut workspace.worst, + ), + (true, 1) => scan_point_pairs::( + arch, + input, + output.as_mut_slice(), + norms, + &mut workspace.worst, + ), + (true, 2) => scan_point_pairs::( + arch, + input, + output.as_mut_slice(), + norms, + &mut workspace.worst, + ), + (true, 3) => scan_point_pairs::( arch, input, output.as_mut_slice(), norms, - !norms.is_empty(), &mut workspace.worst, ), _ => { @@ -280,12 +298,11 @@ fn resize( /// /// `input` supplies square dot products. `norms` contains prepared norm values. #[inline(never)] -fn scan_point_pairs( +fn scan_point_pairs( arch: F::Arch, input: MatrixView<'_, f32>, output: &mut [LeafNeighbor], norms: &[f32], - uses_norms: bool, worst: &mut [f32], ) where F: SIMDVector> + SIMDFloat + std::ops::Div, @@ -302,7 +319,7 @@ fn scan_point_pairs( // itself to the neighbor list of source zero. for source in 1..point_count { let source_start = source * point_count; - let source_norm = if uses_norms { + let source_norm = if USES_NORMS { F::splat(arch, norms[source]) } else { F::default(arch) @@ -316,7 +333,7 @@ fn scan_point_pairs( while target < full { // 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_norms = if uses_norms { + let target_norms = if USES_NORMS { // SAFETY: the full target chunk is below `source < point_count`. // `prepare_workspace` created one norm for each point. unsafe { F::load_simd(arch, norms.as_ptr().add(target)) } @@ -370,7 +387,7 @@ fn scan_point_pairs( 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_norm, target_norm) = if uses_norms { + let (source_norm, target_norm) = if USES_NORMS { // SAFETY: `target < source < point_count == norms.len()`. (norms[source], unsafe { *norms.get_unchecked(target) }) } else { diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 0cb9c1b227..3167562ce4 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -124,14 +124,31 @@ where } workspace.prepare(fanout)?; - select_point_leaders::( - arch, - metric, - input.dots, - norms, - output, - &mut workspace.tracker, - ) + match metric { + Metric::L2 => select_point_leaders::( + arch, + input.dots, + norms, + output, + &mut workspace.tracker, + ), + Metric::Cosine => select_point_leaders::( + arch, + input.dots, + norms, + output, + &mut workspace.tracker, + ), + Metric::CosineNormalized | Metric::InnerProduct => { + select_point_leaders::( + arch, + input.dots, + norms, + output, + &mut workspace.tracker, + ) + } + } } /// This function checks row counts, leader IDs, fanout, and norm lengths. @@ -203,9 +220,8 @@ fn check_length( /// /// `tracker` stores the retained center-column IDs and scores for the current /// point. The function resets this state before it processes another point. -fn select_point_leaders( +fn select_point_leaders( arch: F::Arch, - metric: Metric, dots: MatrixView<'_, f32>, norms: PartitionNorms<'_>, mut output: MutMatrixView<'_, u32>, @@ -219,8 +235,6 @@ where { let leader_count = dots.ncols(); let fanout = output.ncols(); - let uses_point_norm = metric == Metric::Cosine; - let uses_leader_norm = matches!(metric, Metric::L2 | Metric::Cosine); // Reset the tracker for each point. No assignment state can pass from one // output row to another. for (point, (point_dots, point_output)) in dots @@ -229,7 +243,7 @@ where .enumerate() { tracker.fill((u32::MAX, f32::INFINITY)); - let point_norm = if uses_point_norm { + let point_norm = if USES_POINT_NORMS { norms.point_norms[point] } else { 0.0 @@ -242,7 +256,7 @@ where for base in (0..full).step_by(F::LANES) { // SAFETY: `base + F::LANES <= full <= point_dots.len()`. let point_dots = unsafe { F::load_simd(arch, point_dots.as_ptr().add(base)) }; - let leader_norms = if uses_leader_norm { + let leader_norms = if USES_LEADER_NORMS { // SAFETY: `validate` established one norm value per leader. // `base + F::LANES <= full <= leader_count`. unsafe { F::load_simd(arch, norms.leader_norms.as_ptr().add(base)) } @@ -259,7 +273,7 @@ where // Use scalar formulas for the tail. A padded SIMD load can read past the // norm slice and can change L2 rounding. for (leader, &dot) in point_dots.iter().enumerate().skip(full) { - let leader_norm = if uses_leader_norm { + let leader_norm = if USES_LEADER_NORMS { norms.leader_norms[leader] } else { 0.0 From 6c3d9129d2e7856ade6d123669b04a8a5f10b81b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:27:40 +0000 Subject: [PATCH 70/80] perf(pipnn): reuse each leaf source norm --- diskann/src/graph/pipnn/leaf_kernel.rs | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index b1141360ca..354a884757 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -319,13 +319,9 @@ fn scan_point_pairs( // itself to the neighbor list of source zero. for source in 1..point_count { let source_start = source * point_count; - let source_norm = if USES_NORMS { - F::splat(arch, norms[source]) - } else { - F::default(arch) - }; - // SAFETY: `validate` and `prepare_workspace` established - // `source < point_count == worst.len()`. + let source_norm = if USES_NORMS { norms[source] } else { 0.0 }; + let source_norms = F::splat(arch, source_norm); + // SAFETY: `nearest_neighbors` created one threshold for each point. let mut source_worst = unsafe { *worst_ptr.add(source) }; let mut target = 0; let full = source / F::LANES * F::LANES; @@ -335,17 +331,16 @@ fn scan_point_pairs( let pair_dots = unsafe { F::load_simd(arch, dots.as_ptr().add(source_start + target)) }; let target_norms = if USES_NORMS { // SAFETY: the full target chunk is below `source < point_count`. - // `prepare_workspace` created one norm for each point. unsafe { F::load_simd(arch, norms.as_ptr().add(target)) } } else { F::default(arch) }; - let distances = M::leaf_distance(arch, pair_dots, source_norm, target_norms); + let distances = M::leaf_distance(arch, pair_dots, source_norms, target_norms); // 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 is below `source < point_count`. - // `prepare_workspace` created one threshold for each point. + // `nearest_neighbors` created one threshold for each point. 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()); @@ -387,11 +382,11 @@ fn scan_point_pairs( 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_norm, target_norm) = if USES_NORMS { + let target_norm = if USES_NORMS { // SAFETY: `target < source < point_count == norms.len()`. - (norms[source], unsafe { *norms.get_unchecked(target) }) + unsafe { *norms.get_unchecked(target) } } else { - (0.0, 0.0) + 0.0 }; let distance = M::leaf_distance_scalar(dot, source_norm, target_norm); if distance < source_worst { From 3ff5e3cf6465d38ebeb28e1ad35685956a66f632 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:43:13 +0000 Subject: [PATCH 71/80] refactor(pipnn): name prepared norm layouts --- diskann/src/graph/pipnn/leaf_kernel.rs | 8 ++++---- diskann/src/graph/pipnn/partition_kernel.rs | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 354a884757..8a691cd88f 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -298,7 +298,7 @@ fn resize( /// /// `input` supplies square dot products. `norms` contains prepared norm values. #[inline(never)] -fn scan_point_pairs( +fn scan_point_pairs( arch: F::Arch, input: MatrixView<'_, f32>, output: &mut [LeafNeighbor], @@ -319,7 +319,7 @@ fn scan_point_pairs( // itself to the neighbor list of source zero. for source in 1..point_count { let source_start = source * point_count; - let source_norm = if USES_NORMS { norms[source] } else { 0.0 }; + let source_norm = if HAS_NORMS { norms[source] } else { 0.0 }; let source_norms = F::splat(arch, source_norm); // SAFETY: `nearest_neighbors` created one threshold for each point. let mut source_worst = unsafe { *worst_ptr.add(source) }; @@ -329,7 +329,7 @@ fn scan_point_pairs( while target < full { // 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_norms = if USES_NORMS { + let target_norms = if HAS_NORMS { // SAFETY: the full target chunk is below `source < point_count`. unsafe { F::load_simd(arch, norms.as_ptr().add(target)) } } else { @@ -382,7 +382,7 @@ fn scan_point_pairs( 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 target_norm = if USES_NORMS { + let target_norm = if HAS_NORMS { // SAFETY: `target < source < point_count == norms.len()`. unsafe { *norms.get_unchecked(target) } } else { diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 3167562ce4..0307c8c091 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -220,7 +220,7 @@ fn check_length( /// /// `tracker` stores the retained center-column IDs and scores for the current /// point. The function resets this state before it processes another point. -fn select_point_leaders( +fn select_point_leaders( arch: F::Arch, dots: MatrixView<'_, f32>, norms: PartitionNorms<'_>, @@ -243,7 +243,7 @@ where .enumerate() { tracker.fill((u32::MAX, f32::INFINITY)); - let point_norm = if USES_POINT_NORMS { + let point_norm = if HAS_POINT_NORMS { norms.point_norms[point] } else { 0.0 @@ -256,7 +256,7 @@ where for base in (0..full).step_by(F::LANES) { // SAFETY: `base + F::LANES <= full <= point_dots.len()`. let point_dots = unsafe { F::load_simd(arch, point_dots.as_ptr().add(base)) }; - let leader_norms = if USES_LEADER_NORMS { + let leader_norms = if HAS_LEADER_NORMS { // SAFETY: `validate` established one norm value per leader. // `base + F::LANES <= full <= leader_count`. unsafe { F::load_simd(arch, norms.leader_norms.as_ptr().add(base)) } @@ -273,7 +273,7 @@ where // Use scalar formulas for the tail. A padded SIMD load can read past the // norm slice and can change L2 rounding. for (leader, &dot) in point_dots.iter().enumerate().skip(full) { - let leader_norm = if USES_LEADER_NORMS { + let leader_norm = if HAS_LEADER_NORMS { norms.leader_norms[leader] } else { 0.0 From 057873bc0fe6029be88d407495b1f63f267dfdd5 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:23:14 +0000 Subject: [PATCH 72/80] refactor(pipnn): make metric types prepare norms --- diskann/src/graph/pipnn/kernel_metric.rs | 12 +++- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 45 ++++++++++-- .../graph/pipnn/kernel_metric/partition.rs | 69 +++++++++++++++++-- diskann/src/graph/pipnn/leaf_kernel.rs | 6 +- diskann/src/graph/pipnn/partition_kernel.rs | 10 +-- 5 files changed, 120 insertions(+), 22 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index c6dcf71fe8..4c832f3106 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -8,8 +8,10 @@ mod leaf; mod partition; -pub(super) use leaf::LeafKernelMetric; -pub(super) use partition::PartitionKernelMetric; +pub(super) use leaf::LeafMetric; +pub(super) use partition::PartitionMetric; + +use std::collections::TryReserveError; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; @@ -18,6 +20,12 @@ pub(super) struct Cosine; pub(super) struct CosineNormalized; pub(super) struct InnerProduct; +pub(super) fn resize_norms(norms: &mut Vec, len: usize) -> Result<(), TryReserveError> { + norms.try_reserve(len.saturating_sub(norms.len()))?; + norms.resize(len, 0.0); + Ok(()) +} + /// This function converts a squared norm to a norm. /// /// It maps subnormal values to zero and preserves NaN. diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 045d1178e7..7cc3996247 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -3,13 +3,24 @@ * Licensed under the MIT license. */ +use std::collections::TryReserveError; + +use diskann_utils::views::MatrixView; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; -use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar}; +use super::{ + Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar, + norm_from_squared, resize_norms, +}; /// Leaf formulas return ascending distances. /// L2 uses squared norms. Cosine uses norms. Other metrics ignore norms. -pub(in super::super) trait LeafKernelMetric: Send + Sync + 'static { +pub(in super::super) trait LeafMetric: Send + Sync + 'static { + fn prepare_norms(_: MatrixView<'_, f32>, norms: &mut Vec) -> Result<(), TryReserveError> { + norms.clear(); + Ok(()) + } + fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, @@ -37,7 +48,18 @@ fn clamp_nonnegative_scalar(distance: f32) -> f32 { if distance < 0.0 { 0.0 } else { distance } } -impl LeafKernelMetric for L2 { +impl LeafMetric for L2 { + fn prepare_norms( + dots: MatrixView<'_, f32>, + norms: &mut Vec, + ) -> Result<(), TryReserveError> { + resize_norms(norms, dots.nrows())?; + for (point, norm) in norms.iter_mut().enumerate() { + *norm = dots[(point, point)]; + } + Ok(()) + } + #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -53,7 +75,18 @@ impl LeafKernelMetric for L2 { } } -impl LeafKernelMetric for Cosine { +impl LeafMetric for Cosine { + fn prepare_norms( + dots: MatrixView<'_, f32>, + norms: &mut Vec, + ) -> Result<(), TryReserveError> { + resize_norms(norms, dots.nrows())?; + for (point, norm) in norms.iter_mut().enumerate() { + *norm = norm_from_squared(dots[(point, point)]); + } + Ok(()) + } + #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -69,7 +102,7 @@ impl LeafKernelMetric for Cosine { } } -impl LeafKernelMetric for CosineNormalized { +impl LeafMetric for CosineNormalized { #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F where @@ -85,7 +118,7 @@ impl LeafKernelMetric for CosineNormalized { } } -impl LeafKernelMetric for InnerProduct { +impl LeafMetric for InnerProduct { #[inline(always)] fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F where diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index b672c2063d..4583a35344 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -3,13 +3,36 @@ * Licensed under the MIT license. */ +use std::collections::TryReserveError; + +use diskann_utils::views::MatrixView; +use diskann_vector::{Norm, norm::FastL2NormSquared}; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; -use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar}; +use super::{ + Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar, + norm_from_squared, resize_norms, +}; /// Partition formulas return ascending scores. /// L2 uses squared leader norms. Cosine uses point and leader norms. -pub(in super::super) trait PartitionKernelMetric: Send + Sync + 'static { +pub(in super::super) trait PartitionMetric: Send + Sync + 'static { + fn prepare_point_norms( + _: MatrixView<'_, f32>, + norms: &mut Vec, + ) -> Result<(), TryReserveError> { + norms.clear(); + Ok(()) + } + + fn prepare_leader_norms( + _: MatrixView<'_, f32>, + norms: &mut Vec, + ) -> Result<(), TryReserveError> { + norms.clear(); + Ok(()) + } + fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, @@ -18,7 +41,18 @@ pub(in super::super) trait PartitionKernelMetric: Send + Sync + 'static { fn partition_ranking_scalar(dot: f32, point_norm: f32, leader_norm: f32) -> f32; } -impl PartitionKernelMetric for L2 { +impl PartitionMetric for L2 { + fn prepare_leader_norms( + leader_values: MatrixView<'_, f32>, + norms: &mut Vec, + ) -> Result<(), TryReserveError> { + resize_norms(norms, leader_values.nrows())?; + for (norm, leader) in norms.iter_mut().zip(leader_values.row_iter()) { + *norm = leader.iter().map(|value| value * value).sum(); + } + Ok(()) + } + #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, leader_norm: F) -> F where @@ -37,7 +71,30 @@ impl PartitionKernelMetric for L2 { } } -impl PartitionKernelMetric for Cosine { +impl PartitionMetric for Cosine { + fn prepare_point_norms( + point_values: MatrixView<'_, f32>, + norms: &mut Vec, + ) -> Result<(), TryReserveError> { + resize_norms(norms, point_values.nrows())?; + for (norm, point) in norms.iter_mut().zip(point_values.row_iter()) { + *norm = norm_from_squared(FastL2NormSquared.evaluate(point)); + } + Ok(()) + } + + fn prepare_leader_norms( + leader_values: MatrixView<'_, f32>, + norms: &mut Vec, + ) -> Result<(), TryReserveError> { + resize_norms(norms, leader_values.nrows())?; + for (norm, leader) in norms.iter_mut().zip(leader_values.row_iter()) { + let squared_norm = leader.iter().map(|value| value * value).sum(); + *norm = norm_from_squared(squared_norm); + } + Ok(()) + } + #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F where @@ -53,7 +110,7 @@ impl PartitionKernelMetric for Cosine { } } -impl PartitionKernelMetric for CosineNormalized { +impl PartitionMetric for CosineNormalized { #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, _: F) -> F where @@ -69,7 +126,7 @@ impl PartitionKernelMetric for CosineNormalized { } } -impl PartitionKernelMetric for InnerProduct { +impl PartitionMetric for InnerProduct { #[inline(always)] fn partition_ranking(arch: F::Arch, dot: F, _: F, _: F) -> F where diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 8a691cd88f..2a09a4b167 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -24,7 +24,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; -use super::kernel_metric::LeafKernelMetric; +use super::kernel_metric::LeafMetric; /// Largest leaf-local neighbor count supported by the fixed insertion kernel. pub(super) const MAX_LEAF_NEIGHBORS: usize = 3; @@ -149,7 +149,7 @@ where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, - M: LeafKernelMetric, + M: LeafMetric, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { validate(input, norms, &output)?; @@ -307,7 +307,7 @@ fn scan_point_pairs( ) where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, - M: LeafKernelMetric, + M: LeafMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { let (output, _) = output.as_chunks_mut::(); diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 0307c8c091..c9cb07c8c8 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -23,7 +23,7 @@ use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, }; -use super::kernel_metric::PartitionKernelMetric; +use super::kernel_metric::PartitionMetric; /// Reusable nearest-center state for one partition worker. /// @@ -115,7 +115,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, { let norms = validate(metric, input, &output)?; let fanout = output.ncols(); @@ -230,7 +230,7 @@ fn select_point_leaders> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, - M: PartitionKernelMetric, + M: PartitionMetric, u64: From<<::BitMask as SIMDMask>::Underlying>, { let leader_count = dots.ncols(); @@ -412,7 +412,7 @@ fn dispatch_nearest_leaders( #[cfg(test)] mod tests { use super::super::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, PartitionKernelMetric, + Cosine, CosineNormalized, InnerProduct, L2, PartitionMetric, }; use super::*; @@ -436,7 +436,7 @@ mod tests { // This oracle checks SIMD chunking, scalar tails, and tracker order. It uses // the scalar ranking formula for metric `M`. - fn scalar_traversal_reference( + fn scalar_traversal_reference( metric: Metric, input: PartitionInput<'_>, fanout: usize, From c06d42c74c01fdf5a69677c01f99e06086ac40dc Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:37:41 +0000 Subject: [PATCH 73/80] refactor(pipnn): move norm policy into metric types --- diskann/src/graph/pipnn/partition_kernel.rs | 84 +++++++++------------ 1 file changed, 35 insertions(+), 49 deletions(-) diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index c9cb07c8c8..a9f22c3c06 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -18,6 +18,7 @@ //! leaders. Equal scores keep sampled-leader order. NaN is not rankable. use diskann_utils::views::{MatrixView, MutMatrixView}; +#[cfg(test)] use diskann_vector::distance::Metric; use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, @@ -105,7 +106,6 @@ pub(super) enum PartitionKernelError { /// It also returns an error when fewer than `fanout` scores are rankable. pub(super) fn nearest_leaders( arch: A, - metric: Metric, input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, workspace: &mut PartitionKernelWorkspace, @@ -117,43 +117,47 @@ where u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, M: PartitionMetric, { - let norms = validate(metric, input, &output)?; + let norms = validate(input, &output)?; let fanout = output.ncols(); if fanout == 0 || input.dots.nrows() == 0 { return Ok(()); } workspace.prepare(fanout)?; - match metric { - Metric::L2 => select_point_leaders::( + match (norms.point_norms.is_empty(), norms.leader_norms.is_empty()) { + (true, true) => select_point_leaders::( arch, input.dots, norms, output, &mut workspace.tracker, ), - Metric::Cosine => select_point_leaders::( + (true, false) => select_point_leaders::( + arch, + input.dots, + norms, + output, + &mut workspace.tracker, + ), + (false, false) => select_point_leaders::( + arch, + input.dots, + norms, + output, + &mut workspace.tracker, + ), + (false, true) => select_point_leaders::( arch, input.dots, norms, output, &mut workspace.tracker, ), - Metric::CosineNormalized | Metric::InnerProduct => { - select_point_leaders::( - arch, - input.dots, - norms, - output, - &mut workspace.tracker, - ) - } } } /// This function checks row counts, leader IDs, fanout, and norm lengths. fn validate<'a>( - metric: Metric, input: PartitionInput<'a>, output: &MutMatrixView<'_, u32>, ) -> Result, PartitionKernelError> { @@ -178,21 +182,12 @@ fn validate<'a>( }); } - let (point_norm_count, leader_norm_count) = match metric { - Metric::L2 => (0, leader_count), - Metric::Cosine => (point_count, leader_count), - Metric::CosineNormalized | Metric::InnerProduct => (0, 0), - }; - check_length( - "point norms", - input.norms.point_norms.len(), - point_norm_count, - )?; - check_length( - "leader norms", - input.norms.leader_norms.len(), - leader_norm_count, - )?; + if !input.norms.point_norms.is_empty() { + check_length("point norms", input.norms.point_norms.len(), point_count)?; + } + if !input.norms.leader_norms.is_empty() { + check_length("leader norms", input.norms.leader_norms.len(), leader_count)?; + } Ok(input.norms) } @@ -364,30 +359,19 @@ where use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => { - nearest_leaders::(arch, Metric::L2, call.input, call.output, call.workspace) + Metric::L2 => nearest_leaders::(arch, call.input, call.output, call.workspace), + Metric::Cosine => { + nearest_leaders::(arch, call.input, call.output, call.workspace) } - Metric::Cosine => nearest_leaders::( - arch, - Metric::Cosine, - call.input, - call.output, - call.workspace, - ), Metric::CosineNormalized => nearest_leaders::( arch, - Metric::CosineNormalized, - call.input, - call.output, - call.workspace, - ), - Metric::InnerProduct => nearest_leaders::( - arch, - Metric::InnerProduct, call.input, call.output, call.workspace, ), + Metric::InnerProduct => { + nearest_leaders::(arch, call.input, call.output, call.workspace) + } } } } @@ -416,6 +400,7 @@ mod tests { }; use super::*; + use diskann_vector::distance::Metric; fn test_input<'a>( _metric: Metric, @@ -884,11 +869,12 @@ mod integration_tests { }) ); + let short_norms = [0.0; 2]; let wrong_norms = PartitionInput { dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), norms: PartitionNorms { point_norms: &[], - leader_norms: &[], + leader_norms: &short_norms, }, }; assert_eq!( @@ -896,7 +882,7 @@ mod integration_tests { Err(PartitionKernelError::InvalidBufferLength { buffer: "leader norms", expected: 3, - actual: 0, + actual: 2, }) ); From bb25e6f745ed0a88148a527178c6cefcfd947e29 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:58:52 +0000 Subject: [PATCH 74/80] test(utils): cover zero-wrapping matrix area --- diskann-utils/src/views.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs index 1f0cfdd4d3..e38b5b6c58 100644 --- a/diskann-utils/src/views.rs +++ b/diskann-utils/src/views.rs @@ -1054,7 +1054,7 @@ mod tests { "tried to construct a matrix view with 5 rows and 4 columns over a slice of length 12" ); - assert!(MatrixView::try_from(&[] as &[usize], usize::MAX, 2).is_err()); + assert!(MatrixView::try_from(&[] as &[usize], usize::MAX / 2 + 1, 2).is_err()); } #[test] From 45364af12f35173d37fd48a0b244ceaa7c95673f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:18:22 +0000 Subject: [PATCH 75/80] refactor(pipnn): give metrics exact ranking inputs --- diskann/src/graph/pipnn/kernel_metric.rs | 20 +- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 103 ++-- .../graph/pipnn/kernel_metric/partition.rs | 173 ++++--- diskann/src/graph/pipnn/leaf_kernel.rs | 159 +++++-- diskann/src/graph/pipnn/mod.rs | 6 +- diskann/src/graph/pipnn/partition_kernel.rs | 443 ++++++++++-------- 6 files changed, 543 insertions(+), 361 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 4c832f3106..52f03ea9cf 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -13,6 +13,7 @@ pub(super) use partition::PartitionMetric; use std::collections::TryReserveError; +use diskann_utils::views::MatrixView; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; pub(super) struct L2; @@ -20,6 +21,18 @@ pub(super) struct Cosine; pub(super) struct CosineNormalized; pub(super) struct InnerProduct; +pub(super) struct NormPreparation<'a, 'b> { + pub(super) values: MatrixView<'a, f32>, + pub(super) norms: &'b mut Vec, +} + +/// Prepared norms for one point stripe and its sampled leaders. +#[derive(Clone, Copy, Debug)] +pub(super) struct PartitionNorms<'a> { + pub(super) point_norms: &'a [f32], + pub(super) leader_norms: &'a [f32], +} + pub(super) fn resize_norms(norms: &mut Vec, len: usize) -> Result<(), TryReserveError> { norms.try_reserve(len.saturating_sub(norms.len()))?; norms.resize(len, 0.0); @@ -38,12 +51,12 @@ pub(super) fn norm_from_squared(squared_norm: f32) -> f32 { } } -/// This function computes cosine distance with the DiskANN zero-norm and NaN rules. +/// Compute SIMD cosine distance with the DiskANN zero-norm and NaN rules. /// /// Each lane contains one point pair. A zero norm produces zero similarity. A /// NaN norm remains NaN unless the other norm is zero. #[inline(always)] -pub(super) fn cosine_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F +pub(super) fn cosine_distance_simd(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, @@ -59,8 +72,9 @@ where one - cosine } +/// Compute one cosine distance with the DiskANN zero-norm and NaN rules. #[inline(always)] -pub(super) fn cosine_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32 { +pub(super) fn cosine_distance_single(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 { diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 7cc3996247..af8d042b80 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -5,33 +5,32 @@ use std::collections::TryReserveError; -use diskann_utils::views::MatrixView; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; use super::{ - Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar, - norm_from_squared, resize_norms, + Cosine, CosineNormalized, InnerProduct, L2, NormPreparation, cosine_distance_simd, + cosine_distance_single, norm_from_squared, resize_norms, }; /// Leaf formulas return ascending distances. /// L2 uses squared norms. Cosine uses norms. Other metrics ignore norms. pub(in super::super) trait LeafMetric: Send + Sync + 'static { - fn prepare_norms(_: MatrixView<'_, f32>, norms: &mut Vec) -> Result<(), TryReserveError> { - norms.clear(); - Ok(()) - } + /// Prepare one contiguous metric-specific norm for each leaf-local point. + fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError>; - fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F + /// Compute distances for one complete SIMD group. + fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; - fn leaf_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32; + /// Compute one distance outside the complete SIMD prefix. + fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32; } -/// This function clamps negative SIMD roundoff to zero and preserves NaN lanes. +/// Clamp negative SIMD roundoff to zero and preserve NaN lanes. #[inline(always)] -fn clamp_nonnegative(arch: F::Arch, distance: F) -> F +fn clamp_nonnegative_simd(arch: F::Arch, distance: F) -> F where F: SIMDVector + SIMDFloat, F::Mask: SIMDSelect, @@ -42,94 +41,112 @@ where .select(zero.max_simd(distance), distance) } -/// This function clamps negative scalar roundoff to zero and preserves NaN. +/// Clamp negative roundoff to zero and preserve NaN. #[inline(always)] -fn clamp_nonnegative_scalar(distance: f32) -> f32 { +fn clamp_nonnegative_single(distance: f32) -> f32 { if distance < 0.0 { 0.0 } else { distance } } impl LeafMetric for L2 { - fn prepare_norms( - dots: MatrixView<'_, f32>, - norms: &mut Vec, - ) -> Result<(), TryReserveError> { - resize_norms(norms, dots.nrows())?; - for (point, norm) in norms.iter_mut().enumerate() { - *norm = dots[(point, point)]; + fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + resize_norms(preparation.norms, preparation.values.nrows())?; + for (point, norm) in preparation.norms.iter_mut().enumerate() { + *norm = preparation.values[(point, point)]; } Ok(()) } #[inline(always)] - fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F + fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - clamp_nonnegative(arch, source_norm + target_norm - F::splat(arch, 2.0) * dot) + clamp_nonnegative_simd( + arch, + source_norms + target_norms - F::splat(arch, 2.0) * dot_products, + ) } #[inline(always)] - fn leaf_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32 { - clamp_nonnegative_scalar(source_norm + target_norm - 2.0 * dot) + fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32 { + clamp_nonnegative_single(source_norm + target_norm - 2.0 * dot_product) } } impl LeafMetric for Cosine { - fn prepare_norms( - dots: MatrixView<'_, f32>, - norms: &mut Vec, - ) -> Result<(), TryReserveError> { - resize_norms(norms, dots.nrows())?; - for (point, norm) in norms.iter_mut().enumerate() { - *norm = norm_from_squared(dots[(point, point)]); + fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + resize_norms(preparation.norms, preparation.values.nrows())?; + for (point, norm) in preparation.norms.iter_mut().enumerate() { + *norm = norm_from_squared(preparation.values[(point, point)]); } Ok(()) } #[inline(always)] - fn leaf_distance(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F + fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - clamp_nonnegative(arch, cosine_distance(arch, dot, source_norm, target_norm)) + clamp_nonnegative_simd( + arch, + cosine_distance_simd(arch, dot_products, source_norms, target_norms), + ) } #[inline(always)] - fn leaf_distance_scalar(dot: f32, source_norm: f32, target_norm: f32) -> f32 { - clamp_nonnegative_scalar(cosine_distance_scalar(dot, source_norm, target_norm)) + fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32 { + clamp_nonnegative_single(cosine_distance_single( + dot_product, + source_norm, + target_norm, + )) } } impl LeafMetric for CosineNormalized { + fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + preparation.norms.clear(); + Ok(()) + } + #[inline(always)] - fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F + fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - clamp_nonnegative(arch, F::splat(arch, 1.0) - dot) + let _ = (source_norms, target_norms); + clamp_nonnegative_simd(arch, F::splat(arch, 1.0) - dot_products) } #[inline(always)] - fn leaf_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { - clamp_nonnegative_scalar(1.0 - dot) + fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32 { + let _ = (source_norm, target_norm); + clamp_nonnegative_single(1.0 - dot_product) } } impl LeafMetric for InnerProduct { + fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + preparation.norms.clear(); + Ok(()) + } + #[inline(always)] - fn leaf_distance(arch: F::Arch, dot: F, _: F, _: F) -> F + fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - F::default(arch) - dot + let _ = (source_norms, target_norms); + F::default(arch) - dot_products } #[inline(always)] - fn leaf_distance_scalar(dot: f32, _: f32, _: f32) -> f32 { - -dot + fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32 { + let _ = (source_norm, target_norm); + -dot_product } } diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 4583a35344..bfe99faffd 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -5,90 +5,100 @@ use std::collections::TryReserveError; -use diskann_utils::views::MatrixView; use diskann_vector::{Norm, norm::FastL2NormSquared}; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; use super::{ - Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, cosine_distance_scalar, - norm_from_squared, resize_norms, + Cosine, CosineNormalized, InnerProduct, L2, NormPreparation, cosine_distance_simd, + cosine_distance_single, norm_from_squared, resize_norms, }; -/// Partition formulas return ascending scores. +/// Partition formulas return ascending rankings. /// L2 uses squared leader norms. Cosine uses point and leader norms. pub(in super::super) trait PartitionMetric: Send + Sync + 'static { - fn prepare_point_norms( - _: MatrixView<'_, f32>, - norms: &mut Vec, - ) -> Result<(), TryReserveError> { - norms.clear(); - Ok(()) - } - - fn prepare_leader_norms( - _: MatrixView<'_, f32>, - norms: &mut Vec, - ) -> Result<(), TryReserveError> { - norms.clear(); - Ok(()) - } - - fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F + /// Prepare one norm value for each point in the active stripe. + fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError>; + + /// Prepare one norm value for each sampled leader. + fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError>; + + /// Compute rankings for one complete SIMD group. + fn partition_ranking_simd( + arch: F::Arch, + dot_products: F, + point_norms: F, + leader_norms: F, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; - fn partition_ranking_scalar(dot: f32, point_norm: f32, leader_norm: f32) -> f32; + /// Compute one ranking outside the complete SIMD prefix. + fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32; } impl PartitionMetric for L2 { - fn prepare_leader_norms( - leader_values: MatrixView<'_, f32>, - norms: &mut Vec, - ) -> Result<(), TryReserveError> { - resize_norms(norms, leader_values.nrows())?; - for (norm, leader) in norms.iter_mut().zip(leader_values.row_iter()) { + fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + preparation.norms.clear(); + Ok(()) + } + + fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + resize_norms(preparation.norms, preparation.values.nrows())?; + for (norm, leader) in preparation + .norms + .iter_mut() + .zip(preparation.values.row_iter()) + { *norm = leader.iter().map(|value| value * value).sum(); } Ok(()) } #[inline(always)] - fn partition_ranking(arch: F::Arch, dot: F, _: F, leader_norm: F) -> F + fn partition_ranking_simd( + arch: F::Arch, + dot_products: F, + point_norms: F, + leader_norms: F, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - // The point norm is constant for all leaders. Fused arithmetic defines - // the ranking order for complete SIMD groups. - F::splat(arch, -2.0).mul_add_simd(dot, leader_norm) + let _ = point_norms; + // Fused arithmetic defines the ranking order for complete SIMD groups. + F::splat(arch, -2.0).mul_add_simd(dot_products, leader_norms) } #[inline(always)] - fn partition_ranking_scalar(dot: f32, _: f32, leader_norm: f32) -> f32 { - // Non-fused arithmetic defines the ranking order for the scalar tail. - leader_norm - 2.0 * dot + fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32 { + let _ = point_norm; + // Non-fused arithmetic defines the ranking order outside the SIMD prefix. + leader_norm - 2.0 * dot_product } } impl PartitionMetric for Cosine { - fn prepare_point_norms( - point_values: MatrixView<'_, f32>, - norms: &mut Vec, - ) -> Result<(), TryReserveError> { - resize_norms(norms, point_values.nrows())?; - for (norm, point) in norms.iter_mut().zip(point_values.row_iter()) { + fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + resize_norms(preparation.norms, preparation.values.nrows())?; + for (norm, point) in preparation + .norms + .iter_mut() + .zip(preparation.values.row_iter()) + { *norm = norm_from_squared(FastL2NormSquared.evaluate(point)); } Ok(()) } - fn prepare_leader_norms( - leader_values: MatrixView<'_, f32>, - norms: &mut Vec, - ) -> Result<(), TryReserveError> { - resize_norms(norms, leader_values.nrows())?; - for (norm, leader) in norms.iter_mut().zip(leader_values.row_iter()) { + fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + resize_norms(preparation.norms, preparation.values.nrows())?; + for (norm, leader) in preparation + .norms + .iter_mut() + .zip(preparation.values.row_iter()) + { let squared_norm = leader.iter().map(|value| value * value).sum(); *norm = norm_from_squared(squared_norm); } @@ -96,49 +106,88 @@ impl PartitionMetric for Cosine { } #[inline(always)] - fn partition_ranking(arch: F::Arch, dot: F, point_norm: F, leader_norm: F) -> F + fn partition_ranking_simd( + arch: F::Arch, + dot_products: F, + point_norms: F, + leader_norms: F, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - cosine_distance(arch, dot, point_norm, leader_norm) + cosine_distance_simd(arch, dot_products, point_norms, leader_norms) } #[inline(always)] - fn partition_ranking_scalar(dot: f32, point_norm: f32, leader_norm: f32) -> f32 { - cosine_distance_scalar(dot, point_norm, leader_norm) + fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32 { + cosine_distance_single(dot_product, point_norm, leader_norm) } } impl PartitionMetric for CosineNormalized { + fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + preparation.norms.clear(); + Ok(()) + } + + fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + preparation.norms.clear(); + Ok(()) + } + #[inline(always)] - fn partition_ranking(arch: F::Arch, dot: F, _: F, _: F) -> F + fn partition_ranking_simd( + arch: F::Arch, + dot_products: F, + point_norms: F, + leader_norms: F, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - F::splat(arch, 1.0) - dot + let _ = (point_norms, leader_norms); + F::splat(arch, 1.0) - dot_products } #[inline(always)] - fn partition_ranking_scalar(dot: f32, _: f32, _: f32) -> f32 { - 1.0 - dot + fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32 { + let _ = (point_norm, leader_norm); + 1.0 - dot_product } } impl PartitionMetric for InnerProduct { + fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + preparation.norms.clear(); + Ok(()) + } + + fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { + preparation.norms.clear(); + Ok(()) + } + #[inline(always)] - fn partition_ranking(arch: F::Arch, dot: F, _: F, _: F) -> F + fn partition_ranking_simd( + arch: F::Arch, + dot_products: F, + point_norms: F, + leader_norms: F, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - F::default(arch) - dot + let _ = (point_norms, leader_norms); + F::default(arch) - dot_products } #[inline(always)] - fn partition_ranking_scalar(dot: f32, _: f32, _: f32) -> f32 { - -dot + fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32 { + let _ = (point_norm, leader_norm); + -dot_product } } @@ -147,11 +196,11 @@ mod tests { use super::*; #[test] - fn l2_scalar_ranking_preserves_non_fused_rounding() { - let scalar = L2::partition_ranking_scalar(f32::MAX, 0.0, f32::MAX); + fn l2_ranking_preserves_non_fused_arithmetic() { + let ranking = L2::partition_ranking_single(f32::MAX, 0.0, f32::MAX); let fused = (-2.0f32).mul_add(f32::MAX, f32::MAX); - assert_eq!(scalar, f32::NEG_INFINITY); + assert_eq!(ranking, f32::NEG_INFINITY); assert_eq!(fused, -f32::MAX); } } diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 2a09a4b167..45028551c9 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -14,12 +14,12 @@ //! [`MAX_LEAF_NEIGHBORS`]. Positive widths use fixed arrays. //! //! Strict comparisons keep scan order for equal distances. They do not rank NaN. -//! All supported metrics use the same scalar and SIMD traversal. +//! All supported metrics use the same SIMD-group and single-value traversal. //! //! The caller supplies concrete architecture `A` and metric `M`. The function //! checks all shapes and local-ID bounds before it changes workspace or uses an -//! unchecked SIMD load. [`LeafKernelWorkspace`] stores reusable norms and -//! rejection thresholds. +//! unchecked SIMD load. [`LeafKernelWorkspace`] stores reusable rejection +//! thresholds. use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; @@ -168,46 +168,46 @@ where workspace.worst.fill(f32::INFINITY); match (norms.is_empty(), neighbor_count) { - (false, 1) => scan_point_pairs::( + (false, 1) => scan_point_pairs::( arch, input, output.as_mut_slice(), - norms, + PreparedLeafNorms(norms), &mut workspace.worst, ), - (false, 2) => scan_point_pairs::( + (false, 2) => scan_point_pairs::( arch, input, output.as_mut_slice(), - norms, + PreparedLeafNorms(norms), &mut workspace.worst, ), - (false, 3) => scan_point_pairs::( + (false, 3) => scan_point_pairs::( arch, input, output.as_mut_slice(), - norms, + PreparedLeafNorms(norms), &mut workspace.worst, ), - (true, 1) => scan_point_pairs::( + (true, 1) => scan_point_pairs::( arch, input, output.as_mut_slice(), - norms, + EmptyLeafNorms, &mut workspace.worst, ), - (true, 2) => scan_point_pairs::( + (true, 2) => scan_point_pairs::( arch, input, output.as_mut_slice(), - norms, + EmptyLeafNorms, &mut workspace.worst, ), - (true, 3) => scan_point_pairs::( + (true, 3) => scan_point_pairs::( arch, input, output.as_mut_slice(), - norms, + EmptyLeafNorms, &mut workspace.worst, ), _ => { @@ -291,23 +291,94 @@ fn resize( Ok(()) } +/// Provide norm values for one leaf scan. +trait LeafNormAccess +where + F: SIMDVector, +{ + /// Repeat one point norm in all SIMD lanes. + fn repeat_simd(self, arch: F::Arch, point: usize) -> F; + + /// Load one complete SIMD group of point norms. + fn load_simd(self, arch: F::Arch, first_point: usize) -> F; + + /// Read one point norm. + fn read(self, point: usize) -> f32; +} + +/// Prepared norm values for all points in one leaf. +#[derive(Clone, Copy)] +struct PreparedLeafNorms<'a>(&'a [f32]); + +impl LeafNormAccess for PreparedLeafNorms<'_> +where + F: SIMDVector, +{ + #[inline(always)] + fn repeat_simd(self, arch: F::Arch, point: usize) -> F { + F::splat(arch, self.0[point]) + } + + #[inline(always)] + fn load_simd(self, arch: F::Arch, first_point: usize) -> F { + let last_point = first_point + F::LANES; + let norm_group = &self.0[first_point..last_point]; + + // SAFETY: `norm_group` contains one complete SIMD group. + unsafe { F::load_simd(arch, norm_group.as_ptr()) } + } + + #[inline(always)] + fn read(self, point: usize) -> f32 { + self.0[point] + } +} + +/// Zero norm values for a metric that does not use leaf norms. +#[derive(Clone, Copy)] +struct EmptyLeafNorms; + +impl LeafNormAccess for EmptyLeafNorms +where + F: SIMDVector, +{ + #[inline(always)] + fn repeat_simd(self, arch: F::Arch, point: usize) -> F { + let _ = point; + F::default(arch) + } + + #[inline(always)] + fn load_simd(self, arch: F::Arch, first_point: usize) -> F { + let _ = first_point; + F::default(arch) + } + + #[inline(always)] + fn read(self, point: usize) -> f32 { + let _ = point; + 0.0 + } +} + /// Select neighbors from all unordered point pairs in one leaf. /// /// The function reads the strict lower triangle once. It offers each distance to -/// both endpoint lists. SIMD groups and the scalar tail preserve pair scan order. +/// both endpoint lists. SIMD groups and single values preserve pair scan order. /// -/// `input` supplies square dot products. `norms` contains prepared norm values. +/// `input` supplies square dot products. `norms` supplies metric norm values. #[inline(never)] -fn scan_point_pairs( +fn scan_point_pairs( arch: F::Arch, input: MatrixView<'_, f32>, output: &mut [LeafNeighbor], - norms: &[f32], + norms: R, worst: &mut [f32], ) where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: LeafMetric, + R: LeafNormAccess + Copy, u64: From<<::BitMask as SIMDMask>::Underlying>, { let (output, _) = output.as_chunks_mut::(); @@ -319,27 +390,22 @@ fn scan_point_pairs( // itself to the neighbor list of source zero. for source in 1..point_count { let source_start = source * point_count; - let source_norm = if HAS_NORMS { norms[source] } else { 0.0 }; - let source_norms = F::splat(arch, source_norm); + let source_norms = norms.repeat_simd(arch, source); + let source_norm = norms.read(source); // SAFETY: `nearest_neighbors` created one threshold for each point. let mut source_worst = unsafe { *worst_ptr.add(source) }; let mut target = 0; let full = source / F::LANES * F::LANES; while target < full { - // SAFETY: the full chunk is contained in this source's strict-lower prefix. + // SAFETY: The full chunk is in this source's strict-lower prefix. let pair_dots = unsafe { F::load_simd(arch, dots.as_ptr().add(source_start + target)) }; - let target_norms = if HAS_NORMS { - // SAFETY: the full target chunk is below `source < point_count`. - unsafe { F::load_simd(arch, norms.as_ptr().add(target)) } - } else { - F::default(arch) - }; - let distances = M::leaf_distance(arch, pair_dots, source_norms, target_norms); + let target_norms = norms.load_simd(arch, target); + let distances = M::leaf_distance_simd(arch, pair_dots, source_norms, target_norms); // 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 is below `source < point_count`. + // SAFETY: The full target chunk is below `source < point_count`. // `nearest_neighbors` created one threshold for each point. let target_worst = unsafe { F::load_simd(arch, worst_ptr.add(target)) }; let target_eligible = distances.lt_simd(target_worst); @@ -380,15 +446,10 @@ fn scan_point_pairs( } while target < source { - // SAFETY: the scalar target remains in this source's strict-lower prefix. + // SAFETY: The target is in this source's strict-lower prefix. let dot = unsafe { *dots.get_unchecked(source_start + target) }; - let target_norm = if HAS_NORMS { - // SAFETY: `target < source < point_count == norms.len()`. - unsafe { *norms.get_unchecked(target) } - } else { - 0.0 - }; - let distance = M::leaf_distance_scalar(dot, source_norm, target_norm); + let target_norm = norms.read(target); + let distance = M::leaf_distance_single(dot, source_norm, target_norm); if distance < source_worst { source_worst = insert_fixed_neighbor(&mut output[source], target as u32, distance); } @@ -534,16 +595,24 @@ fn prepared_test_norms( metric: diskann_vector::distance::Metric, input: MatrixView<'_, f32>, ) -> Vec { + use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2, NormPreparation}; use diskann_vector::distance::Metric; + fn prepare(input: MatrixView<'_, f32>) -> Vec { + let mut norms = Vec::new(); + M::prepare_leaf_norms(NormPreparation { + values: input, + norms: &mut norms, + }) + .unwrap(); + norms + } + match metric { - Metric::L2 => (0..input.nrows()) - .map(|point| input[(point, point)]) - .collect(), - Metric::Cosine => (0..input.nrows()) - .map(|point| super::kernel_metric::norm_from_squared(input[(point, point)])) - .collect(), - Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), + Metric::L2 => prepare::(input), + Metric::Cosine => prepare::(input), + Metric::CosineNormalized => prepare::(input), + Metric::InnerProduct => prepare::(input), } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 6da4dd1e14..676c00415f 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -6,15 +6,15 @@ //! Numerical kernels for PiPNN graph construction. //! //! [`partition_kernel`] converts point-to-leader dot products into sorted leader -//! positions. The output width sets the fanout. One workspace stores the -//! runtime-sized tracker and reuses it for each point. +//! positions. The output width sets the fanout. A scratch vector stores the +//! ranked leaders and reuses its allocation for each point. //! //! [`leaf_kernel`] reads a lower-triangular Gram matrix. It evaluates each point //! pair once and updates both points. Each point retains at most three local //! neighbors. //! //! `kernel_metric` defines metric markers and shared math. Separate leaf and -//! partition traits define each kernel's scalar and SIMD formulas. +//! partition traits define norm preparation and ranking formulas. //! //! The graph builder selects architecture `A` and metric `M` once. It passes //! these concrete types to both kernels. diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index a9f22c3c06..1d02179dc2 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -24,33 +24,7 @@ use diskann_wide::{ Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, }; -use super::kernel_metric::PartitionMetric; - -/// Reusable nearest-center state for one partition worker. -/// -/// Each entry contains a sampled leader's matrix-column ID and its metric score. -#[derive(Debug, Default)] -pub(super) struct PartitionKernelWorkspace { - tracker: Vec<(u32, f32)>, -} - -impl PartitionKernelWorkspace { - fn prepare(&mut self, fanout: usize) -> Result<(), PartitionKernelError> { - let additional = fanout.saturating_sub(self.tracker.len()); - self.tracker - .try_reserve(additional) - .map_err(|_| PartitionKernelError::Allocation { additional })?; - self.tracker.resize(fanout, (u32::MAX, f32::INFINITY)); - Ok(()) - } -} - -/// L2 uses squared leader norms. Cosine uses point and leader norms. -#[derive(Clone, Copy, Debug)] -pub(super) struct PartitionNorms<'a> { - pub(super) point_norms: &'a [f32], - pub(super) leader_norms: &'a [f32], -} +use super::kernel_metric::{PartitionMetric, PartitionNorms}; /// Dot products between assigned points and sampled partition centers. /// @@ -84,8 +58,8 @@ pub(super) enum PartitionKernelError { /// The requested fanout exceeds the available leader count. #[error("invalid fanout {fanout}: must not exceed {leader_count} leaders")] InvalidFanout { fanout: usize, leader_count: usize }, - /// Reusable tracker storage could not be reserved. - #[error("failed to reserve {additional} partition tracker entries")] + /// Reusable ranked-leader storage could not be reserved. + #[error("failed to reserve {additional} partition ranked-leader entries")] Allocation { additional: usize }, /// Leader positions cannot be represented as `u32`. #[error("leader count {0} exceeds the u32 position limit")] @@ -108,7 +82,7 @@ pub(super) fn nearest_leaders( arch: A, input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, - workspace: &mut PartitionKernelWorkspace, + ranked_leaders: &mut Vec<(u32, f32)>, ) -> Result<(), PartitionKernelError> where A: Architecture, @@ -123,35 +97,44 @@ where return Ok(()); } - workspace.prepare(fanout)?; + let additional = fanout.saturating_sub(ranked_leaders.len()); + ranked_leaders + .try_reserve(additional) + .map_err(|_| PartitionKernelError::Allocation { additional })?; + ranked_leaders.resize(fanout, (u32::MAX, f32::INFINITY)); + match (norms.point_norms.is_empty(), norms.leader_norms.is_empty()) { - (true, true) => select_point_leaders::( + (false, false) => select_point_leaders::( arch, input.dots, - norms, + PreparedNorms(norms.point_norms), + PreparedNorms(norms.leader_norms), output, - &mut workspace.tracker, + ranked_leaders, ), - (true, false) => select_point_leaders::( + (false, true) => select_point_leaders::( arch, input.dots, - norms, + PreparedNorms(norms.point_norms), + EmptyNorms, output, - &mut workspace.tracker, + ranked_leaders, ), - (false, false) => select_point_leaders::( + (true, false) => select_point_leaders::( arch, input.dots, - norms, + EmptyNorms, + PreparedNorms(norms.leader_norms), output, - &mut workspace.tracker, + ranked_leaders, ), - (false, true) => select_point_leaders::( + (true, true) => select_point_leaders::( arch, input.dots, - norms, + EmptyNorms, + EmptyNorms, output, - &mut workspace.tracker, + ranked_leaders, ), } } @@ -182,125 +165,179 @@ fn validate<'a>( }); } - if !input.norms.point_norms.is_empty() { - check_length("point norms", input.norms.point_norms.len(), point_count)?; - } - if !input.norms.leader_norms.is_empty() { - check_length("leader norms", input.norms.leader_norms.len(), leader_count)?; - } + check_norm_count("point norms", input.norms.point_norms, point_count)?; + check_norm_count("leader norms", input.norms.leader_norms, leader_count)?; Ok(input.norms) } -fn check_length( +/// Check one optional norm buffer. +/// +/// An empty slice means that the active metric does not use this norm. +fn check_norm_count( buffer: &'static str, - actual: usize, + norms: &[f32], expected: usize, ) -> Result<(), PartitionKernelError> { - if actual == expected { + if norms.is_empty() || norms.len() == expected { Ok(()) } else { Err(PartitionKernelError::InvalidBufferLength { buffer, expected, - actual, + actual: norms.len(), }) } } +/// Provide norm values for partition ranking. +trait NormValues +where + F: SIMDVector, +{ + /// Repeat one norm in all SIMD lanes. + fn repeat_simd(self, arch: F::Arch, point: usize) -> F; + + /// Load one complete SIMD group of norms. + fn load_simd(self, arch: F::Arch, first_point: usize) -> F; + + /// Read one norm. + fn read(self, point: usize) -> f32; +} + +/// Prepared norm values for points or sampled leaders. +#[derive(Clone, Copy)] +struct PreparedNorms<'a>(&'a [f32]); + +impl NormValues for PreparedNorms<'_> +where + F: SIMDVector, +{ + #[inline(always)] + fn repeat_simd(self, arch: F::Arch, point: usize) -> F { + F::splat(arch, self.0[point]) + } + + #[inline(always)] + fn load_simd(self, arch: F::Arch, first_point: usize) -> F { + let last_point = first_point + F::LANES; + let norm_group = &self.0[first_point..last_point]; + + // SAFETY: `norm_group` contains one complete SIMD group. + unsafe { F::load_simd(arch, norm_group.as_ptr()) } + } + + #[inline(always)] + fn read(self, point: usize) -> f32 { + self.0[point] + } +} + +/// Zero norm values for a metric that does not use one norm type. +#[derive(Clone, Copy)] +struct EmptyNorms; + +impl NormValues for EmptyNorms +where + F: SIMDVector, +{ + #[inline(always)] + fn repeat_simd(self, arch: F::Arch, point: usize) -> F { + let _ = point; + F::default(arch) + } + + #[inline(always)] + fn load_simd(self, arch: F::Arch, first_point: usize) -> F { + let _ = first_point; + F::default(arch) + } + + #[inline(always)] + fn read(self, point: usize) -> f32 { + let _ = point; + 0.0 + } +} + /// Rank sampled partition centers for each assigned point. /// /// The function converts point-to-leader dot products to metric `M` scores. It /// keeps the nearest `output.ncols()` centers in sampled-leader order for ties. /// NaN and positive infinity are not rankable. /// -/// `tracker` stores the retained center-column IDs and scores for the current +/// `ranked_leaders` stores the retained center-column IDs and scores for the current /// point. The function resets this state before it processes another point. -fn select_point_leaders( +fn select_point_leaders( arch: F::Arch, dots: MatrixView<'_, f32>, - norms: PartitionNorms<'_>, + point_norms: P, + leader_norms: L, mut output: MutMatrixView<'_, u32>, - tracker: &mut [(u32, f32)], + ranked_leaders: &mut [(u32, f32)], ) -> Result<(), PartitionKernelError> where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: PartitionMetric, + P: NormValues + Copy, + L: NormValues + Copy, u64: From<<::BitMask as SIMDMask>::Underlying>, { let leader_count = dots.ncols(); let fanout = output.ncols(); - // Reset the tracker for each point. No assignment state can pass from one - // output row to another. + // Reset the retained leaders for each point. No assignment state can pass + // from one output row to another. for (point, (point_dots, point_output)) in dots .row_iter() .zip(output.as_mut_slice().chunks_exact_mut(fanout)) .enumerate() { - tracker.fill((u32::MAX, f32::INFINITY)); - let point_norm = if HAS_POINT_NORMS { - norms.point_norms[point] - } else { - 0.0 - }; - let point_norm_vector = F::splat(arch, point_norm); - // Process all complete SIMD groups first. The scalar tail uses the - // metric's scalar operation order. + ranked_leaders.fill((u32::MAX, f32::INFINITY)); + let point_norm_values = point_norms.repeat_simd(arch, point); + let point_norm = point_norms.read(point); + // Process all complete SIMD groups first. Single-value rankings use the + // non-vector operation order. let full = leader_count / F::LANES * F::LANES; - for base in (0..full).step_by(F::LANES) { - // SAFETY: `base + F::LANES <= full <= point_dots.len()`. - let point_dots = unsafe { F::load_simd(arch, point_dots.as_ptr().add(base)) }; - let leader_norms = if HAS_LEADER_NORMS { - // SAFETY: `validate` established one norm value per leader. - // `base + F::LANES <= full <= leader_count`. - unsafe { F::load_simd(arch, norms.leader_norms.as_ptr().add(base)) } - } else { - F::default(arch) - }; - insert_leader_lanes( - M::partition_ranking(arch, point_dots, point_norm_vector, leader_norms), - base, - tracker, - ); + for first_leader in (0..full).step_by(F::LANES) { + // SAFETY: `first_leader + F::LANES <= full <= point_dots.len()`. + let dot = unsafe { F::load_simd(arch, point_dots.as_ptr().add(first_leader)) }; + let leader_norm_values = leader_norms.load_simd(arch, first_leader); + let scores = + M::partition_ranking_simd(arch, dot, point_norm_values, leader_norm_values); + insert_leader_lanes(scores, first_leader, ranked_leaders); } - // Use scalar formulas for the tail. A padded SIMD load can read past the - // norm slice and can change L2 rounding. + // Use single-value formulas because SIMD padding can change L2 rounding. for (leader, &dot) in point_dots.iter().enumerate().skip(full) { - let leader_norm = if HAS_LEADER_NORMS { - norms.leader_norms[leader] - } else { - 0.0 - }; + let leader_norm = leader_norms.read(leader); insert_leader( - tracker, + ranked_leaders, leader as u32, - M::partition_ranking_scalar(dot, point_norm, leader_norm), + M::partition_ranking_single(dot, point_norm, leader_norm), ); } - if tracker[fanout - 1].0 == u32::MAX { + if ranked_leaders[fanout - 1].0 == u32::MAX { return Err(PartitionKernelError::InsufficientRankableLeaders { point, fanout }); } - // Scatter needs the sampled-center column IDs. Metric scores remain in - // the worker workspace. - for (destination, &(leader, _)) in point_output.iter_mut().zip(tracker.iter()) { + // Scatter needs sampled-center column IDs. Scores remain in scratch. + for (destination, &(leader, _)) in point_output.iter_mut().zip(ranked_leaders.iter()) { *destination = leader; } } Ok(()) } -/// Offer one SIMD group of sampled centers to the current point's tracker. +/// Offer one SIMD group of sampled centers to the current point's ranked_leaders. /// /// `first_leader` is the matrix-column ID of the first lane. Lanes enter in /// sampled-leader order, which preserves tie order. -fn insert_leader_lanes(scores: F, first_leader: usize, tracker: &mut [(u32, f32)]) +fn insert_leader_lanes(scores: F, first_leader: usize, ranked_leaders: &mut [(u32, f32)]) where F: SIMDVector> + SIMDPartialOrd, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let threshold = F::splat(scores.arch(), tracker[tracker.len() - 1].1); + let threshold = F::splat(scores.arch(), ranked_leaders[ranked_leaders.len() - 1].1); let eligible = scores.lt_simd(threshold); if eligible.none() { return; @@ -311,26 +348,26 @@ where while lanes != 0 { let lane = lanes.trailing_zeros() as usize; lanes &= lanes - 1; - insert_leader(tracker, (first_leader + lane) as u32, values[lane]); + insert_leader(ranked_leaders, (first_leader + lane) as u32, values[lane]); } } /// Insert one sampled partition center into the current point's retained set. /// -/// `leader` is the center's column ID in the point-to-leader matrix. `tracker` +/// `leader` is the center's column ID in the point-to-leader matrix. `ranked_leaders` /// stores retained centers in nearest-first order. Equal scores and NaN do not /// enter, so sampled-leader order resolves ties. #[inline(always)] -fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, score: f32) { - let threshold = tracker.len() - 1; - if score.partial_cmp(&tracker[threshold].1) != Some(std::cmp::Ordering::Less) { +fn insert_leader(ranked_leaders: &mut [(u32, f32)], leader: u32, score: f32) { + let threshold = ranked_leaders.len() - 1; + if score.partial_cmp(&ranked_leaders[threshold].1) != Some(std::cmp::Ordering::Less) { return; } - tracker[threshold] = (leader, score); + ranked_leaders[threshold] = (leader, score); let mut slot = threshold; - while slot > 0 && tracker[slot].1 < tracker[slot - 1].1 { - tracker.swap(slot, slot - 1); + while slot > 0 && ranked_leaders[slot].1 < ranked_leaders[slot - 1].1 { + ranked_leaders.swap(slot, slot - 1); slot -= 1; } } @@ -339,7 +376,7 @@ fn insert_leader(tracker: &mut [(u32, f32)], leader: u32, score: f32) { struct DispatchedPartitionCall<'a> { input: PartitionInput<'a>, output: MutMatrixView<'a, u32>, - workspace: &'a mut PartitionKernelWorkspace, + ranked_leaders: &'a mut Vec<(u32, f32)>, } #[cfg(test)] @@ -359,19 +396,24 @@ where use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => nearest_leaders::(arch, call.input, call.output, call.workspace), + Metric::L2 => { + nearest_leaders::(arch, call.input, call.output, call.ranked_leaders) + } Metric::Cosine => { - nearest_leaders::(arch, call.input, call.output, call.workspace) + nearest_leaders::(arch, call.input, call.output, call.ranked_leaders) } Metric::CosineNormalized => nearest_leaders::( arch, call.input, call.output, - call.workspace, + call.ranked_leaders, + ), + Metric::InnerProduct => nearest_leaders::( + arch, + call.input, + call.output, + call.ranked_leaders, ), - Metric::InnerProduct => { - nearest_leaders::(arch, call.input, call.output, call.workspace) - } } } } @@ -381,14 +423,14 @@ fn dispatch_nearest_leaders( metric: Metric, input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, - workspace: &mut PartitionKernelWorkspace, + ranked_leaders: &mut Vec<(u32, f32)>, ) -> Result<(), PartitionKernelError> { diskann_wide::arch::dispatch1_no_features( DispatchPartitionForTest(metric), DispatchedPartitionCall { input, output, - workspace, + ranked_leaders, }, ) } @@ -403,7 +445,6 @@ mod tests { use diskann_vector::distance::Metric; fn test_input<'a>( - _metric: Metric, dots: &'a [f32], point_count: usize, leader_count: usize, @@ -419,10 +460,9 @@ mod tests { } } - // This oracle checks SIMD chunking, scalar tails, and tracker order. It uses - // the scalar ranking formula for metric `M`. - fn scalar_traversal_reference( - metric: Metric, + // This oracle checks SIMD groups, single values, and retained-leader order. + // It uses the single-value ranking formula for metric `M`. + fn ranking_reference( input: PartitionInput<'_>, fanout: usize, output: &mut [u32], @@ -434,40 +474,48 @@ mod tests { .enumerate() { let point_norm = input.norms.point_norms.get(point).copied().unwrap_or(0.0); - let mut tracker = vec![(u32::MAX, f32::INFINITY); fanout]; + let mut ranked_leaders = vec![(u32::MAX, f32::INFINITY); fanout]; for (leader, &dot) in point_dots.iter().enumerate() { - let leader_norm = if matches!(metric, Metric::L2 | Metric::Cosine) { - input.norms.leader_norms[leader] - } else { - 0.0 - }; insert_leader( - &mut tracker, + &mut ranked_leaders, leader as u32, - M::partition_ranking_scalar(dot, point_norm, leader_norm), + M::partition_ranking_single( + dot, + point_norm, + input.norms.leader_norms.get(leader).copied().unwrap_or(0.0), + ), ); } - for (destination, &(leader, _)) in point_output.iter_mut().zip(&tracker) { + for (destination, &(leader, _)) in point_output.iter_mut().zip(&ranked_leaders) { *destination = leader; } } } + fn single_ranking( + dot_product: f32, + point_norms: &[f32], + leader_norms: &[f32], + ) -> f32 { + M::partition_ranking_single( + dot_product, + point_norms.first().copied().unwrap_or(0.0), + leader_norms.first().copied().unwrap_or(0.0), + ) + } + #[test] - fn scalar_ranking_matches_metric_contract() { - assert_eq!(L2::partition_ranking_scalar(2.0, 0.0, 9.0), 5.0); - assert_eq!( - CosineNormalized::partition_ranking_scalar(0.25, 0.0, 0.0), - 0.75 - ); - assert_eq!(InnerProduct::partition_ranking_scalar(3.0, 0.0, 0.0), -3.0); - assert_eq!(Cosine::partition_ranking_scalar(4.0, 2.0, 4.0), 0.5); - assert_eq!(Cosine::partition_ranking_scalar(4.0, 0.0, 4.0), 1.0); - assert!(Cosine::partition_ranking_scalar(1.0, f32::NAN, 1.0).is_nan()); + fn single_ranking_matches_metric_contract() { + assert_eq!(single_ranking::(2.0, &[], &[9.0]), 5.0); + assert_eq!(single_ranking::(0.25, &[], &[]), 0.75); + assert_eq!(single_ranking::(3.0, &[], &[]), -3.0); + assert_eq!(single_ranking::(4.0, &[2.0], &[4.0]), 0.5); + assert_eq!(single_ranking::(4.0, &[0.0], &[4.0]), 1.0); + assert!(single_ranking::(1.0, &[f32::NAN], &[1.0]).is_nan()); } #[test] - fn cosine_special_norms_match_scalar_and_dispatched_kernel() { + fn cosine_special_norms_match_single_and_dispatched_kernel() { let leader_count = 17; let point_norms = [0.0, 0.0, f32::MIN_POSITIVE.sqrt(), f32::NAN]; let dots = vec![1.0; point_norms.len() * leader_count]; @@ -479,7 +527,6 @@ mod tests { f32::NAN, ]); let input = test_input( - Metric::Cosine, &dots, point_norms.len(), leader_count, @@ -487,13 +534,13 @@ mod tests { &leader_norms, ); let mut expected = vec![u32::MAX; point_norms.len() * 2]; - scalar_traversal_reference::(Metric::Cosine, input, 2, &mut expected); + ranking_reference::(input, 2, &mut expected); let mut actual = vec![u32::MAX; point_norms.len() * 2]; dispatch_nearest_leaders( Metric::Cosine, input, MutMatrixView::try_from(actual.as_mut_slice(), point_norms.len(), 2).unwrap(), - &mut PartitionKernelWorkspace::default(), + &mut Vec::new(), ) .unwrap(); @@ -503,26 +550,42 @@ mod tests { } #[test] - fn scalar_topk_orders_candidates_and_preserves_ties() { - let mut tracker = vec![(u32::MAX, f32::INFINITY); 4]; + fn topk_orders_candidates_and_preserves_ties() { + let mut ranked_leaders = vec![(u32::MAX, f32::INFINITY); 4]; for (leader, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 1.0)] { - insert_leader(&mut tracker, leader, distance); + insert_leader(&mut ranked_leaders, leader, distance); } - insert_leader(&mut tracker, 5, f32::NAN); + insert_leader(&mut ranked_leaders, 5, f32::NAN); - assert_eq!(tracker[..], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); + assert_eq!(ranked_leaders[..], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); } #[test] - fn workspace_reuses_runtime_fanout_capacity() { - let mut workspace = PartitionKernelWorkspace::default(); - workspace.prepare(32).unwrap(); - let allocation = workspace.tracker.as_ptr(); + fn ranked_leaders_reuses_runtime_fanout_capacity() { + let dots = [0.0; 32]; + let input = test_input(&dots, 1, 32, &[], &[]); + let mut ranked_leaders = Vec::new(); + let mut wide_output = [u32::MAX; 32]; + dispatch_nearest_leaders( + Metric::InnerProduct, + input, + MutMatrixView::try_from(&mut wide_output[..], 1, 32).unwrap(), + &mut ranked_leaders, + ) + .unwrap(); + let allocation = ranked_leaders.as_ptr(); - workspace.prepare(3).unwrap(); + let mut narrow_output = [u32::MAX; 3]; + dispatch_nearest_leaders( + Metric::InnerProduct, + input, + MutMatrixView::try_from(&mut narrow_output[..], 1, 3).unwrap(), + &mut ranked_leaders, + ) + .unwrap(); - assert_eq!(workspace.tracker.as_ptr(), allocation); - assert_eq!(workspace.tracker.len(), 3); + assert_eq!(ranked_leaders.as_ptr(), allocation); + assert_eq!(ranked_leaders.len(), 3); } } #[cfg(test)] @@ -532,15 +595,11 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod integration_tests { - use super::{ - PartitionInput, PartitionKernelError, PartitionKernelWorkspace, PartitionNorms, - dispatch_nearest_leaders, - }; + use super::{PartitionInput, PartitionKernelError, PartitionNorms, dispatch_nearest_leaders}; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; fn test_input<'a>( - _metric: Metric, dots: &'a [f32], point_count: usize, leader_count: usize, @@ -656,7 +715,7 @@ mod integration_tests { metric, input, MutMatrixView::try_from(output.as_mut_slice(), input.dots.nrows(), fanout).unwrap(), - &mut PartitionKernelWorkspace::default(), + &mut Vec::new(), )?; Ok(output) } @@ -671,7 +730,7 @@ mod integration_tests { ] { for leader_count in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let (dots, point_norms, leader_norms) = differential_data(metric, leader_count); - let input = test_input(metric, &dots, 2, leader_count, &point_norms, &leader_norms); + let input = test_input(&dots, 2, leader_count, &point_norms, &leader_norms); for fanout in [1, 2, 16, 17, 32] { if fanout >= leader_count { continue; @@ -696,18 +755,13 @@ mod integration_tests { let norms = [0.0, 1.0, 4.0, 9.0]; assert_eq!( - run_partition_kernel( - Metric::L2, - test_input(Metric::L2, &dots, 2, 4, &[], &norms), - 2 - ) - .unwrap(), + run_partition_kernel(Metric::L2, test_input(&dots, 2, 4, &[], &norms), 2).unwrap(), [0, 1, 2, 1] ); } #[test] - fn l2_scalar_tail_can_outrank_a_fused_simd_lane() { + fn l2_single_can_outrank_a_fused_simd_lane() { let mut dots = [0.0; 17]; dots[0] = f32::MAX; dots[16] = f32::MAX; @@ -716,7 +770,7 @@ mod integration_tests { assert_eq!( run_partition_kernel( Metric::L2, - test_input(Metric::L2, &dots, 1, 17, &[], &leader_squared_norms,), + test_input(&dots, 1, 17, &[], &leader_squared_norms,), 1, ) .unwrap(), @@ -745,7 +799,7 @@ mod integration_tests { assert_eq!( run_partition_kernel( metric, - test_input(metric, &dots, 2, 3, point_norms, leader_norms), + test_input(&dots, 2, 3, point_norms, leader_norms), 2, ) .unwrap(), @@ -760,7 +814,7 @@ mod integration_tests { assert_eq!( run_partition_kernel( Metric::Cosine, - test_input(Metric::Cosine, &[100.0, -100.0], 1, 2, &[0.0], &[1.0, 1.0]), + test_input(&[100.0, -100.0], 1, 2, &[0.0], &[1.0, 1.0]), 2, ) .unwrap(), @@ -773,12 +827,8 @@ mod integration_tests { let mut dots = [0.0; 8]; dots[7] = -f32::MAX; assert_eq!( - run_partition_kernel( - Metric::InnerProduct, - test_input(Metric::InnerProduct, &dots, 1, 8, &[], &[]), - 8 - ) - .unwrap(), + run_partition_kernel(Metric::InnerProduct, test_input(&dots, 1, 8, &[], &[]), 8) + .unwrap(), [0, 1, 2, 3, 4, 5, 6, 7] ); } @@ -788,7 +838,7 @@ mod integration_tests { assert_eq!( run_partition_kernel( Metric::InnerProduct, - test_input(Metric::InnerProduct, &[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), + test_input(&[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), 2, ) .unwrap(), @@ -801,7 +851,7 @@ mod integration_tests { assert_eq!( run_partition_kernel( Metric::InnerProduct, - test_input(Metric::InnerProduct, &[f32::NAN, 3.0], 1, 2, &[], &[]), + test_input(&[f32::NAN, 3.0], 1, 2, &[], &[]), 2, ), Err(PartitionKernelError::InsufficientRankableLeaders { @@ -813,21 +863,16 @@ mod integration_tests { #[test] fn accepts_empty_points_zero_fanout_and_largest_leader_id() { + run_partition_kernel(Metric::InnerProduct, test_input(&[], 0, 3, &[], &[]), 2).unwrap(); run_partition_kernel( Metric::InnerProduct, - test_input(Metric::InnerProduct, &[], 0, 3, &[], &[]), - 2, - ) - .unwrap(); - run_partition_kernel( - Metric::InnerProduct, - test_input(Metric::InnerProduct, &[1.0, 2.0, 3.0], 1, 3, &[], &[]), + test_input(&[1.0, 2.0, 3.0], 1, 3, &[], &[]), 0, ) .unwrap(); run_partition_kernel( Metric::InnerProduct, - test_input(Metric::InnerProduct, &[], 0, u32::MAX as usize, &[], &[]), + test_input(&[], 0, u32::MAX as usize, &[], &[]), 0, ) .unwrap(); @@ -836,14 +881,7 @@ mod integration_tests { assert_eq!( run_partition_kernel( Metric::InnerProduct, - test_input( - Metric::InnerProduct, - &[], - 0, - u32::MAX as usize + 1, - &[], - &[], - ), + test_input(&[], 0, u32::MAX as usize + 1, &[], &[],), 0, ), Err(PartitionKernelError::TooManyLeaders(u32::MAX as usize + 1)) @@ -853,14 +891,14 @@ mod integration_tests { #[test] fn rejects_wrong_output_norms_and_fanout() { let dots = [0.0; 6]; - let valid_input = test_input(Metric::InnerProduct, &dots, 2, 3, &[], &[]); + let valid_input = test_input(&dots, 2, 3, &[], &[]); let mut wrong_output = [u32::MAX; 3]; assert_eq!( dispatch_nearest_leaders( Metric::InnerProduct, valid_input, MutMatrixView::try_from(&mut wrong_output[..], 1, 3).unwrap(), - &mut PartitionKernelWorkspace::default(), + &mut Vec::new(), ), Err(PartitionKernelError::InvalidOutputShape { expected_rows: 2, @@ -885,7 +923,6 @@ mod integration_tests { actual: 2, }) ); - assert_eq!( run_partition_kernel(Metric::InnerProduct, valid_input, 4), Err(PartitionKernelError::InvalidFanout { @@ -896,11 +933,7 @@ mod integration_tests { let one = [0.0]; assert_eq!( - run_partition_kernel( - Metric::InnerProduct, - test_input(Metric::InnerProduct, &one, 1, 1, &[], &[]), - 2, - ), + run_partition_kernel(Metric::InnerProduct, test_input(&one, 1, 1, &[], &[]), 2,), Err(PartitionKernelError::InvalidFanout { fanout: 2, leader_count: 1, From 8913a727d5e51dd20413116fa97fa77b08058ef2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 12 Aug 2026 05:57:56 +0000 Subject: [PATCH 76/80] refactor(pipnn): deepen numerical kernels Move GEMM, norm preparation, and reusable scratch behind the leaf and partition interfaces. Align fused L2 and bounded cosine semantics with DiskANN distance kernels. --- Cargo.lock | 1 + diskann/Cargo.toml | 5 +- diskann/src/graph/pipnn/kernel_metric.rs | 52 +- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 182 +++-- .../graph/pipnn/kernel_metric/partition.rs | 224 +++--- diskann/src/graph/pipnn/leaf_kernel.rs | 761 +++++++----------- diskann/src/graph/pipnn/partition_kernel.rs | 615 +++++--------- 7 files changed, 757 insertions(+), 1083 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2588eee92a..da640e005a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -438,6 +438,7 @@ dependencies = [ "anyhow", "bytemuck", "dashmap", + "diskann-linalg", "diskann-utils", "diskann-vector", "diskann-wide", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 72911c910f..7cd2a0df09 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -30,6 +30,7 @@ diskann-wide = { workspace = true } # Optional Dependencies dashmap = { workspace = true, optional = true } +diskann-linalg = { workspace = true, optional = true } [dev-dependencies] futures-util = { workspace = true, default-features = false } @@ -56,8 +57,8 @@ panic = "warn" [features] default = ["tracing"] -# Enable PiPNN batch graph construction. -pipnn = [] +# Enable PiPNN numerical kernels. +pipnn = ["dep:diskann-linalg"] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 52f03ea9cf..d3869cdbba 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -11,9 +11,6 @@ mod partition; pub(super) use leaf::LeafMetric; pub(super) use partition::PartitionMetric; -use std::collections::TryReserveError; - -use diskann_utils::views::MatrixView; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; pub(super) struct L2; @@ -21,11 +18,6 @@ pub(super) struct Cosine; pub(super) struct CosineNormalized; pub(super) struct InnerProduct; -pub(super) struct NormPreparation<'a, 'b> { - pub(super) values: MatrixView<'a, f32>, - pub(super) norms: &'b mut Vec, -} - /// Prepared norms for one point stripe and its sampled leaders. #[derive(Clone, Copy, Debug)] pub(super) struct PartitionNorms<'a> { @@ -33,28 +25,10 @@ pub(super) struct PartitionNorms<'a> { pub(super) leader_norms: &'a [f32], } -pub(super) fn resize_norms(norms: &mut Vec, len: usize) -> Result<(), TryReserveError> { - norms.try_reserve(len.saturating_sub(norms.len()))?; - norms.resize(len, 0.0); - Ok(()) -} - -/// This function converts a squared norm to a norm. -/// -/// It maps subnormal values to zero and preserves NaN. -#[inline(always)] -pub(super) fn norm_from_squared(squared_norm: f32) -> f32 { - if squared_norm < f32::MIN_POSITIVE { - 0.0 - } else { - squared_norm.sqrt() - } -} - /// Compute SIMD cosine distance with the DiskANN zero-norm and NaN rules. /// -/// Each lane contains one point pair. A zero norm produces zero similarity. A -/// NaN norm remains NaN unless the other norm is zero. +/// Each lane contains one point pair. A zero norm produces zero similarity. +/// Finite similarity is clamped to the cosine range before distance conversion. #[inline(always)] pub(super) fn cosine_distance_simd(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where @@ -69,7 +43,8 @@ where 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 + let negative_one = F::splat(arch, -1.0); + one - negative_one.max_simd(cosine.min_simd(one)) } /// Compute one cosine distance with the DiskANN zero-norm and NaN rules. @@ -78,22 +53,7 @@ pub(super) fn cosine_distance_single(dot: f32, source_norm: f32, target_norm: f3 if source_norm < f32::MIN_POSITIVE.sqrt() || target_norm < f32::MIN_POSITIVE.sqrt() { 1.0 } else { - 1.0 - dot / (source_norm * target_norm) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn norm_from_squared_applies_zero_threshold_without_erasing_nan() { - assert_eq!(norm_from_squared(-0.0).to_bits(), 0.0f32.to_bits()); - assert_eq!(norm_from_squared(f32::MIN_POSITIVE / 2.0), 0.0); - assert_eq!( - norm_from_squared(f32::MIN_POSITIVE), - f32::MIN_POSITIVE.sqrt() - ); - assert!(norm_from_squared(f32::NAN).is_nan()); + let cosine = dot / (source_norm * target_norm); + 1.0 - (-1.0_f32).max(1.0_f32.min(cosine)) } } diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index af8d042b80..2ac408e206 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -3,150 +3,192 @@ * Licensed under the MIT license. */ -use std::collections::TryReserveError; - +use diskann_utils::views::MatrixView; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; use super::{ - Cosine, CosineNormalized, InnerProduct, L2, NormPreparation, cosine_distance_simd, - cosine_distance_single, norm_from_squared, resize_norms, + Cosine, CosineNormalized, InnerProduct, L2, cosine_distance_simd, cosine_distance_single, }; -/// Leaf formulas return ascending distances. -/// L2 uses squared norms. Cosine uses norms. Other metrics ignore norms. +/// Compute leaf distances for one concrete metric. pub(in super::super) trait LeafMetric: Send + Sync + 'static { /// Prepare one contiguous metric-specific norm for each leaf-local point. - fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError>; + fn prepare_leaf_norms(_dots: MatrixView<'_, f32>, norms: &mut Vec) { + norms.clear(); + } + + /// Prepare one source norm for reuse across SIMD target groups. + #[inline(always)] + fn source_simd(arch: F::Arch, _norms: &[f32], _source: usize) -> F + where + F: SIMDVector, + { + F::default(arch) + } + + /// Prepare one source norm for reuse across single target values. + #[inline(always)] + fn source_single(_norms: &[f32], _source: usize) -> f32 { + 0.0 + } /// Compute distances for one complete SIMD group. - fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F + fn distances_simd( + arch: F::Arch, + norms: &[f32], + source_norms: F, + dot_products: F, + first_target: usize, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; /// Compute one distance outside the complete SIMD prefix. - fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32; + fn distance_single(norms: &[f32], source_norm: f32, dot_product: f32, target: usize) -> f32; } -/// Clamp negative SIMD roundoff to zero and preserve NaN lanes. +/// Load one complete SIMD group of prepared norms. #[inline(always)] -fn clamp_nonnegative_simd(arch: F::Arch, distance: F) -> F +fn load_norms_simd(arch: F::Arch, norms: &[f32], first_norm: usize) -> F where - F: SIMDVector + SIMDFloat, - F::Mask: SIMDSelect, + F: SIMDVector, { - let zero = F::default(arch); - distance - .eq_simd(distance) - .select(zero.max_simd(distance), distance) -} + let last_norm = first_norm + F::LANES; + let norm_group = &norms[first_norm..last_norm]; -/// Clamp negative roundoff to zero and preserve NaN. -#[inline(always)] -fn clamp_nonnegative_single(distance: f32) -> f32 { - if distance < 0.0 { 0.0 } else { distance } + // SAFETY: `norm_group` contains one complete SIMD group. + unsafe { F::load_simd(arch, norm_group.as_ptr()) } } impl LeafMetric for L2 { - fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - resize_norms(preparation.norms, preparation.values.nrows())?; - for (point, norm) in preparation.norms.iter_mut().enumerate() { - *norm = preparation.values[(point, point)]; + fn prepare_leaf_norms(dots: MatrixView<'_, f32>, norms: &mut Vec) { + norms.resize(dots.nrows(), 0.0); + for (point, norm) in norms.iter_mut().enumerate() { + *norm = dots[(point, point)]; } - Ok(()) } #[inline(always)] - fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F + fn source_simd(arch: F::Arch, norms: &[f32], source: usize) -> F + where + F: SIMDVector, + { + F::splat(arch, norms[source]) + } + + #[inline(always)] + fn source_single(norms: &[f32], source: usize) -> f32 { + norms[source] + } + + #[inline(always)] + fn distances_simd( + arch: F::Arch, + norms: &[f32], + source_norms: F, + dot_products: F, + first_target: usize, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - clamp_nonnegative_simd( - arch, - source_norms + target_norms - F::splat(arch, 2.0) * dot_products, - ) + let target_norms = load_norms_simd::(arch, norms, first_target); + (F::splat(arch, -2.0).mul_add_simd(dot_products, source_norms) + target_norms) + .max_simd(F::default(arch)) } #[inline(always)] - fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32 { - clamp_nonnegative_single(source_norm + target_norm - 2.0 * dot_product) + fn distance_single(norms: &[f32], source_norm: f32, dot_product: f32, target: usize) -> f32 { + ((-2.0_f32).mul_add(dot_product, source_norm) + norms[target]).max(0.0) } } impl LeafMetric for Cosine { - fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - resize_norms(preparation.norms, preparation.values.nrows())?; - for (point, norm) in preparation.norms.iter_mut().enumerate() { - *norm = norm_from_squared(preparation.values[(point, point)]); + fn prepare_leaf_norms(dots: MatrixView<'_, f32>, norms: &mut Vec) { + norms.resize(dots.nrows(), 0.0); + for (point, norm) in norms.iter_mut().enumerate() { + *norm = dots[(point, point)].sqrt(); } - Ok(()) } #[inline(always)] - fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F + fn source_simd(arch: F::Arch, norms: &[f32], source: usize) -> F + where + F: SIMDVector, + { + F::splat(arch, norms[source]) + } + + #[inline(always)] + fn source_single(norms: &[f32], source: usize) -> f32 { + norms[source] + } + + #[inline(always)] + fn distances_simd( + arch: F::Arch, + norms: &[f32], + source_norms: F, + dot_products: F, + first_target: usize, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - clamp_nonnegative_simd( - arch, - cosine_distance_simd(arch, dot_products, source_norms, target_norms), - ) + let target_norms = load_norms_simd::(arch, norms, first_target); + cosine_distance_simd(arch, dot_products, source_norms, target_norms) + .max_simd(F::default(arch)) } #[inline(always)] - fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32 { - clamp_nonnegative_single(cosine_distance_single( - dot_product, - source_norm, - target_norm, - )) + fn distance_single(norms: &[f32], source_norm: f32, dot_product: f32, target: usize) -> f32 { + cosine_distance_single(dot_product, source_norm, norms[target]).max(0.0) } } impl LeafMetric for CosineNormalized { - fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - preparation.norms.clear(); - Ok(()) - } - #[inline(always)] - fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F + fn distances_simd( + arch: F::Arch, + _norms: &[f32], + _source_norms: F, + dot_products: F, + _first_target: usize, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - let _ = (source_norms, target_norms); - clamp_nonnegative_simd(arch, F::splat(arch, 1.0) - dot_products) + F::splat(arch, 1.0) - dot_products } #[inline(always)] - fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32 { - let _ = (source_norm, target_norm); - clamp_nonnegative_single(1.0 - dot_product) + fn distance_single(_norms: &[f32], _source_norm: f32, dot_product: f32, _target: usize) -> f32 { + 1.0 - dot_product } } impl LeafMetric for InnerProduct { - fn prepare_leaf_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - preparation.norms.clear(); - Ok(()) - } - #[inline(always)] - fn leaf_distance_simd(arch: F::Arch, dot_products: F, source_norms: F, target_norms: F) -> F + fn distances_simd( + arch: F::Arch, + _norms: &[f32], + _source_norms: F, + dot_products: F, + _first_target: usize, + ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - let _ = (source_norms, target_norms); F::default(arch) - dot_products } #[inline(always)] - fn leaf_distance_single(dot_product: f32, source_norm: f32, target_norm: f32) -> f32 { - let _ = (source_norm, target_norm); + fn distance_single(_norms: &[f32], _source_norm: f32, dot_product: f32, _target: usize) -> f32 { -dot_product } } diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index bfe99faffd..ad237ec01a 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -3,204 +3,216 @@ * Licensed under the MIT license. */ -use std::collections::TryReserveError; - +use diskann_utils::views::MatrixView; use diskann_vector::{Norm, norm::FastL2NormSquared}; use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; use super::{ - Cosine, CosineNormalized, InnerProduct, L2, NormPreparation, cosine_distance_simd, - cosine_distance_single, norm_from_squared, resize_norms, + Cosine, CosineNormalized, InnerProduct, L2, PartitionNorms, cosine_distance_simd, + cosine_distance_single, }; -/// Partition formulas return ascending rankings. -/// L2 uses squared leader norms. Cosine uses point and leader norms. +/// Compute partition rankings for one concrete metric. pub(in super::super) trait PartitionMetric: Send + Sync + 'static { /// Prepare one norm value for each point in the active stripe. - fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError>; + fn prepare_point_norms(_points: MatrixView<'_, f32>, norms: &mut Vec) { + norms.clear(); + } /// Prepare one norm value for each sampled leader. - fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError>; + fn prepare_leader_norms(_leaders: MatrixView<'_, f32>, norms: &mut Vec) { + norms.clear(); + } + + /// Prepare one point norm for reuse across SIMD leader groups. + #[inline(always)] + fn point_simd(arch: F::Arch, _norms: PartitionNorms<'_>, _point: usize) -> F + where + F: SIMDVector, + { + F::default(arch) + } + + /// Prepare one point norm for reuse across single leader values. + #[inline(always)] + fn point_single(_norms: PartitionNorms<'_>, _point: usize) -> f32 { + 0.0 + } /// Compute rankings for one complete SIMD group. - fn partition_ranking_simd( + fn rankings_simd( arch: F::Arch, - dot_products: F, + norms: PartitionNorms<'_>, point_norms: F, - leader_norms: F, + dot_products: F, + first_leader: usize, ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect; /// Compute one ranking outside the complete SIMD prefix. - fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32; + fn ranking_single( + norms: PartitionNorms<'_>, + point_norm: f32, + dot_product: f32, + leader: usize, + ) -> f32; +} + +/// Load one complete SIMD group of prepared norms. +#[inline(always)] +fn load_norms_simd(arch: F::Arch, norms: &[f32], first_norm: usize) -> F +where + F: SIMDVector, +{ + let last_norm = first_norm + F::LANES; + let norm_group = &norms[first_norm..last_norm]; + + // SAFETY: `norm_group` contains one complete SIMD group. + unsafe { F::load_simd(arch, norm_group.as_ptr()) } } impl PartitionMetric for L2 { - fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - preparation.norms.clear(); - Ok(()) - } - - fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - resize_norms(preparation.norms, preparation.values.nrows())?; - for (norm, leader) in preparation - .norms - .iter_mut() - .zip(preparation.values.row_iter()) - { + fn prepare_leader_norms(leaders: MatrixView<'_, f32>, norms: &mut Vec) { + norms.resize(leaders.nrows(), 0.0); + for (norm, leader) in norms.iter_mut().zip(leaders.row_iter()) { *norm = leader.iter().map(|value| value * value).sum(); } - Ok(()) } #[inline(always)] - fn partition_ranking_simd( + fn rankings_simd( arch: F::Arch, + norms: PartitionNorms<'_>, + _point_norms: F, dot_products: F, - point_norms: F, - leader_norms: F, + first_leader: usize, ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - let _ = point_norms; - // Fused arithmetic defines the ranking order for complete SIMD groups. + let leader_norms = load_norms_simd::(arch, norms.leader_norms, first_leader); F::splat(arch, -2.0).mul_add_simd(dot_products, leader_norms) } #[inline(always)] - fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32 { - let _ = point_norm; - // Non-fused arithmetic defines the ranking order outside the SIMD prefix. - leader_norm - 2.0 * dot_product + fn ranking_single( + norms: PartitionNorms<'_>, + _point_norm: f32, + dot_product: f32, + leader: usize, + ) -> f32 { + (-2.0_f32).mul_add(dot_product, norms.leader_norms[leader]) } } impl PartitionMetric for Cosine { - fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - resize_norms(preparation.norms, preparation.values.nrows())?; - for (norm, point) in preparation - .norms - .iter_mut() - .zip(preparation.values.row_iter()) - { - *norm = norm_from_squared(FastL2NormSquared.evaluate(point)); + fn prepare_point_norms(points: MatrixView<'_, f32>, norms: &mut Vec) { + norms.resize(points.nrows(), 0.0); + for (norm, point) in norms.iter_mut().zip(points.row_iter()) { + *norm = FastL2NormSquared.evaluate(point).sqrt(); } - Ok(()) - } - - fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - resize_norms(preparation.norms, preparation.values.nrows())?; - for (norm, leader) in preparation - .norms - .iter_mut() - .zip(preparation.values.row_iter()) - { - let squared_norm = leader.iter().map(|value| value * value).sum(); - *norm = norm_from_squared(squared_norm); + } + + fn prepare_leader_norms(leaders: MatrixView<'_, f32>, norms: &mut Vec) { + norms.resize(leaders.nrows(), 0.0); + for (norm, leader) in norms.iter_mut().zip(leaders.row_iter()) { + *norm = leader.iter().map(|value| value * value).sum::().sqrt(); } - Ok(()) } #[inline(always)] - fn partition_ranking_simd( + fn point_simd(arch: F::Arch, norms: PartitionNorms<'_>, point: usize) -> F + where + F: SIMDVector, + { + F::splat(arch, norms.point_norms[point]) + } + + #[inline(always)] + fn point_single(norms: PartitionNorms<'_>, point: usize) -> f32 { + norms.point_norms[point] + } + + #[inline(always)] + fn rankings_simd( arch: F::Arch, - dot_products: F, + norms: PartitionNorms<'_>, point_norms: F, - leader_norms: F, + dot_products: F, + first_leader: usize, ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { + let leader_norms = load_norms_simd::(arch, norms.leader_norms, first_leader); cosine_distance_simd(arch, dot_products, point_norms, leader_norms) } #[inline(always)] - fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32 { - cosine_distance_single(dot_product, point_norm, leader_norm) + fn ranking_single( + norms: PartitionNorms<'_>, + point_norm: f32, + dot_product: f32, + leader: usize, + ) -> f32 { + cosine_distance_single(dot_product, point_norm, norms.leader_norms[leader]) } } impl PartitionMetric for CosineNormalized { - fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - preparation.norms.clear(); - Ok(()) - } - - fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - preparation.norms.clear(); - Ok(()) - } - #[inline(always)] - fn partition_ranking_simd( + fn rankings_simd( arch: F::Arch, + _norms: PartitionNorms<'_>, + _point_norms: F, dot_products: F, - point_norms: F, - leader_norms: F, + _first_leader: usize, ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - let _ = (point_norms, leader_norms); F::splat(arch, 1.0) - dot_products } #[inline(always)] - fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32 { - let _ = (point_norm, leader_norm); + fn ranking_single( + _norms: PartitionNorms<'_>, + _point_norm: f32, + dot_product: f32, + _leader: usize, + ) -> f32 { 1.0 - dot_product } } impl PartitionMetric for InnerProduct { - fn prepare_point_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - preparation.norms.clear(); - Ok(()) - } - - fn prepare_leader_norms(preparation: NormPreparation<'_, '_>) -> Result<(), TryReserveError> { - preparation.norms.clear(); - Ok(()) - } - #[inline(always)] - fn partition_ranking_simd( + fn rankings_simd( arch: F::Arch, + _norms: PartitionNorms<'_>, + _point_norms: F, dot_products: F, - point_norms: F, - leader_norms: F, + _first_leader: usize, ) -> F where F: SIMDVector + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, { - let _ = (point_norms, leader_norms); F::default(arch) - dot_products } #[inline(always)] - fn partition_ranking_single(dot_product: f32, point_norm: f32, leader_norm: f32) -> f32 { - let _ = (point_norm, leader_norm); + fn ranking_single( + _norms: PartitionNorms<'_>, + _point_norm: f32, + dot_product: f32, + _leader: usize, + ) -> f32 { -dot_product } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn l2_ranking_preserves_non_fused_arithmetic() { - let ranking = L2::partition_ranking_single(f32::MAX, 0.0, f32::MAX); - let fused = (-2.0f32).mul_add(f32::MAX, f32::MAX); - - assert_eq!(ranking, f32::NEG_INFINITY); - assert_eq!(fused, -f32::MAX); - } -} diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 45028551c9..a3c9b6570b 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -3,32 +3,29 @@ * Licensed under the MIT license. */ -//! Leaf-local top-k selection from a lower-triangular Gram matrix. +//! Leaf-local top-k selection from packed `f32` point vectors. //! -//! The input is an `n × n` [`MatrixView`] from `sgemm_aat_lower`. The diagonal -//! contains metric norms. The kernel reads only the strict lower triangle. It -//! evaluates each point pair once and updates both points. +//! The kernel computes the lower-triangular Gram matrix and metric-specific +//! norms. Its ranking loop reads each strict-lower point pair once and updates +//! both points. //! //! The output is an `n × k` matrix of sorted [`LeafNeighbor`] values. Each target -//! is a position in the leaf. The kernel supports `k` from zero through -//! [`MAX_LEAF_NEIGHBORS`]. Positive widths use fixed arrays. +//! is a position in the leaf. Widths 1 through 3 use fixed insertion. Larger +//! widths use the runtime insertion loop. //! //! Strict comparisons keep scan order for equal distances. They do not rank NaN. //! All supported metrics use the same SIMD-group and single-value traversal. //! -//! The caller supplies concrete architecture `A` and metric `M`. The function -//! checks all shapes and local-ID bounds before it changes workspace or uses an -//! unchecked SIMD load. [`LeafKernelWorkspace`] stores reusable rejection -//! thresholds. +//! The caller supplies concrete architecture `A` and metric `M`. The private +//! dot ranker receives the square matrix created by this module. +//! [`LeafKernelWorkspace`] stores reusable numerical scratch. +use crate::{ANNError, ANNResult}; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; use super::kernel_metric::LeafMetric; -/// Largest leaf-local neighbor count supported by the fixed insertion kernel. -pub(super) const MAX_LEAF_NEIGHBORS: usize = 3; - /// One leaf-local neighbor and its metric distance. #[derive(Clone, Copy, Debug, PartialEq)] pub(super) struct LeafNeighbor { @@ -54,28 +51,17 @@ impl Default for LeafNeighbor { } } -/// Reusable temporary storage for leaf top-k selection. +/// Reusable storage for one leaf numerical pipeline. #[derive(Debug, Default)] pub(super) struct LeafKernelWorkspace { + dot_scratch: Vec, + norm_scratch: Vec, worst: Vec, } -/// Validation or allocation error returned by [`nearest_neighbors`]. +/// Validation error returned by the dot-ranking loop. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub(super) enum LeafKernelError { - /// The point count cannot be represented in leaf-local `u32` positions. - #[error("point count {0} exceeds the u32 position limit")] - TooManyPoints(usize), - /// The dot-product matrix is not square. - #[error("leaf dot-product matrix must be square, got {rows} x {cols}")] - NonSquareDots { rows: usize, cols: usize }, - /// The output matrix does not have one row per input point. - #[error("invalid output row count: expected {expected}, got {actual} with {columns} columns")] - InvalidOutputRows { - expected: usize, - actual: usize, - columns: usize, - }, /// A source requests more neighbors than the leaf or fixed kernel supports. #[error("invalid leaf neighbor count {neighbors} for {points} points; maximum is {maximum}")] InvalidNeighborCount { @@ -83,67 +69,65 @@ pub(super) enum LeafKernelError { neighbors: usize, maximum: usize, }, - /// The prepared norm count does not match the point count. - #[error("invalid leaf norm count: expected {expected}, got {actual}")] - InvalidNormCount { expected: usize, actual: usize }, - /// Temporary storage could not be reserved. - #[error("failed to reserve {additional} values for {buffer}")] - Allocation { - buffer: &'static str, - additional: usize, - }, - /// A source did not contain enough rankable targets to fill its output. - #[error("source {source_index} has fewer than {neighbors} rankable leaf neighbors")] - InsufficientRankableNeighbors { - source_index: usize, - neighbors: usize, - }, } /// Return the non-self neighbor count for one leaf. /// /// `points` is the number of points in the leaf. `requested_k` is the configured -/// neighbor count. The result is `min(requested_k, points - 1)`. The function -/// rejects a value above [`MAX_LEAF_NEIGHBORS`]. -/// -/// # Errors +/// neighbor count. The result is `min(requested_k, points - 1)`. /// -/// Returns [`LeafKernelError::TooManyPoints`] when leaf-local positions cannot -/// fit in `u32`, or [`LeafKernelError::InvalidNeighborCount`] when `requested_k` -/// exceeds [`MAX_LEAF_NEIGHBORS`]. -pub(super) fn leaf_neighbor_count( - points: usize, - requested_k: usize, -) -> Result { - if points > u32::MAX as usize { - return Err(LeafKernelError::TooManyPoints(points)); - } - if requested_k > MAX_LEAF_NEIGHBORS { - return Err(LeafKernelError::InvalidNeighborCount { - points, - neighbors: requested_k, - maximum: MAX_LEAF_NEIGHBORS, - }); - } - Ok(requested_k.min(points.saturating_sub(1))) +pub(super) fn leaf_neighbor_count(points: usize, requested_k: usize) -> usize { + requested_k.min(points.saturating_sub(1)) } -/// Select the nearest non-self positions for each point in a leaf. -/// -/// `output` has one row for each input point. Its column count requests the -/// neighbor count. The function checks this shape and the supported count before -/// it changes output. Equal distances keep pair scan order. +/// Compute local nearest neighbors for one packed leaf matrix. /// /// # Errors /// -/// Returns [`LeafKernelError`] for an invalid shape or count. It also returns an -/// error for allocation failure or insufficient rankable neighbors. -pub(super) fn nearest_neighbors( +/// Returns an error for invalid linear-algebra input or output width. +pub(super) fn select_leaf_neighbors( + arch: A, + points: MatrixView<'_, f32>, + output: MutMatrixView<'_, LeafNeighbor>, + workspace: &mut LeafKernelWorkspace, +) -> ANNResult<()> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + M: LeafMetric, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + let point_count = points.nrows(); + let dot_count = point_count * point_count; + let LeafKernelWorkspace { + dot_scratch, + norm_scratch, + worst, + } = workspace; + if dot_scratch.len() < dot_count { + dot_scratch.resize(dot_count, 0.0); + } + diskann_linalg::sgemm_aat_lower( + point_count, + points.ncols(), + points.as_slice(), + &mut dot_scratch[..dot_count], + ) + .map_err(ANNError::new)?; + let dots = MatrixView::try_from(&dot_scratch[..dot_count], point_count, point_count) + .map_err(|error| ANNError::new(error.as_static()))?; + M::prepare_leaf_norms(dots, norm_scratch); + rank_leaf_dots::(arch, dots, norm_scratch, output, worst).map_err(ANNError::new) +} + +/// Rank a prepared lower-triangular Gram matrix. +fn rank_leaf_dots( arch: A, input: MatrixView<'_, f32>, norms: &[f32], mut output: MutMatrixView<'_, LeafNeighbor>, - workspace: &mut LeafKernelWorkspace, + worst: &mut Vec, ) -> Result<(), LeafKernelError> where A: Architecture, @@ -152,120 +136,41 @@ where M: LeafMetric, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { - validate(input, norms, &output)?; + validate_neighbor_count(input, &output)?; let neighbor_count = output.ncols(); if neighbor_count == 0 { return Ok(()); } - resize( - "worst distances", - &mut workspace.worst, - input.nrows(), - f32::INFINITY, - )?; + worst.resize(input.nrows(), f32::INFINITY); output.as_mut_slice().fill(LeafNeighbor::default()); - workspace.worst.fill(f32::INFINITY); + worst.fill(f32::INFINITY); - match (norms.is_empty(), neighbor_count) { - (false, 1) => scan_point_pairs::( - arch, - input, - output.as_mut_slice(), - PreparedLeafNorms(norms), - &mut workspace.worst, - ), - (false, 2) => scan_point_pairs::( - arch, - input, - output.as_mut_slice(), - PreparedLeafNorms(norms), - &mut workspace.worst, - ), - (false, 3) => scan_point_pairs::( - arch, - input, - output.as_mut_slice(), - PreparedLeafNorms(norms), - &mut workspace.worst, - ), - (true, 1) => scan_point_pairs::( - arch, - input, - output.as_mut_slice(), - EmptyLeafNorms, - &mut workspace.worst, - ), - (true, 2) => scan_point_pairs::( - arch, - input, - output.as_mut_slice(), - EmptyLeafNorms, - &mut workspace.worst, - ), - (true, 3) => scan_point_pairs::( + match neighbor_count { + 1 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), + 2 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), + 3 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), + _ => scan_runtime_width::( arch, input, + norms, output.as_mut_slice(), - EmptyLeafNorms, - &mut workspace.worst, + neighbor_count, + worst, ), - _ => { - return Err(LeafKernelError::InvalidNeighborCount { - points: input.nrows(), - neighbors: neighbor_count, - maximum: MAX_LEAF_NEIGHBORS, - }); - } - } - if let Some(source) = output - .as_slice() - .chunks_exact(neighbor_count) - .position(|neighbors| neighbors[neighbor_count - 1].target == u32::MAX) - { - return Err(LeafKernelError::InsufficientRankableNeighbors { - source_index: source, - neighbors: neighbor_count, - }); } Ok(()) } /// Check the safety conditions for the SIMD kernel. /// -/// The matrix views already prove their backing lengths. This function checks -/// that the dot matrix is square. It also checks local-ID range and output width. -/// An error occurs before the kernel changes output or workspace. -fn validate( +/// Check the output width against the number of non-self points. +fn validate_neighbor_count( input: MatrixView<'_, f32>, - norms: &[f32], output: &MutMatrixView<'_, LeafNeighbor>, ) -> Result<(), LeafKernelError> { let point_count = input.nrows(); - let dot_columns = input.ncols(); - if point_count > u32::MAX as usize { - return Err(LeafKernelError::TooManyPoints(point_count)); - } - if point_count != dot_columns { - return Err(LeafKernelError::NonSquareDots { - rows: point_count, - cols: dot_columns, - }); - } - if !norms.is_empty() && norms.len() != point_count { - return Err(LeafKernelError::InvalidNormCount { - expected: point_count, - actual: norms.len(), - }); - } - if output.nrows() != point_count { - return Err(LeafKernelError::InvalidOutputRows { - expected: point_count, - actual: output.nrows(), - columns: output.ncols(), - }); - } - let maximum_neighbors = point_count.saturating_sub(1).min(MAX_LEAF_NEIGHBORS); + let maximum_neighbors = point_count.saturating_sub(1); let neighbor_count = output.ncols(); if neighbor_count > maximum_neighbors { return Err(LeafKernelError::InvalidNeighborCount { @@ -277,136 +182,83 @@ fn validate( Ok(()) } -fn resize( - buffer: &'static str, - values: &mut Vec, - len: usize, - value: T, -) -> Result<(), LeafKernelError> { - let additional = len.saturating_sub(values.len()); - values - .try_reserve(additional) - .map_err(|_| LeafKernelError::Allocation { buffer, additional })?; - values.resize(len, value); - Ok(()) -} - -/// Provide norm values for one leaf scan. -trait LeafNormAccess -where - F: SIMDVector, -{ - /// Repeat one point norm in all SIMD lanes. - fn repeat_simd(self, arch: F::Arch, point: usize) -> F; - - /// Load one complete SIMD group of point norms. - fn load_simd(self, arch: F::Arch, first_point: usize) -> F; - - /// Read one point norm. - fn read(self, point: usize) -> f32; -} - -/// Prepared norm values for all points in one leaf. -#[derive(Clone, Copy)] -struct PreparedLeafNorms<'a>(&'a [f32]); - -impl LeafNormAccess for PreparedLeafNorms<'_> -where - F: SIMDVector, +/// Select neighbors with a fixed output width. +fn scan_fixed_width( + arch: F::Arch, + input: MatrixView<'_, f32>, + norms: &[f32], + output: &mut [LeafNeighbor], + worst: &mut [f32], +) where + F: SIMDVector> + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + M: LeafMetric, + u64: From<<::BitMask as SIMDMask>::Underlying>, { - #[inline(always)] - fn repeat_simd(self, arch: F::Arch, point: usize) -> F { - F::splat(arch, self.0[point]) - } - - #[inline(always)] - fn load_simd(self, arch: F::Arch, first_point: usize) -> F { - let last_point = first_point + F::LANES; - let norm_group = &self.0[first_point..last_point]; - - // SAFETY: `norm_group` contains one complete SIMD group. - unsafe { F::load_simd(arch, norm_group.as_ptr()) } - } - - #[inline(always)] - fn read(self, point: usize) -> f32 { - self.0[point] - } + let (rows, _) = output.as_chunks_mut::(); + scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { + insert_fixed_neighbor(&mut rows[source], target, distance) + }); } -/// Zero norm values for a metric that does not use leaf norms. -#[derive(Clone, Copy)] -struct EmptyLeafNorms; - -impl LeafNormAccess for EmptyLeafNorms -where - F: SIMDVector, +/// Select neighbors with a runtime output width. +fn scan_runtime_width( + arch: F::Arch, + input: MatrixView<'_, f32>, + norms: &[f32], + output: &mut [LeafNeighbor], + width: usize, + worst: &mut [f32], +) where + F: SIMDVector> + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + M: LeafMetric, + u64: From<<::BitMask as SIMDMask>::Underlying>, { - #[inline(always)] - fn repeat_simd(self, arch: F::Arch, point: usize) -> F { - let _ = point; - F::default(arch) - } - - #[inline(always)] - fn load_simd(self, arch: F::Arch, first_point: usize) -> F { - let _ = first_point; - F::default(arch) - } - - #[inline(always)] - fn read(self, point: usize) -> f32 { - let _ = point; - 0.0 - } + scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { + let first = source * width; + insert_runtime_neighbor(&mut output[first..first + width], target, distance) + }); } /// Select neighbors from all unordered point pairs in one leaf. /// /// The function reads the strict lower triangle once. It offers each distance to /// both endpoint lists. SIMD groups and single values preserve pair scan order. -/// -/// `input` supplies square dot products. `norms` supplies metric norm values. #[inline(never)] -fn scan_point_pairs( +fn scan_point_pairs( arch: F::Arch, input: MatrixView<'_, f32>, - output: &mut [LeafNeighbor], - norms: R, + norms: &[f32], worst: &mut [f32], + mut insert: I, ) where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: LeafMetric, - R: LeafNormAccess + Copy, + I: FnMut(usize, u32, f32) -> f32, u64: From<<::BitMask as SIMDMask>::Underlying>, { - let (output, _) = output.as_chunks_mut::(); let point_count = input.nrows(); let dots = input.as_slice(); let worst_ptr = worst.as_mut_ptr(); - // Source zero has no earlier target. Each source after zero can still add - // itself to the neighbor list of source zero. for source in 1..point_count { let source_start = source * point_count; - let source_norms = norms.repeat_simd(arch, source); - let source_norm = norms.read(source); - // SAFETY: `nearest_neighbors` created one threshold for each point. + let source_simd = M::source_simd::(arch, norms, source); + let source_single = M::source_single(norms, source); + // SAFETY: `rank_leaf_dots` created one threshold for each point. let mut source_worst = unsafe { *worst_ptr.add(source) }; let mut target = 0; - let full = source / F::LANES * F::LANES; - - while target < full { - // SAFETY: The full chunk is in this source's strict-lower prefix. - let pair_dots = unsafe { F::load_simd(arch, dots.as_ptr().add(source_start + target)) }; - let target_norms = norms.load_simd(arch, target); - let distances = M::leaf_distance_simd(arch, pair_dots, source_norms, target_norms); - // Every pair may improve the current source and its earlier target. - // Derive both masks before either endpoint mutates its threshold. + let simd_prefix = source - source % F::LANES; + + while target < simd_prefix { + // SAFETY: This complete SIMD group is in the strict-lower prefix. + let dot_products = + unsafe { F::load_simd(arch, dots.as_ptr().add(source_start + target)) }; + let distances = M::distances_simd::(arch, norms, source_simd, dot_products, target); let source_eligible = distances.lt_simd(F::splat(arch, source_worst)); - // SAFETY: The full target chunk is below `source < point_count`. - // `nearest_neighbors` created one threshold for each point. + // SAFETY: The complete target group is below `source < point_count`. 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()); @@ -420,11 +272,7 @@ fn scan_point_pairs( source_bits &= source_bits - 1; let distance = values[lane]; if distance < source_worst { - source_worst = insert_fixed_neighbor( - &mut output[source], - (target + lane) as u32, - distance, - ); + source_worst = insert(source, (target + lane) as u32, distance); } } @@ -433,11 +281,7 @@ fn scan_point_pairs( let lane = target_bits.trailing_zeros() as usize; target_bits &= target_bits - 1; let target_source = target + lane; - let new_worst = insert_fixed_neighbor( - &mut output[target_source], - source as u32, - values[lane], - ); + let new_worst = insert(target_source, source as u32, values[lane]); // SAFETY: `target_source < source < worst.len()`. unsafe { *worst_ptr.add(target_source) = new_worst }; } @@ -447,16 +291,15 @@ fn scan_point_pairs( while target < source { // SAFETY: The target is in this source's strict-lower prefix. - let dot = unsafe { *dots.get_unchecked(source_start + target) }; - let target_norm = norms.read(target); - let distance = M::leaf_distance_single(dot, source_norm, target_norm); + let dot_product = unsafe { *dots.get_unchecked(source_start + target) }; + let distance = M::distance_single(norms, source_single, dot_product, target); if distance < source_worst { - source_worst = insert_fixed_neighbor(&mut output[source], target as u32, distance); + source_worst = insert(source, target as u32, distance); } // SAFETY: `target < source < worst.len()`. let target_worst = unsafe { *worst_ptr.add(target) }; if distance < target_worst { - let new_worst = insert_fixed_neighbor(&mut output[target], source as u32, distance); + let new_worst = insert(target, source as u32, distance); // SAFETY: `target < source < worst.len()`. unsafe { *worst_ptr.add(target) = new_worst }; } @@ -467,11 +310,7 @@ fn scan_point_pairs( } } -/// Insert one target point into a source point's retained neighbor set. -/// -/// `N` is the configured leaf neighbor count. The candidate is closer than the -/// current farthest neighbor. Equal distances keep pair scan order. The function -/// returns the new farthest retained distance. +/// Insert one target into a fixed-width retained neighbor set. #[inline(always)] fn insert_fixed_neighbor( neighbors: &mut [LeafNeighbor; N], @@ -479,39 +318,47 @@ fn insert_fixed_neighbor( distance: f32, ) -> f32 { let entry = LeafNeighbor::new(target, distance); - match N { - 1 => { + if N == 1 { + neighbors[0] = entry; + return distance; + } + if N == 2 { + let first = neighbors[0]; + if distance < first.distance { neighbors[0] = entry; - distance - } - 2 => { - let first = neighbors[0]; - if distance < first.distance { - neighbors[0] = entry; - neighbors[1] = first; - first.distance - } else { - neighbors[1] = entry; - distance - } + neighbors[1] = first; + return first.distance; } - 3 => { - let (first, second) = (neighbors[0], neighbors[1]); - if distance < first.distance { - neighbors[0] = entry; - neighbors[1] = first; - neighbors[2] = second; - } else if distance < second.distance { - neighbors[1] = entry; - neighbors[2] = second; - } else { - neighbors[2] = entry; - return distance; - } - second.distance - } - _ => f32::INFINITY, + neighbors[1] = entry; + return distance; + } + + let (first, second) = (neighbors[0], neighbors[1]); + if distance < first.distance { + neighbors[0] = entry; + neighbors[1] = first; + neighbors[2] = second; + } else if distance < second.distance { + neighbors[1] = entry; + neighbors[2] = second; + } else { + neighbors[2] = entry; + return distance; } + second.distance +} + +/// Insert one target into a runtime-width retained neighbor set. +#[inline(always)] +fn insert_runtime_neighbor(neighbors: &mut [LeafNeighbor], target: u32, distance: f32) -> f32 { + let last = neighbors.len() - 1; + let mut slot = last; + while slot > 0 && distance < neighbors[slot - 1].distance { + neighbors[slot] = neighbors[slot - 1]; + slot -= 1; + } + neighbors[slot] = LeafNeighbor::new(target, distance); + neighbors[last].distance } #[cfg(test)] @@ -539,33 +386,33 @@ where use diskann_vector::distance::Metric; match self.0 { - Metric::L2 => nearest_neighbors::( + Metric::L2 => rank_leaf_dots::( arch, call.input, call.norms, call.output, - call.workspace, + &mut call.workspace.worst, ), - Metric::Cosine => nearest_neighbors::( + Metric::Cosine => rank_leaf_dots::( arch, call.input, call.norms, call.output, - call.workspace, + &mut call.workspace.worst, ), - Metric::CosineNormalized => nearest_neighbors::( + Metric::CosineNormalized => rank_leaf_dots::( arch, call.input, call.norms, call.output, - call.workspace, + &mut call.workspace.worst, ), - Metric::InnerProduct => nearest_neighbors::( + Metric::InnerProduct => rank_leaf_dots::( arch, call.input, call.norms, call.output, - call.workspace, + &mut call.workspace.worst, ), } } @@ -595,16 +442,12 @@ fn prepared_test_norms( metric: diskann_vector::distance::Metric, input: MatrixView<'_, f32>, ) -> Vec { - use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2, NormPreparation}; + use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; use diskann_vector::distance::Metric; fn prepare(input: MatrixView<'_, f32>) -> Vec { let mut norms = Vec::new(); - M::prepare_leaf_norms(NormPreparation { - values: input, - norms: &mut norms, - }) - .unwrap(); + M::prepare_leaf_norms(input, &mut norms); norms } @@ -664,23 +507,11 @@ mod tests { } #[test] - fn neighbor_count_clamps_to_non_self_neighbors_and_rejects_large_k() { - assert_eq!(leaf_neighbor_count(0, 3).unwrap(), 0); - assert_eq!(leaf_neighbor_count(1, 3).unwrap(), 0); - assert_eq!(leaf_neighbor_count(4, 3).unwrap(), 3); - assert_eq!( - leaf_neighbor_count(4, 4), - Err(LeafKernelError::InvalidNeighborCount { - points: 4, - neighbors: 4, - maximum: MAX_LEAF_NEIGHBORS, - }) - ); - #[cfg(target_pointer_width = "64")] - assert_eq!( - leaf_neighbor_count(u32::MAX as usize + 1, 1), - Err(LeafKernelError::TooManyPoints(u32::MAX as usize + 1)) - ); + fn neighbor_count_clamps_to_non_self_neighbors() { + assert_eq!(leaf_neighbor_count(0, 3), 0); + assert_eq!(leaf_neighbor_count(1, 3), 0); + assert_eq!(leaf_neighbor_count(4, 4), 3); + assert_eq!(leaf_neighbor_count(8, 5), 5); } #[test] @@ -691,7 +522,7 @@ mod tests { let norms = prepared_test_norms(Metric::L2, input); let mut workspace = LeafKernelWorkspace::default(); - for neighbor_count in [1, 3, 2] { + for neighbor_count in [1, 3, 5, 2] { let mut output = vec![LeafNeighbor::default(); points * neighbor_count]; dispatch_nearest_neighbors( Metric::L2, @@ -705,6 +536,53 @@ mod tests { } } + #[test] + fn vector_pipeline_selects_exact_l2_neighbors_and_reuses_workspace() { + use super::super::kernel_metric::L2; + + let values = [0.0, 1.0, 3.0, 10.0]; + let points = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let mut output = [LeafNeighbor::default(); 8]; + let mut workspace = LeafKernelWorkspace::default(); + select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + points, + MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), + &mut workspace, + ) + .unwrap(); + + assert_eq!( + output, + [ + LeafNeighbor::new(1, 1.0), + LeafNeighbor::new(2, 9.0), + LeafNeighbor::new(0, 1.0), + LeafNeighbor::new(2, 4.0), + LeafNeighbor::new(1, 4.0), + LeafNeighbor::new(0, 9.0), + LeafNeighbor::new(2, 49.0), + LeafNeighbor::new(1, 81.0), + ] + ); + + let dot_scratch = workspace.dot_scratch.as_ptr(); + let norm_scratch = workspace.norm_scratch.as_ptr(); + let worst = workspace.worst.as_ptr(); + let mut smaller_output = [LeafNeighbor::default(); 6]; + select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + MatrixView::try_from(&values[..3], 3, 1).unwrap(), + MutMatrixView::try_from(&mut smaller_output[..], 3, 2).unwrap(), + &mut workspace, + ) + .unwrap(); + + assert_eq!(workspace.dot_scratch.as_ptr(), dot_scratch); + assert_eq!(workspace.norm_scratch.as_ptr(), norm_scratch); + assert_eq!(workspace.worst.as_ptr(), worst); + } + #[test] fn workspace_can_shrink_and_grow_between_calls() { let mut workspace = LeafKernelWorkspace::default(); @@ -735,8 +613,8 @@ mod integration_tests { use std::cmp::Ordering; use super::{ - LeafKernelError, LeafKernelWorkspace, LeafNeighbor, MAX_LEAF_NEIGHBORS, - dispatch_nearest_neighbors, leaf_neighbor_count, prepared_test_norms, + LeafKernelError, LeafKernelWorkspace, LeafNeighbor, dispatch_nearest_neighbors, + leaf_neighbor_count, prepared_test_norms, }; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -822,8 +700,8 @@ mod integration_tests { 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[source] + norms[target] - 2.0 * dot), - Metric::CosineNormalized => clamp(1.0 - dot), + Metric::L2 => clamp((-2.0_f32).mul_add(dot, norms[source]) + norms[target]), + Metric::CosineNormalized => 1.0 - dot, Metric::InnerProduct => -dot, Metric::Cosine => { let denominator = norms[source] * norms[target]; @@ -832,7 +710,7 @@ mod integration_tests { } else { dot / denominator }; - clamp(1.0 - similarity) + 1.0 - (-1.0_f32).max(1.0_f32.min(similarity)) } }; if distance.partial_cmp(&f32::INFINITY) == Some(Ordering::Less) { @@ -856,7 +734,7 @@ mod integration_tests { requested_k: usize, metric: Metric, ) -> (usize, Vec) { - let leaf_k = leaf_neighbor_count(points, requested_k).unwrap(); + let leaf_k = leaf_neighbor_count(points, requested_k); let input = test_input(dots, points); let norms = prepared_test_norms(metric, input); let mut output = vec![LeafNeighbor::default(); points * leaf_k]; @@ -881,7 +759,7 @@ mod integration_tests { ] { for points in SIMD_BOUNDARY_POINTS { let dots = differential_dots(metric, points); - for requested_k in [1, 2, 3] { + for requested_k in [1, 2, 3, 4, 7] { let expected = brute_force_reference(&dots, points, requested_k, metric); let actual = run_leaf_kernel(&dots, points, requested_k, metric).1; assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); @@ -953,7 +831,44 @@ mod integration_tests { } #[test] - fn clamps_negative_distances_and_preserves_cosine_extremes() { + fn l2_fma_avoids_intermediate_overflow_in_scalar_and_simd_paths() { + let dot = f32::from_bits(f32::MAX.to_bits() - 1); + let expected = (-2.0_f32).mul_add(dot, f32::MAX) + f32::MAX; + assert!(expected.is_finite() && expected > 0.0); + + let scalar = [f32::MAX, 0.0, dot, f32::MAX]; + let scalar_output = run_leaf_kernel(&scalar, 2, 1, Metric::L2).1; + assert_eq!(scalar_output[0].distance.to_bits(), expected.to_bits()); + + let points = 17; + let mut simd = vec![0.0; points * points]; + for point in 0..points { + simd[point * points + point] = f32::MAX; + } + simd[16 * points] = dot; + let simd_output = run_leaf_kernel(&simd, points, 1, Metric::L2).1; + assert_eq!(simd_output[16].target, 0); + assert_eq!(simd_output[16].distance.to_bits(), expected.to_bits()); + } + + #[test] + fn cosine_clamps_simd_similarity_to_metric_range() { + let points = 17; + let mut dots = vec![0.0; points * points]; + for point in 0..points { + dots[point * points + point] = 1.0; + } + dots[16 * points] = 1.000_001; + dots[16 * points + 1] = -1.000_001; + + let output = run_leaf_kernel(&dots, points, 16, Metric::Cosine).1; + let source = &output[16 * 16..17 * 16]; + assert_eq!(source[0], LeafNeighbor::new(0, 0.0)); + assert_eq!(source[15], LeafNeighbor::new(1, 2.0)); + } + + #[test] + fn clamps_leaf_distances_and_cosine_similarity() { #[rustfmt::skip] let out_of_range = [1.0, 0.0, 2.0, 1.0]; assert_eq!( @@ -962,7 +877,7 @@ mod integration_tests { ); assert_eq!( run_leaf_kernel(&out_of_range, 2, 1, Metric::CosineNormalized).1[0].distance, - 0.0 + -1.0 ); assert_eq!( run_leaf_kernel(&out_of_range, 2, 1, Metric::Cosine).1[0].distance, @@ -973,7 +888,7 @@ mod integration_tests { let opposite = [1.0, 0.0, -2.0, 1.0]; assert_eq!( run_leaf_kernel(&opposite, 2, 1, Metric::Cosine).1[0].distance, - 3.0 + 2.0 ); let subnormal = [f32::MIN_POSITIVE / 2.0, 0.0, 1.0, 1.0]; @@ -995,9 +910,8 @@ mod integration_tests { let mut dots = vec![0.0; points * points]; dots[3 * points] = -f32::MAX; - let (leaf_k, output) = - run_leaf_kernel(&dots, points, MAX_LEAF_NEIGHBORS, Metric::InnerProduct); - assert_eq!(leaf_k, MAX_LEAF_NEIGHBORS); + let (leaf_k, output) = run_leaf_kernel(&dots, points, 3, Metric::InnerProduct); + assert_eq!(leaf_k, 3); assert_eq!( output[3 * leaf_k + leaf_k - 1], LeafNeighbor::new(0, f32::MAX) @@ -1005,7 +919,7 @@ mod integration_tests { } #[test] - fn every_metric_ignores_nan_pairs() { + fn leaf_metrics_define_nan_candidate_behavior() { #[rustfmt::skip] let dots = [ 1.0, 0.0, 0.0, @@ -1013,42 +927,19 @@ mod integration_tests { 0.5, 0.25, 1.0, ]; - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { + for metric in [Metric::L2, Metric::Cosine] { + let output = run_leaf_kernel(&dots, 3, 1, metric).1; + assert_eq!(output[0], LeafNeighbor::new(1, 0.0), "metric {metric:?}"); + assert_eq!(output[1], LeafNeighbor::new(0, 0.0), "metric {metric:?}"); + } + + for metric in [Metric::CosineNormalized, Metric::InnerProduct] { let output = run_leaf_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_sources_with_too_few_rankable_neighbors() { - let dots = [1.0, 0.0, f32::NAN, 1.0]; - let mut output = [LeafNeighbor::default(); 2]; - let input = test_input(&dots, 2); - let norms = prepared_test_norms(Metric::L2, input); - let error = dispatch_nearest_neighbors( - Metric::L2, - input, - &norms, - MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::default(), - ) - .unwrap_err(); - - assert_eq!( - error, - LeafKernelError::InsufficientRankableNeighbors { - source_index: 0, - neighbors: 1 - } - ); - } - #[test] fn clamps_k_to_available_non_self_neighbors() { #[rustfmt::skip] @@ -1057,7 +948,7 @@ mod integration_tests { 0.0, 1.0, 3.0, 0.0, 0.0, 1.0, ]; - let (leaf_k, output) = run_leaf_kernel(&dots, 3, MAX_LEAF_NEIGHBORS, Metric::L2); + let (leaf_k, output) = run_leaf_kernel(&dots, 3, 3, Metric::L2); assert_eq!(leaf_k, 2); for (source, neighbors) in output.chunks_exact(leaf_k).enumerate() { @@ -1081,55 +972,10 @@ mod integration_tests { } #[test] - fn rejects_non_square_input_and_invalid_output_dimensions() { - let dots = [0.0; 6]; - let non_square = MatrixView::try_from(&dots[..], 2, 3).unwrap(); - let mut output = [LeafNeighbor::default(); 2]; - assert_eq!( - dispatch_nearest_neighbors( - Metric::L2, - non_square, - &[], - MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::default(), - ), - Err(LeafKernelError::NonSquareDots { rows: 2, cols: 3 }) - ); - + fn rejects_invalid_neighbor_counts() { let square = [0.0; 9]; let square_input = test_input(&square, 3); let square_norms = prepared_test_norms(Metric::L2, square_input); - let mut valid_output = [LeafNeighbor::default(); 3]; - assert_eq!( - dispatch_nearest_neighbors( - Metric::L2, - square_input, - &square_norms[..2], - MutMatrixView::try_from(&mut valid_output[..], 3, 1).unwrap(), - &mut LeafKernelWorkspace::default(), - ), - Err(LeafKernelError::InvalidNormCount { - expected: 3, - actual: 2, - }) - ); - - let mut wrong_rows = [LeafNeighbor::default(); 2]; - assert_eq!( - dispatch_nearest_neighbors( - Metric::L2, - square_input, - &square_norms, - MutMatrixView::try_from(&mut wrong_rows[..], 2, 1).unwrap(), - &mut LeafKernelWorkspace::default(), - ), - Err(LeafKernelError::InvalidOutputRows { - expected: 3, - actual: 2, - columns: 1, - }) - ); - let mut too_many = [LeafNeighbor::default(); 9]; assert_eq!( dispatch_nearest_neighbors( @@ -1149,39 +995,36 @@ mod integration_tests { let square = [0.0; 25]; let square_input = test_input(&square, 5); let square_norms = prepared_test_norms(Metric::L2, square_input); - let mut too_wide = [LeafNeighbor::default(); 20]; + let mut too_wide = [LeafNeighbor::default(); 25]; assert_eq!( dispatch_nearest_neighbors( Metric::L2, square_input, &square_norms, - MutMatrixView::try_from(&mut too_wide[..], 5, 4).unwrap(), + MutMatrixView::try_from(&mut too_wide[..], 5, 5).unwrap(), &mut LeafKernelWorkspace::default(), ), Err(LeafKernelError::InvalidNeighborCount { points: 5, - neighbors: 4, - maximum: MAX_LEAF_NEIGHBORS, + neighbors: 5, + maximum: 4, }) ); } #[test] - fn cosine_zero_norm_masks_nan_norm_at_simd_boundaries() { - for points in [9, 17] { - let mut dots = vec![0.0; points * points]; - for source in 1..points { - dots[source * points + source] = f32::NAN; - } + fn normalized_cosine_keeps_nan_non_rankable_in_scalar_and_simd_paths() { + let scalar = [1.0, 0.0, f32::NAN, 1.0]; + let scalar_output = run_leaf_kernel(&scalar, 2, 1, Metric::CosineNormalized).1; + assert_eq!(scalar_output[0], LeafNeighbor::default()); - let output = run_leaf_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}, source={source}" - ); - } + let points = 17; + let mut dots = vec![0.0; points * points]; + for point in 0..points { + dots[point * points + point] = 1.0; } + dots[16 * points] = f32::NAN; + let simd_output = run_leaf_kernel(&dots, points, 1, Metric::CosineNormalized).1; + assert_eq!(simd_output[16], LeafNeighbor::new(1, 1.0)); } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 1d02179dc2..17c88a4a7b 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -6,17 +6,16 @@ //! Select partition centers for PiPNN point assignment. //! //! A leader is a sampled dataset point that represents one child partition. -//! Each input row contains dot products from one assigned point to all sampled -//! leaders. Each output row contains the nearest leader-column IDs. The scatter -//! step uses each column ID as a child-partition ID. -//! -//! The caller supplies concrete architecture `A` and metric `M`. The function -//! checks row counts, norm units, norm lengths, fanout, and leader-ID range. -//! These checks occur before output changes or unchecked SIMD loads. +//! The kernel prepares reusable leader norms, computes point-to-leader dot +//! products, and returns nearest leader-column IDs for partition scatter. //! //! L2 omits the assigned point's norm because it is constant across all sampled //! leaders. Equal scores keep sampled-leader order. NaN is not rankable. +use std::marker::PhantomData; + +use crate::{ANNError, ANNResult}; +use diskann_linalg::Transpose; use diskann_utils::views::{MatrixView, MutMatrixView}; #[cfg(test)] use diskann_vector::distance::Metric; @@ -26,306 +25,179 @@ use diskann_wide::{ use super::kernel_metric::{PartitionMetric, PartitionNorms}; +/// Sampled leader vectors with metric-specific reusable norms. +pub(super) struct PreparedLeaders<'a, M> { + leader_values: MatrixView<'a, f32>, + leader_norms: Vec, + metric: PhantomData, +} + +impl<'a, M> PreparedLeaders<'a, M> +where + M: PartitionMetric, +{ + /// Prepare leader state for all point stripes in one partition split. + pub(super) fn new(leader_values: MatrixView<'a, f32>) -> Self { + let mut leader_norms = Vec::new(); + M::prepare_leader_norms(leader_values, &mut leader_norms); + Self { + leader_values, + leader_norms, + metric: PhantomData, + } + } + + pub(super) fn len(&self) -> usize { + self.leader_values.nrows() + } +} + +/// Reusable storage for one point-stripe numerical pipeline. +#[derive(Default)] +pub(super) struct PartitionKernelWorkspace { + dot_scratch: Vec, + point_norm_scratch: Vec, + ranked_leader_scratch: Vec<(u32, f32)>, +} + /// Dot products between assigned points and sampled partition centers. /// /// Each row is one point being assigned. Each column is one sampled leader. /// [`Self::norms`] supplies the norm layout for metric `M`. #[derive(Clone, Copy, Debug)] -pub(super) struct PartitionInput<'a> { - pub(super) dots: MatrixView<'a, f32>, - pub(super) norms: PartitionNorms<'a>, -} - -/// Validation error returned by [`nearest_leaders`]. -#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] -pub(super) enum PartitionKernelError { - /// 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 { - expected_rows: usize, - actual_rows: usize, - actual_cols: usize, - }, - /// A metric-specific norm slice has the wrong length. - #[error("invalid {buffer} length: expected {expected}, got {actual}")] - InvalidBufferLength { - buffer: &'static str, - expected: usize, - actual: usize, - }, - /// The requested fanout exceeds the available leader count. - #[error("invalid fanout {fanout}: must not exceed {leader_count} leaders")] - InvalidFanout { fanout: usize, leader_count: usize }, - /// Reusable ranked-leader storage could not be reserved. - #[error("failed to reserve {additional} partition ranked-leader entries")] - Allocation { additional: usize }, - /// Leader positions cannot be represented as `u32`. - #[error("leader count {0} exceeds the u32 position limit")] - TooManyLeaders(usize), - /// A point did not contain enough rankable leaders to fill its output. - #[error("point {point} has fewer than {fanout} rankable leaders")] - InsufficientRankableLeaders { point: usize, fanout: usize }, +struct PartitionInput<'a> { + dots: MatrixView<'a, f32>, + norms: PartitionNorms<'a>, } -/// Select the nearest sampled partition centers for each input point. -/// -/// The output width is the fanout. Each output value is a leader's column ID in -/// `input.dots`. Partition scatter uses that ID to select a child partition. +/// Assign one packed point stripe to prepared partition leaders. /// /// # Errors /// -/// The function returns an error for an invalid shape, norm input, fanout, or allocation. -/// It also returns an error when fewer than `fanout` scores are rankable. -pub(super) fn nearest_leaders( +/// Returns an error for invalid GEMM input. +pub(super) fn assign_leaders( + arch: A, + points: MatrixView<'_, f32>, + leaders: &PreparedLeaders<'_, M>, + output: MutMatrixView<'_, u32>, + workspace: &mut PartitionKernelWorkspace, +) -> ANNResult<()> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: PartitionMetric, +{ + let point_count = points.nrows(); + let leader_count = leaders.len(); + let dot_count = point_count * leader_count; + let PartitionKernelWorkspace { + dot_scratch, + point_norm_scratch, + ranked_leader_scratch, + } = workspace; + if dot_scratch.len() < dot_count { + dot_scratch.resize(dot_count, 0.0); + } + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + point_count, + leader_count, + points.ncols(), + 1.0, + points.as_slice(), + leaders.leader_values.as_slice(), + None, + &mut dot_scratch[..dot_count], + ) + .map_err(ANNError::new)?; + M::prepare_point_norms(points, point_norm_scratch); + let dots = MatrixView::try_from(&dot_scratch[..dot_count], point_count, leader_count) + .map_err(|error| ANNError::new(error.as_static()))?; + rank_leader_dots::( + arch, + PartitionInput { + dots, + norms: PartitionNorms { + point_norms: point_norm_scratch, + leader_norms: &leaders.leader_norms, + }, + }, + output, + ranked_leader_scratch, + ); + Ok(()) +} + +/// Rank prepared point-to-leader dot products. +fn rank_leader_dots( arch: A, input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, ranked_leaders: &mut Vec<(u32, f32)>, -) -> Result<(), PartitionKernelError> -where +) where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, M: PartitionMetric, { - let norms = validate(input, &output)?; let fanout = output.ncols(); if fanout == 0 || input.dots.nrows() == 0 { - return Ok(()); + return; } - let additional = fanout.saturating_sub(ranked_leaders.len()); - ranked_leaders - .try_reserve(additional) - .map_err(|_| PartitionKernelError::Allocation { additional })?; ranked_leaders.resize(fanout, (u32::MAX, f32::INFINITY)); - - match (norms.point_norms.is_empty(), norms.leader_norms.is_empty()) { - (false, false) => select_point_leaders::( - arch, - input.dots, - PreparedNorms(norms.point_norms), - PreparedNorms(norms.leader_norms), - output, - ranked_leaders, - ), - (false, true) => select_point_leaders::( - arch, - input.dots, - PreparedNorms(norms.point_norms), - EmptyNorms, - output, - ranked_leaders, - ), - (true, false) => select_point_leaders::( - arch, - input.dots, - EmptyNorms, - PreparedNorms(norms.leader_norms), - output, - ranked_leaders, - ), - (true, true) => select_point_leaders::( - arch, - input.dots, - EmptyNorms, - EmptyNorms, - output, - ranked_leaders, - ), - } -} - -/// This function checks row counts, leader IDs, fanout, and norm lengths. -fn validate<'a>( - input: PartitionInput<'a>, - output: &MutMatrixView<'_, u32>, -) -> Result, PartitionKernelError> { - let point_count = input.dots.nrows(); - let leader_count = input.dots.ncols(); - let fanout = output.ncols(); - - if output.nrows() != point_count { - return Err(PartitionKernelError::InvalidOutputShape { - expected_rows: point_count, - actual_rows: output.nrows(), - actual_cols: output.ncols(), - }); - } - if leader_count > u32::MAX as usize { - return Err(PartitionKernelError::TooManyLeaders(leader_count)); - } - if fanout > leader_count { - return Err(PartitionKernelError::InvalidFanout { - fanout, - leader_count, - }); - } - - check_norm_count("point norms", input.norms.point_norms, point_count)?; - check_norm_count("leader norms", input.norms.leader_norms, leader_count)?; - Ok(input.norms) -} - -/// Check one optional norm buffer. -/// -/// An empty slice means that the active metric does not use this norm. -fn check_norm_count( - buffer: &'static str, - norms: &[f32], - expected: usize, -) -> Result<(), PartitionKernelError> { - if norms.is_empty() || norms.len() == expected { - Ok(()) - } else { - Err(PartitionKernelError::InvalidBufferLength { - buffer, - expected, - actual: norms.len(), - }) - } -} - -/// Provide norm values for partition ranking. -trait NormValues -where - F: SIMDVector, -{ - /// Repeat one norm in all SIMD lanes. - fn repeat_simd(self, arch: F::Arch, point: usize) -> F; - - /// Load one complete SIMD group of norms. - fn load_simd(self, arch: F::Arch, first_point: usize) -> F; - - /// Read one norm. - fn read(self, point: usize) -> f32; -} - -/// Prepared norm values for points or sampled leaders. -#[derive(Clone, Copy)] -struct PreparedNorms<'a>(&'a [f32]); - -impl NormValues for PreparedNorms<'_> -where - F: SIMDVector, -{ - #[inline(always)] - fn repeat_simd(self, arch: F::Arch, point: usize) -> F { - F::splat(arch, self.0[point]) - } - - #[inline(always)] - fn load_simd(self, arch: F::Arch, first_point: usize) -> F { - let last_point = first_point + F::LANES; - let norm_group = &self.0[first_point..last_point]; - - // SAFETY: `norm_group` contains one complete SIMD group. - unsafe { F::load_simd(arch, norm_group.as_ptr()) } - } - - #[inline(always)] - fn read(self, point: usize) -> f32 { - self.0[point] - } -} - -/// Zero norm values for a metric that does not use one norm type. -#[derive(Clone, Copy)] -struct EmptyNorms; - -impl NormValues for EmptyNorms -where - F: SIMDVector, -{ - #[inline(always)] - fn repeat_simd(self, arch: F::Arch, point: usize) -> F { - let _ = point; - F::default(arch) - } - - #[inline(always)] - fn load_simd(self, arch: F::Arch, first_point: usize) -> F { - let _ = first_point; - F::default(arch) - } - - #[inline(always)] - fn read(self, point: usize) -> f32 { - let _ = point; - 0.0 - } + select_point_leaders::(arch, input.dots, input.norms, output, ranked_leaders); } /// Rank sampled partition centers for each assigned point. /// -/// The function converts point-to-leader dot products to metric `M` scores. It -/// keeps the nearest `output.ncols()` centers in sampled-leader order for ties. -/// NaN and positive infinity are not rankable. -/// -/// `ranked_leaders` stores the retained center-column IDs and scores for the current -/// point. The function resets this state before it processes another point. -fn select_point_leaders( +/// The function keeps nearest-first order for every point. Full SIMD groups use +/// metric-specific formulas. Remaining leaders use the matching single formula. +fn select_point_leaders( arch: F::Arch, dots: MatrixView<'_, f32>, - point_norms: P, - leader_norms: L, + norms: PartitionNorms<'_>, mut output: MutMatrixView<'_, u32>, ranked_leaders: &mut [(u32, f32)], -) -> Result<(), PartitionKernelError> -where +) where F: SIMDVector> + SIMDFloat + std::ops::Div, F::Mask: SIMDSelect, M: PartitionMetric, - P: NormValues + Copy, - L: NormValues + Copy, u64: From<<::BitMask as SIMDMask>::Underlying>, { let leader_count = dots.ncols(); let fanout = output.ncols(); - // Reset the retained leaders for each point. No assignment state can pass - // from one output row to another. + for (point, (point_dots, point_output)) in dots .row_iter() .zip(output.as_mut_slice().chunks_exact_mut(fanout)) .enumerate() { ranked_leaders.fill((u32::MAX, f32::INFINITY)); - let point_norm_values = point_norms.repeat_simd(arch, point); - let point_norm = point_norms.read(point); - // Process all complete SIMD groups first. Single-value rankings use the - // non-vector operation order. - let full = leader_count / F::LANES * F::LANES; - - for first_leader in (0..full).step_by(F::LANES) { - // SAFETY: `first_leader + F::LANES <= full <= point_dots.len()`. - let dot = unsafe { F::load_simd(arch, point_dots.as_ptr().add(first_leader)) }; - let leader_norm_values = leader_norms.load_simd(arch, first_leader); - let scores = - M::partition_ranking_simd(arch, dot, point_norm_values, leader_norm_values); - insert_leader_lanes(scores, first_leader, ranked_leaders); + let point_simd = M::point_simd::(arch, norms, point); + let point_single = M::point_single(norms, point); + let simd_prefix = leader_count - leader_count % F::LANES; + + for first_leader in (0..simd_prefix).step_by(F::LANES) { + // SAFETY: This group is inside the point's leader row. + let dot_products = unsafe { F::load_simd(arch, point_dots.as_ptr().add(first_leader)) }; + let rankings = + M::rankings_simd::(arch, norms, point_simd, dot_products, first_leader); + insert_leader_lanes(rankings, first_leader, ranked_leaders); } - // Use single-value formulas because SIMD padding can change L2 rounding. - for (leader, &dot) in point_dots.iter().enumerate().skip(full) { - let leader_norm = leader_norms.read(leader); - insert_leader( - ranked_leaders, - leader as u32, - M::partition_ranking_single(dot, point_norm, leader_norm), - ); - } - if ranked_leaders[fanout - 1].0 == u32::MAX { - return Err(PartitionKernelError::InsufficientRankableLeaders { point, fanout }); + for (leader, &dot_product) in point_dots.iter().enumerate().skip(simd_prefix) { + let ranking = M::ranking_single(norms, point_single, dot_product, leader); + insert_leader(ranked_leaders, leader as u32, ranking); } - // Scatter needs sampled-center column IDs. Scores remain in scratch. for (destination, &(leader, _)) in point_output.iter_mut().zip(ranked_leaders.iter()) { *destination = leader; } } - Ok(()) } /// Offer one SIMD group of sampled centers to the current point's ranked_leaders. @@ -383,32 +255,30 @@ struct DispatchedPartitionCall<'a> { struct DispatchPartitionForTest(Metric); #[cfg(test)] -impl - diskann_wide::arch::Target1, DispatchedPartitionCall<'_>> - for DispatchPartitionForTest +impl diskann_wide::arch::Target1> for DispatchPartitionForTest where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { - fn run(self, arch: A, call: DispatchedPartitionCall<'_>) -> Result<(), PartitionKernelError> { + fn run(self, arch: A, call: DispatchedPartitionCall<'_>) { use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { Metric::L2 => { - nearest_leaders::(arch, call.input, call.output, call.ranked_leaders) + rank_leader_dots::(arch, call.input, call.output, call.ranked_leaders) } Metric::Cosine => { - nearest_leaders::(arch, call.input, call.output, call.ranked_leaders) + rank_leader_dots::(arch, call.input, call.output, call.ranked_leaders) } - Metric::CosineNormalized => nearest_leaders::( + Metric::CosineNormalized => rank_leader_dots::( arch, call.input, call.output, call.ranked_leaders, ), - Metric::InnerProduct => nearest_leaders::( + Metric::InnerProduct => rank_leader_dots::( arch, call.input, call.output, @@ -424,7 +294,7 @@ fn dispatch_nearest_leaders( input: PartitionInput<'_>, output: MutMatrixView<'_, u32>, ranked_leaders: &mut Vec<(u32, f32)>, -) -> Result<(), PartitionKernelError> { +) { diskann_wide::arch::dispatch1_no_features( DispatchPartitionForTest(metric), DispatchedPartitionCall { @@ -473,17 +343,13 @@ mod tests { .zip(output.chunks_exact_mut(fanout)) .enumerate() { - let point_norm = input.norms.point_norms.get(point).copied().unwrap_or(0.0); + let point_norm = M::point_single(input.norms, point); let mut ranked_leaders = vec![(u32::MAX, f32::INFINITY); fanout]; for (leader, &dot) in point_dots.iter().enumerate() { insert_leader( &mut ranked_leaders, leader as u32, - M::partition_ranking_single( - dot, - point_norm, - input.norms.leader_norms.get(leader).copied().unwrap_or(0.0), - ), + M::ranking_single(input.norms, point_norm, dot, leader), ); } for (destination, &(leader, _)) in point_output.iter_mut().zip(&ranked_leaders) { @@ -497,11 +363,11 @@ mod tests { point_norms: &[f32], leader_norms: &[f32], ) -> f32 { - M::partition_ranking_single( - dot_product, - point_norms.first().copied().unwrap_or(0.0), - leader_norms.first().copied().unwrap_or(0.0), - ) + let norms = PartitionNorms { + point_norms, + leader_norms, + }; + M::ranking_single(norms, M::point_single(norms, 0), dot_product, 0) } #[test] @@ -510,8 +376,10 @@ mod tests { assert_eq!(single_ranking::(0.25, &[], &[]), 0.75); assert_eq!(single_ranking::(3.0, &[], &[]), -3.0); assert_eq!(single_ranking::(4.0, &[2.0], &[4.0]), 0.5); + assert_eq!(single_ranking::(5.0, &[2.0], &[2.0]), 0.0); + assert_eq!(single_ranking::(-5.0, &[2.0], &[2.0]), 2.0); assert_eq!(single_ranking::(4.0, &[0.0], &[4.0]), 1.0); - assert!(single_ranking::(1.0, &[f32::NAN], &[1.0]).is_nan()); + assert_eq!(single_ranking::(1.0, &[f32::NAN], &[1.0]), 0.0); } #[test] @@ -541,12 +409,11 @@ mod tests { input, MutMatrixView::try_from(actual.as_mut_slice(), point_norms.len(), 2).unwrap(), &mut Vec::new(), - ) - .unwrap(); + ); assert_eq!(actual, expected); assert_eq!(&actual[..4], &[0, 1, 0, 1]); - assert_eq!(&actual[6..], &[0, 1]); + assert_eq!(&actual[6..], &[2, 3]); } #[test] @@ -560,6 +427,48 @@ mod tests { assert_eq!(ranked_leaders[..], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); } + #[test] + fn vector_pipeline_assigns_cosine_leaders_and_reuses_workspace() { + let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; + let leaders = + PreparedLeaders::::new(MatrixView::try_from(&leader_values[..], 3, 2).unwrap()); + let point_values = [0.9, 0.1, -0.8, 0.2]; + let points = MatrixView::try_from(&point_values[..], 2, 2).unwrap(); + let mut output = [u32::MAX; 4]; + let mut workspace = PartitionKernelWorkspace::default(); + assign_leaders::<_, Cosine>( + diskann_wide::ARCH, + points, + &leaders, + MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), + &mut workspace, + ) + .unwrap(); + + assert_eq!(output, [0, 1, 2, 1]); + let dot_scratch = workspace.dot_scratch.as_ptr(); + let point_norm_scratch = workspace.point_norm_scratch.as_ptr(); + let ranked_leader_scratch = workspace.ranked_leader_scratch.as_ptr(); + + let mut smaller_output = [u32::MAX; 2]; + assign_leaders::<_, Cosine>( + diskann_wide::ARCH, + MatrixView::try_from(&point_values[..2], 1, 2).unwrap(), + &leaders, + MutMatrixView::try_from(&mut smaller_output[..], 1, 2).unwrap(), + &mut workspace, + ) + .unwrap(); + + assert_eq!(smaller_output, [0, 1]); + assert_eq!(workspace.dot_scratch.as_ptr(), dot_scratch); + assert_eq!(workspace.point_norm_scratch.as_ptr(), point_norm_scratch); + assert_eq!( + workspace.ranked_leader_scratch.as_ptr(), + ranked_leader_scratch + ); + } + #[test] fn ranked_leaders_reuses_runtime_fanout_capacity() { let dots = [0.0; 32]; @@ -571,8 +480,7 @@ mod tests { input, MutMatrixView::try_from(&mut wide_output[..], 1, 32).unwrap(), &mut ranked_leaders, - ) - .unwrap(); + ); let allocation = ranked_leaders.as_ptr(); let mut narrow_output = [u32::MAX; 3]; @@ -581,8 +489,7 @@ mod tests { input, MutMatrixView::try_from(&mut narrow_output[..], 1, 3).unwrap(), &mut ranked_leaders, - ) - .unwrap(); + ); assert_eq!(ranked_leaders.as_ptr(), allocation); assert_eq!(ranked_leaders.len(), 3); @@ -595,8 +502,8 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod integration_tests { - use super::{PartitionInput, PartitionKernelError, PartitionNorms, dispatch_nearest_leaders}; - use diskann_utils::views::{MatrixView, MutMatrixView}; + use super::{PartitionInput, PartitionNorms, dispatch_nearest_leaders}; + use diskann_utils::views::{Matrix, MatrixView}; use diskann_vector::distance::Metric; fn test_input<'a>( @@ -635,14 +542,15 @@ mod integration_tests { .filter_map(|(leader, &dot)| { let leader_norm = leader_norms.get(leader).copied().unwrap_or(0.0); let score = match metric { - Metric::L2 => leader_norm - 2.0 * dot, + Metric::L2 => (-2.0_f32).mul_add(dot, leader_norm), Metric::CosineNormalized => 1.0 - dot, Metric::InnerProduct => -dot, Metric::Cosine => { 1.0 - if point_norm == 0.0 || leader_norm == 0.0 { 0.0 } else { - dot / (point_norm * leader_norm) + let cosine = dot / (point_norm * leader_norm); + (-1.0_f32).max(1.0_f32.min(cosine)) } } }; @@ -705,19 +613,10 @@ mod integration_tests { (dots, point_norms, leader_norms) } - fn run_partition_kernel( - metric: Metric, - input: PartitionInput<'_>, - fanout: usize, - ) -> Result, PartitionKernelError> { - let mut output = vec![u32::MAX; input.dots.nrows() * fanout]; - dispatch_nearest_leaders( - metric, - input, - MutMatrixView::try_from(output.as_mut_slice(), input.dots.nrows(), fanout).unwrap(), - &mut Vec::new(), - )?; - Ok(output) + fn run_partition_kernel(metric: Metric, input: PartitionInput<'_>, fanout: usize) -> Vec { + let mut output = Matrix::new(u32::MAX, input.dots.nrows(), fanout); + dispatch_nearest_leaders(metric, input, output.as_mut_view(), &mut Vec::new()); + output.into_inner().into_vec() } #[test] @@ -736,7 +635,7 @@ mod integration_tests { continue; } assert_eq!( - run_partition_kernel(metric, input, fanout).unwrap(), + run_partition_kernel(metric, input, fanout), brute_force_reference(input, fanout, metric), "{metric:?}, leaders={leader_count}, k={fanout}" ); @@ -755,13 +654,13 @@ mod integration_tests { let norms = [0.0, 1.0, 4.0, 9.0]; assert_eq!( - run_partition_kernel(Metric::L2, test_input(&dots, 2, 4, &[], &norms), 2).unwrap(), + run_partition_kernel(Metric::L2, test_input(&dots, 2, 4, &[], &norms), 2), [0, 1, 2, 1] ); } #[test] - fn l2_single_can_outrank_a_fused_simd_lane() { + fn l2_single_matches_fused_simd_ranking() { let mut dots = [0.0; 17]; dots[0] = f32::MAX; dots[16] = f32::MAX; @@ -772,9 +671,8 @@ mod integration_tests { Metric::L2, test_input(&dots, 1, 17, &[], &leader_squared_norms,), 1, - ) - .unwrap(), - [16] + ), + [0] ); } @@ -801,8 +699,7 @@ mod integration_tests { metric, test_input(&dots, 2, 3, point_norms, leader_norms), 2, - ) - .unwrap(), + ), expected, "metric {metric:?}" ); @@ -816,8 +713,7 @@ mod integration_tests { Metric::Cosine, test_input(&[100.0, -100.0], 1, 2, &[0.0], &[1.0, 1.0]), 2, - ) - .unwrap(), + ), [0, 1] ); } @@ -827,8 +723,7 @@ mod integration_tests { let mut dots = [0.0; 8]; dots[7] = -f32::MAX; assert_eq!( - run_partition_kernel(Metric::InnerProduct, test_input(&dots, 1, 8, &[], &[]), 8) - .unwrap(), + run_partition_kernel(Metric::InnerProduct, test_input(&dots, 1, 8, &[], &[]), 8), [0, 1, 2, 3, 4, 5, 6, 7] ); } @@ -840,104 +735,24 @@ mod integration_tests { Metric::InnerProduct, test_input(&[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), 2, - ) - .unwrap(), + ), [1, 2] ); } #[test] - fn rejects_points_with_too_few_rankable_leaders() { - assert_eq!( - run_partition_kernel( - Metric::InnerProduct, - test_input(&[f32::NAN, 3.0], 1, 2, &[], &[]), - 2, - ), - Err(PartitionKernelError::InsufficientRankableLeaders { - point: 0, - fanout: 2, - }) + fn accepts_empty_points_and_zero_fanout() { + assert!( + run_partition_kernel(Metric::InnerProduct, test_input(&[], 0, 3, &[], &[]), 2) + .is_empty() ); - } - - #[test] - fn accepts_empty_points_zero_fanout_and_largest_leader_id() { - run_partition_kernel(Metric::InnerProduct, test_input(&[], 0, 3, &[], &[]), 2).unwrap(); - run_partition_kernel( - Metric::InnerProduct, - test_input(&[1.0, 2.0, 3.0], 1, 3, &[], &[]), - 0, - ) - .unwrap(); - run_partition_kernel( - Metric::InnerProduct, - test_input(&[], 0, u32::MAX as usize, &[], &[]), - 0, - ) - .unwrap(); - - #[cfg(target_pointer_width = "64")] - assert_eq!( + assert!( run_partition_kernel( Metric::InnerProduct, - test_input(&[], 0, u32::MAX as usize + 1, &[], &[],), + test_input(&[1.0, 2.0, 3.0], 1, 3, &[], &[]), 0, - ), - Err(PartitionKernelError::TooManyLeaders(u32::MAX as usize + 1)) - ); - } - - #[test] - fn rejects_wrong_output_norms_and_fanout() { - let dots = [0.0; 6]; - let valid_input = test_input(&dots, 2, 3, &[], &[]); - let mut wrong_output = [u32::MAX; 3]; - assert_eq!( - dispatch_nearest_leaders( - Metric::InnerProduct, - valid_input, - MutMatrixView::try_from(&mut wrong_output[..], 1, 3).unwrap(), - &mut Vec::new(), - ), - Err(PartitionKernelError::InvalidOutputShape { - expected_rows: 2, - actual_rows: 1, - actual_cols: 3, - }) - ); - - let short_norms = [0.0; 2]; - let wrong_norms = PartitionInput { - dots: MatrixView::try_from(&dots[..], 2, 3).unwrap(), - norms: PartitionNorms { - point_norms: &[], - leader_norms: &short_norms, - }, - }; - assert_eq!( - run_partition_kernel(Metric::L2, wrong_norms, 2), - Err(PartitionKernelError::InvalidBufferLength { - buffer: "leader norms", - expected: 3, - actual: 2, - }) - ); - assert_eq!( - run_partition_kernel(Metric::InnerProduct, valid_input, 4), - Err(PartitionKernelError::InvalidFanout { - fanout: 4, - leader_count: 3, - }) - ); - - let one = [0.0]; - assert_eq!( - run_partition_kernel(Metric::InnerProduct, test_input(&one, 1, 1, &[], &[]), 2,), - Err(PartitionKernelError::InvalidFanout { - fanout: 2, - leader_count: 1, - }) + ) + .is_empty() ); } } From 01e5148f96ce0c67b6d24c3fc5f1f1049132c83a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:40:54 +0000 Subject: [PATCH 77/80] fix(pipnn): mark unrankable kernel slots --- diskann/src/graph/pipnn/leaf_kernel.rs | 8 ++++++- diskann/src/graph/pipnn/partition_kernel.rs | 25 +++++++++++++++++---- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index a3c9b6570b..20a707317b 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -14,7 +14,8 @@ //! widths use the runtime insertion loop. //! //! Strict comparisons keep scan order for equal distances. They do not rank NaN. -//! All supported metrics use the same SIMD-group and single-value traversal. +//! An unfilled output slot contains [`LeafNeighbor::default`]. All supported +//! metrics use the same SIMD-group and single-value traversal. //! //! The caller supplies concrete architecture `A` and metric `M`. The private //! dot ranker receives the square matrix created by this module. @@ -43,6 +44,11 @@ impl LeafNeighbor { pub(super) const fn new(target: u32, distance: f32) -> Self { Self { target, distance } } + + /// Return true when this slot contains a rankable leaf-local target. + pub(super) const fn is_assigned(self) -> bool { + self.target != u32::MAX + } } impl Default for LeafNeighbor { diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 17c88a4a7b..77fe3c3d23 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -10,7 +10,8 @@ //! products, and returns nearest leader-column IDs for partition scatter. //! //! L2 omits the assigned point's norm because it is constant across all sampled -//! leaders. Equal scores keep sampled-leader order. NaN is not rankable. +//! leaders. Equal scores keep sampled-leader order. NaN is not rankable. An +//! unfilled output slot contains [`UNASSIGNED_LEADER`]. use std::marker::PhantomData; @@ -25,6 +26,9 @@ use diskann_wide::{ use super::kernel_metric::{PartitionMetric, PartitionNorms}; +/// No sampled partition center was rankable for this output slot. +pub(super) const UNASSIGNED_LEADER: u32 = u32::MAX; + /// Sampled leader vectors with metric-specific reusable norms. pub(super) struct PreparedLeaders<'a, M> { leader_values: MatrixView<'a, f32>, @@ -72,6 +76,9 @@ struct PartitionInput<'a> { /// Assign one packed point stripe to prepared partition leaders. /// +/// A point can have fewer assignments than the output width. Each remaining +/// slot contains [`UNASSIGNED_LEADER`]. +/// /// # Errors /// /// Returns an error for invalid GEMM input. @@ -149,7 +156,7 @@ fn rank_leader_dots( return; } - ranked_leaders.resize(fanout, (u32::MAX, f32::INFINITY)); + ranked_leaders.resize(fanout, (UNASSIGNED_LEADER, f32::INFINITY)); select_point_leaders::(arch, input.dots, input.norms, output, ranked_leaders); } @@ -177,7 +184,7 @@ fn select_point_leaders( .zip(output.as_mut_slice().chunks_exact_mut(fanout)) .enumerate() { - ranked_leaders.fill((u32::MAX, f32::INFINITY)); + ranked_leaders.fill((UNASSIGNED_LEADER, f32::INFINITY)); let point_simd = M::point_simd::(arch, norms, point); let point_single = M::point_single(norms, point); let simd_prefix = leader_count - leader_count % F::LANES; @@ -502,7 +509,7 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod integration_tests { - use super::{PartitionInput, PartitionNorms, dispatch_nearest_leaders}; + use super::{PartitionInput, PartitionNorms, UNASSIGNED_LEADER, dispatch_nearest_leaders}; use diskann_utils::views::{Matrix, MatrixView}; use diskann_vector::distance::Metric; @@ -644,6 +651,16 @@ mod integration_tests { } } + #[test] + fn non_rankable_leaders_leave_unassigned_slots() { + let dots = [-0.25, f32::NAN]; + + assert_eq!( + run_partition_kernel(Metric::InnerProduct, test_input(&dots, 1, 2, &[], &[]), 2), + [0, UNASSIGNED_LEADER] + ); + } + #[test] fn l2_keeps_the_first_leader_when_boundary_distances_tie() { #[rustfmt::skip] From 0f1e47bb174a30b8bf94fde66a038d3e4356c91a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:00:26 +0000 Subject: [PATCH 78/80] refactor(pipnn): centralize SIMD schema --- diskann-wide/src/traits.rs | 2 +- diskann/src/graph/pipnn/kernel_metric.rs | 13 +- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 127 ++++++++++-------- .../graph/pipnn/kernel_metric/partition.rs | 117 +++++++++------- diskann/src/graph/pipnn/leaf_kernel.rs | 79 +++++------ diskann/src/graph/pipnn/mod.rs | 2 + diskann/src/graph/pipnn/partition_kernel.rs | 54 +++----- diskann/src/graph/pipnn/simd.rs | 61 +++++++++ 8 files changed, 271 insertions(+), 184 deletions(-) create mode 100644 diskann/src/graph/pipnn/simd.rs diff --git a/diskann-wide/src/traits.rs b/diskann-wide/src/traits.rs index 09150f0c7d..a233622615 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. diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index d3869cdbba..be6a287e8e 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -11,7 +11,7 @@ mod partition; pub(super) use leaf::LeafMetric; pub(super) use partition::PartitionMetric; -use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; +use super::simd::PiPNNSIMDVector; pub(super) struct L2; pub(super) struct Cosine; @@ -32,8 +32,7 @@ pub(super) struct PartitionNorms<'a> { #[inline(always)] pub(super) fn cosine_distance_simd(arch: F::Arch, dot: F, source_norm: F, target_norm: F) -> F where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + F: PiPNNSIMDVector, { let zero = F::default(arch); let one = F::splat(arch, 1.0); @@ -41,8 +40,12 @@ where 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)); + let safe_denominator = F::select(source_zero, one, F::select(target_zero, one, denominator)); + let cosine = F::select( + source_zero, + zero, + F::select(target_zero, zero, dot / safe_denominator), + ); let negative_one = F::splat(arch, -1.0); one - negative_one.max_simd(cosine.min_simd(one)) } diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 2ac408e206..2a558baac5 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -4,14 +4,20 @@ */ use diskann_utils::views::MatrixView; -use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; +use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDVector}; +use super::super::simd::{PiPNNSIMDSchema, PiPNNSIMDVector}; use super::{ Cosine, CosineNormalized, InnerProduct, L2, cosine_distance_simd, cosine_distance_single, }; /// Compute leaf distances for one concrete metric. pub(in super::super) trait LeafMetric: Send + Sync + 'static { + /// SIMD representation for leaf distance scores. + type Simd: PiPNNSIMDVector + where + A: PiPNNSIMDSchema; + /// Prepare one contiguous metric-specific norm for each leaf-local point. fn prepare_leaf_norms(_dots: MatrixView<'_, f32>, norms: &mut Vec) { norms.clear(); @@ -19,11 +25,11 @@ pub(in super::super) trait LeafMetric: Send + Sync + 'static { /// Prepare one source norm for reuse across SIMD target groups. #[inline(always)] - fn source_simd(arch: F::Arch, _norms: &[f32], _source: usize) -> F + fn source_simd(arch: A, _norms: &[f32], _source: usize) -> Self::Simd where - F: SIMDVector, + A: PiPNNSIMDSchema, { - F::default(arch) + Self::Simd::::default(arch) } /// Prepare one source norm for reuse across single target values. @@ -33,16 +39,15 @@ pub(in super::super) trait LeafMetric: Send + Sync + 'static { } /// Compute distances for one complete SIMD group. - fn distances_simd( - arch: F::Arch, + fn distances_simd( + arch: A, norms: &[f32], - source_norms: F, - dot_products: F, + source_norms: Self::Simd, + dot_products: Self::Simd, first_target: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect; + A: PiPNNSIMDSchema; /// Compute one distance outside the complete SIMD prefix. fn distance_single(norms: &[f32], source_norm: f32, dot_product: f32, target: usize) -> f32; @@ -52,7 +57,7 @@ pub(in super::super) trait LeafMetric: Send + Sync + 'static { #[inline(always)] fn load_norms_simd(arch: F::Arch, norms: &[f32], first_norm: usize) -> F where - F: SIMDVector, + F: PiPNNSIMDVector, { let last_norm = first_norm + F::LANES; let norm_group = &norms[first_norm..last_norm]; @@ -62,6 +67,11 @@ where } impl LeafMetric for L2 { + type Simd + = A::LeafScore + where + A: PiPNNSIMDSchema; + fn prepare_leaf_norms(dots: MatrixView<'_, f32>, norms: &mut Vec) { norms.resize(dots.nrows(), 0.0); for (point, norm) in norms.iter_mut().enumerate() { @@ -70,11 +80,11 @@ impl LeafMetric for L2 { } #[inline(always)] - fn source_simd(arch: F::Arch, norms: &[f32], source: usize) -> F + fn source_simd(arch: A, norms: &[f32], source: usize) -> Self::Simd where - F: SIMDVector, + A: PiPNNSIMDSchema, { - F::splat(arch, norms[source]) + Self::Simd::::splat(arch, norms[source]) } #[inline(always)] @@ -83,20 +93,19 @@ impl LeafMetric for L2 { } #[inline(always)] - fn distances_simd( - arch: F::Arch, + fn distances_simd( + arch: A, norms: &[f32], - source_norms: F, - dot_products: F, + source_norms: Self::Simd, + dot_products: Self::Simd, first_target: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, { - let target_norms = load_norms_simd::(arch, norms, first_target); - (F::splat(arch, -2.0).mul_add_simd(dot_products, source_norms) + target_norms) - .max_simd(F::default(arch)) + let target_norms = load_norms_simd::>(arch, norms, first_target); + (Self::Simd::::splat(arch, -2.0).mul_add_simd(dot_products, source_norms) + target_norms) + .max_simd(Self::Simd::::default(arch)) } #[inline(always)] @@ -106,6 +115,11 @@ impl LeafMetric for L2 { } impl LeafMetric for Cosine { + type Simd + = A::LeafScore + where + A: PiPNNSIMDSchema; + fn prepare_leaf_norms(dots: MatrixView<'_, f32>, norms: &mut Vec) { norms.resize(dots.nrows(), 0.0); for (point, norm) in norms.iter_mut().enumerate() { @@ -114,11 +128,11 @@ impl LeafMetric for Cosine { } #[inline(always)] - fn source_simd(arch: F::Arch, norms: &[f32], source: usize) -> F + fn source_simd(arch: A, norms: &[f32], source: usize) -> Self::Simd where - F: SIMDVector, + A: PiPNNSIMDSchema, { - F::splat(arch, norms[source]) + Self::Simd::::splat(arch, norms[source]) } #[inline(always)] @@ -127,20 +141,19 @@ impl LeafMetric for Cosine { } #[inline(always)] - fn distances_simd( - arch: F::Arch, + fn distances_simd( + arch: A, norms: &[f32], - source_norms: F, - dot_products: F, + source_norms: Self::Simd, + dot_products: Self::Simd, first_target: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, { - let target_norms = load_norms_simd::(arch, norms, first_target); + let target_norms = load_norms_simd::>(arch, norms, first_target); cosine_distance_simd(arch, dot_products, source_norms, target_norms) - .max_simd(F::default(arch)) + .max_simd(Self::Simd::::default(arch)) } #[inline(always)] @@ -150,19 +163,23 @@ impl LeafMetric for Cosine { } impl LeafMetric for CosineNormalized { + type Simd + = A::LeafScore + where + A: PiPNNSIMDSchema; + #[inline(always)] - fn distances_simd( - arch: F::Arch, + fn distances_simd( + arch: A, _norms: &[f32], - _source_norms: F, - dot_products: F, + _source_norms: Self::Simd, + dot_products: Self::Simd, _first_target: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, { - F::splat(arch, 1.0) - dot_products + Self::Simd::::splat(arch, 1.0) - dot_products } #[inline(always)] @@ -172,19 +189,23 @@ impl LeafMetric for CosineNormalized { } impl LeafMetric for InnerProduct { + type Simd + = A::LeafScore + where + A: PiPNNSIMDSchema; + #[inline(always)] - fn distances_simd( - arch: F::Arch, + fn distances_simd( + arch: A, _norms: &[f32], - _source_norms: F, - dot_products: F, + _source_norms: Self::Simd, + dot_products: Self::Simd, _first_target: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, { - F::default(arch) - dot_products + Self::Simd::::default(arch) - dot_products } #[inline(always)] diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index ad237ec01a..0490b46692 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -5,8 +5,9 @@ use diskann_utils::views::MatrixView; use diskann_vector::{Norm, norm::FastL2NormSquared}; -use diskann_wide::{SIMDFloat, SIMDSelect, SIMDVector}; +use diskann_wide::{SIMDMulAdd, SIMDVector}; +use super::super::simd::{PiPNNSIMDSchema, PiPNNSIMDVector}; use super::{ Cosine, CosineNormalized, InnerProduct, L2, PartitionNorms, cosine_distance_simd, cosine_distance_single, @@ -14,6 +15,11 @@ use super::{ /// Compute partition rankings for one concrete metric. pub(in super::super) trait PartitionMetric: Send + Sync + 'static { + /// SIMD representation for partition ranking scores. + type Simd: PiPNNSIMDVector + where + A: PiPNNSIMDSchema; + /// Prepare one norm value for each point in the active stripe. fn prepare_point_norms(_points: MatrixView<'_, f32>, norms: &mut Vec) { norms.clear(); @@ -26,11 +32,11 @@ pub(in super::super) trait PartitionMetric: Send + Sync + 'static { /// Prepare one point norm for reuse across SIMD leader groups. #[inline(always)] - fn point_simd(arch: F::Arch, _norms: PartitionNorms<'_>, _point: usize) -> F + fn point_simd(arch: A, _norms: PartitionNorms<'_>, _point: usize) -> Self::Simd where - F: SIMDVector, + A: PiPNNSIMDSchema, { - F::default(arch) + Self::Simd::::default(arch) } /// Prepare one point norm for reuse across single leader values. @@ -40,16 +46,15 @@ pub(in super::super) trait PartitionMetric: Send + Sync + 'static { } /// Compute rankings for one complete SIMD group. - fn rankings_simd( - arch: F::Arch, + fn rankings_simd( + arch: A, norms: PartitionNorms<'_>, - point_norms: F, - dot_products: F, + point_norms: Self::Simd, + dot_products: Self::Simd, first_leader: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect; + A: PiPNNSIMDSchema; /// Compute one ranking outside the complete SIMD prefix. fn ranking_single( @@ -64,7 +69,7 @@ pub(in super::super) trait PartitionMetric: Send + Sync + 'static { #[inline(always)] fn load_norms_simd(arch: F::Arch, norms: &[f32], first_norm: usize) -> F where - F: SIMDVector, + F: PiPNNSIMDVector, { let last_norm = first_norm + F::LANES; let norm_group = &norms[first_norm..last_norm]; @@ -74,6 +79,11 @@ where } impl PartitionMetric for L2 { + type Simd + = A::PartitionScore + where + A: PiPNNSIMDSchema; + fn prepare_leader_norms(leaders: MatrixView<'_, f32>, norms: &mut Vec) { norms.resize(leaders.nrows(), 0.0); for (norm, leader) in norms.iter_mut().zip(leaders.row_iter()) { @@ -82,19 +92,18 @@ impl PartitionMetric for L2 { } #[inline(always)] - fn rankings_simd( - arch: F::Arch, + fn rankings_simd( + arch: A, norms: PartitionNorms<'_>, - _point_norms: F, - dot_products: F, + _point_norms: Self::Simd, + dot_products: Self::Simd, first_leader: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, { - let leader_norms = load_norms_simd::(arch, norms.leader_norms, first_leader); - F::splat(arch, -2.0).mul_add_simd(dot_products, leader_norms) + let leader_norms = load_norms_simd::>(arch, norms.leader_norms, first_leader); + Self::Simd::::splat(arch, -2.0).mul_add_simd(dot_products, leader_norms) } #[inline(always)] @@ -109,6 +118,11 @@ impl PartitionMetric for L2 { } impl PartitionMetric for Cosine { + type Simd + = A::PartitionScore + where + A: PiPNNSIMDSchema; + fn prepare_point_norms(points: MatrixView<'_, f32>, norms: &mut Vec) { norms.resize(points.nrows(), 0.0); for (norm, point) in norms.iter_mut().zip(points.row_iter()) { @@ -124,11 +138,11 @@ impl PartitionMetric for Cosine { } #[inline(always)] - fn point_simd(arch: F::Arch, norms: PartitionNorms<'_>, point: usize) -> F + fn point_simd(arch: A, norms: PartitionNorms<'_>, point: usize) -> Self::Simd where - F: SIMDVector, + A: PiPNNSIMDSchema, { - F::splat(arch, norms.point_norms[point]) + Self::Simd::::splat(arch, norms.point_norms[point]) } #[inline(always)] @@ -137,18 +151,17 @@ impl PartitionMetric for Cosine { } #[inline(always)] - fn rankings_simd( - arch: F::Arch, + fn rankings_simd( + arch: A, norms: PartitionNorms<'_>, - point_norms: F, - dot_products: F, + point_norms: Self::Simd, + dot_products: Self::Simd, first_leader: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, { - let leader_norms = load_norms_simd::(arch, norms.leader_norms, first_leader); + let leader_norms = load_norms_simd::>(arch, norms.leader_norms, first_leader); cosine_distance_simd(arch, dot_products, point_norms, leader_norms) } @@ -164,19 +177,23 @@ impl PartitionMetric for Cosine { } impl PartitionMetric for CosineNormalized { + type Simd + = A::PartitionScore + where + A: PiPNNSIMDSchema; + #[inline(always)] - fn rankings_simd( - arch: F::Arch, + fn rankings_simd( + arch: A, _norms: PartitionNorms<'_>, - _point_norms: F, - dot_products: F, + _point_norms: Self::Simd, + dot_products: Self::Simd, _first_leader: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, { - F::splat(arch, 1.0) - dot_products + Self::Simd::::splat(arch, 1.0) - dot_products } #[inline(always)] @@ -191,19 +208,23 @@ impl PartitionMetric for CosineNormalized { } impl PartitionMetric for InnerProduct { + type Simd + = A::PartitionScore + where + A: PiPNNSIMDSchema; + #[inline(always)] - fn rankings_simd( - arch: F::Arch, + fn rankings_simd( + arch: A, _norms: PartitionNorms<'_>, - _point_norms: F, - dot_products: F, + _point_norms: Self::Simd, + dot_products: Self::Simd, _first_leader: usize, - ) -> F + ) -> Self::Simd where - F: SIMDVector + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, { - F::default(arch) - dot_products + Self::Simd::::default(arch) - dot_products } #[inline(always)] diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 20a707317b..39660ae441 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -23,9 +23,12 @@ use crate::{ANNError, ANNResult}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; +use diskann_wide::{SIMDPartialOrd, SIMDVector}; -use super::kernel_metric::LeafMetric; +use super::{ + kernel_metric::LeafMetric, + simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, +}; /// One leaf-local neighbor and its metric distance. #[derive(Clone, Copy, Debug, PartialEq)] @@ -98,11 +101,8 @@ pub(super) fn select_leaf_neighbors( workspace: &mut LeafKernelWorkspace, ) -> ANNResult<()> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, + A: PiPNNSIMDSchema, M: LeafMetric, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { let point_count = points.nrows(); let dot_count = point_count * point_count; @@ -136,11 +136,8 @@ fn rank_leaf_dots( worst: &mut Vec, ) -> Result<(), LeafKernelError> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, + A: PiPNNSIMDSchema, M: LeafMetric, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, { validate_neighbor_count(input, &output)?; let neighbor_count = output.ncols(); @@ -153,10 +150,10 @@ where worst.fill(f32::INFINITY); match neighbor_count { - 1 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), - 2 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), - 3 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), - _ => scan_runtime_width::( + 1 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), + 2 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), + 3 => scan_fixed_width::(arch, input, norms, output.as_mut_slice(), worst), + _ => scan_runtime_width::( arch, input, norms, @@ -189,39 +186,35 @@ fn validate_neighbor_count( } /// Select neighbors with a fixed output width. -fn scan_fixed_width( - arch: F::Arch, +fn scan_fixed_width( + arch: A, input: MatrixView<'_, f32>, norms: &[f32], output: &mut [LeafNeighbor], worst: &mut [f32], ) where - F: SIMDVector> + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, M: LeafMetric, - u64: From<<::BitMask as SIMDMask>::Underlying>, { let (rows, _) = output.as_chunks_mut::(); - scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { + scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { insert_fixed_neighbor(&mut rows[source], target, distance) }); } /// Select neighbors with a runtime output width. -fn scan_runtime_width( - arch: F::Arch, +fn scan_runtime_width( + arch: A, input: MatrixView<'_, f32>, norms: &[f32], output: &mut [LeafNeighbor], width: usize, worst: &mut [f32], ) where - F: SIMDVector> + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, M: LeafMetric, - u64: From<<::BitMask as SIMDMask>::Underlying>, { - scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { + scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { let first = source * width; insert_runtime_neighbor(&mut output[first..first + width], target, distance) }); @@ -232,18 +225,16 @@ fn scan_runtime_width( /// The function reads the strict lower triangle once. It offers each distance to /// both endpoint lists. SIMD groups and single values preserve pair scan order. #[inline(never)] -fn scan_point_pairs( - arch: F::Arch, +fn scan_point_pairs( + arch: A, input: MatrixView<'_, f32>, norms: &[f32], worst: &mut [f32], mut insert: I, ) where - F: SIMDVector> + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, M: LeafMetric, I: FnMut(usize, u32, f32) -> f32, - u64: From<<::BitMask as SIMDMask>::Underlying>, { let point_count = input.nrows(); let dots = input.as_slice(); @@ -251,27 +242,28 @@ fn scan_point_pairs( for source in 1..point_count { let source_start = source * point_count; - let source_simd = M::source_simd::(arch, norms, source); + let source_simd = M::source_simd(arch, norms, source); let source_single = M::source_single(norms, source); // SAFETY: `rank_leaf_dots` created one threshold for each point. let mut source_worst = unsafe { *worst_ptr.add(source) }; let mut target = 0; - let simd_prefix = source - source % F::LANES; + let simd_prefix = source - source % M::Simd::::LANES; while target < simd_prefix { // SAFETY: This complete SIMD group is in the strict-lower prefix. let dot_products = - unsafe { F::load_simd(arch, dots.as_ptr().add(source_start + target)) }; - let distances = M::distances_simd::(arch, norms, source_simd, dot_products, target); - let source_eligible = distances.lt_simd(F::splat(arch, source_worst)); + unsafe { M::Simd::::load_simd(arch, dots.as_ptr().add(source_start + target)) }; + let distances = M::distances_simd(arch, norms, source_simd, dot_products, target); + let source_eligible = distances.lt_simd(M::Simd::::splat(arch, source_worst)); // SAFETY: The complete target group is below `source < point_count`. - let target_worst = unsafe { F::load_simd(arch, worst_ptr.add(target)) }; + let target_worst = unsafe { M::Simd::::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()); + let source_bits = M::Simd::::active_lanes(source_eligible); + let target_bits = M::Simd::::active_lanes(target_eligible); if source_bits | target_bits != 0 { - let values: [f32; 16] = distances.to_array(); + let values = distances.to_array(); + let values = values.as_ref(); let mut source_bits = source_bits; while source_bits != 0 { let lane = source_bits.trailing_zeros() as usize; @@ -292,7 +284,7 @@ fn scan_point_pairs( unsafe { *worst_ptr.add(target_source) = new_worst }; } } - target += F::LANES; + target += M::Simd::::LANES; } while target < source { @@ -382,10 +374,7 @@ struct DispatchLeafForTest(diskann_vector::distance::Metric); impl diskann_wide::arch::Target1, DispatchedLeafCall<'_>> for DispatchLeafForTest where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, { fn run(self, arch: A, call: DispatchedLeafCall<'_>) -> Result<(), LeafKernelError> { use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 676c00415f..092ee7549a 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -24,6 +24,8 @@ //! output and workspace. #[allow(dead_code)] mod kernel_metric; +#[allow(dead_code)] +mod simd; #[allow(dead_code)] mod leaf_kernel; diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 77fe3c3d23..58580e53ef 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -20,11 +20,12 @@ use diskann_linalg::Transpose; use diskann_utils::views::{MatrixView, MutMatrixView}; #[cfg(test)] use diskann_vector::distance::Metric; -use diskann_wide::{ - Architecture, Const, SIMDFloat, SIMDMask, SIMDPartialOrd, SIMDSelect, SIMDVector, -}; +use diskann_wide::{SIMDMask, SIMDVector}; -use super::kernel_metric::{PartitionMetric, PartitionNorms}; +use super::{ + kernel_metric::{PartitionMetric, PartitionNorms}, + simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, +}; /// No sampled partition center was rankable for this output slot. pub(super) const UNASSIGNED_LEADER: u32 = u32::MAX; @@ -90,10 +91,7 @@ pub(super) fn assign_leaders( workspace: &mut PartitionKernelWorkspace, ) -> ANNResult<()> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, { let point_count = points.nrows(); @@ -145,10 +143,7 @@ fn rank_leader_dots( output: MutMatrixView<'_, u32>, ranked_leaders: &mut Vec<(u32, f32)>, ) where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, { let fanout = output.ncols(); @@ -157,24 +152,22 @@ fn rank_leader_dots( } ranked_leaders.resize(fanout, (UNASSIGNED_LEADER, f32::INFINITY)); - select_point_leaders::(arch, input.dots, input.norms, output, ranked_leaders); + select_point_leaders::(arch, input.dots, input.norms, output, ranked_leaders); } /// Rank sampled partition centers for each assigned point. /// /// The function keeps nearest-first order for every point. Full SIMD groups use /// metric-specific formulas. Remaining leaders use the matching single formula. -fn select_point_leaders( - arch: F::Arch, +fn select_point_leaders( + arch: A, dots: MatrixView<'_, f32>, norms: PartitionNorms<'_>, mut output: MutMatrixView<'_, u32>, ranked_leaders: &mut [(u32, f32)], ) where - F: SIMDVector> + SIMDFloat + std::ops::Div, - F::Mask: SIMDSelect, + A: PiPNNSIMDSchema, M: PartitionMetric, - u64: From<<::BitMask as SIMDMask>::Underlying>, { let leader_count = dots.ncols(); let fanout = output.ncols(); @@ -185,15 +178,15 @@ fn select_point_leaders( .enumerate() { ranked_leaders.fill((UNASSIGNED_LEADER, f32::INFINITY)); - let point_simd = M::point_simd::(arch, norms, point); + let point_simd = M::point_simd(arch, norms, point); let point_single = M::point_single(norms, point); - let simd_prefix = leader_count - leader_count % F::LANES; + let simd_prefix = leader_count - leader_count % M::Simd::::LANES; - for first_leader in (0..simd_prefix).step_by(F::LANES) { + for first_leader in (0..simd_prefix).step_by(M::Simd::::LANES) { // SAFETY: This group is inside the point's leader row. - let dot_products = unsafe { F::load_simd(arch, point_dots.as_ptr().add(first_leader)) }; - let rankings = - M::rankings_simd::(arch, norms, point_simd, dot_products, first_leader); + let dot_products = + unsafe { M::Simd::::load_simd(arch, point_dots.as_ptr().add(first_leader)) }; + let rankings = M::rankings_simd(arch, norms, point_simd, dot_products, first_leader); insert_leader_lanes(rankings, first_leader, ranked_leaders); } @@ -213,8 +206,7 @@ fn select_point_leaders( /// sampled-leader order, which preserves tie order. fn insert_leader_lanes(scores: F, first_leader: usize, ranked_leaders: &mut [(u32, f32)]) where - F: SIMDVector> + SIMDPartialOrd, - u64: From<<::BitMask as SIMDMask>::Underlying>, + F: PiPNNSIMDVector, { let threshold = F::splat(scores.arch(), ranked_leaders[ranked_leaders.len() - 1].1); let eligible = scores.lt_simd(threshold); @@ -222,8 +214,9 @@ where return; } - let values: [f32; 16] = scores.to_array(); - let mut lanes = u64::from(eligible.bitmask().to_underlying()); + let values = scores.to_array(); + let values = values.as_ref(); + let mut lanes = F::active_lanes(eligible); while lanes != 0 { let lane = lanes.trailing_zeros() as usize; lanes &= lanes - 1; @@ -264,10 +257,7 @@ struct DispatchPartitionForTest(Metric); #[cfg(test)] impl diskann_wide::arch::Target1> for DispatchPartitionForTest where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, { fn run(self, arch: A, call: DispatchedPartitionCall<'_>) { use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs new file mode 100644 index 0000000000..42f6efe057 --- /dev/null +++ b/diskann/src/graph/pipnn/simd.rs @@ -0,0 +1,61 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! SIMD schema for PiPNN numerical kernels. + +use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; + +/// Default SIMD representation used by every PiPNN numerical stage. +/// +/// This alias is the single build-time width selection. +type DefaultVector = ::f32x16; + +/// Operations required by PiPNN SIMD vectors. +pub(super) trait PiPNNSIMDVector: + SIMDVector + SIMDFloat + std::ops::Div +{ + /// Return one bit for each selected lane. + fn active_lanes(mask: Self::Mask) -> u64; + + /// Select one value from each pair of lanes. + fn select(mask: Self::Mask, if_true: Self, if_false: Self) -> Self; +} + +impl PiPNNSIMDVector for F +where + F: SIMDVector + SIMDFloat + std::ops::Div, + F::Mask: SIMDSelect, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + #[inline(always)] + fn active_lanes(mask: Self::Mask) -> u64 { + u64::from(mask.bitmask().to_underlying()) + } + + #[inline(always)] + fn select(mask: Self::Mask, if_true: Self, if_false: Self) -> Self { + mask.select(if_true, if_false) + } +} + +/// Stage-specific SIMD representations for one architecture. +pub(super) trait PiPNNSIMDSchema: Architecture { + /// SIMD representation for leaf distance scores. + type LeafScore: PiPNNSIMDVector; + /// SIMD representation for partition ranking scores. + type PartitionScore: PiPNNSIMDVector; + /// SIMD representation for relative-hash sketch comparisons. + type HashScore: PiPNNSIMDVector; +} + +impl PiPNNSIMDSchema for A +where + A: Architecture, + DefaultVector: PiPNNSIMDVector, +{ + type LeafScore = DefaultVector; + type PartitionScore = DefaultVector; + type HashScore = DefaultVector; +} From 286f55837223c8525b5a3653982d2bfce36e6eb4 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 21 Aug 2026 08:22:02 +0000 Subject: [PATCH 79/80] refactor(pipnn): specialize sorted insertion --- diskann/src/graph/pipnn/leaf_kernel.rs | 154 ++++++++++++++++--------- 1 file changed, 102 insertions(+), 52 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 39660ae441..a6060d773f 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -195,10 +195,11 @@ fn scan_fixed_width( ) where A: PiPNNSIMDSchema, M: LeafMetric, + [LeafNeighbor; N]: SortedInsert, { let (rows, _) = output.as_chunks_mut::(); scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { - insert_fixed_neighbor(&mut rows[source], target, distance) + insert_neighbor(&mut rows[source], target, distance) }); } @@ -216,7 +217,7 @@ fn scan_runtime_width( { scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { let first = source * width; - insert_runtime_neighbor(&mut output[first..first + width], target, distance) + insert_neighbor(&mut output[first..first + width], target, distance) }); } @@ -308,55 +309,79 @@ fn scan_point_pairs( } } -/// Insert one target into a fixed-width retained neighbor set. -#[inline(always)] -fn insert_fixed_neighbor( - neighbors: &mut [LeafNeighbor; N], - target: u32, - distance: f32, -) -> f32 { - let entry = LeafNeighbor::new(target, distance); - if N == 1 { - neighbors[0] = entry; - return distance; +/// Insert one value that passed the retained set threshold. +trait SortedInsert { + fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T; +} + +impl SortedInsert for [T; 1] { + #[inline(always)] + fn insert_sorted_by(&mut self, value: T, _precedes: impl Fn(T, T) -> bool) -> T { + self[0] = value; + value + } +} + +impl SortedInsert for [T; 2] { + #[inline(always)] + fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + let first = self[0]; + if precedes(value, first) { + self[0] = value; + self[1] = first; + first + } else { + self[1] = value; + value + } } - if N == 2 { - let first = neighbors[0]; - if distance < first.distance { - neighbors[0] = entry; - neighbors[1] = first; - return first.distance; +} + +impl SortedInsert for [T; 3] { + #[inline(always)] + fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + let (first, second) = (self[0], self[1]); + if precedes(value, first) { + self[0] = value; + self[1] = first; + self[2] = second; + second + } else if precedes(value, second) { + self[1] = value; + self[2] = second; + second + } else { + self[2] = value; + value } - neighbors[1] = entry; - return distance; } +} - let (first, second) = (neighbors[0], neighbors[1]); - if distance < first.distance { - neighbors[0] = entry; - neighbors[1] = first; - neighbors[2] = second; - } else if distance < second.distance { - neighbors[1] = entry; - neighbors[2] = second; - } else { - neighbors[2] = entry; - return distance; +impl SortedInsert for [T] { + #[inline(always)] + fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + let last = self.len() - 1; + let mut slot = last; + while slot > 0 && precedes(value, self[slot - 1]) { + self[slot] = self[slot - 1]; + slot -= 1; + } + self[slot] = value; + self[last] } - second.distance } -/// Insert one target into a runtime-width retained neighbor set. #[inline(always)] -fn insert_runtime_neighbor(neighbors: &mut [LeafNeighbor], target: u32, distance: f32) -> f32 { - let last = neighbors.len() - 1; - let mut slot = last; - while slot > 0 && distance < neighbors[slot - 1].distance { - neighbors[slot] = neighbors[slot - 1]; - slot -= 1; - } - neighbors[slot] = LeafNeighbor::new(target, distance); - neighbors[last].distance +fn insert_neighbor(neighbors: &mut R, target: u32, distance: f32) -> f32 +where + R: SortedInsert + ?Sized, +{ + neighbors + .insert_sorted_by( + LeafNeighbor::new(target, distance), + |candidate, retained| candidate.distance < retained.distance, + ) + .distance } #[cfg(test)] @@ -480,25 +505,50 @@ mod tests { } #[test] - fn fixed_insertion_orders_candidates() { - let mut output = [LeafNeighbor::default(); 3]; - let mut worst = f32::INFINITY; - - for (target, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 0.5)] { - if distance < worst { - worst = insert_fixed_neighbor(&mut output, target, distance); + fn sorted_insert_specializations_order_candidates() { + fn retain(neighbors: &mut R) -> f32 + where + R: SortedInsert + ?Sized, + { + let mut worst = f32::INFINITY; + for (target, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 0.5)] { + if distance < worst { + worst = insert_neighbor(neighbors, target, distance); + } } + worst } + let mut one = [LeafNeighbor::default(); 1]; + assert_eq!(retain(&mut one), 0.5); + assert_eq!(one, [LeafNeighbor::new(4, 0.5)]); + + let mut two = [LeafNeighbor::default(); 2]; + assert_eq!(retain(&mut two), 1.0); + assert_eq!(two, [LeafNeighbor::new(4, 0.5), LeafNeighbor::new(1, 1.0)]); + + let mut three = [LeafNeighbor::default(); 3]; + assert_eq!(retain(&mut three), 2.0); assert_eq!( - output, + three, + [ + LeafNeighbor::new(4, 0.5), + LeafNeighbor::new(1, 1.0), + LeafNeighbor::new(3, 2.0), + ] + ); + + let mut runtime = [LeafNeighbor::default(); 4]; + assert_eq!(retain(runtime.as_mut_slice()), 3.0); + assert_eq!( + runtime, [ LeafNeighbor::new(4, 0.5), LeafNeighbor::new(1, 1.0), LeafNeighbor::new(3, 2.0), + LeafNeighbor::new(2, 3.0), ] ); - assert_eq!(worst, 2.0); } #[test] From 1fb768baafcbc196d19ce151b6db6c39d8313754 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:06:11 +0000 Subject: [PATCH 80/80] test(pipnn): clarify numerical contracts Use named cases and derived expected values so each failure identifies a ranking behavior that can change graph output. --- Cargo.lock | 1 + diskann/Cargo.toml | 1 + diskann/src/graph/pipnn/kernel_metric.rs | 82 + diskann/src/graph/pipnn/kernel_metric/leaf.rs | 133 ++ .../graph/pipnn/kernel_metric/partition.rs | 161 ++ diskann/src/graph/pipnn/leaf_kernel.rs | 1394 ++++++++++------- diskann/src/graph/pipnn/partition_kernel.rs | 767 ++++----- 7 files changed, 1506 insertions(+), 1033 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index da640e005a..40240e77c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -449,6 +449,7 @@ dependencies = [ "pin-project", "rand", "relative-path 2.0.1", + "rstest", "serde", "serde_json", "thiserror 2.0.17", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 7cd2a0df09..c0016416dc 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -37,6 +37,7 @@ futures-util = { workspace = true, default-features = false } pin-project.workspace = true rand.workspace = true relative-path = "2.0.1" +rstest.workspace = true serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "sync"] } diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index be6a287e8e..91f40fd69b 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -60,3 +60,85 @@ pub(super) fn cosine_distance_single(dot: f32, source_norm: f32, target_norm: f3 1.0 - (-1.0_f32).max(1.0_f32.min(cosine)) } } + +#[cfg(test)] +mod tests { + use super::cosine_distance_single; + + mod cosine_distance_single_tests { + use super::cosine_distance_single; + + #[test] + fn cosine_zero_source_norm_produces_unit_distance() { + // Given + let source_norm = 0.0; + let zero_norm_similarity = 0.0; + let expected_one_minus_zero_similarity = 1.0 - zero_norm_similarity; + + // When + let actual_distance = cosine_distance_single(100.0, source_norm, 2.0); + + // Then + assert_eq!(actual_distance, expected_one_minus_zero_similarity); + } + + #[test] + fn cosine_zero_target_norm_produces_unit_distance() { + // Given + let target_norm = 0.0; + let zero_norm_similarity = 0.0; + let expected_one_minus_zero_similarity = 1.0 - zero_norm_similarity; + + // When + let actual_distance = cosine_distance_single(-100.0, 2.0, target_norm); + + // Then + assert_eq!(actual_distance, expected_one_minus_zero_similarity); + } + + #[test] + fn cosine_similarity_above_one_clamps_to_zero_distance() { + // Given + let dot_product_above_valid_similarity = 5.0; + let maximum_cosine_similarity = 1.0; + let expected_one_minus_maximum_similarity = 1.0 - maximum_cosine_similarity; + + // When + let actual_distance = + cosine_distance_single(dot_product_above_valid_similarity, 2.0, 2.0); + + // Then + assert_eq!(actual_distance, expected_one_minus_maximum_similarity); + } + + #[test] + fn cosine_similarity_below_negative_one_clamps_to_distance_two() { + // Given + let dot_product_below_valid_similarity = -5.0; + let minimum_cosine_similarity = -1.0; + let expected_one_minus_minimum_similarity = 1.0 - minimum_cosine_similarity; + + // When + let actual_distance = + cosine_distance_single(dot_product_below_valid_similarity, 2.0, 2.0); + + // Then + assert_eq!(actual_distance, expected_one_minus_minimum_similarity); + } + + #[test] + fn cosine_nan_similarity_follows_the_diskann_min_max_rule() { + // Given + let nan_dot_product = f32::NAN; + // `f32::min` keeps its finite operand when the other operand is NaN. + let finite_operand_selected_by_min = 1.0; + let expected_one_minus_selected_operand = 1.0 - finite_operand_selected_by_min; + + // When + let actual_distance = cosine_distance_single(nan_dot_product, 1.0, 1.0); + + // Then + assert_eq!(actual_distance, expected_one_minus_selected_operand); + } + } +} diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 2a558baac5..844e84b401 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -213,3 +213,136 @@ impl LeafMetric for InnerProduct { -dot_product } } + +#[cfg(test)] +#[allow( + clippy::unwrap_used, + reason = "deterministic test matrices must abort on invalid setup" +)] +mod tests { + use super::*; + + mod prepare_leaf_norms_tests { + use super::*; + + #[test] + fn l2_leaf_norms_equal_the_gram_diagonal() { + // Given + let first_squared_norm = 4.0_f32; + let second_squared_norm = 9.0_f32; + let lower_gram_values = [first_squared_norm, 0.0, 0.0, second_squared_norm]; + let lower_gram = MatrixView::try_from(&lower_gram_values[..], 2, 2).unwrap(); + let expected_gram_diagonal = [first_squared_norm, second_squared_norm]; + let mut actual_norms = Vec::new(); + + // When + L2::prepare_leaf_norms(lower_gram, &mut actual_norms); + + // Then + assert_eq!(actual_norms, expected_gram_diagonal); + } + + #[test] + fn cosine_leaf_norms_equal_square_roots_of_the_gram_diagonal() { + // Given + let first_squared_norm = 4.0_f32; + let second_squared_norm = 9.0_f32; + let lower_gram_values = [first_squared_norm, 0.0, 0.0, second_squared_norm]; + let lower_gram = MatrixView::try_from(&lower_gram_values[..], 2, 2).unwrap(); + let expected_square_roots_of_diagonal = + [first_squared_norm.sqrt(), second_squared_norm.sqrt()]; + let mut actual_norms = Vec::new(); + + // When + Cosine::prepare_leaf_norms(lower_gram, &mut actual_norms); + + // Then + assert_eq!(actual_norms, expected_square_roots_of_diagonal); + } + } + + mod distance_single_tests { + use super::*; + + #[test] + fn l2_distance_equals_squared_norm_sum_minus_twice_the_dot_product() { + // Given + let source_squared_norm = 4.0; + let target_squared_norm = 9.0; + let dot_product = 6.0; + let squared_norms = [source_squared_norm, target_squared_norm]; + let expected_squared_l2_distance = + source_squared_norm + target_squared_norm - 2.0 * dot_product; + + // When + let actual_distance = + L2::distance_single(&squared_norms, source_squared_norm, dot_product, 1); + + // Then + assert_eq!(actual_distance, expected_squared_l2_distance); + } + + #[test] + fn l2_clamps_negative_roundoff_to_zero() { + // Given + let source_squared_norm = 1.0; + let target_squared_norm = 1.0; + let dot_product_above_exact_norm = 1.000_001; + let squared_norms = [source_squared_norm, target_squared_norm]; + let expected_non_negative_distance = 0.0; + + // When + let actual_distance = L2::distance_single( + &squared_norms, + source_squared_norm, + dot_product_above_exact_norm, + 1, + ); + + // Then + assert_eq!(actual_distance, expected_non_negative_distance); + } + + #[test] + fn cosine_distance_equals_one_minus_dot_over_norm_product() { + // Given + let source_norm = 2.0; + let target_norm = 4.0; + let dot_product = 4.0; + let norms = [source_norm, target_norm]; + let expected_one_minus_normalized_dot = 1.0 - dot_product / (source_norm * target_norm); + + // When + let actual_distance = Cosine::distance_single(&norms, source_norm, dot_product, 1); + + // Then + assert_eq!(actual_distance, expected_one_minus_normalized_dot); + } + + #[test] + fn normalized_cosine_distance_is_one_minus_the_dot_product() { + // Given + let dot_product = 0.25; + let expected_one_minus_dot = 1.0 - dot_product; + + // When + let actual_distance = CosineNormalized::distance_single(&[], 0.0, dot_product, 0); + + // Then + assert_eq!(actual_distance, expected_one_minus_dot); + } + + #[test] + fn inner_product_distance_is_the_negative_dot_product() { + // Given + let dot_product = 3.0; + let expected_negative_dot = -dot_product; + + // When + let actual_distance = InnerProduct::distance_single(&[], 0.0, dot_product, 0); + + // Then + assert_eq!(actual_distance, expected_negative_dot); + } + } +} diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 0490b46692..05b54fe9b3 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -237,3 +237,164 @@ impl PartitionMetric for InnerProduct { -dot_product } } + +#[cfg(test)] +#[allow( + clippy::unwrap_used, + reason = "deterministic test matrices must abort on invalid setup" +)] +mod tests { + use super::*; + + fn rank_single_leader( + dot_product: f32, + point_norm: f32, + leader_norm: f32, + ) -> f32 { + let point_norms = [point_norm]; + let leader_norms = [leader_norm]; + let norms = PartitionNorms { + point_norms: &point_norms, + leader_norms: &leader_norms, + }; + M::ranking_single(norms, M::point_single(norms, 0), dot_product, 0) + } + + mod prepare_leader_norms_tests { + use super::*; + + #[test] + fn l2_leader_norm_is_the_sum_of_squared_components() { + // Given + let first_leader = [1.0_f32, 2.0]; + let second_leader = [3.0_f32, 4.0]; + let leader_values = [ + first_leader[0], + first_leader[1], + second_leader[0], + second_leader[1], + ]; + let leaders = MatrixView::try_from(&leader_values[..], 2, 2).unwrap(); + let expected_row_squared_norms = [ + first_leader[0].powi(2) + first_leader[1].powi(2), + second_leader[0].powi(2) + second_leader[1].powi(2), + ]; + let mut actual_norms = Vec::new(); + + // When + L2::prepare_leader_norms(leaders, &mut actual_norms); + + // Then + assert_eq!(actual_norms, expected_row_squared_norms); + } + + #[test] + fn cosine_leader_norm_is_the_square_root_of_the_squared_component_sum() { + // Given + let first_leader = [1.0_f32, 2.0]; + let second_leader = [3.0_f32, 4.0]; + let leader_values = [ + first_leader[0], + first_leader[1], + second_leader[0], + second_leader[1], + ]; + let leaders = MatrixView::try_from(&leader_values[..], 2, 2).unwrap(); + let expected_row_norms = [ + (first_leader[0].powi(2) + first_leader[1].powi(2)).sqrt(), + (second_leader[0].powi(2) + second_leader[1].powi(2)).sqrt(), + ]; + let mut actual_norms = Vec::new(); + + // When + Cosine::prepare_leader_norms(leaders, &mut actual_norms); + + // Then + assert_eq!(actual_norms, expected_row_norms); + } + } + + #[test] + fn cosine_point_norm_is_the_square_root_of_the_squared_component_sum() { + // Given + let first_point = [1.0_f32, 2.0]; + let second_point = [3.0_f32, 4.0]; + let point_values = [ + first_point[0], + first_point[1], + second_point[0], + second_point[1], + ]; + let points = MatrixView::try_from(&point_values[..], 2, 2).unwrap(); + let expected_row_norms = [ + (first_point[0].powi(2) + first_point[1].powi(2)).sqrt(), + (second_point[0].powi(2) + second_point[1].powi(2)).sqrt(), + ]; + let mut actual_norms = Vec::new(); + + // When + Cosine::prepare_point_norms(points, &mut actual_norms); + + // Then + assert_eq!(actual_norms, expected_row_norms); + } + + mod ranking_single_tests { + use super::*; + + #[test] + fn l2_ranking_equals_leader_squared_norm_minus_twice_the_dot_product() { + // Given + let dot_product = 2.0; + let leader_squared_norm = 9.0; + let expected_leader_norm_minus_twice_dot = leader_squared_norm - 2.0 * dot_product; + + // When + let actual_ranking = rank_single_leader::(dot_product, 0.0, leader_squared_norm); + + // Then + assert_eq!(actual_ranking, expected_leader_norm_minus_twice_dot); + } + + #[test] + fn cosine_ranking_equals_one_minus_dot_over_norm_product() { + // Given + let dot_product = 4.0; + let point_norm = 2.0; + let leader_norm = 4.0; + let expected_one_minus_normalized_dot = 1.0 - dot_product / (point_norm * leader_norm); + + // When + let actual_ranking = rank_single_leader::(dot_product, point_norm, leader_norm); + + // Then + assert_eq!(actual_ranking, expected_one_minus_normalized_dot); + } + + #[test] + fn normalized_cosine_ranking_is_one_minus_the_dot_product() { + // Given + let dot_product = 0.25; + let expected_one_minus_dot = 1.0 - dot_product; + + // When + let actual_ranking = rank_single_leader::(dot_product, 0.0, 0.0); + + // Then + assert_eq!(actual_ranking, expected_one_minus_dot); + } + + #[test] + fn inner_product_ranking_is_the_negative_dot_product() { + // Given + let dot_product = 3.0; + let expected_negative_dot = -dot_product; + + // When + let actual_ranking = rank_single_leader::(dot_product, 0.0, 0.0); + + // Then + assert_eq!(actual_ranking, expected_negative_dot); + } + } +} diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index a6060d773f..49a4426631 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -309,7 +309,10 @@ fn scan_point_pairs( } } -/// Insert one value that passed the retained set threshold. +/// Insert one value that precedes the current last retained value. +/// +/// The caller checks eligibility before insertion. The returned value is the new +/// last retained value for the next eligibility check. trait SortedInsert { fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T; } @@ -371,6 +374,9 @@ impl SortedInsert for [T] { } } +/// Insert one candidate that is nearer than the current farthest neighbor. +/// +/// Return the new farthest retained distance for the next candidate check. #[inline(always)] fn insert_neighbor(neighbors: &mut R, target: u32, distance: f32) -> f32 where @@ -385,691 +391,893 @@ where } #[cfg(test)] -struct DispatchedLeafCall<'a> { - input: MatrixView<'a, f32>, - norms: &'a [f32], - output: MutMatrixView<'a, LeafNeighbor>, - workspace: &'a mut LeafKernelWorkspace, -} - -#[cfg(test)] -struct DispatchLeafForTest(diskann_vector::distance::Metric); - -#[cfg(test)] -impl diskann_wide::arch::Target1, DispatchedLeafCall<'_>> - for DispatchLeafForTest -where - A: PiPNNSIMDSchema, -{ - fn run(self, arch: A, call: DispatchedLeafCall<'_>) -> Result<(), LeafKernelError> { - use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; - use diskann_vector::distance::Metric; - - match self.0 { - Metric::L2 => rank_leaf_dots::( - arch, - call.input, - call.norms, - call.output, - &mut call.workspace.worst, - ), - Metric::Cosine => rank_leaf_dots::( - arch, - call.input, - call.norms, - call.output, - &mut call.workspace.worst, - ), - Metric::CosineNormalized => rank_leaf_dots::( - arch, - call.input, - call.norms, - call.output, - &mut call.workspace.worst, - ), - Metric::InnerProduct => rank_leaf_dots::( - arch, - call.input, - call.norms, - call.output, - &mut call.workspace.worst, - ), - } - } -} - -#[cfg(test)] -fn dispatch_nearest_neighbors( - metric: diskann_vector::distance::Metric, - input: MatrixView<'_, f32>, - norms: &[f32], - output: MutMatrixView<'_, LeafNeighbor>, - workspace: &mut LeafKernelWorkspace, -) -> Result<(), LeafKernelError> { - diskann_wide::arch::dispatch1_no_features( - DispatchLeafForTest(metric), - DispatchedLeafCall { - input, - norms, - output, - workspace, - }, - ) -} +mod tests { + use std::cmp::Ordering; -#[cfg(test)] -fn prepared_test_norms( - metric: diskann_vector::distance::Metric, - input: MatrixView<'_, f32>, -) -> Vec { - use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use super::*; + use crate::graph::pipnn::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; + use diskann_wide::arch::{self, Target1}; - fn prepare(input: MatrixView<'_, f32>) -> Vec { - let mut norms = Vec::new(); - M::prepare_leaf_norms(input, &mut norms); - norms - } - - match metric { - Metric::L2 => prepare::(input), - Metric::Cosine => prepare::(input), - Metric::CosineNormalized => prepare::(input), - Metric::InnerProduct => prepare::(input), + struct KernelCall<'a> { + input: MatrixView<'a, f32>, + norms: &'a [f32], + output: MutMatrixView<'a, LeafNeighbor>, + workspace: &'a mut LeafKernelWorkspace, } -} -#[cfg(test)] -mod tests { - use super::*; - use diskann_vector::distance::Metric; - - fn test_dots(metric: Metric, points: usize) -> Vec { - let mut dots = vec![f32::NAN; points * points]; - for source in 0..points { - dots[source * points + source] = if metric == Metric::Cosine && source == 0 { - 0.0 - } else { - 1.0 + (source % 5) as f32 - }; - for target in 0..source { - dots[source * points + target] = - (((source * 17 + target * 11) % 23) as f32 - 11.0) * 0.03125; + struct DispatchMetric(Metric); + + impl Target1, KernelCall<'_>> for DispatchMetric + where + A: PiPNNSIMDSchema, + { + fn run(self, arch: A, call: KernelCall<'_>) -> Result<(), LeafKernelError> { + match self.0 { + Metric::L2 => rank_leaf_dots::( + arch, + call.input, + call.norms, + call.output, + &mut call.workspace.worst, + ), + Metric::Cosine => rank_leaf_dots::( + arch, + call.input, + call.norms, + call.output, + &mut call.workspace.worst, + ), + Metric::CosineNormalized => rank_leaf_dots::( + arch, + call.input, + call.norms, + call.output, + &mut call.workspace.worst, + ), + Metric::InnerProduct => rank_leaf_dots::( + arch, + call.input, + call.norms, + call.output, + &mut call.workspace.worst, + ), } } - dots } - fn test_input(dots: &[f32], points: usize) -> MatrixView<'_, f32> { + fn lower_gram_view(dots: &[f32], points: usize) -> MatrixView<'_, f32> { MatrixView::try_from(dots, points, points).unwrap() } - #[test] - fn sorted_insert_specializations_order_candidates() { - fn retain(neighbors: &mut R) -> f32 - where - R: SortedInsert + ?Sized, - { - let mut worst = f32::INFINITY; - for (target, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 0.5)] { - if distance < worst { - worst = insert_neighbor(neighbors, target, distance); - } - } - worst + fn metric_norms(metric: Metric, lower_gram: MatrixView<'_, f32>) -> Vec { + fn prepare(lower_gram: MatrixView<'_, f32>) -> Vec { + let mut norms = Vec::new(); + M::prepare_leaf_norms(lower_gram, &mut norms); + norms } - let mut one = [LeafNeighbor::default(); 1]; - assert_eq!(retain(&mut one), 0.5); - assert_eq!(one, [LeafNeighbor::new(4, 0.5)]); - - let mut two = [LeafNeighbor::default(); 2]; - assert_eq!(retain(&mut two), 1.0); - assert_eq!(two, [LeafNeighbor::new(4, 0.5), LeafNeighbor::new(1, 1.0)]); - - let mut three = [LeafNeighbor::default(); 3]; - assert_eq!(retain(&mut three), 2.0); - assert_eq!( - three, - [ - LeafNeighbor::new(4, 0.5), - LeafNeighbor::new(1, 1.0), - LeafNeighbor::new(3, 2.0), - ] - ); - - let mut runtime = [LeafNeighbor::default(); 4]; - assert_eq!(retain(runtime.as_mut_slice()), 3.0); - assert_eq!( - runtime, - [ - LeafNeighbor::new(4, 0.5), - LeafNeighbor::new(1, 1.0), - LeafNeighbor::new(3, 2.0), - LeafNeighbor::new(2, 3.0), - ] - ); - } - - #[test] - fn neighbor_count_clamps_to_non_self_neighbors() { - assert_eq!(leaf_neighbor_count(0, 3), 0); - assert_eq!(leaf_neighbor_count(1, 3), 0); - assert_eq!(leaf_neighbor_count(4, 4), 3); - assert_eq!(leaf_neighbor_count(8, 5), 5); - } - - #[test] - fn kernel_accepts_different_neighbor_counts() { - let points = 7; - let dots = test_dots(Metric::L2, points); - let input = test_input(&dots, points); - let norms = prepared_test_norms(Metric::L2, input); - let mut workspace = LeafKernelWorkspace::default(); - - for neighbor_count in [1, 3, 5, 2] { - let mut output = vec![LeafNeighbor::default(); points * neighbor_count]; - dispatch_nearest_neighbors( - Metric::L2, - input, - &norms, - MutMatrixView::try_from(output.as_mut_slice(), points, neighbor_count).unwrap(), - &mut workspace, - ) - .unwrap(); - assert!(output.iter().all(|neighbor| neighbor.target != u32::MAX)); + match metric { + Metric::L2 => prepare::(lower_gram), + Metric::Cosine => prepare::(lower_gram), + Metric::CosineNormalized => prepare::(lower_gram), + Metric::InnerProduct => prepare::(lower_gram), } } - #[test] - fn vector_pipeline_selects_exact_l2_neighbors_and_reuses_workspace() { - use super::super::kernel_metric::L2; - - let values = [0.0, 1.0, 3.0, 10.0]; - let points = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let mut output = [LeafNeighbor::default(); 8]; - let mut workspace = LeafKernelWorkspace::default(); - select_leaf_neighbors::<_, L2>( - diskann_wide::ARCH, - points, - MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), - &mut workspace, - ) - .unwrap(); - - assert_eq!( - output, - [ - LeafNeighbor::new(1, 1.0), - LeafNeighbor::new(2, 9.0), - LeafNeighbor::new(0, 1.0), - LeafNeighbor::new(2, 4.0), - LeafNeighbor::new(1, 4.0), - LeafNeighbor::new(0, 9.0), - LeafNeighbor::new(2, 49.0), - LeafNeighbor::new(1, 81.0), - ] - ); - - let dot_scratch = workspace.dot_scratch.as_ptr(); - let norm_scratch = workspace.norm_scratch.as_ptr(); - let worst = workspace.worst.as_ptr(); - let mut smaller_output = [LeafNeighbor::default(); 6]; - select_leaf_neighbors::<_, L2>( - diskann_wide::ARCH, - MatrixView::try_from(&values[..3], 3, 1).unwrap(), - MutMatrixView::try_from(&mut smaller_output[..], 3, 2).unwrap(), - &mut workspace, + fn rank_neighbors_with_workspace( + metric: Metric, + dots: &[f32], + points: usize, + requested_k: usize, + workspace: &mut LeafKernelWorkspace, + ) -> (usize, Vec) { + let leaf_k = leaf_neighbor_count(points, requested_k); + let lower_gram = lower_gram_view(dots, points); + let norms = metric_norms(metric, lower_gram); + let mut output = vec![LeafNeighbor::default(); points * leaf_k]; + arch::dispatch1_no_features( + DispatchMetric(metric), + KernelCall { + input: lower_gram, + norms: &norms, + output: MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), + workspace, + }, ) .unwrap(); - - assert_eq!(workspace.dot_scratch.as_ptr(), dot_scratch); - assert_eq!(workspace.norm_scratch.as_ptr(), norm_scratch); - assert_eq!(workspace.worst.as_ptr(), worst); + (leaf_k, output) } - #[test] - fn workspace_can_shrink_and_grow_between_calls() { - let mut workspace = LeafKernelWorkspace::default(); - for points in [17, 7, 17] { - let dots = test_dots(Metric::L2, points); - let mut output = vec![LeafNeighbor::default(); points * 2]; - let input = test_input(&dots, points); - let norms = prepared_test_norms(Metric::L2, input); - dispatch_nearest_neighbors( - Metric::L2, - input, - &norms, - MutMatrixView::try_from(output.as_mut_slice(), points, 2).unwrap(), - &mut workspace, - ) - .unwrap(); - assert!(output.iter().all(|neighbor| neighbor.target != u32::MAX)); - } + fn rank_neighbors( + metric: Metric, + dots: &[f32], + points: usize, + requested_k: usize, + ) -> (usize, Vec) { + rank_neighbors_with_workspace( + metric, + dots, + points, + requested_k, + &mut LeafKernelWorkspace::default(), + ) } -} -#[cfg(test)] -#[allow( - clippy::expect_used, - clippy::unwrap_used, - reason = "deterministic test fixture construction must abort on invalid setup" -)] -mod integration_tests { - use std::cmp::Ordering; - use super::{ - LeafKernelError, LeafKernelWorkspace, LeafNeighbor, dispatch_nearest_neighbors, - leaf_neighbor_count, prepared_test_norms, - }; - use diskann_utils::views::{MatrixView, MutMatrixView}; - use diskann_vector::distance::Metric; - - 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; - const SOURCE_MIXER: usize = 17; - const TARGET_MIXER: usize = 11; - const MIX_MODULUS: usize = 23; - const MIX_CENTER: f32 = 11.0; - const DOT_FACTOR: f32 = 1.0 / 32.0; - const TIED_TARGETS: [usize; 2] = [1, 2]; - - fn differential_dots(metric: Metric, points: usize) -> Vec { - let mut dots = vec![f32::NAN; points * points]; - for source in 0..points { - dots[source * points + source] = - if metric == Metric::Cosine && source == ZERO_NORM_POSITION { - 0.0 - } else if source == DISTINCT_NORM_POSITION { - 2.0 - } else { - 1.0 + (source % NORM_PERIOD) as f32 - }; - for target in 0..source { - let pair = ((source * SOURCE_MIXER + target * TARGET_MIXER) % MIX_MODULUS) as f32 - - MIX_CENTER; - dots[source * points + target] = if TIED_TARGETS.contains(&target) { - 0.5 + fn reference_distance( + metric: Metric, + dot: f32, + source_diagonal: f32, + target_diagonal: f32, + ) -> f32 { + match metric { + Metric::L2 => ((-2.0_f32).mul_add(dot, source_diagonal) + target_diagonal).max(0.0), + Metric::CosineNormalized => 1.0 - dot, + Metric::InnerProduct => -dot, + Metric::Cosine => { + let source_norm = source_diagonal.sqrt(); + let target_norm = target_diagonal.sqrt(); + if source_norm < f32::MIN_POSITIVE.sqrt() || target_norm < f32::MIN_POSITIVE.sqrt() + { + 1.0 } else { - pair * DOT_FACTOR - }; + let similarity = dot / (source_norm * target_norm); + 1.0 - similarity.clamp(-1.0, 1.0) + } } } - dots } - fn test_input(dots: &[f32], points: usize) -> MatrixView<'_, f32> { - MatrixView::try_from(dots, points, points).unwrap() - } - - fn brute_force_reference( + fn reference_neighbors( + metric: Metric, 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(|source| { - let diagonal = dots[source * points + source]; - if metric == Metric::Cosine { - if diagonal < f32::MIN_POSITIVE { - 0.0 - } else { - diagonal.sqrt() - } - } else { - diagonal - } - }) - .collect(); - for source in 0..points { - let mut candidates = Vec::with_capacity(points - 1); + let mut candidates = Vec::with_capacity(points.saturating_sub(1)); for target in 0..points { - if target == source { + if source == target { continue; } - let (lower_source, lower_target) = if source > target { + let (row, column) = if source > target { (source, target) } else { (target, source) }; - 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((-2.0_f32).mul_add(dot, norms[source]) + norms[target]), - Metric::CosineNormalized => 1.0 - dot, - Metric::InnerProduct => -dot, - Metric::Cosine => { - let denominator = norms[source] * norms[target]; - let similarity = if denominator == 0.0 { - 0.0 - } else { - dot / denominator - }; - 1.0 - (-1.0_f32).max(1.0_f32.min(similarity)) - } - }; + let distance = reference_distance( + metric, + dots[row * points + column], + dots[source * points + source], + dots[target * points + target], + ); if distance.partial_cmp(&f32::INFINITY) == Some(Ordering::Less) { candidates.push(LeafNeighbor::new(target as u32, distance)); } } - candidates.sort_by(|left, right| { - left.distance - .partial_cmp(&right.distance) - .expect("NaN distances were filtered") - }); - let count = candidates.len().min(leaf_k); - output[source * leaf_k..source * leaf_k + count].copy_from_slice(&candidates[..count]); + candidates.sort_by(|left, right| left.distance.total_cmp(&right.distance)); + let retained = candidates.len().min(leaf_k); + output[source * leaf_k..source * leaf_k + retained] + .copy_from_slice(&candidates[..retained]); } output } - fn run_leaf_kernel( - dots: &[f32], - points: usize, - requested_k: usize, - metric: Metric, - ) -> (usize, Vec) { - let leaf_k = leaf_neighbor_count(points, requested_k); - let input = test_input(dots, points); - let norms = prepared_test_norms(metric, input); - let mut output = vec![LeafNeighbor::default(); points * leaf_k]; - dispatch_nearest_neighbors( - metric, - input, - &norms, - MutMatrixView::try_from(output.as_mut_slice(), points, leaf_k).unwrap(), - &mut LeafKernelWorkspace::default(), - ) - .unwrap(); - (leaf_k, output) + /// Build a unit-diagonal lower Gram matrix for the lane-boundary sweep. + /// Similarity decreases as the point-index separation increases. + fn index_distance_lower_gram(points: usize) -> Vec { + let mut dots = vec![f32::NAN; points * points]; + for source in 0..points { + dots[source * points + source] = 1.0; + for target in 0..source { + let separation = (source - target) as f32; + dots[source * points + target] = 1.0 - separation / points as f32; + } + } + dots } - #[test] - fn dispatched_kernel_matches_reference_across_simd_width_boundaries() { - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - for points in SIMD_BOUNDARY_POINTS { - let dots = differential_dots(metric, points); - for requested_k in [1, 2, 3, 4, 7] { - let expected = brute_force_reference(&dots, points, requested_k, metric); - let actual = run_leaf_kernel(&dots, points, requested_k, metric).1; - assert_eq!(actual, expected, "{metric:?}, n={points}, k={requested_k}"); - } - } + fn square_matrix_with_constant_diagonal(points: usize, diagonal: f32) -> Vec { + let mut values = vec![0.0; points * points]; + for point in 0..points { + values[point * points + point] = diagonal; } + values } - #[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, - ]; + mod insert_neighbor_tests { + use super::*; + + #[test] + fn one_slot_insertion_replaces_the_retained_neighbor() { + // Given + let retained_neighbor = LeafNeighbor::new(1, 4.0); + let nearer_candidate = LeafNeighbor::new(2, 2.0); + let expected_neighbors = [nearer_candidate]; + let mut actual_neighbors = [retained_neighbor]; + + // When + insert_neighbor( + &mut actual_neighbors, + nearer_candidate.target, + nearer_candidate.distance, + ); - assert_eq!( - run_leaf_kernel(&dots, 4, 2, Metric::L2).1, - [ - 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), - ] - ); - } + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } - #[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, - ]; - for (metric, expected) in [ - (Metric::L2, [1, 2, 1]), - (Metric::Cosine, [1, 2, 1]), - (Metric::CosineNormalized, [1, 2, 1]), - (Metric::InnerProduct, [1, 2, 1]), - ] { - let positions: Vec<_> = run_leaf_kernel(&dots, 3, 1, metric) - .1 - .iter() - .map(|neighbor| neighbor.target) - .collect(); - assert_eq!(positions, expected, "metric {metric:?}"); + #[test] + fn two_slot_insertion_places_a_nearer_candidate_first() { + // Given + let nearest = LeafNeighbor::new(1, 1.0); + let farthest = LeafNeighbor::new(2, 3.0); + let nearer_candidate = LeafNeighbor::new(3, 0.5); + let expected_neighbors = [nearer_candidate, nearest]; + let mut actual_neighbors = [nearest, farthest]; + + // When + insert_neighbor( + &mut actual_neighbors, + nearer_candidate.target, + nearer_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); } - } - #[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, - ]; + #[test] + fn two_slot_insertion_places_a_middle_distance_last() { + // Given + let nearest = LeafNeighbor::new(1, 1.0); + let farthest = LeafNeighbor::new(2, 3.0); + let eligible_candidate = LeafNeighbor::new(3, 2.0); + let expected_neighbors = [nearest, eligible_candidate]; + let mut actual_neighbors = [nearest, farthest]; + + // When + insert_neighbor( + &mut actual_neighbors, + eligible_candidate.target, + eligible_candidate.distance, + ); - let output = run_leaf_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)); - } + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } - #[test] - fn l2_fma_avoids_intermediate_overflow_in_scalar_and_simd_paths() { - let dot = f32::from_bits(f32::MAX.to_bits() - 1); - let expected = (-2.0_f32).mul_add(dot, f32::MAX) + f32::MAX; - assert!(expected.is_finite() && expected > 0.0); + #[test] + fn three_slot_insertion_places_the_nearest_candidate_first() { + // Given + let nearest = LeafNeighbor::new(1, 1.0); + let middle = LeafNeighbor::new(2, 2.0); + let farthest = LeafNeighbor::new(3, 4.0); + let nearer_candidate = LeafNeighbor::new(4, 0.5); + let expected_neighbors = [nearer_candidate, nearest, middle]; + let mut actual_neighbors = [nearest, middle, farthest]; + + // When + insert_neighbor( + &mut actual_neighbors, + nearer_candidate.target, + nearer_candidate.distance, + ); - let scalar = [f32::MAX, 0.0, dot, f32::MAX]; - let scalar_output = run_leaf_kernel(&scalar, 2, 1, Metric::L2).1; - assert_eq!(scalar_output[0].distance.to_bits(), expected.to_bits()); + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } - let points = 17; - let mut simd = vec![0.0; points * points]; - for point in 0..points { - simd[point * points + point] = f32::MAX; + #[test] + fn three_slot_insertion_places_a_middle_candidate_between_neighbors() { + // Given + let nearest = LeafNeighbor::new(1, 1.0); + let middle = LeafNeighbor::new(2, 2.0); + let farthest = LeafNeighbor::new(3, 4.0); + let middle_candidate = LeafNeighbor::new(4, 1.5); + let expected_neighbors = [nearest, middle_candidate, middle]; + let mut actual_neighbors = [nearest, middle, farthest]; + + // When + insert_neighbor( + &mut actual_neighbors, + middle_candidate.target, + middle_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); } - simd[16 * points] = dot; - let simd_output = run_leaf_kernel(&simd, points, 1, Metric::L2).1; - assert_eq!(simd_output[16].target, 0); - assert_eq!(simd_output[16].distance.to_bits(), expected.to_bits()); - } - #[test] - fn cosine_clamps_simd_similarity_to_metric_range() { - let points = 17; - let mut dots = vec![0.0; points * points]; - for point in 0..points { - dots[point * points + point] = 1.0; + #[test] + fn three_slot_insertion_replaces_the_farthest_neighbor() { + // Given + let nearest = LeafNeighbor::new(1, 1.0); + let middle = LeafNeighbor::new(2, 2.0); + let farthest = LeafNeighbor::new(3, 4.0); + let eligible_candidate = LeafNeighbor::new(4, 3.0); + let expected_neighbors = [nearest, middle, eligible_candidate]; + let mut actual_neighbors = [nearest, middle, farthest]; + + // When + insert_neighbor( + &mut actual_neighbors, + eligible_candidate.target, + eligible_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); } - dots[16 * points] = 1.000_001; - dots[16 * points + 1] = -1.000_001; - let output = run_leaf_kernel(&dots, points, 16, Metric::Cosine).1; - let source = &output[16 * 16..17 * 16]; - assert_eq!(source[0], LeafNeighbor::new(0, 0.0)); - assert_eq!(source[15], LeafNeighbor::new(1, 2.0)); + #[test] + fn runtime_width_insertion_shifts_only_the_later_neighbors() { + // Given + let first = LeafNeighbor::new(1, 1.0); + let second = LeafNeighbor::new(2, 2.0); + let third = LeafNeighbor::new(3, 3.0); + let fourth = LeafNeighbor::new(4, 5.0); + let candidate = LeafNeighbor::new(5, 2.5); + let expected_neighbors = [first, second, candidate, third]; + let mut actual_neighbors = [first, second, third, fourth]; + + // When + insert_neighbor( + actual_neighbors.as_mut_slice(), + candidate.target, + candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[test] + fn sorted_insertion_preserves_existing_order_for_equal_distances() { + // Given + let nearest = LeafNeighbor::new(1, 1.0); + let existing_tie = LeafNeighbor::new(2, 2.0); + let farthest = LeafNeighbor::new(3, 4.0); + let tied_candidate = LeafNeighbor::new(4, 2.0); + let expected_neighbors = [nearest, existing_tie, tied_candidate]; + let mut actual_neighbors = [nearest, existing_tie, farthest]; + + // When + insert_neighbor( + &mut actual_neighbors, + tied_candidate.target, + tied_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } } - #[test] - fn clamps_leaf_distances_and_cosine_similarity() { - #[rustfmt::skip] - let out_of_range = [1.0, 0.0, 2.0, 1.0]; - assert_eq!( - run_leaf_kernel(&out_of_range, 2, 1, Metric::L2).1[0].distance, - 0.0 - ); - assert_eq!( - run_leaf_kernel(&out_of_range, 2, 1, Metric::CosineNormalized).1[0].distance, - -1.0 - ); - assert_eq!( - run_leaf_kernel(&out_of_range, 2, 1, Metric::Cosine).1[0].distance, - 0.0 - ); + mod leaf_neighbor_count_tests { + use super::leaf_neighbor_count; - #[rustfmt::skip] - let opposite = [1.0, 0.0, -2.0, 1.0]; - assert_eq!( - run_leaf_kernel(&opposite, 2, 1, Metric::Cosine).1[0].distance, - 2.0 - ); - - let subnormal = [f32::MIN_POSITIVE / 2.0, 0.0, 1.0, 1.0]; - assert_eq!( - run_leaf_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_leaf_kernel(&minimum_normal, 2, 1, Metric::Cosine).1[0].distance, - 0.0 - ); + #[test] + fn empty_leaf_cannot_retain_neighbors() { + // Given + let point_count = 0; + let requested_k = 3; + let expected_neighbor_count = 0; + + // When + let actual_neighbor_count = leaf_neighbor_count(point_count, requested_k); + + // Then + assert_eq!(actual_neighbor_count, expected_neighbor_count); + } + + #[test] + fn singleton_leaf_cannot_retain_its_source_point() { + // Given + let point_count = 1; + let requested_k = 3; + let expected_non_self_neighbor_count = 0; + + // When + let actual_neighbor_count = leaf_neighbor_count(point_count, requested_k); + + // Then + assert_eq!(actual_neighbor_count, expected_non_self_neighbor_count); + } + + #[test] + fn requested_k_above_available_neighbors_is_clamped() { + // Given + let point_count = 4; + let requested_k = 4; + let expected_all_non_self_neighbors = point_count - 1; + + // When + let actual_neighbor_count = leaf_neighbor_count(point_count, requested_k); + + // Then + assert_eq!(actual_neighbor_count, expected_all_non_self_neighbors); + } + + #[test] + fn requested_k_within_available_neighbors_is_unchanged() { + // Given + let point_count = 8; + let requested_k = 5; + let expected_requested_neighbor_count = requested_k; + + // When + let actual_neighbor_count = leaf_neighbor_count(point_count, requested_k); + + // Then + assert_eq!(actual_neighbor_count, expected_requested_neighbor_count); + } } - #[test] - fn finite_max_distance_fills_the_final_fixed_slot() { - let points = 4; - let mut dots = vec![0.0; points * points]; - dots[3 * points] = -f32::MAX; - - let (leaf_k, output) = run_leaf_kernel(&dots, points, 3, Metric::InnerProduct); - assert_eq!(leaf_k, 3); - assert_eq!( - output[3 * leaf_k + leaf_k - 1], - LeafNeighbor::new(0, f32::MAX) - ); + mod select_leaf_neighbors_tests { + use super::*; + + #[test] + fn l2_pipeline_orders_neighbors_by_squared_distance() { + // Given + let values = [0.0_f32, 1.0, 3.0, 10.0]; + let points = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let expected_neighbors = [ + LeafNeighbor::new(1, (values[0] - values[1]).powi(2)), + LeafNeighbor::new(2, (values[0] - values[2]).powi(2)), + LeafNeighbor::new(0, (values[1] - values[0]).powi(2)), + LeafNeighbor::new(2, (values[1] - values[2]).powi(2)), + LeafNeighbor::new(1, (values[2] - values[1]).powi(2)), + LeafNeighbor::new(0, (values[2] - values[0]).powi(2)), + LeafNeighbor::new(2, (values[3] - values[2]).powi(2)), + LeafNeighbor::new(1, (values[3] - values[1]).powi(2)), + ]; + let mut actual_neighbors = [LeafNeighbor::default(); 8]; + + // When + select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + points, + MutMatrixView::try_from(&mut actual_neighbors[..], 4, 2).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap(); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[test] + fn reused_workspace_matches_fresh_neighbor_selection() { + // Given + let values = [0.0_f32, 1.0, 3.0, 10.0]; + let smaller_points = MatrixView::try_from(&values[..3], 3, 1).unwrap(); + let mut reused_workspace = LeafKernelWorkspace::default(); + let mut discarded_large_output = [LeafNeighbor::default(); 8]; + select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + MatrixView::try_from(&values[..], 4, 1).unwrap(), + MutMatrixView::try_from(&mut discarded_large_output[..], 4, 2).unwrap(), + &mut reused_workspace, + ) + .unwrap(); + let mut expected_neighbors_from_fresh_workspace = [LeafNeighbor::default(); 6]; + select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + smaller_points, + MutMatrixView::try_from(&mut expected_neighbors_from_fresh_workspace[..], 3, 2) + .unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap(); + + // When + let mut actual_neighbors_from_reused_workspace = [LeafNeighbor::default(); 6]; + select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + smaller_points, + MutMatrixView::try_from(&mut actual_neighbors_from_reused_workspace[..], 3, 2) + .unwrap(), + &mut reused_workspace, + ) + .unwrap(); + + // Then + assert_eq!( + actual_neighbors_from_reused_workspace, + expected_neighbors_from_fresh_workspace + ); + } } - #[test] - fn leaf_metrics_define_nan_candidate_behavior() { + mod rank_leaf_dots_tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::two_points_fixed_one(2, 1)] + #[case::scalar_fixed_two(7, 2)] + #[case::lane_minus_one_fixed_three(15, 3)] + #[case::one_complete_lane_fixed_three(16, 3)] + #[case::lane_plus_one_runtime_width(17, 4)] + #[case::two_lanes_minus_one_runtime_width(31, 7)] + #[case::two_complete_lanes_runtime_width(32, 7)] + #[case::two_lanes_plus_one_runtime_width(33, 7)] + #[case::four_complete_lanes_runtime_width(64, 7)] + #[case::sixteen_complete_lanes_runtime_width(256, 7)] + #[case::maximum_leaf_size_runtime_width(512, 7)] + #[trace] + fn dispatched_leaf_ranking_matches_scalar_reference_across_lane_boundaries( + #[values( + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct + )] + metric: Metric, + #[case] point_count: usize, + #[case] requested_k: usize, + ) { + // Given + let dots = index_distance_lower_gram(point_count); + let expected_neighbors = reference_neighbors(metric, &dots, point_count, requested_k); + + // When + let actual_neighbors = rank_neighbors(metric, &dots, point_count, requested_k).1; + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + const UNIT_SQUARED_NORM: f32 = 1.0; + const POINT_0_1_DOT: f32 = 0.0; + const POINT_0_2_DOT: f32 = -1.0; + const POINT_1_2_DOT: f32 = 0.5; + #[rustfmt::skip] + const THREE_POINT_LOWER_GRAM: [f32; 9] = [ + UNIT_SQUARED_NORM, f32::NAN, f32::NAN, + POINT_0_1_DOT, UNIT_SQUARED_NORM, f32::NAN, + POINT_0_2_DOT, POINT_1_2_DOT, UNIT_SQUARED_NORM, + ]; + + fn rank_three_point_fixture(metric: Metric) -> Vec { + rank_neighbors(metric, &THREE_POINT_LOWER_GRAM, 3, 1).1 + } + + #[test] + fn l2_selects_the_nearest_target_from_the_lower_gram_triangle() { + // Given + let expected_neighbors = [ + LeafNeighbor::new(1, 2.0 * UNIT_SQUARED_NORM - 2.0 * POINT_0_1_DOT), + LeafNeighbor::new(2, 2.0 * UNIT_SQUARED_NORM - 2.0 * POINT_1_2_DOT), + LeafNeighbor::new(1, 2.0 * UNIT_SQUARED_NORM - 2.0 * POINT_1_2_DOT), + ]; + + // When + let actual_neighbors = rank_three_point_fixture(Metric::L2); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[test] + fn cosine_selects_the_nearest_target_from_the_lower_gram_triangle() { + // Given + let expected_neighbors = [ + LeafNeighbor::new(1, 1.0 - POINT_0_1_DOT), + LeafNeighbor::new(2, 1.0 - POINT_1_2_DOT), + LeafNeighbor::new(1, 1.0 - POINT_1_2_DOT), + ]; + + // When + let actual_neighbors = rank_three_point_fixture(Metric::Cosine); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[test] + fn normalized_cosine_selects_the_largest_dot_product() { + // Given + let expected_neighbors = [ + LeafNeighbor::new(1, 1.0 - POINT_0_1_DOT), + LeafNeighbor::new(2, 1.0 - POINT_1_2_DOT), + LeafNeighbor::new(1, 1.0 - POINT_1_2_DOT), + ]; + + // When + let actual_neighbors = rank_three_point_fixture(Metric::CosineNormalized); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[test] + fn inner_product_selects_the_largest_dot_product() { + // Given + let expected_neighbors = [ + LeafNeighbor::new(1, -POINT_0_1_DOT), + LeafNeighbor::new(2, -POINT_1_2_DOT), + LeafNeighbor::new(1, -POINT_1_2_DOT), + ]; + + // When + let actual_neighbors = rank_three_point_fixture(Metric::InnerProduct); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[test] + fn equal_l2_distances_keep_target_scan_order() { + // Given + let unit_squared_norm = 1.0; + let tied_dot_product = 0.0; + let expected_tied_distance = 2.0 * unit_squared_norm - 2.0 * tied_dot_product; + #[rustfmt::skip] let dots = [ - 1.0, 0.0, 0.0, - f32::NAN, 1.0, 0.0, - 0.5, 0.25, 1.0, + unit_squared_norm, f32::NAN, f32::NAN, f32::NAN, + tied_dot_product, unit_squared_norm, f32::NAN, f32::NAN, + tied_dot_product, tied_dot_product, unit_squared_norm, f32::NAN, + tied_dot_product, tied_dot_product, tied_dot_product, unit_squared_norm, ]; + let expected_neighbors_in_scan_order = [ + LeafNeighbor::new(1, expected_tied_distance), + LeafNeighbor::new(2, expected_tied_distance), + LeafNeighbor::new(0, expected_tied_distance), + LeafNeighbor::new(2, expected_tied_distance), + LeafNeighbor::new(0, expected_tied_distance), + LeafNeighbor::new(1, expected_tied_distance), + LeafNeighbor::new(0, expected_tied_distance), + LeafNeighbor::new(1, expected_tied_distance), + ]; + + // When + let actual_neighbors = rank_neighbors(Metric::L2, &dots, 4, 2).1; + + // Then + assert_eq!(actual_neighbors, expected_neighbors_in_scan_order); + } - for metric in [Metric::L2, Metric::Cosine] { - let output = run_leaf_kernel(&dots, 3, 1, metric).1; - assert_eq!(output[0], LeafNeighbor::new(1, 0.0), "metric {metric:?}"); - assert_eq!(output[1], LeafNeighbor::new(0, 0.0), "metric {metric:?}"); + #[test] + fn scalar_l2_distance_stays_finite_when_twice_the_dot_product_overflows() { + // Given + let dot_product = f32::from_bits(f32::MAX.to_bits() - 1); + let unfused_twice_dot_product = 2.0 * dot_product; + let expected_fused_distance = (-2.0_f32).mul_add(dot_product, f32::MAX) + f32::MAX; + let dots = [f32::MAX, 0.0, dot_product, f32::MAX]; + + // When + let actual_neighbors = rank_neighbors(Metric::L2, &dots, 2, 1).1; + + // Then + assert!(unfused_twice_dot_product.is_infinite()); + assert!(expected_fused_distance.is_finite() && expected_fused_distance > 0.0); + assert_eq!( + actual_neighbors[0].distance.to_bits(), + expected_fused_distance.to_bits() + ); } - for metric in [Metric::CosineNormalized, Metric::InnerProduct] { - let output = run_leaf_kernel(&dots, 3, 1, metric).1; - assert_eq!(output[0].target, 2, "metric {metric:?}"); - assert_eq!(output[1].target, 2, "metric {metric:?}"); + #[test] + fn simd_l2_distance_stays_finite_when_twice_the_dot_product_overflows() { + // Given + let dot_product = f32::from_bits(f32::MAX.to_bits() - 1); + let unfused_twice_dot_product = 2.0 * dot_product; + let expected_fused_distance = (-2.0_f32).mul_add(dot_product, f32::MAX) + f32::MAX; + let expected_simd_neighbor_target = 0; + let points = 17; + let mut dots = square_matrix_with_constant_diagonal(points, f32::MAX); + dots[16 * points] = dot_product; + + // When + let actual_neighbors = rank_neighbors(Metric::L2, &dots, points, 1).1; + + // Then + assert!(unfused_twice_dot_product.is_infinite()); + assert_eq!(actual_neighbors[16].target, expected_simd_neighbor_target); + assert_eq!( + actual_neighbors[16].distance.to_bits(), + expected_fused_distance.to_bits() + ); } - } - #[test] - fn clamps_k_to_available_non_self_neighbors() { - #[rustfmt::skip] + #[test] + fn cosine_zero_norm_produces_unit_distance() { + // Given + #[rustfmt::skip] let dots = [ - 1.0, 3.0, 3.0, - 0.0, 1.0, 3.0, - 0.0, 0.0, 1.0, + 0.0, 99.0, 99.0, + 0.0, 1.0, 99.0, + 0.0, 0.0, 1.0, ]; - let (leaf_k, output) = run_leaf_kernel(&dots, 3, 3, Metric::L2); - - 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) + let expected_zero_norm_neighbors = + [LeafNeighbor::new(1, 1.0), LeafNeighbor::new(2, 1.0)]; + + // When + let actual_neighbors = rank_neighbors(Metric::Cosine, &dots, 3, 2).1; + + // Then + assert_eq!(&actual_neighbors[..2], &expected_zero_norm_neighbors); + } + + #[test] + fn cosine_similarity_above_one_clamps_to_zero_distance() { + // Given + let dots = [1.0, 0.0, 2.0, 1.0]; + let maximum_cosine_similarity = 1.0; + let expected_one_minus_maximum_similarity = 1.0 - maximum_cosine_similarity; + + // When + let actual_neighbors = rank_neighbors(Metric::Cosine, &dots, 2, 1).1; + + // Then + assert_eq!( + actual_neighbors[0].distance, + expected_one_minus_maximum_similarity ); } - } - #[test] - fn accepts_empty_singleton_and_zero_k_inputs() { - 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_leaf_kernel(dots, points, requested_k, metric).0, 0); + #[test] + fn cosine_similarity_below_negative_one_clamps_to_distance_two() { + // Given + let dots = [1.0, 0.0, -2.0, 1.0]; + let minimum_cosine_similarity = -1.0; + let expected_one_minus_minimum_similarity = 1.0 - minimum_cosine_similarity; + + // When + let actual_neighbors = rank_neighbors(Metric::Cosine, &dots, 2, 1).1; + + // Then + assert_eq!( + actual_neighbors[0].distance, + expected_one_minus_minimum_similarity + ); } - } - #[test] - fn rejects_invalid_neighbor_counts() { - let square = [0.0; 9]; - let square_input = test_input(&square, 3); - let square_norms = prepared_test_norms(Metric::L2, square_input); - let mut too_many = [LeafNeighbor::default(); 9]; - assert_eq!( - dispatch_nearest_neighbors( - Metric::L2, - square_input, - &square_norms, - MutMatrixView::try_from(&mut too_many[..], 3, 3).unwrap(), - &mut LeafKernelWorkspace::default(), - ), - Err(LeafKernelError::InvalidNeighborCount { + #[test] + fn cosine_subnormal_norm_is_treated_as_zero() { + // Given + let dots = [f32::MIN_POSITIVE / 2.0, 0.0, 1.0, 1.0]; + let zero_norm_similarity = 0.0; + let expected_one_minus_zero_similarity = 1.0 - zero_norm_similarity; + + // When + let actual_neighbors = rank_neighbors(Metric::Cosine, &dots, 2, 1).1; + + // Then + assert_eq!( + actual_neighbors[0].distance, + expected_one_minus_zero_similarity + ); + } + + #[test] + fn f32_max_distance_is_still_a_rankable_neighbor() { + // Given + let points = 4; + let expected_leaf_k = 3; + let expected_last_neighbor = LeafNeighbor::new(0, f32::MAX); + let mut dots = vec![0.0; points * points]; + dots[3 * points] = -f32::MAX; + + // When + let (actual_leaf_k, actual_neighbors) = + rank_neighbors(Metric::InnerProduct, &dots, points, expected_leaf_k); + + // Then + assert_eq!(actual_leaf_k, expected_leaf_k); + assert_eq!( + actual_neighbors[3 * actual_leaf_k + actual_leaf_k - 1], + expected_last_neighbor + ); + } + + #[test] + fn scalar_nan_distance_leaves_the_neighbor_slot_unassigned() { + // Given + let dots = [1.0, 0.0, f32::NAN, 1.0]; + let expected_unassigned_neighbors = [LeafNeighbor::default(), LeafNeighbor::default()]; + + // When + let actual_neighbors = rank_neighbors(Metric::CosineNormalized, &dots, 2, 1).1; + + // Then + assert_eq!(actual_neighbors, expected_unassigned_neighbors); + } + + #[test] + fn simd_nan_distance_cannot_replace_a_finite_neighbor() { + // Given + let points = 17; + let mut dots = square_matrix_with_constant_diagonal(points, 1.0); + dots[16 * points] = f32::NAN; + let expected_finite_neighbor = LeafNeighbor::new(1, 1.0); + + // When + let actual_neighbors = rank_neighbors(Metric::CosineNormalized, &dots, points, 1).1; + + // Then + assert_eq!(actual_neighbors[16], expected_finite_neighbor); + } + + #[test] + fn empty_leaf_has_no_neighbors() { + // Given + let dots = []; + let expected_zero_neighbor_width = 0; + let expected_no_neighbors: [LeafNeighbor; 0] = []; + + // When + let (actual_leaf_k, actual_neighbors) = rank_neighbors(Metric::L2, &dots, 0, 2); + + // Then + assert_eq!(actual_leaf_k, expected_zero_neighbor_width); + assert_eq!(actual_neighbors, expected_no_neighbors); + } + + #[test] + fn singleton_leaf_has_no_neighbors() { + // Given + let dots = [4.0]; + let expected_zero_neighbor_width = 0; + let expected_no_neighbors: [LeafNeighbor; 0] = []; + + // When + let (actual_leaf_k, actual_neighbors) = rank_neighbors(Metric::Cosine, &dots, 1, 2); + + // Then + assert_eq!(actual_leaf_k, expected_zero_neighbor_width); + assert_eq!(actual_neighbors, expected_no_neighbors); + } + + #[test] + fn zero_requested_k_has_no_neighbors() { + // Given + let dots = [1.0, 0.0, 0.0, 1.0]; + let expected_zero_neighbor_width = 0; + let expected_no_neighbors: [LeafNeighbor; 0] = []; + + // When + let (actual_leaf_k, actual_neighbors) = + rank_neighbors(Metric::InnerProduct, &dots, 2, 0); + + // Then + assert_eq!(actual_leaf_k, expected_zero_neighbor_width); + assert_eq!(actual_neighbors, expected_no_neighbors); + } + + #[test] + fn neighbor_width_equal_to_point_count_is_rejected() { + // Given + let dots = [0.0; 9]; + let input = lower_gram_view(&dots, 3); + let norms = metric_norms(Metric::L2, input); + let expected_error = LeafKernelError::InvalidNeighborCount { points: 3, neighbors: 3, maximum: 2, - }) - ); - - let square = [0.0; 25]; - let square_input = test_input(&square, 5); - let square_norms = prepared_test_norms(Metric::L2, square_input); - let mut too_wide = [LeafNeighbor::default(); 25]; - assert_eq!( - dispatch_nearest_neighbors( - Metric::L2, - square_input, - &square_norms, - MutMatrixView::try_from(&mut too_wide[..], 5, 5).unwrap(), - &mut LeafKernelWorkspace::default(), - ), - Err(LeafKernelError::InvalidNeighborCount { - points: 5, - neighbors: 5, - maximum: 4, - }) - ); - } - - #[test] - fn normalized_cosine_keeps_nan_non_rankable_in_scalar_and_simd_paths() { - let scalar = [1.0, 0.0, f32::NAN, 1.0]; - let scalar_output = run_leaf_kernel(&scalar, 2, 1, Metric::CosineNormalized).1; - assert_eq!(scalar_output[0], LeafNeighbor::default()); + }; + let mut output = [LeafNeighbor::default(); 9]; + + // When + let actual_error = arch::dispatch1_no_features( + DispatchMetric(Metric::L2), + KernelCall { + input, + norms: &norms, + output: MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), + workspace: &mut LeafKernelWorkspace::default(), + }, + ) + .unwrap_err(); - let points = 17; - let mut dots = vec![0.0; points * points]; - for point in 0..points { - dots[point * points + point] = 1.0; + // Then + assert_eq!(actual_error, expected_error); } - dots[16 * points] = f32::NAN; - let simd_output = run_leaf_kernel(&dots, points, 1, Metric::CosineNormalized).1; - assert_eq!(simd_output[16], LeafNeighbor::new(1, 1.0)); } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 58580e53ef..ba03067cb6 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -18,8 +18,6 @@ use std::marker::PhantomData; use crate::{ANNError, ANNResult}; use diskann_linalg::Transpose; use diskann_utils::views::{MatrixView, MutMatrixView}; -#[cfg(test)] -use diskann_vector::distance::Metric; use diskann_wide::{SIMDMask, SIMDVector}; use super::{ @@ -245,73 +243,53 @@ fn insert_leader(ranked_leaders: &mut [(u32, f32)], leader: u32, score: f32) { } #[cfg(test)] -struct DispatchedPartitionCall<'a> { - input: PartitionInput<'a>, - output: MutMatrixView<'a, u32>, - ranked_leaders: &'a mut Vec<(u32, f32)>, -} +mod tests { + use super::*; + use crate::graph::pipnn::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use diskann_utils::views::{Matrix, MatrixView, MutMatrixView}; + use diskann_vector::distance::Metric; + use diskann_wide::arch::{self, Target1}; -#[cfg(test)] -struct DispatchPartitionForTest(Metric); + struct KernelCall<'a> { + input: PartitionInput<'a>, + output: MutMatrixView<'a, u32>, + ranked_leaders: &'a mut Vec<(u32, f32)>, + } -#[cfg(test)] -impl diskann_wide::arch::Target1> for DispatchPartitionForTest -where - A: PiPNNSIMDSchema, -{ - fn run(self, arch: A, call: DispatchedPartitionCall<'_>) { - use super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + struct DispatchMetric(Metric); - match self.0 { - Metric::L2 => { - rank_leader_dots::(arch, call.input, call.output, call.ranked_leaders) - } - Metric::Cosine => { - rank_leader_dots::(arch, call.input, call.output, call.ranked_leaders) + impl Target1> for DispatchMetric + where + A: PiPNNSIMDSchema, + { + fn run(self, arch: A, call: KernelCall<'_>) { + match self.0 { + Metric::L2 => { + rank_leader_dots::(arch, call.input, call.output, call.ranked_leaders) + } + Metric::Cosine => rank_leader_dots::( + arch, + call.input, + call.output, + call.ranked_leaders, + ), + Metric::CosineNormalized => rank_leader_dots::( + arch, + call.input, + call.output, + call.ranked_leaders, + ), + Metric::InnerProduct => rank_leader_dots::( + arch, + call.input, + call.output, + call.ranked_leaders, + ), } - Metric::CosineNormalized => rank_leader_dots::( - arch, - call.input, - call.output, - call.ranked_leaders, - ), - Metric::InnerProduct => rank_leader_dots::( - arch, - call.input, - call.output, - call.ranked_leaders, - ), } } -} -#[cfg(test)] -fn dispatch_nearest_leaders( - metric: Metric, - input: PartitionInput<'_>, - output: MutMatrixView<'_, u32>, - ranked_leaders: &mut Vec<(u32, f32)>, -) { - diskann_wide::arch::dispatch1_no_features( - DispatchPartitionForTest(metric), - DispatchedPartitionCall { - input, - output, - ranked_leaders, - }, - ) -} - -#[cfg(test)] -mod tests { - use super::super::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, PartitionMetric, - }; - - use super::*; - use diskann_vector::distance::Metric; - - fn test_input<'a>( + fn partition_input<'a>( dots: &'a [f32], point_count: usize, leader_count: usize, @@ -327,439 +305,348 @@ mod tests { } } - // This oracle checks SIMD groups, single values, and retained-leader order. - // It uses the single-value ranking formula for metric `M`. - fn ranking_reference( + fn rank_partition_leaders( + metric: Metric, input: PartitionInput<'_>, fanout: usize, - output: &mut [u32], - ) { - for (point, (point_dots, point_output)) in input - .dots - .row_iter() - .zip(output.chunks_exact_mut(fanout)) - .enumerate() - { - let point_norm = M::point_single(input.norms, point); - let mut ranked_leaders = vec![(u32::MAX, f32::INFINITY); fanout]; - for (leader, &dot) in point_dots.iter().enumerate() { - insert_leader( - &mut ranked_leaders, - leader as u32, - M::ranking_single(input.norms, point_norm, dot, leader), - ); - } - for (destination, &(leader, _)) in point_output.iter_mut().zip(&ranked_leaders) { - *destination = leader; - } - } - } - - fn single_ranking( - dot_product: f32, - point_norms: &[f32], - leader_norms: &[f32], - ) -> f32 { - let norms = PartitionNorms { - point_norms, - leader_norms, - }; - M::ranking_single(norms, M::point_single(norms, 0), dot_product, 0) - } - - #[test] - fn single_ranking_matches_metric_contract() { - assert_eq!(single_ranking::(2.0, &[], &[9.0]), 5.0); - assert_eq!(single_ranking::(0.25, &[], &[]), 0.75); - assert_eq!(single_ranking::(3.0, &[], &[]), -3.0); - assert_eq!(single_ranking::(4.0, &[2.0], &[4.0]), 0.5); - assert_eq!(single_ranking::(5.0, &[2.0], &[2.0]), 0.0); - assert_eq!(single_ranking::(-5.0, &[2.0], &[2.0]), 2.0); - assert_eq!(single_ranking::(4.0, &[0.0], &[4.0]), 1.0); - assert_eq!(single_ranking::(1.0, &[f32::NAN], &[1.0]), 0.0); - } - - #[test] - fn cosine_special_norms_match_single_and_dispatched_kernel() { - let leader_count = 17; - let point_norms = [0.0, 0.0, f32::MIN_POSITIVE.sqrt(), f32::NAN]; - let dots = vec![1.0; point_norms.len() * leader_count]; - let mut leader_norms = vec![1.0; leader_count]; - leader_norms[..4].copy_from_slice(&[ - 0.0, - f32::MIN_POSITIVE.sqrt() / 2.0, - f32::MIN_POSITIVE.sqrt(), - f32::NAN, - ]); - let input = test_input( - &dots, - point_norms.len(), - leader_count, - &point_norms, - &leader_norms, - ); - let mut expected = vec![u32::MAX; point_norms.len() * 2]; - ranking_reference::(input, 2, &mut expected); - let mut actual = vec![u32::MAX; point_norms.len() * 2]; - dispatch_nearest_leaders( - Metric::Cosine, - input, - MutMatrixView::try_from(actual.as_mut_slice(), point_norms.len(), 2).unwrap(), - &mut Vec::new(), - ); - - assert_eq!(actual, expected); - assert_eq!(&actual[..4], &[0, 1, 0, 1]); - assert_eq!(&actual[6..], &[2, 3]); - } - - #[test] - fn topk_orders_candidates_and_preserves_ties() { - let mut ranked_leaders = vec![(u32::MAX, f32::INFINITY); 4]; - for (leader, distance) in [(0, 4.0), (1, 1.0), (2, 3.0), (3, 2.0), (4, 1.0)] { - insert_leader(&mut ranked_leaders, leader, distance); - } - insert_leader(&mut ranked_leaders, 5, f32::NAN); - - assert_eq!(ranked_leaders[..], [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]); - } - - #[test] - fn vector_pipeline_assigns_cosine_leaders_and_reuses_workspace() { - let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; - let leaders = - PreparedLeaders::::new(MatrixView::try_from(&leader_values[..], 3, 2).unwrap()); - let point_values = [0.9, 0.1, -0.8, 0.2]; - let points = MatrixView::try_from(&point_values[..], 2, 2).unwrap(); - let mut output = [u32::MAX; 4]; - let mut workspace = PartitionKernelWorkspace::default(); - assign_leaders::<_, Cosine>( - diskann_wide::ARCH, - points, - &leaders, - MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), - &mut workspace, - ) - .unwrap(); - - assert_eq!(output, [0, 1, 2, 1]); - let dot_scratch = workspace.dot_scratch.as_ptr(); - let point_norm_scratch = workspace.point_norm_scratch.as_ptr(); - let ranked_leader_scratch = workspace.ranked_leader_scratch.as_ptr(); - - let mut smaller_output = [u32::MAX; 2]; - assign_leaders::<_, Cosine>( - diskann_wide::ARCH, - MatrixView::try_from(&point_values[..2], 1, 2).unwrap(), - &leaders, - MutMatrixView::try_from(&mut smaller_output[..], 1, 2).unwrap(), - &mut workspace, - ) - .unwrap(); - - assert_eq!(smaller_output, [0, 1]); - assert_eq!(workspace.dot_scratch.as_ptr(), dot_scratch); - assert_eq!(workspace.point_norm_scratch.as_ptr(), point_norm_scratch); - assert_eq!( - workspace.ranked_leader_scratch.as_ptr(), - ranked_leader_scratch - ); - } - - #[test] - fn ranked_leaders_reuses_runtime_fanout_capacity() { - let dots = [0.0; 32]; - let input = test_input(&dots, 1, 32, &[], &[]); - let mut ranked_leaders = Vec::new(); - let mut wide_output = [u32::MAX; 32]; - dispatch_nearest_leaders( - Metric::InnerProduct, - input, - MutMatrixView::try_from(&mut wide_output[..], 1, 32).unwrap(), - &mut ranked_leaders, - ); - let allocation = ranked_leaders.as_ptr(); - - let mut narrow_output = [u32::MAX; 3]; - dispatch_nearest_leaders( - Metric::InnerProduct, - input, - MutMatrixView::try_from(&mut narrow_output[..], 1, 3).unwrap(), - &mut ranked_leaders, + ) -> Vec { + let mut output = Matrix::new(u32::MAX, input.dots.nrows(), fanout); + arch::dispatch1_no_features( + DispatchMetric(metric), + KernelCall { + input, + output: output.as_mut_view(), + ranked_leaders: &mut Vec::new(), + }, ); - - assert_eq!(ranked_leaders.as_ptr(), allocation); - assert_eq!(ranked_leaders.len(), 3); + output.into_inner().into_vec() } -} -#[cfg(test)] -#[allow( - clippy::expect_used, - clippy::unwrap_used, - reason = "deterministic test fixture construction must abort on invalid setup" -)] -mod integration_tests { - use super::{PartitionInput, PartitionNorms, UNASSIGNED_LEADER, dispatch_nearest_leaders}; - use diskann_utils::views::{Matrix, MatrixView}; - use diskann_vector::distance::Metric; - fn test_input<'a>( - dots: &'a [f32], - point_count: usize, - leader_count: usize, - point_norms: &'a [f32], - leader_norms: &'a [f32], - ) -> PartitionInput<'a> { - PartitionInput { - dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), - norms: PartitionNorms { - point_norms, - leader_norms, - }, + fn reference_score(metric: Metric, dot: f32, point_norm: f32, leader_norm: f32) -> f32 { + match metric { + Metric::L2 => (-2.0_f32).mul_add(dot, leader_norm), + Metric::CosineNormalized => 1.0 - dot, + Metric::InnerProduct => -dot, + Metric::Cosine => { + if point_norm < f32::MIN_POSITIVE.sqrt() || leader_norm < f32::MIN_POSITIVE.sqrt() { + 1.0 + } else { + 1.0 - (dot / (point_norm * leader_norm)).clamp(-1.0, 1.0) + } + } } } - 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_norms = input.norms.point_norms; - let leader_norms = input.norms.leader_norms; - let mut assignments = vec![u32::MAX; point_count * fanout]; - for (point, (point_dots, point_assignments)) in input + fn reference_assignments(metric: Metric, input: PartitionInput<'_>, fanout: usize) -> Vec { + let mut output = vec![UNASSIGNED_LEADER; input.dots.nrows() * fanout]; + for (point, (dots, assignments)) in input .dots - .as_slice() - .chunks_exact(leader_count) - .zip(assignments.chunks_exact_mut(fanout)) + .row_iter() + .zip(output.chunks_exact_mut(fanout)) .enumerate() { - let point_norm = point_norms.get(point).copied().unwrap_or(0.0); - let mut candidates: Vec<_> = point_dots + let point_norm = input.norms.point_norms.get(point).copied().unwrap_or(0.0); + let mut candidates: Vec<_> = dots .iter() .enumerate() .filter_map(|(leader, &dot)| { - let leader_norm = leader_norms.get(leader).copied().unwrap_or(0.0); - let score = match metric { - Metric::L2 => (-2.0_f32).mul_add(dot, leader_norm), - Metric::CosineNormalized => 1.0 - dot, - Metric::InnerProduct => -dot, - Metric::Cosine => { - 1.0 - if point_norm == 0.0 || leader_norm == 0.0 { - 0.0 - } else { - let cosine = dot / (point_norm * leader_norm); - (-1.0_f32).max(1.0_f32.min(cosine)) - } - } - }; + let leader_norm = input.norms.leader_norms.get(leader).copied().unwrap_or(0.0); + let score = reference_score(metric, dot, point_norm, leader_norm); (score.partial_cmp(&f32::INFINITY) == Some(std::cmp::Ordering::Less)) .then_some((leader as u32, score)) }) .collect(); - candidates.sort_by(|left, right| left.1.partial_cmp(&right.1).unwrap()); - for (destination, (leader, _)) in point_assignments.iter_mut().zip(candidates) { + candidates.sort_by(|left, right| left.1.total_cmp(&right.1)); + for (destination, (leader, _)) in assignments.iter_mut().zip(candidates) { *destination = leader; } } - assignments + output } - fn differential_data(metric: Metric, leader_count: usize) -> (Vec, Vec, Vec) { - let dots = (0..2 * leader_count) - .map(|index| { - 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 == leader_count { - f32::NAN - } else { - base * 0.25 - } - }) - .collect(); + /// Build two score rows whose preferred leader direction is opposite. + /// The fractional scores cross complete SIMD groups and scalar tails. + fn lane_boundary_fixture( + metric: Metric, + leader_count: usize, + ) -> (Vec, Vec, Vec) { + let mut dots = Vec::with_capacity(2 * leader_count); + for point in 0..2 { + for leader in 0..leader_count { + let fraction = leader as f32 / leader_count as f32; + dots.push(if point == 0 { fraction } else { 1.0 - fraction }); + } + } let point_norms = if metric == Metric::Cosine { - vec![0.0, 4.0] + vec![1.0; 2] } else { Vec::new() }; let leader_norms = match metric { - Metric::Cosine => (0..leader_count) - .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..leader_count) - .map(|leader| { - let norm = if leader == 2 || leader == 3 { - 3.0 - } else { - leader as f32 + 1.0 - }; - norm * norm - }) - .collect(), + Metric::L2 | Metric::Cosine => vec![1.0; leader_count], Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), }; (dots, point_norms, leader_norms) } - fn run_partition_kernel(metric: Metric, input: PartitionInput<'_>, fanout: usize) -> Vec { - let mut output = Matrix::new(u32::MAX, input.dots.nrows(), fanout); - dispatch_nearest_leaders(metric, input, output.as_mut_view(), &mut Vec::new()); - output.into_inner().into_vec() - } + mod insert_leader_tests { + use super::*; - #[test] - fn dispatched_kernel_matches_reference_across_simd_width_boundaries() { - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - for leader_count in [2, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let (dots, point_norms, leader_norms) = differential_data(metric, leader_count); - let input = test_input(&dots, 2, leader_count, &point_norms, &leader_norms); - for fanout in [1, 2, 16, 17, 32] { - if fanout >= leader_count { - continue; - } - assert_eq!( - run_partition_kernel(metric, input, fanout), - brute_force_reference(input, fanout, metric), - "{metric:?}, leaders={leader_count}, k={fanout}" - ); - } - } + #[test] + fn topk_keeps_nearest_first_order_and_scan_order_ties() { + // Given + let expected_ranked_leaders = [(1, 1.0), (4, 1.0), (3, 2.0), (2, 3.0)]; + let mut ranked_leaders = vec![(UNASSIGNED_LEADER, f32::INFINITY); 4]; + + // When + insert_leader(&mut ranked_leaders, 0, 4.0); + insert_leader(&mut ranked_leaders, 1, 1.0); + insert_leader(&mut ranked_leaders, 2, 3.0); + insert_leader(&mut ranked_leaders, 3, 2.0); + insert_leader(&mut ranked_leaders, 4, 1.0); + + // Then + assert_eq!(ranked_leaders, expected_ranked_leaders); } - } - #[test] - fn non_rankable_leaders_leave_unassigned_slots() { - let dots = [-0.25, f32::NAN]; + #[test] + fn nan_score_does_not_enter_the_topk() { + // Given + let expected_ranked_leaders = [(0, 0.25), (UNASSIGNED_LEADER, f32::INFINITY)]; + let mut ranked_leaders = vec![(UNASSIGNED_LEADER, f32::INFINITY); 2]; - assert_eq!( - run_partition_kernel(Metric::InnerProduct, test_input(&dots, 1, 2, &[], &[]), 2), - [0, UNASSIGNED_LEADER] - ); - } + // When + insert_leader(&mut ranked_leaders, 0, 0.25); + insert_leader(&mut ranked_leaders, 1, f32::NAN); - #[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_partition_kernel(Metric::L2, test_input(&dots, 2, 4, &[], &norms), 2), - [0, 1, 2, 1] - ); + // Then + assert_eq!(ranked_leaders, expected_ranked_leaders); + } } - #[test] - fn l2_single_matches_fused_simd_ranking() { - let mut dots = [0.0; 17]; - dots[0] = f32::MAX; - dots[16] = f32::MAX; - let leader_squared_norms = [f32::MAX; 17]; + mod assign_leaders_tests { + use super::*; - assert_eq!( - run_partition_kernel( - Metric::L2, - test_input(&dots, 1, 17, &[], &leader_squared_norms,), - 1, - ), - [0] - ); - } + #[test] + fn cosine_pipeline_assigns_each_point_to_its_nearest_leaders() { + // Given + let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; + let leaders = PreparedLeaders::::new( + MatrixView::try_from(&leader_values[..], 3, 2).unwrap(), + ); + let point_values = [0.9, 0.1, -0.8, 0.2]; + let points = MatrixView::try_from(&point_values[..], 2, 2).unwrap(); + let expected_leaders_by_descending_cosine_similarity = [0, 1, 2, 1]; + let mut actual_assignments = [UNASSIGNED_LEADER; 4]; + + // When + assign_leaders::<_, Cosine>( + diskann_wide::ARCH, + points, + &leaders, + MutMatrixView::try_from(&mut actual_assignments[..], 2, 2).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap(); - #[test] - fn supports_every_partition_metric() { - #[rustfmt::skip] - let dots = [ - 1.0, 0.0, -1.0, - 2.0, 6.0, 0.0, - ]; - for (metric, point_norms, leader_norms, 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]), - ] { + // Then assert_eq!( - run_partition_kernel( - metric, - test_input(&dots, 2, 3, point_norms, leader_norms), - 2, - ), - expected, - "metric {metric:?}" + actual_assignments, + expected_leaders_by_descending_cosine_similarity + ); + } + + #[test] + fn reused_workspace_matches_fresh_leader_assignment() { + // Given + let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; + let leaders = PreparedLeaders::::new( + MatrixView::try_from(&leader_values[..], 3, 2).unwrap(), + ); + let point_values = [0.9, 0.1, -0.8, 0.2]; + let smaller_points = MatrixView::try_from(&point_values[..2], 1, 2).unwrap(); + let mut reused_workspace = PartitionKernelWorkspace::default(); + let mut discarded_large_output = [UNASSIGNED_LEADER; 4]; + assign_leaders::<_, Cosine>( + diskann_wide::ARCH, + MatrixView::try_from(&point_values[..], 2, 2).unwrap(), + &leaders, + MutMatrixView::try_from(&mut discarded_large_output[..], 2, 2).unwrap(), + &mut reused_workspace, + ) + .unwrap(); + let mut expected_assignments_from_fresh_workspace = [UNASSIGNED_LEADER; 2]; + assign_leaders::<_, Cosine>( + diskann_wide::ARCH, + smaller_points, + &leaders, + MutMatrixView::try_from(&mut expected_assignments_from_fresh_workspace[..], 1, 2) + .unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap(); + + // When + let mut actual_assignments_from_reused_workspace = [UNASSIGNED_LEADER; 2]; + assign_leaders::<_, Cosine>( + diskann_wide::ARCH, + smaller_points, + &leaders, + MutMatrixView::try_from(&mut actual_assignments_from_reused_workspace[..], 1, 2) + .unwrap(), + &mut reused_workspace, + ) + .unwrap(); + + // Then + assert_eq!( + actual_assignments_from_reused_workspace, + expected_assignments_from_fresh_workspace ); } } - #[test] - fn cosine_treats_a_zero_norm_as_zero_similarity() { - assert_eq!( - run_partition_kernel( + mod rank_leader_dots_tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::two_leaders_fanout_one(2, 1)] + #[case::scalar_fanout_two(7, 2)] + #[case::lane_minus_one(15, 3)] + #[case::one_complete_lane(16, 3)] + #[case::lane_plus_one(17, 4)] + #[case::two_lanes_minus_one(31, 7)] + #[case::two_complete_lanes(32, 7)] + #[case::two_lanes_plus_one(33, 7)] + #[trace] + fn dispatched_partition_ranking_matches_scalar_reference_across_lane_boundaries( + #[values( + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct + )] + metric: Metric, + #[case] leader_count: usize, + #[case] fanout: usize, + ) { + // Given + let (dots, point_norms, leader_norms) = lane_boundary_fixture(metric, leader_count); + let input = partition_input(&dots, 2, leader_count, &point_norms, &leader_norms); + let expected_assignments = reference_assignments(metric, input, fanout); + + // When + let actual_assignments = rank_partition_leaders(metric, input, fanout); + + // Then + assert_eq!(actual_assignments, expected_assignments); + } + + #[test] + fn equal_l2_scores_keep_sampled_leader_order() { + // Given + let dots = [0.0, 0.0, 0.0, 0.0]; + let leader_squared_norms = [1.0, 1.0, 1.0, 1.0]; + let expected_sampled_leader_order = [0, 1]; + + // When + let actual_assignments = rank_partition_leaders( + Metric::L2, + partition_input(&dots, 1, 4, &[], &leader_squared_norms), + 2, + ); + + // Then + assert_eq!(actual_assignments, expected_sampled_leader_order); + } + + #[test] + fn cosine_zero_norm_keeps_sampled_leader_order() { + // Given + let dots = [100.0, -100.0]; + let point_norms = [0.0]; + let leader_norms = [1.0, 1.0]; + let expected_sampled_leader_order = [0, 1]; + + // When + let actual_assignments = rank_partition_leaders( Metric::Cosine, - test_input(&[100.0, -100.0], 1, 2, &[0.0], &[1.0, 1.0]), + partition_input(&dots, 1, 2, &point_norms, &leader_norms), 2, - ), - [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_partition_kernel(Metric::InnerProduct, test_input(&dots, 1, 8, &[], &[]), 8), - [0, 1, 2, 3, 4, 5, 6, 7] - ); - } + // Then + assert_eq!(actual_assignments, expected_sampled_leader_order); + } + + #[test] + fn f32_max_score_is_still_a_rankable_leader() { + // Given + let mut dots = [0.0; 8]; + dots[7] = -f32::MAX; + let expected_all_leaders_in_scan_order = [0, 1, 2, 3, 4, 5, 6, 7]; + + // When + let actual_assignments = rank_partition_leaders( + Metric::InnerProduct, + partition_input(&dots, 1, 8, &[], &[]), + 8, + ); + + // Then + assert_eq!(actual_assignments, expected_all_leaders_in_scan_order); + } + + #[test] + fn nan_leader_does_not_displace_finite_leaders() { + // Given + let dots = [f32::NAN, 3.0, 2.0]; + let expected_finite_leaders = [1, 2]; - #[test] - fn ignores_nan_distances_without_displacing_finite_leaders() { - assert_eq!( - run_partition_kernel( + // When + let actual_assignments = rank_partition_leaders( Metric::InnerProduct, - test_input(&[f32::NAN, 3.0, 2.0], 1, 3, &[], &[]), + partition_input(&dots, 1, 3, &[], &[]), 2, - ), - [1, 2] - ); - } + ); - #[test] - fn accepts_empty_points_and_zero_fanout() { - assert!( - run_partition_kernel(Metric::InnerProduct, test_input(&[], 0, 3, &[], &[]), 2) - .is_empty() - ); - assert!( - run_partition_kernel( + // Then + assert_eq!(actual_assignments, expected_finite_leaders); + } + + #[test] + fn empty_point_matrix_produces_no_assignments() { + // Given + let dots = []; + let expected_no_assignments: [u32; 0] = []; + + // When + let actual_assignments = rank_partition_leaders( Metric::InnerProduct, - test_input(&[1.0, 2.0, 3.0], 1, 3, &[], &[]), + partition_input(&dots, 0, 3, &[], &[]), + 2, + ); + + // Then + assert_eq!(actual_assignments, expected_no_assignments); + } + + #[test] + fn zero_fanout_produces_no_assignments() { + // Given + let dots = [1.0, 2.0, 3.0]; + let expected_no_assignments: [u32; 0] = []; + + // When + let actual_assignments = rank_partition_leaders( + Metric::InnerProduct, + partition_input(&dots, 1, 3, &[], &[]), 0, - ) - .is_empty() - ); + ); + + // Then + assert_eq!(actual_assignments, expected_no_assignments); + } } }