From ae6886ba9a1b1e1886ba0f66191acf2f1a4d3126 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 001/229] 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 8da96e4b786db19d2dc9b96fe50c5f0b66ec5997 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 002/229] 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 0c862c6a08..4ea0f6e561 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.56.0" diff --git a/Cargo.toml b/Cargo.toml index 59e02da4ba..670db975eb 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.56.0" } # Algorithm diskann = { path = "diskann", version = "0.56.0" } +diskann-pipnn = { path = "diskann-pipnn", version = "0.56.0" } # Providers diskann-providers = { path = "diskann-providers", default-features = false, version = "0.56.0" } diskann-inmem = { path = "diskann-inmem", default-features = false, version = "0.56.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 15d8365d5b397c8f2b08cf92b7ba95b18e56ee5c 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 003/229] 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 ab6989c53f..792e17bc7e 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 5b56fdae7f16927af31fd7f769ff56ef10af9709 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 004/229] 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 846bf597c2dfedb74f39f758835ed5aacf2e0cbc 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 005/229] 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 35f65ca85660942d70b08d04e374a28409fcc9b7 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 006/229] 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 b25d701c90aa62e8d10e2e3a67f2ba1be0c837be 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 007/229] 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 7ba81b6779f633d4d03f7efefd3028863ed90256 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 008/229] 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 0090fead171c4df67c6ddd6494d9a8fed48c9886 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 009/229] 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 03277f649b8c5d7b0ff8b30c4a1138f5bbec01fc 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 010/229] 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 218cc380d4d698cf2b50b8cab67a1950a012516e 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 011/229] 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 b7cd61dc3815d41c273f77e7db50b7436f280995 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 012/229] 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 71539697e708fee0ab40a45fd6c55d1cce7b253e 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 013/229] 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 4ea0f6e561..9aea8d728e 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 ad2d7dc34799d1a4b9eddf42585abdfd58385b4c 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 014/229] 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 7025c22e6ce8d86f70bf4f56813b1e9acdfd511e 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 015/229] 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 d16d9468c40a8e6d88f4cb653a356407605c1fe1 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 016/229] 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 da45d8ebd2914e71c38fe08fab6347d9824052bc 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 017/229] 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 699d11b3b491d5545df9991a0224d2f42f59a0bd 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 018/229] 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 400e91ad697f43e48d5eabbfb0b7dc0fe0f8b0b9 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 019/229] 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 afe1d386ff42c3d398d63d35aad5173fdb497117 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 020/229] 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 d4a1230be3f318018a09b70f86d8b92ce76e1f27 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 021/229] 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 85ec4cc8314cec17d49a5a77bccb75477d85a890 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 022/229] 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 792e17bc7e..049719fe91 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 9aea8d728e..0c862c6a08 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.56.0" diff --git a/Cargo.toml b/Cargo.toml index 670db975eb..59e02da4ba 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.56.0" } # Algorithm diskann = { path = "diskann", version = "0.56.0" } -diskann-pipnn = { path = "diskann-pipnn", version = "0.56.0" } # Providers diskann-providers = { path = "diskann-providers", default-features = false, version = "0.56.0" } diskann-inmem = { path = "diskann-inmem", default-features = false, version = "0.56.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 e3d20a67ceb0cfa8f4d64b2f833faf3f864a7ad7 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 023/229] 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 1bf66b28b8507a9ae65e9e22d49b3abdf92dabde 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 024/229] 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 44fb69a55722109da5ed04831f5bd1295e50ab75 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 025/229] 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 0c862c6a08..fd93eef60a 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 318f8ea845e8fa3e298b33a5a0583103907307a4 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 026/229] 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 049719fe91..da5d553650 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 c13226def11646c9fa7351edbac1bba79ef635d9 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 027/229] 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 22ca49274ff5d62dbe8bcb28de5729fb426d362d 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 028/229] 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 da5d553650..096b567979 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 db4b71179bccbada6577f830fed3e170f8221c31 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 029/229] 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 096b567979..f857d27a17 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 fd93eef60a..0c862c6a08 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 6dfc07b5c5ccfecffb3e800ee173c915ffb7e4f8 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 030/229] 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 8ce46fcc7ee315f57e0bc5129a342eaf4df217d2 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 031/229] 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 1d49dab7b139ddfd2eaa2af8e221dac2329673e1 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 032/229] 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 ad19c826d58a57a21b2230efb50d27827566c5fb 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 033/229] 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 9628de57b27087fc88b86166f1120a64c418fb1c 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 034/229] 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 88d06c1a060fa6446fbaf1e528db97969c87b0f3 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 035/229] 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 6f23dee774d6e8812b42e324198eeab19f025ecd 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 036/229] 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 88f6f7286760d9406111022e77de2e7d89cbd77e 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 037/229] 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 815c2a4a0dfa8f8faa6d2d491bdb55a35dff8cca 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 038/229] 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 bf989e2a731509a74b3e4dbcd37c796b877b06de 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 039/229] 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 51684ad4cac330848b9b9b0a792e7a86a7f7b860 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 040/229] 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 98be0dccb59bab2cfa4f99ef2461e99d553710ae 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 041/229] 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 96ceb24d644641c858da39942d6e0ac218fbfe8b 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 042/229] 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 465a96be91c114bbbb856a5fed9131f683a5fe1a 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 043/229] 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 32d6524a83971b742bb6391777961de89655a04f 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 044/229] 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 038c0697f72d189bb0b3a340b5f5b735ed0ebd9c 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 045/229] 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 df9061aa725625522117c23a6e055245bcb1edaf 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 046/229] 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 aa948a66259808adfb9260fe789802a1d57a5320 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 047/229] 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 163cff77b3dd1808e6a472aba6f7564ca71e35ee 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 048/229] 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 f94581792289012a8f225357c4bfc128c03f3363 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 049/229] 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 6d0dc2b78ebacee0bbeeb71a55a4bdf462a266f1 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 050/229] 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 b7e4db796f611c28611a623e077b72d818fb87a2 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 051/229] 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 bf7cead073157f69642a5be78d092ca6a7c0edcf 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 052/229] 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 fd0f4e2dab442bd5cd7a3357857efec13725af53 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 053/229] 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 6671cc858f75582314a891d176b0bacc140f629b 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 054/229] 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 b40d54ce44fa147407f712aa635af5249953e726 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 055/229] 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 181338809a82ccd55537d3f0adfc91df04985389 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 056/229] 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 fe5a29c574fdf0bc0c99025b3211d3998a1345af 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 057/229] 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 e7074379021397192f4259c816503aabb51207ce 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 058/229] 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 a1bdc73502697c55ed78d735439c6639373979a9 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 059/229] 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 1f8ae265f6f3183922747935e50142f6bc6082af 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 060/229] 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 617e74641daf5053dade42efd157ce7524c515c2 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 061/229] 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 05ad4da337fccd1e9bda571ba168f21436a71a0c 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 062/229] 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 dc5cdba72338121da6a66fa834ff8db688b66f1b 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 063/229] 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 e32b0a3f7247fd9719921292c1f88ad9cc3bdc30 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 064/229] 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 00b69ce0184e630e440d8e3282c7f14d1cd7c653 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 065/229] 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 21885e69a9ed2d7f02cc327060d47a9c0420da56 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 066/229] 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 803f9816dfd0774a216bf22f2f251d1a7ca3ea76 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 067/229] 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 b28d67a70aa02634bc7d4821adf5e2ed690b083f 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 068/229] 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 afbb5769efc2d457186cafa818dcc20e4fce350d 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 069/229] 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 17aea03a35d88f937b703d5e015c7ec0cfcd2506 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 070/229] 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 0dccf08daea5b863b91c939f661c230645169202 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 071/229] 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 a3396a51b88b6e37ef0a95f1f8a13469cab41fe0 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 072/229] 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 9d4342af273aaace2106a0a337bab256a665f15c 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 073/229] 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 52bc12af6ff5ad4111027dd1eb24ac8227c3e4d0 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 074/229] 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 cc82f9c20f477226eafbcdc25cfaa47933299086 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 075/229] 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 09704bda3abe43247ba739e23b3ceab7f87a17f6 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 076/229] 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 0c862c6a08..b9b05d26fe 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 38cd3b51dd6b24b977b276c7d4d63bbdb0263633 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 077/229] 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 b946acefafcc1416f2cb15ebaeab1095933377a0 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 078/229] 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 52973f5487c37e8b985c53300c4e1cea7c12097e 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 079/229] 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 f55a5d0a94e409de0f7f98b831e242ed6f05d268 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 080/229] 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 b9b05d26fe..85a46b7e09 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); + } } } From 2de5c8641a21f6bf724c1aca053fd0954cc5806e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 24 Aug 2026 08:48:03 +0000 Subject: [PATCH 081/229] test(pipnn): use semantic numerical fixtures Derive Gram, dot, and norm inputs from source vectors and name each test by the ranking behavior it protects. --- diskann/src/graph/pipnn/kernel_metric.rs | 22 +- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 51 +++-- .../graph/pipnn/kernel_metric/partition.rs | 14 +- diskann/src/graph/pipnn/leaf_kernel.rs | 207 ++++++++++-------- diskann/src/graph/pipnn/partition_kernel.rs | 65 ++++-- 5 files changed, 215 insertions(+), 144 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 91f40fd69b..365cbc10b1 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -69,65 +69,65 @@ mod tests { use super::cosine_distance_single; #[test] - fn cosine_zero_source_norm_produces_unit_distance() { + fn 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); + let actual_distance = cosine_distance_single(0.0, source_norm, 2.0); // Then assert_eq!(actual_distance, expected_one_minus_zero_similarity); } #[test] - fn cosine_zero_target_norm_produces_unit_distance() { + fn 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); + let actual_distance = cosine_distance_single(0.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() { + fn similarity_above_one_clamps_to_zero_distance() { // Given - let dot_product_above_valid_similarity = 5.0; + let dot_product_just_above_norm_product = 4.000_001; 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); + cosine_distance_single(dot_product_just_above_norm_product, 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() { + fn similarity_below_negative_one_clamps_to_distance_two() { // Given - let dot_product_below_valid_similarity = -5.0; + let dot_product_just_below_negative_norm_product = -4.000_001; 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); + cosine_distance_single(dot_product_just_below_negative_norm_product, 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() { + fn nan_similarity_clamps_to_zero_distance() { // Given let nan_dot_product = f32::NAN; // `f32::min` keeps its finite operand when the other operand is NaN. diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 844e84b401..cb2c7d5502 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -226,38 +226,45 @@ mod tests { use super::*; #[test] - fn l2_leaf_norms_equal_the_gram_diagonal() { + fn leaf_norms_equal_point_self_dot_products_with_l2() { // 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 first_point = [2.0_f32, 1.0]; + let second_point = [1.0_f32, 3.0]; + let first_self_dot = first_point[0] * first_point[0] + first_point[1] * first_point[1]; + let cross_dot = first_point[0] * second_point[0] + first_point[1] * second_point[1]; + let second_self_dot = + second_point[0] * second_point[0] + second_point[1] * second_point[1]; + let gram_values = [first_self_dot, cross_dot, cross_dot, second_self_dot]; + let gram = MatrixView::try_from(&gram_values[..], 2, 2).unwrap(); + let expected_point_self_dots = [first_self_dot, second_self_dot]; let mut actual_norms = Vec::new(); // When - L2::prepare_leaf_norms(lower_gram, &mut actual_norms); + L2::prepare_leaf_norms(gram, &mut actual_norms); // Then - assert_eq!(actual_norms, expected_gram_diagonal); + assert_eq!(actual_norms, expected_point_self_dots); } #[test] - fn cosine_leaf_norms_equal_square_roots_of_the_gram_diagonal() { + fn leaf_norms_equal_point_l2_norms_with_cosine() { // 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 first_point = [2.0_f32, 1.0]; + let second_point = [1.0_f32, 3.0]; + let first_self_dot = first_point[0] * first_point[0] + first_point[1] * first_point[1]; + let cross_dot = first_point[0] * second_point[0] + first_point[1] * second_point[1]; + let second_self_dot = + second_point[0] * second_point[0] + second_point[1] * second_point[1]; + let gram_values = [first_self_dot, cross_dot, cross_dot, second_self_dot]; + let gram = MatrixView::try_from(&gram_values[..], 2, 2).unwrap(); + let expected_point_l2_norms = [first_self_dot.sqrt(), second_self_dot.sqrt()]; let mut actual_norms = Vec::new(); // When - Cosine::prepare_leaf_norms(lower_gram, &mut actual_norms); + Cosine::prepare_leaf_norms(gram, &mut actual_norms); // Then - assert_eq!(actual_norms, expected_square_roots_of_diagonal); + assert_eq!(actual_norms, expected_point_l2_norms); } } @@ -265,7 +272,7 @@ mod tests { use super::*; #[test] - fn l2_distance_equals_squared_norm_sum_minus_twice_the_dot_product() { + fn distance_equals_squared_norm_sum_minus_twice_the_dot_product_with_l2() { // Given let source_squared_norm = 4.0; let target_squared_norm = 9.0; @@ -283,7 +290,7 @@ mod tests { } #[test] - fn l2_clamps_negative_roundoff_to_zero() { + fn negative_roundoff_is_clamped_to_zero_with_l2() { // Given let source_squared_norm = 1.0; let target_squared_norm = 1.0; @@ -304,7 +311,7 @@ mod tests { } #[test] - fn cosine_distance_equals_one_minus_dot_over_norm_product() { + fn distance_equals_one_minus_dot_over_norm_product_with_cosine() { // Given let source_norm = 2.0; let target_norm = 4.0; @@ -320,7 +327,7 @@ mod tests { } #[test] - fn normalized_cosine_distance_is_one_minus_the_dot_product() { + fn distance_equals_one_minus_the_dot_product_with_normalized_cosine() { // Given let dot_product = 0.25; let expected_one_minus_dot = 1.0 - dot_product; @@ -333,7 +340,7 @@ mod tests { } #[test] - fn inner_product_distance_is_the_negative_dot_product() { + fn distance_equals_the_negative_dot_product_with_inner_product() { // Given let dot_product = 3.0; let expected_negative_dot = -dot_product; diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 05b54fe9b3..b5ab332ffc 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -264,7 +264,7 @@ mod tests { use super::*; #[test] - fn l2_leader_norm_is_the_sum_of_squared_components() { + fn leader_norm_equals_the_sum_of_squared_components_with_l2() { // Given let first_leader = [1.0_f32, 2.0]; let second_leader = [3.0_f32, 4.0]; @@ -289,7 +289,7 @@ mod tests { } #[test] - fn cosine_leader_norm_is_the_square_root_of_the_squared_component_sum() { + fn leader_norm_equals_the_l2_norm_with_cosine() { // Given let first_leader = [1.0_f32, 2.0]; let second_leader = [3.0_f32, 4.0]; @@ -315,7 +315,7 @@ mod tests { } #[test] - fn cosine_point_norm_is_the_square_root_of_the_squared_component_sum() { + fn point_norm_equals_the_l2_norm_with_cosine() { // Given let first_point = [1.0_f32, 2.0]; let second_point = [3.0_f32, 4.0]; @@ -343,7 +343,7 @@ mod tests { use super::*; #[test] - fn l2_ranking_equals_leader_squared_norm_minus_twice_the_dot_product() { + fn ranking_equals_leader_squared_norm_minus_twice_the_dot_product_with_l2() { // Given let dot_product = 2.0; let leader_squared_norm = 9.0; @@ -357,7 +357,7 @@ mod tests { } #[test] - fn cosine_ranking_equals_one_minus_dot_over_norm_product() { + fn ranking_equals_one_minus_dot_over_norm_product_with_cosine() { // Given let dot_product = 4.0; let point_norm = 2.0; @@ -372,7 +372,7 @@ mod tests { } #[test] - fn normalized_cosine_ranking_is_one_minus_the_dot_product() { + fn ranking_equals_one_minus_the_dot_product_with_normalized_cosine() { // Given let dot_product = 0.25; let expected_one_minus_dot = 1.0 - dot_product; @@ -385,7 +385,7 @@ mod tests { } #[test] - fn inner_product_ranking_is_the_negative_dot_product() { + fn ranking_equals_the_negative_dot_product_with_inner_product() { // Given let dot_product = 3.0; let expected_negative_dot = -dot_product; diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 49a4426631..2e9903f025 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -566,26 +566,40 @@ mod tests { 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]; + /// Build a Gram matrix from points on the line `x = 1`. + /// + /// Normalized cosine receives unit vectors because that metric assumes + /// normalized source data. Other metrics receive the original vectors. + fn lane_boundary_gram_from_point_vectors(metric: Metric, points: usize) -> Vec { + let denominator = (points + 1) as f32; + let point_vectors: Vec<_> = (0..points) + .map(|point| { + let vector = [1.0, (point + 1) as f32 / denominator]; + if metric == Metric::CosineNormalized { + let norm = vector[0].hypot(vector[1]); + [vector[0] / norm, vector[1] / norm] + } else { + vector + } + }) + .collect(); + let mut gram = vec![0.0; 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; + for target in 0..points { + gram[source * points + target] = point_vectors[source][0] + * point_vectors[target][0] + + point_vectors[source][1] * point_vectors[target][1]; } } - dots + gram } - fn square_matrix_with_constant_diagonal(points: usize, diagonal: f32) -> Vec { - let mut values = vec![0.0; points * points]; + fn gram_with_uniform_self_dots(points: usize, self_dot: f32) -> Vec { + let mut gram = vec![0.0; points * points]; for point in 0..points { - values[point * points + point] = diagonal; + gram[point * points + point] = self_dot; } - values + gram } mod insert_neighbor_tests { @@ -821,7 +835,7 @@ mod tests { use super::*; #[test] - fn l2_pipeline_orders_neighbors_by_squared_distance() { + fn orders_neighbors_by_squared_distance_with_l2() { // Given let values = [0.0_f32, 1.0, 3.0, 10.0]; let points = MatrixView::try_from(&values[..], 4, 1).unwrap(); @@ -922,7 +936,7 @@ mod tests { #[case] requested_k: usize, ) { // Given - let dots = index_distance_lower_gram(point_count); + let dots = lane_boundary_gram_from_point_vectors(metric, point_count); let expected_neighbors = reference_neighbors(metric, &dots, point_count, requested_k); // When @@ -932,99 +946,111 @@ mod tests { 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; + const POINT_ZERO: [f32; 2] = [1.0, 0.0]; + const POINT_ONE: [f32; 2] = [0.0, 1.0]; + const POINT_TWO: [f32; 2] = [0.6, 0.8]; - #[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 point_dot(left: [f32; 2], right: [f32; 2]) -> f32 { + left[0] * right[0] + left[1] * right[1] + } - fn rank_three_point_fixture(metric: Metric) -> Vec { - rank_neighbors(metric, &THREE_POINT_LOWER_GRAM, 3, 1).1 + fn three_unit_point_gram() -> [f32; 9] { + let points = [POINT_ZERO, POINT_ONE, POINT_TWO]; + std::array::from_fn(|index| { + let source = index / points.len(); + let target = index % points.len(); + point_dot(points[source], points[target]) + }) + } + + fn rank_three_unit_points(metric: Metric) -> Vec { + rank_neighbors(metric, &three_unit_point_gram(), 3, 1).1 } #[test] - fn l2_selects_the_nearest_target_from_the_lower_gram_triangle() { + fn selects_the_smallest_squared_distance_neighbor_for_each_point_with_l2() { // Given + let point_zero_two_distance = point_dot(POINT_ZERO, POINT_ZERO) + + point_dot(POINT_TWO, POINT_TWO) + - 2.0 * point_dot(POINT_ZERO, POINT_TWO); + let point_one_two_distance = point_dot(POINT_ONE, POINT_ONE) + + point_dot(POINT_TWO, POINT_TWO) + - 2.0 * point_dot(POINT_ONE, POINT_TWO); 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), + LeafNeighbor::new(2, point_zero_two_distance), + LeafNeighbor::new(2, point_one_two_distance), + LeafNeighbor::new(1, point_one_two_distance), ]; // When - let actual_neighbors = rank_three_point_fixture(Metric::L2); + let actual_neighbors = rank_three_unit_points(Metric::L2); // Then assert_eq!(actual_neighbors, expected_neighbors); } #[test] - fn cosine_selects_the_nearest_target_from_the_lower_gram_triangle() { + fn selects_the_highest_similarity_neighbor_for_each_point_with_cosine() { // 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), + LeafNeighbor::new(2, 1.0 - point_dot(POINT_ZERO, POINT_TWO)), + LeafNeighbor::new(2, 1.0 - point_dot(POINT_ONE, POINT_TWO)), + LeafNeighbor::new(1, 1.0 - point_dot(POINT_ONE, POINT_TWO)), ]; // When - let actual_neighbors = rank_three_point_fixture(Metric::Cosine); + let actual_neighbors = rank_three_unit_points(Metric::Cosine); // Then assert_eq!(actual_neighbors, expected_neighbors); } #[test] - fn normalized_cosine_selects_the_largest_dot_product() { + fn selects_the_highest_dot_product_neighbor_for_each_point_with_normalized_cosine() { // 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), + LeafNeighbor::new(2, 1.0 - point_dot(POINT_ZERO, POINT_TWO)), + LeafNeighbor::new(2, 1.0 - point_dot(POINT_ONE, POINT_TWO)), + LeafNeighbor::new(1, 1.0 - point_dot(POINT_ONE, POINT_TWO)), ]; // When - let actual_neighbors = rank_three_point_fixture(Metric::CosineNormalized); + let actual_neighbors = rank_three_unit_points(Metric::CosineNormalized); // Then assert_eq!(actual_neighbors, expected_neighbors); } #[test] - fn inner_product_selects_the_largest_dot_product() { + fn selects_the_highest_dot_product_neighbor_for_each_point_with_inner_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), + LeafNeighbor::new(2, -point_dot(POINT_ZERO, POINT_TWO)), + LeafNeighbor::new(2, -point_dot(POINT_ONE, POINT_TWO)), + LeafNeighbor::new(1, -point_dot(POINT_ONE, POINT_TWO)), ]; // When - let actual_neighbors = rank_three_point_fixture(Metric::InnerProduct); + let actual_neighbors = rank_three_unit_points(Metric::InnerProduct); // Then assert_eq!(actual_neighbors, expected_neighbors); } #[test] - fn equal_l2_distances_keep_target_scan_order() { + fn equal_distances_keep_target_scan_order_with_l2() { // 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; + // This is the Gram matrix of four orthogonal unit vectors. #[rustfmt::skip] - let dots = [ - 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 gram = [ + unit_squared_norm, tied_dot_product, tied_dot_product, tied_dot_product, + tied_dot_product, unit_squared_norm, tied_dot_product, tied_dot_product, + tied_dot_product, tied_dot_product, unit_squared_norm, tied_dot_product, + 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), @@ -1037,22 +1063,22 @@ mod tests { ]; // When - let actual_neighbors = rank_neighbors(Metric::L2, &dots, 4, 2).1; + let actual_neighbors = rank_neighbors(Metric::L2, &gram, 4, 2).1; // Then assert_eq!(actual_neighbors, expected_neighbors_in_scan_order); } #[test] - fn scalar_l2_distance_stays_finite_when_twice_the_dot_product_overflows() { + fn scalar_distance_stays_finite_when_twice_the_dot_product_overflows_with_l2() { // 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]; + let gram = [f32::MAX, dot_product, dot_product, f32::MAX]; // When - let actual_neighbors = rank_neighbors(Metric::L2, &dots, 2, 1).1; + let actual_neighbors = rank_neighbors(Metric::L2, &gram, 2, 1).1; // Then assert!(unfused_twice_dot_product.is_infinite()); @@ -1064,18 +1090,19 @@ mod tests { } #[test] - fn simd_l2_distance_stays_finite_when_twice_the_dot_product_overflows() { + fn simd_distance_stays_finite_when_twice_the_dot_product_overflows_with_l2() { // 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; + let mut gram = gram_with_uniform_self_dots(points, f32::MAX); + gram[16 * points] = dot_product; + gram[16] = dot_product; // When - let actual_neighbors = rank_neighbors(Metric::L2, &dots, points, 1).1; + let actual_neighbors = rank_neighbors(Metric::L2, &gram, points, 1).1; // Then assert!(unfused_twice_dot_product.is_infinite()); @@ -1087,33 +1114,36 @@ mod tests { } #[test] - fn cosine_zero_norm_produces_unit_distance() { + fn zero_norm_produces_unit_distance_with_cosine() { // Given + // This is the Gram matrix of one zero vector and two orthogonal unit vectors. #[rustfmt::skip] - let dots = [ - 0.0, 99.0, 99.0, - 0.0, 1.0, 99.0, - 0.0, 0.0, 1.0, - ]; + let gram = [ + 0.0, 0.0, 0.0, + 0.0, 1.0, 0.0, + 0.0, 0.0, 1.0, + ]; 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; + let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, 3, 2).1; // Then assert_eq!(&actual_neighbors[..2], &expected_zero_norm_neighbors); } #[test] - fn cosine_similarity_above_one_clamps_to_zero_distance() { + fn similarity_above_one_clamps_to_zero_distance_with_cosine() { // Given - let dots = [1.0, 0.0, 2.0, 1.0]; + // A small excess models dot-product roundoff above cosine similarity one. + let rounded_dot_product = 1.000_001; + let gram = [1.0, rounded_dot_product, rounded_dot_product, 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; + let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, 2, 1).1; // Then assert_eq!( @@ -1123,14 +1153,16 @@ mod tests { } #[test] - fn cosine_similarity_below_negative_one_clamps_to_distance_two() { + fn similarity_below_negative_one_clamps_to_distance_two_with_cosine() { // Given - let dots = [1.0, 0.0, -2.0, 1.0]; + // A small excess models dot-product roundoff below cosine similarity minus one. + let rounded_dot_product = -1.000_001; + let gram = [1.0, rounded_dot_product, rounded_dot_product, 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; + let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, 2, 1).1; // Then assert_eq!( @@ -1140,14 +1172,15 @@ mod tests { } #[test] - fn cosine_subnormal_norm_is_treated_as_zero() { + fn subnormal_norm_is_treated_as_zero_with_cosine() { // Given - let dots = [f32::MIN_POSITIVE / 2.0, 0.0, 1.0, 1.0]; + let subnormal_self_dot = f32::MIN_POSITIVE / 2.0; + let gram = [subnormal_self_dot, 0.0, 0.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; + let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, 2, 1).1; // Then assert_eq!( @@ -1162,12 +1195,13 @@ mod tests { 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; + let mut gram = gram_with_uniform_self_dots(points, f32::MAX); + gram[3 * points] = -f32::MAX; + gram[3] = -f32::MAX; // When let (actual_leaf_k, actual_neighbors) = - rank_neighbors(Metric::InnerProduct, &dots, points, expected_leaf_k); + rank_neighbors(Metric::InnerProduct, &gram, points, expected_leaf_k); // Then assert_eq!(actual_leaf_k, expected_leaf_k); @@ -1180,11 +1214,11 @@ mod tests { #[test] fn scalar_nan_distance_leaves_the_neighbor_slot_unassigned() { // Given - let dots = [1.0, 0.0, f32::NAN, 1.0]; + let gram = [1.0, f32::NAN, f32::NAN, 1.0]; let expected_unassigned_neighbors = [LeafNeighbor::default(), LeafNeighbor::default()]; // When - let actual_neighbors = rank_neighbors(Metric::CosineNormalized, &dots, 2, 1).1; + let actual_neighbors = rank_neighbors(Metric::CosineNormalized, &gram, 2, 1).1; // Then assert_eq!(actual_neighbors, expected_unassigned_neighbors); @@ -1194,12 +1228,13 @@ mod tests { 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 mut gram = gram_with_uniform_self_dots(points, 1.0); + gram[16 * points] = f32::NAN; + gram[16] = f32::NAN; let expected_finite_neighbor = LeafNeighbor::new(1, 1.0); // When - let actual_neighbors = rank_neighbors(Metric::CosineNormalized, &dots, points, 1).1; + let actual_neighbors = rank_neighbors(Metric::CosineNormalized, &gram, points, 1).1; // Then assert_eq!(actual_neighbors[16], expected_finite_neighbor); diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index ba03067cb6..275053f832 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -364,26 +364,52 @@ mod tests { output } - /// Build two score rows whose preferred leader direction is opposite. - /// The fractional scores cross complete SIMD groups and scalar tails. - fn lane_boundary_fixture( + /// Build ranking input from two axis points and leaders between those axes. + /// + /// The first point prefers early leaders. The second point prefers late leaders. + fn lane_boundary_input_from_point_and_leader_vectors( 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 points = [[1.0_f32, 0.0], [0.0, 1.0]]; + let denominator = (leader_count + 1) as f32; + let leaders: Vec<_> = (0..leader_count) + .map(|leader| { + let second_component = (leader + 1) as f32 / denominator; + let vector = [1.0 - second_component, second_component]; + if metric == Metric::CosineNormalized { + let norm = vector[0].hypot(vector[1]); + [vector[0] / norm, vector[1] / norm] + } else { + vector + } + }) + .collect(); + let dots = points + .iter() + .flat_map(|point| { + leaders + .iter() + .map(|leader| point[0] * leader[0] + point[1] * leader[1]) + }) + .collect(); let point_norms = if metric == Metric::Cosine { - vec![1.0; 2] + points + .iter() + .map(|point| point[0].hypot(point[1])) + .collect() } else { Vec::new() }; let leader_norms = match metric { - Metric::L2 | Metric::Cosine => vec![1.0; leader_count], + Metric::L2 => leaders + .iter() + .map(|leader| leader[0] * leader[0] + leader[1] * leader[1]) + .collect(), + Metric::Cosine => leaders + .iter() + .map(|leader| leader[0].hypot(leader[1])) + .collect(), Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), }; (dots, point_norms, leader_norms) @@ -428,7 +454,7 @@ mod tests { use super::*; #[test] - fn cosine_pipeline_assigns_each_point_to_its_nearest_leaders() { + fn assigns_each_point_to_highest_similarity_leaders_with_cosine() { // Given let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; let leaders = PreparedLeaders::::new( @@ -532,7 +558,8 @@ mod tests { #[case] fanout: usize, ) { // Given - let (dots, point_norms, leader_norms) = lane_boundary_fixture(metric, leader_count); + let (dots, point_norms, leader_norms) = + lane_boundary_input_from_point_and_leader_vectors(metric, leader_count); let input = partition_input(&dots, 2, leader_count, &point_norms, &leader_norms); let expected_assignments = reference_assignments(metric, input, fanout); @@ -544,7 +571,7 @@ mod tests { } #[test] - fn equal_l2_scores_keep_sampled_leader_order() { + fn equal_scores_keep_sampled_leader_order_with_l2() { // Given let dots = [0.0, 0.0, 0.0, 0.0]; let leader_squared_norms = [1.0, 1.0, 1.0, 1.0]; @@ -562,9 +589,9 @@ mod tests { } #[test] - fn cosine_zero_norm_keeps_sampled_leader_order() { + fn zero_norm_keeps_sampled_leader_order_with_cosine() { // Given - let dots = [100.0, -100.0]; + let dots = [0.0, 0.0]; let point_norms = [0.0]; let leader_norms = [1.0, 1.0]; let expected_sampled_leader_order = [0, 1]; @@ -583,8 +610,10 @@ mod tests { #[test] fn f32_max_score_is_still_a_rankable_leader() { // Given + let maximum_rankable_score = f32::MAX; + let dot_product_that_produces_it = -maximum_rankable_score; let mut dots = [0.0; 8]; - dots[7] = -f32::MAX; + dots[7] = dot_product_that_produces_it; let expected_all_leaders_in_scan_order = [0, 1, 2, 3, 4, 5, 6, 7]; // When From d7865522d060544fb048eb628353ae16708fcb8b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:05:26 +0000 Subject: [PATCH 082/229] refactor(wide): keep lane arrays read-only --- diskann-wide/src/traits.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-wide/src/traits.rs b/diskann-wide/src/traits.rs index a233622615..da66fe9167 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: AsRef<[T]>; } /// Map scalar + lengths to arrays. From 17ea40f599c0f5f73ae86beedc2b405769db3e07 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:15:11 +0000 Subject: [PATCH 083/229] test(pipnn): clarify norm test names --- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 4 ++-- diskann/src/graph/pipnn/kernel_metric/partition.rs | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index cb2c7d5502..6823d8fdce 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -226,7 +226,7 @@ mod tests { use super::*; #[test] - fn leaf_norms_equal_point_self_dot_products_with_l2() { + fn returns_the_squared_euclidean_norm_of_each_point_for_l2() { // Given let first_point = [2.0_f32, 1.0]; let second_point = [1.0_f32, 3.0]; @@ -247,7 +247,7 @@ mod tests { } #[test] - fn leaf_norms_equal_point_l2_norms_with_cosine() { + fn returns_the_euclidean_norm_of_each_point_for_cosine() { // Given let first_point = [2.0_f32, 1.0]; let second_point = [1.0_f32, 3.0]; diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index b5ab332ffc..5df9171ad4 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -264,7 +264,7 @@ mod tests { use super::*; #[test] - fn leader_norm_equals_the_sum_of_squared_components_with_l2() { + fn returns_the_squared_euclidean_norm_of_each_leader_for_l2() { // Given let first_leader = [1.0_f32, 2.0]; let second_leader = [3.0_f32, 4.0]; @@ -289,7 +289,7 @@ mod tests { } #[test] - fn leader_norm_equals_the_l2_norm_with_cosine() { + fn returns_the_euclidean_norm_of_each_leader_for_cosine() { // Given let first_leader = [1.0_f32, 2.0]; let second_leader = [3.0_f32, 4.0]; @@ -315,7 +315,7 @@ mod tests { } #[test] - fn point_norm_equals_the_l2_norm_with_cosine() { + fn returns_the_euclidean_norm_of_each_point_for_cosine() { // Given let first_point = [1.0_f32, 2.0]; let second_point = [3.0_f32, 4.0]; From a60773cb08638716125c5c10e12bd021f730d29b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 24 Aug 2026 10:46:10 +0000 Subject: [PATCH 084/229] test(pipnn): cover SIMD and insertion contracts Add direct SIMD metric coverage and make the caller-owned insertion eligibility precondition explicit in names, docs, and tests. --- diskann/src/graph/pipnn/kernel_metric.rs | 84 ++++++++- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 81 ++++++++ .../graph/pipnn/kernel_metric/partition.rs | 81 ++++++++ diskann/src/graph/pipnn/leaf_kernel.rs | 173 ++++++++++++++---- diskann/src/graph/pipnn/partition_kernel.rs | 20 ++ 5 files changed, 399 insertions(+), 40 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 365cbc10b1..1a5a47c975 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -63,7 +63,28 @@ pub(super) fn cosine_distance_single(dot: f32, source_norm: f32, target_norm: f3 #[cfg(test)] mod tests { - use super::cosine_distance_single; + use super::super::simd::PiPNNSIMDSchema; + use super::{cosine_distance_simd, cosine_distance_single}; + use diskann_wide::{ARCH, SIMDVector, arch::Current}; + + fn simd_cosine_distances( + dot_products: [f32; 16], + source_norms: [f32; 16], + target_norms: [f32; 16], + ) -> [f32; 16] { + type TestVector = ::LeafScore; + assert_eq!(TestVector::LANES, 16); + // SAFETY: each input and output contains exactly one complete SIMD vector. + unsafe { + let dots = TestVector::load_simd(ARCH, dot_products.as_ptr()); + let sources = TestVector::load_simd(ARCH, source_norms.as_ptr()); + let targets = TestVector::load_simd(ARCH, target_norms.as_ptr()); + let distances = cosine_distance_simd(ARCH, dots, sources, targets); + let mut output = [0.0; 16]; + distances.store_simd(output.as_mut_ptr()); + output + } + } mod cosine_distance_single_tests { use super::cosine_distance_single; @@ -141,4 +162,65 @@ mod tests { assert_eq!(actual_distance, expected_one_minus_selected_operand); } } + + mod cosine_distance_simd_tests { + use super::simd_cosine_distances; + + #[test] + fn zero_source_norm_produces_unit_distance_in_every_lane() { + let actual_distances = simd_cosine_distances([0.0; 16], [0.0; 16], [2.0; 16]); + assert_eq!(actual_distances, [1.0; 16]); + } + + #[test] + fn zero_target_norm_produces_unit_distance_in_every_lane() { + let actual_distances = simd_cosine_distances([0.0; 16], [2.0; 16], [0.0; 16]); + assert_eq!(actual_distances, [1.0; 16]); + } + + #[test] + fn subnormal_norm_produces_unit_distance_in_every_lane() { + let subnormal_norm = f32::MIN_POSITIVE.sqrt() / 2.0; + let actual_distances = + simd_cosine_distances([0.0; 16], [2.0; 16], [subnormal_norm; 16]); + assert_eq!(actual_distances, [1.0; 16]); + } + + #[test] + fn roundoff_past_similarity_bounds_is_clamped_in_every_lane() { + let dot_above_norm_product = 4.000_001; + let dot_below_negative_norm_product = -4.000_001; + let dot_products = [ + dot_above_norm_product, + dot_above_norm_product, + dot_above_norm_product, + dot_above_norm_product, + dot_above_norm_product, + dot_above_norm_product, + dot_above_norm_product, + dot_above_norm_product, + dot_below_negative_norm_product, + dot_below_negative_norm_product, + dot_below_negative_norm_product, + dot_below_negative_norm_product, + dot_below_negative_norm_product, + dot_below_negative_norm_product, + dot_below_negative_norm_product, + dot_below_negative_norm_product, + ]; + let expected_distances = [ + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, + ]; + + let actual_distances = simd_cosine_distances(dot_products, [2.0; 16], [2.0; 16]); + + assert_eq!(actual_distances, expected_distances); + } + + #[test] + fn nan_similarity_clamps_to_zero_distance_in_every_lane() { + let actual_distances = simd_cosine_distances([f32::NAN; 16], [1.0; 16], [1.0; 16]); + assert_eq!(actual_distances, [0.0; 16]); + } + } } diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 6823d8fdce..7a176c96af 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -221,6 +221,25 @@ impl LeafMetric for InnerProduct { )] mod tests { use super::*; + use diskann_wide::{ARCH, SIMDVector, arch::Current}; + + fn simd_distances( + norms: &[f32], + source: usize, + dot_products: [f32; 16], + first_target: usize, + ) -> [f32; 16] { + assert_eq!(M::Simd::::LANES, 16); + let source_norms = M::source_simd(ARCH, norms, source); + // SAFETY: the input and output contain exactly one complete SIMD vector. + unsafe { + let dots = M::Simd::::load_simd(ARCH, dot_products.as_ptr()); + let distances = M::distances_simd(ARCH, norms, source_norms, dots, first_target); + let mut output = [0.0; 16]; + distances.store_simd(output.as_mut_ptr()); + output + } + } mod prepare_leaf_norms_tests { use super::*; @@ -352,4 +371,66 @@ mod tests { assert_eq!(actual_distance, expected_negative_dot); } } + + mod distances_simd_tests { + use super::*; + + #[test] + fn every_lane_uses_squared_norm_sum_minus_twice_the_dot_product_with_l2() { + let source_squared_norm = 4.0; + let target_squared_norm = 9.0; + let dot_product = 6.0; + let expected_distance = source_squared_norm + target_squared_norm - 2.0 * dot_product; + let mut squared_norms = [target_squared_norm; 17]; + squared_norms[0] = source_squared_norm; + + let actual_distances = simd_distances::(&squared_norms, 0, [dot_product; 16], 1); + + assert_eq!(actual_distances, [expected_distance; 16]); + } + + #[test] + fn nan_dot_products_clamp_to_zero_in_every_lane_with_l2() { + let squared_norms = [1.0; 17]; + let expected_non_negative_distance = 0.0; + + let actual_distances = simd_distances::(&squared_norms, 0, [f32::NAN; 16], 1); + + assert_eq!(actual_distances, [expected_non_negative_distance; 16]); + } + + #[test] + fn every_lane_uses_one_minus_normalized_dot_with_cosine() { + let source_norm = 2.0; + let target_norm = 4.0; + let dot_product = 4.0; + let expected_distance = 1.0 - dot_product / (source_norm * target_norm); + let mut norms = [target_norm; 17]; + norms[0] = source_norm; + + let actual_distances = simd_distances::(&norms, 0, [dot_product; 16], 1); + + assert_eq!(actual_distances, [expected_distance; 16]); + } + + #[test] + fn every_lane_uses_one_minus_dot_product_with_normalized_cosine() { + let dot_product = 0.25; + let expected_distance = 1.0 - dot_product; + + let actual_distances = simd_distances::(&[], 0, [dot_product; 16], 0); + + assert_eq!(actual_distances, [expected_distance; 16]); + } + + #[test] + fn every_lane_uses_negative_dot_product_with_inner_product() { + let dot_product = 3.0; + let expected_distance = -dot_product; + + let actual_distances = simd_distances::(&[], 0, [dot_product; 16], 0); + + assert_eq!(actual_distances, [expected_distance; 16]); + } + } } diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 5df9171ad4..2946e43fed 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -245,6 +245,30 @@ impl PartitionMetric for InnerProduct { )] mod tests { use super::*; + use diskann_wide::{ARCH, SIMDVector, arch::Current}; + + fn simd_rankings( + point_norms: &[f32], + leader_norms: &[f32], + point: usize, + dot_products: [f32; 16], + first_leader: usize, + ) -> [f32; 16] { + assert_eq!(M::Simd::::LANES, 16); + let norms = PartitionNorms { + point_norms, + leader_norms, + }; + let point_norm = M::point_simd(ARCH, norms, point); + // SAFETY: the input and output contain exactly one complete SIMD vector. + unsafe { + let dots = M::Simd::::load_simd(ARCH, dot_products.as_ptr()); + let rankings = M::rankings_simd(ARCH, norms, point_norm, dots, first_leader); + let mut output = [0.0; 16]; + rankings.store_simd(output.as_mut_ptr()); + output + } + } fn rank_single_leader( dot_product: f32, @@ -397,4 +421,61 @@ mod tests { assert_eq!(actual_ranking, expected_negative_dot); } } + + mod rankings_simd_tests { + use super::*; + + #[test] + fn every_lane_uses_leader_squared_norm_minus_twice_the_dot_product_with_l2() { + let leader_squared_norm = 9.0; + let dot_product = 2.0; + let expected_ranking = leader_squared_norm - 2.0 * dot_product; + let actual_rankings = + simd_rankings::(&[], &[leader_squared_norm; 16], 0, [dot_product; 16], 0); + assert_eq!(actual_rankings, [expected_ranking; 16]); + } + + #[test] + fn every_lane_uses_one_minus_normalized_dot_with_cosine() { + let point_norm = 2.0; + let leader_norm = 4.0; + let dot_product = 4.0; + let expected_ranking = 1.0 - dot_product / (point_norm * leader_norm); + let actual_rankings = + simd_rankings::(&[point_norm], &[leader_norm; 16], 0, [dot_product; 16], 0); + assert_eq!(actual_rankings, [expected_ranking; 16]); + } + + #[test] + fn every_lane_uses_one_minus_dot_product_with_normalized_cosine() { + let dot_product = 0.25; + let expected_ranking = 1.0 - dot_product; + let actual_rankings = + simd_rankings::(&[], &[], 0, [dot_product; 16], 0); + assert_eq!(actual_rankings, [expected_ranking; 16]); + } + + #[test] + fn every_lane_uses_negative_dot_product_with_inner_product() { + let dot_product = 3.0; + let expected_ranking = -dot_product; + let actual_rankings = simd_rankings::(&[], &[], 0, [dot_product; 16], 0); + assert_eq!(actual_rankings, [expected_ranking; 16]); + } + + #[test] + fn every_lane_stays_finite_when_twice_the_dot_product_overflows_with_l2() { + let dot_product = f32::from_bits(f32::MAX.to_bits() - 1); + let leader_squared_norm = f32::MAX; + let unfused_twice_dot_product = 2.0 * dot_product; + let expected_fused_ranking = (-2.0_f32).mul_add(dot_product, leader_squared_norm); + + let actual_rankings = + simd_rankings::(&[], &[leader_squared_norm; 16], 0, [dot_product; 16], 0); + + assert!(unfused_twice_dot_product.is_infinite()); + assert!(expected_fused_ranking.is_finite()); + assert_eq!(actual_rankings, [expected_fused_ranking; 16]); + } + } } diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 2e9903f025..4a29fd389b 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -199,7 +199,7 @@ fn scan_fixed_width( { let (rows, _) = output.as_chunks_mut::(); scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { - insert_neighbor(&mut rows[source], target, distance) + insert_eligible_neighbor(&mut rows[source], target, distance) }); } @@ -217,7 +217,7 @@ fn scan_runtime_width( { scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { let first = source * width; - insert_neighbor(&mut output[first..first + width], target, distance) + insert_eligible_neighbor(&mut output[first..first + width], target, distance) }); } @@ -309,17 +309,19 @@ fn scan_point_pairs( } } -/// Insert one value that precedes the current last retained value. +/// Insert one value that the caller has already found eligible. /// -/// The caller checks eligibility before insertion. The returned value is the new -/// last retained value for the next eligibility check. +/// The caller must prove that `value` precedes the current last retained value. +/// This method intentionally does not repeat that check. A caller that violates +/// the precondition replaces a valid retained value and corrupts the top-k set. +/// The return value is the new last retained value. trait SortedInsert { - fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T; + fn insert_eligible_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 { + fn insert_eligible_sorted_by(&mut self, value: T, _precedes: impl Fn(T, T) -> bool) -> T { self[0] = value; value } @@ -327,7 +329,7 @@ impl SortedInsert for [T; 1] { impl SortedInsert for [T; 2] { #[inline(always)] - fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + fn insert_eligible_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { let first = self[0]; if precedes(value, first) { self[0] = value; @@ -342,7 +344,7 @@ impl SortedInsert for [T; 2] { impl SortedInsert for [T; 3] { #[inline(always)] - fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + fn insert_eligible_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; @@ -362,7 +364,7 @@ impl SortedInsert for [T; 3] { impl SortedInsert for [T] { #[inline(always)] - fn insert_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + fn insert_eligible_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]) { @@ -374,16 +376,18 @@ impl SortedInsert for [T] { } } -/// Insert one candidate that is nearer than the current farthest neighbor. +/// Insert one candidate that the caller has already found nearer than the current farthest. /// -/// Return the new farthest retained distance for the next candidate check. +/// This function intentionally does not reject an ineligible candidate. The pair +/// scan owns the eligibility check so it can filter SIMD lanes before insertion. +/// The return value is the new farthest retained distance. #[inline(always)] -fn insert_neighbor(neighbors: &mut R, target: u32, distance: f32) -> f32 +fn insert_eligible_neighbor(neighbors: &mut R, target: u32, distance: f32) -> f32 where R: SortedInsert + ?Sized, { neighbors - .insert_sorted_by( + .insert_eligible_sorted_by( LeafNeighbor::new(target, distance), |candidate, retained| candidate.distance < retained.distance, ) @@ -602,11 +606,11 @@ mod tests { gram } - mod insert_neighbor_tests { + mod insert_eligible_neighbor_tests { use super::*; #[test] - fn one_slot_insertion_replaces_the_retained_neighbor() { + fn nearer_candidate_replaces_the_only_retained_neighbor() { // Given let retained_neighbor = LeafNeighbor::new(1, 4.0); let nearer_candidate = LeafNeighbor::new(2, 2.0); @@ -614,7 +618,7 @@ mod tests { let mut actual_neighbors = [retained_neighbor]; // When - insert_neighbor( + insert_eligible_neighbor( &mut actual_neighbors, nearer_candidate.target, nearer_candidate.distance, @@ -625,7 +629,27 @@ mod tests { } #[test] - fn two_slot_insertion_places_a_nearer_candidate_first() { + fn direct_call_does_not_recheck_candidate_eligibility() { + // Given: deliberately bypass the pair scan's eligibility check. + let nearest = LeafNeighbor::new(1, 1.0); + let current_farthest = LeafNeighbor::new(2, 3.0); + let ineligible_farther_candidate = LeafNeighbor::new(3, 5.0); + let expected_unchecked_result = [nearest, ineligible_farther_candidate]; + let mut actual_neighbors = [nearest, current_farthest]; + + // When + insert_eligible_neighbor( + &mut actual_neighbors, + ineligible_farther_candidate.target, + ineligible_farther_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_unchecked_result); + } + + #[test] + fn nearer_candidate_moves_to_the_front_of_two_retained_neighbors() { // Given let nearest = LeafNeighbor::new(1, 1.0); let farthest = LeafNeighbor::new(2, 3.0); @@ -634,7 +658,7 @@ mod tests { let mut actual_neighbors = [nearest, farthest]; // When - insert_neighbor( + insert_eligible_neighbor( &mut actual_neighbors, nearer_candidate.target, nearer_candidate.distance, @@ -645,7 +669,7 @@ mod tests { } #[test] - fn two_slot_insertion_places_a_middle_distance_last() { + fn middle_distance_candidate_replaces_the_farther_of_two_neighbors() { // Given let nearest = LeafNeighbor::new(1, 1.0); let farthest = LeafNeighbor::new(2, 3.0); @@ -654,7 +678,7 @@ mod tests { let mut actual_neighbors = [nearest, farthest]; // When - insert_neighbor( + insert_eligible_neighbor( &mut actual_neighbors, eligible_candidate.target, eligible_candidate.distance, @@ -665,7 +689,7 @@ mod tests { } #[test] - fn three_slot_insertion_places_the_nearest_candidate_first() { + fn nearest_candidate_moves_to_the_front_of_three_retained_neighbors() { // Given let nearest = LeafNeighbor::new(1, 1.0); let middle = LeafNeighbor::new(2, 2.0); @@ -675,7 +699,7 @@ mod tests { let mut actual_neighbors = [nearest, middle, farthest]; // When - insert_neighbor( + insert_eligible_neighbor( &mut actual_neighbors, nearer_candidate.target, nearer_candidate.distance, @@ -686,7 +710,7 @@ mod tests { } #[test] - fn three_slot_insertion_places_a_middle_candidate_between_neighbors() { + fn middle_candidate_is_inserted_between_three_retained_neighbors() { // Given let nearest = LeafNeighbor::new(1, 1.0); let middle = LeafNeighbor::new(2, 2.0); @@ -696,7 +720,7 @@ mod tests { let mut actual_neighbors = [nearest, middle, farthest]; // When - insert_neighbor( + insert_eligible_neighbor( &mut actual_neighbors, middle_candidate.target, middle_candidate.distance, @@ -707,7 +731,7 @@ mod tests { } #[test] - fn three_slot_insertion_replaces_the_farthest_neighbor() { + fn closer_candidate_replaces_the_farthest_of_three_neighbors() { // Given let nearest = LeafNeighbor::new(1, 1.0); let middle = LeafNeighbor::new(2, 2.0); @@ -717,7 +741,7 @@ mod tests { let mut actual_neighbors = [nearest, middle, farthest]; // When - insert_neighbor( + insert_eligible_neighbor( &mut actual_neighbors, eligible_candidate.target, eligible_candidate.distance, @@ -728,7 +752,7 @@ mod tests { } #[test] - fn runtime_width_insertion_shifts_only_the_later_neighbors() { + fn middle_candidate_shifts_only_farther_runtime_neighbors() { // Given let first = LeafNeighbor::new(1, 1.0); let second = LeafNeighbor::new(2, 2.0); @@ -739,7 +763,7 @@ mod tests { let mut actual_neighbors = [first, second, third, fourth]; // When - insert_neighbor( + insert_eligible_neighbor( actual_neighbors.as_mut_slice(), candidate.target, candidate.distance, @@ -750,7 +774,7 @@ mod tests { } #[test] - fn sorted_insertion_preserves_existing_order_for_equal_distances() { + fn equal_distance_candidate_stays_after_the_existing_neighbor() { // Given let nearest = LeafNeighbor::new(1, 1.0); let existing_tie = LeafNeighbor::new(2, 2.0); @@ -760,7 +784,7 @@ mod tests { let mut actual_neighbors = [nearest, existing_tie, farthest]; // When - insert_neighbor( + insert_eligible_neighbor( &mut actual_neighbors, tied_candidate.target, tied_candidate.distance, @@ -775,7 +799,7 @@ mod tests { use super::leaf_neighbor_count; #[test] - fn empty_leaf_cannot_retain_neighbors() { + fn returns_zero_when_the_leaf_is_empty() { // Given let point_count = 0; let requested_k = 3; @@ -789,7 +813,7 @@ mod tests { } #[test] - fn singleton_leaf_cannot_retain_its_source_point() { + fn returns_zero_when_the_leaf_contains_only_the_source() { // Given let point_count = 1; let requested_k = 3; @@ -803,7 +827,7 @@ mod tests { } #[test] - fn requested_k_above_available_neighbors_is_clamped() { + fn returns_the_non_self_point_count_when_requested_k_is_larger() { // Given let point_count = 4; let requested_k = 4; @@ -817,7 +841,7 @@ mod tests { } #[test] - fn requested_k_within_available_neighbors_is_unchanged() { + fn returns_requested_k_when_enough_non_self_points_exist() { // Given let point_count = 8; let requested_k = 5; @@ -864,6 +888,35 @@ mod tests { assert_eq!(actual_neighbors, expected_neighbors); } + #[test] + fn later_farther_candidate_cannot_replace_the_retained_neighbor() { + // Given + let point_values = [0.0_f32, 10.0, 1.0]; + let source = 2; + let first_scanned_target = 0_u32; + let later_farther_target = 1_usize; + let expected_distance = + (point_values[source] - point_values[first_scanned_target as usize]).powi(2); + let later_distance = + (point_values[source] - point_values[later_farther_target]).powi(2); + let expected_nearest_neighbor = + LeafNeighbor::new(first_scanned_target, expected_distance); + assert!(later_distance > expected_distance); + let mut actual_neighbors = [LeafNeighbor::default(); 3]; + + // When + select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + MatrixView::try_from(&point_values[..], 3, 1).unwrap(), + MutMatrixView::try_from(&mut actual_neighbors[..], 3, 1).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap(); + + // Then + assert_eq!(actual_neighbors[source], expected_nearest_neighbor); + } + #[test] fn reused_workspace_matches_fresh_neighbor_selection() { // Given @@ -1133,6 +1186,21 @@ mod tests { assert_eq!(&actual_neighbors[..2], &expected_zero_norm_neighbors); } + #[test] + fn zero_target_norm_remains_rankable_in_a_complete_simd_group_with_cosine() { + // Given + let points = 17; + let mut gram = gram_with_uniform_self_dots(points, 1.0); + gram[0] = 0.0; + let expected_zero_norm_neighbor = LeafNeighbor::new(0, 1.0); + + // When + let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, points, 1).1; + + // Then + assert_eq!(actual_neighbors[16], expected_zero_norm_neighbor); + } + #[test] fn similarity_above_one_clamps_to_zero_distance_with_cosine() { // Given @@ -1224,17 +1292,44 @@ mod tests { assert_eq!(actual_neighbors, expected_unassigned_neighbors); } - #[test] - fn simd_nan_distance_cannot_replace_a_finite_neighbor() { + #[rstest] + #[case::normalized_cosine(Metric::CosineNormalized, 1.0)] + #[case::inner_product(Metric::InnerProduct, -0.0)] + fn simd_nan_distance_cannot_replace_a_finite_neighbor( + #[case] metric: Metric, + #[case] expected_distance: f32, + ) { // Given let points = 17; let mut gram = gram_with_uniform_self_dots(points, 1.0); gram[16 * points] = f32::NAN; gram[16] = f32::NAN; - let expected_finite_neighbor = LeafNeighbor::new(1, 1.0); + let expected_finite_neighbor = LeafNeighbor::new(1, expected_distance); + + // When + let actual_neighbors = rank_neighbors(metric, &gram, points, 1).1; + + // Then + assert_eq!(actual_neighbors[16], expected_finite_neighbor); + } + + #[rstest] + #[case::l2(Metric::L2, 2.0)] + #[case::normalized_cosine(Metric::CosineNormalized, 1.0)] + #[case::inner_product(Metric::InnerProduct, -0.0)] + fn simd_positive_infinity_cannot_fill_a_neighbor_slot( + #[case] metric: Metric, + #[case] expected_distance: f32, + ) { + // Given: negative-infinite dot products produce positive-infinite scores here. + let points = 17; + let mut gram = gram_with_uniform_self_dots(points, 1.0); + gram[16 * points] = f32::NEG_INFINITY; + gram[16] = f32::NEG_INFINITY; + let expected_finite_neighbor = LeafNeighbor::new(1, expected_distance); // When - let actual_neighbors = rank_neighbors(Metric::CosineNormalized, &gram, points, 1).1; + let actual_neighbors = rank_neighbors(metric, &gram, points, 1).1; // Then assert_eq!(actual_neighbors[16], expected_finite_neighbor); diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 275053f832..573b2bfb2e 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -607,6 +607,26 @@ mod tests { assert_eq!(actual_assignments, expected_sampled_leader_order); } + #[test] + fn zero_point_norm_keeps_first_leader_in_a_complete_simd_group_with_cosine() { + // Given + let leader_count = 17; + let dots = [0.0; 17]; + let point_norms = [0.0]; + let leader_norms = [1.0; 17]; + let expected_first_leader = [0]; + + // When + let actual_assignment = rank_partition_leaders( + Metric::Cosine, + partition_input(&dots, 1, leader_count, &point_norms, &leader_norms), + 1, + ); + + // Then + assert_eq!(actual_assignment, expected_first_leader); + } + #[test] fn f32_max_score_is_still_a_rankable_leader() { // Given From fca286add5d63466ba3dd5af0ad7a46c878a29ec Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:22:49 +0000 Subject: [PATCH 085/229] refactor(pipnn): use one SIMD vector type --- diskann/src/graph/pipnn/kernel_metric.rs | 24 +++-- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 100 +++++++----------- .../graph/pipnn/kernel_metric/partition.rs | 92 ++++++---------- diskann/src/graph/pipnn/leaf_kernel.rs | 14 +-- diskann/src/graph/pipnn/partition_kernel.rs | 6 +- diskann/src/graph/pipnn/simd.rs | 14 +-- 6 files changed, 101 insertions(+), 149 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 1a5a47c975..c67245c77a 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -72,18 +72,20 @@ mod tests { source_norms: [f32; 16], target_norms: [f32; 16], ) -> [f32; 16] { - type TestVector = ::LeafScore; - assert_eq!(TestVector::LANES, 16); - // SAFETY: each input and output contains exactly one complete SIMD vector. - unsafe { - let dots = TestVector::load_simd(ARCH, dot_products.as_ptr()); - let sources = TestVector::load_simd(ARCH, source_norms.as_ptr()); - let targets = TestVector::load_simd(ARCH, target_norms.as_ptr()); - let distances = cosine_distance_simd(ARCH, dots, sources, targets); - let mut output = [0.0; 16]; - distances.store_simd(output.as_mut_ptr()); - output + type TestVector = ::Vector; + assert_eq!(16 % TestVector::LANES, 0); + let mut output = [0.0; 16]; + for first in (0..16).step_by(TestVector::LANES) { + // SAFETY: each offset starts one complete SIMD group in every array. + unsafe { + let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); + let sources = TestVector::load_simd(ARCH, source_norms.as_ptr().add(first)); + let targets = TestVector::load_simd(ARCH, target_norms.as_ptr().add(first)); + cosine_distance_simd(ARCH, dots, sources, targets) + .store_simd(output.as_mut_ptr().add(first)); + } } + output } mod cosine_distance_single_tests { diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 7a176c96af..78927279fc 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -13,11 +13,6 @@ use super::{ /// 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(); @@ -25,11 +20,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: A, _norms: &[f32], _source: usize) -> Self::Simd + fn source_simd(arch: A, _norms: &[f32], _source: usize) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::default(arch) + A::Vector::default(arch) } /// Prepare one source norm for reuse across single target values. @@ -42,10 +37,10 @@ pub(in super::super) trait LeafMetric: Send + Sync + 'static { fn distances_simd( arch: A, norms: &[f32], - source_norms: Self::Simd, - dot_products: Self::Simd, + source_norms: A::Vector, + dot_products: A::Vector, first_target: usize, - ) -> Self::Simd + ) -> A::Vector where A: PiPNNSIMDSchema; @@ -67,11 +62,6 @@ 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() { @@ -80,11 +70,11 @@ impl LeafMetric for L2 { } #[inline(always)] - fn source_simd(arch: A, norms: &[f32], source: usize) -> Self::Simd + fn source_simd(arch: A, norms: &[f32], source: usize) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::splat(arch, norms[source]) + A::Vector::splat(arch, norms[source]) } #[inline(always)] @@ -96,16 +86,16 @@ impl LeafMetric for L2 { fn distances_simd( arch: A, norms: &[f32], - source_norms: Self::Simd, - dot_products: Self::Simd, + source_norms: A::Vector, + dot_products: A::Vector, first_target: usize, - ) -> Self::Simd + ) -> A::Vector where A: PiPNNSIMDSchema, { - 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)) + let target_norms = load_norms_simd::(arch, norms, first_target); + (A::Vector::splat(arch, -2.0).mul_add_simd(dot_products, source_norms) + target_norms) + .max_simd(A::Vector::default(arch)) } #[inline(always)] @@ -115,11 +105,6 @@ 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() { @@ -128,11 +113,11 @@ impl LeafMetric for Cosine { } #[inline(always)] - fn source_simd(arch: A, norms: &[f32], source: usize) -> Self::Simd + fn source_simd(arch: A, norms: &[f32], source: usize) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::splat(arch, norms[source]) + A::Vector::splat(arch, norms[source]) } #[inline(always)] @@ -144,16 +129,16 @@ impl LeafMetric for Cosine { fn distances_simd( arch: A, norms: &[f32], - source_norms: Self::Simd, - dot_products: Self::Simd, + source_norms: A::Vector, + dot_products: A::Vector, first_target: usize, - ) -> Self::Simd + ) -> A::Vector where 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(Self::Simd::::default(arch)) + .max_simd(A::Vector::default(arch)) } #[inline(always)] @@ -163,23 +148,18 @@ impl LeafMetric for Cosine { } impl LeafMetric for CosineNormalized { - type Simd - = A::LeafScore - where - A: PiPNNSIMDSchema; - #[inline(always)] fn distances_simd( arch: A, _norms: &[f32], - _source_norms: Self::Simd, - dot_products: Self::Simd, + _source_norms: A::Vector, + dot_products: A::Vector, _first_target: usize, - ) -> Self::Simd + ) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::splat(arch, 1.0) - dot_products + A::Vector::splat(arch, 1.0) - dot_products } #[inline(always)] @@ -189,23 +169,18 @@ impl LeafMetric for CosineNormalized { } impl LeafMetric for InnerProduct { - type Simd - = A::LeafScore - where - A: PiPNNSIMDSchema; - #[inline(always)] fn distances_simd( arch: A, _norms: &[f32], - _source_norms: Self::Simd, - dot_products: Self::Simd, + _source_norms: A::Vector, + dot_products: A::Vector, _first_target: usize, - ) -> Self::Simd + ) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::default(arch) - dot_products + A::Vector::default(arch) - dot_products } #[inline(always)] @@ -229,16 +204,19 @@ mod tests { dot_products: [f32; 16], first_target: usize, ) -> [f32; 16] { - assert_eq!(M::Simd::::LANES, 16); + type TestVector = ::Vector; + assert_eq!(16 % TestVector::LANES, 0); let source_norms = M::source_simd(ARCH, norms, source); - // SAFETY: the input and output contain exactly one complete SIMD vector. - unsafe { - let dots = M::Simd::::load_simd(ARCH, dot_products.as_ptr()); - let distances = M::distances_simd(ARCH, norms, source_norms, dots, first_target); - let mut output = [0.0; 16]; - distances.store_simd(output.as_mut_ptr()); - output + let mut output = [0.0; 16]; + for first in (0..16).step_by(TestVector::LANES) { + // SAFETY: each offset starts one complete SIMD group in both arrays. + unsafe { + let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); + M::distances_simd(ARCH, norms, source_norms, dots, first_target + first) + .store_simd(output.as_mut_ptr().add(first)); + } } + output } mod prepare_leaf_norms_tests { diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 2946e43fed..288474cbe7 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -15,11 +15,6 @@ 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(); @@ -32,11 +27,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: A, _norms: PartitionNorms<'_>, _point: usize) -> Self::Simd + fn point_simd(arch: A, _norms: PartitionNorms<'_>, _point: usize) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::default(arch) + A::Vector::default(arch) } /// Prepare one point norm for reuse across single leader values. @@ -49,10 +44,10 @@ pub(in super::super) trait PartitionMetric: Send + Sync + 'static { fn rankings_simd( arch: A, norms: PartitionNorms<'_>, - point_norms: Self::Simd, - dot_products: Self::Simd, + point_norms: A::Vector, + dot_products: A::Vector, first_leader: usize, - ) -> Self::Simd + ) -> A::Vector where A: PiPNNSIMDSchema; @@ -79,11 +74,6 @@ 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()) { @@ -95,15 +85,15 @@ impl PartitionMetric for L2 { fn rankings_simd( arch: A, norms: PartitionNorms<'_>, - _point_norms: Self::Simd, - dot_products: Self::Simd, + _point_norms: A::Vector, + dot_products: A::Vector, first_leader: usize, - ) -> Self::Simd + ) -> A::Vector where A: PiPNNSIMDSchema, { - 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) + let leader_norms = load_norms_simd::(arch, norms.leader_norms, first_leader); + A::Vector::splat(arch, -2.0).mul_add_simd(dot_products, leader_norms) } #[inline(always)] @@ -118,11 +108,6 @@ 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()) { @@ -138,11 +123,11 @@ impl PartitionMetric for Cosine { } #[inline(always)] - fn point_simd(arch: A, norms: PartitionNorms<'_>, point: usize) -> Self::Simd + fn point_simd(arch: A, norms: PartitionNorms<'_>, point: usize) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::splat(arch, norms.point_norms[point]) + A::Vector::splat(arch, norms.point_norms[point]) } #[inline(always)] @@ -154,14 +139,14 @@ impl PartitionMetric for Cosine { fn rankings_simd( arch: A, norms: PartitionNorms<'_>, - point_norms: Self::Simd, - dot_products: Self::Simd, + point_norms: A::Vector, + dot_products: A::Vector, first_leader: usize, - ) -> Self::Simd + ) -> A::Vector where 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) } @@ -177,23 +162,18 @@ impl PartitionMetric for Cosine { } impl PartitionMetric for CosineNormalized { - type Simd - = A::PartitionScore - where - A: PiPNNSIMDSchema; - #[inline(always)] fn rankings_simd( arch: A, _norms: PartitionNorms<'_>, - _point_norms: Self::Simd, - dot_products: Self::Simd, + _point_norms: A::Vector, + dot_products: A::Vector, _first_leader: usize, - ) -> Self::Simd + ) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::splat(arch, 1.0) - dot_products + A::Vector::splat(arch, 1.0) - dot_products } #[inline(always)] @@ -208,23 +188,18 @@ impl PartitionMetric for CosineNormalized { } impl PartitionMetric for InnerProduct { - type Simd - = A::PartitionScore - where - A: PiPNNSIMDSchema; - #[inline(always)] fn rankings_simd( arch: A, _norms: PartitionNorms<'_>, - _point_norms: Self::Simd, - dot_products: Self::Simd, + _point_norms: A::Vector, + dot_products: A::Vector, _first_leader: usize, - ) -> Self::Simd + ) -> A::Vector where A: PiPNNSIMDSchema, { - Self::Simd::::default(arch) - dot_products + A::Vector::default(arch) - dot_products } #[inline(always)] @@ -254,20 +229,23 @@ mod tests { dot_products: [f32; 16], first_leader: usize, ) -> [f32; 16] { - assert_eq!(M::Simd::::LANES, 16); + type TestVector = ::Vector; + assert_eq!(16 % TestVector::LANES, 0); let norms = PartitionNorms { point_norms, leader_norms, }; let point_norm = M::point_simd(ARCH, norms, point); - // SAFETY: the input and output contain exactly one complete SIMD vector. - unsafe { - let dots = M::Simd::::load_simd(ARCH, dot_products.as_ptr()); - let rankings = M::rankings_simd(ARCH, norms, point_norm, dots, first_leader); - let mut output = [0.0; 16]; - rankings.store_simd(output.as_mut_ptr()); - output + let mut output = [0.0; 16]; + for first in (0..16).step_by(TestVector::LANES) { + // SAFETY: each offset starts one complete SIMD group in both arrays. + unsafe { + let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); + M::rankings_simd(ARCH, norms, point_norm, dots, first_leader + first) + .store_simd(output.as_mut_ptr().add(first)); + } } + output } fn rank_single_leader( diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 4a29fd389b..5d870a94a5 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -248,19 +248,19 @@ fn scan_point_pairs( // 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 % M::Simd::::LANES; + let simd_prefix = source - source % A::Vector::LANES; while target < simd_prefix { // SAFETY: This complete SIMD group is in the strict-lower prefix. let dot_products = - unsafe { M::Simd::::load_simd(arch, dots.as_ptr().add(source_start + target)) }; + unsafe { A::Vector::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)); + let source_eligible = distances.lt_simd(A::Vector::splat(arch, source_worst)); // SAFETY: The complete target group is below `source < point_count`. - let target_worst = unsafe { M::Simd::::load_simd(arch, worst_ptr.add(target)) }; + let target_worst = unsafe { A::Vector::load_simd(arch, worst_ptr.add(target)) }; let target_eligible = distances.lt_simd(target_worst); - let source_bits = M::Simd::::active_lanes(source_eligible); - let target_bits = M::Simd::::active_lanes(target_eligible); + let source_bits = A::Vector::active_lanes(source_eligible); + let target_bits = A::Vector::active_lanes(target_eligible); if source_bits | target_bits != 0 { let values = distances.to_array(); @@ -285,7 +285,7 @@ fn scan_point_pairs( unsafe { *worst_ptr.add(target_source) = new_worst }; } } - target += M::Simd::::LANES; + target += A::Vector::LANES; } while target < source { diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 573b2bfb2e..1f573ca585 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -178,12 +178,12 @@ fn select_point_leaders( 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 % M::Simd::::LANES; + let simd_prefix = leader_count - leader_count % A::Vector::LANES; - for first_leader in (0..simd_prefix).step_by(M::Simd::::LANES) { + for first_leader in (0..simd_prefix).step_by(A::Vector::LANES) { // SAFETY: This group is inside the point's leader row. let dot_products = - unsafe { M::Simd::::load_simd(arch, point_dots.as_ptr().add(first_leader)) }; + unsafe { A::Vector::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); } diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 42f6efe057..33e0fe1a19 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -40,14 +40,10 @@ where } } -/// Stage-specific SIMD representations for one architecture. +/// PiPNN SIMD representation 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; + /// SIMD vector used by every numerical stage. + type Vector: PiPNNSIMDVector; } impl PiPNNSIMDSchema for A @@ -55,7 +51,5 @@ where A: Architecture, DefaultVector: PiPNNSIMDVector, { - type LeafScore = DefaultVector; - type PartitionScore = DefaultVector; - type HashScore = DefaultVector; + type Vector = DefaultVector; } From 89d607117aac3b458298c7d6d332e640cf005a2b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:53:16 +0000 Subject: [PATCH 086/229] test(pipnn): expose production kernel calls Keep runtime-dispatch adapters only for SIMD checks and pass direct production calls through setup closures for ordinary behavior tests. --- diskann/src/graph/pipnn/kernel_metric.rs | 58 +- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 53 +- .../graph/pipnn/kernel_metric/partition.rs | 104 +-- diskann/src/graph/pipnn/leaf_kernel.rs | 714 ++++++++++++------ diskann/src/graph/pipnn/partition_kernel.rs | 465 +++++++----- 5 files changed, 880 insertions(+), 514 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index c67245c77a..c5f9d9e21f 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -63,29 +63,33 @@ pub(super) fn cosine_distance_single(dot: f32, source_norm: f32, target_norm: f3 #[cfg(test)] mod tests { - use super::super::simd::PiPNNSIMDSchema; - use super::{cosine_distance_simd, cosine_distance_single}; - use diskann_wide::{ARCH, SIMDVector, arch::Current}; - - fn simd_cosine_distances( - dot_products: [f32; 16], - source_norms: [f32; 16], - target_norms: [f32; 16], - ) -> [f32; 16] { - type TestVector = ::Vector; - assert_eq!(16 % TestVector::LANES, 0); - let mut output = [0.0; 16]; - for first in (0..16).step_by(TestVector::LANES) { - // SAFETY: each offset starts one complete SIMD group in every array. - unsafe { - let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); - let sources = TestVector::load_simd(ARCH, source_norms.as_ptr().add(first)); - let targets = TestVector::load_simd(ARCH, target_norms.as_ptr().add(first)); - cosine_distance_simd(ARCH, dots, sources, targets) - .store_simd(output.as_mut_ptr().add(first)); + use super::cosine_distance_single; + + mod test_support { + use super::super::super::simd::PiPNNSIMDSchema; + use super::super::cosine_distance_simd; + use diskann_wide::{ARCH, SIMDVector, arch::Current}; + + pub(super) fn run_cosine_distance_simd( + dot_products: [f32; 16], + source_norms: [f32; 16], + target_norms: [f32; 16], + ) -> [f32; 16] { + type TestVector = ::Vector; + assert_eq!(16 % TestVector::LANES, 0); + let mut output = [0.0; 16]; + for first in (0..16).step_by(TestVector::LANES) { + // SAFETY: each offset starts one complete SIMD group in every array. + unsafe { + let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); + let sources = TestVector::load_simd(ARCH, source_norms.as_ptr().add(first)); + let targets = TestVector::load_simd(ARCH, target_norms.as_ptr().add(first)); + cosine_distance_simd(ARCH, dots, sources, targets) + .store_simd(output.as_mut_ptr().add(first)); + } } + output } - output } mod cosine_distance_single_tests { @@ -166,17 +170,17 @@ mod tests { } mod cosine_distance_simd_tests { - use super::simd_cosine_distances; + use super::test_support::run_cosine_distance_simd; #[test] fn zero_source_norm_produces_unit_distance_in_every_lane() { - let actual_distances = simd_cosine_distances([0.0; 16], [0.0; 16], [2.0; 16]); + let actual_distances = run_cosine_distance_simd([0.0; 16], [0.0; 16], [2.0; 16]); assert_eq!(actual_distances, [1.0; 16]); } #[test] fn zero_target_norm_produces_unit_distance_in_every_lane() { - let actual_distances = simd_cosine_distances([0.0; 16], [2.0; 16], [0.0; 16]); + let actual_distances = run_cosine_distance_simd([0.0; 16], [2.0; 16], [0.0; 16]); assert_eq!(actual_distances, [1.0; 16]); } @@ -184,7 +188,7 @@ mod tests { fn subnormal_norm_produces_unit_distance_in_every_lane() { let subnormal_norm = f32::MIN_POSITIVE.sqrt() / 2.0; let actual_distances = - simd_cosine_distances([0.0; 16], [2.0; 16], [subnormal_norm; 16]); + run_cosine_distance_simd([0.0; 16], [2.0; 16], [subnormal_norm; 16]); assert_eq!(actual_distances, [1.0; 16]); } @@ -214,14 +218,14 @@ mod tests { 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, ]; - let actual_distances = simd_cosine_distances(dot_products, [2.0; 16], [2.0; 16]); + let actual_distances = run_cosine_distance_simd(dot_products, [2.0; 16], [2.0; 16]); assert_eq!(actual_distances, expected_distances); } #[test] fn nan_similarity_clamps_to_zero_distance_in_every_lane() { - let actual_distances = simd_cosine_distances([f32::NAN; 16], [1.0; 16], [1.0; 16]); + let actual_distances = run_cosine_distance_simd([f32::NAN; 16], [1.0; 16], [1.0; 16]); assert_eq!(actual_distances, [0.0; 16]); } } diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 78927279fc..eb01d8dc90 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -196,27 +196,31 @@ impl LeafMetric for InnerProduct { )] mod tests { use super::*; - use diskann_wide::{ARCH, SIMDVector, arch::Current}; - fn simd_distances( - norms: &[f32], - source: usize, - dot_products: [f32; 16], - first_target: usize, - ) -> [f32; 16] { - type TestVector = ::Vector; - assert_eq!(16 % TestVector::LANES, 0); - let source_norms = M::source_simd(ARCH, norms, source); - let mut output = [0.0; 16]; - for first in (0..16).step_by(TestVector::LANES) { - // SAFETY: each offset starts one complete SIMD group in both arrays. - unsafe { - let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); - M::distances_simd(ARCH, norms, source_norms, dots, first_target + first) - .store_simd(output.as_mut_ptr().add(first)); + mod test_support { + use super::*; + use diskann_wide::{ARCH, SIMDVector, arch::Current}; + + pub(super) fn run_distances_simd( + norms: &[f32], + source: usize, + dot_products: [f32; 16], + first_target: usize, + ) -> [f32; 16] { + type TestVector = ::Vector; + assert_eq!(16 % TestVector::LANES, 0); + let source_norms = M::source_simd(ARCH, norms, source); + let mut output = [0.0; 16]; + for first in (0..16).step_by(TestVector::LANES) { + // SAFETY: each offset starts one complete SIMD group in both arrays. + unsafe { + let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); + M::distances_simd(ARCH, norms, source_norms, dots, first_target + first) + .store_simd(output.as_mut_ptr().add(first)); + } } + output } - output } mod prepare_leaf_norms_tests { @@ -351,6 +355,7 @@ mod tests { } mod distances_simd_tests { + use super::test_support::run_distances_simd; use super::*; #[test] @@ -362,7 +367,8 @@ mod tests { let mut squared_norms = [target_squared_norm; 17]; squared_norms[0] = source_squared_norm; - let actual_distances = simd_distances::(&squared_norms, 0, [dot_product; 16], 1); + let actual_distances = + run_distances_simd::(&squared_norms, 0, [dot_product; 16], 1); assert_eq!(actual_distances, [expected_distance; 16]); } @@ -372,7 +378,7 @@ mod tests { let squared_norms = [1.0; 17]; let expected_non_negative_distance = 0.0; - let actual_distances = simd_distances::(&squared_norms, 0, [f32::NAN; 16], 1); + let actual_distances = run_distances_simd::(&squared_norms, 0, [f32::NAN; 16], 1); assert_eq!(actual_distances, [expected_non_negative_distance; 16]); } @@ -386,7 +392,7 @@ mod tests { let mut norms = [target_norm; 17]; norms[0] = source_norm; - let actual_distances = simd_distances::(&norms, 0, [dot_product; 16], 1); + let actual_distances = run_distances_simd::(&norms, 0, [dot_product; 16], 1); assert_eq!(actual_distances, [expected_distance; 16]); } @@ -396,7 +402,8 @@ mod tests { let dot_product = 0.25; let expected_distance = 1.0 - dot_product; - let actual_distances = simd_distances::(&[], 0, [dot_product; 16], 0); + let actual_distances = + run_distances_simd::(&[], 0, [dot_product; 16], 0); assert_eq!(actual_distances, [expected_distance; 16]); } @@ -406,7 +413,7 @@ mod tests { let dot_product = 3.0; let expected_distance = -dot_product; - let actual_distances = simd_distances::(&[], 0, [dot_product; 16], 0); + let actual_distances = run_distances_simd::(&[], 0, [dot_product; 16], 0); assert_eq!(actual_distances, [expected_distance; 16]); } diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 288474cbe7..c1fd9c2852 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -220,46 +220,50 @@ impl PartitionMetric for InnerProduct { )] mod tests { use super::*; - use diskann_wide::{ARCH, SIMDVector, arch::Current}; - fn simd_rankings( - point_norms: &[f32], - leader_norms: &[f32], - point: usize, - dot_products: [f32; 16], - first_leader: usize, - ) -> [f32; 16] { - type TestVector = ::Vector; - assert_eq!(16 % TestVector::LANES, 0); - let norms = PartitionNorms { - point_norms, - leader_norms, - }; - let point_norm = M::point_simd(ARCH, norms, point); - let mut output = [0.0; 16]; - for first in (0..16).step_by(TestVector::LANES) { - // SAFETY: each offset starts one complete SIMD group in both arrays. - unsafe { - let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); - M::rankings_simd(ARCH, norms, point_norm, dots, first_leader + first) - .store_simd(output.as_mut_ptr().add(first)); + mod test_support { + use super::*; + use diskann_wide::{ARCH, SIMDVector, arch::Current}; + + pub(super) fn run_rankings_simd( + point_norms: &[f32], + leader_norms: &[f32], + point: usize, + dot_products: [f32; 16], + first_leader: usize, + ) -> [f32; 16] { + type TestVector = ::Vector; + assert_eq!(16 % TestVector::LANES, 0); + let norms = PartitionNorms { + point_norms, + leader_norms, + }; + let point_norm = M::point_simd(ARCH, norms, point); + let mut output = [0.0; 16]; + for first in (0..16).step_by(TestVector::LANES) { + // SAFETY: each offset starts one complete SIMD group in both arrays. + unsafe { + let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); + M::rankings_simd(ARCH, norms, point_norm, dots, first_leader + first) + .store_simd(output.as_mut_ptr().add(first)); + } } + output } - output - } - 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) + pub(super) fn run_ranking_single( + 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 { @@ -342,6 +346,7 @@ mod tests { } mod ranking_single_tests { + use super::test_support::run_ranking_single; use super::*; #[test] @@ -352,7 +357,7 @@ mod tests { 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); + let actual_ranking = run_ranking_single::(dot_product, 0.0, leader_squared_norm); // Then assert_eq!(actual_ranking, expected_leader_norm_minus_twice_dot); @@ -367,7 +372,7 @@ mod tests { 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); + let actual_ranking = run_ranking_single::(dot_product, point_norm, leader_norm); // Then assert_eq!(actual_ranking, expected_one_minus_normalized_dot); @@ -380,7 +385,7 @@ mod tests { let expected_one_minus_dot = 1.0 - dot_product; // When - let actual_ranking = rank_single_leader::(dot_product, 0.0, 0.0); + let actual_ranking = run_ranking_single::(dot_product, 0.0, 0.0); // Then assert_eq!(actual_ranking, expected_one_minus_dot); @@ -393,7 +398,7 @@ mod tests { let expected_negative_dot = -dot_product; // When - let actual_ranking = rank_single_leader::(dot_product, 0.0, 0.0); + let actual_ranking = run_ranking_single::(dot_product, 0.0, 0.0); // Then assert_eq!(actual_ranking, expected_negative_dot); @@ -401,6 +406,7 @@ mod tests { } mod rankings_simd_tests { + use super::test_support::run_rankings_simd; use super::*; #[test] @@ -409,7 +415,7 @@ mod tests { let dot_product = 2.0; let expected_ranking = leader_squared_norm - 2.0 * dot_product; let actual_rankings = - simd_rankings::(&[], &[leader_squared_norm; 16], 0, [dot_product; 16], 0); + run_rankings_simd::(&[], &[leader_squared_norm; 16], 0, [dot_product; 16], 0); assert_eq!(actual_rankings, [expected_ranking; 16]); } @@ -419,8 +425,13 @@ mod tests { let leader_norm = 4.0; let dot_product = 4.0; let expected_ranking = 1.0 - dot_product / (point_norm * leader_norm); - let actual_rankings = - simd_rankings::(&[point_norm], &[leader_norm; 16], 0, [dot_product; 16], 0); + let actual_rankings = run_rankings_simd::( + &[point_norm], + &[leader_norm; 16], + 0, + [dot_product; 16], + 0, + ); assert_eq!(actual_rankings, [expected_ranking; 16]); } @@ -429,7 +440,7 @@ mod tests { let dot_product = 0.25; let expected_ranking = 1.0 - dot_product; let actual_rankings = - simd_rankings::(&[], &[], 0, [dot_product; 16], 0); + run_rankings_simd::(&[], &[], 0, [dot_product; 16], 0); assert_eq!(actual_rankings, [expected_ranking; 16]); } @@ -437,7 +448,8 @@ mod tests { fn every_lane_uses_negative_dot_product_with_inner_product() { let dot_product = 3.0; let expected_ranking = -dot_product; - let actual_rankings = simd_rankings::(&[], &[], 0, [dot_product; 16], 0); + let actual_rankings = + run_rankings_simd::(&[], &[], 0, [dot_product; 16], 0); assert_eq!(actual_rankings, [expected_ranking; 16]); } @@ -449,7 +461,7 @@ mod tests { let expected_fused_ranking = (-2.0_f32).mul_add(dot_product, leader_squared_norm); let actual_rankings = - simd_rankings::(&[], &[leader_squared_norm; 16], 0, [dot_product; 16], 0); + run_rankings_simd::(&[], &[leader_squared_norm; 16], 0, [dot_product; 16], 0); assert!(unfused_twice_dot_product.is_infinite()); assert!(expected_fused_ranking.is_finite()); diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 5d870a94a5..dc01fb0c86 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -396,214 +396,261 @@ where #[cfg(test)] mod tests { - use std::cmp::Ordering; - 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}; - struct KernelCall<'a> { - input: MatrixView<'a, f32>, - norms: &'a [f32], - output: MutMatrixView<'a, LeafNeighbor>, - workspace: &'a mut LeafKernelWorkspace, - } + mod test_support { + use std::cmp::Ordering; + + use super::*; + use diskann_wide::arch::{self, Target1}; - 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, - ), + struct KernelCall<'a> { + input: MatrixView<'a, f32>, + norms: &'a [f32], + output: MutMatrixView<'a, LeafNeighbor>, + workspace: &'a mut LeafKernelWorkspace, + } + + 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, + ), + } } } - } - fn lower_gram_view(dots: &[f32], points: usize) -> MatrixView<'_, f32> { - MatrixView::try_from(dots, points, points).unwrap() - } + fn lower_gram_view(dots: &[f32], points: usize) -> MatrixView<'_, f32> { + MatrixView::try_from(dots, points, points).unwrap() + } - fn metric_norms(metric: Metric, lower_gram: MatrixView<'_, f32>) -> Vec { - fn prepare(lower_gram: MatrixView<'_, f32>) -> Vec { + 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 + } + + match metric { + Metric::L2 => prepare::(lower_gram), + Metric::Cosine => prepare::(lower_gram), + Metric::CosineNormalized => prepare::(lower_gram), + Metric::InnerProduct => prepare::(lower_gram), + } + } + + pub(super) fn with_rank_leaf_dots_fixture( + dots: &[f32], + points: usize, + output_width: usize, + run: F, + ) -> Result, LeafKernelError> + where + M: LeafMetric, + F: FnOnce( + MatrixView<'_, f32>, + &[f32], + MutMatrixView<'_, LeafNeighbor>, + &mut Vec, + ) -> Result<(), LeafKernelError>, + { + let input = lower_gram_view(dots, points); let mut norms = Vec::new(); - M::prepare_leaf_norms(lower_gram, &mut norms); - norms + M::prepare_leaf_norms(input, &mut norms); + let mut output = vec![LeafNeighbor::default(); points * output_width]; + let output_view = + MutMatrixView::try_from(output.as_mut_slice(), points, output_width).unwrap(); + run(input, &norms, output_view, &mut Vec::new())?; + Ok(output) } - match metric { - Metric::L2 => prepare::(lower_gram), - Metric::Cosine => prepare::(lower_gram), - Metric::CosineNormalized => prepare::(lower_gram), - Metric::InnerProduct => prepare::(lower_gram), + /// Run the production ranker through runtime architecture and metric dispatch. + fn run_rank_leaf_dots_with_output_width( + metric: Metric, + dots: &[f32], + points: usize, + output_width: usize, + workspace: &mut LeafKernelWorkspace, + ) -> Result, LeafKernelError> { + let lower_gram = lower_gram_view(dots, points); + let norms = metric_norms(metric, lower_gram); + let mut output = vec![LeafNeighbor::default(); points * output_width]; + arch::dispatch1_no_features( + DispatchMetric(metric), + KernelCall { + input: lower_gram, + norms: &norms, + output: MutMatrixView::try_from(output.as_mut_slice(), points, output_width) + .unwrap(), + workspace, + }, + )?; + Ok(output) } - } - 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(); - (leaf_k, output) - } + fn run_rank_leaf_dots_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 output = + run_rank_leaf_dots_with_output_width(metric, dots, points, leaf_k, workspace) + .expect("valid leaf neighbor width"); + (leaf_k, output) + } - 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(), - ) - } + pub(super) fn run_rank_leaf_dots( + metric: Metric, + dots: &[f32], + points: usize, + requested_k: usize, + ) -> (usize, Vec) { + run_rank_leaf_dots_with_workspace( + metric, + dots, + points, + requested_k, + &mut LeafKernelWorkspace::default(), + ) + } - 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 { - let similarity = dot / (source_norm * target_norm); - 1.0 - similarity.clamp(-1.0, 1.0) + 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 { + let similarity = dot / (source_norm * target_norm); + 1.0 - similarity.clamp(-1.0, 1.0) + } } } } - } - fn reference_neighbors( - metric: Metric, - dots: &[f32], - points: usize, - requested_k: usize, - ) -> Vec { - let leaf_k = requested_k.min(points.saturating_sub(1)); - let mut output = vec![LeafNeighbor::default(); points * leaf_k]; - for source in 0..points { - let mut candidates = Vec::with_capacity(points.saturating_sub(1)); - for target in 0..points { - if source == target { - continue; - } - let (row, column) = if source > target { - (source, target) - } else { - (target, source) - }; - 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)); + pub(super) fn reference_neighbors( + metric: Metric, + dots: &[f32], + points: usize, + requested_k: usize, + ) -> Vec { + let leaf_k = requested_k.min(points.saturating_sub(1)); + let mut output = vec![LeafNeighbor::default(); points * leaf_k]; + for source in 0..points { + let mut candidates = Vec::with_capacity(points.saturating_sub(1)); + for target in 0..points { + if source == target { + continue; + } + let (row, column) = if source > target { + (source, target) + } else { + (target, source) + }; + 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.total_cmp(&right.distance)); + let retained = candidates.len().min(leaf_k); + output[source * leaf_k..source * leaf_k + retained] + .copy_from_slice(&candidates[..retained]); } - 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 } - output - } - /// Build a Gram matrix from points on the line `x = 1`. - /// - /// Normalized cosine receives unit vectors because that metric assumes - /// normalized source data. Other metrics receive the original vectors. - fn lane_boundary_gram_from_point_vectors(metric: Metric, points: usize) -> Vec { - let denominator = (points + 1) as f32; - let point_vectors: Vec<_> = (0..points) - .map(|point| { - let vector = [1.0, (point + 1) as f32 / denominator]; - if metric == Metric::CosineNormalized { - let norm = vector[0].hypot(vector[1]); - [vector[0] / norm, vector[1] / norm] - } else { - vector + /// Build a Gram matrix from points on the line `x = 1`. + /// + /// Normalized cosine receives unit vectors because that metric assumes + /// normalized source data. Other metrics receive the original vectors. + pub(super) fn lane_boundary_gram_from_point_vectors( + metric: Metric, + points: usize, + ) -> Vec { + let denominator = (points + 1) as f32; + let point_vectors: Vec<_> = (0..points) + .map(|point| { + let vector = [1.0, (point + 1) as f32 / denominator]; + if metric == Metric::CosineNormalized { + let norm = vector[0].hypot(vector[1]); + [vector[0] / norm, vector[1] / norm] + } else { + vector + } + }) + .collect(); + let mut gram = vec![0.0; points * points]; + for source in 0..points { + for target in 0..points { + gram[source * points + target] = point_vectors[source][0] + * point_vectors[target][0] + + point_vectors[source][1] * point_vectors[target][1]; } - }) - .collect(); - let mut gram = vec![0.0; points * points]; - for source in 0..points { - for target in 0..points { - gram[source * points + target] = point_vectors[source][0] - * point_vectors[target][0] - + point_vectors[source][1] * point_vectors[target][1]; } + gram } - gram - } - fn gram_with_uniform_self_dots(points: usize, self_dot: f32) -> Vec { - let mut gram = vec![0.0; points * points]; - for point in 0..points { - gram[point * points + point] = self_dot; + pub(super) fn gram_with_uniform_self_dots(points: usize, self_dot: f32) -> Vec { + let mut gram = vec![0.0; points * points]; + for point in 0..points { + gram[point * points + point] = self_dot; + } + gram } - gram } mod insert_eligible_neighbor_tests { @@ -961,6 +1008,7 @@ mod tests { } mod rank_leaf_dots_tests { + use super::test_support::*; use super::*; use rstest::rstest; @@ -993,7 +1041,7 @@ mod tests { let expected_neighbors = reference_neighbors(metric, &dots, point_count, requested_k); // When - let actual_neighbors = rank_neighbors(metric, &dots, point_count, requested_k).1; + let actual_neighbors = run_rank_leaf_dots(metric, &dots, point_count, requested_k).1; // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1016,10 +1064,6 @@ mod tests { }) } - fn rank_three_unit_points(metric: Metric) -> Vec { - rank_neighbors(metric, &three_unit_point_gram(), 3, 1).1 - } - #[test] fn selects_the_smallest_squared_distance_neighbor_for_each_point_with_l2() { // Given @@ -1034,9 +1078,24 @@ mod tests { LeafNeighbor::new(2, point_one_two_distance), LeafNeighbor::new(1, point_one_two_distance), ]; + let gram = three_unit_point_gram(); // When - let actual_neighbors = rank_three_unit_points(Metric::L2); + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 3, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, L2>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1050,9 +1109,24 @@ mod tests { LeafNeighbor::new(2, 1.0 - point_dot(POINT_ONE, POINT_TWO)), LeafNeighbor::new(1, 1.0 - point_dot(POINT_ONE, POINT_TWO)), ]; + let gram = three_unit_point_gram(); // When - let actual_neighbors = rank_three_unit_points(Metric::Cosine); + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 3, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, Cosine>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1066,9 +1140,24 @@ mod tests { LeafNeighbor::new(2, 1.0 - point_dot(POINT_ONE, POINT_TWO)), LeafNeighbor::new(1, 1.0 - point_dot(POINT_ONE, POINT_TWO)), ]; + let gram = three_unit_point_gram(); // When - let actual_neighbors = rank_three_unit_points(Metric::CosineNormalized); + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 3, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, CosineNormalized>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1082,9 +1171,24 @@ mod tests { LeafNeighbor::new(2, -point_dot(POINT_ONE, POINT_TWO)), LeafNeighbor::new(1, -point_dot(POINT_ONE, POINT_TWO)), ]; + let gram = three_unit_point_gram(); // When - let actual_neighbors = rank_three_unit_points(Metric::InnerProduct); + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 3, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, InnerProduct>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1116,7 +1220,21 @@ mod tests { ]; // When - let actual_neighbors = rank_neighbors(Metric::L2, &gram, 4, 2).1; + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 4, + 2, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, L2>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors_in_scan_order); @@ -1131,7 +1249,21 @@ mod tests { let gram = [f32::MAX, dot_product, dot_product, f32::MAX]; // When - let actual_neighbors = rank_neighbors(Metric::L2, &gram, 2, 1).1; + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 2, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, L2>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert!(unfused_twice_dot_product.is_infinite()); @@ -1155,7 +1287,7 @@ mod tests { gram[16] = dot_product; // When - let actual_neighbors = rank_neighbors(Metric::L2, &gram, points, 1).1; + let actual_neighbors = run_rank_leaf_dots(Metric::L2, &gram, points, 1).1; // Then assert!(unfused_twice_dot_product.is_infinite()); @@ -1180,7 +1312,21 @@ mod tests { [LeafNeighbor::new(1, 1.0), LeafNeighbor::new(2, 1.0)]; // When - let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, 3, 2).1; + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 3, + 2, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, Cosine>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!(&actual_neighbors[..2], &expected_zero_norm_neighbors); @@ -1195,7 +1341,7 @@ mod tests { let expected_zero_norm_neighbor = LeafNeighbor::new(0, 1.0); // When - let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, points, 1).1; + let actual_neighbors = run_rank_leaf_dots(Metric::Cosine, &gram, points, 1).1; // Then assert_eq!(actual_neighbors[16], expected_zero_norm_neighbor); @@ -1211,7 +1357,21 @@ mod tests { let expected_one_minus_maximum_similarity = 1.0 - maximum_cosine_similarity; // When - let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, 2, 1).1; + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 2, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, Cosine>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!( @@ -1230,7 +1390,21 @@ mod tests { let expected_one_minus_minimum_similarity = 1.0 - minimum_cosine_similarity; // When - let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, 2, 1).1; + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 2, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, Cosine>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!( @@ -1248,7 +1422,21 @@ mod tests { let expected_one_minus_zero_similarity = 1.0 - zero_norm_similarity; // When - let actual_neighbors = rank_neighbors(Metric::Cosine, &gram, 2, 1).1; + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 2, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, Cosine>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!( @@ -1268,13 +1456,25 @@ mod tests { gram[3] = -f32::MAX; // When - let (actual_leaf_k, actual_neighbors) = - rank_neighbors(Metric::InnerProduct, &gram, points, expected_leaf_k); + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + points, + expected_leaf_k, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, InnerProduct>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then - assert_eq!(actual_leaf_k, expected_leaf_k); assert_eq!( - actual_neighbors[3 * actual_leaf_k + actual_leaf_k - 1], + actual_neighbors[3 * expected_leaf_k + expected_leaf_k - 1], expected_last_neighbor ); } @@ -1286,7 +1486,21 @@ mod tests { let expected_unassigned_neighbors = [LeafNeighbor::default(), LeafNeighbor::default()]; // When - let actual_neighbors = rank_neighbors(Metric::CosineNormalized, &gram, 2, 1).1; + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + 2, + 1, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, CosineNormalized>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // Then assert_eq!(actual_neighbors, expected_unassigned_neighbors); @@ -1307,7 +1521,7 @@ mod tests { let expected_finite_neighbor = LeafNeighbor::new(1, expected_distance); // When - let actual_neighbors = rank_neighbors(metric, &gram, points, 1).1; + let actual_neighbors = run_rank_leaf_dots(metric, &gram, points, 1).1; // Then assert_eq!(actual_neighbors[16], expected_finite_neighbor); @@ -1329,7 +1543,7 @@ mod tests { let expected_finite_neighbor = LeafNeighbor::new(1, expected_distance); // When - let actual_neighbors = rank_neighbors(metric, &gram, points, 1).1; + let actual_neighbors = run_rank_leaf_dots(metric, &gram, points, 1).1; // Then assert_eq!(actual_neighbors[16], expected_finite_neighbor); @@ -1338,70 +1552,116 @@ mod tests { #[test] fn empty_leaf_has_no_neighbors() { // Given - let dots = []; + let empty_point_count = 0; + let empty_gram = []; 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); + let actual_neighbors = with_rank_leaf_dots_fixture::( + &empty_gram, + empty_point_count, + expected_zero_neighbor_width, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, L2>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // 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 singleton_point_count = 1; + let singleton_gram = [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); + let actual_neighbors = with_rank_leaf_dots_fixture::( + &singleton_gram, + singleton_point_count, + expected_zero_neighbor_width, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, Cosine>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // 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 point_count = 2; + let zero_requested_k = 0; + let gram = [1.0, 0.0, 0.0, 1.0]; let expected_no_neighbors: [LeafNeighbor; 0] = []; // When - let (actual_leaf_k, actual_neighbors) = - rank_neighbors(Metric::InnerProduct, &dots, 2, 0); + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + point_count, + zero_requested_k, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, InnerProduct>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); // 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 point_count = 3; + let invalid_neighbor_width = point_count; + let maximum_non_self_width = point_count - 1; + let gram = [0.0; 9]; let expected_error = LeafKernelError::InvalidNeighborCount { - points: 3, - neighbors: 3, - maximum: 2, + points: point_count, + neighbors: invalid_neighbor_width, + maximum: maximum_non_self_width, }; - 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(), + let actual_error = with_rank_leaf_dots_fixture::( + &gram, + point_count, + invalid_neighbor_width, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, L2>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) }, ) .unwrap_err(); diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 1f573ca585..467bef3788 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -248,171 +248,199 @@ mod tests { 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}; - struct KernelCall<'a> { - input: PartitionInput<'a>, - output: MutMatrixView<'a, u32>, - ranked_leaders: &'a mut Vec<(u32, f32)>, - } + mod test_support { + use super::*; + use diskann_wide::arch::{self, Target1}; - struct DispatchMetric(Metric); + struct KernelCall<'a> { + input: PartitionInput<'a>, + output: MutMatrixView<'a, u32>, + ranked_leaders: &'a mut Vec<(u32, f32)>, + } - 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) + struct DispatchMetric(Metric); + + 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::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, - ), } } - } - fn partition_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, - }, + pub(super) fn partition_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 rank_partition_leaders( - metric: Metric, - input: PartitionInput<'_>, - fanout: usize, - ) -> 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(), - }, - ); - output.into_inner().into_vec() - } + pub(super) fn with_rank_leader_dots_fixture( + input: PartitionInput<'_>, + nearest_leader_count: usize, + run: F, + ) -> Vec + where + F: FnOnce(PartitionInput<'_>, MutMatrixView<'_, u32>, &mut Vec<(u32, f32)>), + { + let mut output = Matrix::new(u32::MAX, input.dots.nrows(), nearest_leader_count); + run(input, output.as_mut_view(), &mut Vec::new()); + output.into_inner().into_vec() + } - 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) + /// Run the production ranker through runtime architecture and metric dispatch. + pub(super) fn run_rank_leader_dots( + metric: Metric, + input: PartitionInput<'_>, + nearest_leader_count: usize, + ) -> Vec { + let mut output = Matrix::new(u32::MAX, input.dots.nrows(), nearest_leader_count); + arch::dispatch1_no_features( + DispatchMetric(metric), + KernelCall { + input, + output: output.as_mut_view(), + ranked_leaders: &mut Vec::new(), + }, + ); + output.into_inner().into_vec() + } + + 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 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 - .row_iter() - .zip(output.chunks_exact_mut(fanout)) - .enumerate() - { - let point_norm = input.norms.point_norms.get(point).copied().unwrap_or(0.0); - let mut candidates: Vec<_> = dots - .iter() + pub(super) fn reference_assignments( + metric: Metric, + input: PartitionInput<'_>, + nearest_leader_count: usize, + ) -> Vec { + let mut output = vec![UNASSIGNED_LEADER; input.dots.nrows() * nearest_leader_count]; + for (point, (dots, assignments)) in input + .dots + .row_iter() + .zip(output.chunks_exact_mut(nearest_leader_count)) .enumerate() - .filter_map(|(leader, &dot)| { - 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.total_cmp(&right.1)); - for (destination, (leader, _)) in assignments.iter_mut().zip(candidates) { - *destination = leader; + { + 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 = + 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.total_cmp(&right.1)); + for (destination, (leader, _)) in assignments.iter_mut().zip(candidates) { + *destination = leader; + } } + output } - output - } - /// Build ranking input from two axis points and leaders between those axes. - /// - /// The first point prefers early leaders. The second point prefers late leaders. - fn lane_boundary_input_from_point_and_leader_vectors( - metric: Metric, - leader_count: usize, - ) -> (Vec, Vec, Vec) { - let points = [[1.0_f32, 0.0], [0.0, 1.0]]; - let denominator = (leader_count + 1) as f32; - let leaders: Vec<_> = (0..leader_count) - .map(|leader| { - let second_component = (leader + 1) as f32 / denominator; - let vector = [1.0 - second_component, second_component]; - if metric == Metric::CosineNormalized { - let norm = vector[0].hypot(vector[1]); - [vector[0] / norm, vector[1] / norm] - } else { - vector - } - }) - .collect(); - let dots = points - .iter() - .flat_map(|point| { - leaders - .iter() - .map(|leader| point[0] * leader[0] + point[1] * leader[1]) - }) - .collect(); - let point_norms = if metric == Metric::Cosine { - points - .iter() - .map(|point| point[0].hypot(point[1])) - .collect() - } else { - Vec::new() - }; - let leader_norms = match metric { - Metric::L2 => leaders - .iter() - .map(|leader| leader[0] * leader[0] + leader[1] * leader[1]) - .collect(), - Metric::Cosine => leaders + /// Build ranking input from two axis points and leaders between those axes. + /// + /// The first point prefers early leaders. The second point prefers late leaders. + pub(super) fn lane_boundary_input_from_point_and_leader_vectors( + metric: Metric, + leader_count: usize, + ) -> (Vec, Vec, Vec) { + let points = [[1.0_f32, 0.0], [0.0, 1.0]]; + let denominator = (leader_count + 1) as f32; + let leaders: Vec<_> = (0..leader_count) + .map(|leader| { + let second_component = (leader + 1) as f32 / denominator; + let vector = [1.0 - second_component, second_component]; + if metric == Metric::CosineNormalized { + let norm = vector[0].hypot(vector[1]); + [vector[0] / norm, vector[1] / norm] + } else { + vector + } + }) + .collect(); + let dots = points .iter() - .map(|leader| leader[0].hypot(leader[1])) - .collect(), - Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), - }; - (dots, point_norms, leader_norms) + .flat_map(|point| { + leaders + .iter() + .map(|leader| point[0] * leader[0] + point[1] * leader[1]) + }) + .collect(); + let point_norms = if metric == Metric::Cosine { + points + .iter() + .map(|point| point[0].hypot(point[1])) + .collect() + } else { + Vec::new() + }; + let leader_norms = match metric { + Metric::L2 => leaders + .iter() + .map(|leader| leader[0] * leader[0] + leader[1] * leader[1]) + .collect(), + Metric::Cosine => leaders + .iter() + .map(|leader| leader[0].hypot(leader[1])) + .collect(), + Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), + }; + (dots, point_norms, leader_norms) + } } mod insert_leader_tests { @@ -533,12 +561,13 @@ mod tests { } mod rank_leader_dots_tests { + use super::test_support::*; use super::*; use rstest::rstest; #[rstest] - #[case::two_leaders_fanout_one(2, 1)] - #[case::scalar_fanout_two(7, 2)] + #[case::two_leaders_select_one(2, 1)] + #[case::scalar_select_two(7, 2)] #[case::lane_minus_one(15, 3)] #[case::one_complete_lane(16, 3)] #[case::lane_plus_one(17, 4)] @@ -555,16 +584,16 @@ mod tests { )] metric: Metric, #[case] leader_count: usize, - #[case] fanout: usize, + #[case] nearest_leader_count: usize, ) { // Given let (dots, point_norms, leader_norms) = lane_boundary_input_from_point_and_leader_vectors(metric, leader_count); let input = partition_input(&dots, 2, leader_count, &point_norms, &leader_norms); - let expected_assignments = reference_assignments(metric, input, fanout); + let expected_assignments = reference_assignments(metric, input, nearest_leader_count); // When - let actual_assignments = rank_partition_leaders(metric, input, fanout); + let actual_assignments = run_rank_leader_dots(metric, input, nearest_leader_count); // Then assert_eq!(actual_assignments, expected_assignments); @@ -573,15 +602,23 @@ mod tests { #[test] fn equal_scores_keep_sampled_leader_order_with_l2() { // Given + let point_count = 1; + let leader_count = 4; + let nearest_leader_count = 2; 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]; + let input = + partition_input(&dots, point_count, leader_count, &[], &leader_squared_norms); + // When - let actual_assignments = rank_partition_leaders( - Metric::L2, - partition_input(&dots, 1, 4, &[], &leader_squared_norms), - 2, + let actual_assignments = with_rank_leader_dots_fixture( + input, + nearest_leader_count, + |input, output, ranked_leaders| { + rank_leader_dots::<_, L2>(diskann_wide::ARCH, input, output, ranked_leaders); + }, ); // Then @@ -591,16 +628,34 @@ mod tests { #[test] fn zero_norm_keeps_sampled_leader_order_with_cosine() { // Given + let point_count = 1; + let leader_count = 2; + let nearest_leader_count = 2; let dots = [0.0, 0.0]; let point_norms = [0.0]; let leader_norms = [1.0, 1.0]; let expected_sampled_leader_order = [0, 1]; + let input = partition_input( + &dots, + point_count, + leader_count, + &point_norms, + &leader_norms, + ); + // When - let actual_assignments = rank_partition_leaders( - Metric::Cosine, - partition_input(&dots, 1, 2, &point_norms, &leader_norms), - 2, + let actual_assignments = with_rank_leader_dots_fixture( + input, + nearest_leader_count, + |input, output, ranked_leaders| { + rank_leader_dots::<_, Cosine>( + diskann_wide::ARCH, + input, + output, + ranked_leaders, + ); + }, ); // Then @@ -610,17 +665,25 @@ mod tests { #[test] fn zero_point_norm_keeps_first_leader_in_a_complete_simd_group_with_cosine() { // Given + let point_count = 1; let leader_count = 17; + let nearest_leader_count = 1; let dots = [0.0; 17]; let point_norms = [0.0]; let leader_norms = [1.0; 17]; let expected_first_leader = [0]; // When - let actual_assignment = rank_partition_leaders( + let actual_assignment = run_rank_leader_dots( Metric::Cosine, - partition_input(&dots, 1, leader_count, &point_norms, &leader_norms), - 1, + partition_input( + &dots, + point_count, + leader_count, + &point_norms, + &leader_norms, + ), + nearest_leader_count, ); // Then @@ -630,17 +693,29 @@ mod tests { #[test] fn f32_max_score_is_still_a_rankable_leader() { // Given + let point_count = 1; + let leader_count = 8; + let nearest_leader_count = leader_count; let maximum_rankable_score = f32::MAX; let dot_product_that_produces_it = -maximum_rankable_score; let mut dots = [0.0; 8]; dots[7] = dot_product_that_produces_it; let expected_all_leaders_in_scan_order = [0, 1, 2, 3, 4, 5, 6, 7]; + let input = partition_input(&dots, point_count, leader_count, &[], &[]); + // When - let actual_assignments = rank_partition_leaders( - Metric::InnerProduct, - partition_input(&dots, 1, 8, &[], &[]), - 8, + let actual_assignments = with_rank_leader_dots_fixture( + input, + nearest_leader_count, + |input, output, ranked_leaders| { + rank_leader_dots::<_, InnerProduct>( + diskann_wide::ARCH, + input, + output, + ranked_leaders, + ); + }, ); // Then @@ -650,14 +725,26 @@ mod tests { #[test] fn nan_leader_does_not_displace_finite_leaders() { // Given + let point_count = 1; + let leader_count = 3; + let nearest_leader_count = 2; let dots = [f32::NAN, 3.0, 2.0]; let expected_finite_leaders = [1, 2]; + let input = partition_input(&dots, point_count, leader_count, &[], &[]); + // When - let actual_assignments = rank_partition_leaders( - Metric::InnerProduct, - partition_input(&dots, 1, 3, &[], &[]), - 2, + let actual_assignments = with_rank_leader_dots_fixture( + input, + nearest_leader_count, + |input, output, ranked_leaders| { + rank_leader_dots::<_, InnerProduct>( + diskann_wide::ARCH, + input, + output, + ranked_leaders, + ); + }, ); // Then @@ -667,31 +754,27 @@ mod tests { #[test] fn empty_point_matrix_produces_no_assignments() { // Given - let dots = []; + let empty_point_count = 0; + let leader_count = 3; + let nearest_leader_count = 2; + let no_dot_products = []; let expected_no_assignments: [u32; 0] = []; - // When - let actual_assignments = rank_partition_leaders( - Metric::InnerProduct, - 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] = []; + let input = + partition_input(&no_dot_products, empty_point_count, leader_count, &[], &[]); // When - let actual_assignments = rank_partition_leaders( - Metric::InnerProduct, - partition_input(&dots, 1, 3, &[], &[]), - 0, + let actual_assignments = with_rank_leader_dots_fixture( + input, + nearest_leader_count, + |input, output, ranked_leaders| { + rank_leader_dots::<_, InnerProduct>( + diskann_wide::ARCH, + input, + output, + ranked_leaders, + ); + }, ); // Then From c35b9445178913e1e5abbd0fdd9f3554b0382c99 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:41:02 +0000 Subject: [PATCH 087/229] fix(pipnn): align metric edge cases Keep NaN scores non-rankable and make SIMD InnerProduct match scalar unary negation for signed zero. --- diskann/src/graph/pipnn/kernel_metric.rs | 24 ++++-------- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 39 ++++++++++++++----- .../graph/pipnn/kernel_metric/partition.rs | 16 +++++++- 3 files changed, 52 insertions(+), 27 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index c5f9d9e21f..3d4b644f04 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -47,7 +47,8 @@ where 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)) + let distance = one - negative_one.max_simd(cosine.min_simd(one)); + F::select(cosine.ne_simd(cosine), cosine, distance) } /// Compute one cosine distance with the DiskANN zero-norm and NaN rules. @@ -57,7 +58,7 @@ pub(super) fn cosine_distance_single(dot: f32, source_norm: f32, target_norm: f3 1.0 } else { let cosine = dot / (source_norm * target_norm); - 1.0 - (-1.0_f32).max(1.0_f32.min(cosine)) + 1.0 - cosine.clamp(-1.0, 1.0) } } @@ -154,18 +155,9 @@ mod tests { } #[test] - fn nan_similarity_clamps_to_zero_distance() { - // 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); + fn nan_similarity_remains_nan() { + let actual_distance = cosine_distance_single(f32::NAN, 1.0, 1.0); + assert!(actual_distance.is_nan()); } } @@ -224,9 +216,9 @@ mod tests { } #[test] - fn nan_similarity_clamps_to_zero_distance_in_every_lane() { + fn nan_similarity_remains_nan_in_every_lane() { let actual_distances = run_cosine_distance_simd([f32::NAN; 16], [1.0; 16], [1.0; 16]); - assert_eq!(actual_distances, [0.0; 16]); + assert!(actual_distances.into_iter().all(f32::is_nan)); } } } diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index eb01d8dc90..9890ab537b 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -4,7 +4,7 @@ */ use diskann_utils::views::MatrixView; -use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDVector}; +use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDPartialEq, SIMDVector}; use super::super::simd::{PiPNNSIMDSchema, PiPNNSIMDVector}; use super::{ @@ -94,13 +94,20 @@ impl LeafMetric for L2 { A: PiPNNSIMDSchema, { let target_norms = load_norms_simd::(arch, norms, first_target); - (A::Vector::splat(arch, -2.0).mul_add_simd(dot_products, source_norms) + target_norms) - .max_simd(A::Vector::default(arch)) + let distances = + A::Vector::splat(arch, -2.0).mul_add_simd(dot_products, source_norms) + target_norms; + let non_negative = distances.max_simd(A::Vector::default(arch)); + A::Vector::select(distances.ne_simd(distances), distances, non_negative) } #[inline(always)] 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) + let distance = (-2.0_f32).mul_add(dot_product, source_norm) + norms[target]; + if distance.is_nan() { + distance + } else { + distance.max(0.0) + } } } @@ -138,12 +145,11 @@ impl LeafMetric for Cosine { { let target_norms = load_norms_simd::(arch, norms, first_target); cosine_distance_simd(arch, dot_products, source_norms, target_norms) - .max_simd(A::Vector::default(arch)) } #[inline(always)] 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) + cosine_distance_single(dot_product, source_norm, norms[target]) } } @@ -180,7 +186,7 @@ impl LeafMetric for InnerProduct { where A: PiPNNSIMDSchema, { - A::Vector::default(arch) - dot_products + A::Vector::splat(arch, -1.0) * dot_products } #[inline(always)] @@ -374,13 +380,12 @@ mod tests { } #[test] - fn nan_dot_products_clamp_to_zero_in_every_lane_with_l2() { + fn nan_dot_products_remain_nan_in_every_lane_with_l2() { let squared_norms = [1.0; 17]; - let expected_non_negative_distance = 0.0; let actual_distances = run_distances_simd::(&squared_norms, 0, [f32::NAN; 16], 1); - assert_eq!(actual_distances, [expected_non_negative_distance; 16]); + assert!(actual_distances.into_iter().all(f32::is_nan)); } #[test] @@ -417,5 +422,19 @@ mod tests { assert_eq!(actual_distances, [expected_distance; 16]); } + + #[test] + fn signed_zero_bits_match_scalar_negation_with_inner_product() { + let dot_products: [f32; 16] = [ + 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, + -0.0, + ]; + let expected_bits = dot_products.map(|dot| (-dot).to_bits()); + + let actual_bits = + run_distances_simd::(&[], 0, dot_products, 0).map(f32::to_bits); + + assert_eq!(actual_bits, expected_bits); + } } } diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index c1fd9c2852..fd6e57679f 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -199,7 +199,7 @@ impl PartitionMetric for InnerProduct { where A: PiPNNSIMDSchema, { - A::Vector::default(arch) - dot_products + A::Vector::splat(arch, -1.0) * dot_products } #[inline(always)] @@ -453,6 +453,20 @@ mod tests { assert_eq!(actual_rankings, [expected_ranking; 16]); } + #[test] + fn signed_zero_bits_match_scalar_negation_with_inner_product() { + let dot_products: [f32; 16] = [ + 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, + -0.0, + ]; + let expected_bits = dot_products.map(|dot| (-dot).to_bits()); + + let actual_bits = + run_rankings_simd::(&[], &[], 0, dot_products, 0).map(f32::to_bits); + + assert_eq!(actual_bits, expected_bits); + } + #[test] fn every_lane_stays_finite_when_twice_the_dot_product_overflows_with_l2() { let dot_product = f32::from_bits(f32::MAX.to_bits() - 1); From 816a34fe5613e2735cea3fb6024ffe1e390d2624 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:10:26 +0000 Subject: [PATCH 088/229] docs(pipnn): correct leaf width contract --- 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 092ee7549a..32588f73c4 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -10,8 +10,8 @@ //! 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. +//! pair one time. It updates both points. Each point keeps the requested number +//! of neighbors. This number cannot exceed the number of other points in the leaf. //! //! `kernel_metric` defines metric markers and shared math. Separate leaf and //! partition traits define norm preparation and ranking formulas. From 6844cfa33dcbbb6e675991fcb166741f6a74ac5f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:33:25 +0000 Subject: [PATCH 089/229] ci(pipnn): run kernels under Miri --- .github/workflows/nightly.yml | 12 ++++++++++++ diskann/src/graph/pipnn/kernel_metric/partition.rs | 12 ++++++------ diskann/src/graph/pipnn/leaf_kernel.rs | 5 +++++ 3 files changed, 23 insertions(+), 6 deletions(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 8103743955..87991e638c 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -135,3 +135,15 @@ jobs: run: cargo +nightly miri nextest run --locked --package diskann-quantization env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance + + - name: PiPNN numerical pointer boundaries + env: + MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance + run: | + cargo +nightly miri test --locked -p diskann-wide --lib div_ + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::kernel_metric + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::leaf_kernel::tests::rank_leaf_dots_tests + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::partition_kernel::tests::rank_leader_dots_tests diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index fd6e57679f..277ece78d5 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -282,8 +282,8 @@ mod tests { ]; 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), + first_leader[0] * first_leader[0] + first_leader[1] * first_leader[1], + second_leader[0] * second_leader[0] + second_leader[1] * second_leader[1], ]; let mut actual_norms = Vec::new(); @@ -307,8 +307,8 @@ mod tests { ]; 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(), + (first_leader[0] * first_leader[0] + first_leader[1] * first_leader[1]).sqrt(), + (second_leader[0] * second_leader[0] + second_leader[1] * second_leader[1]).sqrt(), ]; let mut actual_norms = Vec::new(); @@ -333,8 +333,8 @@ mod tests { ]; 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(), + (first_point[0] * first_point[0] + first_point[1] * first_point[1]).sqrt(), + (second_point[0] * second_point[0] + second_point[1] * second_point[1]).sqrt(), ]; let mut actual_norms = Vec::new(); diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index dc01fb0c86..a9023209c7 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -1036,6 +1036,11 @@ mod tests { #[case] point_count: usize, #[case] requested_k: usize, ) { + // Miri covers the pointer boundaries in the smaller lane cases. + if cfg!(miri) && point_count > 64 { + return; + } + // Given let dots = lane_boundary_gram_from_point_vectors(metric, point_count); let expected_neighbors = reference_neighbors(metric, &dots, point_count, requested_k); From 1f0a5f1df1413dde876d2ddbebd2b91ff343fadb Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 25 Aug 2026 16:58:56 +0000 Subject: [PATCH 090/229] test(pipnn): cover kernel behavior contracts Cover exceptional metric values, exact norm boundaries, all-ineligible SIMD groups, and zero-width partition output. --- diskann/src/graph/pipnn/kernel_metric.rs | 66 +++++++++++++++++-- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 47 +++++++++++++ diskann/src/graph/pipnn/leaf_kernel.rs | 36 ++++++++++ diskann/src/graph/pipnn/partition_kernel.rs | 62 ++++++++++++++++- 4 files changed, 203 insertions(+), 8 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs index 3d4b644f04..333fb11eff 100644 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -97,14 +97,14 @@ mod tests { use super::cosine_distance_single; #[test] - fn zero_source_norm_produces_unit_distance() { + fn zero_source_norm_takes_precedence_over_nan_dot_product() { // 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(0.0, source_norm, 2.0); + let actual_distance = cosine_distance_single(f32::NAN, source_norm, 2.0); // Then assert_eq!(actual_distance, expected_one_minus_zero_similarity); @@ -124,6 +124,21 @@ mod tests { assert_eq!(actual_distance, expected_one_minus_zero_similarity); } + #[test] + fn minimum_normal_norm_uses_normalized_similarity() { + // Given + let source_norm = f32::MIN_POSITIVE.sqrt(); + let target_norm = 1.0; + let dot_product = source_norm / 2.0; + let expected_distance = 1.0 - dot_product / (source_norm * target_norm); + + // When + let actual_distance = cosine_distance_single(dot_product, source_norm, target_norm); + + // Then + assert_eq!(actual_distance, expected_distance); + } + #[test] fn similarity_above_one_clamps_to_zero_distance() { // Given @@ -159,14 +174,26 @@ mod tests { let actual_distance = cosine_distance_single(f32::NAN, 1.0, 1.0); assert!(actual_distance.is_nan()); } + + #[test] + fn nan_source_norm_produces_nan_distance() { + let actual_distance = cosine_distance_single(0.0, f32::NAN, 1.0); + assert!(actual_distance.is_nan()); + } + + #[test] + fn nan_target_norm_produces_nan_distance() { + let actual_distance = cosine_distance_single(0.0, 1.0, f32::NAN); + assert!(actual_distance.is_nan()); + } } mod cosine_distance_simd_tests { use super::test_support::run_cosine_distance_simd; #[test] - fn zero_source_norm_produces_unit_distance_in_every_lane() { - let actual_distances = run_cosine_distance_simd([0.0; 16], [0.0; 16], [2.0; 16]); + fn zero_source_norm_takes_precedence_over_nan_dot_product_in_every_lane() { + let actual_distances = run_cosine_distance_simd([f32::NAN; 16], [0.0; 16], [2.0; 16]); assert_eq!(actual_distances, [1.0; 16]); } @@ -185,10 +212,24 @@ mod tests { } #[test] - fn roundoff_past_similarity_bounds_is_clamped_in_every_lane() { + fn minimum_normal_norm_uses_normalized_similarity_in_every_lane() { + let source_norm = f32::MIN_POSITIVE.sqrt(); + let target_norm = 1.0; + let dot_product = source_norm / 2.0; + let expected_distance = 1.0 - dot_product / (source_norm * target_norm); + + let actual_distances = + run_cosine_distance_simd([dot_product; 16], [source_norm; 16], [target_norm; 16]); + + assert_eq!(actual_distances, [expected_distance; 16]); + } + + #[test] + fn similarity_outside_bounds_is_clamped_in_every_lane() { let dot_above_norm_product = 4.000_001; let dot_below_negative_norm_product = -4.000_001; let dot_products = [ + f32::INFINITY, dot_above_norm_product, dot_above_norm_product, dot_above_norm_product, @@ -196,8 +237,7 @@ mod tests { dot_above_norm_product, dot_above_norm_product, dot_above_norm_product, - dot_above_norm_product, - dot_below_negative_norm_product, + f32::NEG_INFINITY, dot_below_negative_norm_product, dot_below_negative_norm_product, dot_below_negative_norm_product, @@ -220,5 +260,17 @@ mod tests { let actual_distances = run_cosine_distance_simd([f32::NAN; 16], [1.0; 16], [1.0; 16]); assert!(actual_distances.into_iter().all(f32::is_nan)); } + + #[test] + fn nan_source_norm_produces_nan_distance_in_every_lane() { + let actual_distances = run_cosine_distance_simd([0.0; 16], [f32::NAN; 16], [1.0; 16]); + assert!(actual_distances.into_iter().all(f32::is_nan)); + } + + #[test] + fn nan_target_norm_produces_nan_distance_in_every_lane() { + let actual_distances = run_cosine_distance_simd([0.0; 16], [1.0; 16], [f32::NAN; 16]); + assert!(actual_distances.into_iter().all(f32::is_nan)); + } } } diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs index 9890ab537b..d87a16c426 100644 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -317,6 +317,21 @@ mod tests { assert_eq!(actual_distance, expected_non_negative_distance); } + #[test] + fn nan_dot_product_produces_nan_distance_with_l2() { + // Given + let source_squared_norm = 1.0; + let target_squared_norm = 1.0; + let squared_norms = [source_squared_norm, target_squared_norm]; + + // When + let actual_distance = + L2::distance_single(&squared_norms, source_squared_norm, f32::NAN, 1); + + // Then + assert!(actual_distance.is_nan()); + } + #[test] fn distance_equals_one_minus_dot_over_norm_product_with_cosine() { // Given @@ -333,6 +348,18 @@ mod tests { assert_eq!(actual_distance, expected_one_minus_normalized_dot); } + #[test] + fn nan_dot_product_produces_nan_distance_with_cosine() { + // Given + let norms = [1.0, 1.0]; + + // When + let actual_distance = Cosine::distance_single(&norms, 1.0, f32::NAN, 1); + + // Then + assert!(actual_distance.is_nan()); + } + #[test] fn distance_equals_one_minus_the_dot_product_with_normalized_cosine() { // Given @@ -379,6 +406,17 @@ mod tests { assert_eq!(actual_distances, [expected_distance; 16]); } + #[test] + fn negative_roundoff_is_clamped_to_zero_in_every_lane_with_l2() { + let squared_norms = [1.0; 17]; + let dot_product_above_exact_norm = 1.000_001; + + let actual_distances = + run_distances_simd::(&squared_norms, 0, [dot_product_above_exact_norm; 16], 1); + + assert_eq!(actual_distances, [0.0; 16]); + } + #[test] fn nan_dot_products_remain_nan_in_every_lane_with_l2() { let squared_norms = [1.0; 17]; @@ -402,6 +440,15 @@ mod tests { assert_eq!(actual_distances, [expected_distance; 16]); } + #[test] + fn nan_dot_products_remain_nan_in_every_lane_with_cosine() { + let norms = [1.0; 17]; + + let actual_distances = run_distances_simd::(&norms, 0, [f32::NAN; 16], 1); + + assert!(actual_distances.into_iter().all(f32::is_nan)); + } + #[test] fn every_lane_uses_one_minus_dot_product_with_normalized_cosine() { let dot_product = 0.25; diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index a9023209c7..0479f5645a 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -1511,7 +1511,43 @@ mod tests { assert_eq!(actual_neighbors, expected_unassigned_neighbors); } + #[test] + fn complete_simd_group_without_eligible_distances_leaves_source_unassigned_with_l2() { + // Given + let point_count = 17; + let source = 16; + let requested_k = 1; + let mut gram = gram_with_uniform_self_dots(point_count, 1.0); + gram[source * point_count..source * point_count + source].fill(f32::NEG_INFINITY); + let expected_unassigned_neighbor = LeafNeighbor::default(); + + // When + let actual_neighbors = with_rank_leaf_dots_fixture::( + &gram, + point_count, + requested_k, + |input, norms, output, farthest_distances| { + rank_leaf_dots::<_, L2>( + diskann_wide::ARCH, + input, + norms, + output, + farthest_distances, + ) + }, + ) + .unwrap(); + + // Then + assert_eq!( + actual_neighbors[source * requested_k], + expected_unassigned_neighbor + ); + } + #[rstest] + #[case::l2(Metric::L2, 2.0)] + #[case::cosine(Metric::Cosine, 1.0)] #[case::normalized_cosine(Metric::CosineNormalized, 1.0)] #[case::inner_product(Metric::InnerProduct, -0.0)] fn simd_nan_distance_cannot_replace_a_finite_neighbor( diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 467bef3788..75e0f1a710 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -723,7 +723,7 @@ mod tests { } #[test] - fn nan_leader_does_not_displace_finite_leaders() { + fn nan_leader_does_not_displace_finite_leaders_with_inner_product() { // Given let point_count = 1; let leader_count = 3; @@ -751,6 +751,66 @@ mod tests { assert_eq!(actual_assignments, expected_finite_leaders); } + #[test] + fn nan_leader_does_not_displace_finite_leaders_with_cosine() { + // Given + let point_count = 1; + let leader_count = 3; + let nearest_leader_count = 2; + let dots = [f32::NAN, 0.75, 0.5]; + let point_norms = [1.0]; + let leader_norms = [1.0; 3]; + let expected_finite_leaders = [1, 2]; + + let input = partition_input( + &dots, + point_count, + leader_count, + &point_norms, + &leader_norms, + ); + + // When + let actual_assignments = with_rank_leader_dots_fixture( + input, + nearest_leader_count, + |input, output, ranked_leaders| { + rank_leader_dots::<_, Cosine>( + diskann_wide::ARCH, + input, + output, + ranked_leaders, + ); + }, + ); + + // Then + assert_eq!(actual_assignments, expected_finite_leaders); + } + + #[test] + fn zero_output_width_leaves_ranked_leaders_unchanged() { + // Given + let dot_products = [0.0]; + let leader_squared_norms = [1.0]; + let input = partition_input(&dot_products, 1, 1, &[], &leader_squared_norms); + let mut no_assignments = []; + let output = MutMatrixView::try_from(&mut no_assignments[..], 1, 0).unwrap(); + let expected_ranked_leaders = vec![(7, 0.25)]; + let mut actual_ranked_leaders = expected_ranked_leaders.clone(); + + // When + rank_leader_dots::<_, L2>( + diskann_wide::ARCH, + input, + output, + &mut actual_ranked_leaders, + ); + + // Then + assert_eq!(actual_ranked_leaders, expected_ranked_leaders); + } + #[test] fn empty_point_matrix_produces_no_assignments() { // Given From 056bef7714423c604148ca7669eafe1fb89c50f7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 26 Aug 2026 06:01:32 +0000 Subject: [PATCH 091/229] refactor(pipnn): localize SIMD lane bounds --- diskann-wide/src/traits.rs | 2 +- diskann/src/graph/pipnn/leaf_kernel.rs | 2 +- diskann/src/graph/pipnn/partition_kernel.rs | 2 +- diskann/src/graph/pipnn/simd.rs | 18 +++++++++++++++--- 4 files changed, 18 insertions(+), 6 deletions(-) diff --git a/diskann-wide/src/traits.rs b/diskann-wide/src/traits.rs index da66fe9167..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]>; + 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 0479f5645a..71a8b15556 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -263,7 +263,7 @@ fn scan_point_pairs( let target_bits = A::Vector::active_lanes(target_eligible); if source_bits | target_bits != 0 { - let values = distances.to_array(); + let values = distances.to_lane_array(); let values = values.as_ref(); let mut source_bits = source_bits; while source_bits != 0 { diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 75e0f1a710..a41f675cb5 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -212,7 +212,7 @@ where return; } - let values = scores.to_array(); + let values = scores.to_lane_array(); let values = values.as_ref(); let mut lanes = F::active_lanes(eligible); while lanes != 0 { diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 33e0fe1a19..621114c761 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -5,7 +5,9 @@ //! SIMD schema for PiPNN numerical kernels. -use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector}; +use diskann_wide::{ + Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, SupportedLaneCount, +}; /// Default SIMD representation used by every PiPNN numerical stage. /// @@ -16,6 +18,9 @@ type DefaultVector = ::f32x16; pub(super) trait PiPNNSIMDVector: SIMDVector + SIMDFloat + std::ops::Div { + /// Convert the SIMD value to a readable lane array. + fn to_lane_array(self) -> impl AsRef<[f32]>; + /// Return one bit for each selected lane. fn active_lanes(mask: Self::Mask) -> u64; @@ -23,12 +28,19 @@ pub(super) trait PiPNNSIMDVector: fn select(mask: Self::Mask, if_true: Self, if_false: Self) -> Self; } -impl PiPNNSIMDVector for F +impl PiPNNSIMDVector for F where - F: SIMDVector + SIMDFloat + std::ops::Div, + F: SIMDVector> + SIMDFloat + std::ops::Div, + Const: SupportedLaneCount, F::Mask: SIMDSelect, u64: From<<::BitMask as SIMDMask>::Underlying>, { + #[inline(always)] + fn to_lane_array(self) -> impl AsRef<[f32]> { + let values: [f32; N] = self.to_array(); + values + } + #[inline(always)] fn active_lanes(mask: Self::Mask) -> u64 { u64::from(mask.bitmask().to_underlying()) From 53700500b40f4bf400a66734e44bd5b2d83a7eb7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:30:09 +0000 Subject: [PATCH 092/229] fix(pipnn): retarget worker SIMD kernels --- diskann/src/graph/pipnn/leaf_kernel.rs | 18 ++++++++++++------ diskann/src/graph/pipnn/partition_kernel.rs | 6 +++++- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 71a8b15556..e5c723e544 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -198,8 +198,11 @@ fn scan_fixed_width( [LeafNeighbor; N]: SortedInsert, { let (rows, _) = output.as_chunks_mut::(); - scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { - insert_eligible_neighbor(&mut rows[source], target, distance) + // Rayon outlines leaf workers. Reapply target features before the SIMD scan. + arch.run(move || { + scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { + insert_eligible_neighbor(&mut rows[source], target, distance) + }); }); } @@ -215,9 +218,12 @@ fn scan_runtime_width( A: PiPNNSIMDSchema, M: LeafMetric, { - scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { - let first = source * width; - insert_eligible_neighbor(&mut output[first..first + width], target, distance) + // Rayon outlines leaf workers. Reapply target features before the SIMD scan. + arch.run(move || { + scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { + let first = source * width; + insert_eligible_neighbor(&mut output[first..first + width], target, distance) + }); }); } @@ -225,7 +231,7 @@ 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)] +#[inline(always)] fn scan_point_pairs( arch: A, input: MatrixView<'_, f32>, diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index a41f675cb5..b0e8fc5539 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -150,13 +150,17 @@ fn rank_leader_dots( } ranked_leaders.resize(fanout, (UNASSIGNED_LEADER, f32::INFINITY)); - select_point_leaders::(arch, input.dots, input.norms, output, ranked_leaders); + // Rayon outlines stripe workers. Reapply target features before ranking leaders. + arch.run(move || { + 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. +#[inline(always)] fn select_point_leaders( arch: A, dots: MatrixView<'_, f32>, From 01b764977ed0239219c154bc16678eb08504a8d5 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:30:09 +0000 Subject: [PATCH 093/229] fix(pipnn): unify cosine norm reduction --- .../graph/pipnn/kernel_metric/partition.rs | 37 +++++++++++++++---- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs index 277ece78d5..cc185193e2 100644 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -73,6 +73,14 @@ where unsafe { F::load_simd(arch, norm_group.as_ptr()) } } +/// Prepare cosine norms with one reduction order for points and leaders. +fn prepare_cosine_norms(vectors: MatrixView<'_, f32>, norms: &mut Vec) { + norms.resize(vectors.nrows(), 0.0); + for (norm, vector) in norms.iter_mut().zip(vectors.row_iter()) { + *norm = FastL2NormSquared.evaluate(vector).sqrt(); + } +} + impl PartitionMetric for L2 { fn prepare_leader_norms(leaders: MatrixView<'_, f32>, norms: &mut Vec) { norms.resize(leaders.nrows(), 0.0); @@ -109,17 +117,11 @@ impl PartitionMetric for L2 { impl PartitionMetric for Cosine { 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(); - } + prepare_cosine_norms(points, norms); } 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(); - } + prepare_cosine_norms(leaders, norms); } #[inline(always)] @@ -345,6 +347,25 @@ mod tests { assert_eq!(actual_norms, expected_row_norms); } + #[test] + fn same_vector_has_identical_point_and_leader_norm_bits_with_cosine() { + // Given + const REASSOCIATION_BOUNDARY_DIMENSIONS: usize = 129; + let mut vector = [1.0_f32; REASSOCIATION_BOUNDARY_DIMENSIONS]; + vector[0] = 4096.0; + let vectors = + MatrixView::try_from(&vector[..], 1, REASSOCIATION_BOUNDARY_DIMENSIONS).unwrap(); + let mut point_norms = Vec::new(); + let mut leader_norms = Vec::new(); + + // When + Cosine::prepare_point_norms(vectors, &mut point_norms); + Cosine::prepare_leader_norms(vectors, &mut leader_norms); + + // Then + assert_eq!(point_norms[0].to_bits(), leader_norms[0].to_bits()); + } + mod ranking_single_tests { use super::test_support::run_ranking_single; use super::*; From c45218167dd28104ff6226f68db9efcfaf1ca94a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 3 Sep 2026 06:41:05 +0000 Subject: [PATCH 094/229] test(linalg): assert zero-row success --- diskann-linalg/src/lib.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index 778fb94cf8..b5cd083b2a 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -721,8 +721,18 @@ mod sgemm_aat_lower_tests { } #[test] - fn accepts_a_matrix_with_no_rows() { - sgemm_aat_lower(0, 3, &[], &mut []).unwrap(); + fn zero_row_matrix_returns_success() { + // Given + let zero_rows = 0; + let dimensions = 3; + let no_input = []; + let mut no_output = []; + + // When + let result = sgemm_aat_lower(zero_rows, dimensions, &no_input, &mut no_output); + + // Then + assert!(result.is_ok()); } #[test] From 57a2b7af1a4ffee94d9d2952c5d59b25f507e580 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 3 Sep 2026 06:59:37 +0000 Subject: [PATCH 095/229] test(linalg): clarify zero-dimension result --- diskann-linalg/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index b5cd083b2a..2cff82707a 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -736,7 +736,7 @@ mod sgemm_aat_lower_tests { } #[test] - fn zero_inner_dimension_zeros_only_the_lower_triangle() { + fn zero_inner_dimension_produces_zero_lower_triangle() { let untouched = -123.0; let mut c = [untouched; 9]; From 6d825d352e52624073fafbec8a66f64a2dd602df Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 4 Sep 2026 06:00:27 +0000 Subject: [PATCH 096/229] refactor(pipnn): simplify metric kernels Metric modules own GEMM setup and leader caches. Kernels consume flat ranking distances without metric-specific norm plumbing. --- .github/workflows/nightly.yml | 9 +- diskann-linalg/src/faer.rs | 32 +- diskann-linalg/src/lib.rs | 243 +++--- 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 - diskann/src/graph/pipnn/kernel_metric.rs | 276 ------- diskann/src/graph/pipnn/kernel_metric/leaf.rs | 487 ------------ .../graph/pipnn/kernel_metric/partition.rs | 506 ------------ diskann/src/graph/pipnn/leaf_kernel.rs | 736 ++++++------------ diskann/src/graph/pipnn/leaf_metric.rs | 300 +++++++ diskann/src/graph/pipnn/mod.rs | 152 +++- diskann/src/graph/pipnn/partition_kernel.rs | 477 ++++-------- diskann/src/graph/pipnn/partition_metric.rs | 505 ++++++++++++ diskann/src/graph/pipnn/simd.rs | 19 +- 23 files changed, 1517 insertions(+), 2295 deletions(-) delete mode 100644 diskann/src/graph/pipnn/kernel_metric.rs delete mode 100644 diskann/src/graph/pipnn/kernel_metric/leaf.rs delete mode 100644 diskann/src/graph/pipnn/kernel_metric/partition.rs create mode 100644 diskann/src/graph/pipnn/leaf_metric.rs create mode 100644 diskann/src/graph/pipnn/partition_metric.rs diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 87991e638c..c29d60275f 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -140,10 +140,11 @@ jobs: env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance run: | - cargo +nightly miri test --locked -p diskann-wide --lib div_ cargo +nightly miri test --locked -p diskann --features pipnn --lib \ - graph::pipnn::kernel_metric + graph::pipnn::leaf_metric cargo +nightly miri test --locked -p diskann --features pipnn --lib \ - graph::pipnn::leaf_kernel::tests::rank_leaf_dots_tests + graph::pipnn::partition_metric cargo +nightly miri test --locked -p diskann --features pipnn --lib \ - graph::pipnn::partition_kernel::tests::rank_leader_dots_tests + graph::pipnn::leaf_kernel::tests::rank_leaf_distances_tests + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::partition_kernel::tests::rank_leader_distances_tests diff --git a/diskann-linalg/src/faer.rs b/diskann-linalg/src/faer.rs index f4f498a13b..b21638e3da 100644 --- a/diskann-linalg/src/faer.rs +++ b/diskann-linalg/src/faer.rs @@ -53,14 +53,28 @@ pub(super) fn sgemm_impl( faer::linalg::matmul::matmul(c, beta, a, b, alpha, Par::Seq) } -/// Compute the lower triangle of `A * Aᵀ` with Faer. -/// -/// Leaf selection reads each symmetric pair once. It does not read the upper -/// triangle. `BlockStructure::TriangularLower` prevents writes to that triangle. +/// Replace the lower triangle of `C` with `alpha * A * Aᵀ`. +pub(super) fn sgemm_aat_lower_impl(m: usize, k: usize, alpha: f32, a: &[f32], c: &mut [f32]) { + sgemm_aat_lower_operation(faer::Accum::Replace, m, k, alpha, a, c); +} + +/// Add the lower triangle of `alpha * A * Aᵀ` to `C`. +pub(super) fn sgemm_aat_lower_add_impl(m: usize, k: usize, alpha: f32, a: &[f32], c: &mut [f32]) { + sgemm_aat_lower_operation(faer::Accum::Add, m, k, alpha, a, c); +} + +/// Apply one lower-triangle accumulation mode. /// -/// `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]) { +/// The public wrapper validates both matrix shapes. Faer writes the diagonal +/// and lower triangle. Faer does not read or write the upper triangle. +fn sgemm_aat_lower_operation( + accumulation: faer::Accum, + m: usize, + k: usize, + alpha: f32, + a: &[f32], + c: &mut [f32], +) { use faer::linalg::matmul::triangular::{matmul, BlockStructure}; let a = faer::mat::MatRef::from_row_major_slice(a, m, k); @@ -70,12 +84,12 @@ pub(super) fn sgemm_aat_lower_impl(m: usize, k: usize, a: &[f32], c: &mut [f32]) matmul( c, BlockStructure::TriangularLower, - faer::Accum::Replace, + accumulation, a, BlockStructure::Rectangular, at, BlockStructure::Rectangular, - 1.0, + alpha, Par::Seq, ); } diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index 2cff82707a..94582f555a 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -189,21 +189,49 @@ pub fn sgemm( Ok(()) } -/// Compute the lower triangle of $C = A A^\mathsf{T}$. +/// Replace the lower triangle of `C` with $\alpha A A^\mathsf{T}$. /// -/// `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. +/// `A` is a dense row-major $m \times k$ matrix. The function does not change +/// the upper triangle of `C`. Use this operation when `C` has no prior term. /// /// # Errors /// /// 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> { +pub fn sgemm_aat_lower( + m: usize, + k: usize, + alpha: f32, + 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); + faer::sgemm_aat_lower_impl(m, k, alpha, a, c); + Ok(()) +} + +/// Add the lower triangle of $\alpha A A^\mathsf{T}$ to `C`. +/// +/// `A` is a dense row-major $m \times k$ matrix. The function changes only the +/// lower triangle of `C`. Use this operation when `C` contains an initial term. +/// +/// # Errors +/// +/// 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_add( + m: usize, + k: usize, + alpha: f32, + 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_add_impl(m, k, alpha, a, c); Ok(()) } @@ -691,124 +719,153 @@ mod tests { } } #[cfg(test)] -#[allow( - clippy::expect_used, - clippy::unwrap_used, - reason = "deterministic test fixture construction must abort on invalid setup" -)] +#[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] mod sgemm_aat_lower_tests { - use super::{sgemm_aat_lower, MatrixName, SgemmError}; + use super::{sgemm_aat_lower, sgemm_aat_lower_add, 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, + fn replace_writes_the_scaled_lower_triangle_and_preserves_the_upper_triangle() { + // Given + let first_row = [1.0_f32, 2.0]; + let second_row = [3.0_f32, 4.0]; + let row_count = 2; + let dimension_count = first_row.len(); + let matrix = [first_row[0], first_row[1], second_row[0], second_row[1]]; + let scale = -1.0_f32; + let first_self_dot = first_row[0].mul_add(first_row[0], first_row[1] * first_row[1]); + let cross_dot = first_row[0].mul_add(second_row[0], first_row[1] * second_row[1]); + let second_self_dot = second_row[0].mul_add(second_row[0], second_row[1] * second_row[1]); + let upper_triangle_sentinel = -123.0_f32; + let expected = [ + scale * first_self_dot, + upper_triangle_sentinel, + scale * cross_dot, + scale * second_self_dot, ]; - let untouched = -123.0; - let mut c = [untouched; 9]; + let mut actual = [upper_triangle_sentinel; 4]; - sgemm_aat_lower(3, 2, &a, &mut c).unwrap(); + // When + sgemm_aat_lower(row_count, dimension_count, scale, &matrix, &mut actual).unwrap(); - #[rustfmt::skip] - assert_eq!(c, [ - 5.0, untouched, untouched, - 11.0, 25.0, untouched, - 17.0, 39.0, 61.0, - ]); + // Then + assert_eq!(actual, expected); } #[test] - fn zero_row_matrix_returns_success() { + fn add_accumulates_the_scaled_lower_triangle_and_preserves_the_upper_triangle() { // Given - let zero_rows = 0; - let dimensions = 3; - let no_input = []; - let mut no_output = []; + let first_row = [1.0_f32, 2.0]; + let second_row = [3.0_f32, 4.0]; + let row_count = 2; + let dimension_count = first_row.len(); + let matrix = [first_row[0], first_row[1], second_row[0], second_row[1]]; + let scale = -2.0_f32; + let first_self_dot = first_row[0].mul_add(first_row[0], first_row[1] * first_row[1]); + let cross_dot = first_row[0].mul_add(second_row[0], first_row[1] * second_row[1]); + let second_self_dot = second_row[0].mul_add(second_row[0], second_row[1] * second_row[1]); + let initial_first = 10.0_f32; + let initial_cross = 20.0_f32; + let initial_second = 30.0_f32; + let upper_triangle_sentinel = -123.0_f32; + let expected = [ + scale.mul_add(first_self_dot, initial_first), + upper_triangle_sentinel, + scale.mul_add(cross_dot, initial_cross), + scale.mul_add(second_self_dot, initial_second), + ]; + let mut actual = [ + initial_first, + upper_triangle_sentinel, + initial_cross, + initial_second, + ]; // When - let result = sgemm_aat_lower(zero_rows, dimensions, &no_input, &mut no_output); + sgemm_aat_lower_add(row_count, dimension_count, scale, &matrix, &mut actual).unwrap(); // Then - assert!(result.is_ok()); + assert_eq!(actual, expected); } #[test] - fn zero_inner_dimension_produces_zero_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]; + fn invalid_input_length_returns_the_a_matrix_error() { + // Given + let row_count = 2; + let dimension_count = 2; + let invalid_input = [0.0_f32; 3]; + let mut output = [0.0_f32; 4]; + let expected = SgemmError::InvalidMatrixDimensions { + matrix_name: MatrixName::A, + expected_rows: row_count, + expected_cols: dimension_count, + actual_len: invalid_input.len(), + }; - let error = sgemm_aat_lower(2, 2, &[0.0; 3], &mut c).unwrap_err(); + // When + let actual = sgemm_aat_lower(row_count, dimension_count, 1.0, &invalid_input, &mut output) + .unwrap_err(); - assert_eq!( - error, - SgemmError::InvalidMatrixDimensions { - matrix_name: MatrixName::A, - expected_rows: 2, - expected_cols: 2, - actual_len: 3, - } - ); + // Then + assert_eq!(actual, expected); } #[test] - fn rejects_invalid_output_dimensions() { - let mut c = [0.0; 3]; + fn invalid_output_length_returns_the_c_matrix_error() { + // Given + let row_count = 2; + let dimension_count = 2; + let input = [0.0_f32; 4]; + let mut invalid_output = [0.0_f32; 3]; + let expected = SgemmError::InvalidMatrixDimensions { + matrix_name: MatrixName::C, + expected_rows: row_count, + expected_cols: row_count, + actual_len: invalid_output.len(), + }; - let error = sgemm_aat_lower(2, 2, &[0.0; 4], &mut c).unwrap_err(); + // When + let actual = + sgemm_aat_lower_add(row_count, dimension_count, 1.0, &input, &mut invalid_output) + .unwrap_err(); - assert_eq!( - error, - SgemmError::InvalidMatrixDimensions { - matrix_name: MatrixName::C, - expected_rows: 2, - expected_cols: 2, - actual_len: 3, - } - ); + // Then + assert_eq!(actual, expected); } #[test] - fn rejects_input_size_overflow() { - let error = sgemm_aat_lower(usize::MAX, 2, &[], &mut []).unwrap_err(); + fn input_size_overflow_returns_the_a_matrix_error() { + // Given + let row_count = usize::MAX; + let dimension_count = 2; + let expected = SgemmError::DimensionOverflow { + matrix_name: MatrixName::A, + rows: row_count, + cols: dimension_count, + }; - assert_eq!( - error, - SgemmError::DimensionOverflow { - matrix_name: MatrixName::A, - rows: usize::MAX, - cols: 2, - } - ); + // When + let actual = sgemm_aat_lower(row_count, dimension_count, 1.0, &[], &mut []).unwrap_err(); + + // Then + assert_eq!(actual, expected); } #[test] - fn rejects_output_size_overflow() { - let error = sgemm_aat_lower(usize::MAX, 0, &[], &mut []).unwrap_err(); + fn output_size_overflow_returns_the_c_matrix_error() { + // Given + let row_count = usize::MAX; + let dimension_count = 0; + let expected = SgemmError::DimensionOverflow { + matrix_name: MatrixName::C, + rows: row_count, + cols: row_count, + }; - assert_eq!( - error, - SgemmError::DimensionOverflow { - matrix_name: MatrixName::C, - rows: usize::MAX, - cols: usize::MAX, - } - ); + // When + let actual = + sgemm_aat_lower_add(row_count, dimension_count, 1.0, &[], &mut []).unwrap_err(); + + // Then + assert_eq!(actual, expected); } } diff --git a/diskann-wide/src/arch/aarch64/f32x2_.rs b/diskann-wide/src/arch/aarch64/f32x2_.rs index f0b7431a87..318227ca5d 100644 --- a/diskann-wide/src/arch/aarch64/f32x2_.rs +++ b/diskann-wide/src/arch/aarch64/f32x2_.rs @@ -31,7 +31,6 @@ 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!( @@ -91,7 +90,6 @@ 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 83779dacf0..82cf391076 100644 --- a/diskann-wide/src/arch/aarch64/f32x4_.rs +++ b/diskann-wide/src/arch/aarch64/f32x4_.rs @@ -32,7 +32,6 @@ 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); @@ -188,7 +187,6 @@ 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 b93c861e7a..836193a452 100644 --- a/diskann-wide/src/arch/x86_64/v3/f32x16_.rs +++ b/diskann-wide/src/arch/x86_64/v3/f32x16_.rs @@ -54,7 +54,6 @@ 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 45cc64a4a2..60ffa4477c 100644 --- a/diskann-wide/src/arch/x86_64/v3/f32x4_.rs +++ b/diskann-wide/src/arch/x86_64/v3/f32x4_.rs @@ -31,7 +31,6 @@ 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)] @@ -254,7 +253,6 @@ 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 48ecea7aee..054b249e8f 100644 --- a/diskann-wide/src/arch/x86_64/v3/f32x8_.rs +++ b/diskann-wide/src/arch/x86_64/v3/f32x8_.rs @@ -33,7 +33,6 @@ 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)] @@ -267,7 +266,6 @@ 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 ff59119992..d38465f906 100644 --- a/diskann-wide/src/arch/x86_64/v4/f32x16_.rs +++ b/diskann-wide/src/arch/x86_64/v4/f32x16_.rs @@ -57,7 +57,6 @@ 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)] @@ -241,7 +240,6 @@ 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 7028b2fdcf..328dba4d26 100644 --- a/diskann-wide/src/arch/x86_64/v4/f32x4_.rs +++ b/diskann-wide/src/arch/x86_64/v4/f32x4_.rs @@ -33,7 +33,6 @@ 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)] @@ -211,7 +210,6 @@ 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 d38de49de8..3158ffc1dd 100644 --- a/diskann-wide/src/arch/x86_64/v4/f32x8_.rs +++ b/diskann-wide/src/arch/x86_64/v4/f32x8_.rs @@ -36,7 +36,6 @@ 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)] @@ -207,7 +206,6 @@ 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 d6adcb7b13..30d08e6cb1 100644 --- a/diskann-wide/src/doubled.rs +++ b/diskann-wide/src/doubled.rs @@ -205,15 +205,6 @@ 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 1e8d707769..507dc1288f 100644 --- a/diskann-wide/src/emulated.rs +++ b/diskann-wide/src/emulated.rs @@ -199,15 +199,6 @@ 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 @@ -911,10 +902,6 @@ 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 69bba4bc41..fc15661f1e 100644 --- a/diskann-wide/src/test_utils/ops.rs +++ b/diskann-wide/src/test_utils/ops.rs @@ -425,38 +425,6 @@ 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! { @@ -1173,7 +1141,6 @@ 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; diff --git a/diskann/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs deleted file mode 100644 index 333fb11eff..0000000000 --- a/diskann/src/graph/pipnn/kernel_metric.rs +++ /dev/null @@ -1,276 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! This module provides metric markers and shared numerical functions. - -mod leaf; -mod partition; - -pub(super) use leaf::LeafMetric; -pub(super) use partition::PartitionMetric; - -use super::simd::PiPNNSIMDVector; - -pub(super) struct L2; -pub(super) struct Cosine; -pub(super) struct CosineNormalized; -pub(super) struct InnerProduct; - -/// 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], -} - -/// Compute SIMD cosine distance with the DiskANN zero-norm and NaN rules. -/// -/// 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 - F: PiPNNSIMDVector, -{ - let zero = F::default(arch); - let one = F::splat(arch, 1.0); - let minimum_norm = F::splat(arch, f32::MIN_POSITIVE.sqrt()); - 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 = 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); - let distance = one - negative_one.max_simd(cosine.min_simd(one)); - F::select(cosine.ne_simd(cosine), cosine, distance) -} - -/// Compute one cosine distance with the DiskANN zero-norm and NaN rules. -#[inline(always)] -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 { - let cosine = dot / (source_norm * target_norm); - 1.0 - cosine.clamp(-1.0, 1.0) - } -} - -#[cfg(test)] -mod tests { - use super::cosine_distance_single; - - mod test_support { - use super::super::super::simd::PiPNNSIMDSchema; - use super::super::cosine_distance_simd; - use diskann_wide::{ARCH, SIMDVector, arch::Current}; - - pub(super) fn run_cosine_distance_simd( - dot_products: [f32; 16], - source_norms: [f32; 16], - target_norms: [f32; 16], - ) -> [f32; 16] { - type TestVector = ::Vector; - assert_eq!(16 % TestVector::LANES, 0); - let mut output = [0.0; 16]; - for first in (0..16).step_by(TestVector::LANES) { - // SAFETY: each offset starts one complete SIMD group in every array. - unsafe { - let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); - let sources = TestVector::load_simd(ARCH, source_norms.as_ptr().add(first)); - let targets = TestVector::load_simd(ARCH, target_norms.as_ptr().add(first)); - cosine_distance_simd(ARCH, dots, sources, targets) - .store_simd(output.as_mut_ptr().add(first)); - } - } - output - } - } - - mod cosine_distance_single_tests { - use super::cosine_distance_single; - - #[test] - fn zero_source_norm_takes_precedence_over_nan_dot_product() { - // 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(f32::NAN, source_norm, 2.0); - - // Then - assert_eq!(actual_distance, expected_one_minus_zero_similarity); - } - - #[test] - fn 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(0.0, 2.0, target_norm); - - // Then - assert_eq!(actual_distance, expected_one_minus_zero_similarity); - } - - #[test] - fn minimum_normal_norm_uses_normalized_similarity() { - // Given - let source_norm = f32::MIN_POSITIVE.sqrt(); - let target_norm = 1.0; - let dot_product = source_norm / 2.0; - let expected_distance = 1.0 - dot_product / (source_norm * target_norm); - - // When - let actual_distance = cosine_distance_single(dot_product, source_norm, target_norm); - - // Then - assert_eq!(actual_distance, expected_distance); - } - - #[test] - fn similarity_above_one_clamps_to_zero_distance() { - // Given - let dot_product_just_above_norm_product = 4.000_001; - 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_just_above_norm_product, 2.0, 2.0); - - // Then - assert_eq!(actual_distance, expected_one_minus_maximum_similarity); - } - - #[test] - fn similarity_below_negative_one_clamps_to_distance_two() { - // Given - let dot_product_just_below_negative_norm_product = -4.000_001; - 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_just_below_negative_norm_product, 2.0, 2.0); - - // Then - assert_eq!(actual_distance, expected_one_minus_minimum_similarity); - } - - #[test] - fn nan_similarity_remains_nan() { - let actual_distance = cosine_distance_single(f32::NAN, 1.0, 1.0); - assert!(actual_distance.is_nan()); - } - - #[test] - fn nan_source_norm_produces_nan_distance() { - let actual_distance = cosine_distance_single(0.0, f32::NAN, 1.0); - assert!(actual_distance.is_nan()); - } - - #[test] - fn nan_target_norm_produces_nan_distance() { - let actual_distance = cosine_distance_single(0.0, 1.0, f32::NAN); - assert!(actual_distance.is_nan()); - } - } - - mod cosine_distance_simd_tests { - use super::test_support::run_cosine_distance_simd; - - #[test] - fn zero_source_norm_takes_precedence_over_nan_dot_product_in_every_lane() { - let actual_distances = run_cosine_distance_simd([f32::NAN; 16], [0.0; 16], [2.0; 16]); - assert_eq!(actual_distances, [1.0; 16]); - } - - #[test] - fn zero_target_norm_produces_unit_distance_in_every_lane() { - let actual_distances = run_cosine_distance_simd([0.0; 16], [2.0; 16], [0.0; 16]); - assert_eq!(actual_distances, [1.0; 16]); - } - - #[test] - fn subnormal_norm_produces_unit_distance_in_every_lane() { - let subnormal_norm = f32::MIN_POSITIVE.sqrt() / 2.0; - let actual_distances = - run_cosine_distance_simd([0.0; 16], [2.0; 16], [subnormal_norm; 16]); - assert_eq!(actual_distances, [1.0; 16]); - } - - #[test] - fn minimum_normal_norm_uses_normalized_similarity_in_every_lane() { - let source_norm = f32::MIN_POSITIVE.sqrt(); - let target_norm = 1.0; - let dot_product = source_norm / 2.0; - let expected_distance = 1.0 - dot_product / (source_norm * target_norm); - - let actual_distances = - run_cosine_distance_simd([dot_product; 16], [source_norm; 16], [target_norm; 16]); - - assert_eq!(actual_distances, [expected_distance; 16]); - } - - #[test] - fn similarity_outside_bounds_is_clamped_in_every_lane() { - let dot_above_norm_product = 4.000_001; - let dot_below_negative_norm_product = -4.000_001; - let dot_products = [ - f32::INFINITY, - dot_above_norm_product, - dot_above_norm_product, - dot_above_norm_product, - dot_above_norm_product, - dot_above_norm_product, - dot_above_norm_product, - dot_above_norm_product, - f32::NEG_INFINITY, - dot_below_negative_norm_product, - dot_below_negative_norm_product, - dot_below_negative_norm_product, - dot_below_negative_norm_product, - dot_below_negative_norm_product, - dot_below_negative_norm_product, - dot_below_negative_norm_product, - ]; - let expected_distances = [ - 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, - ]; - - let actual_distances = run_cosine_distance_simd(dot_products, [2.0; 16], [2.0; 16]); - - assert_eq!(actual_distances, expected_distances); - } - - #[test] - fn nan_similarity_remains_nan_in_every_lane() { - let actual_distances = run_cosine_distance_simd([f32::NAN; 16], [1.0; 16], [1.0; 16]); - assert!(actual_distances.into_iter().all(f32::is_nan)); - } - - #[test] - fn nan_source_norm_produces_nan_distance_in_every_lane() { - let actual_distances = run_cosine_distance_simd([0.0; 16], [f32::NAN; 16], [1.0; 16]); - assert!(actual_distances.into_iter().all(f32::is_nan)); - } - - #[test] - fn nan_target_norm_produces_nan_distance_in_every_lane() { - let actual_distances = run_cosine_distance_simd([0.0; 16], [1.0; 16], [f32::NAN; 16]); - assert!(actual_distances.into_iter().all(f32::is_nan)); - } - } -} diff --git a/diskann/src/graph/pipnn/kernel_metric/leaf.rs b/diskann/src/graph/pipnn/kernel_metric/leaf.rs deleted file mode 100644 index d87a16c426..0000000000 --- a/diskann/src/graph/pipnn/kernel_metric/leaf.rs +++ /dev/null @@ -1,487 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use diskann_utils::views::MatrixView; -use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDPartialEq, 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 { - /// Prepare one contiguous metric-specific norm for each leaf-local point. - 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: A, _norms: &[f32], _source: usize) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::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 distances_simd( - arch: A, - norms: &[f32], - source_norms: A::Vector, - dot_products: A::Vector, - first_target: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema; - - /// Compute one distance outside the complete SIMD prefix. - fn distance_single(norms: &[f32], source_norm: f32, dot_product: f32, target: 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: PiPNNSIMDVector, -{ - 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 LeafMetric for L2 { - 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)]; - } - } - - #[inline(always)] - fn source_simd(arch: A, norms: &[f32], source: usize) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::splat(arch, norms[source]) - } - - #[inline(always)] - fn source_single(norms: &[f32], source: usize) -> f32 { - norms[source] - } - - #[inline(always)] - fn distances_simd( - arch: A, - norms: &[f32], - source_norms: A::Vector, - dot_products: A::Vector, - first_target: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema, - { - let target_norms = load_norms_simd::(arch, norms, first_target); - let distances = - A::Vector::splat(arch, -2.0).mul_add_simd(dot_products, source_norms) + target_norms; - let non_negative = distances.max_simd(A::Vector::default(arch)); - A::Vector::select(distances.ne_simd(distances), distances, non_negative) - } - - #[inline(always)] - fn distance_single(norms: &[f32], source_norm: f32, dot_product: f32, target: usize) -> f32 { - let distance = (-2.0_f32).mul_add(dot_product, source_norm) + norms[target]; - if distance.is_nan() { - distance - } else { - distance.max(0.0) - } - } -} - -impl LeafMetric for Cosine { - 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(); - } - } - - #[inline(always)] - fn source_simd(arch: A, norms: &[f32], source: usize) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::splat(arch, norms[source]) - } - - #[inline(always)] - fn source_single(norms: &[f32], source: usize) -> f32 { - norms[source] - } - - #[inline(always)] - fn distances_simd( - arch: A, - norms: &[f32], - source_norms: A::Vector, - dot_products: A::Vector, - first_target: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema, - { - let target_norms = load_norms_simd::(arch, norms, first_target); - cosine_distance_simd(arch, dot_products, source_norms, target_norms) - } - - #[inline(always)] - fn distance_single(norms: &[f32], source_norm: f32, dot_product: f32, target: usize) -> f32 { - cosine_distance_single(dot_product, source_norm, norms[target]) - } -} - -impl LeafMetric for CosineNormalized { - #[inline(always)] - fn distances_simd( - arch: A, - _norms: &[f32], - _source_norms: A::Vector, - dot_products: A::Vector, - _first_target: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::splat(arch, 1.0) - dot_products - } - - #[inline(always)] - fn distance_single(_norms: &[f32], _source_norm: f32, dot_product: f32, _target: usize) -> f32 { - 1.0 - dot_product - } -} - -impl LeafMetric for InnerProduct { - #[inline(always)] - fn distances_simd( - arch: A, - _norms: &[f32], - _source_norms: A::Vector, - dot_products: A::Vector, - _first_target: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::splat(arch, -1.0) * dot_products - } - - #[inline(always)] - fn distance_single(_norms: &[f32], _source_norm: f32, dot_product: f32, _target: usize) -> f32 { - -dot_product - } -} - -#[cfg(test)] -#[allow( - clippy::unwrap_used, - reason = "deterministic test matrices must abort on invalid setup" -)] -mod tests { - use super::*; - - mod test_support { - use super::*; - use diskann_wide::{ARCH, SIMDVector, arch::Current}; - - pub(super) fn run_distances_simd( - norms: &[f32], - source: usize, - dot_products: [f32; 16], - first_target: usize, - ) -> [f32; 16] { - type TestVector = ::Vector; - assert_eq!(16 % TestVector::LANES, 0); - let source_norms = M::source_simd(ARCH, norms, source); - let mut output = [0.0; 16]; - for first in (0..16).step_by(TestVector::LANES) { - // SAFETY: each offset starts one complete SIMD group in both arrays. - unsafe { - let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); - M::distances_simd(ARCH, norms, source_norms, dots, first_target + first) - .store_simd(output.as_mut_ptr().add(first)); - } - } - output - } - } - - mod prepare_leaf_norms_tests { - use super::*; - - #[test] - fn returns_the_squared_euclidean_norm_of_each_point_for_l2() { - // Given - let first_point = [2.0_f32, 1.0]; - let second_point = [1.0_f32, 3.0]; - let first_self_dot = first_point[0] * first_point[0] + first_point[1] * first_point[1]; - let cross_dot = first_point[0] * second_point[0] + first_point[1] * second_point[1]; - let second_self_dot = - second_point[0] * second_point[0] + second_point[1] * second_point[1]; - let gram_values = [first_self_dot, cross_dot, cross_dot, second_self_dot]; - let gram = MatrixView::try_from(&gram_values[..], 2, 2).unwrap(); - let expected_point_self_dots = [first_self_dot, second_self_dot]; - let mut actual_norms = Vec::new(); - - // When - L2::prepare_leaf_norms(gram, &mut actual_norms); - - // Then - assert_eq!(actual_norms, expected_point_self_dots); - } - - #[test] - fn returns_the_euclidean_norm_of_each_point_for_cosine() { - // Given - let first_point = [2.0_f32, 1.0]; - let second_point = [1.0_f32, 3.0]; - let first_self_dot = first_point[0] * first_point[0] + first_point[1] * first_point[1]; - let cross_dot = first_point[0] * second_point[0] + first_point[1] * second_point[1]; - let second_self_dot = - second_point[0] * second_point[0] + second_point[1] * second_point[1]; - let gram_values = [first_self_dot, cross_dot, cross_dot, second_self_dot]; - let gram = MatrixView::try_from(&gram_values[..], 2, 2).unwrap(); - let expected_point_l2_norms = [first_self_dot.sqrt(), second_self_dot.sqrt()]; - let mut actual_norms = Vec::new(); - - // When - Cosine::prepare_leaf_norms(gram, &mut actual_norms); - - // Then - assert_eq!(actual_norms, expected_point_l2_norms); - } - } - - mod distance_single_tests { - use super::*; - - #[test] - fn distance_equals_squared_norm_sum_minus_twice_the_dot_product_with_l2() { - // 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 negative_roundoff_is_clamped_to_zero_with_l2() { - // 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 nan_dot_product_produces_nan_distance_with_l2() { - // Given - let source_squared_norm = 1.0; - let target_squared_norm = 1.0; - let squared_norms = [source_squared_norm, target_squared_norm]; - - // When - let actual_distance = - L2::distance_single(&squared_norms, source_squared_norm, f32::NAN, 1); - - // Then - assert!(actual_distance.is_nan()); - } - - #[test] - fn distance_equals_one_minus_dot_over_norm_product_with_cosine() { - // 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 nan_dot_product_produces_nan_distance_with_cosine() { - // Given - let norms = [1.0, 1.0]; - - // When - let actual_distance = Cosine::distance_single(&norms, 1.0, f32::NAN, 1); - - // Then - assert!(actual_distance.is_nan()); - } - - #[test] - fn distance_equals_one_minus_the_dot_product_with_normalized_cosine() { - // 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 distance_equals_the_negative_dot_product_with_inner_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); - } - } - - mod distances_simd_tests { - use super::test_support::run_distances_simd; - use super::*; - - #[test] - fn every_lane_uses_squared_norm_sum_minus_twice_the_dot_product_with_l2() { - let source_squared_norm = 4.0; - let target_squared_norm = 9.0; - let dot_product = 6.0; - let expected_distance = source_squared_norm + target_squared_norm - 2.0 * dot_product; - let mut squared_norms = [target_squared_norm; 17]; - squared_norms[0] = source_squared_norm; - - let actual_distances = - run_distances_simd::(&squared_norms, 0, [dot_product; 16], 1); - - assert_eq!(actual_distances, [expected_distance; 16]); - } - - #[test] - fn negative_roundoff_is_clamped_to_zero_in_every_lane_with_l2() { - let squared_norms = [1.0; 17]; - let dot_product_above_exact_norm = 1.000_001; - - let actual_distances = - run_distances_simd::(&squared_norms, 0, [dot_product_above_exact_norm; 16], 1); - - assert_eq!(actual_distances, [0.0; 16]); - } - - #[test] - fn nan_dot_products_remain_nan_in_every_lane_with_l2() { - let squared_norms = [1.0; 17]; - - let actual_distances = run_distances_simd::(&squared_norms, 0, [f32::NAN; 16], 1); - - assert!(actual_distances.into_iter().all(f32::is_nan)); - } - - #[test] - fn every_lane_uses_one_minus_normalized_dot_with_cosine() { - let source_norm = 2.0; - let target_norm = 4.0; - let dot_product = 4.0; - let expected_distance = 1.0 - dot_product / (source_norm * target_norm); - let mut norms = [target_norm; 17]; - norms[0] = source_norm; - - let actual_distances = run_distances_simd::(&norms, 0, [dot_product; 16], 1); - - assert_eq!(actual_distances, [expected_distance; 16]); - } - - #[test] - fn nan_dot_products_remain_nan_in_every_lane_with_cosine() { - let norms = [1.0; 17]; - - let actual_distances = run_distances_simd::(&norms, 0, [f32::NAN; 16], 1); - - assert!(actual_distances.into_iter().all(f32::is_nan)); - } - - #[test] - fn every_lane_uses_one_minus_dot_product_with_normalized_cosine() { - let dot_product = 0.25; - let expected_distance = 1.0 - dot_product; - - let actual_distances = - run_distances_simd::(&[], 0, [dot_product; 16], 0); - - assert_eq!(actual_distances, [expected_distance; 16]); - } - - #[test] - fn every_lane_uses_negative_dot_product_with_inner_product() { - let dot_product = 3.0; - let expected_distance = -dot_product; - - let actual_distances = run_distances_simd::(&[], 0, [dot_product; 16], 0); - - assert_eq!(actual_distances, [expected_distance; 16]); - } - - #[test] - fn signed_zero_bits_match_scalar_negation_with_inner_product() { - let dot_products: [f32; 16] = [ - 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, - -0.0, - ]; - let expected_bits = dot_products.map(|dot| (-dot).to_bits()); - - let actual_bits = - run_distances_simd::(&[], 0, dot_products, 0).map(f32::to_bits); - - assert_eq!(actual_bits, expected_bits); - } - } -} diff --git a/diskann/src/graph/pipnn/kernel_metric/partition.rs b/diskann/src/graph/pipnn/kernel_metric/partition.rs deleted file mode 100644 index cc185193e2..0000000000 --- a/diskann/src/graph/pipnn/kernel_metric/partition.rs +++ /dev/null @@ -1,506 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use diskann_utils::views::MatrixView; -use diskann_vector::{Norm, norm::FastL2NormSquared}; -use diskann_wide::{SIMDMulAdd, SIMDVector}; - -use super::super::simd::{PiPNNSIMDSchema, PiPNNSIMDVector}; -use super::{ - Cosine, CosineNormalized, InnerProduct, L2, PartitionNorms, cosine_distance_simd, - cosine_distance_single, -}; - -/// 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(_points: MatrixView<'_, f32>, norms: &mut Vec) { - norms.clear(); - } - - /// Prepare one norm value for each sampled leader. - 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: A, _norms: PartitionNorms<'_>, _point: usize) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::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 rankings_simd( - arch: A, - norms: PartitionNorms<'_>, - point_norms: A::Vector, - dot_products: A::Vector, - first_leader: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema; - - /// Compute one ranking outside the complete SIMD prefix. - 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: PiPNNSIMDVector, -{ - 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()) } -} - -/// Prepare cosine norms with one reduction order for points and leaders. -fn prepare_cosine_norms(vectors: MatrixView<'_, f32>, norms: &mut Vec) { - norms.resize(vectors.nrows(), 0.0); - for (norm, vector) in norms.iter_mut().zip(vectors.row_iter()) { - *norm = FastL2NormSquared.evaluate(vector).sqrt(); - } -} - -impl PartitionMetric for L2 { - 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(); - } - } - - #[inline(always)] - fn rankings_simd( - arch: A, - norms: PartitionNorms<'_>, - _point_norms: A::Vector, - dot_products: A::Vector, - first_leader: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema, - { - let leader_norms = load_norms_simd::(arch, norms.leader_norms, first_leader); - A::Vector::splat(arch, -2.0).mul_add_simd(dot_products, leader_norms) - } - - #[inline(always)] - 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(points: MatrixView<'_, f32>, norms: &mut Vec) { - prepare_cosine_norms(points, norms); - } - - fn prepare_leader_norms(leaders: MatrixView<'_, f32>, norms: &mut Vec) { - prepare_cosine_norms(leaders, norms); - } - - #[inline(always)] - fn point_simd(arch: A, norms: PartitionNorms<'_>, point: usize) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::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: A, - norms: PartitionNorms<'_>, - point_norms: A::Vector, - dot_products: A::Vector, - first_leader: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema, - { - 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 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 { - #[inline(always)] - fn rankings_simd( - arch: A, - _norms: PartitionNorms<'_>, - _point_norms: A::Vector, - dot_products: A::Vector, - _first_leader: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::splat(arch, 1.0) - dot_products - } - - #[inline(always)] - fn ranking_single( - _norms: PartitionNorms<'_>, - _point_norm: f32, - dot_product: f32, - _leader: usize, - ) -> f32 { - 1.0 - dot_product - } -} - -impl PartitionMetric for InnerProduct { - #[inline(always)] - fn rankings_simd( - arch: A, - _norms: PartitionNorms<'_>, - _point_norms: A::Vector, - dot_products: A::Vector, - _first_leader: usize, - ) -> A::Vector - where - A: PiPNNSIMDSchema, - { - A::Vector::splat(arch, -1.0) * dot_products - } - - #[inline(always)] - fn ranking_single( - _norms: PartitionNorms<'_>, - _point_norm: f32, - dot_product: f32, - _leader: usize, - ) -> f32 { - -dot_product - } -} - -#[cfg(test)] -#[allow( - clippy::unwrap_used, - reason = "deterministic test matrices must abort on invalid setup" -)] -mod tests { - use super::*; - - mod test_support { - use super::*; - use diskann_wide::{ARCH, SIMDVector, arch::Current}; - - pub(super) fn run_rankings_simd( - point_norms: &[f32], - leader_norms: &[f32], - point: usize, - dot_products: [f32; 16], - first_leader: usize, - ) -> [f32; 16] { - type TestVector = ::Vector; - assert_eq!(16 % TestVector::LANES, 0); - let norms = PartitionNorms { - point_norms, - leader_norms, - }; - let point_norm = M::point_simd(ARCH, norms, point); - let mut output = [0.0; 16]; - for first in (0..16).step_by(TestVector::LANES) { - // SAFETY: each offset starts one complete SIMD group in both arrays. - unsafe { - let dots = TestVector::load_simd(ARCH, dot_products.as_ptr().add(first)); - M::rankings_simd(ARCH, norms, point_norm, dots, first_leader + first) - .store_simd(output.as_mut_ptr().add(first)); - } - } - output - } - - pub(super) fn run_ranking_single( - 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 returns_the_squared_euclidean_norm_of_each_leader_for_l2() { - // 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] * first_leader[0] + first_leader[1] * first_leader[1], - second_leader[0] * second_leader[0] + second_leader[1] * second_leader[1], - ]; - 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 returns_the_euclidean_norm_of_each_leader_for_cosine() { - // 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] * first_leader[0] + first_leader[1] * first_leader[1]).sqrt(), - (second_leader[0] * second_leader[0] + second_leader[1] * second_leader[1]).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 returns_the_euclidean_norm_of_each_point_for_cosine() { - // 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] * first_point[0] + first_point[1] * first_point[1]).sqrt(), - (second_point[0] * second_point[0] + second_point[1] * second_point[1]).sqrt(), - ]; - let mut actual_norms = Vec::new(); - - // When - Cosine::prepare_point_norms(points, &mut actual_norms); - - // Then - assert_eq!(actual_norms, expected_row_norms); - } - - #[test] - fn same_vector_has_identical_point_and_leader_norm_bits_with_cosine() { - // Given - const REASSOCIATION_BOUNDARY_DIMENSIONS: usize = 129; - let mut vector = [1.0_f32; REASSOCIATION_BOUNDARY_DIMENSIONS]; - vector[0] = 4096.0; - let vectors = - MatrixView::try_from(&vector[..], 1, REASSOCIATION_BOUNDARY_DIMENSIONS).unwrap(); - let mut point_norms = Vec::new(); - let mut leader_norms = Vec::new(); - - // When - Cosine::prepare_point_norms(vectors, &mut point_norms); - Cosine::prepare_leader_norms(vectors, &mut leader_norms); - - // Then - assert_eq!(point_norms[0].to_bits(), leader_norms[0].to_bits()); - } - - mod ranking_single_tests { - use super::test_support::run_ranking_single; - use super::*; - - #[test] - fn ranking_equals_leader_squared_norm_minus_twice_the_dot_product_with_l2() { - // 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 = run_ranking_single::(dot_product, 0.0, leader_squared_norm); - - // Then - assert_eq!(actual_ranking, expected_leader_norm_minus_twice_dot); - } - - #[test] - fn ranking_equals_one_minus_dot_over_norm_product_with_cosine() { - // 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 = run_ranking_single::(dot_product, point_norm, leader_norm); - - // Then - assert_eq!(actual_ranking, expected_one_minus_normalized_dot); - } - - #[test] - fn ranking_equals_one_minus_the_dot_product_with_normalized_cosine() { - // Given - let dot_product = 0.25; - let expected_one_minus_dot = 1.0 - dot_product; - - // When - let actual_ranking = run_ranking_single::(dot_product, 0.0, 0.0); - - // Then - assert_eq!(actual_ranking, expected_one_minus_dot); - } - - #[test] - fn ranking_equals_the_negative_dot_product_with_inner_product() { - // Given - let dot_product = 3.0; - let expected_negative_dot = -dot_product; - - // When - let actual_ranking = run_ranking_single::(dot_product, 0.0, 0.0); - - // Then - assert_eq!(actual_ranking, expected_negative_dot); - } - } - - mod rankings_simd_tests { - use super::test_support::run_rankings_simd; - use super::*; - - #[test] - fn every_lane_uses_leader_squared_norm_minus_twice_the_dot_product_with_l2() { - let leader_squared_norm = 9.0; - let dot_product = 2.0; - let expected_ranking = leader_squared_norm - 2.0 * dot_product; - let actual_rankings = - run_rankings_simd::(&[], &[leader_squared_norm; 16], 0, [dot_product; 16], 0); - assert_eq!(actual_rankings, [expected_ranking; 16]); - } - - #[test] - fn every_lane_uses_one_minus_normalized_dot_with_cosine() { - let point_norm = 2.0; - let leader_norm = 4.0; - let dot_product = 4.0; - let expected_ranking = 1.0 - dot_product / (point_norm * leader_norm); - let actual_rankings = run_rankings_simd::( - &[point_norm], - &[leader_norm; 16], - 0, - [dot_product; 16], - 0, - ); - assert_eq!(actual_rankings, [expected_ranking; 16]); - } - - #[test] - fn every_lane_uses_one_minus_dot_product_with_normalized_cosine() { - let dot_product = 0.25; - let expected_ranking = 1.0 - dot_product; - let actual_rankings = - run_rankings_simd::(&[], &[], 0, [dot_product; 16], 0); - assert_eq!(actual_rankings, [expected_ranking; 16]); - } - - #[test] - fn every_lane_uses_negative_dot_product_with_inner_product() { - let dot_product = 3.0; - let expected_ranking = -dot_product; - let actual_rankings = - run_rankings_simd::(&[], &[], 0, [dot_product; 16], 0); - assert_eq!(actual_rankings, [expected_ranking; 16]); - } - - #[test] - fn signed_zero_bits_match_scalar_negation_with_inner_product() { - let dot_products: [f32; 16] = [ - 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, -0.0, 0.0, - -0.0, - ]; - let expected_bits = dot_products.map(|dot| (-dot).to_bits()); - - let actual_bits = - run_rankings_simd::(&[], &[], 0, dot_products, 0).map(f32::to_bits); - - assert_eq!(actual_bits, expected_bits); - } - - #[test] - fn every_lane_stays_finite_when_twice_the_dot_product_overflows_with_l2() { - let dot_product = f32::from_bits(f32::MAX.to_bits() - 1); - let leader_squared_norm = f32::MAX; - let unfused_twice_dot_product = 2.0 * dot_product; - let expected_fused_ranking = (-2.0_f32).mul_add(dot_product, leader_squared_norm); - - let actual_rankings = - run_rankings_simd::(&[], &[leader_squared_norm; 16], 0, [dot_product; 16], 0); - - assert!(unfused_twice_dot_product.is_infinite()); - assert!(expected_fused_ranking.is_finite()); - assert_eq!(actual_rankings, [expected_fused_ranking; 16]); - } - } -} diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index e5c723e544..1c92696704 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -5,9 +5,8 @@ //! Leaf-local top-k selection from packed `f32` point vectors. //! -//! 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 metric fills a flattened lower-triangular distance buffer. The kernel +//! 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. Widths 1 through 3 use fixed insertion. Larger @@ -17,8 +16,7 @@ //! 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. +//! The caller supplies concrete architecture `A` and metric `M`. //! [`LeafKernelWorkspace`] stores reusable numerical scratch. use crate::{ANNError, ANNResult}; @@ -26,7 +24,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{SIMDPartialOrd, SIMDVector}; use super::{ - kernel_metric::LeafMetric, + leaf_metric::LeafMetric, simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, }; @@ -42,8 +40,8 @@ pub(super) struct LeafNeighbor { impl LeafNeighbor { /// Construct a leaf-local neighbor. /// - /// `target` is a position in the leaf. `distance` is its score relative to - /// the source of the output row. + /// `target` is a position in the leaf. `distance` is measured from the + /// source point of the output row. pub(super) const fn new(target: u32, distance: f32) -> Self { Self { target, distance } } @@ -63,12 +61,11 @@ impl Default for LeafNeighbor { /// Reusable storage for one leaf numerical pipeline. #[derive(Debug, Default)] pub(super) struct LeafKernelWorkspace { - dot_scratch: Vec, - norm_scratch: Vec, + distance_scratch: Vec, worst: Vec, } -/// Validation error returned by the dot-ranking loop. +/// Validation error returned by the distance-ranking loop. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub(super) enum LeafKernelError { /// A source requests more neighbors than the leaf or fixed kernel supports. @@ -105,58 +102,75 @@ where M: LeafMetric, { let point_count = points.nrows(); - let dot_count = point_count * point_count; + let distance_count = point_count * point_count; let LeafKernelWorkspace { - dot_scratch, - norm_scratch, + distance_scratch, worst, } = workspace; - if dot_scratch.len() < dot_count { - dot_scratch.resize(dot_count, 0.0); + if distance_scratch.len() < distance_count { + distance_scratch.resize(distance_count, 0.0); } - diskann_linalg::sgemm_aat_lower( + M::compute_distances(arch, points, &mut distance_scratch[..distance_count])?; + rank_leaf_distances( + arch, + &distance_scratch[..distance_count], point_count, - points.ncols(), - points.as_slice(), - &mut dot_scratch[..dot_count], + output, + worst, ) - .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) + .map_err(ANNError::new) } -/// Rank a prepared lower-triangular Gram matrix. -fn rank_leaf_dots( +/// Rank one flattened lower-triangle buffer. +/// +/// `distance_flatten` contains `point_count * point_count` elements. The metric +/// initializes each strict-lower entry. The kernel does not read the upper triangle. +fn rank_leaf_distances( arch: A, - input: MatrixView<'_, f32>, - norms: &[f32], + distance_flatten: &[f32], + point_count: usize, mut output: MutMatrixView<'_, LeafNeighbor>, worst: &mut Vec, ) -> Result<(), LeafKernelError> where A: PiPNNSIMDSchema, - M: LeafMetric, { - validate_neighbor_count(input, &output)?; + validate_neighbor_count(point_count, &output)?; let neighbor_count = output.ncols(); if neighbor_count == 0 { return Ok(()); } - worst.resize(input.nrows(), f32::INFINITY); + worst.resize(point_count, f32::INFINITY); output.as_mut_slice().fill(LeafNeighbor::default()); 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, + distance_flatten, + point_count, + output.as_mut_slice(), + worst, + ), + 2 => scan_fixed_width::( + arch, + distance_flatten, + point_count, + output.as_mut_slice(), + worst, + ), + 3 => scan_fixed_width::( + arch, + distance_flatten, + point_count, + output.as_mut_slice(), + worst, + ), + _ => scan_runtime_width( arch, - input, - norms, + distance_flatten, + point_count, output.as_mut_slice(), neighbor_count, worst, @@ -169,10 +183,9 @@ where /// /// Check the output width against the number of non-self points. fn validate_neighbor_count( - input: MatrixView<'_, f32>, + point_count: usize, output: &MutMatrixView<'_, LeafNeighbor>, ) -> Result<(), LeafKernelError> { - let point_count = input.nrows(); let maximum_neighbors = point_count.saturating_sub(1); let neighbor_count = output.ncols(); if neighbor_count > maximum_neighbors { @@ -186,96 +199,101 @@ fn validate_neighbor_count( } /// Select neighbors with a fixed output width. -fn scan_fixed_width( +fn scan_fixed_width( arch: A, - input: MatrixView<'_, f32>, - norms: &[f32], + distance_flatten: &[f32], + point_count: usize, output: &mut [LeafNeighbor], worst: &mut [f32], ) where A: PiPNNSIMDSchema, - M: LeafMetric, [LeafNeighbor; N]: SortedInsert, { let (rows, _) = output.as_chunks_mut::(); // Rayon outlines leaf workers. Reapply target features before the SIMD scan. arch.run(move || { - scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { - insert_eligible_neighbor(&mut rows[source], target, distance) - }); + scan_point_pairs( + arch, + distance_flatten, + point_count, + worst, + |source, target, distance| { + insert_eligible_neighbor(&mut rows[source], target, distance) + }, + ); }); } /// Select neighbors with a runtime output width. -fn scan_runtime_width( +fn scan_runtime_width( arch: A, - input: MatrixView<'_, f32>, - norms: &[f32], + distance_flatten: &[f32], + point_count: usize, output: &mut [LeafNeighbor], width: usize, worst: &mut [f32], ) where A: PiPNNSIMDSchema, - M: LeafMetric, { // Rayon outlines leaf workers. Reapply target features before the SIMD scan. arch.run(move || { - scan_point_pairs::(arch, input, norms, worst, |source, target, distance| { - let first = source * width; - insert_eligible_neighbor(&mut output[first..first + width], target, distance) - }); + scan_point_pairs( + arch, + distance_flatten, + point_count, + worst, + |source, target, distance| { + let first = source * width; + insert_eligible_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. +/// both endpoint lists. SIMD groups and scalar tails preserve pair scan order. #[inline(always)] -fn scan_point_pairs( +fn scan_point_pairs( arch: A, - input: MatrixView<'_, f32>, - norms: &[f32], + distance_flatten: &[f32], + point_count: usize, worst: &mut [f32], mut insert: I, ) where A: PiPNNSIMDSchema, - M: LeafMetric, I: FnMut(usize, u32, f32) -> f32, { - let point_count = input.nrows(); - let dots = input.as_slice(); let worst_ptr = worst.as_mut_ptr(); for source in 1..point_count { let source_start = source * point_count; - 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. + // SAFETY: `rank_leaf_distances` created one threshold for each point. let mut source_worst = unsafe { *worst_ptr.add(source) }; let mut target = 0; let simd_prefix = source - source % A::Vector::LANES; while target < simd_prefix { // SAFETY: This complete SIMD group is in the strict-lower prefix. - let dot_products = - unsafe { A::Vector::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(A::Vector::splat(arch, source_worst)); + let distance_group = unsafe { + A::Vector::load_simd(arch, distance_flatten.as_ptr().add(source_start + target)) + }; + let source_eligible = distance_group.lt_simd(A::Vector::splat(arch, source_worst)); // SAFETY: The complete target group is below `source < point_count`. let target_worst = unsafe { A::Vector::load_simd(arch, worst_ptr.add(target)) }; - let target_eligible = distances.lt_simd(target_worst); + let target_eligible = distance_group.lt_simd(target_worst); let source_bits = A::Vector::active_lanes(source_eligible); let target_bits = A::Vector::active_lanes(target_eligible); if source_bits | target_bits != 0 { - let values = distances.to_lane_array(); - let values = values.as_ref(); + let distance_lanes = distance_group.to_lane_array(); + let distance_lanes = distance_lanes.as_ref(); 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]; + let distance = distance_lanes[lane]; if distance < source_worst { source_worst = insert(source, (target + lane) as u32, distance); } @@ -286,7 +304,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(target_source, source as u32, values[lane]); + let new_worst = insert(target_source, source as u32, distance_lanes[lane]); // SAFETY: `target_source < source < worst.len()`. unsafe { *worst_ptr.add(target_source) = new_worst }; } @@ -296,8 +314,7 @@ fn scan_point_pairs( while target < source { // SAFETY: The target is in this source's strict-lower prefix. - let dot_product = unsafe { *dots.get_unchecked(source_start + target) }; - let distance = M::distance_single(norms, source_single, dot_product, target); + let distance = unsafe { *distance_flatten.get_unchecked(source_start + target) }; if distance < source_worst { source_worst = insert(source, target as u32, distance); } @@ -403,7 +420,7 @@ where #[cfg(test)] mod tests { use super::*; - use crate::graph::pipnn::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2}; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -413,122 +430,126 @@ mod tests { use super::*; use diskann_wide::arch::{self, Target1}; + pub(super) trait TestMetric { + const METRIC: Metric; + } + + impl TestMetric for L2 { + const METRIC: Metric = Metric::L2; + } + + impl TestMetric for Cosine { + const METRIC: Metric = Metric::Cosine; + } + + impl TestMetric for CosineNormalized { + const METRIC: Metric = Metric::CosineNormalized; + } + + impl TestMetric for InnerProduct { + const METRIC: Metric = Metric::InnerProduct; + } + struct KernelCall<'a> { - input: MatrixView<'a, f32>, - norms: &'a [f32], + distance_flatten: &'a [f32], + point_count: usize, output: MutMatrixView<'a, LeafNeighbor>, - workspace: &'a mut LeafKernelWorkspace, + worst: &'a mut Vec, } - struct DispatchMetric(Metric); + struct RankDistances; - impl Target1, KernelCall<'_>> for DispatchMetric + impl Target1, KernelCall<'_>> for RankDistances 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, - ), - } + rank_leaf_distances( + arch, + call.distance_flatten, + call.point_count, + call.output, + call.worst, + ) } } - fn lower_gram_view(dots: &[f32], points: usize) -> MatrixView<'_, f32> { - MatrixView::try_from(dots, points, points).unwrap() - } - - 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 - } - + fn reference_distance( + metric: Metric, + dot: f32, + source_diagonal: f32, + target_diagonal: f32, + ) -> f32 { match metric { - Metric::L2 => prepare::(lower_gram), - Metric::Cosine => prepare::(lower_gram), - Metric::CosineNormalized => prepare::(lower_gram), - Metric::InnerProduct => prepare::(lower_gram), + Metric::L2 => (-2.0_f32).mul_add(dot, source_diagonal) + target_diagonal, + Metric::CosineNormalized => -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 { + 1.0 - (dot / (source_norm * target_norm)).clamp(-1.0, 1.0) + } + } } } - pub(super) fn with_rank_leaf_dots_fixture( - dots: &[f32], - points: usize, - output_width: usize, - run: F, - ) -> Result, LeafKernelError> - where - M: LeafMetric, - F: FnOnce( - MatrixView<'_, f32>, - &[f32], - MutMatrixView<'_, LeafNeighbor>, - &mut Vec, - ) -> Result<(), LeafKernelError>, - { - let input = lower_gram_view(dots, points); - let mut norms = Vec::new(); - M::prepare_leaf_norms(input, &mut norms); - let mut output = vec![LeafNeighbor::default(); points * output_width]; - let output_view = - MutMatrixView::try_from(output.as_mut_slice(), points, output_width).unwrap(); - run(input, &norms, output_view, &mut Vec::new())?; - Ok(output) + fn distance_flatten(metric: Metric, dots: &[f32], points: usize) -> Vec { + let mut distances = vec![0.0; points * points]; + for source in 0..points { + for target in 0..=source { + distances[source * points + target] = reference_distance( + metric, + dots[source * points + target], + dots[source * points + source], + dots[target * points + target], + ); + } + } + distances } - /// Run the production ranker through runtime architecture and metric dispatch. - fn run_rank_leaf_dots_with_output_width( + fn rank_with_output_width( metric: Metric, dots: &[f32], points: usize, output_width: usize, workspace: &mut LeafKernelWorkspace, ) -> Result, LeafKernelError> { - let lower_gram = lower_gram_view(dots, points); - let norms = metric_norms(metric, lower_gram); + let distance_flatten = distance_flatten(metric, dots, points); let mut output = vec![LeafNeighbor::default(); points * output_width]; arch::dispatch1_no_features( - DispatchMetric(metric), + RankDistances, KernelCall { - input: lower_gram, - norms: &norms, + distance_flatten: &distance_flatten, + point_count: points, output: MutMatrixView::try_from(output.as_mut_slice(), points, output_width) .unwrap(), - workspace, + worst: &mut workspace.worst, }, )?; Ok(output) } - fn run_rank_leaf_dots_with_workspace( + pub(super) fn rank_distance_fixture( + dots: &[f32], + points: usize, + output_width: usize, + ) -> Result, LeafKernelError> { + rank_with_output_width( + M::METRIC, + dots, + points, + output_width, + &mut LeafKernelWorkspace::default(), + ) + } + + fn run_with_workspace( metric: Metric, dots: &[f32], points: usize, @@ -536,19 +557,18 @@ mod tests { workspace: &mut LeafKernelWorkspace, ) -> (usize, Vec) { let leaf_k = leaf_neighbor_count(points, requested_k); - let output = - run_rank_leaf_dots_with_output_width(metric, dots, points, leaf_k, workspace) - .expect("valid leaf neighbor width"); + let output = rank_with_output_width(metric, dots, points, leaf_k, workspace) + .expect("valid leaf neighbor width"); (leaf_k, output) } - pub(super) fn run_rank_leaf_dots( + pub(super) fn run_rank_leaf_distances( metric: Metric, dots: &[f32], points: usize, requested_k: usize, ) -> (usize, Vec) { - run_rank_leaf_dots_with_workspace( + run_with_workspace( metric, dots, points, @@ -557,31 +577,6 @@ mod tests { ) } - 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 { - let similarity = dot / (source_norm * target_norm); - 1.0 - similarity.clamp(-1.0, 1.0) - } - } - } - } - pub(super) fn reference_neighbors( metric: Metric, dots: &[f32], @@ -620,9 +615,6 @@ mod tests { } /// Build a Gram matrix from points on the line `x = 1`. - /// - /// Normalized cosine receives unit vectors because that metric assumes - /// normalized source data. Other metrics receive the original vectors. pub(super) fn lane_boundary_gram_from_point_vectors( metric: Metric, points: usize, @@ -851,20 +843,6 @@ mod tests { mod leaf_neighbor_count_tests { use super::leaf_neighbor_count; - #[test] - fn returns_zero_when_the_leaf_is_empty() { - // 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 returns_zero_when_the_leaf_contains_only_the_source() { // Given @@ -1013,7 +991,7 @@ mod tests { } } - mod rank_leaf_dots_tests { + mod rank_leaf_distances_tests { use super::test_support::*; use super::*; use rstest::rstest; @@ -1052,7 +1030,8 @@ mod tests { let expected_neighbors = reference_neighbors(metric, &dots, point_count, requested_k); // When - let actual_neighbors = run_rank_leaf_dots(metric, &dots, point_count, requested_k).1; + let actual_neighbors = + run_rank_leaf_distances(metric, &dots, point_count, requested_k).1; // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1092,21 +1071,7 @@ mod tests { let gram = three_unit_point_gram(); // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 3, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, L2>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 3, 1).unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1123,21 +1088,7 @@ mod tests { let gram = three_unit_point_gram(); // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 3, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, Cosine>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 3, 1).unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1147,28 +1098,14 @@ mod tests { fn selects_the_highest_dot_product_neighbor_for_each_point_with_normalized_cosine() { // Given let expected_neighbors = [ - LeafNeighbor::new(2, 1.0 - point_dot(POINT_ZERO, POINT_TWO)), - LeafNeighbor::new(2, 1.0 - point_dot(POINT_ONE, POINT_TWO)), - LeafNeighbor::new(1, 1.0 - point_dot(POINT_ONE, POINT_TWO)), + LeafNeighbor::new(2, -point_dot(POINT_ZERO, POINT_TWO)), + LeafNeighbor::new(2, -point_dot(POINT_ONE, POINT_TWO)), + LeafNeighbor::new(1, -point_dot(POINT_ONE, POINT_TWO)), ]; let gram = three_unit_point_gram(); // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 3, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, CosineNormalized>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 3, 1).unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1185,21 +1122,7 @@ mod tests { let gram = three_unit_point_gram(); // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 3, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, InnerProduct>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 3, 1).unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -1231,26 +1154,36 @@ mod tests { ]; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 4, - 2, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, L2>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 4, 2).unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors_in_scan_order); } + #[test] + fn finite_negative_l2_ranking_remains_rankable() { + // Given + let point_count = 2; + let nearest_neighbor_count = 1; + let self_dot = 1.0_f32; + let dot_roundoff = f32::EPSILON; + let cross_dot = self_dot + dot_roundoff; + let expected_distance = (-2.0_f32).mul_add(cross_dot, self_dot) + self_dot; + let gram = [self_dot, cross_dot, cross_dot, self_dot]; + let expected = [ + LeafNeighbor::new(1, expected_distance), + LeafNeighbor::new(0, expected_distance), + ]; + + // When + let actual = + rank_distance_fixture::(&gram, point_count, nearest_neighbor_count).unwrap(); + + // Then + assert!(expected_distance < 0.0); + assert_eq!(actual, expected); + } + #[test] fn scalar_distance_stays_finite_when_twice_the_dot_product_overflows_with_l2() { // Given @@ -1260,21 +1193,7 @@ mod tests { let gram = [f32::MAX, dot_product, dot_product, f32::MAX]; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 2, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, L2>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); // Then assert!(unfused_twice_dot_product.is_infinite()); @@ -1298,7 +1217,7 @@ mod tests { gram[16] = dot_product; // When - let actual_neighbors = run_rank_leaf_dots(Metric::L2, &gram, points, 1).1; + let actual_neighbors = run_rank_leaf_distances(Metric::L2, &gram, points, 1).1; // Then assert!(unfused_twice_dot_product.is_infinite()); @@ -1323,21 +1242,7 @@ mod tests { [LeafNeighbor::new(1, 1.0), LeafNeighbor::new(2, 1.0)]; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 3, - 2, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, Cosine>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 3, 2).unwrap(); // Then assert_eq!(&actual_neighbors[..2], &expected_zero_norm_neighbors); @@ -1352,7 +1257,7 @@ mod tests { let expected_zero_norm_neighbor = LeafNeighbor::new(0, 1.0); // When - let actual_neighbors = run_rank_leaf_dots(Metric::Cosine, &gram, points, 1).1; + let actual_neighbors = run_rank_leaf_distances(Metric::Cosine, &gram, points, 1).1; // Then assert_eq!(actual_neighbors[16], expected_zero_norm_neighbor); @@ -1368,21 +1273,7 @@ mod tests { let expected_one_minus_maximum_similarity = 1.0 - maximum_cosine_similarity; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 2, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, Cosine>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); // Then assert_eq!( @@ -1401,21 +1292,7 @@ mod tests { let expected_one_minus_minimum_similarity = 1.0 - minimum_cosine_similarity; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 2, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, Cosine>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); // Then assert_eq!( @@ -1433,21 +1310,7 @@ mod tests { let expected_one_minus_zero_similarity = 1.0 - zero_norm_similarity; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 2, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, Cosine>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); // Then assert_eq!( @@ -1467,21 +1330,8 @@ mod tests { gram[3] = -f32::MAX; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - points, - expected_leaf_k, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, InnerProduct>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = + rank_distance_fixture::(&gram, points, expected_leaf_k).unwrap(); // Then assert_eq!( @@ -1497,21 +1347,7 @@ mod tests { let expected_unassigned_neighbors = [LeafNeighbor::default(), LeafNeighbor::default()]; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - 2, - 1, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, CosineNormalized>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); // Then assert_eq!(actual_neighbors, expected_unassigned_neighbors); @@ -1528,21 +1364,8 @@ mod tests { let expected_unassigned_neighbor = LeafNeighbor::default(); // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - point_count, - requested_k, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, L2>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); + let actual_neighbors = + rank_distance_fixture::(&gram, point_count, requested_k).unwrap(); // Then assert_eq!( @@ -1554,7 +1377,7 @@ mod tests { #[rstest] #[case::l2(Metric::L2, 2.0)] #[case::cosine(Metric::Cosine, 1.0)] - #[case::normalized_cosine(Metric::CosineNormalized, 1.0)] + #[case::normalized_cosine(Metric::CosineNormalized, -0.0)] #[case::inner_product(Metric::InnerProduct, -0.0)] fn simd_nan_distance_cannot_replace_a_finite_neighbor( #[case] metric: Metric, @@ -1568,7 +1391,7 @@ mod tests { let expected_finite_neighbor = LeafNeighbor::new(1, expected_distance); // When - let actual_neighbors = run_rank_leaf_dots(metric, &gram, points, 1).1; + let actual_neighbors = run_rank_leaf_distances(metric, &gram, points, 1).1; // Then assert_eq!(actual_neighbors[16], expected_finite_neighbor); @@ -1576,13 +1399,13 @@ mod tests { #[rstest] #[case::l2(Metric::L2, 2.0)] - #[case::normalized_cosine(Metric::CosineNormalized, 1.0)] + #[case::normalized_cosine(Metric::CosineNormalized, -0.0)] #[case::inner_product(Metric::InnerProduct, -0.0)] fn simd_positive_infinity_cannot_fill_a_neighbor_slot( #[case] metric: Metric, #[case] expected_distance: f32, ) { - // Given: negative-infinite dot products produce positive-infinite scores here. + // Given: negative-infinite dot products produce positive-infinite distances here. let points = 17; let mut gram = gram_with_uniform_self_dots(points, 1.0); gram[16 * points] = f32::NEG_INFINITY; @@ -1590,41 +1413,12 @@ mod tests { let expected_finite_neighbor = LeafNeighbor::new(1, expected_distance); // When - let actual_neighbors = run_rank_leaf_dots(metric, &gram, points, 1).1; + let actual_neighbors = run_rank_leaf_distances(metric, &gram, points, 1).1; // Then assert_eq!(actual_neighbors[16], expected_finite_neighbor); } - #[test] - fn empty_leaf_has_no_neighbors() { - // Given - let empty_point_count = 0; - let empty_gram = []; - let expected_zero_neighbor_width = 0; - let expected_no_neighbors: [LeafNeighbor; 0] = []; - - // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &empty_gram, - empty_point_count, - expected_zero_neighbor_width, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, L2>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); - - // Then - assert_eq!(actual_neighbors, expected_no_neighbors); - } - #[test] fn singleton_leaf_has_no_neighbors() { // Given @@ -1634,48 +1428,10 @@ mod tests { let expected_no_neighbors: [LeafNeighbor; 0] = []; // When - let actual_neighbors = with_rank_leaf_dots_fixture::( + let actual_neighbors = rank_distance_fixture::( &singleton_gram, singleton_point_count, expected_zero_neighbor_width, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, Cosine>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap(); - - // Then - assert_eq!(actual_neighbors, expected_no_neighbors); - } - - #[test] - fn zero_requested_k_has_no_neighbors() { - // Given - let point_count = 2; - let zero_requested_k = 0; - let gram = [1.0, 0.0, 0.0, 1.0]; - let expected_no_neighbors: [LeafNeighbor; 0] = []; - - // When - let actual_neighbors = with_rank_leaf_dots_fixture::( - &gram, - point_count, - zero_requested_k, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, InnerProduct>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, ) .unwrap(); @@ -1697,21 +1453,9 @@ mod tests { }; // When - let actual_error = with_rank_leaf_dots_fixture::( - &gram, - point_count, - invalid_neighbor_width, - |input, norms, output, farthest_distances| { - rank_leaf_dots::<_, L2>( - diskann_wide::ARCH, - input, - norms, - output, - farthest_distances, - ) - }, - ) - .unwrap_err(); + let actual_error = + rank_distance_fixture::(&gram, point_count, invalid_neighbor_width) + .unwrap_err(); // Then assert_eq!(actual_error, expected_error); diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs new file mode 100644 index 0000000000..50223e0997 --- /dev/null +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -0,0 +1,300 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Build ranking distances for one PiPNN leaf. +//! +//! A ranking distance preserves nearest-first order. It does not need to equal +//! the mathematical metric distance. The leaf kernel reads only the lower triangle. + +use crate::{ANNError, ANNResult}; +use diskann_utils::views::MatrixView; +use diskann_wide::SIMDVector; + +use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, simd::PiPNNSIMDSchema}; + +/// Fill one flattened lower-triangular ranking buffer. +/// +/// An implementation initializes the diagonal and lower triangle. The upper +/// triangle stays unspecified. The input matrix has one point in each row. +pub(super) trait LeafMetric: Send + Sync + 'static { + /// Compute ranking distances for all unordered point pairs. + /// + /// `storage` has `points.nrows() * points.nrows()` elements. + fn compute_distances( + arch: A, + points: MatrixView<'_, f32>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema; +} + +/// Fill each lower-triangle entry with the two endpoint squared norms. +/// +/// Each SIMD operation copies one full group of norm sums. The scalar loop +/// copies the remaining entries. The GEMM operation adds the dot-product term. +fn fill_lower_norm_sums(arch: A, storage: &mut [f32], points: usize, norms: &[f32]) +where + A: PiPNNSIMDSchema, +{ + for source in 0..points { + let row = &mut storage[source * points..source * points + source + 1]; + let simd_end = row.len() - row.len() % A::Vector::LANES; + let source_norm = A::Vector::splat(arch, norms[source]); + for target in (0..simd_end).step_by(A::Vector::LANES) { + // SAFETY: both offsets start complete SIMD groups in live slices. + unsafe { + let target_norms = A::Vector::load_simd(arch, norms.as_ptr().add(target)); + (source_norm + target_norms).store_simd(row.as_mut_ptr().add(target)); + } + } + for target in simd_end..row.len() { + row[target] = norms[source] + norms[target]; + } + } +} + +impl LeafMetric for L2 { + fn compute_distances( + arch: A, + points: MatrixView<'_, f32>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema, + { + let point_count = points.nrows(); + // The expanded L2 formula is `||x||² + ||y||² - 2(x·y)`. + let squared_norms: Vec = points + .row_iter() + .map(|point| point.iter().map(|value| value * value).sum()) + .collect(); + arch.run(|| fill_lower_norm_sums(arch, storage, point_count, &squared_norms)); + diskann_linalg::sgemm_aat_lower_add( + point_count, + points.ncols(), + -2.0, + points.as_slice(), + storage, + ) + .map_err(ANNError::new)?; + Ok(()) + } +} + +impl LeafMetric for Cosine { + fn compute_distances( + _arch: A, + points: MatrixView<'_, f32>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema, + { + let point_count = points.nrows(); + // The diagonal supplies each point norm after GEMM computes all dots. + diskann_linalg::sgemm_aat_lower( + point_count, + points.ncols(), + 1.0, + points.as_slice(), + storage, + ) + .map_err(ANNError::new)?; + let norms: Vec = (0..point_count) + .map(|point| storage[point * point_count + point].sqrt()) + .collect(); + // Convert each lower-triangle dot to the bounded cosine distance. + for source in 0..point_count { + for target in 0..=source { + let index = source * point_count + target; + storage[index] = cosine_distance(storage[index], norms[source], norms[target]); + } + } + Ok(()) + } +} + +impl LeafMetric for CosineNormalized { + fn compute_distances( + _arch: A, + points: MatrixView<'_, f32>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema, + { + let point_count = points.nrows(); + // The constant in `1 - dot` does not change nearest-first order. + diskann_linalg::sgemm_aat_lower( + point_count, + points.ncols(), + -1.0, + points.as_slice(), + storage, + ) + .map_err(ANNError::new)?; + Ok(()) + } +} + +impl LeafMetric for InnerProduct { + fn compute_distances( + _arch: A, + points: MatrixView<'_, f32>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema, + { + let point_count = points.nrows(); + // DiskANN ranks inner products in descending order through `-dot`. + diskann_linalg::sgemm_aat_lower( + point_count, + points.ncols(), + -1.0, + points.as_slice(), + storage, + ) + .map_err(ANNError::new)?; + Ok(()) + } +} + +#[cfg(test)] +#[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] +mod tests { + use super::*; + use diskann_wide::ARCH; + use rstest::rstest; + + const POINT_COUNT: usize = 2; + const DIMENSION_COUNT: usize = 2; + const FIRST_POINT: usize = 0; + const SECOND_POINT: usize = 1; + const STALE_DISTANCE: f32 = 99.0; + const FLOAT_TOLERANCE: f32 = 1.0e-6; + + fn compute_pair_ranking( + first_point: [f32; DIMENSION_COUNT], + second_point: [f32; DIMENSION_COUNT], + ) -> f32 { + let point_values = [ + first_point[0], + first_point[1], + second_point[0], + second_point[1], + ]; + let points = MatrixView::try_from(&point_values[..], POINT_COUNT, DIMENSION_COUNT).unwrap(); + let mut storage = [STALE_DISTANCE; POINT_COUNT * POINT_COUNT]; + + M::compute_distances(ARCH, points, &mut storage).unwrap(); + + storage[SECOND_POINT * POINT_COUNT + FIRST_POINT] + } + + mod compute_distances_tests { + use super::*; + + #[test] + fn squared_l2_ranking_equals_the_sum_of_squared_coordinate_differences() { + // Given + let first_point = [3.0_f32, 4.0]; + let second_point = [0.0_f32, 4.0]; + let x_difference = first_point[0] - second_point[0]; + let y_difference = first_point[1] - second_point[1]; + let expected = x_difference.mul_add(x_difference, y_difference * y_difference); + + // When + let actual = compute_pair_ranking::(first_point, second_point); + + // Then + assert_eq!(actual, expected); + } + + #[test] + fn cosine_ranking_equals_one_minus_normalized_similarity() { + // Given + let first_point = [2.0_f32, 0.0]; + let second_point = [1.0_f32, 1.0]; + let dot = first_point[0].mul_add(second_point[0], first_point[1] * second_point[1]); + let first_norm = first_point[0].hypot(first_point[1]); + let second_norm = second_point[0].hypot(second_point[1]); + let expected = 1.0 - dot / (first_norm * second_norm); + + // When + let actual = compute_pair_ranking::(first_point, second_point); + + // Then + assert!( + (actual - expected).abs() <= FLOAT_TOLERANCE, + "actual {actual} differs from expected {expected}" + ); + } + + #[test] + fn cosine_ranking_is_one_when_one_point_has_zero_norm() { + // Given + let zero_point = [0.0_f32, 0.0]; + let unit_point = [1.0_f32, 0.0]; + let zero_similarity = 0.0_f32; + let expected = 1.0 - zero_similarity; + + // When + let actual = compute_pair_ranking::(zero_point, unit_point); + + // Then + assert_eq!(actual, expected); + } + + #[test] + fn normalized_cosine_ranking_equals_the_negative_dot_product() { + // Given + let first_point = [1.0_f32, 0.0]; + let second_point = [0.6_f32, 0.8]; + let dot = first_point[0].mul_add(second_point[0], first_point[1] * second_point[1]); + let expected = -dot; + + // When + let actual = compute_pair_ranking::(first_point, second_point); + + // Then + assert_eq!(actual, expected); + } + + #[test] + fn inner_product_ranking_equals_the_negative_dot_product() { + // Given + let first_point = [2.0_f32, -1.0]; + let second_point = [3.0_f32, 4.0]; + let dot = first_point[0].mul_add(second_point[0], first_point[1] * second_point[1]); + let expected = -dot; + + // When + let actual = compute_pair_ranking::(first_point, second_point); + + // Then + assert_eq!(actual, expected); + } + + #[rstest] + #[case::positive_zero(0.0)] + #[case::negative_zero(-0.0)] + #[trace] + fn inner_product_ranking_matches_negated_scalar_dot_bits(#[case] zero_coordinate: f32) { + // Given + let first_point = [zero_coordinate, 0.0]; + let second_point = [1.0_f32, 0.0]; + let dot = first_point[0].mul_add(second_point[0], first_point[1] * second_point[1]); + let expected = -dot; + + // When + let actual = compute_pair_ranking::(first_point, second_point); + + // Then + assert_eq!(actual.to_bits(), expected.to_bits()); + } + } +} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 32588f73c4..874fddaf5e 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,29 +5,139 @@ //! Numerical kernels for PiPNN graph construction. //! -//! [`partition_kernel`] converts point-to-leader dot products into sorted leader -//! 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 one time. It updates both points. Each point keeps the requested number -//! of neighbors. This number cannot exceed the number of other points in the leaf. -//! -//! `kernel_metric` defines metric markers and shared math. Separate leaf and -//! 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. -//! -//! 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)] -mod kernel_metric; -#[allow(dead_code)] -mod simd; +//! Metric modules fill flattened distance buffers. Kernel modules traverse those +//! buffers and retain nearest neighbors. Callers select architecture `A` and +//! metric `M` once for each graph build. #[allow(dead_code)] mod leaf_kernel; #[allow(dead_code)] +mod leaf_metric; +#[allow(dead_code)] mod partition_kernel; +#[allow(dead_code)] +mod partition_metric; +#[allow(dead_code)] +mod simd; + +pub(super) struct L2; +pub(super) struct Cosine; +pub(super) struct CosineNormalized; +pub(super) struct InnerProduct; + +/// Convert one dot product and two norms to cosine distance. +/// +/// Treat a zero or subnormal norm as zero similarity. This rule takes precedence +/// over the dot value. Clamp finite similarity to the cosine range. Otherwise, +/// a NaN input produces a NaN distance. +#[inline(always)] +fn cosine_distance(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 / (source_norm * target_norm)).clamp(-1.0, 1.0) + } +} + +#[cfg(test)] +mod cosine_distance_contract_tests { + use super::cosine_distance; + use rstest::rstest; + + mod cosine_distance_tests { + use super::*; + + #[test] + fn zero_norm_takes_precedence_over_a_nan_dot_product() { + // Given + let dot = f32::NAN; + let source_norm = 0.0_f32; + let target_norm = 1.0_f32; + let zero_similarity = 0.0_f32; + let expected = 1.0 - zero_similarity; + + // When + let actual = cosine_distance(dot, source_norm, target_norm); + + // Then + assert_eq!(actual, expected); + } + + #[rstest] + #[case::zero_source(0.0, 1.0)] + #[case::zero_target(1.0, 0.0)] + #[case::subnormal_source(f32::MIN_POSITIVE.sqrt() / 2.0, 1.0)] + #[case::subnormal_target(1.0, f32::MIN_POSITIVE.sqrt() / 2.0)] + #[trace] + fn zero_or_subnormal_norm_produces_unit_distance( + #[case] source_norm: f32, + #[case] target_norm: f32, + ) { + // Given + let dot = 0.0_f32; + let zero_similarity = 0.0_f32; + let expected = 1.0 - zero_similarity; + + // When + let actual = cosine_distance(dot, source_norm, target_norm); + + // Then + assert_eq!(actual, expected); + } + + #[test] + fn minimum_normal_norm_uses_normalized_similarity() { + // Given + let source_norm = f32::MIN_POSITIVE.sqrt(); + let target_norm = 1.0_f32; + let expected_similarity = 0.5_f32; + let dot = expected_similarity * source_norm * target_norm; + let expected = 1.0 - expected_similarity; + + // When + let actual = cosine_distance(dot, source_norm, target_norm); + + // Then + assert_eq!(actual, expected); + } + + #[rstest] + #[case::above_one(1.0)] + #[case::below_negative_one(-1.0)] + #[trace] + fn finite_similarity_outside_the_cosine_range_is_clamped(#[case] bounded_similarity: f32) { + // Given + let source_norm = 2.0_f32; + let target_norm = 2.0_f32; + let norm_product = source_norm * target_norm; + let rounding_excess = f32::EPSILON * norm_product; + let dot = bounded_similarity * (norm_product + rounding_excess); + let expected = 1.0 - bounded_similarity; + + // When + let actual = cosine_distance(dot, source_norm, target_norm); + + // Then + assert_eq!(actual, expected); + } + + #[rstest] + #[case::nan_dot(f32::NAN, 1.0, 1.0)] + #[case::nan_source_norm(0.0, f32::NAN, 1.0)] + #[case::nan_target_norm(0.0, 1.0, f32::NAN)] + #[trace] + fn nan_without_a_zero_norm_produces_nan_distance( + #[case] dot: f32, + #[case] source_norm: f32, + #[case] target_norm: f32, + ) { + // Given: the case supplies one NaN. The other values do not select the zero-norm rule. + + // When + let actual = cosine_distance(dot, source_norm, target_norm); + + // Then + assert!(actual.is_nan()); + } + } +} diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index b0e8fc5539..0da345ce06 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -6,74 +6,33 @@ //! Select partition centers for PiPNN point assignment. //! //! A leader is a sampled dataset point that represents one child partition. -//! The kernel prepares reusable leader norms, computes point-to-leader dot -//! products, and returns nearest leader-column IDs for partition scatter. +//! The metric builds final point-to-leader distances. The kernel 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. An +//! leaders. Equal distances keep sampled-leader order. NaN is not rankable. An //! unfilled output slot contains [`UNASSIGNED_LEADER`]. -use std::marker::PhantomData; - -use crate::{ANNError, ANNResult}; -use diskann_linalg::Transpose; +use crate::ANNResult; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{SIMDMask, SIMDVector}; use super::{ - kernel_metric::{PartitionMetric, PartitionNorms}, + partition_metric::PartitionMetric, simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, }; /// 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>, - 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, + distance_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)] -struct PartitionInput<'a> { - dots: MatrixView<'a, f32>, - norms: PartitionNorms<'a>, -} - -/// Assign one packed point stripe to prepared partition leaders. +/// Assign one packed point stripe to metric-owned partition leaders. /// /// A point can have fewer assignments than the output width. Each remaining /// slot contains [`UNASSIGNED_LEADER`]. @@ -84,7 +43,7 @@ struct PartitionInput<'a> { pub(super) fn assign_leaders( arch: A, points: MatrixView<'_, f32>, - leaders: &PreparedLeaders<'_, M>, + leaders: &M::Leaders<'_>, output: MutMatrixView<'_, u32>, workspace: &mut PartitionKernelWorkspace, ) -> ANNResult<()> @@ -93,108 +52,96 @@ where M: PartitionMetric, { let point_count = points.nrows(); - let leader_count = leaders.len(); - let dot_count = point_count * leader_count; + let leader_count = M::leader_count(leaders); + let distance_count = point_count * leader_count; let PartitionKernelWorkspace { - dot_scratch, - point_norm_scratch, + distance_scratch, ranked_leader_scratch, } = workspace; - if dot_scratch.len() < dot_count { - dot_scratch.resize(dot_count, 0.0); + if distance_scratch.len() < distance_count { + distance_scratch.resize(distance_count, 0.0); } - diskann_linalg::sgemm( - Transpose::None, - Transpose::Ordinary, + M::compute_distances( + arch, + points, + leaders, + &mut distance_scratch[..distance_count], + )?; + rank_leader_distances( + arch, + &distance_scratch[..distance_count], 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( +/// Rank final point-to-leader distances. +/// +/// `distance_flatten` contains `point_count * leader_count` elements. Each row +/// stores all leader distances for one point. The caller supplies at least one +/// point and one output column. +fn rank_leader_distances( arch: A, - input: PartitionInput<'_>, + distance_flatten: &[f32], + point_count: usize, + leader_count: usize, output: MutMatrixView<'_, u32>, ranked_leaders: &mut Vec<(u32, f32)>, ) where A: PiPNNSIMDSchema, - M: PartitionMetric, { let fanout = output.ncols(); - if fanout == 0 || input.dots.nrows() == 0 { - return; - } - ranked_leaders.resize(fanout, (UNASSIGNED_LEADER, f32::INFINITY)); // Rayon outlines stripe workers. Reapply target features before ranking leaders. arch.run(move || { - select_point_leaders::(arch, input.dots, input.norms, output, ranked_leaders); + select_point_leaders( + arch, + distance_flatten, + point_count, + leader_count, + 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. #[inline(always)] -fn select_point_leaders( +fn select_point_leaders( arch: A, - dots: MatrixView<'_, f32>, - norms: PartitionNorms<'_>, + distance_flatten: &[f32], + point_count: usize, + leader_count: usize, mut output: MutMatrixView<'_, u32>, ranked_leaders: &mut [(u32, f32)], ) where A: PiPNNSIMDSchema, - M: PartitionMetric, { - let leader_count = dots.ncols(); let fanout = output.ncols(); - for (point, (point_dots, point_output)) in dots - .row_iter() - .zip(output.as_mut_slice().chunks_exact_mut(fanout)) + for (point, point_output) in output + .as_mut_slice() + .chunks_exact_mut(fanout) + .take(point_count) .enumerate() { + let row_start = point * leader_count; + let point_distances = &distance_flatten[row_start..row_start + leader_count]; 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 % A::Vector::LANES; for first_leader in (0..simd_prefix).step_by(A::Vector::LANES) { // SAFETY: This group is inside the point's leader row. - let dot_products = - unsafe { A::Vector::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); + let distance_group = + unsafe { A::Vector::load_simd(arch, point_distances.as_ptr().add(first_leader)) }; + insert_leader_lanes(distance_group, first_leader, ranked_leaders); } - 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); + for (leader, &distance) in point_distances.iter().enumerate().skip(simd_prefix) { + insert_leader(ranked_leaders, leader as u32, distance); } for (destination, &(leader, _)) in point_output.iter_mut().zip(ranked_leaders.iter()) { *destination = leader; @@ -206,39 +153,43 @@ fn select_point_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, ranked_leaders: &mut [(u32, f32)]) +fn insert_leader_lanes(distances: F, first_leader: usize, ranked_leaders: &mut [(u32, f32)]) where F: PiPNNSIMDVector, { - let threshold = F::splat(scores.arch(), ranked_leaders[ranked_leaders.len() - 1].1); - let eligible = scores.lt_simd(threshold); + let threshold = F::splat(distances.arch(), ranked_leaders[ranked_leaders.len() - 1].1); + let eligible = distances.lt_simd(threshold); if eligible.none() { return; } - let values = scores.to_lane_array(); - let values = values.as_ref(); + let distance_lanes = distances.to_lane_array(); + let distance_lanes = distance_lanes.as_ref(); let mut lanes = F::active_lanes(eligible); while lanes != 0 { let lane = lanes.trailing_zeros() as usize; lanes &= lanes - 1; - insert_leader(ranked_leaders, (first_leader + lane) as u32, values[lane]); + insert_leader( + ranked_leaders, + (first_leader + lane) as u32, + distance_lanes[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. `ranked_leaders` -/// stores retained centers in nearest-first order. Equal scores and NaN do not +/// stores retained centers in nearest-first order. Equal distances and NaN do not /// enter, so sampled-leader order resolves ties. #[inline(always)] -fn insert_leader(ranked_leaders: &mut [(u32, f32)], leader: u32, score: f32) { +fn insert_leader(ranked_leaders: &mut [(u32, f32)], leader: u32, distance: f32) { let threshold = ranked_leaders.len() - 1; - if score.partial_cmp(&ranked_leaders[threshold].1) != Some(std::cmp::Ordering::Less) { + if distance.partial_cmp(&ranked_leaders[threshold].1) != Some(std::cmp::Ordering::Less) { return; } - ranked_leaders[threshold] = (leader, score); + ranked_leaders[threshold] = (leader, distance); let mut slot = threshold; while slot > 0 && ranked_leaders[slot].1 < ranked_leaders[slot - 1].1 { ranked_leaders.swap(slot, slot - 1); @@ -249,7 +200,7 @@ fn insert_leader(ranked_leaders: &mut [(u32, f32)], leader: u32, score: f32) { #[cfg(test)] mod tests { use super::*; - use crate::graph::pipnn::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use crate::graph::pipnn::Cosine; use diskann_utils::views::{Matrix, MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; @@ -257,45 +208,36 @@ mod tests { use super::*; use diskann_wide::arch::{self, Target1}; + #[derive(Clone, Copy)] + pub(super) struct DotFixture<'a> { + dots: MatrixView<'a, f32>, + point_norms: &'a [f32], + leader_norms: &'a [f32], + } + struct KernelCall<'a> { - input: PartitionInput<'a>, + distance_flatten: &'a [f32], + point_count: usize, + leader_count: usize, output: MutMatrixView<'a, u32>, ranked_leaders: &'a mut Vec<(u32, f32)>, } - struct DispatchMetric(Metric); + struct RankDistances; - impl Target1> for DispatchMetric + impl Target1> for RankDistances 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, - ), - } + rank_leader_distances( + arch, + call.distance_flatten, + call.point_count, + call.leader_count, + call.output, + call.ranked_leaders, + ); } } @@ -305,51 +247,18 @@ mod tests { leader_count: usize, point_norms: &'a [f32], leader_norms: &'a [f32], - ) -> PartitionInput<'a> { - PartitionInput { + ) -> DotFixture<'a> { + DotFixture { dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), - norms: PartitionNorms { - point_norms, - leader_norms, - }, + point_norms, + leader_norms, } } - pub(super) fn with_rank_leader_dots_fixture( - input: PartitionInput<'_>, - nearest_leader_count: usize, - run: F, - ) -> Vec - where - F: FnOnce(PartitionInput<'_>, MutMatrixView<'_, u32>, &mut Vec<(u32, f32)>), - { - let mut output = Matrix::new(u32::MAX, input.dots.nrows(), nearest_leader_count); - run(input, output.as_mut_view(), &mut Vec::new()); - output.into_inner().into_vec() - } - - /// Run the production ranker through runtime architecture and metric dispatch. - pub(super) fn run_rank_leader_dots( - metric: Metric, - input: PartitionInput<'_>, - nearest_leader_count: usize, - ) -> Vec { - let mut output = Matrix::new(u32::MAX, input.dots.nrows(), nearest_leader_count); - arch::dispatch1_no_features( - DispatchMetric(metric), - KernelCall { - input, - output: output.as_mut_view(), - ranked_leaders: &mut Vec::new(), - }, - ); - output.into_inner().into_vec() - } - - fn reference_score(metric: Metric, dot: f32, point_norm: f32, leader_norm: f32) -> f32 { + fn reference_distance(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::CosineNormalized => -dot, Metric::InnerProduct => -dot, Metric::Cosine => { if point_norm < f32::MIN_POSITIVE.sqrt() @@ -363,9 +272,47 @@ mod tests { } } + fn distance_flatten(metric: Metric, input: DotFixture<'_>) -> Vec { + input + .dots + .row_iter() + .enumerate() + .flat_map(|(point, dots)| { + dots.iter().enumerate().map(move |(leader, &dot)| { + reference_distance( + metric, + dot, + input.point_norms.get(point).copied().unwrap_or(0.0), + input.leader_norms.get(leader).copied().unwrap_or(0.0), + ) + }) + }) + .collect() + } + + pub(super) fn rank_distance_fixture( + metric: Metric, + input: DotFixture<'_>, + nearest_leader_count: usize, + ) -> Vec { + let distance_flatten = distance_flatten(metric, input); + let mut output = Matrix::new(u32::MAX, input.dots.nrows(), nearest_leader_count); + arch::dispatch1_no_features( + RankDistances, + KernelCall { + distance_flatten: &distance_flatten, + point_count: input.dots.nrows(), + leader_count: input.dots.ncols(), + output: output.as_mut_view(), + ranked_leaders: &mut Vec::new(), + }, + ); + output.into_inner().into_vec() + } + pub(super) fn reference_assignments( metric: Metric, - input: PartitionInput<'_>, + input: DotFixture<'_>, nearest_leader_count: usize, ) -> Vec { let mut output = vec![UNASSIGNED_LEADER; input.dots.nrows() * nearest_leader_count]; @@ -375,16 +322,15 @@ mod tests { .zip(output.chunks_exact_mut(nearest_leader_count)) .enumerate() { - let point_norm = input.norms.point_norms.get(point).copied().unwrap_or(0.0); + let point_norm = input.point_norms.get(point).copied().unwrap_or(0.0); let mut candidates: Vec<_> = dots .iter() .enumerate() .filter_map(|(leader, &dot)| { - 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)) + let leader_norm = input.leader_norms.get(leader).copied().unwrap_or(0.0); + let distance = reference_distance(metric, dot, point_norm, leader_norm); + (distance.partial_cmp(&f32::INFINITY) == Some(std::cmp::Ordering::Less)) + .then_some((leader as u32, distance)) }) .collect(); candidates.sort_by(|left, right| left.1.total_cmp(&right.1)); @@ -396,8 +342,6 @@ mod tests { } /// Build ranking input from two axis points and leaders between those axes. - /// - /// The first point prefers early leaders. The second point prefers late leaders. pub(super) fn lane_boundary_input_from_point_and_leader_vectors( metric: Metric, leader_count: usize, @@ -468,7 +412,7 @@ mod tests { } #[test] - fn nan_score_does_not_enter_the_topk() { + fn nan_distance_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]; @@ -489,9 +433,8 @@ mod tests { fn assigns_each_point_to_highest_similarity_leaders_with_cosine() { // 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 leaders = + Cosine::create_leaders(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]; @@ -518,9 +461,8 @@ mod tests { 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 leaders = + Cosine::create_leaders(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(); @@ -564,7 +506,7 @@ mod tests { } } - mod rank_leader_dots_tests { + mod rank_leader_distances_tests { use super::test_support::*; use super::*; use rstest::rstest; @@ -597,14 +539,14 @@ mod tests { let expected_assignments = reference_assignments(metric, input, nearest_leader_count); // When - let actual_assignments = run_rank_leader_dots(metric, input, nearest_leader_count); + let actual_assignments = rank_distance_fixture(metric, input, nearest_leader_count); // Then assert_eq!(actual_assignments, expected_assignments); } #[test] - fn equal_scores_keep_sampled_leader_order_with_l2() { + fn equal_distances_keep_sampled_leader_order_with_l2() { // Given let point_count = 1; let leader_count = 4; @@ -617,13 +559,7 @@ mod tests { partition_input(&dots, point_count, leader_count, &[], &leader_squared_norms); // When - let actual_assignments = with_rank_leader_dots_fixture( - input, - nearest_leader_count, - |input, output, ranked_leaders| { - rank_leader_dots::<_, L2>(diskann_wide::ARCH, input, output, ranked_leaders); - }, - ); + let actual_assignments = rank_distance_fixture(Metric::L2, input, nearest_leader_count); // Then assert_eq!(actual_assignments, expected_sampled_leader_order); @@ -649,18 +585,8 @@ mod tests { ); // When - let actual_assignments = with_rank_leader_dots_fixture( - input, - nearest_leader_count, - |input, output, ranked_leaders| { - rank_leader_dots::<_, Cosine>( - diskann_wide::ARCH, - input, - output, - ranked_leaders, - ); - }, - ); + let actual_assignments = + rank_distance_fixture(Metric::Cosine, input, nearest_leader_count); // Then assert_eq!(actual_assignments, expected_sampled_leader_order); @@ -678,7 +604,7 @@ mod tests { let expected_first_leader = [0]; // When - let actual_assignment = run_rank_leader_dots( + let actual_assignment = rank_distance_fixture( Metric::Cosine, partition_input( &dots, @@ -695,13 +621,13 @@ mod tests { } #[test] - fn f32_max_score_is_still_a_rankable_leader() { + fn f32_max_distance_is_still_rankable() { // Given let point_count = 1; let leader_count = 8; let nearest_leader_count = leader_count; - let maximum_rankable_score = f32::MAX; - let dot_product_that_produces_it = -maximum_rankable_score; + let maximum_rankable_distance = f32::MAX; + let dot_product_that_produces_it = -maximum_rankable_distance; let mut dots = [0.0; 8]; dots[7] = dot_product_that_produces_it; let expected_all_leaders_in_scan_order = [0, 1, 2, 3, 4, 5, 6, 7]; @@ -709,18 +635,8 @@ mod tests { let input = partition_input(&dots, point_count, leader_count, &[], &[]); // When - let actual_assignments = with_rank_leader_dots_fixture( - input, - nearest_leader_count, - |input, output, ranked_leaders| { - rank_leader_dots::<_, InnerProduct>( - diskann_wide::ARCH, - input, - output, - ranked_leaders, - ); - }, - ); + let actual_assignments = + rank_distance_fixture(Metric::InnerProduct, input, nearest_leader_count); // Then assert_eq!(actual_assignments, expected_all_leaders_in_scan_order); @@ -738,18 +654,8 @@ mod tests { let input = partition_input(&dots, point_count, leader_count, &[], &[]); // When - let actual_assignments = with_rank_leader_dots_fixture( - input, - nearest_leader_count, - |input, output, ranked_leaders| { - rank_leader_dots::<_, InnerProduct>( - diskann_wide::ARCH, - input, - output, - ranked_leaders, - ); - }, - ); + let actual_assignments = + rank_distance_fixture(Metric::InnerProduct, input, nearest_leader_count); // Then assert_eq!(actual_assignments, expected_finite_leaders); @@ -775,74 +681,11 @@ mod tests { ); // When - let actual_assignments = with_rank_leader_dots_fixture( - input, - nearest_leader_count, - |input, output, ranked_leaders| { - rank_leader_dots::<_, Cosine>( - diskann_wide::ARCH, - input, - output, - ranked_leaders, - ); - }, - ); + let actual_assignments = + rank_distance_fixture(Metric::Cosine, input, nearest_leader_count); // Then assert_eq!(actual_assignments, expected_finite_leaders); } - - #[test] - fn zero_output_width_leaves_ranked_leaders_unchanged() { - // Given - let dot_products = [0.0]; - let leader_squared_norms = [1.0]; - let input = partition_input(&dot_products, 1, 1, &[], &leader_squared_norms); - let mut no_assignments = []; - let output = MutMatrixView::try_from(&mut no_assignments[..], 1, 0).unwrap(); - let expected_ranked_leaders = vec![(7, 0.25)]; - let mut actual_ranked_leaders = expected_ranked_leaders.clone(); - - // When - rank_leader_dots::<_, L2>( - diskann_wide::ARCH, - input, - output, - &mut actual_ranked_leaders, - ); - - // Then - assert_eq!(actual_ranked_leaders, expected_ranked_leaders); - } - - #[test] - fn empty_point_matrix_produces_no_assignments() { - // Given - let empty_point_count = 0; - let leader_count = 3; - let nearest_leader_count = 2; - let no_dot_products = []; - let expected_no_assignments: [u32; 0] = []; - - let input = - partition_input(&no_dot_products, empty_point_count, leader_count, &[], &[]); - - // When - let actual_assignments = with_rank_leader_dots_fixture( - input, - nearest_leader_count, - |input, output, ranked_leaders| { - rank_leader_dots::<_, InnerProduct>( - diskann_wide::ARCH, - input, - output, - ranked_leaders, - ); - }, - ); - - // Then - assert_eq!(actual_assignments, expected_no_assignments); - } } } diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs new file mode 100644 index 0000000000..77a313ac2a --- /dev/null +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -0,0 +1,505 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Build ranking distances from point stripes to partition leaders. +//! +//! A ranking distance preserves nearest-first order. One leader set serves all +//! point stripes in a partition split. + +use std::sync::OnceLock; + +use crate::{ANNError, ANNResult}; +use diskann_linalg::Transpose; +use diskann_utils::views::MatrixView; +use diskann_vector::{Norm, norm::FastL2NormSquared}; +use diskann_wide::SIMDVector; + +use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, simd::PiPNNSIMDSchema}; + +/// Store leader values with immutable metric data. +/// +/// The metric selects `Cache`. An L2 cache stores squared norms. A cosine cache +/// stores norms. `OnceLock` lets concurrent point stripes initialize data once. +pub(super) struct PartitionLeaders<'a, Cache> { + values: MatrixView<'a, f32>, + cache: Cache, +} + +/// Fill one flattened point-to-leader ranking buffer. +/// +/// The associated leader type hides metric data from the caller. The caller +/// creates one value and shares it across all point stripes. +pub(super) trait PartitionMetric: Send + Sync + 'static { + /// Leader values and immutable metric data for one partition split. + type Leaders<'a>: Sync; + + /// Bind one non-empty leader matrix to this metric. + /// + /// Partitioning creates at least one leader before it calls this function. + fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a>; + + /// Return the number of leaders in a metric-owned leader set. + fn leader_count(leaders: &Self::Leaders<'_>) -> usize; + + /// Compute one row-major point-to-leader ranking buffer. + /// + /// `storage` has `points.nrows() * leader_count` elements. + fn compute_distances( + arch: A, + points: MatrixView<'_, f32>, + leaders: &Self::Leaders<'_>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema; +} + +/// Compute L2 squared norms with the established sequential reduction order. +/// +/// Small rounding differences can change a leader tie. Keep this order equal to +/// the pre-GEMM implementation. +fn l2_squared_norms(vectors: MatrixView<'_, f32>) -> Vec { + vectors + .row_iter() + .map(|vector| vector.iter().map(|value| value * value).sum()) + .collect() +} + +/// Compute cosine norms with the same reduction for points and leaders. +fn cosine_norms(vectors: MatrixView<'_, f32>) -> Vec { + vectors + .row_iter() + .map(|vector| FastL2NormSquared.evaluate(vector).sqrt()) + .collect() +} + +/// Copy one leader value into the same column of each point row. +/// +/// L2 uses this operation to initialize every row with leader squared norms. +/// Full SIMD groups use dispatched vector stores. The scalar tail copies the rest. +fn fill_rows(arch: A, rows: &mut [f32], row_values: &[f32]) +where + A: PiPNNSIMDSchema, +{ + let row_length = row_values.len(); + let simd_end = row_length - row_length % A::Vector::LANES; + for row in rows.chunks_exact_mut(row_length) { + for start in (0..simd_end).step_by(A::Vector::LANES) { + // SAFETY: both offsets start complete SIMD groups in live slices. + unsafe { + A::Vector::load_simd(arch, row_values.as_ptr().add(start)) + .store_simd(row.as_mut_ptr().add(start)); + } + } + row[simd_end..].copy_from_slice(&row_values[simd_end..]); + } +} + +impl PartitionMetric for L2 { + type Leaders<'a> = PartitionLeaders<'a, OnceLock>>; + + fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { + PartitionLeaders { + values, + cache: OnceLock::new(), + } + } + + fn leader_count(leaders: &Self::Leaders<'_>) -> usize { + leaders.values.nrows() + } + + fn compute_distances( + arch: A, + points: MatrixView<'_, f32>, + leaders: &Self::Leaders<'_>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema, + { + // The point norm is constant across a point row. It cannot change ranking. + let leader_norms = leaders + .cache + .get_or_init(|| l2_squared_norms(leaders.values)); + arch.run(|| fill_rows(arch, storage, leader_norms)); + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + points.nrows(), + leaders.values.nrows(), + points.ncols(), + -2.0, + points.as_slice(), + leaders.values.as_slice(), + Some(1.0), + storage, + ) + .map_err(ANNError::new)?; + Ok(()) + } +} + +impl PartitionMetric for Cosine { + type Leaders<'a> = PartitionLeaders<'a, OnceLock>>; + + fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { + PartitionLeaders { + values, + cache: OnceLock::new(), + } + } + + fn leader_count(leaders: &Self::Leaders<'_>) -> usize { + leaders.values.nrows() + } + + fn compute_distances( + _arch: A, + points: MatrixView<'_, f32>, + leaders: &Self::Leaders<'_>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema, + { + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + points.nrows(), + leaders.values.nrows(), + points.ncols(), + 1.0, + points.as_slice(), + leaders.values.as_slice(), + None, + storage, + ) + .map_err(ANNError::new)?; + let point_norms = cosine_norms(points); + let leader_norms = leaders.cache.get_or_init(|| cosine_norms(leaders.values)); + let leader_count = leaders.values.nrows(); + // Convert each dot to cosine distance. Reuse leader norms across stripes. + for (row, &point_norm) in storage + .chunks_exact_mut(leader_count) + .zip(point_norms.iter()) + { + for (distance, &leader_norm) in row.iter_mut().zip(leader_norms.iter()) { + *distance = cosine_distance(*distance, point_norm, leader_norm); + } + } + Ok(()) + } +} + +// Normalized cosine and inner product have the same ranking expression. +// Both metrics rank candidates with `-dot`. +impl PartitionMetric for CosineNormalized { + type Leaders<'a> = PartitionLeaders<'a, ()>; + + fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { + PartitionLeaders { values, cache: () } + } + + fn leader_count(leaders: &Self::Leaders<'_>) -> usize { + leaders.values.nrows() + } + + fn compute_distances( + _arch: A, + points: MatrixView<'_, f32>, + leaders: &Self::Leaders<'_>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema, + { + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + points.nrows(), + leaders.values.nrows(), + points.ncols(), + -1.0, + points.as_slice(), + leaders.values.as_slice(), + None, + storage, + ) + .map_err(ANNError::new)?; + Ok(()) + } +} + +impl PartitionMetric for InnerProduct { + type Leaders<'a> = PartitionLeaders<'a, ()>; + + fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { + PartitionLeaders { values, cache: () } + } + + fn leader_count(leaders: &Self::Leaders<'_>) -> usize { + leaders.values.nrows() + } + + fn compute_distances( + _arch: A, + points: MatrixView<'_, f32>, + leaders: &Self::Leaders<'_>, + storage: &mut [f32], + ) -> ANNResult<()> + where + A: PiPNNSIMDSchema, + { + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + points.nrows(), + leaders.values.nrows(), + points.ncols(), + -1.0, + points.as_slice(), + leaders.values.as_slice(), + None, + storage, + ) + .map_err(ANNError::new)?; + Ok(()) + } +} + +#[cfg(test)] +#[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] +mod tests { + use super::*; + use diskann_wide::ARCH; + + const DIMENSION_COUNT: usize = 2; + const STALE_DISTANCE: f32 = 99.0; + const FLOAT_TOLERANCE: f32 = 1.0e-6; + + fn matrix(values: &[f32], rows: usize) -> MatrixView<'_, f32> { + MatrixView::try_from(values, rows, DIMENSION_COUNT).unwrap() + } + + fn compute_one_ranking( + point: [f32; DIMENSION_COUNT], + leader: [f32; DIMENSION_COUNT], + ) -> f32 { + let leaders = M::create_leaders(matrix(&leader, 1)); + let mut storage = [STALE_DISTANCE]; + + M::compute_distances(ARCH, matrix(&point, 1), &leaders, &mut storage).unwrap(); + + storage[0] + } + + mod compute_distances_tests { + use super::*; + + #[test] + fn squared_l2_ranking_omits_the_point_norm() { + // Given + let point = [0.0_f32, 4.0]; + let leader = [3.0_f32, 4.0]; + let leader_squared_norm = leader[0].mul_add(leader[0], leader[1] * leader[1]); + let dot = point[0].mul_add(leader[0], point[1] * leader[1]); + let expected = (-2.0_f32).mul_add(dot, leader_squared_norm); + + // When + let actual = compute_one_ranking::(point, leader); + + // Then + assert_eq!(actual, expected); + } + + #[test] + fn l2_leader_cache_preserves_the_sequential_reduction_order() { + fn next_regression_value(state: &mut u64) -> f32 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 + } + + // Given + const DIMENSIONS: usize = 129; + const REGRESSION_SEED: u64 = 0x3a85_f952_c718_6e49; + let mut state = REGRESSION_SEED; + let _point_values: Vec = (0..DIMENSIONS) + .map(|_| next_regression_value(&mut state)) + .collect(); + let leader_values: Vec = (0..DIMENSIONS) + .map(|_| next_regression_value(&mut state)) + .collect(); + let expected: f32 = leader_values.iter().map(|value| value * value).sum(); + let reassociated = FastL2NormSquared.evaluate(leader_values.as_slice()); + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), 1, DIMENSIONS).unwrap(); + let leaders = L2::create_leaders(leader_matrix); + let zero_point = vec![0.0_f32; DIMENSIONS]; + let point_matrix = MatrixView::try_from(zero_point.as_slice(), 1, DIMENSIONS).unwrap(); + let mut distance = [STALE_DISTANCE]; + + // When + L2::compute_distances(ARCH, point_matrix, &leaders, &mut distance).unwrap(); + let actual = leaders.cache.get().unwrap()[0]; + + // Then + assert_ne!(expected.to_bits(), reassociated.to_bits()); + assert_eq!(actual.to_bits(), expected.to_bits()); + } + + #[test] + fn cosine_ranking_equals_one_minus_normalized_similarity() { + // Given + let point = [2.0_f32, 0.0]; + let leader = [1.0_f32, 1.0]; + let dot = point[0].mul_add(leader[0], point[1] * leader[1]); + let point_norm = point[0].hypot(point[1]); + let leader_norm = leader[0].hypot(leader[1]); + let expected = 1.0 - dot / (point_norm * leader_norm); + + // When + let actual = compute_one_ranking::(point, leader); + + // Then + assert!( + (actual - expected).abs() <= FLOAT_TOLERANCE, + "actual {actual} differs from expected {expected}" + ); + } + + #[test] + fn normalized_cosine_ranking_equals_the_negative_dot_product() { + // Given + let point = [1.0_f32, 0.0]; + let leader = [0.6_f32, 0.8]; + let dot = point[0].mul_add(leader[0], point[1] * leader[1]); + let expected = -dot; + + // When + let actual = compute_one_ranking::(point, leader); + + // Then + assert_eq!(actual, expected); + } + + #[test] + fn inner_product_ranking_equals_the_negative_dot_product() { + // Given + let point = [2.0_f32, -1.0]; + let leader = [3.0_f32, 4.0]; + let dot = point[0].mul_add(leader[0], point[1] * leader[1]); + let expected = -dot; + + // When + let actual = compute_one_ranking::(point, leader); + + // Then + assert_eq!(actual, expected); + } + + #[test] + fn reused_l2_leaders_match_fresh_leaders_for_a_new_point_stripe() { + // Given + let leader_values = [3.0_f32, 4.0, 0.0, 2.0]; + let leader_count = leader_values.len() / DIMENSION_COUNT; + let first_point = [1.0_f32, 0.0]; + let second_point = [0.0_f32, 1.0]; + let first_leader_squared_norm = + leader_values[0].mul_add(leader_values[0], leader_values[1] * leader_values[1]); + let second_leader_squared_norm = + leader_values[2].mul_add(leader_values[2], leader_values[3] * leader_values[3]); + let first_dot = + second_point[0].mul_add(leader_values[0], second_point[1] * leader_values[1]); + let second_dot = + second_point[0].mul_add(leader_values[2], second_point[1] * leader_values[3]); + let expected = [ + (-2.0_f32).mul_add(first_dot, first_leader_squared_norm), + (-2.0_f32).mul_add(second_dot, second_leader_squared_norm), + ]; + let reused_leaders = L2::create_leaders(matrix(&leader_values, leader_count)); + let fresh_leaders = L2::create_leaders(matrix(&leader_values, leader_count)); + let mut discarded_first_output = [STALE_DISTANCE; 2]; + L2::compute_distances( + ARCH, + matrix(&first_point, 1), + &reused_leaders, + &mut discarded_first_output, + ) + .unwrap(); + let mut reused_output = [STALE_DISTANCE; 2]; + let mut fresh_output = [STALE_DISTANCE; 2]; + + // When + L2::compute_distances( + ARCH, + matrix(&second_point, 1), + &reused_leaders, + &mut reused_output, + ) + .unwrap(); + L2::compute_distances( + ARCH, + matrix(&second_point, 1), + &fresh_leaders, + &mut fresh_output, + ) + .unwrap(); + + // Then + assert_eq!(reused_output, expected); + assert_eq!(fresh_output, expected); + } + + #[test] + fn reused_cosine_leaders_match_fresh_leaders_for_a_new_point_stripe() { + // Given + let leader_values = [1.0_f32, 0.0, 0.0, 2.0]; + let leader_count = leader_values.len() / DIMENSION_COUNT; + let first_point = [1.0_f32, 0.0]; + let second_point = [0.0_f32, 1.0]; + let orthogonal_similarity = 0.0_f32; + let equal_direction_similarity = 1.0_f32; + let expected = [ + 1.0 - orthogonal_similarity, + 1.0 - equal_direction_similarity, + ]; + let reused_leaders = Cosine::create_leaders(matrix(&leader_values, leader_count)); + let fresh_leaders = Cosine::create_leaders(matrix(&leader_values, leader_count)); + let mut discarded_first_output = [STALE_DISTANCE; 2]; + Cosine::compute_distances( + ARCH, + matrix(&first_point, 1), + &reused_leaders, + &mut discarded_first_output, + ) + .unwrap(); + let mut reused_output = [STALE_DISTANCE; 2]; + let mut fresh_output = [STALE_DISTANCE; 2]; + + // When + Cosine::compute_distances( + ARCH, + matrix(&second_point, 1), + &reused_leaders, + &mut reused_output, + ) + .unwrap(); + Cosine::compute_distances( + ARCH, + matrix(&second_point, 1), + &fresh_leaders, + &mut fresh_output, + ) + .unwrap(); + + // Then + assert_eq!(reused_output, expected); + assert_eq!(fresh_output, expected); + } + } +} diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 621114c761..4742bc385b 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -5,9 +5,7 @@ //! SIMD schema for PiPNN numerical kernels. -use diskann_wide::{ - Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, SupportedLaneCount, -}; +use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDVector, SupportedLaneCount}; /// Default SIMD representation used by every PiPNN numerical stage. /// @@ -15,24 +13,18 @@ use diskann_wide::{ type DefaultVector = ::f32x16; /// Operations required by PiPNN SIMD vectors. -pub(super) trait PiPNNSIMDVector: - SIMDVector + SIMDFloat + std::ops::Div -{ +pub(super) trait PiPNNSIMDVector: SIMDVector + SIMDFloat { /// Convert the SIMD value to a readable lane array. fn to_lane_array(self) -> impl AsRef<[f32]>; /// 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: SIMDVector> + SIMDFloat, Const: SupportedLaneCount, - F::Mask: SIMDSelect, u64: From<<::BitMask as SIMDMask>::Underlying>, { #[inline(always)] @@ -45,11 +37,6 @@ where 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) - } } /// PiPNN SIMD representation for one architecture. From b46ecdd80d7783d53d9a3c6fe1f3fd840a702809 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 4 Sep 2026 08:01:27 +0000 Subject: [PATCH 097/229] refactor(pipnn): remove metric architecture input Portable scalar prefill loops auto-vectorize without exposing SIMD types through metric interfaces. --- diskann/src/graph/pipnn/leaf_kernel.rs | 2 +- diskann/src/graph/pipnn/leaf_metric.rs | 86 ++++------------- diskann/src/graph/pipnn/mod.rs | 5 +- diskann/src/graph/pipnn/partition_kernel.rs | 7 +- diskann/src/graph/pipnn/partition_metric.rs | 101 +++++--------------- diskann/src/graph/pipnn/simd.rs | 4 +- 6 files changed, 49 insertions(+), 156 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 1c92696704..3dfcdfebcc 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -110,7 +110,7 @@ where if distance_scratch.len() < distance_count { distance_scratch.resize(distance_count, 0.0); } - M::compute_distances(arch, points, &mut distance_scratch[..distance_count])?; + M::compute_distances(points, &mut distance_scratch[..distance_count])?; rank_leaf_distances( arch, &distance_scratch[..distance_count], diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index 50223e0997..30737eb89d 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -10,9 +10,8 @@ use crate::{ANNError, ANNResult}; use diskann_utils::views::MatrixView; -use diskann_wide::SIMDVector; -use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, simd::PiPNNSIMDSchema}; +use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; /// Fill one flattened lower-triangular ranking buffer. /// @@ -22,56 +21,27 @@ pub(super) trait LeafMetric: Send + Sync + 'static { /// Compute ranking distances for all unordered point pairs. /// /// `storage` has `points.nrows() * points.nrows()` elements. - fn compute_distances( - arch: A, - points: MatrixView<'_, f32>, - storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema; -} - -/// Fill each lower-triangle entry with the two endpoint squared norms. -/// -/// Each SIMD operation copies one full group of norm sums. The scalar loop -/// copies the remaining entries. The GEMM operation adds the dot-product term. -fn fill_lower_norm_sums(arch: A, storage: &mut [f32], points: usize, norms: &[f32]) -where - A: PiPNNSIMDSchema, -{ - for source in 0..points { - let row = &mut storage[source * points..source * points + source + 1]; - let simd_end = row.len() - row.len() % A::Vector::LANES; - let source_norm = A::Vector::splat(arch, norms[source]); - for target in (0..simd_end).step_by(A::Vector::LANES) { - // SAFETY: both offsets start complete SIMD groups in live slices. - unsafe { - let target_norms = A::Vector::load_simd(arch, norms.as_ptr().add(target)); - (source_norm + target_norms).store_simd(row.as_mut_ptr().add(target)); - } - } - for target in simd_end..row.len() { - row[target] = norms[source] + norms[target]; - } - } + fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()>; } impl LeafMetric for L2 { - fn compute_distances( - arch: A, - points: MatrixView<'_, f32>, - storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema, - { + fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { let point_count = points.nrows(); // The expanded L2 formula is `||x||² + ||y||² - 2(x·y)`. let squared_norms: Vec = points .row_iter() .map(|point| point.iter().map(|value| value * value).sum()) .collect(); - arch.run(|| fill_lower_norm_sums(arch, storage, point_count, &squared_norms)); + // Initialize the first two terms before GEMM adds the dot-product term. + for source in 0..point_count { + let row = &mut storage[source * point_count..source * point_count + source + 1]; + let source_norm = squared_norms[source]; + let mut target = 0; + while target < row.len() { + row[target] = source_norm + squared_norms[target]; + target += 1; + } + } diskann_linalg::sgemm_aat_lower_add( point_count, points.ncols(), @@ -85,14 +55,7 @@ impl LeafMetric for L2 { } impl LeafMetric for Cosine { - fn compute_distances( - _arch: A, - points: MatrixView<'_, f32>, - storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema, - { + fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { let point_count = points.nrows(); // The diagonal supplies each point norm after GEMM computes all dots. diskann_linalg::sgemm_aat_lower( @@ -118,14 +81,7 @@ impl LeafMetric for Cosine { } impl LeafMetric for CosineNormalized { - fn compute_distances( - _arch: A, - points: MatrixView<'_, f32>, - storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema, - { + fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { let point_count = points.nrows(); // The constant in `1 - dot` does not change nearest-first order. diskann_linalg::sgemm_aat_lower( @@ -141,14 +97,7 @@ impl LeafMetric for CosineNormalized { } impl LeafMetric for InnerProduct { - fn compute_distances( - _arch: A, - points: MatrixView<'_, f32>, - storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema, - { + fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { let point_count = points.nrows(); // DiskANN ranks inner products in descending order through `-dot`. diskann_linalg::sgemm_aat_lower( @@ -167,7 +116,6 @@ impl LeafMetric for InnerProduct { #[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] mod tests { use super::*; - use diskann_wide::ARCH; use rstest::rstest; const POINT_COUNT: usize = 2; @@ -190,7 +138,7 @@ mod tests { let points = MatrixView::try_from(&point_values[..], POINT_COUNT, DIMENSION_COUNT).unwrap(); let mut storage = [STALE_DISTANCE; POINT_COUNT * POINT_COUNT]; - M::compute_distances(ARCH, points, &mut storage).unwrap(); + M::compute_distances(points, &mut storage).unwrap(); storage[SECOND_POINT * POINT_COUNT + FIRST_POINT] } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 874fddaf5e..de68241ada 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,9 +5,8 @@ //! Numerical kernels for PiPNN graph construction. //! -//! Metric modules fill flattened distance buffers. Kernel modules traverse those -//! buffers and retain nearest neighbors. Callers select architecture `A` and -//! metric `M` once for each graph build. +//! Metric modules fill portable distance buffers. Kernel modules use architecture +//! `A` to traverse those buffers. Callers select metric `M` once per graph build. #[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 0da345ce06..580ce7085f 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -61,12 +61,7 @@ where if distance_scratch.len() < distance_count { distance_scratch.resize(distance_count, 0.0); } - M::compute_distances( - arch, - points, - leaders, - &mut distance_scratch[..distance_count], - )?; + M::compute_distances(points, leaders, &mut distance_scratch[..distance_count])?; rank_leader_distances( arch, &distance_scratch[..distance_count], diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 77a313ac2a..9d33cfbefd 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -14,9 +14,8 @@ use crate::{ANNError, ANNResult}; use diskann_linalg::Transpose; use diskann_utils::views::MatrixView; use diskann_vector::{Norm, norm::FastL2NormSquared}; -use diskann_wide::SIMDVector; -use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance, simd::PiPNNSIMDSchema}; +use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; /// Store leader values with immutable metric data. /// @@ -46,14 +45,11 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { /// Compute one row-major point-to-leader ranking buffer. /// /// `storage` has `points.nrows() * leader_count` elements. - fn compute_distances( - arch: A, + fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema; + ) -> ANNResult<()>; } /// Compute L2 squared norms with the established sequential reduction order. @@ -75,28 +71,6 @@ fn cosine_norms(vectors: MatrixView<'_, f32>) -> Vec { .collect() } -/// Copy one leader value into the same column of each point row. -/// -/// L2 uses this operation to initialize every row with leader squared norms. -/// Full SIMD groups use dispatched vector stores. The scalar tail copies the rest. -fn fill_rows(arch: A, rows: &mut [f32], row_values: &[f32]) -where - A: PiPNNSIMDSchema, -{ - let row_length = row_values.len(); - let simd_end = row_length - row_length % A::Vector::LANES; - for row in rows.chunks_exact_mut(row_length) { - for start in (0..simd_end).step_by(A::Vector::LANES) { - // SAFETY: both offsets start complete SIMD groups in live slices. - unsafe { - A::Vector::load_simd(arch, row_values.as_ptr().add(start)) - .store_simd(row.as_mut_ptr().add(start)); - } - } - row[simd_end..].copy_from_slice(&row_values[simd_end..]); - } -} - impl PartitionMetric for L2 { type Leaders<'a> = PartitionLeaders<'a, OnceLock>>; @@ -111,20 +85,24 @@ impl PartitionMetric for L2 { leaders.values.nrows() } - fn compute_distances( - arch: A, + fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema, - { + ) -> ANNResult<()> { // The point norm is constant across a point row. It cannot change ranking. let leader_norms = leaders .cache .get_or_init(|| l2_squared_norms(leaders.values)); - arch.run(|| fill_rows(arch, storage, leader_norms)); + // Initialize each point row before GEMM adds the dot-product term. + let leader_count = leader_norms.len(); + for row in storage.chunks_exact_mut(leader_count) { + let mut leader = 0; + while leader < leader_count { + row[leader] = leader_norms[leader]; + leader += 1; + } + } diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -156,15 +134,11 @@ impl PartitionMetric for Cosine { leaders.values.nrows() } - fn compute_distances( - _arch: A, + fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema, - { + ) -> ANNResult<()> { diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -207,15 +181,11 @@ impl PartitionMetric for CosineNormalized { leaders.values.nrows() } - fn compute_distances( - _arch: A, + fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema, - { + ) -> ANNResult<()> { diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -244,15 +214,11 @@ impl PartitionMetric for InnerProduct { leaders.values.nrows() } - fn compute_distances( - _arch: A, + fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, storage: &mut [f32], - ) -> ANNResult<()> - where - A: PiPNNSIMDSchema, - { + ) -> ANNResult<()> { diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -274,7 +240,6 @@ impl PartitionMetric for InnerProduct { #[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] mod tests { use super::*; - use diskann_wide::ARCH; const DIMENSION_COUNT: usize = 2; const STALE_DISTANCE: f32 = 99.0; @@ -291,7 +256,7 @@ mod tests { let leaders = M::create_leaders(matrix(&leader, 1)); let mut storage = [STALE_DISTANCE]; - M::compute_distances(ARCH, matrix(&point, 1), &leaders, &mut storage).unwrap(); + M::compute_distances(matrix(&point, 1), &leaders, &mut storage).unwrap(); storage[0] } @@ -344,7 +309,7 @@ mod tests { let mut distance = [STALE_DISTANCE]; // When - L2::compute_distances(ARCH, point_matrix, &leaders, &mut distance).unwrap(); + L2::compute_distances(point_matrix, &leaders, &mut distance).unwrap(); let actual = leaders.cache.get().unwrap()[0]; // Then @@ -425,7 +390,6 @@ mod tests { let fresh_leaders = L2::create_leaders(matrix(&leader_values, leader_count)); let mut discarded_first_output = [STALE_DISTANCE; 2]; L2::compute_distances( - ARCH, matrix(&first_point, 1), &reused_leaders, &mut discarded_first_output, @@ -436,19 +400,13 @@ mod tests { // When L2::compute_distances( - ARCH, matrix(&second_point, 1), &reused_leaders, &mut reused_output, ) .unwrap(); - L2::compute_distances( - ARCH, - matrix(&second_point, 1), - &fresh_leaders, - &mut fresh_output, - ) - .unwrap(); + L2::compute_distances(matrix(&second_point, 1), &fresh_leaders, &mut fresh_output) + .unwrap(); // Then assert_eq!(reused_output, expected); @@ -472,7 +430,6 @@ mod tests { let fresh_leaders = Cosine::create_leaders(matrix(&leader_values, leader_count)); let mut discarded_first_output = [STALE_DISTANCE; 2]; Cosine::compute_distances( - ARCH, matrix(&first_point, 1), &reused_leaders, &mut discarded_first_output, @@ -483,19 +440,13 @@ mod tests { // When Cosine::compute_distances( - ARCH, matrix(&second_point, 1), &reused_leaders, &mut reused_output, ) .unwrap(); - Cosine::compute_distances( - ARCH, - matrix(&second_point, 1), - &fresh_leaders, - &mut fresh_output, - ) - .unwrap(); + Cosine::compute_distances(matrix(&second_point, 1), &fresh_leaders, &mut fresh_output) + .unwrap(); // Then assert_eq!(reused_output, expected); diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 4742bc385b..edbc57d6b7 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -7,7 +7,7 @@ use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDVector, SupportedLaneCount}; -/// Default SIMD representation used by every PiPNN numerical stage. +/// Default SIMD representation used by both PiPNN ranking kernels. /// /// This alias is the single build-time width selection. type DefaultVector = ::f32x16; @@ -41,7 +41,7 @@ where /// PiPNN SIMD representation for one architecture. pub(super) trait PiPNNSIMDSchema: Architecture { - /// SIMD vector used by every numerical stage. + /// SIMD vector used by both ranking kernels. type Vector: PiPNNSIMDVector; } From b3aade5f44829c0ce3a29c21368050fdab5eb779 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 7 Sep 2026 03:21:48 +0000 Subject: [PATCH 098/229] refactor(pipnn): simplify numerical seams Test ranking directly from distances and keep metric checks at their production seam. Prepare leader norms before stripe workers and retain fixed-width neighbor insertion without a generic comparator. Equal distances need no fixed candidate order. Accept either GEMM zero sign instead of requiring scalar bit patterns across architectures. --- diskann/src/graph/pipnn/leaf_kernel.rs | 741 +++++--------------- diskann/src/graph/pipnn/leaf_metric.rs | 72 +- diskann/src/graph/pipnn/partition_kernel.rs | 367 +++------- diskann/src/graph/pipnn/partition_metric.rs | 148 ++-- 4 files changed, 361 insertions(+), 967 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 3dfcdfebcc..c6c0441897 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -12,7 +12,7 @@ //! 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. +//! Equal distances can select either candidate. The kernel does not rank NaN. //! An unfilled output slot contains [`LeafNeighbor::default`]. All supported //! metrics use the same SIMD-group and single-value traversal. //! @@ -207,7 +207,7 @@ fn scan_fixed_width( worst: &mut [f32], ) where A: PiPNNSIMDSchema, - [LeafNeighbor; N]: SortedInsert, + [LeafNeighbor; N]: NeighborInsert, { let (rows, _) = output.as_chunks_mut::(); // Rayon outlines leaf workers. Reapply target features before the SIMD scan. @@ -253,7 +253,7 @@ fn scan_runtime_width( /// 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 scalar tails preserve pair scan order. +/// both endpoint lists. #[inline(always)] fn scan_point_pairs( arch: A, @@ -332,69 +332,67 @@ fn scan_point_pairs( } } -/// Insert one value that the caller has already found eligible. +/// Insert a neighbor that is closer than the current farthest neighbor. /// -/// The caller must prove that `value` precedes the current last retained value. -/// This method intentionally does not repeat that check. A caller that violates -/// the precondition replaces a valid retained value and corrupts the top-k set. -/// The return value is the new last retained value. -trait SortedInsert { - fn insert_eligible_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T; +/// The pair scan checks eligibility before insertion. The retained neighbors +/// stay in distance order. The result is the new farthest neighbor. +trait NeighborInsert { + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor; } -impl SortedInsert for [T; 1] { +impl NeighborInsert for [LeafNeighbor; 1] { #[inline(always)] - fn insert_eligible_sorted_by(&mut self, value: T, _precedes: impl Fn(T, T) -> bool) -> T { - self[0] = value; - value + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor { + self[0] = candidate; + candidate } } -impl SortedInsert for [T; 2] { +impl NeighborInsert for [LeafNeighbor; 2] { #[inline(always)] - fn insert_eligible_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor { let first = self[0]; - if precedes(value, first) { - self[0] = value; + if candidate.distance < first.distance { + self[0] = candidate; self[1] = first; first } else { - self[1] = value; - value + self[1] = candidate; + candidate } } } -impl SortedInsert for [T; 3] { +impl NeighborInsert for [LeafNeighbor; 3] { #[inline(always)] - fn insert_eligible_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor { let (first, second) = (self[0], self[1]); - if precedes(value, first) { - self[0] = value; + if candidate.distance < first.distance { + self[0] = candidate; self[1] = first; self[2] = second; second - } else if precedes(value, second) { - self[1] = value; + } else if candidate.distance < second.distance { + self[1] = candidate; self[2] = second; second } else { - self[2] = value; - value + self[2] = candidate; + candidate } } } -impl SortedInsert for [T] { +impl NeighborInsert for [LeafNeighbor] { #[inline(always)] - fn insert_eligible_sorted_by(&mut self, value: T, precedes: impl Fn(T, T) -> bool) -> T { + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor { let last = self.len() - 1; let mut slot = last; - while slot > 0 && precedes(value, self[slot - 1]) { + while slot > 0 && candidate.distance < self[slot - 1].distance { self[slot] = self[slot - 1]; slot -= 1; } - self[slot] = value; + self[slot] = candidate; self[last] } } @@ -407,22 +405,18 @@ impl SortedInsert for [T] { #[inline(always)] fn insert_eligible_neighbor(neighbors: &mut R, target: u32, distance: f32) -> f32 where - R: SortedInsert + ?Sized, + R: NeighborInsert + ?Sized, { neighbors - .insert_eligible_sorted_by( - LeafNeighbor::new(target, distance), - |candidate, retained| candidate.distance < retained.distance, - ) + .insert_eligible(LeafNeighbor::new(target, distance)) .distance } #[cfg(test)] mod tests { use super::*; - use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2}; + use crate::graph::pipnn::L2; use diskann_utils::views::{MatrixView, MutMatrixView}; - use diskann_vector::distance::Metric; mod test_support { use std::cmp::Ordering; @@ -430,26 +424,6 @@ mod tests { use super::*; use diskann_wide::arch::{self, Target1}; - pub(super) trait TestMetric { - const METRIC: Metric; - } - - impl TestMetric for L2 { - const METRIC: Metric = Metric::L2; - } - - impl TestMetric for Cosine { - const METRIC: Metric = Metric::Cosine; - } - - impl TestMetric for CosineNormalized { - const METRIC: Metric = Metric::CosineNormalized; - } - - impl TestMetric for InnerProduct { - const METRIC: Metric = Metric::InnerProduct; - } - struct KernelCall<'a> { distance_flatten: &'a [f32], point_count: usize, @@ -474,180 +448,68 @@ mod tests { } } - 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, - Metric::CosineNormalized => -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 { - 1.0 - (dot / (source_norm * target_norm)).clamp(-1.0, 1.0) - } - } - } - } - - fn distance_flatten(metric: Metric, dots: &[f32], points: usize) -> Vec { - let mut distances = vec![0.0; points * points]; - for source in 0..points { - for target in 0..=source { - distances[source * points + target] = reference_distance( - metric, - dots[source * points + target], - dots[source * points + source], - dots[target * points + target], - ); - } - } - distances - } - - fn rank_with_output_width( - metric: Metric, - dots: &[f32], + pub(super) fn rank_distance_fixture( + distances: &[f32], points: usize, output_width: usize, - workspace: &mut LeafKernelWorkspace, ) -> Result, LeafKernelError> { - let distance_flatten = distance_flatten(metric, dots, points); let mut output = vec![LeafNeighbor::default(); points * output_width]; arch::dispatch1_no_features( RankDistances, KernelCall { - distance_flatten: &distance_flatten, + distance_flatten: distances, point_count: points, output: MutMatrixView::try_from(output.as_mut_slice(), points, output_width) .unwrap(), - worst: &mut workspace.worst, + worst: &mut Vec::new(), }, )?; Ok(output) } - pub(super) fn rank_distance_fixture( - dots: &[f32], - points: usize, - output_width: usize, - ) -> Result, LeafKernelError> { - rank_with_output_width( - M::METRIC, - dots, - points, - output_width, - &mut LeafKernelWorkspace::default(), - ) - } - - fn run_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 output = rank_with_output_width(metric, dots, points, leaf_k, workspace) - .expect("valid leaf neighbor width"); - (leaf_k, output) - } - - pub(super) fn run_rank_leaf_distances( - metric: Metric, - dots: &[f32], - points: usize, - requested_k: usize, - ) -> (usize, Vec) { - run_with_workspace( - metric, - dots, - points, - requested_k, - &mut LeafKernelWorkspace::default(), - ) - } - pub(super) fn reference_neighbors( - metric: Metric, - dots: &[f32], + distances: &[f32], points: usize, - requested_k: usize, + width: usize, ) -> Vec { - let leaf_k = requested_k.min(points.saturating_sub(1)); - let mut output = vec![LeafNeighbor::default(); points * leaf_k]; + let mut output = vec![LeafNeighbor::default(); points * width]; for source in 0..points { let mut candidates = Vec::with_capacity(points.saturating_sub(1)); for target in 0..points { if source == target { continue; } - let (row, column) = if source > target { - (source, target) - } else { - (target, source) - }; - 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) { + let row = source.max(target); + let column = source.min(target); + let distance = distances[row * points + column]; + if distance < f32::INFINITY { candidates.push(LeafNeighbor::new(target as u32, distance)); } } - 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] + candidates.sort_unstable_by(|left, right| { + left.distance + .partial_cmp(&right.distance) + .unwrap_or(Ordering::Equal) + }); + let retained = candidates.len().min(width); + output[source * width..source * width + retained] .copy_from_slice(&candidates[..retained]); } output } - /// Build a Gram matrix from points on the line `x = 1`. - pub(super) fn lane_boundary_gram_from_point_vectors( - metric: Metric, - points: usize, - ) -> Vec { - let denominator = (points + 1) as f32; - let point_vectors: Vec<_> = (0..points) - .map(|point| { - let vector = [1.0, (point + 1) as f32 / denominator]; - if metric == Metric::CosineNormalized { - let norm = vector[0].hypot(vector[1]); - [vector[0] / norm, vector[1] / norm] - } else { - vector - } - }) - .collect(); - let mut gram = vec![0.0; points * points]; - for source in 0..points { - for target in 0..points { - gram[source * points + target] = point_vectors[source][0] - * point_vectors[target][0] - + point_vectors[source][1] * point_vectors[target][1]; + /// Give each pair a unique distance, with alternating signs in each row. + pub(super) fn lower_triangle_distances(points: usize) -> Vec { + // NaN in unused entries detects reads of the diagonal or upper triangle. + let mut distances = vec![f32::NAN; points * points]; + for source in 1..points { + for target in 0..source { + let distance = (source * points + target + 1) as f32; + distances[source * points + target] = + if target % 2 == 0 { -distance } else { distance }; } } - gram - } - - pub(super) fn gram_with_uniform_self_dots(points: usize, self_dot: f32) -> Vec { - let mut gram = vec![0.0; points * points]; - for point in 0..points { - gram[point * points + point] = self_dot; - } - gram + distances } } @@ -817,27 +679,6 @@ mod tests { // Then assert_eq!(actual_neighbors, expected_neighbors); } - - #[test] - fn equal_distance_candidate_stays_after_the_existing_neighbor() { - // 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_eligible_neighbor( - &mut actual_neighbors, - tied_candidate.target, - tied_candidate.distance, - ); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - } } mod leaf_neighbor_count_tests { @@ -999,24 +840,20 @@ mod tests { #[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::lane_minus_one_fixed_three(16, 3)] + #[case::one_complete_lane_fixed_one(17, 1)] + #[case::one_complete_lane_fixed_two(17, 2)] + #[case::one_complete_lane_fixed_three(17, 3)] + #[case::lane_plus_one_runtime_width(18, 4)] + #[case::all_non_self_neighbors(17, 16)] + #[case::two_lanes_minus_one_runtime_width(32, 7)] + #[case::two_complete_lanes_runtime_width(33, 7)] + #[case::two_lanes_plus_one_runtime_width(34, 7)] + #[case::four_complete_lanes_runtime_width(65, 7)] + #[case::sixteen_complete_lanes_runtime_width(257, 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, ) { @@ -1025,418 +862,166 @@ mod tests { return; } - // Given - let dots = lane_boundary_gram_from_point_vectors(metric, point_count); - let expected_neighbors = reference_neighbors(metric, &dots, point_count, requested_k); + // Given: the last source row has `point_count - 1` distances. + let distances = lower_triangle_distances(point_count); + let expected_neighbors = reference_neighbors(&distances, point_count, requested_k); // When let actual_neighbors = - run_rank_leaf_distances(metric, &dots, point_count, requested_k).1; - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - } - - const POINT_ZERO: [f32; 2] = [1.0, 0.0]; - const POINT_ONE: [f32; 2] = [0.0, 1.0]; - const POINT_TWO: [f32; 2] = [0.6, 0.8]; - - fn point_dot(left: [f32; 2], right: [f32; 2]) -> f32 { - left[0] * right[0] + left[1] * right[1] - } - - fn three_unit_point_gram() -> [f32; 9] { - let points = [POINT_ZERO, POINT_ONE, POINT_TWO]; - std::array::from_fn(|index| { - let source = index / points.len(); - let target = index % points.len(); - point_dot(points[source], points[target]) - }) - } - - #[test] - fn selects_the_smallest_squared_distance_neighbor_for_each_point_with_l2() { - // Given - let point_zero_two_distance = point_dot(POINT_ZERO, POINT_ZERO) - + point_dot(POINT_TWO, POINT_TWO) - - 2.0 * point_dot(POINT_ZERO, POINT_TWO); - let point_one_two_distance = point_dot(POINT_ONE, POINT_ONE) - + point_dot(POINT_TWO, POINT_TWO) - - 2.0 * point_dot(POINT_ONE, POINT_TWO); - let expected_neighbors = [ - LeafNeighbor::new(2, point_zero_two_distance), - LeafNeighbor::new(2, point_one_two_distance), - LeafNeighbor::new(1, point_one_two_distance), - ]; - let gram = three_unit_point_gram(); - - // When - let actual_neighbors = rank_distance_fixture::(&gram, 3, 1).unwrap(); + rank_distance_fixture(&distances, point_count, requested_k).unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); } #[test] - fn selects_the_highest_similarity_neighbor_for_each_point_with_cosine() { + fn only_the_lower_triangle_supplies_neighbor_distances() { // Given - let expected_neighbors = [ - LeafNeighbor::new(2, 1.0 - point_dot(POINT_ZERO, POINT_TWO)), - LeafNeighbor::new(2, 1.0 - point_dot(POINT_ONE, POINT_TWO)), - LeafNeighbor::new(1, 1.0 - point_dot(POINT_ONE, POINT_TWO)), + #[rustfmt::skip] + let distances = [ + f32::NAN, f32::NAN, f32::NAN, + 3.0, f32::NAN, f32::NAN, + 1.0, 2.0, f32::NAN, ]; - let gram = three_unit_point_gram(); - - // When - let actual_neighbors = rank_distance_fixture::(&gram, 3, 1).unwrap(); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - } - - #[test] - fn selects_the_highest_dot_product_neighbor_for_each_point_with_normalized_cosine() { - // Given let expected_neighbors = [ - LeafNeighbor::new(2, -point_dot(POINT_ZERO, POINT_TWO)), - LeafNeighbor::new(2, -point_dot(POINT_ONE, POINT_TWO)), - LeafNeighbor::new(1, -point_dot(POINT_ONE, POINT_TWO)), + LeafNeighbor::new(2, 1.0), + LeafNeighbor::new(2, 2.0), + LeafNeighbor::new(0, 1.0), ]; - let gram = three_unit_point_gram(); // When - let actual_neighbors = rank_distance_fixture::(&gram, 3, 1).unwrap(); + let actual_neighbors = rank_distance_fixture(&distances, 3, 1).unwrap(); // Then assert_eq!(actual_neighbors, expected_neighbors); } - #[test] - fn selects_the_highest_dot_product_neighbor_for_each_point_with_inner_product() { - // Given - let expected_neighbors = [ - LeafNeighbor::new(2, -point_dot(POINT_ZERO, POINT_TWO)), - LeafNeighbor::new(2, -point_dot(POINT_ONE, POINT_TWO)), - LeafNeighbor::new(1, -point_dot(POINT_ONE, POINT_TWO)), - ]; - let gram = three_unit_point_gram(); - - // When - let actual_neighbors = rank_distance_fixture::(&gram, 3, 1).unwrap(); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - } - - #[test] - fn equal_distances_keep_target_scan_order_with_l2() { - // 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; - // This is the Gram matrix of four orthogonal unit vectors. - #[rustfmt::skip] - let gram = [ - unit_squared_norm, tied_dot_product, tied_dot_product, tied_dot_product, - tied_dot_product, unit_squared_norm, tied_dot_product, tied_dot_product, - tied_dot_product, tied_dot_product, unit_squared_norm, tied_dot_product, - 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_distance_fixture::(&gram, 4, 2).unwrap(); - - // Then - assert_eq!(actual_neighbors, expected_neighbors_in_scan_order); - } - - #[test] - fn finite_negative_l2_ranking_remains_rankable() { + #[rstest] + #[case::finite_negative(-f32::EPSILON)] + #[case::negative_infinity(f32::NEG_INFINITY)] + fn negative_distances_remain_rankable(#[case] distance: f32) { // Given - let point_count = 2; - let nearest_neighbor_count = 1; - let self_dot = 1.0_f32; - let dot_roundoff = f32::EPSILON; - let cross_dot = self_dot + dot_roundoff; - let expected_distance = (-2.0_f32).mul_add(cross_dot, self_dot) + self_dot; - let gram = [self_dot, cross_dot, cross_dot, self_dot]; + let distances = [f32::NAN, f32::NAN, distance, f32::NAN]; let expected = [ - LeafNeighbor::new(1, expected_distance), - LeafNeighbor::new(0, expected_distance), + LeafNeighbor::new(1, distance), + LeafNeighbor::new(0, distance), ]; // When - let actual = - rank_distance_fixture::(&gram, point_count, nearest_neighbor_count).unwrap(); + let actual = rank_distance_fixture(&distances, 2, 1).unwrap(); // Then - assert!(expected_distance < 0.0); assert_eq!(actual, expected); } #[test] - fn scalar_distance_stays_finite_when_twice_the_dot_product_overflows_with_l2() { - // 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 gram = [f32::MAX, dot_product, dot_product, f32::MAX]; - - // When - let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); - - // 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() - ); - } - - #[test] - fn simd_distance_stays_finite_when_twice_the_dot_product_overflows_with_l2() { - // 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 gram = gram_with_uniform_self_dots(points, f32::MAX); - gram[16 * points] = dot_product; - gram[16] = dot_product; - - // When - let actual_neighbors = run_rank_leaf_distances(Metric::L2, &gram, 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 zero_norm_produces_unit_distance_with_cosine() { - // Given - // This is the Gram matrix of one zero vector and two orthogonal unit vectors. - #[rustfmt::skip] - let gram = [ - 0.0, 0.0, 0.0, - 0.0, 1.0, 0.0, - 0.0, 0.0, 1.0, - ]; - let expected_zero_norm_neighbors = - [LeafNeighbor::new(1, 1.0), LeafNeighbor::new(2, 1.0)]; - - // When - let actual_neighbors = rank_distance_fixture::(&gram, 3, 2).unwrap(); - - // Then - assert_eq!(&actual_neighbors[..2], &expected_zero_norm_neighbors); - } - - #[test] - fn zero_target_norm_remains_rankable_in_a_complete_simd_group_with_cosine() { - // Given - let points = 17; - let mut gram = gram_with_uniform_self_dots(points, 1.0); - gram[0] = 0.0; - let expected_zero_norm_neighbor = LeafNeighbor::new(0, 1.0); - - // When - let actual_neighbors = run_rank_leaf_distances(Metric::Cosine, &gram, points, 1).1; - - // Then - assert_eq!(actual_neighbors[16], expected_zero_norm_neighbor); - } - - #[test] - fn similarity_above_one_clamps_to_zero_distance_with_cosine() { - // Given - // A small excess models dot-product roundoff above cosine similarity one. - let rounded_dot_product = 1.000_001; - let gram = [1.0, rounded_dot_product, rounded_dot_product, 1.0]; - let maximum_cosine_similarity = 1.0; - let expected_one_minus_maximum_similarity = 1.0 - maximum_cosine_similarity; - - // When - let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); - - // Then - assert_eq!( - actual_neighbors[0].distance, - expected_one_minus_maximum_similarity - ); - } - - #[test] - fn similarity_below_negative_one_clamps_to_distance_two_with_cosine() { - // Given - // A small excess models dot-product roundoff below cosine similarity minus one. - let rounded_dot_product = -1.000_001; - let gram = [1.0, rounded_dot_product, rounded_dot_product, 1.0]; - let minimum_cosine_similarity = -1.0; - let expected_one_minus_minimum_similarity = 1.0 - minimum_cosine_similarity; - - // When - let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); - - // Then - assert_eq!( - actual_neighbors[0].distance, - expected_one_minus_minimum_similarity - ); - } - - #[test] - fn subnormal_norm_is_treated_as_zero_with_cosine() { - // Given - let subnormal_self_dot = f32::MIN_POSITIVE / 2.0; - let gram = [subnormal_self_dot, 0.0, 0.0, 1.0]; - let zero_norm_similarity = 0.0; - let expected_one_minus_zero_similarity = 1.0 - zero_norm_similarity; + fn tied_distances_fill_capacity_with_distinct_non_self_neighbors() { + // Given: each point has three equally distant candidates for two slots. + let point_count = 4; + let width = 2; + let distances = [1.0; 16]; // When - let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); - - // Then - assert_eq!( - actual_neighbors[0].distance, - expected_one_minus_zero_similarity - ); + let actual = rank_distance_fixture(&distances, point_count, width).unwrap(); + + // Then: any two candidates are valid; no tie order is required. + for (source, neighbors) in actual.chunks_exact(width).enumerate() { + assert_eq!(neighbors[0].distance, 1.0); + assert_eq!(neighbors[1].distance, 1.0); + assert_ne!(neighbors[0].target, neighbors[1].target); + for neighbor in neighbors { + assert!(neighbor.target < point_count as u32); + assert_ne!(neighbor.target, source as u32); + } + } } #[test] fn f32_max_distance_is_still_a_rankable_neighbor() { // Given - let points = 4; - let expected_leaf_k = 3; + let point_count = 4; + let width = 3; + let source = 3; + let mut distances = [1.0; 16]; + distances[source * point_count] = f32::MAX; let expected_last_neighbor = LeafNeighbor::new(0, f32::MAX); - let mut gram = gram_with_uniform_self_dots(points, f32::MAX); - gram[3 * points] = -f32::MAX; - gram[3] = -f32::MAX; // When - let actual_neighbors = - rank_distance_fixture::(&gram, points, expected_leaf_k).unwrap(); + let actual = rank_distance_fixture(&distances, point_count, width).unwrap(); // Then - assert_eq!( - actual_neighbors[3 * expected_leaf_k + expected_leaf_k - 1], - expected_last_neighbor - ); + assert_eq!(actual[source * width + width - 1], expected_last_neighbor); } - #[test] - fn scalar_nan_distance_leaves_the_neighbor_slot_unassigned() { + #[rstest] + #[case::nan(f32::NAN)] + #[case::positive_infinity(f32::INFINITY)] + fn non_rankable_distances_leave_neighbor_slots_unassigned( + #[case] distance: f32, + #[values(2, 17, 18)] point_count: usize, + ) { // Given - let gram = [1.0, f32::NAN, f32::NAN, 1.0]; - let expected_unassigned_neighbors = [LeafNeighbor::default(), LeafNeighbor::default()]; + let distances = vec![distance; point_count * point_count]; + let expected = vec![LeafNeighbor::default(); point_count]; // When - let actual_neighbors = rank_distance_fixture::(&gram, 2, 1).unwrap(); + let actual = rank_distance_fixture(&distances, point_count, 1).unwrap(); // Then - assert_eq!(actual_neighbors, expected_unassigned_neighbors); - } - - #[test] - fn complete_simd_group_without_eligible_distances_leaves_source_unassigned_with_l2() { - // Given - let point_count = 17; - let source = 16; - let requested_k = 1; - let mut gram = gram_with_uniform_self_dots(point_count, 1.0); - gram[source * point_count..source * point_count + source].fill(f32::NEG_INFINITY); - let expected_unassigned_neighbor = LeafNeighbor::default(); - - // When - let actual_neighbors = - rank_distance_fixture::(&gram, point_count, requested_k).unwrap(); - - // Then - assert_eq!( - actual_neighbors[source * requested_k], - expected_unassigned_neighbor - ); + assert_eq!(actual, expected); } #[rstest] - #[case::l2(Metric::L2, 2.0)] - #[case::cosine(Metric::Cosine, 1.0)] - #[case::normalized_cosine(Metric::CosineNormalized, -0.0)] - #[case::inner_product(Metric::InnerProduct, -0.0)] - fn simd_nan_distance_cannot_replace_a_finite_neighbor( - #[case] metric: Metric, - #[case] expected_distance: f32, + #[case::nan(f32::NAN)] + #[case::positive_infinity(f32::INFINITY)] + fn non_rankable_candidates_cannot_replace_finite_neighbors( + #[case] distance: f32, + #[values(17, 18)] point_count: usize, ) { - // Given - let points = 17; - let mut gram = gram_with_uniform_self_dots(points, 1.0); - gram[16 * points] = f32::NAN; - gram[16] = f32::NAN; - let expected_finite_neighbor = LeafNeighbor::new(1, expected_distance); + // Given: the last candidate is in a full SIMD group or its scalar tail. + let source = point_count - 1; + let invalid_target = source - 1; + let mut distances = vec![1.0; point_count * point_count]; + distances[source * point_count + invalid_target] = distance; // When - let actual_neighbors = run_rank_leaf_distances(metric, &gram, points, 1).1; + let actual = rank_distance_fixture(&distances, point_count, 1).unwrap(); // Then - assert_eq!(actual_neighbors[16], expected_finite_neighbor); + assert!(actual[source].target < invalid_target as u32); + assert_eq!(actual[source].distance, 1.0); } - #[rstest] - #[case::l2(Metric::L2, 2.0)] - #[case::normalized_cosine(Metric::CosineNormalized, -0.0)] - #[case::inner_product(Metric::InnerProduct, -0.0)] - fn simd_positive_infinity_cannot_fill_a_neighbor_slot( - #[case] metric: Metric, - #[case] expected_distance: f32, - ) { - // Given: negative-infinite dot products produce positive-infinite distances here. - let points = 17; - let mut gram = gram_with_uniform_self_dots(points, 1.0); - gram[16 * points] = f32::NEG_INFINITY; - gram[16] = f32::NEG_INFINITY; - let expected_finite_neighbor = LeafNeighbor::new(1, expected_distance); + #[test] + fn scalar_tail_updates_both_endpoints_after_a_rejected_simd_group() { + // Given: the final pair is the only rankable pair in the leaf. + let point_count = 18; + let source = 17; + let target = 16; + let mut distances = vec![f32::INFINITY; point_count * point_count]; + distances[source * point_count + target] = -1.0; + let mut expected = vec![LeafNeighbor::default(); point_count]; + expected[source] = LeafNeighbor::new(target as u32, -1.0); + expected[target] = LeafNeighbor::new(source as u32, -1.0); // When - let actual_neighbors = run_rank_leaf_distances(metric, &gram, points, 1).1; + let actual = rank_distance_fixture(&distances, point_count, 1).unwrap(); // Then - assert_eq!(actual_neighbors[16], expected_finite_neighbor); + assert_eq!(actual, expected); } #[test] fn singleton_leaf_has_no_neighbors() { // Given - let singleton_point_count = 1; - let singleton_gram = [4.0]; - let expected_zero_neighbor_width = 0; - let expected_no_neighbors: [LeafNeighbor; 0] = []; + let distances = [f32::NAN]; + let expected: [LeafNeighbor; 0] = []; // When - let actual_neighbors = rank_distance_fixture::( - &singleton_gram, - singleton_point_count, - expected_zero_neighbor_width, - ) - .unwrap(); + let actual = rank_distance_fixture(&distances, 1, 0).unwrap(); // Then - assert_eq!(actual_neighbors, expected_no_neighbors); + assert_eq!(actual, expected); } #[test] @@ -1444,18 +1029,16 @@ mod tests { // Given let point_count = 3; let invalid_neighbor_width = point_count; - let maximum_non_self_width = point_count - 1; - let gram = [0.0; 9]; + let distances = [0.0; 9]; let expected_error = LeafKernelError::InvalidNeighborCount { points: point_count, neighbors: invalid_neighbor_width, - maximum: maximum_non_self_width, + maximum: point_count - 1, }; // When let actual_error = - rank_distance_fixture::(&gram, point_count, invalid_neighbor_width) - .unwrap_err(); + rank_distance_fixture(&distances, point_count, invalid_neighbor_width).unwrap_err(); // Then assert_eq!(actual_error, expected_error); diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index 30737eb89d..c441b94109 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -17,6 +17,7 @@ use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; /// /// An implementation initializes the diagonal and lower triangle. The upper /// triangle stays unspecified. The input matrix has one point in each row. +/// A zero distance can have either sign. Equal distances can select either candidate. pub(super) trait LeafMetric: Send + Sync + 'static { /// Compute ranking distances for all unordered point pairs. /// @@ -98,17 +99,8 @@ impl LeafMetric for CosineNormalized { impl LeafMetric for InnerProduct { fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { - let point_count = points.nrows(); - // DiskANN ranks inner products in descending order through `-dot`. - diskann_linalg::sgemm_aat_lower( - point_count, - points.ncols(), - -1.0, - points.as_slice(), - storage, - ) - .map_err(ANNError::new)?; - Ok(()) + // Both metrics rank with `-dot`. Their graph-pruning policies stay separate. + CosineNormalized::compute_distances(points, storage) } } @@ -182,16 +174,16 @@ mod tests { ); } - #[test] - fn cosine_ranking_is_one_when_one_point_has_zero_norm() { - // Given - let zero_point = [0.0_f32, 0.0]; + #[rstest] + #[case::zero([0.0, 0.0])] + #[case::subnormal([f32::MIN_POSITIVE.sqrt() / 2.0, 0.0])] + fn small_norm_produces_unit_cosine_ranking(#[case] small_point: [f32; DIMENSION_COUNT]) { + // Given: a zero or subnormal norm represents zero similarity. let unit_point = [1.0_f32, 0.0]; - let zero_similarity = 0.0_f32; - let expected = 1.0 - zero_similarity; + let expected = 1.0; // When - let actual = compute_pair_ranking::(zero_point, unit_point); + let actual = compute_pair_ranking::(small_point, unit_point); // Then assert_eq!(actual, expected); @@ -228,21 +220,43 @@ mod tests { } #[rstest] - #[case::positive_zero(0.0)] - #[case::negative_zero(-0.0)] - #[trace] - fn inner_product_ranking_matches_negated_scalar_dot_bits(#[case] zero_coordinate: f32) { - // Given - let first_point = [zero_coordinate, 0.0]; - let second_point = [1.0_f32, 0.0]; - let dot = first_point[0].mul_add(second_point[0], first_point[1] * second_point[1]); - let expected = -dot; + #[case::l2(compute_pair_ranking::)] + #[case::cosine(compute_pair_ranking::)] + #[case::normalized_cosine(compute_pair_ranking::)] + #[case::inner_product(compute_pair_ranking::)] + fn nan_coordinate_produces_nan_ranking( + #[case] compute: fn([f32; DIMENSION_COUNT], [f32; DIMENSION_COUNT]) -> f32, + ) { + // Given: neither vector has zero norm. + let first_point = [f32::NAN, 1.0]; + let second_point = [1.0, 0.0]; // When - let actual = compute_pair_ranking::(first_point, second_point); + let actual = compute(first_point, second_point); // Then - assert_eq!(actual.to_bits(), expected.to_bits()); + assert!(actual.is_nan()); + } + + #[rstest] + #[case::inner_product(InnerProduct::compute_distances)] + #[case::normalized_cosine(CosineNormalized::compute_distances)] + fn orthogonal_vectors_have_zero_ranking( + #[case] compute: fn(MatrixView<'_, f32>, &mut [f32]) -> ANNResult<()>, + #[values(2, 17, 129)] dimensions: usize, + ) { + // Given: the two unit vectors have disjoint nonzero coordinates. + let mut values = vec![0.0; 2 * dimensions]; + values[0] = 1.0; + values[dimensions + 1] = 1.0; + let points = MatrixView::try_from(values.as_slice(), 2, dimensions).unwrap(); + let mut distances = [STALE_DISTANCE; 4]; + + // When + compute(points, &mut distances).unwrap(); + + // Then: the contract permits either zero sign. + assert_eq!(distances[2], 0.0); } } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 580ce7085f..42be8c54d8 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -10,7 +10,7 @@ //! leader-column IDs for partition scatter. //! //! L2 omits the assigned point's norm because it is constant across all sampled -//! leaders. Equal distances keep sampled-leader order. NaN is not rankable. An +//! leaders. Equal distances can select either leader. NaN is not rankable. An //! unfilled output slot contains [`UNASSIGNED_LEADER`]. use crate::ANNResult; @@ -146,8 +146,8 @@ fn select_point_leaders( /// 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. +/// `first_leader` is the matrix-column ID of the first lane. The function reads +/// selected lanes from low to high. fn insert_leader_lanes(distances: F, first_leader: usize, ranked_leaders: &mut [(u32, f32)]) where F: PiPNNSIMDVector, @@ -175,8 +175,8 @@ where /// 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. `ranked_leaders` -/// stores retained centers in nearest-first order. Equal distances and NaN do not -/// enter, so sampled-leader order resolves ties. +/// stores retained centers in nearest-first order. A candidate enters only when +/// its distance is less than the current farthest distance. NaN does not enter. #[inline(always)] fn insert_leader(ranked_leaders: &mut [(u32, f32)], leader: u32, distance: f32) { let threshold = ranked_leaders.len() - 1; @@ -197,19 +197,11 @@ mod tests { use super::*; use crate::graph::pipnn::Cosine; use diskann_utils::views::{Matrix, MatrixView, MutMatrixView}; - use diskann_vector::distance::Metric; mod test_support { use super::*; use diskann_wide::arch::{self, Target1}; - #[derive(Clone, Copy)] - pub(super) struct DotFixture<'a> { - dots: MatrixView<'a, f32>, - point_norms: &'a [f32], - leader_norms: &'a [f32], - } - struct KernelCall<'a> { distance_flatten: &'a [f32], point_count: usize, @@ -236,68 +228,18 @@ mod tests { } } - pub(super) fn partition_input<'a>( - dots: &'a [f32], - point_count: usize, - leader_count: usize, - point_norms: &'a [f32], - leader_norms: &'a [f32], - ) -> DotFixture<'a> { - DotFixture { - dots: MatrixView::try_from(dots, point_count, leader_count).unwrap(), - point_norms, - leader_norms, - } - } - - fn reference_distance(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 => -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 distance_flatten(metric: Metric, input: DotFixture<'_>) -> Vec { - input - .dots - .row_iter() - .enumerate() - .flat_map(|(point, dots)| { - dots.iter().enumerate().map(move |(leader, &dot)| { - reference_distance( - metric, - dot, - input.point_norms.get(point).copied().unwrap_or(0.0), - input.leader_norms.get(leader).copied().unwrap_or(0.0), - ) - }) - }) - .collect() - } - pub(super) fn rank_distance_fixture( - metric: Metric, - input: DotFixture<'_>, + distances: MatrixView<'_, f32>, nearest_leader_count: usize, ) -> Vec { - let distance_flatten = distance_flatten(metric, input); - let mut output = Matrix::new(u32::MAX, input.dots.nrows(), nearest_leader_count); + let mut output = + Matrix::new(UNASSIGNED_LEADER, distances.nrows(), nearest_leader_count); arch::dispatch1_no_features( RankDistances, KernelCall { - distance_flatten: &distance_flatten, - point_count: input.dots.nrows(), - leader_count: input.dots.ncols(), + distance_flatten: distances.as_slice(), + point_count: distances.nrows(), + leader_count: distances.ncols(), output: output.as_mut_view(), ranked_leaders: &mut Vec::new(), }, @@ -306,106 +248,37 @@ mod tests { } pub(super) fn reference_assignments( - metric: Metric, - input: DotFixture<'_>, + distances: MatrixView<'_, f32>, nearest_leader_count: usize, ) -> Vec { - let mut output = vec![UNASSIGNED_LEADER; input.dots.nrows() * nearest_leader_count]; - for (point, (dots, assignments)) in input - .dots + let mut output = vec![UNASSIGNED_LEADER; distances.nrows() * nearest_leader_count]; + for (row, assignments) in distances .row_iter() .zip(output.chunks_exact_mut(nearest_leader_count)) - .enumerate() { - let point_norm = input.point_norms.get(point).copied().unwrap_or(0.0); - let mut candidates: Vec<_> = dots + let mut candidates: Vec<_> = row .iter() .enumerate() - .filter_map(|(leader, &dot)| { - let leader_norm = input.leader_norms.get(leader).copied().unwrap_or(0.0); - let distance = reference_distance(metric, dot, point_norm, leader_norm); - (distance.partial_cmp(&f32::INFINITY) == Some(std::cmp::Ordering::Less)) - .then_some((leader as u32, distance)) + .filter_map(|(leader, &distance)| { + (distance < f32::INFINITY).then_some((leader as u32, distance)) }) .collect(); - candidates.sort_by(|left, right| left.1.total_cmp(&right.1)); + candidates.sort_unstable_by(|left, right| { + left.1 + .partial_cmp(&right.1) + .unwrap_or(std::cmp::Ordering::Equal) + }); for (destination, (leader, _)) in assignments.iter_mut().zip(candidates) { *destination = leader; } } output } - - /// Build ranking input from two axis points and leaders between those axes. - pub(super) fn lane_boundary_input_from_point_and_leader_vectors( - metric: Metric, - leader_count: usize, - ) -> (Vec, Vec, Vec) { - let points = [[1.0_f32, 0.0], [0.0, 1.0]]; - let denominator = (leader_count + 1) as f32; - let leaders: Vec<_> = (0..leader_count) - .map(|leader| { - let second_component = (leader + 1) as f32 / denominator; - let vector = [1.0 - second_component, second_component]; - if metric == Metric::CosineNormalized { - let norm = vector[0].hypot(vector[1]); - [vector[0] / norm, vector[1] / norm] - } else { - vector - } - }) - .collect(); - let dots = points - .iter() - .flat_map(|point| { - leaders - .iter() - .map(|leader| point[0] * leader[0] + point[1] * leader[1]) - }) - .collect(); - let point_norms = if metric == Metric::Cosine { - points - .iter() - .map(|point| point[0].hypot(point[1])) - .collect() - } else { - Vec::new() - }; - let leader_norms = match metric { - Metric::L2 => leaders - .iter() - .map(|leader| leader[0] * leader[0] + leader[1] * leader[1]) - .collect(), - Metric::Cosine => leaders - .iter() - .map(|leader| leader[0].hypot(leader[1])) - .collect(), - Metric::CosineNormalized | Metric::InnerProduct => Vec::new(), - }; - (dots, point_norms, leader_norms) - } } mod insert_leader_tests { use super::*; - #[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 nan_distance_does_not_enter_the_topk() { // Given @@ -512,175 +385,95 @@ mod tests { #[case::lane_minus_one(15, 3)] #[case::one_complete_lane(16, 3)] #[case::lane_plus_one(17, 4)] + #[case::all_leaders(17, 17)] #[case::two_lanes_minus_one(31, 7)] #[case::two_complete_lanes(32, 7)] #[case::two_lanes_plus_one(33, 7)] + #[case::root_partition(1000, 64)] #[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] nearest_leader_count: usize, ) { - // Given - let (dots, point_norms, leader_norms) = - lane_boundary_input_from_point_and_leader_vectors(metric, leader_count); - let input = partition_input(&dots, 2, leader_count, &point_norms, &leader_norms); - let expected_assignments = reference_assignments(metric, input, nearest_leader_count); - - // When - let actual_assignments = rank_distance_fixture(metric, input, nearest_leader_count); - - // Then - assert_eq!(actual_assignments, expected_assignments); - } - - #[test] - fn equal_distances_keep_sampled_leader_order_with_l2() { - // Given - let point_count = 1; - let leader_count = 4; - let nearest_leader_count = 2; - 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]; - - let input = - partition_input(&dots, point_count, leader_count, &[], &leader_squared_norms); - - // When - let actual_assignments = rank_distance_fixture(Metric::L2, input, nearest_leader_count); - - // Then - assert_eq!(actual_assignments, expected_sampled_leader_order); - } - - #[test] - fn zero_norm_keeps_sampled_leader_order_with_cosine() { - // Given - let point_count = 1; - let leader_count = 2; - let nearest_leader_count = 2; - let dots = [0.0, 0.0]; - let point_norms = [0.0]; - let leader_norms = [1.0, 1.0]; - let expected_sampled_leader_order = [0, 1]; - - let input = partition_input( - &dots, - point_count, - leader_count, - &point_norms, - &leader_norms, - ); - - // When - let actual_assignments = - rank_distance_fixture(Metric::Cosine, input, nearest_leader_count); - - // Then - assert_eq!(actual_assignments, expected_sampled_leader_order); - } - - #[test] - fn zero_point_norm_keeps_first_leader_in_a_complete_simd_group_with_cosine() { - // Given - let point_count = 1; - let leader_count = 17; - let nearest_leader_count = 1; - let dots = [0.0; 17]; - let point_norms = [0.0]; - let leader_norms = [1.0; 17]; - let expected_first_leader = [0]; + // Given: unique distances increase in one row and decrease in the other. + let values: Vec<_> = (0..leader_count) + .chain((0..leader_count).rev()) + .map(|leader| leader as f32 - leader_count as f32 / 2.0) + .collect(); + let distances = MatrixView::try_from(values.as_slice(), 2, leader_count).unwrap(); + let expected = reference_assignments(distances, nearest_leader_count); // When - let actual_assignment = rank_distance_fixture( - Metric::Cosine, - partition_input( - &dots, - point_count, - leader_count, - &point_norms, - &leader_norms, - ), - nearest_leader_count, - ); + let actual = rank_distance_fixture(distances, nearest_leader_count); // Then - assert_eq!(actual_assignment, expected_first_leader); + assert_eq!(actual, expected); } - #[test] - fn f32_max_distance_is_still_rankable() { + #[rstest] + #[case::scalar(4)] + #[case::simd_with_tail(17)] + fn nearer_leader_displaces_one_of_the_tied_candidates(#[case] leader_count: usize) { // Given - let point_count = 1; - let leader_count = 8; - let nearest_leader_count = leader_count; - let maximum_rankable_distance = f32::MAX; - let dot_product_that_produces_it = -maximum_rankable_distance; - let mut dots = [0.0; 8]; - dots[7] = dot_product_that_produces_it; - let expected_all_leaders_in_scan_order = [0, 1, 2, 3, 4, 5, 6, 7]; - - let input = partition_input(&dots, point_count, leader_count, &[], &[]); + let nearest_leader = leader_count - 1; + let mut values = vec![1.0; leader_count]; + values[nearest_leader] = 0.0; + let distances = MatrixView::try_from(values.as_slice(), 1, leader_count).unwrap(); // When - let actual_assignments = - rank_distance_fixture(Metric::InnerProduct, input, nearest_leader_count); + let actual = rank_distance_fixture(distances, 3); - // Then - assert_eq!(actual_assignments, expected_all_leaders_in_scan_order); + // Then: the remaining two slots can contain any distinct tied leaders. + assert_eq!(actual[0], nearest_leader as u32); + assert!(actual[1] < nearest_leader as u32); + assert!(actual[2] < nearest_leader as u32); + assert_ne!(actual[1], actual[2]); } - #[test] - fn nan_leader_does_not_displace_finite_leaders_with_inner_product() { - // Given - let point_count = 1; - let leader_count = 3; - let nearest_leader_count = 2; - let dots = [f32::NAN, 3.0, 2.0]; - let expected_finite_leaders = [1, 2]; - - let input = partition_input(&dots, point_count, leader_count, &[], &[]); + #[rstest] + #[case::nan(f32::NAN)] + #[case::positive_infinity(f32::INFINITY)] + fn non_rankable_distances_leave_only_unfilled_ranks_unassigned( + #[case] unrankable: f32, + #[values(3, 17, 33)] leader_count: usize, + ) { + // Given: only the first row has rankable candidates. + let mut values = vec![unrankable; 2 * leader_count]; + values[leader_count - 2] = -1.0; + values[leader_count - 1] = 2.0; + let distances = MatrixView::try_from(values.as_slice(), 2, leader_count).unwrap(); + let expected = [ + (leader_count - 2) as u32, + (leader_count - 1) as u32, + UNASSIGNED_LEADER, + UNASSIGNED_LEADER, + UNASSIGNED_LEADER, + UNASSIGNED_LEADER, + ]; // When - let actual_assignments = - rank_distance_fixture(Metric::InnerProduct, input, nearest_leader_count); + let actual = rank_distance_fixture(distances, 3); // Then - assert_eq!(actual_assignments, expected_finite_leaders); + assert_eq!(actual, expected); } #[test] - fn nan_leader_does_not_displace_finite_leaders_with_cosine() { - // Given - let point_count = 1; - let leader_count = 3; - let nearest_leader_count = 2; - let dots = [f32::NAN, 0.75, 0.5]; - let point_norms = [1.0]; - let leader_norms = [1.0; 3]; - let expected_finite_leaders = [1, 2]; - - let input = partition_input( - &dots, - point_count, - leader_count, - &point_norms, - &leader_norms, - ); + fn negative_infinity_and_f32_max_remain_rankable() { + // Given: the smallest distance is in the scalar tail. + let mut values = [f32::NAN; 17]; + values[0] = f32::MAX; + values[1] = 0.0; + values[2] = -f32::EPSILON; + values[16] = f32::NEG_INFINITY; + let distances = MatrixView::try_from(&values[..], 1, values.len()).unwrap(); + let expected = [16, 2, 1, 0]; // When - let actual_assignments = - rank_distance_fixture(Metric::Cosine, input, nearest_leader_count); + let actual = rank_distance_fixture(distances, 4); // Then - assert_eq!(actual_assignments, expected_finite_leaders); + assert_eq!(actual, expected); } } } diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 9d33cfbefd..9d4271736f 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -8,8 +8,6 @@ //! A ranking distance preserves nearest-first order. One leader set serves all //! point stripes in a partition split. -use std::sync::OnceLock; - use crate::{ANNError, ANNResult}; use diskann_linalg::Transpose; use diskann_utils::views::MatrixView; @@ -19,11 +17,11 @@ use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; /// Store leader values with immutable metric data. /// -/// The metric selects `Cache`. An L2 cache stores squared norms. A cosine cache -/// stores norms. `OnceLock` lets concurrent point stripes initialize data once. -pub(super) struct PartitionLeaders<'a, Cache> { +/// L2 stores squared norms. Cosine stores norms. Construction computes them once +/// before point stripes share the leader set. Other metrics need no norms. +pub(super) struct PartitionLeaders<'a, Norms> { values: MatrixView<'a, f32>, - cache: Cache, + norms: Norms, } /// Fill one flattened point-to-leader ranking buffer. @@ -45,6 +43,7 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { /// Compute one row-major point-to-leader ranking buffer. /// /// `storage` has `points.nrows() * leader_count` elements. + /// A zero distance can have either sign. Equal distances can select either leader. fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, @@ -52,10 +51,9 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { ) -> ANNResult<()>; } -/// Compute L2 squared norms with the established sequential reduction order. +/// Compute L2 squared norms with sequential accumulation. /// -/// Small rounding differences can change a leader tie. Keep this order equal to -/// the pre-GEMM implementation. +/// This order fixes rounding of the leader term before GEMM adds the dot term. fn l2_squared_norms(vectors: MatrixView<'_, f32>) -> Vec { vectors .row_iter() @@ -72,12 +70,12 @@ fn cosine_norms(vectors: MatrixView<'_, f32>) -> Vec { } impl PartitionMetric for L2 { - type Leaders<'a> = PartitionLeaders<'a, OnceLock>>; + type Leaders<'a> = PartitionLeaders<'a, Vec>; fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { PartitionLeaders { values, - cache: OnceLock::new(), + norms: l2_squared_norms(values), } } @@ -91,9 +89,7 @@ impl PartitionMetric for L2 { storage: &mut [f32], ) -> ANNResult<()> { // The point norm is constant across a point row. It cannot change ranking. - let leader_norms = leaders - .cache - .get_or_init(|| l2_squared_norms(leaders.values)); + let leader_norms = &leaders.norms; // Initialize each point row before GEMM adds the dot-product term. let leader_count = leader_norms.len(); for row in storage.chunks_exact_mut(leader_count) { @@ -121,12 +117,12 @@ impl PartitionMetric for L2 { } impl PartitionMetric for Cosine { - type Leaders<'a> = PartitionLeaders<'a, OnceLock>>; + type Leaders<'a> = PartitionLeaders<'a, Vec>; fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { PartitionLeaders { values, - cache: OnceLock::new(), + norms: cosine_norms(values), } } @@ -153,7 +149,7 @@ impl PartitionMetric for Cosine { ) .map_err(ANNError::new)?; let point_norms = cosine_norms(points); - let leader_norms = leaders.cache.get_or_init(|| cosine_norms(leaders.values)); + let leader_norms = &leaders.norms; let leader_count = leaders.values.nrows(); // Convert each dot to cosine distance. Reuse leader norms across stripes. for (row, &point_norm) in storage @@ -174,7 +170,7 @@ impl PartitionMetric for CosineNormalized { type Leaders<'a> = PartitionLeaders<'a, ()>; fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { - PartitionLeaders { values, cache: () } + PartitionLeaders { values, norms: () } } fn leader_count(leaders: &Self::Leaders<'_>) -> usize { @@ -204,14 +200,14 @@ impl PartitionMetric for CosineNormalized { } impl PartitionMetric for InnerProduct { - type Leaders<'a> = PartitionLeaders<'a, ()>; + type Leaders<'a> = ::Leaders<'a>; fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { - PartitionLeaders { values, cache: () } + CosineNormalized::create_leaders(values) } fn leader_count(leaders: &Self::Leaders<'_>) -> usize { - leaders.values.nrows() + CosineNormalized::leader_count(leaders) } fn compute_distances( @@ -219,20 +215,7 @@ impl PartitionMetric for InnerProduct { leaders: &Self::Leaders<'_>, storage: &mut [f32], ) -> ANNResult<()> { - diskann_linalg::sgemm( - Transpose::None, - Transpose::Ordinary, - points.nrows(), - leaders.values.nrows(), - points.ncols(), - -1.0, - points.as_slice(), - leaders.values.as_slice(), - None, - storage, - ) - .map_err(ANNError::new)?; - Ok(()) + CosineNormalized::compute_distances(points, leaders, storage) } } @@ -261,6 +244,26 @@ mod tests { storage[0] } + mod create_leaders_tests { + use super::*; + + #[test] + fn l2_leader_norms_preserve_the_sequential_reduction_order() { + // Given: each small square is half an ULP at the first squared norm. + // Each sequential addition rounds back to that norm. + let mut values = [1.0_f32; 129]; + values[0] = 4096.0; + let expected_squared_norm = 16_777_216.0_f32; + let matrix = MatrixView::try_from(&values[..], 1, values.len()).unwrap(); + + // When + let leaders = L2::create_leaders(matrix); + + // Then + assert_eq!(leaders.norms, [expected_squared_norm]); + } + } + mod compute_distances_tests { use super::*; @@ -280,43 +283,6 @@ mod tests { assert_eq!(actual, expected); } - #[test] - fn l2_leader_cache_preserves_the_sequential_reduction_order() { - fn next_regression_value(state: &mut u64) -> f32 { - *state ^= *state << 13; - *state ^= *state >> 7; - *state ^= *state << 17; - (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 - } - - // Given - const DIMENSIONS: usize = 129; - const REGRESSION_SEED: u64 = 0x3a85_f952_c718_6e49; - let mut state = REGRESSION_SEED; - let _point_values: Vec = (0..DIMENSIONS) - .map(|_| next_regression_value(&mut state)) - .collect(); - let leader_values: Vec = (0..DIMENSIONS) - .map(|_| next_regression_value(&mut state)) - .collect(); - let expected: f32 = leader_values.iter().map(|value| value * value).sum(); - let reassociated = FastL2NormSquared.evaluate(leader_values.as_slice()); - let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), 1, DIMENSIONS).unwrap(); - let leaders = L2::create_leaders(leader_matrix); - let zero_point = vec![0.0_f32; DIMENSIONS]; - let point_matrix = MatrixView::try_from(zero_point.as_slice(), 1, DIMENSIONS).unwrap(); - let mut distance = [STALE_DISTANCE]; - - // When - L2::compute_distances(point_matrix, &leaders, &mut distance).unwrap(); - let actual = leaders.cache.get().unwrap()[0]; - - // Then - assert_ne!(expected.to_bits(), reassociated.to_bits()); - assert_eq!(actual.to_bits(), expected.to_bits()); - } - #[test] fn cosine_ranking_equals_one_minus_normalized_similarity() { // Given @@ -337,6 +303,44 @@ mod tests { ); } + #[rstest::rstest] + #[case::zero_point([0.0, 0.0], [1.0, 0.0])] + #[case::zero_leader([1.0, 0.0], [0.0, 0.0])] + #[case::subnormal_point([f32::MIN_POSITIVE.sqrt() / 2.0, 0.0], [1.0, 0.0])] + #[case::subnormal_leader([1.0, 0.0], [f32::MIN_POSITIVE.sqrt() / 2.0, 0.0])] + fn small_norm_produces_unit_cosine_ranking( + #[case] point: [f32; DIMENSION_COUNT], + #[case] leader: [f32; DIMENSION_COUNT], + ) { + // Given: a zero or subnormal norm represents zero similarity. + let expected = 1.0; + + // When + let actual = compute_one_ranking::(point, leader); + + // Then + assert_eq!(actual, expected); + } + + #[rstest::rstest] + #[case::l2(compute_one_ranking::)] + #[case::cosine(compute_one_ranking::)] + #[case::normalized_cosine(compute_one_ranking::)] + #[case::inner_product(compute_one_ranking::)] + fn nan_coordinate_produces_nan_ranking( + #[case] compute: fn([f32; DIMENSION_COUNT], [f32; DIMENSION_COUNT]) -> f32, + ) { + // Given: neither vector has zero norm. + let point = [f32::NAN, 1.0]; + let leader = [1.0, 0.0]; + + // When + let actual = compute(point, leader); + + // Then + assert!(actual.is_nan()); + } + #[test] fn normalized_cosine_ranking_equals_the_negative_dot_product() { // Given From 55fd99b33613d2ea9c657e23e6a2c8e5cd73d854 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 7 Sep 2026 04:17:13 +0000 Subject: [PATCH 099/229] refactor(pipnn): simplify leaf kernel flow Reject invalid neighbor widths before distance evaluation or scratch mutation. Keep ranking infallible and return the farthest distance directly from insertion. --- diskann/src/graph/pipnn/leaf_kernel.rs | 245 +++++++++++-------------- 1 file changed, 108 insertions(+), 137 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index c6c0441897..f7978f62f2 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -28,20 +28,22 @@ use super::{ simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, }; -/// One leaf-local neighbor and its metric distance. +/// One leaf-local neighbor and its ranking distance. +/// +/// A ranking distance preserves nearest-first order. It need not equal the +/// metric distance. #[derive(Clone, Copy, Debug, PartialEq)] pub(super) struct LeafNeighbor { /// Target position in the leaf, not a dataset ID. pub(super) target: u32, - /// Distance from the source point to `target`. + /// Ranking distance from the source point to `target`. pub(super) distance: f32, } impl LeafNeighbor { /// Construct a leaf-local neighbor. /// - /// `target` is a position in the leaf. `distance` is measured from the - /// source point of the output row. + /// The output row determines the source point. pub(super) const fn new(target: u32, distance: f32) -> Self { Self { target, distance } } @@ -65,10 +67,10 @@ pub(super) struct LeafKernelWorkspace { worst: Vec, } -/// Validation error returned by the distance-ranking loop. +/// Invalid output width for leaf-neighbor selection. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub(super) enum LeafKernelError { - /// A source requests more neighbors than the leaf or fixed kernel supports. + /// A source requests more neighbors than the leaf has other points. #[error("invalid leaf neighbor count {neighbors} for {points} points; maximum is {maximum}")] InvalidNeighborCount { points: usize, @@ -91,6 +93,7 @@ pub(super) fn leaf_neighbor_count(points: usize, requested_k: usize) -> usize { /// # Errors /// /// Returns an error for invalid linear-algebra input or output width. +/// Invalid output widths leave the output and workspace unchanged. pub(super) fn select_leaf_neighbors( arch: A, points: MatrixView<'_, f32>, @@ -102,6 +105,7 @@ where M: LeafMetric, { let point_count = points.nrows(); + validate_neighbor_count(point_count, &output).map_err(ANNError::new)?; let distance_count = point_count * point_count; let LeafKernelWorkspace { distance_scratch, @@ -117,28 +121,27 @@ where point_count, output, worst, - ) - .map_err(ANNError::new) + ); + Ok(()) } /// Rank one flattened lower-triangle buffer. /// /// `distance_flatten` contains `point_count * point_count` elements. The metric /// initializes each strict-lower entry. The kernel does not read the upper triangle. +/// The caller validates the output width against the non-self point count. fn rank_leaf_distances( arch: A, distance_flatten: &[f32], point_count: usize, mut output: MutMatrixView<'_, LeafNeighbor>, worst: &mut Vec, -) -> Result<(), LeafKernelError> -where +) where A: PiPNNSIMDSchema, { - validate_neighbor_count(point_count, &output)?; let neighbor_count = output.ncols(); if neighbor_count == 0 { - return Ok(()); + return; } worst.resize(point_count, f32::INFINITY); @@ -176,11 +179,8 @@ where worst, ), } - Ok(()) } -/// Check the safety conditions for the SIMD kernel. -/// /// Check the output width against the number of non-self points. fn validate_neighbor_count( point_count: usize, @@ -218,7 +218,7 @@ fn scan_fixed_width( point_count, worst, |source, target, distance| { - insert_eligible_neighbor(&mut rows[source], target, distance) + rows[source].insert_eligible(LeafNeighbor::new(target, distance)) }, ); }); @@ -244,7 +244,7 @@ fn scan_runtime_width( worst, |source, target, distance| { let first = source * width; - insert_eligible_neighbor(&mut output[first..first + width], target, distance) + output[first..first + width].insert_eligible(LeafNeighbor::new(target, distance)) }, ); }); @@ -335,57 +335,57 @@ fn scan_point_pairs( /// Insert a neighbor that is closer than the current farthest neighbor. /// /// The pair scan checks eligibility before insertion. The retained neighbors -/// stay in distance order. The result is the new farthest neighbor. +/// stay in distance order. The result is their new farthest ranking distance. trait NeighborInsert { - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor; + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32; } impl NeighborInsert for [LeafNeighbor; 1] { #[inline(always)] - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor { + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32 { self[0] = candidate; - candidate + candidate.distance } } impl NeighborInsert for [LeafNeighbor; 2] { #[inline(always)] - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor { + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32 { let first = self[0]; if candidate.distance < first.distance { self[0] = candidate; self[1] = first; - first + first.distance } else { self[1] = candidate; - candidate + candidate.distance } } } impl NeighborInsert for [LeafNeighbor; 3] { #[inline(always)] - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor { + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32 { let (first, second) = (self[0], self[1]); if candidate.distance < first.distance { self[0] = candidate; self[1] = first; self[2] = second; - second + second.distance } else if candidate.distance < second.distance { self[1] = candidate; self[2] = second; - second + second.distance } else { self[2] = candidate; - candidate + candidate.distance } } } impl NeighborInsert for [LeafNeighbor] { #[inline(always)] - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> LeafNeighbor { + fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32 { let last = self.len() - 1; let mut slot = last; while slot > 0 && candidate.distance < self[slot - 1].distance { @@ -393,25 +393,10 @@ impl NeighborInsert for [LeafNeighbor] { slot -= 1; } self[slot] = candidate; - self[last] + self[last].distance } } -/// Insert one candidate that the caller has already found nearer than the current farthest. -/// -/// This function intentionally does not reject an ineligible candidate. The pair -/// scan owns the eligibility check so it can filter SIMD lanes before insertion. -/// The return value is the new farthest retained distance. -#[inline(always)] -fn insert_eligible_neighbor(neighbors: &mut R, target: u32, distance: f32) -> f32 -where - R: NeighborInsert + ?Sized, -{ - neighbors - .insert_eligible(LeafNeighbor::new(target, distance)) - .distance -} - #[cfg(test)] mod tests { use super::*; @@ -433,18 +418,18 @@ mod tests { struct RankDistances; - impl Target1, KernelCall<'_>> for RankDistances + impl Target1> for RankDistances where A: PiPNNSIMDSchema, { - fn run(self, arch: A, call: KernelCall<'_>) -> Result<(), LeafKernelError> { + fn run(self, arch: A, call: KernelCall<'_>) { rank_leaf_distances( arch, call.distance_flatten, call.point_count, call.output, call.worst, - ) + ); } } @@ -452,7 +437,7 @@ mod tests { distances: &[f32], points: usize, output_width: usize, - ) -> Result, LeafKernelError> { + ) -> Vec { let mut output = vec![LeafNeighbor::default(); points * output_width]; arch::dispatch1_no_features( RankDistances, @@ -463,8 +448,8 @@ mod tests { .unwrap(), worst: &mut Vec::new(), }, - )?; - Ok(output) + ); + output } pub(super) fn reference_neighbors( @@ -513,7 +498,7 @@ mod tests { } } - mod insert_eligible_neighbor_tests { + mod insert_eligible_tests { use super::*; #[test] @@ -522,37 +507,15 @@ mod tests { let retained_neighbor = LeafNeighbor::new(1, 4.0); let nearer_candidate = LeafNeighbor::new(2, 2.0); let expected_neighbors = [nearer_candidate]; + let expected_farthest = nearer_candidate.distance; let mut actual_neighbors = [retained_neighbor]; // When - insert_eligible_neighbor( - &mut actual_neighbors, - nearer_candidate.target, - nearer_candidate.distance, - ); + let actual_farthest = actual_neighbors.insert_eligible(nearer_candidate); // Then assert_eq!(actual_neighbors, expected_neighbors); - } - - #[test] - fn direct_call_does_not_recheck_candidate_eligibility() { - // Given: deliberately bypass the pair scan's eligibility check. - let nearest = LeafNeighbor::new(1, 1.0); - let current_farthest = LeafNeighbor::new(2, 3.0); - let ineligible_farther_candidate = LeafNeighbor::new(3, 5.0); - let expected_unchecked_result = [nearest, ineligible_farther_candidate]; - let mut actual_neighbors = [nearest, current_farthest]; - - // When - insert_eligible_neighbor( - &mut actual_neighbors, - ineligible_farther_candidate.target, - ineligible_farther_candidate.distance, - ); - - // Then - assert_eq!(actual_neighbors, expected_unchecked_result); + assert_eq!(actual_farthest, expected_farthest); } #[test] @@ -562,17 +525,15 @@ mod tests { let farthest = LeafNeighbor::new(2, 3.0); let nearer_candidate = LeafNeighbor::new(3, 0.5); let expected_neighbors = [nearer_candidate, nearest]; + let expected_farthest = nearest.distance; let mut actual_neighbors = [nearest, farthest]; // When - insert_eligible_neighbor( - &mut actual_neighbors, - nearer_candidate.target, - nearer_candidate.distance, - ); + let actual_farthest = actual_neighbors.insert_eligible(nearer_candidate); // Then assert_eq!(actual_neighbors, expected_neighbors); + assert_eq!(actual_farthest, expected_farthest); } #[test] @@ -582,17 +543,15 @@ mod tests { let farthest = LeafNeighbor::new(2, 3.0); let eligible_candidate = LeafNeighbor::new(3, 2.0); let expected_neighbors = [nearest, eligible_candidate]; + let expected_farthest = eligible_candidate.distance; let mut actual_neighbors = [nearest, farthest]; // When - insert_eligible_neighbor( - &mut actual_neighbors, - eligible_candidate.target, - eligible_candidate.distance, - ); + let actual_farthest = actual_neighbors.insert_eligible(eligible_candidate); // Then assert_eq!(actual_neighbors, expected_neighbors); + assert_eq!(actual_farthest, expected_farthest); } #[test] @@ -603,17 +562,15 @@ mod tests { let farthest = LeafNeighbor::new(3, 4.0); let nearer_candidate = LeafNeighbor::new(4, 0.5); let expected_neighbors = [nearer_candidate, nearest, middle]; + let expected_farthest = middle.distance; let mut actual_neighbors = [nearest, middle, farthest]; // When - insert_eligible_neighbor( - &mut actual_neighbors, - nearer_candidate.target, - nearer_candidate.distance, - ); + let actual_farthest = actual_neighbors.insert_eligible(nearer_candidate); // Then assert_eq!(actual_neighbors, expected_neighbors); + assert_eq!(actual_farthest, expected_farthest); } #[test] @@ -624,17 +581,15 @@ mod tests { let farthest = LeafNeighbor::new(3, 4.0); let middle_candidate = LeafNeighbor::new(4, 1.5); let expected_neighbors = [nearest, middle_candidate, middle]; + let expected_farthest = middle.distance; let mut actual_neighbors = [nearest, middle, farthest]; // When - insert_eligible_neighbor( - &mut actual_neighbors, - middle_candidate.target, - middle_candidate.distance, - ); + let actual_farthest = actual_neighbors.insert_eligible(middle_candidate); // Then assert_eq!(actual_neighbors, expected_neighbors); + assert_eq!(actual_farthest, expected_farthest); } #[test] @@ -645,17 +600,15 @@ mod tests { let farthest = LeafNeighbor::new(3, 4.0); let eligible_candidate = LeafNeighbor::new(4, 3.0); let expected_neighbors = [nearest, middle, eligible_candidate]; + let expected_farthest = eligible_candidate.distance; let mut actual_neighbors = [nearest, middle, farthest]; // When - insert_eligible_neighbor( - &mut actual_neighbors, - eligible_candidate.target, - eligible_candidate.distance, - ); + let actual_farthest = actual_neighbors.insert_eligible(eligible_candidate); // Then assert_eq!(actual_neighbors, expected_neighbors); + assert_eq!(actual_farthest, expected_farthest); } #[test] @@ -667,17 +620,15 @@ mod tests { let fourth = LeafNeighbor::new(4, 5.0); let candidate = LeafNeighbor::new(5, 2.5); let expected_neighbors = [first, second, candidate, third]; + let expected_farthest = third.distance; let mut actual_neighbors = [first, second, third, fourth]; // When - insert_eligible_neighbor( - actual_neighbors.as_mut_slice(), - candidate.target, - candidate.distance, - ); + let actual_farthest = actual_neighbors.as_mut_slice().insert_eligible(candidate); // Then assert_eq!(actual_neighbors, expected_neighbors); + assert_eq!(actual_farthest, expected_farthest); } } @@ -730,6 +681,47 @@ mod tests { mod select_leaf_neighbors_tests { use super::*; + #[test] + fn invalid_neighbor_width_leaves_buffers_unchanged() { + // Given + let values = [0.0_f32, 1.0, 3.0]; + let point_count = values.len(); + let invalid_width = point_count; + let points = MatrixView::try_from(&values[..], point_count, 1).unwrap(); + let expected_output = [LeafNeighbor::default(); 9]; + // A smaller prior leaf forces scratch growth if validation runs too late. + let expected_distances = [99.0; 4]; + let expected_thresholds = [7.0; 2]; + let expected_error = LeafKernelError::InvalidNeighborCount { + points: point_count, + neighbors: invalid_width, + maximum: point_count - 1, + }; + let mut output = expected_output; + let mut workspace = LeafKernelWorkspace { + distance_scratch: expected_distances.to_vec(), + worst: expected_thresholds.to_vec(), + }; + + // When + let error = select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + points, + MutMatrixView::try_from(&mut output[..], point_count, invalid_width).unwrap(), + &mut workspace, + ) + .unwrap_err(); + + // Then + assert_eq!( + error.downcast_ref::(), + Some(&expected_error) + ); + assert_eq!(output, expected_output); + assert_eq!(workspace.distance_scratch, expected_distances); + assert_eq!(workspace.worst, expected_thresholds); + } + #[test] fn orders_neighbors_by_squared_distance_with_l2() { // Given @@ -867,8 +859,7 @@ mod tests { let expected_neighbors = reference_neighbors(&distances, point_count, requested_k); // When - let actual_neighbors = - rank_distance_fixture(&distances, point_count, requested_k).unwrap(); + let actual_neighbors = rank_distance_fixture(&distances, point_count, requested_k); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -890,7 +881,7 @@ mod tests { ]; // When - let actual_neighbors = rank_distance_fixture(&distances, 3, 1).unwrap(); + let actual_neighbors = rank_distance_fixture(&distances, 3, 1); // Then assert_eq!(actual_neighbors, expected_neighbors); @@ -908,7 +899,7 @@ mod tests { ]; // When - let actual = rank_distance_fixture(&distances, 2, 1).unwrap(); + let actual = rank_distance_fixture(&distances, 2, 1); // Then assert_eq!(actual, expected); @@ -922,7 +913,7 @@ mod tests { let distances = [1.0; 16]; // When - let actual = rank_distance_fixture(&distances, point_count, width).unwrap(); + let actual = rank_distance_fixture(&distances, point_count, width); // Then: any two candidates are valid; no tie order is required. for (source, neighbors) in actual.chunks_exact(width).enumerate() { @@ -947,7 +938,7 @@ mod tests { let expected_last_neighbor = LeafNeighbor::new(0, f32::MAX); // When - let actual = rank_distance_fixture(&distances, point_count, width).unwrap(); + let actual = rank_distance_fixture(&distances, point_count, width); // Then assert_eq!(actual[source * width + width - 1], expected_last_neighbor); @@ -965,7 +956,7 @@ mod tests { let expected = vec![LeafNeighbor::default(); point_count]; // When - let actual = rank_distance_fixture(&distances, point_count, 1).unwrap(); + let actual = rank_distance_fixture(&distances, point_count, 1); // Then assert_eq!(actual, expected); @@ -985,7 +976,7 @@ mod tests { distances[source * point_count + invalid_target] = distance; // When - let actual = rank_distance_fixture(&distances, point_count, 1).unwrap(); + let actual = rank_distance_fixture(&distances, point_count, 1); // Then assert!(actual[source].target < invalid_target as u32); @@ -1005,7 +996,7 @@ mod tests { expected[target] = LeafNeighbor::new(source as u32, -1.0); // When - let actual = rank_distance_fixture(&distances, point_count, 1).unwrap(); + let actual = rank_distance_fixture(&distances, point_count, 1); // Then assert_eq!(actual, expected); @@ -1018,30 +1009,10 @@ mod tests { let expected: [LeafNeighbor; 0] = []; // When - let actual = rank_distance_fixture(&distances, 1, 0).unwrap(); + let actual = rank_distance_fixture(&distances, 1, 0); // Then assert_eq!(actual, expected); } - - #[test] - fn neighbor_width_equal_to_point_count_is_rejected() { - // Given - let point_count = 3; - let invalid_neighbor_width = point_count; - let distances = [0.0; 9]; - let expected_error = LeafKernelError::InvalidNeighborCount { - points: point_count, - neighbors: invalid_neighbor_width, - maximum: point_count - 1, - }; - - // When - let actual_error = - rank_distance_fixture(&distances, point_count, invalid_neighbor_width).unwrap_err(); - - // Then - assert_eq!(actual_error, expected_error); - } } } From adf08976d09c2355ea2d346941b9ad5b8b21ce95 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 7 Sep 2026 07:11:13 +0000 Subject: [PATCH 100/229] refactor(pipnn): scan distance matrix views --- diskann/src/graph/pipnn/leaf_kernel.rs | 160 +++++++------------- diskann/src/graph/pipnn/partition_kernel.rs | 80 +++------- 2 files changed, 80 insertions(+), 160 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index f7978f62f2..48e4cdfb9a 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -5,7 +5,7 @@ //! Leaf-local top-k selection from packed `f32` point vectors. //! -//! The metric fills a flattened lower-triangular distance buffer. The kernel +//! The metric fills the lower triangle of a distance matrix. The kernel //! reads each strict-lower point pair once and updates both points. //! //! The output is an `n × k` matrix of sorted [`LeafNeighbor`] values. Each target @@ -114,26 +114,24 @@ where if distance_scratch.len() < distance_count { distance_scratch.resize(distance_count, 0.0); } - M::compute_distances(points, &mut distance_scratch[..distance_count])?; - rank_leaf_distances( - arch, - &distance_scratch[..distance_count], + let mut distances = MutMatrixView::try_from( + &mut distance_scratch[..distance_count], point_count, - output, - worst, - ); + point_count, + ) + .map_err(|error| ANNError::new(error.as_static()))?; + M::compute_distances(points, distances.as_mut_slice())?; + rank_leaf_distances(arch, distances.as_view(), output, worst); Ok(()) } -/// Rank one flattened lower-triangle buffer. +/// Rank the lower triangle of a square distance matrix. /// -/// `distance_flatten` contains `point_count * point_count` elements. The metric -/// initializes each strict-lower entry. The kernel does not read the upper triangle. +/// The metric initializes each strict-lower entry. The kernel does not read the upper triangle. /// The caller validates the output width against the non-self point count. fn rank_leaf_distances( arch: A, - distance_flatten: &[f32], - point_count: usize, + distances: MatrixView<'_, f32>, mut output: MutMatrixView<'_, LeafNeighbor>, worst: &mut Vec, ) where @@ -144,40 +142,15 @@ fn rank_leaf_distances( return; } - worst.resize(point_count, f32::INFINITY); + worst.resize(distances.nrows(), f32::INFINITY); output.as_mut_slice().fill(LeafNeighbor::default()); worst.fill(f32::INFINITY); match neighbor_count { - 1 => scan_fixed_width::( - arch, - distance_flatten, - point_count, - output.as_mut_slice(), - worst, - ), - 2 => scan_fixed_width::( - arch, - distance_flatten, - point_count, - output.as_mut_slice(), - worst, - ), - 3 => scan_fixed_width::( - arch, - distance_flatten, - point_count, - output.as_mut_slice(), - worst, - ), - _ => scan_runtime_width( - arch, - distance_flatten, - point_count, - output.as_mut_slice(), - neighbor_count, - worst, - ), + 1 => scan_fixed_width::(arch, distances, output.as_mut_slice(), worst), + 2 => scan_fixed_width::(arch, distances, output.as_mut_slice(), worst), + 3 => scan_fixed_width::(arch, distances, output.as_mut_slice(), worst), + _ => scan_runtime_width(arch, distances, output, worst), } } @@ -201,8 +174,7 @@ fn validate_neighbor_count( /// Select neighbors with a fixed output width. fn scan_fixed_width( arch: A, - distance_flatten: &[f32], - point_count: usize, + distances: MatrixView<'_, f32>, output: &mut [LeafNeighbor], worst: &mut [f32], ) where @@ -214,11 +186,10 @@ fn scan_fixed_width( arch.run(move || { scan_point_pairs( arch, - distance_flatten, - point_count, + distances, worst, - |source, target, distance| { - rows[source].insert_eligible(LeafNeighbor::new(target, distance)) + |source_idx, target_idx, distance| { + rows[source_idx].insert_eligible(LeafNeighbor::new(target_idx, distance)) }, ); }); @@ -227,10 +198,8 @@ fn scan_fixed_width( /// Select neighbors with a runtime output width. fn scan_runtime_width( arch: A, - distance_flatten: &[f32], - point_count: usize, - output: &mut [LeafNeighbor], - width: usize, + distances: MatrixView<'_, f32>, + mut output: MutMatrixView<'_, LeafNeighbor>, worst: &mut [f32], ) where A: PiPNNSIMDSchema, @@ -239,12 +208,12 @@ fn scan_runtime_width( arch.run(move || { scan_point_pairs( arch, - distance_flatten, - point_count, + distances, worst, - |source, target, distance| { - let first = source * width; - output[first..first + width].insert_eligible(LeafNeighbor::new(target, distance)) + |source_idx, target_idx, distance| { + output + .row_mut(source_idx) + .insert_eligible(LeafNeighbor::new(target_idx, distance)) }, ); }); @@ -255,33 +224,26 @@ fn scan_runtime_width( /// The function reads the strict lower triangle once. It offers each distance to /// both endpoint lists. #[inline(always)] -fn scan_point_pairs( - arch: A, - distance_flatten: &[f32], - point_count: usize, - worst: &mut [f32], - mut insert: I, -) where +fn scan_point_pairs(arch: A, distances: MatrixView<'_, f32>, worst: &mut [f32], mut insert: I) +where A: PiPNNSIMDSchema, I: FnMut(usize, u32, f32) -> f32, { let worst_ptr = worst.as_mut_ptr(); - for source in 1..point_count { - let source_start = source * point_count; + for source_idx in 1..distances.nrows() { + let source_distances = &distances.row(source_idx)[..source_idx]; // SAFETY: `rank_leaf_distances` created one threshold for each point. - let mut source_worst = unsafe { *worst_ptr.add(source) }; - let mut target = 0; - let simd_prefix = source - source % A::Vector::LANES; + let mut source_worst = unsafe { *worst_ptr.add(source_idx) }; + let simd_end = source_distances.len() - source_distances.len() % A::Vector::LANES; - while target < simd_prefix { + for target_idx in (0..simd_end).step_by(A::Vector::LANES) { // SAFETY: This complete SIMD group is in the strict-lower prefix. - let distance_group = unsafe { - A::Vector::load_simd(arch, distance_flatten.as_ptr().add(source_start + target)) - }; + let distance_group = + unsafe { A::Vector::load_simd(arch, source_distances.as_ptr().add(target_idx)) }; let source_eligible = distance_group.lt_simd(A::Vector::splat(arch, source_worst)); - // SAFETY: The complete target group is below `source < point_count`. - let target_worst = unsafe { A::Vector::load_simd(arch, worst_ptr.add(target)) }; + // SAFETY: The complete target group is below `source_idx < worst.len()`. + let target_worst = unsafe { A::Vector::load_simd(arch, worst_ptr.add(target_idx)) }; let target_eligible = distance_group.lt_simd(target_worst); let source_bits = A::Vector::active_lanes(source_eligible); let target_bits = A::Vector::active_lanes(target_eligible); @@ -295,7 +257,7 @@ fn scan_point_pairs( source_bits &= source_bits - 1; let distance = distance_lanes[lane]; if distance < source_worst { - source_worst = insert(source, (target + lane) as u32, distance); + source_worst = insert(source_idx, (target_idx + lane) as u32, distance); } } @@ -303,32 +265,28 @@ fn scan_point_pairs( 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 = insert(target_source, source as u32, distance_lanes[lane]); - // SAFETY: `target_source < source < worst.len()`. - unsafe { *worst_ptr.add(target_source) = new_worst }; + let candidate_idx = target_idx + lane; + let new_worst = insert(candidate_idx, source_idx as u32, distance_lanes[lane]); + // SAFETY: `candidate_idx < source_idx < worst.len()`. + unsafe { *worst_ptr.add(candidate_idx) = new_worst }; } } - target += A::Vector::LANES; } - while target < source { - // SAFETY: The target is in this source's strict-lower prefix. - let distance = unsafe { *distance_flatten.get_unchecked(source_start + target) }; + for (target_idx, &distance) in source_distances.iter().enumerate().skip(simd_end) { if distance < source_worst { - source_worst = insert(source, target as u32, distance); + source_worst = insert(source_idx, target_idx as u32, distance); } - // SAFETY: `target < source < worst.len()`. - let target_worst = unsafe { *worst_ptr.add(target) }; + // SAFETY: `target_idx < source_idx < worst.len()`. + let target_worst = unsafe { *worst_ptr.add(target_idx) }; if distance < target_worst { - let new_worst = insert(target, source as u32, distance); - // SAFETY: `target < source < worst.len()`. - unsafe { *worst_ptr.add(target) = new_worst }; + let new_worst = insert(target_idx, source_idx as u32, distance); + // SAFETY: `target_idx < source_idx < worst.len()`. + unsafe { *worst_ptr.add(target_idx) = new_worst }; } - target += 1; } - // SAFETY: `source < worst.len()`. - unsafe { *worst_ptr.add(source) = source_worst }; + // SAFETY: `source_idx < worst.len()`. + unsafe { *worst_ptr.add(source_idx) = source_worst }; } } @@ -410,8 +368,7 @@ mod tests { use diskann_wide::arch::{self, Target1}; struct KernelCall<'a> { - distance_flatten: &'a [f32], - point_count: usize, + distances: MatrixView<'a, f32>, output: MutMatrixView<'a, LeafNeighbor>, worst: &'a mut Vec, } @@ -423,13 +380,7 @@ mod tests { A: PiPNNSIMDSchema, { fn run(self, arch: A, call: KernelCall<'_>) { - rank_leaf_distances( - arch, - call.distance_flatten, - call.point_count, - call.output, - call.worst, - ); + rank_leaf_distances(arch, call.distances, call.output, call.worst); } } @@ -442,8 +393,7 @@ mod tests { arch::dispatch1_no_features( RankDistances, KernelCall { - distance_flatten: distances, - point_count: points, + distances: MatrixView::try_from(distances, points, points).unwrap(), output: MutMatrixView::try_from(output.as_mut_slice(), points, output_width) .unwrap(), worst: &mut Vec::new(), diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 42be8c54d8..287ed5d32a 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -13,7 +13,7 @@ //! leaders. Equal distances can select either leader. NaN is not rankable. An //! unfilled output slot contains [`UNASSIGNED_LEADER`]. -use crate::ANNResult; +use crate::{ANNError, ANNResult}; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{SIMDMask, SIMDVector}; @@ -62,27 +62,23 @@ where distance_scratch.resize(distance_count, 0.0); } M::compute_distances(points, leaders, &mut distance_scratch[..distance_count])?; - rank_leader_distances( - arch, + let distances = MatrixView::try_from( &distance_scratch[..distance_count], point_count, leader_count, - output, - ranked_leader_scratch, - ); + ) + .map_err(|error| ANNError::new(error.as_static()))?; + rank_leader_distances(arch, distances, output, ranked_leader_scratch); Ok(()) } /// Rank final point-to-leader distances. /// -/// `distance_flatten` contains `point_count * leader_count` elements. Each row -/// stores all leader distances for one point. The caller supplies at least one -/// point and one output column. +/// Each distance row stores all leader distances for one point. The caller +/// supplies at least one point, one leader, and one output column. fn rank_leader_distances( arch: A, - distance_flatten: &[f32], - point_count: usize, - leader_count: usize, + distances: MatrixView<'_, f32>, output: MutMatrixView<'_, u32>, ranked_leaders: &mut Vec<(u32, f32)>, ) where @@ -92,14 +88,7 @@ fn rank_leader_distances( ranked_leaders.resize(fanout, (UNASSIGNED_LEADER, f32::INFINITY)); // Rayon outlines stripe workers. Reapply target features before ranking leaders. arch.run(move || { - select_point_leaders( - arch, - distance_flatten, - point_count, - leader_count, - output, - ranked_leaders, - ); + select_point_leaders(arch, distances, output, ranked_leaders); }); } @@ -107,35 +96,27 @@ fn rank_leader_distances( #[inline(always)] fn select_point_leaders( arch: A, - distance_flatten: &[f32], - point_count: usize, - leader_count: usize, + distances: MatrixView<'_, f32>, mut output: MutMatrixView<'_, u32>, ranked_leaders: &mut [(u32, f32)], ) where A: PiPNNSIMDSchema, { - let fanout = output.ncols(); + let leader_count = distances.ncols(); + let simd_end = leader_count - leader_count % A::Vector::LANES; - for (point, point_output) in output - .as_mut_slice() - .chunks_exact_mut(fanout) - .take(point_count) - .enumerate() - { - let row_start = point * leader_count; - let point_distances = &distance_flatten[row_start..row_start + leader_count]; + for (point_distances, point_output) in distances.row_iter().zip(output.row_iter_mut()) { ranked_leaders.fill((UNASSIGNED_LEADER, f32::INFINITY)); - let simd_prefix = leader_count - leader_count % A::Vector::LANES; - for first_leader in (0..simd_prefix).step_by(A::Vector::LANES) { + for leader_base_idx in (0..simd_end).step_by(A::Vector::LANES) { // SAFETY: This group is inside the point's leader row. - let distance_group = - unsafe { A::Vector::load_simd(arch, point_distances.as_ptr().add(first_leader)) }; - insert_leader_lanes(distance_group, first_leader, ranked_leaders); + let distance_group = unsafe { + A::Vector::load_simd(arch, point_distances.as_ptr().add(leader_base_idx)) + }; + insert_leader_lanes(distance_group, leader_base_idx, ranked_leaders); } - for (leader, &distance) in point_distances.iter().enumerate().skip(simd_prefix) { + for (leader, &distance) in point_distances.iter().enumerate().skip(simd_end) { insert_leader(ranked_leaders, leader as u32, distance); } for (destination, &(leader, _)) in point_output.iter_mut().zip(ranked_leaders.iter()) { @@ -146,9 +127,9 @@ fn select_point_leaders( /// 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. The function reads +/// `leader_base_idx` is the matrix-column index of the first lane. The function reads /// selected lanes from low to high. -fn insert_leader_lanes(distances: F, first_leader: usize, ranked_leaders: &mut [(u32, f32)]) +fn insert_leader_lanes(distances: F, leader_base_idx: usize, ranked_leaders: &mut [(u32, f32)]) where F: PiPNNSIMDVector, { @@ -166,7 +147,7 @@ where lanes &= lanes - 1; insert_leader( ranked_leaders, - (first_leader + lane) as u32, + (leader_base_idx + lane) as u32, distance_lanes[lane], ); } @@ -203,9 +184,7 @@ mod tests { use diskann_wide::arch::{self, Target1}; struct KernelCall<'a> { - distance_flatten: &'a [f32], - point_count: usize, - leader_count: usize, + distances: MatrixView<'a, f32>, output: MutMatrixView<'a, u32>, ranked_leaders: &'a mut Vec<(u32, f32)>, } @@ -217,14 +196,7 @@ mod tests { A: PiPNNSIMDSchema, { fn run(self, arch: A, call: KernelCall<'_>) { - rank_leader_distances( - arch, - call.distance_flatten, - call.point_count, - call.leader_count, - call.output, - call.ranked_leaders, - ); + rank_leader_distances(arch, call.distances, call.output, call.ranked_leaders); } } @@ -237,9 +209,7 @@ mod tests { arch::dispatch1_no_features( RankDistances, KernelCall { - distance_flatten: distances.as_slice(), - point_count: distances.nrows(), - leader_count: distances.ncols(), + distances, output: output.as_mut_view(), ranked_leaders: &mut Vec::new(), }, From b308086662d9e52e417b38d68d939134396d6476 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 7 Sep 2026 12:25:10 +0000 Subject: [PATCH 101/229] refactor(pipnn): share top-k state and trim tests --- diskann/src/graph/pipnn/leaf_kernel.rs | 622 +++----------------- diskann/src/graph/pipnn/leaf_metric.rs | 30 +- diskann/src/graph/pipnn/mod.rs | 137 ++--- diskann/src/graph/pipnn/partition_kernel.rs | 232 ++------ diskann/src/graph/pipnn/partition_metric.rs | 90 +-- diskann/src/graph/pipnn/simd.rs | 54 +- diskann/src/graph/pipnn/topk.rs | 483 +++++++++++++++ 7 files changed, 735 insertions(+), 913 deletions(-) create mode 100644 diskann/src/graph/pipnn/topk.rs diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 48e4cdfb9a..662a37cfdc 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -8,12 +8,11 @@ //! The metric fills the lower triangle of a distance matrix. The kernel //! 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. Widths 1 through 3 use fixed insertion. Larger -//! widths use the runtime insertion loop. +//! The output is an `n × k` matrix of sorted [`Candidate`] values. Each target +//! is a position in the leaf. //! //! Equal distances can select either candidate. The kernel does not rank NaN. -//! An unfilled output slot contains [`LeafNeighbor::default`]. All supported +//! An unfilled output slot contains [`Candidate::default`]. All supported //! metrics use the same SIMD-group and single-value traversal. //! //! The caller supplies concrete architecture `A` and metric `M`. @@ -21,45 +20,13 @@ use crate::{ANNError, ANNResult}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_wide::{SIMDPartialOrd, SIMDVector}; use super::{ leaf_metric::LeafMetric, - simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, + simd::{PiPNNSIMDSchema, distance_blocks}, + topk::{Candidate, with_topk_rows}, }; -/// One leaf-local neighbor and its ranking distance. -/// -/// A ranking distance preserves nearest-first order. It need not equal the -/// metric distance. -#[derive(Clone, Copy, Debug, PartialEq)] -pub(super) struct LeafNeighbor { - /// Target position in the leaf, not a dataset ID. - pub(super) target: u32, - /// Ranking distance from the source point to `target`. - pub(super) distance: f32, -} - -impl LeafNeighbor { - /// Construct a leaf-local neighbor. - /// - /// The output row determines the source point. - 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 { - fn default() -> Self { - Self::new(u32::MAX, f32::INFINITY) - } -} - /// Reusable storage for one leaf numerical pipeline. #[derive(Debug, Default)] pub(super) struct LeafKernelWorkspace { @@ -97,7 +64,7 @@ pub(super) fn leaf_neighbor_count(points: usize, requested_k: usize) -> usize { pub(super) fn select_leaf_neighbors( arch: A, points: MatrixView<'_, f32>, - output: MutMatrixView<'_, LeafNeighbor>, + output: MutMatrixView<'_, Candidate>, workspace: &mut LeafKernelWorkspace, ) -> ANNResult<()> where @@ -125,39 +92,37 @@ where Ok(()) } -/// Rank the lower triangle of a square distance matrix. -/// -/// The metric initializes each strict-lower entry. The kernel does not read the upper triangle. -/// The caller validates the output width against the non-self point count. +/// Traverse the strict lower triangle and update each pair's two candidate rows. fn rank_leaf_distances( arch: A, distances: MatrixView<'_, f32>, - mut output: MutMatrixView<'_, LeafNeighbor>, + output: MutMatrixView<'_, Candidate>, worst: &mut Vec, ) where A: PiPNNSIMDSchema, { - let neighbor_count = output.ncols(); - if neighbor_count == 0 { + if output.ncols() == 0 { return; } - worst.resize(distances.nrows(), f32::INFINITY); - output.as_mut_slice().fill(LeafNeighbor::default()); - worst.fill(f32::INFINITY); - - match neighbor_count { - 1 => scan_fixed_width::(arch, distances, output.as_mut_slice(), worst), - 2 => scan_fixed_width::(arch, distances, output.as_mut_slice(), worst), - 3 => scan_fixed_width::(arch, distances, output.as_mut_slice(), worst), - _ => scan_runtime_width(arch, distances, output, worst), - } + with_topk_rows!(output, worst, |topks| { + // Rayon outlines leaf workers. Reapply target features before the SIMD scan. + arch.run(move || { + for source_idx in 1..distances.nrows() { + let source_distances = &distances.row(source_idx)[..source_idx]; + for block in distance_blocks(arch, source_distances) { + topks.update_one(source_idx, &block); + topks.update_many(source_idx as u32, &block); + } + } + }); + }); } /// Check the output width against the number of non-self points. fn validate_neighbor_count( point_count: usize, - output: &MutMatrixView<'_, LeafNeighbor>, + output: &MutMatrixView<'_, Candidate>, ) -> Result<(), LeafKernelError> { let maximum_neighbors = point_count.saturating_sub(1); let neighbor_count = output.ncols(); @@ -171,190 +136,6 @@ fn validate_neighbor_count( Ok(()) } -/// Select neighbors with a fixed output width. -fn scan_fixed_width( - arch: A, - distances: MatrixView<'_, f32>, - output: &mut [LeafNeighbor], - worst: &mut [f32], -) where - A: PiPNNSIMDSchema, - [LeafNeighbor; N]: NeighborInsert, -{ - let (rows, _) = output.as_chunks_mut::(); - // Rayon outlines leaf workers. Reapply target features before the SIMD scan. - arch.run(move || { - scan_point_pairs( - arch, - distances, - worst, - |source_idx, target_idx, distance| { - rows[source_idx].insert_eligible(LeafNeighbor::new(target_idx, distance)) - }, - ); - }); -} - -/// Select neighbors with a runtime output width. -fn scan_runtime_width( - arch: A, - distances: MatrixView<'_, f32>, - mut output: MutMatrixView<'_, LeafNeighbor>, - worst: &mut [f32], -) where - A: PiPNNSIMDSchema, -{ - // Rayon outlines leaf workers. Reapply target features before the SIMD scan. - arch.run(move || { - scan_point_pairs( - arch, - distances, - worst, - |source_idx, target_idx, distance| { - output - .row_mut(source_idx) - .insert_eligible(LeafNeighbor::new(target_idx, 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. -#[inline(always)] -fn scan_point_pairs(arch: A, distances: MatrixView<'_, f32>, worst: &mut [f32], mut insert: I) -where - A: PiPNNSIMDSchema, - I: FnMut(usize, u32, f32) -> f32, -{ - let worst_ptr = worst.as_mut_ptr(); - - for source_idx in 1..distances.nrows() { - let source_distances = &distances.row(source_idx)[..source_idx]; - // SAFETY: `rank_leaf_distances` created one threshold for each point. - let mut source_worst = unsafe { *worst_ptr.add(source_idx) }; - let simd_end = source_distances.len() - source_distances.len() % A::Vector::LANES; - - for target_idx in (0..simd_end).step_by(A::Vector::LANES) { - // SAFETY: This complete SIMD group is in the strict-lower prefix. - let distance_group = - unsafe { A::Vector::load_simd(arch, source_distances.as_ptr().add(target_idx)) }; - let source_eligible = distance_group.lt_simd(A::Vector::splat(arch, source_worst)); - // SAFETY: The complete target group is below `source_idx < worst.len()`. - let target_worst = unsafe { A::Vector::load_simd(arch, worst_ptr.add(target_idx)) }; - let target_eligible = distance_group.lt_simd(target_worst); - let source_bits = A::Vector::active_lanes(source_eligible); - let target_bits = A::Vector::active_lanes(target_eligible); - - if source_bits | target_bits != 0 { - let distance_lanes = distance_group.to_lane_array(); - let distance_lanes = distance_lanes.as_ref(); - 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 = distance_lanes[lane]; - if distance < source_worst { - source_worst = insert(source_idx, (target_idx + lane) as u32, distance); - } - } - - let mut target_bits = target_bits; - while target_bits != 0 { - let lane = target_bits.trailing_zeros() as usize; - target_bits &= target_bits - 1; - let candidate_idx = target_idx + lane; - let new_worst = insert(candidate_idx, source_idx as u32, distance_lanes[lane]); - // SAFETY: `candidate_idx < source_idx < worst.len()`. - unsafe { *worst_ptr.add(candidate_idx) = new_worst }; - } - } - } - - for (target_idx, &distance) in source_distances.iter().enumerate().skip(simd_end) { - if distance < source_worst { - source_worst = insert(source_idx, target_idx as u32, distance); - } - // SAFETY: `target_idx < source_idx < worst.len()`. - let target_worst = unsafe { *worst_ptr.add(target_idx) }; - if distance < target_worst { - let new_worst = insert(target_idx, source_idx as u32, distance); - // SAFETY: `target_idx < source_idx < worst.len()`. - unsafe { *worst_ptr.add(target_idx) = new_worst }; - } - } - // SAFETY: `source_idx < worst.len()`. - unsafe { *worst_ptr.add(source_idx) = source_worst }; - } -} - -/// Insert a neighbor that is closer than the current farthest neighbor. -/// -/// The pair scan checks eligibility before insertion. The retained neighbors -/// stay in distance order. The result is their new farthest ranking distance. -trait NeighborInsert { - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32; -} - -impl NeighborInsert for [LeafNeighbor; 1] { - #[inline(always)] - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32 { - self[0] = candidate; - candidate.distance - } -} - -impl NeighborInsert for [LeafNeighbor; 2] { - #[inline(always)] - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32 { - let first = self[0]; - if candidate.distance < first.distance { - self[0] = candidate; - self[1] = first; - first.distance - } else { - self[1] = candidate; - candidate.distance - } - } -} - -impl NeighborInsert for [LeafNeighbor; 3] { - #[inline(always)] - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32 { - let (first, second) = (self[0], self[1]); - if candidate.distance < first.distance { - self[0] = candidate; - self[1] = first; - self[2] = second; - second.distance - } else if candidate.distance < second.distance { - self[1] = candidate; - self[2] = second; - second.distance - } else { - self[2] = candidate; - candidate.distance - } - } -} - -impl NeighborInsert for [LeafNeighbor] { - #[inline(always)] - fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32 { - let last = self.len() - 1; - let mut slot = last; - while slot > 0 && candidate.distance < self[slot - 1].distance { - self[slot] = self[slot - 1]; - slot -= 1; - } - self[slot] = candidate; - self[last].distance - } -} - #[cfg(test)] mod tests { use super::*; @@ -369,7 +150,7 @@ mod tests { struct KernelCall<'a> { distances: MatrixView<'a, f32>, - output: MutMatrixView<'a, LeafNeighbor>, + output: MutMatrixView<'a, Candidate>, worst: &'a mut Vec, } @@ -388,8 +169,8 @@ mod tests { distances: &[f32], points: usize, output_width: usize, - ) -> Vec { - let mut output = vec![LeafNeighbor::default(); points * output_width]; + ) -> Vec { + let mut output = vec![Candidate::default(); points * output_width]; arch::dispatch1_no_features( RankDistances, KernelCall { @@ -406,8 +187,8 @@ mod tests { distances: &[f32], points: usize, width: usize, - ) -> Vec { - let mut output = vec![LeafNeighbor::default(); points * width]; + ) -> Vec { + let mut output = vec![Candidate::default(); points * width]; for source in 0..points { let mut candidates = Vec::with_capacity(points.saturating_sub(1)); for target in 0..points { @@ -418,7 +199,7 @@ mod tests { let column = source.min(target); let distance = distances[row * points + column]; if distance < f32::INFINITY { - candidates.push(LeafNeighbor::new(target as u32, distance)); + candidates.push(Candidate::new(target as u32, distance)); } } candidates.sort_unstable_by(|left, right| { @@ -448,183 +229,20 @@ mod tests { } } - mod insert_eligible_tests { - use super::*; - - #[test] - fn nearer_candidate_replaces_the_only_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 expected_farthest = nearer_candidate.distance; - let mut actual_neighbors = [retained_neighbor]; - - // When - let actual_farthest = actual_neighbors.insert_eligible(nearer_candidate); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - assert_eq!(actual_farthest, expected_farthest); - } - - #[test] - fn nearer_candidate_moves_to_the_front_of_two_retained_neighbors() { - // 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 expected_farthest = nearest.distance; - let mut actual_neighbors = [nearest, farthest]; - - // When - let actual_farthest = actual_neighbors.insert_eligible(nearer_candidate); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - assert_eq!(actual_farthest, expected_farthest); - } - - #[test] - fn middle_distance_candidate_replaces_the_farther_of_two_neighbors() { - // 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 expected_farthest = eligible_candidate.distance; - let mut actual_neighbors = [nearest, farthest]; - - // When - let actual_farthest = actual_neighbors.insert_eligible(eligible_candidate); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - assert_eq!(actual_farthest, expected_farthest); - } - - #[test] - fn nearest_candidate_moves_to_the_front_of_three_retained_neighbors() { - // 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 expected_farthest = middle.distance; - let mut actual_neighbors = [nearest, middle, farthest]; - - // When - let actual_farthest = actual_neighbors.insert_eligible(nearer_candidate); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - assert_eq!(actual_farthest, expected_farthest); - } - - #[test] - fn middle_candidate_is_inserted_between_three_retained_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 expected_farthest = middle.distance; - let mut actual_neighbors = [nearest, middle, farthest]; - - // When - let actual_farthest = actual_neighbors.insert_eligible(middle_candidate); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - assert_eq!(actual_farthest, expected_farthest); - } - - #[test] - fn closer_candidate_replaces_the_farthest_of_three_neighbors() { - // 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 expected_farthest = eligible_candidate.distance; - let mut actual_neighbors = [nearest, middle, farthest]; - - // When - let actual_farthest = actual_neighbors.insert_eligible(eligible_candidate); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - assert_eq!(actual_farthest, expected_farthest); - } - - #[test] - fn middle_candidate_shifts_only_farther_runtime_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 expected_farthest = third.distance; - let mut actual_neighbors = [first, second, third, fourth]; - - // When - let actual_farthest = actual_neighbors.as_mut_slice().insert_eligible(candidate); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - assert_eq!(actual_farthest, expected_farthest); - } - } - mod leaf_neighbor_count_tests { use super::leaf_neighbor_count; + use rstest::rstest; - #[test] - fn returns_zero_when_the_leaf_contains_only_the_source() { - // 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 returns_the_non_self_point_count_when_requested_k_is_larger() { - // 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 returns_requested_k_when_enough_non_self_points_exist() { - // 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); + #[rstest] + #[case::singleton(1, 3, 0)] + #[case::clamped(4, 4, 3)] + #[case::requested(8, 5, 5)] + fn count_is_bounded_by_non_self_points( + #[case] points: usize, + #[case] requested: usize, + #[case] expected: usize, + ) { + assert_eq!(leaf_neighbor_count(points, requested), expected); } } @@ -638,7 +256,7 @@ mod tests { let point_count = values.len(); let invalid_width = point_count; let points = MatrixView::try_from(&values[..], point_count, 1).unwrap(); - let expected_output = [LeafNeighbor::default(); 9]; + let expected_output = [Candidate::default(); 9]; // A smaller prior leaf forces scratch growth if validation runs too late. let expected_distances = [99.0; 4]; let expected_thresholds = [7.0; 2]; @@ -673,103 +291,52 @@ mod tests { } #[test] - fn orders_neighbors_by_squared_distance_with_l2() { + fn l2_selection_reuses_output_and_scratch_for_a_smaller_leaf() { // 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)), + Candidate::new(1, (values[0] - values[1]).powi(2)), + Candidate::new(2, (values[0] - values[2]).powi(2)), + Candidate::new(0, (values[1] - values[0]).powi(2)), + Candidate::new(2, (values[1] - values[2]).powi(2)), + Candidate::new(1, (values[2] - values[1]).powi(2)), + Candidate::new(0, (values[2] - values[0]).powi(2)), + Candidate::new(2, (values[3] - values[2]).powi(2)), + Candidate::new(1, (values[3] - values[1]).powi(2)), ]; - let mut actual_neighbors = [LeafNeighbor::default(); 8]; + let mut actual_neighbors = [Candidate::default(); 8]; + let mut workspace = LeafKernelWorkspace::default(); // When select_leaf_neighbors::<_, L2>( diskann_wide::ARCH, points, MutMatrixView::try_from(&mut actual_neighbors[..], 4, 2).unwrap(), - &mut LeafKernelWorkspace::default(), + &mut workspace, ) .unwrap(); - - // Then assert_eq!(actual_neighbors, expected_neighbors); - } - - #[test] - fn later_farther_candidate_cannot_replace_the_retained_neighbor() { - // Given - let point_values = [0.0_f32, 10.0, 1.0]; - let source = 2; - let first_scanned_target = 0_u32; - let later_farther_target = 1_usize; - let expected_distance = - (point_values[source] - point_values[first_scanned_target as usize]).powi(2); - let later_distance = - (point_values[source] - point_values[later_farther_target]).powi(2); - let expected_nearest_neighbor = - LeafNeighbor::new(first_scanned_target, expected_distance); - assert!(later_distance > expected_distance); - let mut actual_neighbors = [LeafNeighbor::default(); 3]; - - // When - select_leaf_neighbors::<_, L2>( - diskann_wide::ARCH, - MatrixView::try_from(&point_values[..], 3, 1).unwrap(), - MutMatrixView::try_from(&mut actual_neighbors[..], 3, 1).unwrap(), - &mut LeafKernelWorkspace::default(), - ) - .unwrap(); - - // Then - assert_eq!(actual_neighbors[source], expected_nearest_neighbor); - } - - #[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]; + // Reuse both buffers with a smaller, different dataset and unfilled slots. + let smaller = [0.0, 2.0, f32::NAN]; select_leaf_neighbors::<_, L2>( diskann_wide::ARCH, - smaller_points, - MutMatrixView::try_from(&mut actual_neighbors_from_reused_workspace[..], 3, 2) - .unwrap(), - &mut reused_workspace, + MatrixView::try_from(&smaller[..], 3, 1).unwrap(), + MutMatrixView::try_from(&mut actual_neighbors[..6], 3, 2).unwrap(), + &mut workspace, ) .unwrap(); - - // Then assert_eq!( - actual_neighbors_from_reused_workspace, - expected_neighbors_from_fresh_workspace + actual_neighbors[..6], + [ + Candidate::new(1, 4.0), + Candidate::default(), + Candidate::new(0, 4.0), + Candidate::default(), + Candidate::default(), + Candidate::default(), + ] ); } } @@ -791,8 +358,6 @@ mod tests { #[case::two_lanes_minus_one_runtime_width(32, 7)] #[case::two_complete_lanes_runtime_width(33, 7)] #[case::two_lanes_plus_one_runtime_width(34, 7)] - #[case::four_complete_lanes_runtime_width(65, 7)] - #[case::sixteen_complete_lanes_runtime_width(257, 7)] #[case::maximum_leaf_size_runtime_width(512, 7)] #[trace] fn dispatched_leaf_ranking_matches_scalar_reference_across_lane_boundaries( @@ -815,38 +380,13 @@ mod tests { assert_eq!(actual_neighbors, expected_neighbors); } - #[test] - fn only_the_lower_triangle_supplies_neighbor_distances() { - // Given - #[rustfmt::skip] - let distances = [ - f32::NAN, f32::NAN, f32::NAN, - 3.0, f32::NAN, f32::NAN, - 1.0, 2.0, f32::NAN, - ]; - let expected_neighbors = [ - LeafNeighbor::new(2, 1.0), - LeafNeighbor::new(2, 2.0), - LeafNeighbor::new(0, 1.0), - ]; - - // When - let actual_neighbors = rank_distance_fixture(&distances, 3, 1); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - } - #[rstest] #[case::finite_negative(-f32::EPSILON)] #[case::negative_infinity(f32::NEG_INFINITY)] fn negative_distances_remain_rankable(#[case] distance: f32) { // Given let distances = [f32::NAN, f32::NAN, distance, f32::NAN]; - let expected = [ - LeafNeighbor::new(1, distance), - LeafNeighbor::new(0, distance), - ]; + let expected = [Candidate::new(1, distance), Candidate::new(0, distance)]; // When let actual = rank_distance_fixture(&distances, 2, 1); @@ -869,10 +409,10 @@ mod tests { for (source, neighbors) in actual.chunks_exact(width).enumerate() { assert_eq!(neighbors[0].distance, 1.0); assert_eq!(neighbors[1].distance, 1.0); - assert_ne!(neighbors[0].target, neighbors[1].target); + assert_ne!(neighbors[0].local_idx, neighbors[1].local_idx); for neighbor in neighbors { - assert!(neighbor.target < point_count as u32); - assert_ne!(neighbor.target, source as u32); + assert!(neighbor.local_idx < point_count as u32); + assert_ne!(neighbor.local_idx, source as u32); } } } @@ -885,7 +425,7 @@ mod tests { let source = 3; let mut distances = [1.0; 16]; distances[source * point_count] = f32::MAX; - let expected_last_neighbor = LeafNeighbor::new(0, f32::MAX); + let expected_last_neighbor = Candidate::new(0, f32::MAX); // When let actual = rank_distance_fixture(&distances, point_count, width); @@ -897,13 +437,11 @@ mod tests { #[rstest] #[case::nan(f32::NAN)] #[case::positive_infinity(f32::INFINITY)] - fn non_rankable_distances_leave_neighbor_slots_unassigned( - #[case] distance: f32, - #[values(2, 17, 18)] point_count: usize, - ) { - // Given + fn non_rankable_distances_leave_neighbor_slots_unassigned(#[case] distance: f32) { + // One leaf exercises scalar rows, a full SIMD group, and its tail. + let point_count = 18; let distances = vec![distance; point_count * point_count]; - let expected = vec![LeafNeighbor::default(); point_count]; + let expected = vec![Candidate::default(); point_count]; // When let actual = rank_distance_fixture(&distances, point_count, 1); @@ -929,7 +467,7 @@ mod tests { let actual = rank_distance_fixture(&distances, point_count, 1); // Then - assert!(actual[source].target < invalid_target as u32); + assert!(actual[source].local_idx < invalid_target as u32); assert_eq!(actual[source].distance, 1.0); } @@ -941,9 +479,9 @@ mod tests { let target = 16; let mut distances = vec![f32::INFINITY; point_count * point_count]; distances[source * point_count + target] = -1.0; - let mut expected = vec![LeafNeighbor::default(); point_count]; - expected[source] = LeafNeighbor::new(target as u32, -1.0); - expected[target] = LeafNeighbor::new(source as u32, -1.0); + let mut expected = vec![Candidate::default(); point_count]; + expected[source] = Candidate::new(target as u32, -1.0); + expected[target] = Candidate::new(source as u32, -1.0); // When let actual = rank_distance_fixture(&distances, point_count, 1); @@ -956,7 +494,7 @@ mod tests { fn singleton_leaf_has_no_neighbors() { // Given let distances = [f32::NAN]; - let expected: [LeafNeighbor; 0] = []; + let expected: [Candidate; 0] = []; // When let actual = rank_distance_fixture(&distances, 1, 0); diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index c441b94109..9a3d73bcef 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -239,24 +239,32 @@ mod tests { } #[rstest] - #[case::inner_product(InnerProduct::compute_distances)] - #[case::normalized_cosine(CosineNormalized::compute_distances)] + #[case::inner_product(compute_pair_ranking::)] + #[case::normalized_cosine(compute_pair_ranking::)] fn orthogonal_vectors_have_zero_ranking( - #[case] compute: fn(MatrixView<'_, f32>, &mut [f32]) -> ANNResult<()>, - #[values(2, 17, 129)] dimensions: usize, + #[case] compute: fn([f32; DIMENSION_COUNT], [f32; DIMENSION_COUNT]) -> f32, ) { - // Given: the two unit vectors have disjoint nonzero coordinates. + // The contract permits either zero sign at both metric entry points. + assert_eq!(compute([1.0, 0.0], [0.0, 1.0]), 0.0); + } + + #[test] + fn inner_product_uses_the_last_coordinate_beyond_a_complete_dimension_block() { + let dimensions = 129; let mut values = vec![0.0; 2 * dimensions]; - values[0] = 1.0; - values[dimensions + 1] = 1.0; + values[dimensions - 1] = 2.0; + values[dimensions] = -1.0; + values[2 * dimensions - 1] = 3.0; let points = MatrixView::try_from(values.as_slice(), 2, dimensions).unwrap(); let mut distances = [STALE_DISTANCE; 4]; - // When - compute(points, &mut distances).unwrap(); + InnerProduct::compute_distances(points, &mut distances).unwrap(); - // Then: the contract permits either zero sign. - assert_eq!(distances[2], 0.0); + // Dot products are 4, 6 and 10; all nonzero pair contribution is in the tail. + assert_eq!( + [distances[0], distances[2], distances[3]], + [-4.0, -6.0, -10.0] + ); } } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index de68241ada..e91ed5fff8 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -18,6 +18,8 @@ mod partition_kernel; mod partition_metric; #[allow(dead_code)] mod simd; +#[allow(dead_code)] +mod topk; pub(super) struct L2; pub(super) struct Cosine; @@ -39,104 +41,57 @@ fn cosine_distance(dot: f32, source_norm: f32, target_norm: f32) -> f32 { } #[cfg(test)] -mod cosine_distance_contract_tests { +mod cosine_distance_tests { use super::cosine_distance; use rstest::rstest; - mod cosine_distance_tests { - use super::*; - - #[test] - fn zero_norm_takes_precedence_over_a_nan_dot_product() { - // Given - let dot = f32::NAN; - let source_norm = 0.0_f32; - let target_norm = 1.0_f32; - let zero_similarity = 0.0_f32; - let expected = 1.0 - zero_similarity; - - // When - let actual = cosine_distance(dot, source_norm, target_norm); - - // Then - assert_eq!(actual, expected); - } - - #[rstest] - #[case::zero_source(0.0, 1.0)] - #[case::zero_target(1.0, 0.0)] - #[case::subnormal_source(f32::MIN_POSITIVE.sqrt() / 2.0, 1.0)] - #[case::subnormal_target(1.0, f32::MIN_POSITIVE.sqrt() / 2.0)] - #[trace] - fn zero_or_subnormal_norm_produces_unit_distance( - #[case] source_norm: f32, - #[case] target_norm: f32, - ) { - // Given - let dot = 0.0_f32; - let zero_similarity = 0.0_f32; - let expected = 1.0 - zero_similarity; - - // When - let actual = cosine_distance(dot, source_norm, target_norm); - - // Then - assert_eq!(actual, expected); - } - - #[test] - fn minimum_normal_norm_uses_normalized_similarity() { - // Given - let source_norm = f32::MIN_POSITIVE.sqrt(); - let target_norm = 1.0_f32; - let expected_similarity = 0.5_f32; - let dot = expected_similarity * source_norm * target_norm; - let expected = 1.0 - expected_similarity; - - // When - let actual = cosine_distance(dot, source_norm, target_norm); - - // Then - assert_eq!(actual, expected); - } - - #[rstest] - #[case::above_one(1.0)] - #[case::below_negative_one(-1.0)] - #[trace] - fn finite_similarity_outside_the_cosine_range_is_clamped(#[case] bounded_similarity: f32) { - // Given - let source_norm = 2.0_f32; - let target_norm = 2.0_f32; - let norm_product = source_norm * target_norm; - let rounding_excess = f32::EPSILON * norm_product; - let dot = bounded_similarity * (norm_product + rounding_excess); - let expected = 1.0 - bounded_similarity; + #[rstest] + #[case::zero_source(0.0, 0.0, 1.0)] + #[case::zero_target(0.0, 1.0, 0.0)] + #[case::subnormal_source(f32::MIN_POSITIVE.sqrt() / 2.0, f32::MIN_POSITIVE.sqrt() / 2.0, 1.0)] + #[case::subnormal_target(f32::MIN_POSITIVE.sqrt() / 2.0, 1.0, f32::MIN_POSITIVE.sqrt() / 2.0)] + #[case::zero_norm_before_nan_dot(f32::NAN, 0.0, 1.0)] + fn small_norm_produces_unit_distance( + #[case] dot: f32, + #[case] source_norm: f32, + #[case] target_norm: f32, + ) { + assert_eq!(cosine_distance(dot, source_norm, target_norm), 1.0); + } - // When - let actual = cosine_distance(dot, source_norm, target_norm); + #[test] + fn minimum_normal_norm_uses_normalized_similarity() { + let source_norm = f32::MIN_POSITIVE.sqrt(); + let expected_similarity = 0.5; + let dot = expected_similarity * source_norm; - // Then - assert_eq!(actual, expected); - } + assert_eq!( + cosine_distance(dot, source_norm, 1.0), + 1.0 - expected_similarity + ); + } - #[rstest] - #[case::nan_dot(f32::NAN, 1.0, 1.0)] - #[case::nan_source_norm(0.0, f32::NAN, 1.0)] - #[case::nan_target_norm(0.0, 1.0, f32::NAN)] - #[trace] - fn nan_without_a_zero_norm_produces_nan_distance( - #[case] dot: f32, - #[case] source_norm: f32, - #[case] target_norm: f32, - ) { - // Given: the case supplies one NaN. The other values do not select the zero-norm rule. + #[rstest] + #[case::above_one(1.0)] + #[case::below_negative_one(-1.0)] + fn finite_similarity_outside_the_cosine_range_is_clamped(#[case] bounded_similarity: f32) { + let norm = 2.0; + let norm_product = norm * norm; + let rounding_excess = f32::EPSILON * norm_product; + let dot = bounded_similarity * (norm_product + rounding_excess); - // When - let actual = cosine_distance(dot, source_norm, target_norm); + assert_eq!(cosine_distance(dot, norm, norm), 1.0 - bounded_similarity); + } - // Then - assert!(actual.is_nan()); - } + #[rstest] + #[case::nan_dot(f32::NAN, 1.0, 1.0)] + #[case::nan_source_norm(0.0, f32::NAN, 1.0)] + #[case::nan_target_norm(0.0, 1.0, f32::NAN)] + fn nan_without_a_zero_norm_produces_nan_distance( + #[case] dot: f32, + #[case] source_norm: f32, + #[case] target_norm: f32, + ) { + assert!(cosine_distance(dot, source_norm, target_norm).is_nan()); } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 287ed5d32a..00d62cbfac 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -15,21 +15,21 @@ use crate::{ANNError, ANNResult}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_wide::{SIMDMask, SIMDVector}; use super::{ partition_metric::PartitionMetric, - simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, + simd::{PiPNNSIMDSchema, distance_blocks}, + topk::{Candidate, UNASSIGNED, with_topk_rows}, }; /// No sampled partition center was rankable for this output slot. -pub(super) const UNASSIGNED_LEADER: u32 = u32::MAX; +pub(super) const UNASSIGNED_LEADER: u32 = UNASSIGNED; /// Reusable storage for one point-stripe numerical pipeline. #[derive(Default)] pub(super) struct PartitionKernelWorkspace { distance_scratch: Vec, - ranked_leader_scratch: Vec<(u32, f32)>, + ranked_leader_scratch: Vec, } /// Assign one packed point stripe to metric-owned partition leaders. @@ -79,98 +79,34 @@ where fn rank_leader_distances( arch: A, distances: MatrixView<'_, f32>, - output: MutMatrixView<'_, u32>, - ranked_leaders: &mut Vec<(u32, f32)>, + mut output: MutMatrixView<'_, u32>, + ranked_leaders: &mut Vec, ) where A: PiPNNSIMDSchema, { let fanout = output.ncols(); - ranked_leaders.resize(fanout, (UNASSIGNED_LEADER, f32::INFINITY)); + ranked_leaders.resize(fanout, Candidate::default()); // Rayon outlines stripe workers. Reapply target features before ranking leaders. arch.run(move || { - select_point_leaders(arch, distances, output, ranked_leaders); - }); -} - -/// Rank sampled partition centers for each assigned point. -#[inline(always)] -fn select_point_leaders( - arch: A, - distances: MatrixView<'_, f32>, - mut output: MutMatrixView<'_, u32>, - ranked_leaders: &mut [(u32, f32)], -) where - A: PiPNNSIMDSchema, -{ - let leader_count = distances.ncols(); - let simd_end = leader_count - leader_count % A::Vector::LANES; - - for (point_distances, point_output) in distances.row_iter().zip(output.row_iter_mut()) { - ranked_leaders.fill((UNASSIGNED_LEADER, f32::INFINITY)); - - for leader_base_idx in (0..simd_end).step_by(A::Vector::LANES) { - // SAFETY: This group is inside the point's leader row. - let distance_group = unsafe { - A::Vector::load_simd(arch, point_distances.as_ptr().add(leader_base_idx)) - }; - insert_leader_lanes(distance_group, leader_base_idx, ranked_leaders); - } - - for (leader, &distance) in point_distances.iter().enumerate().skip(simd_end) { - insert_leader(ranked_leaders, leader as u32, distance); - } - for (destination, &(leader, _)) in point_output.iter_mut().zip(ranked_leaders.iter()) { - *destination = leader; - } - } -} - -/// Offer one SIMD group of sampled centers to the current point's ranked_leaders. -/// -/// `leader_base_idx` is the matrix-column index of the first lane. The function reads -/// selected lanes from low to high. -fn insert_leader_lanes(distances: F, leader_base_idx: usize, ranked_leaders: &mut [(u32, f32)]) -where - F: PiPNNSIMDVector, -{ - let threshold = F::splat(distances.arch(), ranked_leaders[ranked_leaders.len() - 1].1); - let eligible = distances.lt_simd(threshold); - if eligible.none() { - return; - } - - let distance_lanes = distances.to_lane_array(); - let distance_lanes = distance_lanes.as_ref(); - let mut lanes = F::active_lanes(eligible); - while lanes != 0 { - let lane = lanes.trailing_zeros() as usize; - lanes &= lanes - 1; - insert_leader( - ranked_leaders, - (leader_base_idx + lane) as u32, - distance_lanes[lane], + let mut worst = [f32::INFINITY]; + with_topk_rows!( + MutMatrixView::row_vector(ranked_leaders.as_mut_slice()), + &mut worst[..], + |topks| { + for (point_distances, point_output) in + distances.row_iter().zip(output.row_iter_mut()) + { + topks.reset(); + for block in distance_blocks(arch, point_distances) { + topks.update_one(0, &block); + } + for (destination, leader) in point_output.iter_mut().zip(topks.row(0)) { + *destination = leader.local_idx; + } + } + } ); - } -} - -/// 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. `ranked_leaders` -/// stores retained centers in nearest-first order. A candidate enters only when -/// its distance is less than the current farthest distance. NaN does not enter. -#[inline(always)] -fn insert_leader(ranked_leaders: &mut [(u32, f32)], leader: u32, distance: f32) { - let threshold = ranked_leaders.len() - 1; - if distance.partial_cmp(&ranked_leaders[threshold].1) != Some(std::cmp::Ordering::Less) { - return; - } - - ranked_leaders[threshold] = (leader, distance); - let mut slot = threshold; - while slot > 0 && ranked_leaders[slot].1 < ranked_leaders[slot - 1].1 { - ranked_leaders.swap(slot, slot - 1); - slot -= 1; - } + }); } #[cfg(test)] @@ -186,7 +122,7 @@ mod tests { struct KernelCall<'a> { distances: MatrixView<'a, f32>, output: MutMatrixView<'a, u32>, - ranked_leaders: &'a mut Vec<(u32, f32)>, + ranked_leaders: &'a mut Vec, } struct RankDistances; @@ -204,8 +140,8 @@ mod tests { distances: MatrixView<'_, f32>, nearest_leader_count: usize, ) -> Vec { - let mut output = - Matrix::new(UNASSIGNED_LEADER, distances.nrows(), nearest_leader_count); + // A previous valid leader must be overwritten even when no candidate is rankable. + let mut output = Matrix::new(0, distances.nrows(), nearest_leader_count); arch::dispatch1_no_features( RankDistances, KernelCall { @@ -216,59 +152,13 @@ mod tests { ); output.into_inner().into_vec() } - - pub(super) fn reference_assignments( - distances: MatrixView<'_, f32>, - nearest_leader_count: usize, - ) -> Vec { - let mut output = vec![UNASSIGNED_LEADER; distances.nrows() * nearest_leader_count]; - for (row, assignments) in distances - .row_iter() - .zip(output.chunks_exact_mut(nearest_leader_count)) - { - let mut candidates: Vec<_> = row - .iter() - .enumerate() - .filter_map(|(leader, &distance)| { - (distance < f32::INFINITY).then_some((leader as u32, distance)) - }) - .collect(); - candidates.sort_unstable_by(|left, right| { - left.1 - .partial_cmp(&right.1) - .unwrap_or(std::cmp::Ordering::Equal) - }); - for (destination, (leader, _)) in assignments.iter_mut().zip(candidates) { - *destination = leader; - } - } - output - } - } - - mod insert_leader_tests { - use super::*; - - #[test] - fn nan_distance_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]; - - // When - insert_leader(&mut ranked_leaders, 0, 0.25); - insert_leader(&mut ranked_leaders, 1, f32::NAN); - - // Then - assert_eq!(ranked_leaders, expected_ranked_leaders); - } } mod assign_leaders_tests { use super::*; #[test] - fn assigns_each_point_to_highest_similarity_leaders_with_cosine() { + fn cosine_assignment_reuses_workspace_and_output_for_a_different_point() { // Given let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; let leaders = @@ -276,7 +166,9 @@ mod tests { 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 expected_reassigned_leaders = [2, 1]; let mut actual_assignments = [UNASSIGNED_LEADER; 4]; + let mut workspace = PartitionKernelWorkspace::default(); // When assign_leaders::<_, Cosine>( @@ -284,7 +176,7 @@ mod tests { points, &leaders, MutMatrixView::try_from(&mut actual_assignments[..], 2, 2).unwrap(), - &mut PartitionKernelWorkspace::default(), + &mut workspace, ) .unwrap(); @@ -293,54 +185,19 @@ mod tests { 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 = - Cosine::create_leaders(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]; + // When: shrink to the second point, whose nearest leader differs from the old prefix. assign_leaders::<_, Cosine>( diskann_wide::ARCH, - smaller_points, + MatrixView::try_from(&point_values[2..], 1, 2).unwrap(), &leaders, - MutMatrixView::try_from(&mut actual_assignments_from_reused_workspace[..], 1, 2) - .unwrap(), - &mut reused_workspace, + MutMatrixView::try_from(&mut actual_assignments[..2], 1, 2).unwrap(), + &mut workspace, ) .unwrap(); // Then - assert_eq!( - actual_assignments_from_reused_workspace, - expected_assignments_from_fresh_workspace - ); + assert_eq!(actual_assignments[..2], expected_reassigned_leaders); } } @@ -361,7 +218,7 @@ mod tests { #[case::two_lanes_plus_one(33, 7)] #[case::root_partition(1000, 64)] #[trace] - fn dispatched_partition_ranking_matches_scalar_reference_across_lane_boundaries( + fn dispatched_partition_ranking_selects_nearest_columns_across_lane_boundaries( #[case] leader_count: usize, #[case] nearest_leader_count: usize, ) { @@ -371,7 +228,10 @@ mod tests { .map(|leader| leader as f32 - leader_count as f32 / 2.0) .collect(); let distances = MatrixView::try_from(values.as_slice(), 2, leader_count).unwrap(); - let expected = reference_assignments(distances, nearest_leader_count); + let expected: Vec<_> = (0..nearest_leader_count) + .chain((leader_count - nearest_leader_count..leader_count).rev()) + .map(|leader| leader as u32) + .collect(); // When let actual = rank_distance_fixture(distances, nearest_leader_count); @@ -403,11 +263,9 @@ mod tests { #[rstest] #[case::nan(f32::NAN)] #[case::positive_infinity(f32::INFINITY)] - fn non_rankable_distances_leave_only_unfilled_ranks_unassigned( - #[case] unrankable: f32, - #[values(3, 17, 33)] leader_count: usize, - ) { - // Given: only the first row has rankable candidates. + fn non_rankable_distances_leave_only_unfilled_ranks_unassigned(#[case] unrankable: f32) { + // Given: only the first row has rankable candidates, spanning SIMD and scalar tail. + let leader_count = 17; let mut values = vec![unrankable; 2 * leader_count]; values[leader_count - 2] = -1.0; values[leader_count - 1] = 2.0; diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 9d4271736f..70b2552013 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -372,89 +372,33 @@ mod tests { } #[test] - fn reused_l2_leaders_match_fresh_leaders_for_a_new_point_stripe() { - // Given + fn l2_reinitializes_every_output_row_for_a_new_point_stripe() { let leader_values = [3.0_f32, 4.0, 0.0, 2.0]; - let leader_count = leader_values.len() / DIMENSION_COUNT; - let first_point = [1.0_f32, 0.0]; - let second_point = [0.0_f32, 1.0]; - let first_leader_squared_norm = - leader_values[0].mul_add(leader_values[0], leader_values[1] * leader_values[1]); - let second_leader_squared_norm = - leader_values[2].mul_add(leader_values[2], leader_values[3] * leader_values[3]); - let first_dot = - second_point[0].mul_add(leader_values[0], second_point[1] * leader_values[1]); - let second_dot = - second_point[0].mul_add(leader_values[2], second_point[1] * leader_values[3]); - let expected = [ - (-2.0_f32).mul_add(first_dot, first_leader_squared_norm), - (-2.0_f32).mul_add(second_dot, second_leader_squared_norm), - ]; - let reused_leaders = L2::create_leaders(matrix(&leader_values, leader_count)); - let fresh_leaders = L2::create_leaders(matrix(&leader_values, leader_count)); - let mut discarded_first_output = [STALE_DISTANCE; 2]; - L2::compute_distances( - matrix(&first_point, 1), - &reused_leaders, - &mut discarded_first_output, - ) - .unwrap(); - let mut reused_output = [STALE_DISTANCE; 2]; - let mut fresh_output = [STALE_DISTANCE; 2]; + let leaders = L2::create_leaders(matrix(&leader_values, 2)); + let mut output = [STALE_DISTANCE; 4]; - // When - L2::compute_distances( - matrix(&second_point, 1), - &reused_leaders, - &mut reused_output, - ) - .unwrap(); - L2::compute_distances(matrix(&second_point, 1), &fresh_leaders, &mut fresh_output) - .unwrap(); + // Leader norms are 25 and 4; each entry is norm - 2 * dot. + L2::compute_distances(matrix(&[1.0, 0.0, 0.0, 1.0], 2), &leaders, &mut output).unwrap(); + assert_eq!(output, [19.0, 4.0, 17.0, 0.0]); - // Then - assert_eq!(reused_output, expected); - assert_eq!(fresh_output, expected); + L2::compute_distances(matrix(&[2.0, 1.0, -1.0, 2.0], 2), &leaders, &mut output) + .unwrap(); + assert_eq!(output, [5.0, 0.0, 15.0, -4.0]); } #[test] - fn reused_cosine_leaders_match_fresh_leaders_for_a_new_point_stripe() { - // Given + fn cosine_overwrites_output_for_a_new_point_stripe() { let leader_values = [1.0_f32, 0.0, 0.0, 2.0]; - let leader_count = leader_values.len() / DIMENSION_COUNT; - let first_point = [1.0_f32, 0.0]; - let second_point = [0.0_f32, 1.0]; - let orthogonal_similarity = 0.0_f32; - let equal_direction_similarity = 1.0_f32; - let expected = [ - 1.0 - orthogonal_similarity, - 1.0 - equal_direction_similarity, - ]; - let reused_leaders = Cosine::create_leaders(matrix(&leader_values, leader_count)); - let fresh_leaders = Cosine::create_leaders(matrix(&leader_values, leader_count)); - let mut discarded_first_output = [STALE_DISTANCE; 2]; - Cosine::compute_distances( - matrix(&first_point, 1), - &reused_leaders, - &mut discarded_first_output, - ) - .unwrap(); - let mut reused_output = [STALE_DISTANCE; 2]; - let mut fresh_output = [STALE_DISTANCE; 2]; + let leaders = Cosine::create_leaders(matrix(&leader_values, 2)); + let mut output = [STALE_DISTANCE; 4]; - // When - Cosine::compute_distances( - matrix(&second_point, 1), - &reused_leaders, - &mut reused_output, - ) - .unwrap(); - Cosine::compute_distances(matrix(&second_point, 1), &fresh_leaders, &mut fresh_output) + Cosine::compute_distances(matrix(&[1.0, 0.0, 0.0, 3.0], 2), &leaders, &mut output) .unwrap(); + assert_eq!(output, [0.0, 1.0, 1.0, 0.0]); - // Then - assert_eq!(reused_output, expected); - assert_eq!(fresh_output, expected); + Cosine::compute_distances(matrix(&[0.0, 4.0, -2.0, 0.0], 2), &leaders, &mut output) + .unwrap(); + assert_eq!(output, [1.0, 0.0, 2.0, 1.0]); } } } diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index edbc57d6b7..964b65dbd8 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -14,9 +14,6 @@ type DefaultVector = ::f32x16; /// Operations required by PiPNN SIMD vectors. pub(super) trait PiPNNSIMDVector: SIMDVector + SIMDFloat { - /// Convert the SIMD value to a readable lane array. - fn to_lane_array(self) -> impl AsRef<[f32]>; - /// Return one bit for each selected lane. fn active_lanes(mask: Self::Mask) -> u64; } @@ -27,12 +24,6 @@ where Const: SupportedLaneCount, u64: From<<::BitMask as SIMDMask>::Underlying>, { - #[inline(always)] - fn to_lane_array(self) -> impl AsRef<[f32]> { - let values: [f32; N] = self.to_array(); - values - } - #[inline(always)] fn active_lanes(mask: Self::Mask) -> u64 { u64::from(mask.bitmask().to_underlying()) @@ -52,3 +43,48 @@ where { type Vector = DefaultVector; } + +/// One group of distances and their column indexes in the supplied slice. +#[derive(Clone, Copy)] +pub(super) enum DistanceBlock<'a, F: PiPNNSIMDVector> { + Simd { + first_idx: usize, + values: F, + lanes: &'a [f32], + }, + Scalar { + idx: usize, + distance: f32, + }, +} + +/// Iterate over distance groups without a callback on the ranking hot path. +/// +/// Scalar lanes borrow the input slice, so sharing a block across updates needs +/// no temporary array. Tail entries follow in column order. The caller keeps the +/// selected architecture's execution scope active. +#[inline(always)] +pub(super) fn distance_blocks( + arch: A, + distances: &[f32], +) -> impl Iterator> { + let simd_end = distances.len() - distances.len() % A::Vector::LANES; + distances[..simd_end] + .chunks_exact(A::Vector::LANES) + .enumerate() + .map(move |(group, lanes)| DistanceBlock::Simd { + first_idx: group * A::Vector::LANES, + // SAFETY: chunks_exact yields one complete SIMD group. + values: unsafe { A::Vector::load_simd(arch, lanes.as_ptr()) }, + lanes, + }) + .chain( + distances[simd_end..] + .iter() + .enumerate() + .map(move |(tail, &distance)| DistanceBlock::Scalar { + idx: simd_end + tail, + distance, + }), + ) +} diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs new file mode 100644 index 0000000000..29a061bc1c --- /dev/null +++ b/diskann/src/graph/pipnn/topk.rs @@ -0,0 +1,483 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Fixed-capacity candidate rows with cached farthest distances. +//! Updates allocate no storage. NaN and positive infinity are not retained. + +use diskann_utils::views::MutMatrixView; + +use super::simd::{DistanceBlock, PiPNNSIMDVector}; + +/// An output slot with no assigned candidate. +pub(super) const UNASSIGNED: u32 = u32::MAX; + +/// A local candidate index paired with its ranking distance. +#[derive(Clone, Copy, Debug, PartialEq)] +pub(super) struct Candidate { + pub(super) local_idx: u32, + pub(super) distance: f32, +} + +impl Candidate { + pub(super) const fn new(local_idx: u32, distance: f32) -> Self { + Self { + local_idx, + distance, + } + } + + pub(super) const fn is_assigned(self) -> bool { + self.local_idx != UNASSIGNED + } +} + +impl Default for Candidate { + fn default() -> Self { + Self::new(UNASSIGNED, f32::INFINITY) + } +} + +// Private dispatch tag; callers supply only the result shape. +pub(super) const RUNTIME_WIDTH: usize = 0; + +/// Run an operation with the top-k implementation selected once from its row width. +/// +/// The body expands at the call site, preserving static specialization without +/// a callback in the inner loop. Like an ordinary block, return and ? affect the caller. +macro_rules! with_topk_rows { + ($rows:expr, $worst:expr, |$topks:ident| $body:block) => {{ + let rows = $rows; + let worst = $worst; + match rows.ncols() { + 1 => { + let mut $topks = $crate::graph::pipnn::topk::TopKRows::<1>::new(rows, worst); + $body + } + 2 => { + let mut $topks = $crate::graph::pipnn::topk::TopKRows::<2>::new(rows, worst); + $body + } + 3 => { + let mut $topks = $crate::graph::pipnn::topk::TopKRows::<3>::new(rows, worst); + $body + } + _ => { + let mut $topks = $crate::graph::pipnn::topk::TopKRows::< + { $crate::graph::pipnn::topk::RUNTIME_WIDTH }, + >::new(rows, worst); + $body + } + } + }}; +} +pub(super) use with_topk_rows; + +/// Borrowed top-k results and one contiguous cached threshold per row. +/// +/// Each row has the same positive capacity. Only this view mutates the rows and +/// their thresholds while borrowed. Callers supply each candidate ID at most once +/// per row; updates do not deduplicate. Equal distances need no fixed tie order. +/// The with_topk_rows macro selects the implementation from the result shape. +pub(super) struct TopKRows<'a, const K: usize> { + rows: MutMatrixView<'a, Candidate>, + worst: &'a mut [f32], +} + +impl<'a, const K: usize> TopKRows<'a, K> { + /// Bind and clear caller-owned result and threshold buffers. + pub(super) fn new(rows: MutMatrixView<'a, Candidate>, worst: &'a mut [f32]) -> Self { + debug_assert!(rows.ncols() > 0, "top-k row capacity must be positive"); + debug_assert!( + K == RUNTIME_WIDTH || K == rows.ncols(), + "top-k compile-time width must match its rows" + ); + debug_assert_eq!( + rows.nrows(), + worst.len(), + "top-k rows and thresholds must match" + ); + let mut topk = Self { rows, worst }; + topk.reset(); + topk + } + + pub(super) fn reset(&mut self) { + self.rows.as_mut_slice().fill(Candidate::default()); + self.worst.fill(f32::INFINITY); + } + + pub(super) fn row(&self, row_idx: usize) -> &[Candidate] { + self.rows.row(row_idx) + } + + /// Offer one candidate and synchronize the row's cached threshold. + #[inline(always)] + pub(super) fn insert(&mut self, row_idx: usize, candidate: Candidate) { + if candidate.distance < self.worst[row_idx] { + self.worst[row_idx] = + insert_eligible::(row_mut::(&mut self.rows, row_idx), candidate); + } + } + + /// Offer the block's candidates to one row using their slice-column indexes. + #[inline(always)] + pub(super) fn update_one( + &mut self, + row_idx: usize, + block: &DistanceBlock<'_, F>, + ) { + match *block { + DistanceBlock::Scalar { idx, distance } => { + self.insert(row_idx, Candidate::new(idx as u32, distance)); + } + DistanceBlock::Simd { + first_idx, + values, + lanes, + } => { + let mut worst = self.worst[row_idx]; + let mut eligible = F::active_lanes(values.lt_simd(F::splat(values.arch(), worst))); + if eligible == 0 { + return; + } + let row = row_mut::(&mut self.rows, row_idx); + while eligible != 0 { + let lane = eligible.trailing_zeros() as usize; + eligible &= eligible - 1; + // Earlier candidates in this same block can lower the threshold. + if lanes[lane] < worst { + worst = insert_eligible::( + row, + Candidate::new((first_idx + lane) as u32, lanes[lane]), + ); + } + } + self.worst[row_idx] = worst; + } + } + } + + /// Offer one candidate to each row identified by the block's slice-column indexes. + #[inline(always)] + pub(super) fn update_many( + &mut self, + candidate_idx: u32, + block: &DistanceBlock<'_, F>, + ) { + match *block { + DistanceBlock::Scalar { idx, distance } => { + self.insert(idx, Candidate::new(candidate_idx, distance)); + } + DistanceBlock::Simd { + first_idx, + values, + lanes, + } => { + let thresholds = &mut self.worst[first_idx..][..F::LANES]; + // SAFETY: the slice above has one threshold for every SIMD lane. + let worst = unsafe { F::load_simd(values.arch(), thresholds.as_ptr()) }; + let mut eligible = F::active_lanes(values.lt_simd(worst)); + while eligible != 0 { + let lane = eligible.trailing_zeros() as usize; + eligible &= eligible - 1; + // Each lane updates a different row once; its threshold is still current. + thresholds[lane] = insert_eligible::( + row_mut::(&mut self.rows, first_idx + lane), + Candidate::new(candidate_idx, lanes[lane]), + ); + } + } + } + } +} + +/// Keep fixed-width row addressing visible to the compiler as well as fixed-width insertion. +#[inline(always)] +fn row_mut<'a, const K: usize>( + rows: &'a mut MutMatrixView<'_, Candidate>, + row_idx: usize, +) -> &'a mut [Candidate] { + if K == RUNTIME_WIDTH { + rows.row_mut(row_idx) + } else { + let (rows, _) = rows.as_mut_slice().as_chunks_mut::(); + &mut rows[row_idx] + } +} + +/// Insert after checking the current threshold; keep the small-capacity specializations. +#[inline(always)] +fn insert_eligible(row: &mut [Candidate], candidate: Candidate) -> f32 { + match K { + 1 => { + row[0] = candidate; + candidate.distance + } + 2 => { + let first = row[0]; + if candidate.distance < first.distance { + row[0] = candidate; + row[1] = first; + first.distance + } else { + row[1] = candidate; + candidate.distance + } + } + 3 => { + let (first, second) = (row[0], row[1]); + if candidate.distance < first.distance { + row[0] = candidate; + row[1] = first; + row[2] = second; + second.distance + } else if candidate.distance < second.distance { + row[1] = candidate; + row[2] = second; + second.distance + } else { + row[2] = candidate; + candidate.distance + } + } + _ => { + let last = row.len() - 1; + let mut slot = last; + while slot > 0 && candidate.distance < row[slot - 1].distance { + row[slot] = row[slot - 1]; + slot -= 1; + } + row[slot] = candidate; + row[last].distance + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + mod insert_tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::one(&[(1, 4.0)], (2, 2.0), &[(2, 2.0)])] + #[case::two_front(&[(1, 1.0), (2, 3.0)], (3, 0.5), &[(3, 0.5), (1, 1.0)])] + #[case::two_back(&[(1, 1.0), (2, 3.0)], (3, 2.0), &[(1, 1.0), (3, 2.0)])] + #[case::three_front(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 0.5), &[(4, 0.5), (1, 1.0), (2, 2.0)])] + #[case::three_middle(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 1.5), &[(1, 1.0), (4, 1.5), (2, 2.0)])] + #[case::three_back(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 3.0), &[(1, 1.0), (2, 2.0), (4, 3.0)])] + #[case::runtime_middle(&[(1, 1.0), (2, 2.0), (3, 3.0), (4, 5.0)], (5, 2.5), &[(1, 1.0), (2, 2.0), (5, 2.5), (3, 3.0)])] + fn insertion_preserves_nearest_order_and_threshold( + #[case] initial: &[(u32, f32)], + #[case] candidate: (u32, f32), + #[case] expected: &[(u32, f32)], + ) { + let expected: Vec<_> = expected + .iter() + .map(|&(idx, distance)| Candidate::new(idx, distance)) + .collect(); + let mut output = vec![Candidate::default(); initial.len()]; + let mut worst = [f32::INFINITY]; + + with_topk_rows!( + MutMatrixView::row_vector(output.as_mut_slice()), + &mut worst[..], + |topk| { + for &(idx, distance) in initial { + topk.insert(0, Candidate::new(idx, distance)); + } + topk.insert(0, Candidate::new(candidate.0, candidate.1)); + + assert_eq!(topk.row(0), expected); + assert_eq!(topk.worst[0], expected.last().unwrap().distance); + } + ); + } + } + + mod update_tests { + use super::*; + use crate::graph::pipnn::simd::distance_blocks; + use rstest::rstest; + + #[test] + fn one_row_rechecks_the_threshold_after_each_insert() { + let mut distances = [f32::INFINITY; 16]; + distances[0] = 1.0; + distances[1] = 9.0; + let mut output = [Candidate::default(); 1]; + let mut worst = [f32::INFINITY]; + let mut topk = TopKRows::::new( + MutMatrixView::row_vector(&mut output[..]), + &mut worst, + ); + + for block in distance_blocks(diskann_wide::ARCH, &distances) { + topk.update_one(0, &block); + } + + assert_eq!(topk.row(0), [Candidate::new(0, 1.0)]); + assert_eq!(topk.worst, [1.0]); + } + + #[test] + fn many_rows_use_their_own_threshold_and_distance() { + let mut distances = [2.0; 16]; + distances[1] = 4.0; + distances[2] = 6.0; + let mut output = [Candidate::default(); 16]; + let mut worst = [f32::INFINITY; 16]; + let mut topk = TopKRows::::new( + MutMatrixView::try_from(&mut output[..], 16, 1).unwrap(), + &mut worst, + ); + topk.insert(0, Candidate::new(80, 1.0)); + topk.insert(1, Candidate::new(81, 5.0)); + topk.insert(2, Candidate::new(82, 5.0)); + + for block in distance_blocks(diskann_wide::ARCH, &distances) { + topk.update_many(42, &block); + } + + assert_eq!(topk.row(0), [Candidate::new(80, 1.0)]); + assert_eq!(topk.row(1), [Candidate::new(42, 4.0)]); + assert_eq!(topk.row(2), [Candidate::new(82, 5.0)]); + assert!( + topk.rows + .row_iter() + .skip(3) + .all(|row| row == [Candidate::new(42, 2.0)]) + ); + assert_eq!(topk.worst[..3], [1.0, 4.0, 5.0]); + assert!(topk.worst[3..].iter().all(|&worst| worst == 2.0)); + } + + #[test] + fn one_and_many_updates_have_independent_acceptance() { + let distances = [2.0; 16]; + let mut output = [Candidate::default(); 17]; + let mut worst = [f32::INFINITY; 17]; + let mut topk = TopKRows::::new( + MutMatrixView::try_from(&mut output[..], 17, 1).unwrap(), + &mut worst, + ); + topk.insert(16, Candidate::new(70, 1.0)); + + for block in distance_blocks(diskann_wide::ARCH, &distances) { + topk.update_one(16, &block); + topk.update_many(16, &block); + } + + assert_eq!(topk.row(16), [Candidate::new(70, 1.0)]); + assert!( + topk.rows + .row_iter() + .take(16) + .all(|row| row == [Candidate::new(16, 2.0)]) + ); + } + + #[rstest] + #[case::nan(f32::NAN)] + #[case::infinity(f32::INFINITY)] + fn non_rankable_candidates_leave_reset_rows_unassigned(#[case] distance: f32) { + let distances = [distance; 16]; + let mut output = [Candidate::new(70, 1.0); 16]; + let mut worst = [1.0; 16]; + let mut topk = TopKRows::::new( + MutMatrixView::try_from(&mut output[..], 16, 1).unwrap(), + &mut worst, + ); + + for block in distance_blocks(diskann_wide::ARCH, &distances) { + topk.update_one(0, &block); + topk.update_many(42, &block); + } + topk.insert(0, Candidate::new(43, distance)); + + assert!( + topk.rows + .as_slice() + .iter() + .all(|&candidate| candidate == Candidate::default()) + ); + assert!(topk.worst.iter().all(|&worst| worst == f32::INFINITY)); + } + + #[test] + fn middle_insertions_keep_nearest_candidates_and_reset_clears_state() { + let mut output = [Candidate::default(); 3]; + let mut worst = [f32::INFINITY]; + let mut topk = TopKRows::::new( + MutMatrixView::row_vector(&mut output[..]), + &mut worst, + ); + for (idx, distance) in [0.0, 4.0, 6.0, 2.0, 3.0].into_iter().enumerate() { + topk.insert(0, Candidate::new(idx as u32, distance)); + } + assert_eq!( + topk.row(0), + [ + Candidate::new(0, 0.0), + Candidate::new(3, 2.0), + Candidate::new(4, 3.0) + ] + ); + assert_eq!(topk.worst, [3.0]); + + topk.reset(); + + assert_eq!(topk.row(0), [Candidate::default(); 3]); + assert_eq!(topk.worst, [f32::INFINITY]); + } + + #[rstest] + #[case::empty(0)] + #[case::lane_minus_one(15)] + #[case::one_lane(16)] + #[case::lane_with_tail(17)] + #[case::two_lanes_with_tail(33)] + fn shared_block_updates_match_scalar_insertion(#[case] count: usize) { + let rows = count + 1; + let source_idx = count; + let distances: Vec<_> = (0..count) + .map(|idx| { + let distance = idx as f32 + 0.25; + if idx % 2 == 0 { -distance } else { distance } + }) + .collect(); + let mut actual_rows = vec![Candidate::default(); rows * 3]; + let mut expected_rows = actual_rows.clone(); + let mut actual_worst = vec![f32::INFINITY; rows]; + let mut expected_worst = actual_worst.clone(); + let mut actual = TopKRows::<3>::new( + MutMatrixView::try_from(actual_rows.as_mut_slice(), rows, 3).unwrap(), + &mut actual_worst, + ); + let mut expected = TopKRows::<3>::new( + MutMatrixView::try_from(expected_rows.as_mut_slice(), rows, 3).unwrap(), + &mut expected_worst, + ); + actual.insert(source_idx, Candidate::new(90, 0.0)); + expected.insert(source_idx, Candidate::new(90, 0.0)); + actual.insert(0, Candidate::new(91, -4.0)); + expected.insert(0, Candidate::new(91, -4.0)); + + for block in distance_blocks(diskann_wide::ARCH, &distances) { + actual.update_one(source_idx, &block); + actual.update_many(42, &block); + } + for (idx, &distance) in distances.iter().enumerate() { + expected.insert(source_idx, Candidate::new(idx as u32, distance)); + expected.insert(idx, Candidate::new(42, distance)); + } + + assert_eq!(actual.rows.as_slice(), expected.rows.as_slice()); + assert_eq!(actual.worst, expected.worst); + } + } +} From 5f85a352163c88567268a98d440ebb57c9764edb Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 8 Sep 2026 07:18:07 +0000 Subject: [PATCH 102/229] refactor(pipnn): simplify top-k insertion Keep constant widths in row dispatch so the compiler can optimize one insertion loop. Test shifts directly and share width-dispatch coverage across scalar and SIMD updates. --- diskann/src/graph/pipnn/topk.rs | 209 ++++++++++++++------------------ 1 file changed, 90 insertions(+), 119 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 29a061bc1c..51f437739f 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -42,10 +42,15 @@ impl Default for Candidate { // Private dispatch tag; callers supply only the result shape. pub(super) const RUNTIME_WIDTH: usize = 0; -/// Run an operation with the top-k implementation selected once from its row width. +/// Select the top-k implementation once from the row width. /// -/// The body expands at the call site, preserving static specialization without -/// a callback in the inner loop. Like an ordinary block, return and ? affect the caller. +/// For widths one to three, `row_mut` exposes a row with a compile-time length. +/// After inlining, the compiler can unroll the insertion loop and remove bounds checks. +/// Other widths use the same insertion loop with a runtime row length. +/// +/// Each branch selects a different `TopKRows` type. +/// Use a code block to keep that type through the candidate loops without a callback. +/// Like an ordinary block, `return` and `?` affect the caller. macro_rules! with_topk_rows { ($rows:expr, $worst:expr, |$topks:ident| $body:block) => {{ let rows = $rows; @@ -116,8 +121,7 @@ impl<'a, const K: usize> TopKRows<'a, K> { #[inline(always)] pub(super) fn insert(&mut self, row_idx: usize, candidate: Candidate) { if candidate.distance < self.worst[row_idx] { - self.worst[row_idx] = - insert_eligible::(row_mut::(&mut self.rows, row_idx), candidate); + self.worst[row_idx] = insert_eligible(row_mut::(&mut self.rows, row_idx), candidate); } } @@ -148,7 +152,7 @@ impl<'a, const K: usize> TopKRows<'a, K> { eligible &= eligible - 1; // Earlier candidates in this same block can lower the threshold. if lanes[lane] < worst { - worst = insert_eligible::( + worst = insert_eligible( row, Candidate::new((first_idx + lane) as u32, lanes[lane]), ); @@ -183,7 +187,7 @@ impl<'a, const K: usize> TopKRows<'a, K> { let lane = eligible.trailing_zeros() as usize; eligible &= eligible - 1; // Each lane updates a different row once; its threshold is still current. - thresholds[lane] = insert_eligible::( + thresholds[lane] = insert_eligible( row_mut::(&mut self.rows, first_idx + lane), Candidate::new(candidate_idx, lanes[lane]), ); @@ -193,7 +197,7 @@ impl<'a, const K: usize> TopKRows<'a, K> { } } -/// Keep fixed-width row addressing visible to the compiler as well as fixed-width insertion. +/// Preserve fixed row lengths for inlined candidate insertion. #[inline(always)] fn row_mut<'a, const K: usize>( rows: &'a mut MutMatrixView<'_, Candidate>, @@ -207,98 +211,82 @@ fn row_mut<'a, const K: usize>( } } -/// Insert after checking the current threshold; keep the small-capacity specializations. +/// Insert after the threshold check; keep candidates in nearest-first order. #[inline(always)] -fn insert_eligible(row: &mut [Candidate], candidate: Candidate) -> f32 { - match K { - 1 => { - row[0] = candidate; - candidate.distance - } - 2 => { - let first = row[0]; - if candidate.distance < first.distance { - row[0] = candidate; - row[1] = first; - first.distance - } else { - row[1] = candidate; - candidate.distance - } - } - 3 => { - let (first, second) = (row[0], row[1]); - if candidate.distance < first.distance { - row[0] = candidate; - row[1] = first; - row[2] = second; - second.distance - } else if candidate.distance < second.distance { - row[1] = candidate; - row[2] = second; - second.distance - } else { - row[2] = candidate; - candidate.distance - } - } - _ => { - let last = row.len() - 1; - let mut slot = last; - while slot > 0 && candidate.distance < row[slot - 1].distance { - row[slot] = row[slot - 1]; - slot -= 1; - } - row[slot] = candidate; - row[last].distance - } +fn insert_eligible(row: &mut [Candidate], candidate: Candidate) -> f32 { + let last = row.len() - 1; + let mut slot = last; + while slot > 0 && candidate.distance < row[slot - 1].distance { + row[slot] = row[slot - 1]; + slot -= 1; } + row[slot] = candidate; + row[last].distance } #[cfg(test)] mod tests { use super::*; - mod insert_tests { + mod insert_eligible_tests { use super::*; use rstest::rstest; #[rstest] - #[case::one(&[(1, 4.0)], (2, 2.0), &[(2, 2.0)])] - #[case::two_front(&[(1, 1.0), (2, 3.0)], (3, 0.5), &[(3, 0.5), (1, 1.0)])] - #[case::two_back(&[(1, 1.0), (2, 3.0)], (3, 2.0), &[(1, 1.0), (3, 2.0)])] - #[case::three_front(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 0.5), &[(4, 0.5), (1, 1.0), (2, 2.0)])] - #[case::three_middle(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 1.5), &[(1, 1.0), (4, 1.5), (2, 2.0)])] - #[case::three_back(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 3.0), &[(1, 1.0), (2, 2.0), (4, 3.0)])] - #[case::runtime_middle(&[(1, 1.0), (2, 2.0), (3, 3.0), (4, 5.0)], (5, 2.5), &[(1, 1.0), (2, 2.0), (5, 2.5), (3, 3.0)])] - fn insertion_preserves_nearest_order_and_threshold( + #[case::singleton(&[(1, 4.0)], (2, 2.0), &[(2, 2.0)])] + #[case::no_shift(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 3.0), &[(1, 1.0), (2, 2.0), (4, 3.0)])] + #[case::one_shift(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 1.5), &[(1, 1.0), (4, 1.5), (2, 2.0)])] + #[case::multiple_shifts(&[(1, 1.0), (2, 2.0), (3, 3.0), (4, 5.0)], (5, 0.5), &[(5, 0.5), (1, 1.0), (2, 2.0), (3, 3.0)])] + #[case::negative_infinity(&[(1, -1.0), (2, 1.0), (UNASSIGNED, f32::INFINITY)], (3, f32::NEG_INFINITY), &[(3, f32::NEG_INFINITY), (1, -1.0), (2, 1.0)])] + #[case::positive_zero(&[(1, 1.0)], (2, 0.0), &[(2, 0.0)])] + #[case::negative_zero(&[(1, 1.0)], (2, -0.0), &[(2, -0.0)])] + fn insertion_keeps_nearest_candidates_and_returns_threshold( #[case] initial: &[(u32, f32)], #[case] candidate: (u32, f32), #[case] expected: &[(u32, f32)], ) { + // Given + let expected_threshold = expected.last().unwrap().1.to_bits(); let expected: Vec<_> = expected + .iter() + .map(|&(idx, distance)| (idx, distance.to_bits())) + .collect(); + let mut row: Vec<_> = initial .iter() .map(|&(idx, distance)| Candidate::new(idx, distance)) .collect(); - let mut output = vec![Candidate::default(); initial.len()]; - let mut worst = [f32::INFINITY]; - with_topk_rows!( - MutMatrixView::row_vector(output.as_mut_slice()), - &mut worst[..], - |topk| { - for &(idx, distance) in initial { - topk.insert(0, Candidate::new(idx, distance)); - } - topk.insert(0, Candidate::new(candidate.0, candidate.1)); + // When + let threshold = insert_eligible(&mut row, Candidate::new(candidate.0, candidate.1)); - assert_eq!(topk.row(0), expected); - assert_eq!(topk.worst[0], expected.last().unwrap().distance); - } - ); + // Then + let actual: Vec<_> = row + .iter() + .map(|candidate| (candidate.local_idx, candidate.distance.to_bits())) + .collect(); + assert_eq!(actual, expected); + assert_eq!(threshold.to_bits(), expected_threshold); } } + #[test] + fn reset_clears_rows_and_thresholds() { + // Given + let mut output = [Candidate::new(10, -1.0), Candidate::new(11, 2.0)]; + let mut worst = [-1.0, 2.0]; + let mut topk = TopKRows::<1> { + rows: MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), + worst: &mut worst, + }; + + // When + topk.reset(); + + // Then + assert_eq!(topk.rows.as_slice(), [Candidate::default(); 2]); + assert_eq!(topk.worst, [f32::INFINITY; 2]); + } + mod update_tests { use super::*; use crate::graph::pipnn::simd::distance_blocks; @@ -408,40 +396,18 @@ mod tests { assert!(topk.worst.iter().all(|&worst| worst == f32::INFINITY)); } - #[test] - fn middle_insertions_keep_nearest_candidates_and_reset_clears_state() { - let mut output = [Candidate::default(); 3]; - let mut worst = [f32::INFINITY]; - let mut topk = TopKRows::::new( - MutMatrixView::row_vector(&mut output[..]), - &mut worst, - ); - for (idx, distance) in [0.0, 4.0, 6.0, 2.0, 3.0].into_iter().enumerate() { - topk.insert(0, Candidate::new(idx as u32, distance)); - } - assert_eq!( - topk.row(0), - [ - Candidate::new(0, 0.0), - Candidate::new(3, 2.0), - Candidate::new(4, 3.0) - ] - ); - assert_eq!(topk.worst, [3.0]); - - topk.reset(); - - assert_eq!(topk.row(0), [Candidate::default(); 3]); - assert_eq!(topk.worst, [f32::INFINITY]); - } - #[rstest] #[case::empty(0)] #[case::lane_minus_one(15)] #[case::one_lane(16)] #[case::lane_with_tail(17)] #[case::two_lanes_with_tail(33)] - fn shared_block_updates_match_scalar_insertion(#[case] count: usize) { + #[trace] + fn shared_block_updates_match_runtime_scalar_insertion( + #[case] count: usize, + #[values(1, 2, 3, 4)] width: usize, + ) { + // Given let rows = count + 1; let source_idx = count; let distances: Vec<_> = (0..count) @@ -450,34 +416,39 @@ mod tests { if idx % 2 == 0 { -distance } else { distance } }) .collect(); - let mut actual_rows = vec![Candidate::default(); rows * 3]; + let mut actual_rows = vec![Candidate::default(); rows * width]; let mut expected_rows = actual_rows.clone(); let mut actual_worst = vec![f32::INFINITY; rows]; let mut expected_worst = actual_worst.clone(); - let mut actual = TopKRows::<3>::new( - MutMatrixView::try_from(actual_rows.as_mut_slice(), rows, 3).unwrap(), - &mut actual_worst, - ); - let mut expected = TopKRows::<3>::new( - MutMatrixView::try_from(expected_rows.as_mut_slice(), rows, 3).unwrap(), + let mut expected = TopKRows::::new( + MutMatrixView::try_from(expected_rows.as_mut_slice(), rows, width).unwrap(), &mut expected_worst, ); - actual.insert(source_idx, Candidate::new(90, 0.0)); expected.insert(source_idx, Candidate::new(90, 0.0)); - actual.insert(0, Candidate::new(91, -4.0)); expected.insert(0, Candidate::new(91, -4.0)); - - for block in distance_blocks(diskann_wide::ARCH, &distances) { - actual.update_one(source_idx, &block); - actual.update_many(42, &block); - } for (idx, &distance) in distances.iter().enumerate() { expected.insert(source_idx, Candidate::new(idx as u32, distance)); expected.insert(idx, Candidate::new(42, distance)); } - assert_eq!(actual.rows.as_slice(), expected.rows.as_slice()); - assert_eq!(actual.worst, expected.worst); + with_topk_rows!( + MutMatrixView::try_from(actual_rows.as_mut_slice(), rows, width).unwrap(), + &mut actual_worst[..], + |actual| { + actual.insert(source_idx, Candidate::new(90, 0.0)); + actual.insert(0, Candidate::new(91, -4.0)); + + // When + for block in distance_blocks(diskann_wide::ARCH, &distances) { + actual.update_one(source_idx, &block); + actual.update_many(42, &block); + } + + // Then + assert_eq!(actual.rows.as_slice(), expected.rows.as_slice()); + assert_eq!(actual.worst, expected.worst); + } + ); } } } From 652bc9e7f0cc2e4402e033c23d986bf033691103 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:22:48 +0000 Subject: [PATCH 103/229] pipnn: assemble direct-candidate graph builder --- Cargo.lock | 1 + diskann/Cargo.toml | 6 +- diskann/src/graph/pipnn/finalization.rs | 123 ++++ diskann/src/graph/pipnn/finalization/tests.rs | 103 +++ diskann/src/graph/pipnn/leaf_build.rs | 347 ++++++++++ diskann/src/graph/pipnn/leaf_build/tests.rs | 370 +++++++++++ diskann/src/graph/pipnn/mod.rs | 210 +++++- diskann/src/graph/pipnn/partitioning.rs | 612 ++++++++++++++++++ diskann/src/graph/pipnn/partitioning/tests.rs | 277 ++++++++ diskann/tests/build_graph.rs | 231 +++++++ diskann/tests/config.rs | 133 ++++ 11 files changed, 2403 insertions(+), 10 deletions(-) create mode 100644 diskann/src/graph/pipnn/finalization.rs create mode 100644 diskann/src/graph/pipnn/finalization/tests.rs create mode 100644 diskann/src/graph/pipnn/leaf_build.rs create mode 100644 diskann/src/graph/pipnn/leaf_build/tests.rs create mode 100644 diskann/src/graph/pipnn/partitioning.rs create mode 100644 diskann/src/graph/pipnn/partitioning/tests.rs create mode 100644 diskann/tests/build_graph.rs create mode 100644 diskann/tests/config.rs diff --git a/Cargo.lock b/Cargo.lock index 85a46b7e09..c8a45250fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -448,6 +448,7 @@ dependencies = [ "num-traits", "pin-project", "rand", + "rayon", "relative-path 2.0.1", "rstest", "serde", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index c0016416dc..2568ac6669 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -14,6 +14,7 @@ targets = ["x86_64-unknown-linux-gnu", "aarch64-pc-windows-msvc", "x86_64-pc-win [dependencies] anyhow.workspace = true bytemuck = { workspace = true, features = ["must_cast"]} +diskann-linalg = { workspace = true, optional = true } diskann-utils = { workspace = true, default-features = false } futures-util = { workspace = true, default-features = false } half = { workspace = true, features = ["bytemuck", "num-traits"] } @@ -22,6 +23,7 @@ half = { workspace = true, features = ["bytemuck", "num-traits"] } hashbrown = { version = "0.16.0", default-features = false, features = ["default-hasher"] } num-traits.workspace = true rand.workspace = true +rayon = { workspace = true, optional = true } thiserror.workspace = true tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } tracing = { workspace = true, optional = true } @@ -58,8 +60,8 @@ panic = "warn" [features] default = ["tracing"] -# Enable PiPNN numerical kernels. -pipnn = ["dep:diskann-linalg"] +# Enable PiPNN batch graph construction. +pipnn = ["dep:diskann-linalg", "dep:rayon", "tracing"] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs new file mode 100644 index 0000000000..dab9419cc8 --- /dev/null +++ b/diskann/src/graph/pipnn/finalization.rs @@ -0,0 +1,123 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Orders complete PiPNN candidate rows and applies shared Vamana RobustPrune. + +use std::convert::Infallible; + +use crate::{ + graph::{prune, AdjacencyList, Config}, + neighbor::Neighbor, + utils::VectorRepr, + ANNError, ANNResult, +}; +use diskann_utils::views::MatrixView; +use diskann_vector::{distance::Metric, DistanceFunction}; +use rayon::prelude::*; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum FinalizationError { + #[error("candidate row count {rows} does not match the dataset point count {points}")] + RowCountMismatch { rows: usize, points: usize }, + #[error("candidate ID {candidate} in row {row} is outside a {points}-point dataset")] + InvalidCandidateId { + row: usize, + candidate: u32, + points: usize, + }, +} + +#[derive(Default)] +struct Workspace { + prune: prune::Scratch, + cache: Vec<(f32, Option)>, +} + +pub(crate) fn prune_overfull( + data: MatrixView<'_, T>, + candidates: Vec>, + graph: &Config, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + validate_candidates(&candidates, data.nrows()).map_err(ANNError::opaque)?; + + let degree = graph.pruned_degree().get(); + let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind(), false); + let distance = T::distance(metric, Some(data.ncols())); + + // build_graph installs the complete call tree in the caller-owned pool. + #[allow(clippy::disallowed_methods)] + candidates + .into_par_iter() + .enumerate() + .map_init(Workspace::default, |workspace, (source, mut row)| { + if row.len() <= degree { + return Ok(row); + } + + let source_id = u32::try_from(source).map_err(ANNError::opaque)?; + let source_vector = data.row(source); + let pool = workspace.prune.candidates_mut(); + pool.clear(); + pool.try_reserve(row.len()).map_err(ANNError::opaque)?; + pool.extend(row.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + let candidate_count = pool.len(); + let mut context = workspace.prune.as_context(candidate_count); + prune::robust_prune( + &mut context, + policy, + &mut workspace.cache, + Some, + |left, right| { + Ok::<_, Infallible>( + distance.evaluate_similarity( + data.row(*left as usize), + data.row(*right as usize), + ), + ) + }, + |id| id == source_id, + ) + .map_err(ANNError::opaque)?; + + row.clear(); + row.extend_from_slice(workspace.prune.neighbors()); + Ok(row) + }) + .collect() +} + +fn validate_candidates( + candidates: &[AdjacencyList], + points: usize, +) -> Result<(), FinalizationError> { + if candidates.len() != points { + return Err(FinalizationError::RowCountMismatch { + rows: candidates.len(), + points, + }); + } + for (row_id, row) in candidates.iter().enumerate() { + if let Some(&candidate) = row.iter().find(|&&id| id as usize >= points) { + return Err(FinalizationError::InvalidCandidateId { + row: row_id, + candidate, + points, + }); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/pipnn/finalization/tests.rs b/diskann/src/graph/pipnn/finalization/tests.rs new file mode 100644 index 0000000000..fde3a6e0ea --- /dev/null +++ b/diskann/src/graph/pipnn/finalization/tests.rs @@ -0,0 +1,103 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use crate::graph::{ + config::{self, MaxDegree}, + AdjacencyList, +}; +use diskann_utils::views::MatrixView; + +use super::*; + +fn graph_config(degree: usize) -> Config { + config::Builder::new_with( + degree, + MaxDegree::same(), + degree, + Metric::L2.into(), + |builder| { + builder.alpha(1.2); + }, + ) + .build() + .unwrap() +} + +fn row(ids: impl IntoIterator) -> AdjacencyList { + AdjacencyList::from_iter_untrusted(ids) +} + +#[test] +fn preserves_rows_within_the_degree_bound() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![row([3, 1]), row([]), row([]), row([])]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert_eq!(&*actual[0], &[1, 3]); +} + +#[test] +fn prunes_an_overfull_row_with_the_vamana_kernel() { + let data = [0.0_f32, 1.0, 2.0, -3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![row([3, 2, 1]), row([]), row([]), row([])]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert!(actual[0].len() <= 2); + assert!(actual[0].contains(1)); +} + +#[test] +fn rejects_invalid_candidate_ids_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![row([1, 3]), row([]), row([])]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::InvalidCandidateId { + row: 0, + candidate: 3, + points: 3, + }) + )); +} + +#[test] +fn rejects_candidate_row_count_mismatch_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![row([]), row([]), row([]), row([])]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::RowCountMismatch { rows: 4, points: 3 }) + )); +} + +#[test] +fn rejects_more_candidates_than_the_shared_position_type_can_represent() { + let count = u16::MAX as usize + 1; + let data = vec![0.0_f32; count + 1]; + let data = MatrixView::try_from(&data[..], count + 1, 1).unwrap(); + let mut candidates = Vec::with_capacity(count + 1); + candidates.push(row(1..=count as u32)); + candidates.resize_with(count + 1, AdjacencyList::new); + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::>(), + Some(prune::RobustPruneError::TooManyCandidates { actual, max }) + if *actual == count && *max == u16::MAX as usize + )); +} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs new file mode 100644 index 0000000000..f5fdb22928 --- /dev/null +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -0,0 +1,347 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Leaf construction and direct candidate accumulation. + +use std::{ + collections::{HashSet, TryReserveError}, + sync::Mutex, +}; + +use crate::{graph::AdjacencyList, utils::VectorRepr}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use rayon::prelude::*; + +use crate::leaf_kernel::{ + nearest_leaf_neighbors, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, +}; + +/// Failure while converting leaves into direct graph candidates. +#[derive(Debug, thiserror::Error)] +pub(crate) enum LeafBuildError { + #[error("leaf build requires at least one dimension")] + EmptyDimensions, + #[error("dataset point count {0} exceeds the u32 ID limit")] + TooManyPoints(usize), + #[error("leaf {leaf} is empty")] + EmptyLeaf { leaf: usize }, + #[error("point ID {point} in leaf {leaf} is outside a {points}-point dataset")] + InvalidPointId { + leaf: usize, + point: u32, + points: usize, + }, + #[error("point ID {point} appears more than once in leaf {leaf}")] + DuplicatePointId { leaf: usize, point: u32 }, + #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] + ShapeOverflow { + leaf: usize, + rows: usize, + columns: usize, + }, + #[error("failed to reserve {additional} values for {buffer}")] + Allocation { + buffer: &'static str, + additional: usize, + #[source] + source: TryReserveError, + }, + #[error("failed to convert point {point} in leaf {leaf}")] + Conversion { + leaf: usize, + point: u32, + #[source] + source: diskann::ANNError, + }, + #[error("lower-AAT failed for leaf {leaf}")] + LowerAat { + leaf: usize, + #[source] + source: diskann_linalg::SgemmError, + }, + #[error("nearest-neighbor selection failed for leaf {leaf}")] + Kernel { + leaf: usize, + #[source] + source: LeafKernelError, + }, + #[error("leaf kernel returned local position {position} for a {points}-point leaf")] + InvalidLocalPosition { position: u32, points: usize }, + #[error("candidate row {point} is poisoned")] + PoisonedCandidateRow { point: u32 }, +} + +#[derive(Default)] +struct LeafBuffers { + points: Vec, + dots: Vec, + nearest: Vec, + local_graph: Vec>, + top_k: LeafTopKWorkspace, + seen_ids: HashSet, +} + +impl LeafBuffers { + fn prepare( + &mut self, + leaf: usize, + points: usize, + dimensions: usize, + k: usize, + ) -> Result { + let point_values = points + .checked_mul(dimensions) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: dimensions, + })?; + let dot_values = points + .checked_mul(points) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: points, + })?; + let actual_k = k.min(points.saturating_sub(1)); + let nearest_values = points + .checked_mul(actual_k) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: actual_k, + })?; + + resize("leaf points", &mut self.points, point_values, 0.0)?; + resize("leaf dot products", &mut self.dots, dot_values, 0.0)?; + resize( + "leaf nearest neighbors", + &mut self.nearest, + nearest_values, + LeafNeighbor::default(), + )?; + if points > self.local_graph.len() { + let additional = points - self.local_graph.len(); + self.local_graph + .try_reserve(additional) + .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; + self.local_graph.resize_with(points, AdjacencyList::new); + } + self.local_graph[..points] + .iter_mut() + .for_each(AdjacencyList::clear); + Ok(actual_k) + } +} + +struct DirectCandidates { + rows: Vec>>, +} + +impl DirectCandidates { + fn new(points: usize) -> Result { + let mut rows = Vec::new(); + rows.try_reserve_exact(points) + .map_err(|source| allocation_error("candidate rows", points, source))?; + rows.resize_with(points, || Mutex::new(AdjacencyList::new())); + Ok(Self { rows }) + } + + fn add_leaf( + &self, + point_ids: &[u32], + local_graph: &[AdjacencyList], + ) -> Result<(), LeafBuildError> { + for (&source, additions) in point_ids.iter().zip(local_graph) { + // Every point ID is validated before leaf-local work begins. + let row = &self.rows[source as usize]; + let mut row = row.lock().map_err(|_| poisoned_row(source))?; + row.extend_from_slice(additions); + } + Ok(()) + } + + fn into_rows(self) -> Result>, LeafBuildError> { + let mut output = Vec::new(); + output + .try_reserve_exact(self.rows.len()) + .map_err(|source| allocation_error("candidate output", self.rows.len(), source))?; + for (point, row) in self.rows.into_iter().enumerate() { + let mut row = row.into_inner().map_err(|_| poisoned_row(point as u32))?; + row.sort(); + output.push(row); + } + Ok(output) + } +} + +/// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. +#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. +pub(crate) fn build_leaf_candidates( + data: MatrixView<'_, T>, + leaves: &[Vec], + k: usize, + metric: Metric, +) -> Result>, LeafBuildError> +where + T: VectorRepr + 'static, +{ + if data.ncols() == 0 { + return Err(LeafBuildError::EmptyDimensions); + } + if data.nrows() > u32::MAX as usize { + return Err(LeafBuildError::TooManyPoints(data.nrows())); + } + + let candidates = DirectCandidates::new(data.nrows())?; + leaves.par_iter().enumerate().try_for_each_init( + LeafBuffers::default, + |buffers, (leaf, point_ids)| { + build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) + }, + )?; + candidates.into_rows() +} + +fn build_leaf( + data: MatrixView<'_, T>, + leaf: usize, + point_ids: &[u32], + k: usize, + metric: Metric, + buffers: &mut LeafBuffers, + candidates: &DirectCandidates, +) -> Result<(), LeafBuildError> +where + T: VectorRepr + 'static, +{ + if point_ids.is_empty() { + return Err(LeafBuildError::EmptyLeaf { leaf }); + } + buffers.seen_ids.clear(); + buffers + .seen_ids + .try_reserve(point_ids.len()) + .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; + for &point in point_ids { + if point as usize >= data.nrows() { + return Err(LeafBuildError::InvalidPointId { + leaf, + point, + points: data.nrows(), + }); + } + if !buffers.seen_ids.insert(point) { + return Err(LeafBuildError::DuplicatePointId { leaf, point }); + } + } + let actual_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), k)?; + if actual_k == 0 { + return Ok(()); + } + + for (&point, output) in point_ids + .iter() + .zip(buffers.points.chunks_exact_mut(data.ncols())) + { + // Point IDs were validated before the zero-k/singleton fast path. + let row = data.row(point as usize); + T::as_f32_into(row, output).map_err(|source| LeafBuildError::Conversion { + leaf, + point, + source: source.into(), + })?; + } + + diskann_linalg::sgemm_aat_lower( + &buffers.points, + point_ids.len(), + data.ncols(), + &mut buffers.dots, + ) + .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; + nearest_leaf_neighbors( + LeafTopK { + dots: &buffers.dots, + points: point_ids.len(), + metric, + }, + k, + &mut buffers.nearest, + &mut buffers.top_k, + ) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + + add_symmetric_edges( + point_ids, + actual_k, + &buffers.nearest, + &mut buffers.local_graph[..point_ids.len()], + )?; + candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) +} + +fn add_symmetric_edges( + point_ids: &[u32], + k: usize, + nearest: &[LeafNeighbor], + local_graph: &mut [AdjacencyList], +) -> Result<(), LeafBuildError> { + for (source, nearest) in nearest.chunks_exact(k).enumerate() { + for neighbor in nearest { + let target = neighbor.position as usize; + let Some(&target_id) = point_ids.get(target) else { + return Err(LeafBuildError::InvalidLocalPosition { + position: neighbor.position, + points: point_ids.len(), + }); + }; + let source_id = point_ids[source]; + if source_id != target_id { + local_graph[source].push(target_id); + local_graph[target].push(source_id); + } + } + } + Ok(()) +} + +fn resize( + buffer: &'static str, + values: &mut Vec, + len: usize, + value: T, +) -> Result<(), LeafBuildError> { + if len > values.len() { + let additional = len - values.len(); + values + .try_reserve(additional) + .map_err(|source| allocation_error(buffer, additional, source))?; + values.resize(len, value); + } else { + values.truncate(len); + } + Ok(()) +} + +fn allocation_error( + buffer: &'static str, + additional: usize, + source: TryReserveError, +) -> LeafBuildError { + LeafBuildError::Allocation { + buffer, + additional, + source, + } +} + +fn poisoned_row(point: u32) -> LeafBuildError { + LeafBuildError::PoisonedCandidateRow { point } +} + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs new file mode 100644 index 0000000000..d3f63263dc --- /dev/null +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -0,0 +1,370 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use half::f16; +use std::collections::BTreeSet; + +use super::{ + add_symmetric_edges, allocation_error, build_leaf_candidates, DirectCandidates, LeafBuffers, + LeafBuildError, +}; + +fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + MatrixView::try_from(data, rows, columns).unwrap() +} + +fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() +} + +fn build( + data: MatrixView<'_, T>, + leaves: &[Vec], + k: usize, + metric: Metric, +) -> Result>, LeafBuildError> +where + T: diskann::utils::VectorRepr + 'static, +{ + pool().install(|| build_leaf_candidates(data, leaves, k, metric)) +} + +fn rows(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() +} + +fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { + let mut graph = vec![BTreeSet::new(); data.len()]; + for (source, left) in data.iter().enumerate() { + let mut nearest: Vec<_> = data + .iter() + .enumerate() + .filter(|(target, _)| *target != source) + .map(|(target, right)| { + let distance = left + .iter() + .zip(right) + .map(|(x, y)| (x - y) * (x - y)) + .sum::(); + (target, distance) + }) + .collect(); + nearest.sort_by(|left, right| { + left.1 + .total_cmp(&right.1) + .then_with(|| left.0.cmp(&right.0)) + }); + for &(target, _) in nearest.iter().take(k) { + graph[source].insert(target as u32); + graph[target].insert(source as u32); + } + } + graph + .into_iter() + .map(|neighbors| neighbors.into_iter().collect()) + .collect() +} + +#[test] +fn leaf_adjacency_matches_an_independent_all_pairs_reference() { + let points = [ + [0.0_f32, 0.0], + [1.0, 0.2], + [3.1, 0.5], + [7.8, 1.4], + [-2.3, 4.1], + [6.7, -3.2], + ]; + let flat: Vec<_> = points.into_iter().flatten().collect(); + + let actual = rows( + build( + view(&flat, points.len(), 2), + &[(0..points.len() as u32).collect()], + 2, + Metric::L2, + ) + .unwrap(), + ); + + assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); +} + +#[test] +fn retains_and_deduplicates_candidates_from_overlapping_leaves() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; + + let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + rows(graph), + [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] + ); +} + +#[test] +fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { + let dimensions = 9; + let mut data = vec![0.0_f32; 10 * dimensions]; + for row in 1..10 { + data[row * dimensions + row - 1] = 1.0; + } + + let graph = build( + view(&data, 10, dimensions), + &[(0..10).collect()], + 1, + Metric::L2, + ) + .unwrap(); + + assert_eq!(&*graph[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); + assert!(graph.iter().enumerate().all(|(source, neighbors)| { + neighbors.iter().all(|&target| target as usize != source) + && neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + })); +} + +#[test] +fn global_id_translation_is_independent_of_leaf_order() { + let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; + let leaves = vec![vec![4, 1, 3]]; + + let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + rows(graph), + [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] + ); +} + +fn assert_source_type(data: &[T]) +where + T: diskann::utils::VectorRepr + 'static, +{ + let leaves = vec![vec![0, 1, 2, 3]]; + let graph = build(view(data, 4, 2), &leaves, 1, Metric::L2).unwrap(); + assert_eq!(rows(graph), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); +} + +#[test] +fn gathers_every_supported_source_type_without_full_dataset_conversion() { + assert_source_type(&[0.0_f32, 0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0]); + assert_source_type(&[0_i8, 0, 1, 0, 2, 0, 3, 0]); + assert_source_type(&[0_u8, 0, 1, 0, 2, 0, 3, 0]); + assert_source_type(&[ + f16::from_f32(0.0), + f16::from_f32(0.0), + f16::from_f32(1.0), + f16::from_f32(0.0), + f16::from_f32(2.0), + f16::from_f32(0.0), + f16::from_f32(3.0), + f16::from_f32(0.0), + ]); +} + +#[test] +fn all_metrics_produce_symmetric_unique_non_self_candidates() { + let data = [1.0_f32, 0.0, 0.8, 0.2, 0.0, 1.0, -1.0, 0.0]; + let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); + for (source, neighbors) in graph.iter().enumerate() { + assert!(neighbors.iter().all(|&target| target as usize != source)); + assert!(neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32))); + assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); + } + } +} + +#[test] +fn parallel_leaf_schedule_does_not_change_candidate_order() { + let data: Vec = (0..64).map(|value| value as f32).collect(); + let leaves: Vec> = (0..32) + .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) + .collect(); + let pool = pool(); + pool.install(|| { + let expected = build_leaf_candidates(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + for _ in 0..8 { + let actual = build_leaf_candidates(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + assert_eq!(actual, expected); + } + }); +} + +#[test] +fn rejects_invalid_shape_inputs_without_panicking() { + let data = [0.0_f32, 1.0]; + let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); + assert!(matches!( + build(no_dimensions, &[], 1, Metric::L2), + Err(LeafBuildError::EmptyDimensions) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), + Err(LeafBuildError::EmptyLeaf { leaf: 0 }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { + leaf: 0, + point: 2, + points: 2 + }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) + )); +} + +#[test] +fn singleton_and_zero_k_leaves_add_no_candidates() { + let data = [0.0_f32, 1.0, 2.0]; + let singleton = build( + view(&data, 3, 1), + &[vec![0], vec![1], vec![2]], + 1, + Metric::L2, + ) + .unwrap(); + let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); + assert!(singleton.iter().chain(&zero_k).all(|row| row.is_empty())); +} + +#[test] +fn reuses_worker_buffers_for_smaller_leaves() { + let mut buffers = LeafBuffers::default(); + buffers.prepare(0, 64, 128, 2).unwrap(); + let points = buffers.points.as_ptr(); + let dots = buffers.dots.as_ptr(); + let nearest = buffers.nearest.as_ptr(); + + buffers.prepare(1, 8, 128, 2).unwrap(); + + assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.dots.as_ptr(), dots); + assert_eq!(buffers.nearest.as_ptr(), nearest); +} + +#[test] +fn reports_shape_overflow_before_allocating() { + let mut buffers = LeafBuffers::default(); + assert!(matches!( + buffers.prepare(7, usize::MAX, 2, 1), + Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) + )); +} + +#[test] +fn rejects_an_invalid_kernel_position() { + let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + let error = add_symmetric_edges( + &[10, 20], + 1, + &[ + crate::leaf_kernel::LeafNeighbor::new(9, 1.0), + crate::leaf_kernel::LeafNeighbor::new(0, 1.0), + ], + &mut graph, + ) + .unwrap_err(); + assert!(matches!( + error, + LeafBuildError::InvalidLocalPosition { + position: 9, + points: 2 + } + )); +} + +#[test] +fn skips_duplicate_global_ids_without_self_edges() { + let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + add_symmetric_edges( + &[7, 7], + 1, + &[ + crate::leaf_kernel::LeafNeighbor::new(1, 0.0), + crate::leaf_kernel::LeafNeighbor::new(0, 0.0), + ], + &mut graph, + ) + .unwrap(); + assert!(graph.iter().all(|row| row.is_empty())); +} + +#[test] +fn poisoned_candidate_rows_return_errors() { + let candidates = DirectCandidates::new(1).unwrap(); + let _ = std::panic::catch_unwind(|| { + let _guard = candidates.rows[0].lock().unwrap(); + panic!("poison candidate row"); + }); + assert!(matches!( + candidates.add_leaf(&[0], &[diskann::graph::AdjacencyList::new()]), + Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + )); + assert!(matches!( + candidates.into_rows(), + Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + )); +} + +#[test] +fn allocation_errors_preserve_buffer_context() { + let mut values = Vec::::new(); + let source = values.try_reserve(usize::MAX).unwrap_err(); + let error = allocation_error("test", 1, source); + assert!(matches!( + error, + LeafBuildError::Allocation { + buffer: "test", + additional: 1, + .. + } + )); +} + +#[test] +fn direct_candidate_accumulator_keeps_unique_sorted_rows() { + let candidates = DirectCandidates::new(2).unwrap(); + candidates + .add_leaf( + &[0, 1], + &[ + diskann::graph::AdjacencyList::from_iter_untrusted([1, 1]), + diskann::graph::AdjacencyList::from_iter_untrusted([0]), + ], + ) + .unwrap(); + assert_eq!(rows(candidates.into_rows().unwrap()), [vec![1], vec![0]]); +} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index e91ed5fff8..b611f467b4 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,24 +3,42 @@ * Licensed under the MIT license. */ -//! Numerical kernels for PiPNN graph construction. +//! Numerical kernels for provider-independent PiPNN graph construction. //! -//! Metric modules fill portable distance buffers. Kernel modules use architecture -//! `A` to traverse those buffers. Callers select metric `M` once per graph build. +//! PiPNN assigns points to overlapping leader partitions. It computes one +//! lower-triangular all-pairs matrix per bounded leaf. It then merges selected +//! leaf neighbors into graph candidates. +//! +//! - [`partition_kernel`] ranks leader-column positions. Output width is runtime +//! fanout. A scratch vector retains ranked leaders and reuses its allocation. +//! - [`leaf_kernel`] scans each pair in the strict lower triangle one time. It +//! updates both endpoints. Each point keeps the requested number of neighbors. +//! This number cannot exceed the number of other points in the leaf. +//! - `leaf_metric` and `partition_metric` fill portable ranking buffers. +//! +//! The graph build selects concrete architecture `A` and metric `M` types once. +//! Metric computation is architecture-neutral. Kernels use `A` for SIMD ranking. -#[allow(dead_code)] +mod finalization; +mod leaf_build; mod leaf_kernel; -#[allow(dead_code)] mod leaf_metric; -#[allow(dead_code)] mod partition_kernel; -#[allow(dead_code)] mod partition_metric; -#[allow(dead_code)] +mod partitioning; mod simd; #[allow(dead_code)] mod topk; +use crate::{ + graph::{AdjacencyList, Config}, + utils::VectorRepr, + ANNError, ANNResult, +}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use rayon::ThreadPool; + pub(super) struct L2; pub(super) struct Cosine; pub(super) struct CosineNormalized; @@ -95,3 +113,179 @@ mod cosine_distance_tests { assert!(cosine_distance(dot, source_norm, target_norm).is_nan()); } } + +/// Configuration of PiPNN's partitioning and local-neighbor algorithm. +/// +/// Graph degree, pruning policy, and alpha belong to DiskANN's graph +/// configuration and are supplied separately through [`PiPNNBuildContext`]. +#[derive(Clone, Debug, PartialEq)] +pub struct PiPNNConfig { + /// Maximum number of points in a leaf. + pub c_max: usize, + /// Minimum leaf size used by global small-leaf merging. + pub c_min: usize, + /// Fraction of a cluster sampled as partition leaders. + pub p_samp: f64, + /// Number of nearest leaders retained at each overlapping partition level. + pub fanout: Vec, + /// Number of nearest neighbors selected within each leaf. + pub k: usize, + /// Number of independent partition passes over the dataset. + pub replicas: usize, +} + +impl PiPNNConfig { + fn validate(&self) -> ANNResult<()> { + if self.c_max == 0 { + return Err(config_error("c_max must be greater than zero")); + } + if self.c_min == 0 { + return Err(config_error("c_min must be greater than zero")); + } + if self.c_min > self.c_max { + return Err(config_error(format!( + "c_min ({}) must not exceed c_max ({})", + self.c_min, self.c_max + ))); + } + if !self.p_samp.is_finite() || !(0.0..=1.0).contains(&self.p_samp) || self.p_samp == 0.0 { + return Err(config_error(format!( + "p_samp ({}) must be finite and in (0, 1]", + self.p_samp + ))); + } + if self.fanout.is_empty() { + return Err(config_error("fanout must not be empty")); + } + if let Some(&fanout) = self + .fanout + .iter() + .find(|&&fanout| !(1..=partition_kernel::MAX_PARTITION_FANOUT).contains(&fanout)) + { + return Err(config_error(format!( + "fanout ({fanout}) must be in [1, {}]", + partition_kernel::MAX_PARTITION_FANOUT + ))); + } + if self.k == 0 { + return Err(config_error("k must be greater than zero")); + } + if self.replicas == 0 { + return Err(config_error("replicas must be greater than zero")); + } + Ok(()) + } +} + +/// Validated, borrowed policy and execution context for one PiPNN graph build. +#[derive(Debug)] +pub struct PiPNNBuildContext<'a> { + pub(crate) config: PiPNNConfig, + pub(crate) graph: &'a Config, + pub(crate) metric: Metric, + pub(crate) pool: &'a ThreadPool, +} + +impl<'a> PiPNNBuildContext<'a> { + /// Validate and combine PiPNN configuration with outer graph policy. + pub fn new( + config: PiPNNConfig, + graph: &'a Config, + metric: Metric, + pool: &'a ThreadPool, + ) -> ANNResult { + config.validate()?; + if !graph.alpha().is_finite() || graph.alpha() < 1.0 { + return Err(config_error(format!( + "graph alpha ({}) must be finite and at least 1", + graph.alpha() + ))); + } + if graph.prune_kind() != metric.into() { + return Err(config_error(format!( + "graph prune kind {:?} is incompatible with metric {metric:?}", + graph.prune_kind() + ))); + } + + Ok(Self { + config, + graph, + metric, + pool, + }) + } +} + +/// Build PiPNN adjacency for real rows in `data`. +/// +/// This is the core algorithm boundary. Search entry-point selection, frozen nodes, +/// providers, serialization, and index writers belong to the outer build pipelines. +pub fn build_graph( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync + 'static, +{ + context.pool.install(|| build_graph_inner(data, context)) +} + +fn build_graph_inner( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync + 'static, +{ + if data.nrows() == 0 { + return Err(ANNError::log_dimension_mismatch_error( + "PiPNN requires at least one data row".into(), + )); + } + if data.ncols() == 0 { + return Err(ANNError::log_dimension_mismatch_error( + "PiPNN requires at least one data dimension".into(), + )); + } + if data.nrows() > u32::MAX as usize { + return Err(config_error(format!( + "dataset row count ({}) exceeds the u32 graph ID limit", + data.nrows() + ))); + } + data.nrows().checked_mul(data.ncols()).ok_or_else(|| { + ANNError::log_dimension_mismatch_error(format!( + "PiPNN dataset shape {} x {} overflows usize", + data.nrows(), + data.ncols() + )) + })?; + let metric = effective_metric::(context.metric); + + let leaves = tracing::info_span!("pipnn.partition") + .in_scope(|| partitioning::partition(data, &context.config, metric))?; + let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::build_leaf_candidates(data, &leaves, context.config.k, metric) + .map_err(ANNError::opaque) + })?; + tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) +} + +fn effective_metric(metric: Metric) -> Metric { + use std::any::TypeId; + + if metric == Metric::CosineNormalized + && (TypeId::of::() == TypeId::of::() || TypeId::of::() == TypeId::of::()) + { + Metric::Cosine + } else { + metric + } +} + +#[track_caller] +fn config_error(message: impl std::fmt::Display) -> ANNError { + ANNError::log_index_config_error("PiPNN".into(), message.to_string()) +} diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs new file mode 100644 index 0000000000..40e321c83d --- /dev/null +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -0,0 +1,612 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Deterministic overlapping partition construction for PiPNN. +//! +//! The stage maps real dataset rows to bounded leaf ID lists. Numerical work +//! reuses the partition kernel and dense GEMM; scratch belongs to the Rayon +//! iterator that uses it, so no thread-local cleanup protocol is required. + +use std::collections::HashSet; + +use crate::{utils::VectorRepr, ANNError, ANNResult}; +use diskann_linalg::Transpose; +use diskann_utils::views::MatrixView; +use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; +use rand::{prelude::IndexedRandom, SeedableRng}; +use rayon::prelude::*; + +use crate::{ + partition_kernel::{nearest_leaders, PartitionTopK}, + PiPNNConfig, +}; + +// Private algorithm and batching constants live together. None are user policy. +const PARTITION_SEED: u64 = 1_000; +const LEADER_CAP: usize = 1_000; +const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 512 * 1024; +const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; +const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; +const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; +const SCATTER_STRIPE_ROWS: usize = 64 * 1024; +const MAX_PARTITION_ITERATIONS: usize = 30; + +/// A partition failure with enough context to diagnose non-progressing input. +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub(crate) enum PartitionError { + #[error("PiPNN cannot partition an empty dataset")] + EmptyDataset, + #[error("PiPNN cannot partition vectors with zero dimensions")] + EmptyDimensions, + #[error("dataset has {0} rows, which exceeds the u32 ID limit")] + TooManyPoints(usize), + #[error("{buffer} shape {rows} x {cols} overflows usize")] + ShapeOverflow { + buffer: &'static str, + rows: usize, + cols: usize, + }, + #[error( + "partition stopped after {limit} iterations with an oversized cluster of size \ + {size} at level {level}" + )] + IterationLimit { + size: usize, + level: usize, + limit: usize, + }, + #[error("partition produced an invalid leaf of size {size}; expected 1..={limit}")] + InvalidLeaf { size: usize, limit: usize }, + #[error("invalid {buffer} length: expected {expected}, got {actual}")] + InvalidBufferLength { + buffer: &'static str, + expected: usize, + actual: usize, + }, + #[error("partition worker did not publish its result")] + MissingWorkerResult, +} + +struct WorkItem { + indices: Vec, + level: usize, + seed: u64, +} + +#[derive(Default)] +struct StripeBuffers { + points: Vec, + dots: Vec, + row_scales: Vec, +} + +/// Partition every configured replica of `data` into leaves no larger than +/// `config.c_max`. The caller is responsible for installing this operation in +/// its build-owned Rayon pool. +pub(crate) fn partition( + data: MatrixView<'_, T>, + config: &PiPNNConfig, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let points = data.nrows(); + if points == 0 { + return Err(ANNError::opaque(PartitionError::EmptyDataset)); + } + if data.ncols() == 0 { + return Err(ANNError::opaque(PartitionError::EmptyDimensions)); + } + if points > u32::MAX as usize { + return Err(ANNError::opaque(PartitionError::TooManyPoints(points))); + } + + let mut leaves = Vec::new(); + for replica in 0..config.replicas { + let seed = mix_seed(PARTITION_SEED, replica as u64); + let mut replica_leaves = partition_replica(data, config, metric, seed)?; + leaves + .try_reserve(replica_leaves.len()) + .map_err(ANNError::opaque)?; + leaves.append(&mut replica_leaves); + } + validate_leaves(&leaves, config.c_max)?; + Ok(leaves) +} + +fn partition_replica( + data: MatrixView<'_, T>, + config: &PiPNNConfig, + metric: Metric, + seed: u64, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let initial_indices = point_ids(data.nrows())?; + if data.nrows() <= config.c_max { + let mut leaves = Vec::new(); + leaves.try_reserve_exact(1).map_err(ANNError::opaque)?; + leaves.push(initial_indices); + return Ok(leaves); + } + + let mut leaves = Vec::new(); + let mut work = Vec::new(); + work.try_reserve_exact(1).map_err(ANNError::opaque)?; + work.push(WorkItem { + indices: initial_indices, + level: 0, + seed, + }); + + for _ in 0..MAX_PARTITION_ITERATIONS { + if work.is_empty() { + return global_merge_small(leaves, config.c_min, config.c_max); + } + + let mut results = Vec::new(); + results + .try_reserve_exact(work.len()) + .map_err(ANNError::opaque)?; + results.resize_with(work.len(), || None); + // build_graph installs this complete private call tree into the + // caller-owned pool; the indexed fill cannot escape that pool. + #[allow(clippy::disallowed_methods)] + results + .par_iter_mut() + .zip(work.into_par_iter()) + .for_each(|(slot, item)| { + *slot = Some(partition_one_level(data, config, metric, item)); + }); + + let mut next_work = Vec::new(); + for result in results { + let (mut pending, mut finished) = + result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??; + next_work + .try_reserve(pending.len()) + .map_err(ANNError::opaque)?; + leaves + .try_reserve(finished.len()) + .map_err(ANNError::opaque)?; + next_work.append(&mut pending); + leaves.append(&mut finished); + } + work = next_work; + } + + if work.is_empty() { + return global_merge_small(leaves, config.c_min, config.c_max); + } + let mut largest = &work[0]; + for item in &work[1..] { + if item.indices.len() > largest.indices.len() { + largest = item; + } + } + Err(ANNError::opaque(PartitionError::IterationLimit { + size: largest.indices.len(), + level: largest.level, + limit: MAX_PARTITION_ITERATIONS, + })) +} + +fn partition_one_level( + data: MatrixView<'_, T>, + config: &PiPNNConfig, + metric: Metric, + item: WorkItem, +) -> ANNResult<(Vec, Vec>)> +where + T: VectorRepr + Send + Sync, +{ + let points = item.indices.len(); + let fanout = config.fanout.get(item.level).copied().unwrap_or(1); + let leaders = sample_leaders( + &item.indices, + config.p_samp, + mix_seed(item.seed, points as u64), + )?; + let clusters = assign_to_leaders(data, &item.indices, &leaders, fanout, metric)?; + + let mut pending = Vec::new(); + let mut finished = Vec::new(); + pending + .try_reserve(clusters.len()) + .map_err(ANNError::opaque)?; + finished + .try_reserve(clusters.len()) + .map_err(ANNError::opaque)?; + let child_seed = mix_seed(item.seed, points as u64); + for cluster in clusters { + if cluster.is_empty() { + continue; + } + if cluster.len() <= config.c_max { + finished.push(cluster); + } else { + pending.push(WorkItem { + indices: cluster, + level: item.level + 1, + seed: child_seed, + }); + } + } + Ok((pending, finished)) +} + +fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResult> { + let count = sample_num_leaders(points.len(), sampling_fraction); + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let mut leaders = Vec::new(); + leaders.try_reserve_exact(count).map_err(ANNError::opaque)?; + leaders.extend(points.choose_multiple(&mut rng, count).copied()); + Ok(leaders) +} + +fn sample_num_leaders(points: usize, sampling_fraction: f64) -> usize { + ((points as f64 * sampling_fraction).ceil() as usize) + .clamp(2, LEADER_CAP) + .min(points) +} + +// A single LCG mixer derives both replica and recursive seeds. Wrapping makes +// the mapping stable across debug/release builds and supported platforms. +fn mix_seed(seed: u64, salt: u64) -> u64 { + seed.wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(salt) +} + +fn assign_to_leaders( + data: MatrixView<'_, T>, + points: &[u32], + leaders: &[u32], + fanout: usize, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let dimensions = data.ncols(); + let leader_values_len = checked_area("leader data", leaders.len(), dimensions)?; + let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; + gather_rows(data, leaders, &mut leader_values)?; + + let mut leader_scales = if matches!(metric, Metric::L2 | Metric::Cosine) { + filled_vec(leaders.len(), 0.0f32)? + } else { + Vec::new() + }; + for (scale, row) in leader_scales + .iter_mut() + .zip(leader_values.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(row); + if metric == Metric::Cosine { + *scale = scale.sqrt(); + } + } + + let fanout = fanout.min(leaders.len()); + let assignment_len = checked_area("partition assignments", points.len(), fanout)?; + let mut assignments = filled_vec(assignment_len, 0u32)?; + let stripe_rows = assignment_stripe_rows(leaders.len()); + let assignment_stripe = checked_area("assignment stripe", stripe_rows, fanout)?; + + // build_graph pins this terminal operation to the caller-owned pool. + #[allow(clippy::disallowed_methods)] + assignments + .par_chunks_mut(assignment_stripe) + .enumerate() + .try_for_each_init(StripeBuffers::default, |buffers, (stripe, output)| { + let first = stripe * stripe_rows; + let rows = output.len() / fanout; + let point_values_len = checked_area("point stripe", rows, dimensions)?; + let dots_len = checked_area("dot-product stripe", rows, leaders.len())?; + resize_fallible(&mut buffers.points, point_values_len, 0.0)?; + resize_fallible(&mut buffers.dots, dots_len, 0.0)?; + gather_rows(data, &points[first..first + rows], &mut buffers.points)?; + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + rows, + leaders.len(), + dimensions, + 1.0, + &buffers.points, + &leader_values, + None, + &mut buffers.dots, + ) + .map_err(ANNError::opaque)?; + + let row_scales = if metric == Metric::Cosine { + resize_fallible(&mut buffers.row_scales, rows, 0.0)?; + for (scale, row) in buffers + .row_scales + .iter_mut() + .zip(buffers.points.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(row); + } + buffers.row_scales.as_slice() + } else { + &[] + }; + nearest_leaders( + PartitionTopK { + dots: &buffers.dots, + rows, + leaders: leaders.len(), + row_scales, + leader_scales: &leader_scales, + metric, + }, + fanout, + output, + ) + .map_err(ANNError::opaque) + })?; + + scatter_assignments(points, &assignments, fanout, leaders.len()) +} + +fn gather_rows(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> +where + T: VectorRepr, +{ + let expected = checked_area("gather output", indices.len(), data.ncols())?; + if output.len() != expected { + return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "gather output", + expected, + actual: output.len(), + })); + } + for (&index, row) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { + T::as_f32_into(data.row(index as usize), row).map_err(Into::::into)?; + } + Ok(()) +} + +fn scatter_assignments( + points: &[u32], + assignments: &[u32], + fanout: usize, + leaders: usize, +) -> ANNResult>> { + if points.len() < PARALLEL_SCATTER_MIN_POINTS { + return scatter_serial(points, assignments, fanout, leaders); + } + + let assignment_stripe = checked_area("scatter assignment stripe", SCATTER_STRIPE_ROWS, fanout)?; + let stripes = points.len().div_ceil(SCATTER_STRIPE_ROWS); + let mut partials = Vec::new(); + partials + .try_reserve_exact(stripes) + .map_err(ANNError::opaque)?; + partials.resize_with(stripes, || None); + // See the pool invariant at the other partition terminal operations. + #[allow(clippy::disallowed_methods)] + partials + .par_iter_mut() + .zip( + points + .par_chunks(SCATTER_STRIPE_ROWS) + .zip(assignments.par_chunks(assignment_stripe)), + ) + .for_each(|(slot, (points, assignments))| { + *slot = Some(scatter_serial(points, assignments, fanout, leaders)); + }); + + let mut locals = Vec::new(); + locals + .try_reserve_exact(stripes) + .map_err(ANNError::opaque)?; + for result in partials { + locals.push(result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??); + } + + let mut sizes = filled_vec(leaders, 0usize)?; + for local in &locals { + for (size, cluster) in sizes.iter_mut().zip(local) { + *size = size.checked_add(cluster.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "cluster size", + rows: *size, + cols: cluster.len(), + }) + })?; + } + } + + let mut clusters = clusters_with_capacities(&sizes)?; + for local in locals { + for (cluster, part) in clusters.iter_mut().zip(local) { + debug_assert!(cluster.capacity().saturating_sub(cluster.len()) >= part.len()); + cluster.extend(part); + } + } + Ok(clusters) +} + +fn scatter_serial( + points: &[u32], + assignments: &[u32], + fanout: usize, + leaders: usize, +) -> ANNResult>> { + let mut sizes = filled_vec(leaders, 0usize)?; + for &leader in assignments { + let Some(size) = sizes.get_mut(leader as usize) else { + return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "leader assignment", + expected: leaders, + actual: leader as usize + 1, + })); + }; + *size = size.checked_add(1).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "cluster size", + rows: *size, + cols: 1, + }) + })?; + } + let mut clusters = clusters_with_capacities(&sizes)?; + for (&point, row) in points.iter().zip(assignments.chunks_exact(fanout)) { + for &leader in row { + clusters[leader as usize].push(point); + } + } + Ok(clusters) +} + +fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { + let mut clusters = Vec::new(); + clusters + .try_reserve_exact(sizes.len()) + .map_err(ANNError::opaque)?; + for &size in sizes { + let mut cluster = Vec::new(); + cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + clusters.push(cluster); + } + Ok(clusters) +} + +fn global_merge_small( + leaves: Vec>, + c_min: usize, + c_max: usize, +) -> ANNResult>> { + let mut merged = Vec::new(); + let mut small_leaves = Vec::new(); + merged.try_reserve(leaves.len()).map_err(ANNError::opaque)?; + small_leaves + .try_reserve(leaves.len()) + .map_err(ANNError::opaque)?; + for leaf in leaves { + if leaf.len() >= c_min { + merged.push(leaf); + } else { + small_leaves.push(leaf); + } + } + if small_leaves.is_empty() { + return Ok(merged); + } + + let mut small = HashSet::new(); + small.try_reserve(c_max).map_err(ANNError::opaque)?; + + for leaf in small_leaves { + let combined = small.len().checked_add(leaf.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "small-leaf merge", + rows: small.len(), + cols: leaf.len(), + }) + })?; + if combined > c_max { + merged.push(drain_sorted(&mut small)?); + } + small.try_reserve(leaf.len()).map_err(ANNError::opaque)?; + small.extend(leaf); + if small.len() >= c_min { + merged.push(drain_sorted(&mut small)?); + } + } + + if !small.is_empty() { + let mut remainder = drain_sorted(&mut small)?; + if remainder.len() < c_min { + if let Some(last) = merged.last_mut() { + remainder.retain(|id| !last.contains(id)); + let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "small-leaf tail merge", + rows: last.len(), + cols: remainder.len(), + }) + })?; + if combined <= c_max { + last.try_reserve(remainder.len()) + .map_err(ANNError::opaque)?; + last.append(&mut remainder); + last.sort_unstable(); + } + } + } + if !remainder.is_empty() { + merged.push(remainder); + } + } + + validate_leaves(&merged, c_max)?; + Ok(merged) +} + +fn drain_sorted(set: &mut HashSet) -> ANNResult> { + let mut values = Vec::new(); + values + .try_reserve_exact(set.len()) + .map_err(ANNError::opaque)?; + values.extend(set.drain()); + values.sort_unstable(); + Ok(values) +} + +fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { + if let Some(leaf) = leaves + .iter() + .find(|leaf| leaf.is_empty() || leaf.len() > c_max) + { + return Err(ANNError::opaque(PartitionError::InvalidLeaf { + size: leaf.len(), + limit: c_max, + })); + } + Ok(()) +} + +fn point_ids(points: usize) -> ANNResult> { + let mut ids = Vec::new(); + ids.try_reserve_exact(points).map_err(ANNError::opaque)?; + ids.extend(0..points as u32); + Ok(ids) +} + +fn filled_vec(len: usize, value: T) -> ANNResult> { + let mut values = Vec::new(); + values.try_reserve_exact(len).map_err(ANNError::opaque)?; + values.resize(len, value); + Ok(values) +} + +fn resize_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { + if len > values.len() { + values + .try_reserve(len - values.len()) + .map_err(ANNError::opaque)?; + } + values.resize(len, value); + Ok(()) +} + +fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { + rows.checked_mul(cols) + .ok_or_else(|| ANNError::opaque(PartitionError::ShapeOverflow { buffer, rows, cols })) +} + +fn assignment_stripe_rows(leaders: usize) -> usize { + (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::())) + .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) +} + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs new file mode 100644 index 0000000000..30eed3a2d1 --- /dev/null +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -0,0 +1,277 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_utils::views::{Matrix, MatrixView}; +use diskann_vector::{distance::Metric, Half}; + +use super::*; + +fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { + PiPNNConfig { + c_max, + c_min, + p_samp: 0.25, + fanout, + k: 2, + replicas, + } +} + +fn clustered_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let row = position / dimensions; + let column = position % dimensions; + position += 1; + (row / 8) as f32 * 10.0 + column as f32 * 0.01 + row as f32 * 0.001 + } + }), + points, + dimensions, + ) +} + +fn directional_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let row = position / dimensions; + let column = position % dimensions; + position += 1; + let angle = std::f32::consts::TAU * row as f32 / points as f32; + match column { + 0 => angle.cos(), + 1 => angle.sin(), + _ => 0.0, + } + } + }), + points, + dimensions, + ) +} + +fn sorted_memberships(leaves: &[Vec]) -> Vec> { + let mut memberships: Vec> = leaves + .iter() + .map(|leaf| { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids + }) + .collect(); + memberships.sort(); + memberships +} + +fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { + assert!(leaves + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max)); + let mut counts = vec![0usize; points]; + for leaf in leaves { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids.len(), leaf.len(), "duplicate ID inside a leaf"); + for &id in leaf { + assert!((id as usize) < points); + counts[id as usize] += 1; + } + } + assert!(counts.iter().all(|&count| count >= replicas)); +} + +#[test] +fn returns_one_leaf_at_and_below_c_max() { + for points in [7, 8] { + let data = clustered_data(points, 3); + let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + assert_eq!(leaves, vec![(0..points as u32).collect::>()]); + } +} + +#[test] +fn partition_is_fixed_seed_deterministic_and_bounded() { + let data = clustered_data(96, 8); + let config = config(4, 16, vec![3, 2], 2); + + let first = partition(data.as_view(), &config, Metric::L2).unwrap(); + let second = partition(data.as_view(), &config, Metric::L2).unwrap(); + + assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); + assert_valid_partition(&first, 96, 16, 2); + assert!(first.iter().map(Vec::len).sum::() > 96 * 2); +} + +#[test] +fn recursion_after_fanout_levels_falls_back_to_one() { + let data = clustered_data(80, 4); + let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + + assert_valid_partition(&leaves, 80, 8, 1); +} + +#[test] +fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { + let data = Matrix::new(1.0f32, 24, 4); + let error = partition(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = error.downcast::().unwrap(); + + assert!(matches!( + error, + PartitionError::IterationLimit { + size: 24, + limit: MAX_PARTITION_ITERATIONS, + .. + } + )); +} + +#[test] +fn global_merge_canonicalizes_small_leaf_membership() { + let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; + + let merged = global_merge_small(leaves, 4, 8).unwrap(); + + assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); +} + +#[test] +fn global_merge_never_overfills_before_reaching_c_min() { + let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; + + let merged = global_merge_small(leaves, 11, 11).unwrap(); + + assert_eq!( + merged, + vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] + ); +} + +#[test] +fn replicas_cover_every_point_once_or_more_per_replica() { + let data = directional_data(72, 5); + let leaves = partition( + data.as_view(), + &config(3, 12, vec![3, 2], 3), + Metric::CosineNormalized, + ) + .unwrap(); + + assert_valid_partition(&leaves, 72, 12, 3); +} + +#[test] +fn supported_source_types_share_partition_contract() { + let f32_data: Vec = (0..64 * 4).map(|value| (value % 23) as f32).collect(); + let half_data: Vec = f32_data.iter().copied().map(Half::from_f32).collect(); + let u8_data: Vec = f32_data.iter().map(|value| *value as u8).collect(); + let i8_data: Vec = u8_data.iter().map(|value| *value as i8 - 11).collect(); + let config = config(2, 16, vec![2, 1], 1); + + let f32_leaves = partition( + MatrixView::try_from(f32_data.as_slice(), 64, 4).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + let half_leaves = partition( + MatrixView::try_from(half_data.as_slice(), 64, 4).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + let u8_leaves = partition( + MatrixView::try_from(u8_data.as_slice(), 64, 4).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + let i8_leaves = partition( + MatrixView::try_from(i8_data.as_slice(), 64, 4).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + + for leaves in [&f32_leaves, &half_leaves, &u8_leaves, &i8_leaves] { + assert_valid_partition(leaves, 64, 16, 1); + } + assert_eq!( + sorted_memberships(&f32_leaves), + sorted_memberships(&half_leaves) + ); + assert_eq!( + sorted_memberships(&f32_leaves), + sorted_memberships(&u8_leaves) + ); + assert_eq!( + sorted_memberships(&u8_leaves), + sorted_memberships(&i8_leaves) + ); +} + +#[test] +fn all_metrics_produce_valid_partitions() { + let data = directional_data(64, 8); + let config = config(2, 20, vec![2], 1); + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let leaves = partition(data.as_view(), &config, metric).unwrap(); + assert_valid_partition(&leaves, 64, 20, 1); + } +} + +#[test] +fn leader_count_is_bounded() { + assert_eq!(sample_num_leaders(1, 1.0), 1); + assert_eq!(sample_num_leaders(10, 0.01), 2); + assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); +} + +#[test] +fn parallel_scatter_matches_serial_order() { + let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); + let assignments: Vec = points + .iter() + .flat_map(|point| [point % 7, (point + 3) % 7]) + .collect(); + + let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); + let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); + + assert_eq!(actual, expected); +} + +#[test] +fn rejects_empty_dataset() { + let data = Matrix::::new(0.0, 0, 4); + let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDataset + ); +} + +#[test] +fn rejects_zero_dimensions() { + let data = Matrix::::new(0.0, 4, 0); + let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDimensions + ); +} diff --git a/diskann/tests/build_graph.rs b/diskann/tests/build_graph.rs new file mode 100644 index 0000000000..0f823778eb --- /dev/null +++ b/diskann/tests/build_graph.rs @@ -0,0 +1,231 @@ +/* + * 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::config::{self, MaxDegree}; +use diskann::graph::pipnn::{build_graph, PiPNNBuildContext, PiPNNConfig}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use half::f16; +use rand::{rngs::StdRng, Rng, SeedableRng}; + +fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 4, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + } +} + +fn graph_config(metric: Metric, degree: usize) -> diskann::graph::Config { + config::Builder::new_with(degree, MaxDegree::same(), 8, metric.into(), |builder| { + builder.alpha(1.2); + }) + .build() + .unwrap() +} + +fn pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() +} + +fn rows(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() +} + +fn assert_graph_invariants( + graph: &[diskann::graph::AdjacencyList], + points: usize, + degree: usize, +) { + assert_eq!(graph.len(), points); + for (source, row) in graph.iter().enumerate() { + assert!(row.len() <= degree); + let mut sorted = row.to_vec(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!(sorted.len(), row.len()); + assert!(row + .iter() + .all(|&id| (id as usize) < points && id as usize != source)); + } +} + +#[test] +fn builds_a_single_leaf_graph_for_real_dataset_ids() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let graph = graph_config(Metric::L2, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); + + let graph = graph_config(Metric::L2, 1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let pruned = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&pruned, 4, 1); + for (source, neighbors) in pruned.iter().enumerate() { + assert_eq!(source.abs_diff(neighbors[0] as usize), 1); + } +} + +#[test] +fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { + let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; + let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); + let graph = graph_config(Metric::L2, 1); + let pool = pool(2); + let config = PiPNNConfig { + c_max: 5, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 4, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&actual, 5, 1); + assert!(actual.iter().all(|row| row.len() == 1)); +} + +#[test] +fn rejects_empty_dataset_dimensions_at_the_public_boundary() { + let graph = graph_config(Metric::L2, 2); + let pool = pool(1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); + let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); + + assert!(build_graph(no_rows, &context).is_err()); + assert!(build_graph(no_columns, &context).is_err()); +} + +#[test] +fn supports_every_source_type_and_metric() { + fn build(values: &[T], metric: Metric) { + let data = MatrixView::try_from(values, 6, 2).unwrap(); + let graph = graph_config(metric, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); + let actual = build_graph(data, &context).unwrap(); + assert_graph_invariants(&actual, 6, 2); + } + + let values = [ + 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, + ]; + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + build(&values, metric); + } + build(&values.map(f16::from_f32), Metric::L2); + build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); + build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); +} + +#[test] +fn integer_normalized_cosine_matches_cosine() { + fn assert_match(values: &[T]) { + let data = MatrixView::try_from(values, 8, 2).unwrap(); + let pool = pool(2); + let build = |metric| { + let graph = graph_config(metric, 2); + let config = PiPNNConfig { + c_max: 8, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + rows(build_graph(data, &context).unwrap()) + }; + assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); + } + + assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); + assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); +} + +#[test] +fn is_deterministic_for_a_fixed_pool_size() { + let data: Vec = (0..96 * 4) + .map(|value| ((value * 17 + 3) % 101) as f32) + .collect(); + let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); + let graph = graph_config(Metric::L2, 8); + let pool = pool(4); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.25, + fanout: vec![3, 2], + k: 3, + replicas: 2, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let first = build_graph(data, &context).unwrap(); + let second = build_graph(data, &context).unwrap(); + + assert_eq!(first, second); + assert_graph_invariants(&first, 96, 8); +} + +#[test] +fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { + let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); + for case in 0..24 { + let points = rng.random_range(4..=32); + let dimensions = rng.random_range(1..=8); + let c_max = rng.random_range(4..=points.min(12)); + let c_min = rng.random_range(1..=c_max); + let degree = rng.random_range(1..=points.min(8)); + let values: Vec = (0..points * dimensions) + .map(|_| rng.random_range(-10.0..10.0)) + .collect(); + let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); + let graph = graph_config(Metric::L2, degree); + let pool = pool(2); + let config = PiPNNConfig { + c_max, + c_min, + p_samp: 0.5, + fanout: vec![2], + k: rng.random_range(1..=3), + replicas: rng.random_range(1..=2), + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context) + .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); + assert_graph_invariants(&actual, points, degree); + } +} diff --git a/diskann/tests/config.rs b/diskann/tests/config.rs new file mode 100644 index 0000000000..7d46e9a2a5 --- /dev/null +++ b/diskann/tests/config.rs @@ -0,0 +1,133 @@ +/* + * 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::config::{self, MaxDegree}; +use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig}; +use diskann_vector::distance::Metric; + +fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 512, + c_min: 64, + p_samp: 0.01, + fanout: vec![10, 3], + k: 2, + replicas: 1, + } +} + +fn graph_config(metric: Metric, alpha: f32) -> diskann::graph::Config { + config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { + builder.alpha(alpha); + }) + .build() + .unwrap() +} + +fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap() +} + +#[test] +fn accepts_the_six_algorithm_parameters_with_outer_graph_policy() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); +} + +#[test] +fn rejects_each_invalid_algorithm_parameter() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + let mut cases = [ + PiPNNConfig { + c_max: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 513, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 0.0, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: -0.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 1.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: f64::NAN, + ..pipnn_config() + }, + PiPNNConfig { + fanout: Vec::new(), + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![1, 0], + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![17], + ..pipnn_config() + }, + PiPNNConfig { + k: 0, + ..pipnn_config() + }, + PiPNNConfig { + replicas: 0, + ..pipnn_config() + }, + ]; + + for config in &mut cases { + let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .expect_err("invalid PiPNN config must be rejected"); + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + } +} + +#[test] +fn rejects_graph_policy_for_a_different_metric() { + let graph = graph_config(Metric::InnerProduct, 1.2); + let pool = pool(); + + let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + assert!(error.to_string().contains("prune kind")); +} + +#[test] +fn rejects_invalid_outer_alpha() { + let pool = pool(); + for alpha in [0.9, f32::NAN, f32::INFINITY] { + let graph = graph_config(Metric::L2, alpha); + let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + assert!(error.to_string().contains("alpha")); + } +} From b0c3561eb63a8f3b7f940e98a63cf3cf701adece Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:27:00 +0000 Subject: [PATCH 104/229] pipnn: cover partition validation boundaries --- diskann/src/graph/pipnn/partitioning/tests.rs | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 30eed3a2d1..13231db797 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -275,3 +275,43 @@ fn rejects_zero_dimensions() { PartitionError::EmptyDimensions ); } + +#[test] +fn rejects_invalid_gather_output_length() { + let data = Matrix::::new(0.0, 2, 2); + let error = gather_rows(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "gather output", + expected: 4, + actual: 3, + } + ); +} + +#[test] +fn rejects_assignment_to_an_unknown_leader() { + let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "leader assignment", + expected: 2, + actual: 3, + } + ); +} + +#[test] +fn rejects_empty_and_oversized_leaves() { + for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { + let error = validate_leaves(&leaves, 2).unwrap_err(); + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidLeaf { size, limit: 2 } + ); + } +} From 7e98ca193f702bf0481a550c0e938118c6cd7e2a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:16:50 +0000 Subject: [PATCH 105/229] pipnn: harden core graph construction --- diskann/src/graph/pipnn/leaf_build.rs | 26 +++++++----------- diskann/src/graph/pipnn/mod.rs | 3 +++ diskann/src/graph/pipnn/partitioning.rs | 21 ++++++--------- diskann/src/graph/pipnn/partitioning/tests.rs | 26 ++++++++++++++++++ diskann/src/graph/pipnn/tests.rs | 27 +++++++++++++++++++ 5 files changed, 74 insertions(+), 29 deletions(-) create mode 100644 diskann/src/graph/pipnn/tests.rs diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index f5fdb22928..91aaa50f62 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -123,13 +123,11 @@ impl LeafBuffers { nearest_values, LeafNeighbor::default(), )?; - if points > self.local_graph.len() { - let additional = points - self.local_graph.len(); - self.local_graph - .try_reserve(additional) - .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; - self.local_graph.resize_with(points, AdjacencyList::new); - } + let additional = points.saturating_sub(self.local_graph.len()); + self.local_graph + .try_reserve(additional) + .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; + self.local_graph.resize_with(points, AdjacencyList::new); self.local_graph[..points] .iter_mut() .for_each(AdjacencyList::clear); @@ -315,15 +313,11 @@ fn resize( len: usize, value: T, ) -> Result<(), LeafBuildError> { - if len > values.len() { - let additional = len - values.len(); - values - .try_reserve(additional) - .map_err(|source| allocation_error(buffer, additional, source))?; - values.resize(len, value); - } else { - values.truncate(len); - } + let additional = len.saturating_sub(values.len()); + values + .try_reserve(additional) + .map_err(|source| allocation_error(buffer, additional, source))?; + values.resize(len, value); Ok(()) } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index b611f467b4..30e8ab0a5b 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -289,3 +289,6 @@ fn effective_metric(metric: Metric) -> Metric { fn config_error(message: impl std::fmt::Display) -> ANNError { ANNError::log_index_config_error("PiPNN".into(), message.to_string()) } + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 40e321c83d..d9384edfef 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -26,11 +26,11 @@ use crate::{ // Private algorithm and batching constants live together. None are user policy. const PARTITION_SEED: u64 = 1_000; const LEADER_CAP: usize = 1_000; -const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 512 * 1024; +const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; -const SCATTER_STRIPE_ROWS: usize = 64 * 1024; +const SCATTER_STRIPE_ROWS: usize = 65_536; const MAX_PARTITION_ITERATIONS: usize = 30; /// A partition failure with enough context to diagnose non-progressing input. @@ -182,12 +182,9 @@ where if work.is_empty() { return global_merge_small(leaves, config.c_min, config.c_max); } - let mut largest = &work[0]; - for item in &work[1..] { - if item.indices.len() > largest.indices.len() { - largest = item; - } - } + let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { + return global_merge_small(leaves, config.c_min, config.c_max); + }; Err(ANNError::opaque(PartitionError::IterationLimit { size: largest.indices.len(), level: largest.level, @@ -589,11 +586,9 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { } fn resize_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { - if len > values.len() { - values - .try_reserve(len - values.len()) - .map_err(ANNError::opaque)?; - } + values + .try_reserve(len.saturating_sub(values.len())) + .map_err(ANNError::opaque)?; values.resize(len, value); Ok(()) } diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 13231db797..d5a6fd81b8 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -154,6 +154,13 @@ fn global_merge_never_overfills_before_reaching_c_min() { ); } +#[test] +fn global_merge_fills_exact_capacity_before_flushing() { + let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + + assert_eq!(merged, vec![vec![0, 1, 2, 3]]); +} + #[test] fn replicas_cover_every_point_once_or_more_per_replica() { let data = directional_data(72, 5); @@ -240,6 +247,25 @@ fn leader_count_is_bounded() { assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); } +#[test] +fn replica_seed_derivation_is_stable_and_distinct() { + assert_eq!(mix_seed(PARTITION_SEED, 0), 9_518_416_997_697_480); + assert_eq!(mix_seed(PARTITION_SEED, 1), 9_518_416_997_697_481); +} + +#[test] +fn leader_assignment_handles_multiple_stripes() { + let points = 2_048; + let data: Vec = (0..points).map(|point| point as f32).collect(); + let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); + let point_ids: Vec = (0..points as u32).collect(); + + let clusters = assign_to_leaders(data, &point_ids, &[0, 2_047], 1, Metric::L2).unwrap(); + + assert_eq!(clusters[0], (0..1_024).collect::>()); + assert_eq!(clusters[1], (1_024..2_048).collect::>()); +} + #[test] fn parallel_scatter_matches_serial_order() { let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); diff --git a/diskann/src/graph/pipnn/tests.rs b/diskann/src/graph/pipnn/tests.rs new file mode 100644 index 0000000000..d5689eede6 --- /dev/null +++ b/diskann/src/graph/pipnn/tests.rs @@ -0,0 +1,27 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; +use half::f16; + +#[test] +fn integer_normalized_cosine_uses_unnormalized_cosine() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let expected = if metric == Metric::CosineNormalized { + Metric::Cosine + } else { + metric + }; + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), metric); + assert_eq!(effective_metric::(metric), metric); + } +} From 2090ca652e7dfa8be5ff9587445057a9d2f25261 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:36:30 +0000 Subject: [PATCH 106/229] docs(pipnn): describe integer cosine policy --- diskann/src/graph/pipnn/mod.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 30e8ab0a5b..974acf95a8 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -221,6 +221,8 @@ impl<'a> PiPNNBuildContext<'a> { /// /// This is the core algorithm boundary. Search entry-point selection, frozen nodes, /// providers, serialization, and index writers belong to the outer build pipelines. +/// For raw `u8` and `i8` rows, `CosineNormalized` is evaluated as `Cosine` because +/// those representations are converted to f32 scratch but are not unit-normalized. pub fn build_graph( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, From 69db8da5813a344ce1f1b1cacbc1087442411971 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:39:54 +0000 Subject: [PATCH 107/229] pipnn: expose core config validation --- diskann/src/graph/pipnn/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 974acf95a8..c9faeed431 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -135,7 +135,8 @@ pub struct PiPNNConfig { } impl PiPNNConfig { - fn validate(&self) -> ANNResult<()> { + /// Validate the algorithm-specific partition and leaf-build parameters. + pub fn validate(&self) -> ANNResult<()> { if self.c_max == 0 { return Err(config_error("c_max must be greater than zero")); } From 3dae642f45dbd6fea9b53ff879cb7132b41e75ad Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 12:23:36 +0000 Subject: [PATCH 108/229] refactor(pipnn): consume leaves during leaf build --- diskann/src/graph/pipnn/leaf_build.rs | 6 +++--- diskann/src/graph/pipnn/leaf_build/tests.rs | 8 +++++--- diskann/src/graph/pipnn/mod.rs | 2 +- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 91aaa50f62..31f1f9e8bd 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -180,7 +180,7 @@ impl DirectCandidates { #[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. pub(crate) fn build_leaf_candidates( data: MatrixView<'_, T>, - leaves: &[Vec], + leaves: Vec>, k: usize, metric: Metric, ) -> Result>, LeafBuildError> @@ -195,10 +195,10 @@ where } let candidates = DirectCandidates::new(data.nrows())?; - leaves.par_iter().enumerate().try_for_each_init( + leaves.into_par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) + build_leaf(data, leaf, &point_ids, k, metric, buffers, &candidates) }, )?; candidates.into_rows() diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index d3f63263dc..efa597dd72 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -33,7 +33,7 @@ fn build( where T: diskann::utils::VectorRepr + 'static, { - pool().install(|| build_leaf_candidates(data, leaves, k, metric)) + pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) } fn rows(graph: Vec>) -> Vec> { @@ -204,9 +204,11 @@ fn parallel_leaf_schedule_does_not_change_candidate_order() { .collect(); let pool = pool(); pool.install(|| { - let expected = build_leaf_candidates(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + let expected = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); for _ in 0..8 { - let actual = build_leaf_candidates(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + let actual = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); assert_eq!(actual, expected); } }); diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index c9faeed431..eb79c0979f 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -269,7 +269,7 @@ where let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition(data, &context.config, metric))?; let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates(data, &leaves, context.config.k, metric) + leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) .map_err(ANNError::opaque) })?; tracing::info_span!("pipnn.finalization") From d0412ec4d1c83bea4b1cf191bfec83181db9e736 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 12:23:36 +0000 Subject: [PATCH 109/229] fix(pipnn): preserve established replica seeds --- diskann/src/graph/pipnn/partitioning.rs | 11 ++++++++--- diskann/src/graph/pipnn/partitioning/tests.rs | 4 ++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index d9384edfef..6c38e80247 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -25,6 +25,7 @@ use crate::{ // Private algorithm and batching constants live together. None are user policy. const PARTITION_SEED: u64 = 1_000; +const REPLICA_SEED_STEP: u64 = 7_919; const LEADER_CAP: usize = 1_000; const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; @@ -106,7 +107,7 @@ where let mut leaves = Vec::new(); for replica in 0..config.replicas { - let seed = mix_seed(PARTITION_SEED, replica as u64); + let seed = replica_seed(replica); let mut replica_leaves = partition_replica(data, config, metric, seed)?; leaves .try_reserve(replica_leaves.len()) @@ -251,8 +252,12 @@ fn sample_num_leaders(points: usize, sampling_fraction: f64) -> usize { .min(points) } -// A single LCG mixer derives both replica and recursive seeds. Wrapping makes -// the mapping stable across debug/release builds and supported platforms. +fn replica_seed(replica: usize) -> u64 { + PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) +} + +// A single LCG mixer derives recursive seeds. Wrapping makes the mapping stable +// across debug/release builds and supported platforms. fn mix_seed(seed: u64, salt: u64) -> u64 { seed.wrapping_mul(6_364_136_223_846_793_005) .wrapping_add(salt) diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index d5a6fd81b8..56ec5f8f2e 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -249,8 +249,8 @@ fn leader_count_is_bounded() { #[test] fn replica_seed_derivation_is_stable_and_distinct() { - assert_eq!(mix_seed(PARTITION_SEED, 0), 9_518_416_997_697_480); - assert_eq!(mix_seed(PARTITION_SEED, 1), 9_518_416_997_697_481); + assert_eq!(replica_seed(0), 1_000); + assert_eq!(replica_seed(1), 8_919); } #[test] From 60d25747aaad45b818ad250c0fa906dd329887c3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:25:52 +0000 Subject: [PATCH 110/229] docs(pipnn): define graph-construction boundary --- diskann-pipnn/README.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 diskann-pipnn/README.md diff --git a/diskann-pipnn/README.md b/diskann-pipnn/README.md new file mode 100644 index 0000000000..4ac1618f03 --- /dev/null +++ b/diskann-pipnn/README.md @@ -0,0 +1,19 @@ +# PiPNN graph construction + +This crate implements the graph-construction stages from [PiPNN: Pick in Partitions for Fast and Accurate ANN Graph Construction](https://arxiv.org/html/2602.21247v1). + +## Boundary + +PiPNN core consumes a dense `MatrixView`, graph policy, and a caller-owned Rayon pool, then returns adjacency lists for the dataset's real point IDs. It does not own start or frozen points, vector or neighbor providers, PQ, disk headers, serialization, or search. Those concerns remain in the outer in-memory and disk pipelines. + +The dense view is intentional: partition assignment and leaf all-pairs kernels operate over the whole source matrix. Materializing provider state inside the algorithm would couple numerical graph construction to storage lifecycle and would require a second dataset copy. Integrations should finish PiPNN scratch before allocating or populating their searchable provider. + +## Policy ownership + +- `PiPNNConfig` owns partition and leaf-selection parameters: leaf bounds, sampling fraction, fanout levels, leaf `k`, and replicas. +- DiskANN graph configuration owns metric, output degree, build-L, alpha, and prune policy. +- Candidate-merging policies are separate validated options; they must not make graph policy fields redundant or silently cap the requested degree. + +## Execution + +A build runs partitioning, leaf construction, candidate merging, then graph finalization. All parallel work executes in the supplied pool. Per-job scratch is initialized through Rayon and is released through normal ownership when its stage completes; the core has no global thread-local buffers or cleanup broadcasts. From d84b805f5c58c389e02421d0f85869a364bbf4bb Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:10:42 +0000 Subject: [PATCH 111/229] refactor(pipnn): own partition-stage configuration --- diskann/src/graph/pipnn/mod.rs | 3 +- diskann/src/graph/pipnn/partitioning.rs | 31 ++++++++++++++++--- diskann/src/graph/pipnn/partitioning/tests.rs | 31 +++++++++---------- 3 files changed, 44 insertions(+), 21 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index eb79c0979f..7f8b00ec3d 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -266,8 +266,9 @@ where })?; let metric = effective_metric::(context.metric); + let partition = partitioning::PartitionConfig::from(&context.config); let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition(data, &context.config, metric))?; + .in_scope(|| partitioning::partition(data, partition, metric))?; let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) .map_err(ANNError::opaque) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 6c38e80247..bbe4a6a92b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -34,6 +34,29 @@ const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const SCATTER_STRIPE_ROWS: usize = 65_536; const MAX_PARTITION_ITERATIONS: usize = 30; +/// Policy owned by the partition stage. Leaf-neighbor and merge settings do +/// not cross this boundary. +#[derive(Clone, Debug)] +pub(crate) struct PartitionConfig { + c_max: usize, + c_min: usize, + p_samp: f64, + fanout: Vec, + replicas: usize, +} + +impl From<&PiPNNConfig> for PartitionConfig { + fn from(config: &PiPNNConfig) -> Self { + Self { + c_max: config.c_max, + c_min: config.c_min, + p_samp: config.p_samp, + fanout: config.fanout.clone(), + replicas: config.replicas, + } + } +} + /// A partition failure with enough context to diagnose non-progressing input. #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum PartitionError { @@ -88,7 +111,7 @@ struct StripeBuffers { /// its build-owned Rayon pool. pub(crate) fn partition( data: MatrixView<'_, T>, - config: &PiPNNConfig, + config: PartitionConfig, metric: Metric, ) -> ANNResult>> where @@ -108,7 +131,7 @@ where let mut leaves = Vec::new(); for replica in 0..config.replicas { let seed = replica_seed(replica); - let mut replica_leaves = partition_replica(data, config, metric, seed)?; + let mut replica_leaves = partition_replica(data, &config, metric, seed)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::opaque)?; @@ -120,7 +143,7 @@ where fn partition_replica( data: MatrixView<'_, T>, - config: &PiPNNConfig, + config: &PartitionConfig, metric: Metric, seed: u64, ) -> ANNResult>> @@ -195,7 +218,7 @@ where fn partition_one_level( data: MatrixView<'_, T>, - config: &PiPNNConfig, + config: &PartitionConfig, metric: Metric, item: WorkItem, ) -> ANNResult<(Vec, Vec>)> diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 56ec5f8f2e..ef7b4b517e 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -8,13 +8,12 @@ use diskann_vector::{distance::Metric, Half}; use super::*; -fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { - PiPNNConfig { +fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PartitionConfig { + PartitionConfig { c_max, c_min, p_samp: 0.25, fanout, - k: 2, replicas, } } @@ -91,7 +90,7 @@ fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, repl fn returns_one_leaf_at_and_below_c_max() { for points in [7, 8] { let data = clustered_data(points, 3); - let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); assert_eq!(leaves, vec![(0..points as u32).collect::>()]); } } @@ -101,8 +100,8 @@ fn partition_is_fixed_seed_deterministic_and_bounded() { let data = clustered_data(96, 8); let config = config(4, 16, vec![3, 2], 2); - let first = partition(data.as_view(), &config, Metric::L2).unwrap(); - let second = partition(data.as_view(), &config, Metric::L2).unwrap(); + let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); + let second = partition(data.as_view(), config, Metric::L2).unwrap(); assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); assert_valid_partition(&first, 96, 16, 2); @@ -112,7 +111,7 @@ fn partition_is_fixed_seed_deterministic_and_bounded() { #[test] fn recursion_after_fanout_levels_falls_back_to_one() { let data = clustered_data(80, 4); - let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); assert_valid_partition(&leaves, 80, 8, 1); } @@ -120,7 +119,7 @@ fn recursion_after_fanout_levels_falls_back_to_one() { #[test] fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { let data = Matrix::new(1.0f32, 24, 4); - let error = partition(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); let error = error.downcast::().unwrap(); assert!(matches!( @@ -166,7 +165,7 @@ fn replicas_cover_every_point_once_or_more_per_replica() { let data = directional_data(72, 5); let leaves = partition( data.as_view(), - &config(3, 12, vec![3, 2], 3), + config(3, 12, vec![3, 2], 3), Metric::CosineNormalized, ) .unwrap(); @@ -184,25 +183,25 @@ fn supported_source_types_share_partition_contract() { let f32_leaves = partition( MatrixView::try_from(f32_data.as_slice(), 64, 4).unwrap(), - &config, + config.clone(), Metric::L2, ) .unwrap(); let half_leaves = partition( MatrixView::try_from(half_data.as_slice(), 64, 4).unwrap(), - &config, + config.clone(), Metric::L2, ) .unwrap(); let u8_leaves = partition( MatrixView::try_from(u8_data.as_slice(), 64, 4).unwrap(), - &config, + config.clone(), Metric::L2, ) .unwrap(); let i8_leaves = partition( MatrixView::try_from(i8_data.as_slice(), 64, 4).unwrap(), - &config, + config, Metric::L2, ) .unwrap(); @@ -235,7 +234,7 @@ fn all_metrics_produce_valid_partitions() { Metric::CosineNormalized, Metric::InnerProduct, ] { - let leaves = partition(data.as_view(), &config, metric).unwrap(); + let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); assert_valid_partition(&leaves, 64, 20, 1); } } @@ -283,7 +282,7 @@ fn parallel_scatter_matches_serial_order() { #[test] fn rejects_empty_dataset() { let data = Matrix::::new(0.0, 0, 4); - let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); assert_eq!( error.downcast::().unwrap(), @@ -294,7 +293,7 @@ fn rejects_empty_dataset() { #[test] fn rejects_zero_dimensions() { let data = Matrix::::new(0.0, 4, 0); - let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); assert_eq!( error.downcast::().unwrap(), From 918eeda33b3a6cd6b772fb253e0a5272a7397017 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:43:33 +0000 Subject: [PATCH 112/229] perf(pipnn): release owned leaves after leaf stage --- diskann/src/graph/pipnn/leaf_build.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 31f1f9e8bd..a0babd3e89 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -195,10 +195,10 @@ where } let candidates = DirectCandidates::new(data.nrows())?; - leaves.into_par_iter().enumerate().try_for_each_init( + leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf(data, leaf, &point_ids, k, metric, buffers, &candidates) + build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) }, )?; candidates.into_rows() From 4b4b755df3fa5012f7f17312fc65cc6836a24cba Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:46:11 +0000 Subject: [PATCH 113/229] docs(pipnn): explain partition stage contract --- diskann/src/graph/pipnn/partitioning.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index bbe4a6a92b..64b8eaeea2 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -106,9 +106,14 @@ struct StripeBuffers { row_scales: Vec, } -/// Partition every configured replica of `data` into leaves no larger than -/// `config.c_max`. The caller is responsible for installing this operation in -/// its build-owned Rayon pool. +/// Partition every configured replica into overlapping bounded leaves. +/// +/// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped +/// to the private leader bound), assigns every point to its nearest `fanout` +/// leaders for the current level, and recurses only on oversized clusters. +/// Levels beyond `fanout.len()` retain one leader assignment. Completed small +/// leaves are merged without exceeding `c_max`; every input point must remain +/// covered once per replica. The caller installs the operation in its pool. pub(crate) fn partition( data: MatrixView<'_, T>, config: PartitionConfig, From 8fb2a1275263e66d3eec99c235a0de081238c5f3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 30 Jul 2026 07:53:25 +0000 Subject: [PATCH 114/229] perf(pipnn): reduce partition and leaf overhead --- diskann/src/graph/pipnn/leaf_build.rs | 67 ++++++++++++++----- diskann/src/graph/pipnn/leaf_build/tests.rs | 7 ++ diskann/src/graph/pipnn/partitioning.rs | 39 +++++++---- diskann/src/graph/pipnn/partitioning/tests.rs | 7 ++ 4 files changed, 88 insertions(+), 32 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index a0babd3e89..6d8a6ffa3a 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -115,14 +115,18 @@ impl LeafBuffers { columns: actual_k, })?; - resize("leaf points", &mut self.points, point_values, 0.0)?; - resize("leaf dot products", &mut self.dots, dot_values, 0.0)?; - resize( + grow("leaf points", &mut self.points, point_values, 0.0)?; + grow("leaf dot products", &mut self.dots, dot_values, 0.0)?; + grow( "leaf nearest neighbors", &mut self.nearest, nearest_values, LeafNeighbor::default(), )?; + Ok(actual_k) + } + + fn prepare_local_graph(&mut self, points: usize) -> Result<(), LeafBuildError> { let additional = points.saturating_sub(self.local_graph.len()); self.local_graph .try_reserve(additional) @@ -131,7 +135,7 @@ impl LeafBuffers { self.local_graph[..points] .iter_mut() .for_each(AdjacencyList::clear); - Ok(actual_k) + Ok(()) } } @@ -219,11 +223,6 @@ where if point_ids.is_empty() { return Err(LeafBuildError::EmptyLeaf { leaf }); } - buffers.seen_ids.clear(); - buffers - .seen_ids - .try_reserve(point_ids.len()) - .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; for &point in point_ids { if point as usize >= data.nrows() { return Err(LeafBuildError::InvalidPointId { @@ -232,8 +231,24 @@ where points: data.nrows(), }); } - if !buffers.seen_ids.insert(point) { - return Err(LeafBuildError::DuplicatePointId { leaf, point }); + } + if point_ids.is_sorted() { + if let Some(pair) = point_ids.windows(2).find(|pair| pair[0] == pair[1]) { + return Err(LeafBuildError::DuplicatePointId { + leaf, + point: pair[0], + }); + } + } else { + buffers.seen_ids.clear(); + buffers + .seen_ids + .try_reserve(point_ids.len()) + .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; + for &point in point_ids { + if !buffers.seen_ids.insert(point) { + return Err(LeafBuildError::DuplicatePointId { leaf, point }); + } } } let actual_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), k)?; @@ -241,11 +256,14 @@ where return Ok(()); } + let point_values = point_ids.len() * data.ncols(); + let dot_values = point_ids.len() * point_ids.len(); + let nearest_values = point_ids.len() * actual_k; + for (&point, output) in point_ids .iter() - .zip(buffers.points.chunks_exact_mut(data.ncols())) + .zip(buffers.points[..point_values].chunks_exact_mut(data.ncols())) { - // Point IDs were validated before the zero-k/singleton fast path. let row = data.row(point as usize); T::as_f32_into(row, output).map_err(|source| LeafBuildError::Conversion { leaf, @@ -255,28 +273,29 @@ where } diskann_linalg::sgemm_aat_lower( - &buffers.points, + &buffers.points[..point_values], point_ids.len(), data.ncols(), - &mut buffers.dots, + &mut buffers.dots[..dot_values], ) .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; nearest_leaf_neighbors( LeafTopK { - dots: &buffers.dots, + dots: &buffers.dots[..dot_values], points: point_ids.len(), metric, }, k, - &mut buffers.nearest, + &mut buffers.nearest[..nearest_values], &mut buffers.top_k, ) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + buffers.prepare_local_graph(point_ids.len())?; add_symmetric_edges( point_ids, actual_k, - &buffers.nearest, + &buffers.nearest[..nearest_values], &mut buffers.local_graph[..point_ids.len()], )?; candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) @@ -307,6 +326,18 @@ fn add_symmetric_edges( Ok(()) } +fn grow( + buffer: &'static str, + values: &mut Vec, + len: usize, + value: T, +) -> Result<(), LeafBuildError> { + if values.len() < len { + resize(buffer, values, len, value)?; + } + Ok(()) +} + fn resize( buffer: &'static str, values: &mut Vec, diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index efa597dd72..8f91ea559c 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -246,6 +246,10 @@ fn rejects_invalid_shape_inputs_without_panicking() { build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![1, 0, 1]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 1 }) + )); } #[test] @@ -275,6 +279,9 @@ fn reuses_worker_buffers_for_smaller_leaves() { assert_eq!(buffers.points.as_ptr(), points); assert_eq!(buffers.dots.as_ptr(), dots); assert_eq!(buffers.nearest.as_ptr(), nearest); + assert_eq!(buffers.points.len(), 64 * 128); + assert_eq!(buffers.dots.len(), 64 * 64); + assert_eq!(buffers.nearest.len(), 64 * 2); } #[test] diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 64b8eaeea2..122261dc91 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -31,7 +31,6 @@ const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; -const SCATTER_STRIPE_ROWS: usize = 65_536; const MAX_PARTITION_ITERATIONS: usize = 30; /// Policy owned by the partition stage. Leaf-neighbor and merge settings do @@ -413,8 +412,9 @@ fn scatter_assignments( return scatter_serial(points, assignments, fanout, leaders); } - let assignment_stripe = checked_area("scatter assignment stripe", SCATTER_STRIPE_ROWS, fanout)?; - let stripes = points.len().div_ceil(SCATTER_STRIPE_ROWS); + let stripe_rows = points.len().div_ceil(rayon::current_num_threads().max(1)); + let assignment_stripe = checked_area("scatter assignment stripe", stripe_rows, fanout)?; + let stripes = points.len().div_ceil(stripe_rows); let mut partials = Vec::new(); partials .try_reserve_exact(stripes) @@ -426,7 +426,7 @@ fn scatter_assignments( .par_iter_mut() .zip( points - .par_chunks(SCATTER_STRIPE_ROWS) + .par_chunks(stripe_rows) .zip(assignments.par_chunks(assignment_stripe)), ) .for_each(|(slot, (points, assignments))| { @@ -454,14 +454,20 @@ fn scatter_assignments( } } - let mut clusters = clusters_with_capacities(&sizes)?; - for local in locals { - for (cluster, part) in clusters.iter_mut().zip(local) { - debug_assert!(cluster.capacity().saturating_sub(cluster.len()) >= part.len()); - cluster.extend(part); - } - } - Ok(clusters) + // See the pool invariant at the other partition terminal operations. + #[allow(clippy::disallowed_methods)] + sizes + .into_par_iter() + .enumerate() + .map(|(leader, size)| { + let mut cluster = Vec::new(); + cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + for local in &locals { + cluster.extend_from_slice(&local[leader]); + } + Ok(cluster) + }) + .collect() } fn scatter_serial( @@ -632,8 +638,13 @@ fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult usize { - (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::())) - .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) + let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::()); + let rows = if rows.is_power_of_two() { + rows + } else { + rows.next_power_of_two() / 2 + }; + rows.clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) } #[cfg(test)] diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index ef7b4b517e..b9714c7f64 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -252,6 +252,13 @@ fn replica_seed_derivation_is_stable_and_distinct() { assert_eq!(replica_seed(1), 8_919); } +#[test] +fn assignment_stripes_use_power_of_two_row_counts() { + assert_eq!(assignment_stripe_rows(1_000), 128); + assert_eq!(assignment_stripe_rows(256), 512); + assert_eq!(assignment_stripe_rows(1), MAX_ASSIGNMENT_STRIPE_ROWS); +} + #[test] fn leader_assignment_handles_multiple_stripes() { let points = 2_048; From 2677d89226a2c7398d73186bd5898cd13ba6361c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 30 Jul 2026 13:25:27 +0000 Subject: [PATCH 115/229] perf(pipnn): reuse partition scratch across work items --- diskann/src/graph/pipnn/partitioning.rs | 192 +++++++++++++----- diskann/src/graph/pipnn/partitioning/tests.rs | 23 ++- 2 files changed, 160 insertions(+), 55 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 122261dc91..2d9c37d29e 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -6,10 +6,11 @@ //! Deterministic overlapping partition construction for PiPNN. //! //! The stage maps real dataset rows to bounded leaf ID lists. Numerical work -//! reuses the partition kernel and dense GEMM; scratch belongs to the Rayon -//! iterator that uses it, so no thread-local cleanup protocol is required. +//! reuses the partition kernel and dense GEMM. A stage-owned pool leases scratch +//! to Rayon chunks and takes it back after each chunk; computation never holds +//! the pool lock, and no thread-local cleanup protocol is required. -use std::collections::HashSet; +use std::{collections::HashSet, sync::Mutex}; use crate::{utils::VectorRepr, ANNError, ANNResult}; use diskann_linalg::Transpose; @@ -105,6 +106,28 @@ struct StripeBuffers { row_scales: Vec, } +#[derive(Default)] +struct StripeBufferPool { + available: Mutex>, +} + +impl StripeBufferPool { + fn take(&self) -> StripeBuffers { + self.available + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .pop() + .unwrap_or_default() + } + + fn put(&self, buffers: StripeBuffers) { + self.available + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(buffers); + } +} + /// Partition every configured replica into overlapping bounded leaves. /// /// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped @@ -133,9 +156,10 @@ where } let mut leaves = Vec::new(); + let stripe_buffers = StripeBufferPool::default(); for replica in 0..config.replicas { let seed = replica_seed(replica); - let mut replica_leaves = partition_replica(data, &config, metric, seed)?; + let mut replica_leaves = partition_replica(data, &config, metric, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::opaque)?; @@ -150,6 +174,7 @@ fn partition_replica( config: &PartitionConfig, metric: Metric, seed: u64, + stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where T: VectorRepr + Send + Sync, @@ -188,7 +213,13 @@ where .par_iter_mut() .zip(work.into_par_iter()) .for_each(|(slot, item)| { - *slot = Some(partition_one_level(data, config, metric, item)); + *slot = Some(partition_one_level( + data, + config, + metric, + item, + stripe_buffers, + )); }); let mut next_work = Vec::new(); @@ -225,6 +256,7 @@ fn partition_one_level( config: &PartitionConfig, metric: Metric, item: WorkItem, + stripe_buffers: &StripeBufferPool, ) -> ANNResult<(Vec, Vec>)> where T: VectorRepr + Send + Sync, @@ -236,7 +268,14 @@ where config.p_samp, mix_seed(item.seed, points as u64), )?; - let clusters = assign_to_leaders(data, &item.indices, &leaders, fanout, metric)?; + let clusters = assign_to_leaders( + data, + &item.indices, + &leaders, + fanout, + metric, + stripe_buffers, + )?; let mut pending = Vec::new(); let mut finished = Vec::new(); @@ -296,6 +335,7 @@ fn assign_to_leaders( leaders: &[u32], fanout: usize, metric: Metric, + stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where T: VectorRepr + Send + Sync, @@ -325,65 +365,109 @@ where let mut assignments = filled_vec(assignment_len, 0u32)?; let stripe_rows = assignment_stripe_rows(leaders.len()); let assignment_stripe = checked_area("assignment stripe", stripe_rows, fanout)?; + let stripes = points.len().div_ceil(stripe_rows); + let worker_stripes = stripes.div_ceil(rayon::current_num_threads().max(1)); + let worker_rows = checked_area("assignment worker", worker_stripes, stripe_rows)?; + let worker_assignment = checked_area("assignment worker", worker_rows, fanout)?; + // Each worker chunk owns one scratch value and reuses it for its stripes. // build_graph pins this terminal operation to the caller-owned pool. #[allow(clippy::disallowed_methods)] assignments - .par_chunks_mut(assignment_stripe) + .par_chunks_mut(worker_assignment) .enumerate() - .try_for_each_init(StripeBuffers::default, |buffers, (stripe, output)| { - let first = stripe * stripe_rows; - let rows = output.len() / fanout; - let point_values_len = checked_area("point stripe", rows, dimensions)?; - let dots_len = checked_area("dot-product stripe", rows, leaders.len())?; - resize_fallible(&mut buffers.points, point_values_len, 0.0)?; - resize_fallible(&mut buffers.dots, dots_len, 0.0)?; - gather_rows(data, &points[first..first + rows], &mut buffers.points)?; - diskann_linalg::sgemm( - Transpose::None, - Transpose::Ordinary, - rows, - leaders.len(), - dimensions, - 1.0, - &buffers.points, - &leader_values, - None, - &mut buffers.dots, - ) - .map_err(ANNError::opaque)?; - - let row_scales = if metric == Metric::Cosine { - resize_fallible(&mut buffers.row_scales, rows, 0.0)?; - for (scale, row) in buffers - .row_scales - .iter_mut() - .zip(buffers.points.chunks_exact(dimensions)) - { - *scale = FastL2NormSquared.evaluate(row); + .try_for_each(|(worker, worker_output)| { + let mut buffers = stripe_buffers.take(); + let result = (|| -> ANNResult<()> { + let worker_first = worker * worker_rows; + for (stripe, output) in worker_output.chunks_mut(assignment_stripe).enumerate() { + let first = worker_first + stripe * stripe_rows; + let rows = output.len() / fanout; + assign_stripe( + data, + &points[first..first + rows], + &leader_values, + &leader_scales, + metric, + fanout, + &mut buffers, + output, + )?; } - buffers.row_scales.as_slice() - } else { - &[] - }; - nearest_leaders( - PartitionTopK { - dots: &buffers.dots, - rows, - leaders: leaders.len(), - row_scales, - leader_scales: &leader_scales, - metric, - }, - fanout, - output, - ) - .map_err(ANNError::opaque) + Ok(()) + })(); + stripe_buffers.put(buffers); + result })?; scatter_assignments(points, &assignments, fanout, leaders.len()) } +#[inline] +#[allow(clippy::too_many_arguments)] +fn assign_stripe( + data: MatrixView<'_, T>, + points: &[u32], + leader_values: &[f32], + leader_scales: &[f32], + metric: Metric, + fanout: usize, + buffers: &mut StripeBuffers, + output: &mut [u32], +) -> ANNResult<()> +where + T: VectorRepr, +{ + let rows = points.len(); + let dimensions = data.ncols(); + let leaders = leader_values.len() / dimensions; + let point_values_len = checked_area("point stripe", rows, dimensions)?; + let dots_len = checked_area("dot-product stripe", rows, leaders)?; + resize_fallible(&mut buffers.points, point_values_len, 0.0)?; + resize_fallible(&mut buffers.dots, dots_len, 0.0)?; + gather_rows(data, points, &mut buffers.points)?; + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + rows, + leaders, + dimensions, + 1.0, + &buffers.points, + leader_values, + None, + &mut buffers.dots, + ) + .map_err(ANNError::opaque)?; + + let row_scales = if metric == Metric::Cosine { + resize_fallible(&mut buffers.row_scales, rows, 0.0)?; + for (scale, row) in buffers + .row_scales + .iter_mut() + .zip(buffers.points.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(row); + } + buffers.row_scales.as_slice() + } else { + &[] + }; + nearest_leaders( + PartitionTopK { + dots: &buffers.dots, + rows, + leaders, + row_scales, + leader_scales, + metric, + }, + fanout, + output, + ) + .map_err(ANNError::opaque) +} + fn gather_rows(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> where T: VectorRepr, diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index b9714c7f64..d3815d533e 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -259,6 +259,19 @@ fn assignment_stripes_use_power_of_two_row_counts() { assert_eq!(assignment_stripe_rows(1), MAX_ASSIGNMENT_STRIPE_ROWS); } +#[test] +fn stripe_buffer_pool_reuses_returned_capacity() { + let pool = StripeBufferPool::default(); + let mut buffers = pool.take(); + buffers.points.resize(16, 0.0); + let points = buffers.points.as_ptr(); + pool.put(buffers); + + let buffers = pool.take(); + assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.points.len(), 16); +} + #[test] fn leader_assignment_handles_multiple_stripes() { let points = 2_048; @@ -266,7 +279,15 @@ fn leader_assignment_handles_multiple_stripes() { let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); let point_ids: Vec = (0..points as u32).collect(); - let clusters = assign_to_leaders(data, &point_ids, &[0, 2_047], 1, Metric::L2).unwrap(); + let clusters = assign_to_leaders( + data, + &point_ids, + &[0, 2_047], + 1, + Metric::L2, + &StripeBufferPool::default(), + ) + .unwrap(); assert_eq!(clusters[0], (0..1_024).collect::>()); assert_eq!(clusters[1], (1_024..2_048).collect::>()); From caa39b6cf925503a7f93d5f2036bc00121cad3f0 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:06:21 +0000 Subject: [PATCH 116/229] docs(pipnn): document core stage invariants --- diskann/src/graph/pipnn/finalization.rs | 24 ++- diskann/src/graph/pipnn/leaf_build.rs | 32 +++- diskann/src/graph/pipnn/leaf_build/tests.rs | 61 ++++--- diskann/src/graph/pipnn/mod.rs | 152 ++++++++++++++++-- diskann/src/graph/pipnn/partitioning.rs | 20 +++ diskann/src/graph/pipnn/partitioning/tests.rs | 112 +++++++------ 6 files changed, 323 insertions(+), 78 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index dab9419cc8..5cae12978a 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -3,7 +3,17 @@ * Licensed under the MIT license. */ -//! Orders complete PiPNN candidate rows and applies shared Vamana RobustPrune. +//! Final graph-degree enforcement through the shared Vamana RobustPrune kernel. +//! +//! Candidate merging may produce more than `R` IDs for a point. This stage first +//! validates every global ID, then processes rows independently in the caller's +//! Rayon pool. Rows already within the degree bound are returned without any +//! distance work. Overfull rows are converted to source-distance candidates, +//! passed through RobustPrune, and rewritten from the selected output. +//! +//! The shared kernel owns occlusion and alpha-round semantics; this adapter owns +//! only contiguous dataset access and distance specialization for the source +//! representation. use std::convert::Infallible; @@ -29,12 +39,19 @@ pub(crate) enum FinalizationError { }, } +/// Per-Rayon-job state retained across rows. +/// +/// `prune` owns candidate/state/output buffers. `cache` stores provider lookup +/// results required by the shared kernel. Reusing both avoids per-node +/// allocations, which would otherwise dominate finalization for millions of +/// short rows. #[derive(Default)] struct Workspace { prune: prune::Scratch, cache: Vec<(f32, Option)>, } +/// Validate candidate IDs and prune only rows whose length exceeds graph degree. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, @@ -56,6 +73,8 @@ where .into_par_iter() .enumerate() .map_init(Workspace::default, |workspace, (source, mut row)| { + // Candidate accumulators already enforce uniqueness. A bounded row + // therefore satisfies the graph policy without distance evaluation. if row.len() <= degree { return Ok(row); } @@ -71,6 +90,9 @@ where distance.evaluate_similarity(source_vector, data.row(candidate as usize)), ) })); + // as_context sorts the active candidate prefix by source distance. + // The callback below is needed only for selected-to-candidate + // occlusion checks; dimension specialization stays in `distance`. let candidate_count = pool.len(); let mut context = workspace.prune.as_context(candidate_count); prune::robust_prune( diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6d8a6ffa3a..f847077ef2 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -3,7 +3,19 @@ * Licensed under the MIT license. */ -//! Leaf construction and direct candidate accumulation. +//! Leaf-local graph construction and candidate accumulation. +//! +//! Partitioning supplies leaves as global point IDs. For each leaf this module: +//! +//! 1. validates IDs and converts only those rows to reusable `f32` scratch; +//! 2. computes the lower triangle of `A · Aᵀ`; +//! 3. runs the dual-endpoint leaf top-k kernel; and +//! 4. translates leaf-local positions back to dataset IDs. +//! +//! The final step merges symmetric adjacency rows under per-point locks because +//! overlapping leaves are processed concurrently. Numeric buffers retain their +//! high-water length; every consumer therefore receives an explicit active +//! prefix rather than treating `Vec::len()` as the current leaf shape. use std::{ collections::{HashSet, TryReserveError}, @@ -74,6 +86,12 @@ pub(crate) enum LeafBuildError { PoisonedCandidateRow { point: u32 }, } +/// Scratch leased to one Rayon job and reused for successive leaves. +/// +/// The three numerical vectors retain their largest observed leaf shape. The +/// adjacency rows are prepared separately because zero-k/singleton leaves never +/// write them, and because later candidate-merging modes do not necessarily use +/// this representation. #[derive(Default)] struct LeafBuffers { points: Vec, @@ -139,6 +157,12 @@ impl LeafBuffers { } } +/// Concurrent accumulator indexed by global dataset ID. +/// +/// A point may appear in several overlapping leaves, so workers lock only the +/// destination row long enough to append one leaf's additions. Sorting and +/// duplicate removal are deferred until all leaves finish; doing either under +/// the lock would lengthen the contended section for no semantic benefit. struct DirectCandidates { rows: Vec>>, } @@ -208,6 +232,12 @@ where candidates.into_rows() } +/// Build and publish one leaf's symmetric nearest-neighbor rows. +/// +/// Validation precedes all dataset indexing. Sorted partition output takes the +/// adjacent-duplicate path, while arbitrary-order callers use `seen_ids`. The +/// active lengths computed after `prepare` must be used for every later slice, +/// because the reusable vectors may still be longer than this leaf. fn build_leaf( data: MatrixView<'_, T>, leaf: usize, diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index 8f91ea559c..0236370dee 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -148,30 +148,55 @@ fn global_id_translation_is_independent_of_leaf_order() { ); } -fn assert_source_type(data: &[T]) +fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> where T: diskann::utils::VectorRepr + 'static, { - let leaves = vec![vec![0, 1, 2, 3]]; - let graph = build(view(data, 4, 2), &leaves, 1, Metric::L2).unwrap(); - assert_eq!(rows(graph), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); + let leaves = vec![(0..points as u32).collect()]; + rows(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) +} + +fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) +where + T: diskann::utils::VectorRepr + 'static, +{ + let points = 8; + // Source dimension controls VectorRepr conversion chunking. Cover tails on + // both sides of 4-, 8-, and 16-element boundaries, then a second 16-lane + // chunk. Input integers remain exact in every tested representation. + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let row = index / dimensions; + let column = index % dimensions; + ((row * 7 + column * 3 + row * column) % 23) as u8 + }) + .collect(); + let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + assert_eq!( + source_graph(&converted, points, dimensions), + source_graph(&f32_data, points, dimensions), + "{label} dimensions={dimensions}" + ); + } +} + +#[test] +fn f16_conversion_matches_f32_across_dimension_boundaries() { + assert_source_conversion_matches_f32("f16", |value| f16::from_f32(value as f32)); +} + +#[test] +fn u8_conversion_matches_f32_across_dimension_boundaries() { + assert_source_conversion_matches_f32("u8", |value| value); } #[test] -fn gathers_every_supported_source_type_without_full_dataset_conversion() { - assert_source_type(&[0.0_f32, 0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0]); - assert_source_type(&[0_i8, 0, 1, 0, 2, 0, 3, 0]); - assert_source_type(&[0_u8, 0, 1, 0, 2, 0, 3, 0]); - assert_source_type(&[ - f16::from_f32(0.0), - f16::from_f32(0.0), - f16::from_f32(1.0), - f16::from_f32(0.0), - f16::from_f32(2.0), - f16::from_f32(0.0), - f16::from_f32(3.0), - f16::from_f32(0.0), - ]); +fn i8_conversion_matches_f32_across_dimension_boundaries() { + // Applying the same translation to every coordinate preserves L2 pair + // ordering while exercising signed conversion. + assert_source_conversion_matches_f32("i8", |value| value as i8 - 11); } #[test] diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 7f8b00ec3d..5f01cbfe7f 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,22 +3,144 @@ * Licensed under the MIT license. */ -//! Numerical kernels for provider-independent PiPNN graph construction. +//! Provider-independent PiPNN graph construction. //! -//! PiPNN assigns points to overlapping leader partitions. It computes one -//! lower-triangular all-pairs matrix per bounded leaf. It then merges selected -//! leaf neighbors into graph candidates. +//! PiPNN means **Pick-in-Partitions Nearest Neighbors**. It 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. This crate constructs that adjacency; it does not execute queries. //! -//! - [`partition_kernel`] ranks leader-column positions. Output width is runtime -//! fanout. A scratch vector retains ranked leaders and reuses its allocation. -//! - [`leaf_kernel`] scans each pair in the strict lower triangle one time. It -//! updates both endpoints. Each point keeps the requested number of neighbors. -//! This number cannot exceed the number of other points in the leaf. -//! - `leaf_metric` and `partition_metric` fill portable ranking buffers. +//! 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: //! -//! The graph build selects concrete architecture `A` and metric `M` types once. -//! Metric computation is architecture-neutral. Kernels use `A` for SIMD ranking. +//! 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 prune.** Candidate edges from overlapping leaves are combined +//! into one unique list per source. Vamana RobustPrune then 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 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 crate owns overlapping partition generation, leaf-local nearest-neighbor +//! construction, candidate merging, and graph-degree finalization. The caller +//! supplies a contiguous dataset view, DiskANN graph policy, and the Rayon pool. +//! Providers, start/frozen points, quantization, persistence, and search remain +//! outside this algorithm boundary. +//! +//! 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. +//! +//! Both handles are prepared once per build metric and reused across stripes or +//! leaves. Each output view supplies call-specific fanout or neighbor width. +//! Preparation selects the runtime architecture and returns a direct function +//! pointer; repeated calls do not repeat ISA or metric dispatch. +//! +//! # Main modules and structures +//! +//! ## Public build API +//! +//! - [`PiPNNConfig`] holds Randomized Ball Carving, fanout, leaf size, local `k`, +//! and replication parameters. +//! - [`PiPNNBuildContext`] validates that algorithm parameters, graph pruning +//! policy, metric, and caller-owned Rayon pool agree. +//! - [`build_graph`] runs the full pipeline over a borrowed row-major dataset and +//! returns one dataset-ID adjacency list per input point. +//! +//! ## [`partition_kernel`] +//! +//! Partition callers compute point-by-leader dots with GEMM. +//! [`partition_kernel::PartitionInput`] bundles that tile with typed +//! [`partition_kernel::PartitionScales`]. A prepared +//! [`partition_kernel::PartitionKernel`] writes sorted leader-local positions to +//! caller-owned output. Fanout is 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 caller-owned output. [`leaf_kernel::leaf_neighbor_count`] derives each +//! leaf's width from point count and requested `k`. Module documentation explains +//! fixed-width selection, `process_pairs`, and stable endpoint insertion. +//! +//! ## Private pipeline stages +//! +//! - `partitioning` recursively samples leaders, invokes the partition kernel, +//! scatters points into overlapping children, and returns bounded leaves. +//! - `leaf_build` gathers each leaf, computes its Gram matrix, invokes the leaf +//! kernel, translates local positions to dataset IDs, and merges candidates. +//! - `finalization` applies shared Vamana RobustPrune to overfull candidate lists. +//! - `kernel_metric` owns norm preparation, exact ranking inputs, and numerical +//! edge cases. The graph build selects one concrete metric for both kernels. +//! +//! # Typical use +//! +//! 1. Construct [`PiPNNConfig`] and DiskANN graph [`Config`]. +//! 2. Create [`PiPNNBuildContext`] with metric and caller-owned Rayon pool. +//! 3. Call [`build_graph`] with one row-major [`MatrixView`] of dataset vectors. +//! 4. The outer index builder chooses start/frozen points and serializes returned +//! adjacency; those policies are intentionally not part of this crate. +//! +//! Stage outputs are owned values. Leaves move into candidate construction; +//! candidate lists move into finalization. Ownership releases each stage's large +//! scratch before the next outer allocation. +//! +//! # 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 itself never names instruction +//! sets; `diskann-wide` owns architecture selection. + +<<<<<<< HEAD +======= +mod kernel_metric; +mod simd; +>>>>>>> a73b9c25c (docs(pipnn): document core stage invariants) mod finalization; mod leaf_build; mod leaf_kernel; @@ -264,15 +386,21 @@ where data.ncols() )) })?; + // Integer source rows are not guaranteed unit-normalized after conversion, + // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); let partition = partitioning::PartitionConfig::from(&context.config); let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition(data, partition, metric))?; + // `leaves` is consumed here. Workers borrow individual ID rows during the + // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) .map_err(ANNError::opaque) })?; + // Finalization consumes candidate rows and reuses their allocations for the + // resulting adjacency where possible. tracing::info_span!("pipnn.finalization") .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 2d9c37d29e..248b8cb1bb 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -106,6 +106,13 @@ struct StripeBuffers { row_scales: Vec, } +/// Stage-owned high-water scratch storage for partition assignment. +/// +/// Rayon initializes `map_init` state per split job, not per physical worker. +/// Large dimensions would therefore allocate and zero the point/dot buffers many +/// times during recursive partitioning. Chunks instead take ownership of one +/// buffer here and return it afterward. The mutex protects only the short +/// pop/push operations; gather, GEMM, and top-k run without holding it. #[derive(Default)] struct StripeBufferPool { available: Mutex>, @@ -329,6 +336,13 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { .wrapping_add(salt) } +/// Assign each point to its nearest `fanout` sampled leaders. +/// +/// Leader rows are gathered once. Point rows are processed in cache-sized +/// stripes, while a worker chunk retains one leased scratch buffer across all of +/// its stripes. The flat assignment matrix preserves point order and is then +/// scattered into per-leader clusters; preserving order is required for fixed +/// seed determinism in later recursion levels. fn assign_to_leaders( data: MatrixView<'_, T>, points: &[u32], @@ -486,6 +500,12 @@ where Ok(()) } +/// Convert the flat point-major assignment matrix into leader-major clusters. +/// +/// Small inputs use one serial exact-capacity pass. Large inputs form at most +/// one partial cluster set per Rayon worker, then merge each leader independently. +/// Concatenating partials in stripe order keeps the same member order as the +/// serial implementation while removing a large serial copy tail. fn scatter_assignments( points: &[u32], assignments: &[u32], diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index d3815d533e..acaffd1c73 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -173,54 +173,60 @@ fn replicas_cover_every_point_once_or_more_per_replica() { assert_valid_partition(&leaves, 72, 12, 3); } +fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) +where + T: diskann::utils::VectorRepr + Send + Sync, +{ + let points = 64; + // Partition gathering converts source rows before GEMM. Exercise conversion + // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let row = index / dimensions; + let column = index % dimensions; + ((row * 5 + column * 7 + row * column) % 23) as u8 + }) + .collect(); + let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + let config = config(2, 16, vec![2, 1], 1); + let expected = partition( + MatrixView::try_from(&f32_data, points, dimensions).unwrap(), + config.clone(), + Metric::L2, + ) + .unwrap(); + let actual = partition( + MatrixView::try_from(&converted, points, dimensions).unwrap(), + config, + Metric::L2, + ) + .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); + + assert_valid_partition(&actual, points, 16, 1); + assert_eq!( + sorted_memberships(&actual), + sorted_memberships(&expected), + "{label} dimensions={dimensions}" + ); + } +} + #[test] -fn supported_source_types_share_partition_contract() { - let f32_data: Vec = (0..64 * 4).map(|value| (value % 23) as f32).collect(); - let half_data: Vec = f32_data.iter().copied().map(Half::from_f32).collect(); - let u8_data: Vec = f32_data.iter().map(|value| *value as u8).collect(); - let i8_data: Vec = u8_data.iter().map(|value| *value as i8 - 11).collect(); - let config = config(2, 16, vec![2, 1], 1); - - let f32_leaves = partition( - MatrixView::try_from(f32_data.as_slice(), 64, 4).unwrap(), - config.clone(), - Metric::L2, - ) - .unwrap(); - let half_leaves = partition( - MatrixView::try_from(half_data.as_slice(), 64, 4).unwrap(), - config.clone(), - Metric::L2, - ) - .unwrap(); - let u8_leaves = partition( - MatrixView::try_from(u8_data.as_slice(), 64, 4).unwrap(), - config.clone(), - Metric::L2, - ) - .unwrap(); - let i8_leaves = partition( - MatrixView::try_from(i8_data.as_slice(), 64, 4).unwrap(), - config, - Metric::L2, - ) - .unwrap(); +fn f16_partition_matches_f32_across_dimension_boundaries() { + assert_partition_conversion_matches_f32("f16", |value| Half::from_f32(value as f32)); +} - for leaves in [&f32_leaves, &half_leaves, &u8_leaves, &i8_leaves] { - assert_valid_partition(leaves, 64, 16, 1); - } - assert_eq!( - sorted_memberships(&f32_leaves), - sorted_memberships(&half_leaves) - ); - assert_eq!( - sorted_memberships(&f32_leaves), - sorted_memberships(&u8_leaves) - ); - assert_eq!( - sorted_memberships(&u8_leaves), - sorted_memberships(&i8_leaves) - ); +#[test] +fn u8_partition_matches_f32_across_dimension_boundaries() { + assert_partition_conversion_matches_f32("u8", |value| value); +} + +#[test] +fn i8_partition_matches_f32_across_dimension_boundaries() { + // The same translation in every coordinate preserves L2 ordering. + assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); } #[test] @@ -272,6 +278,20 @@ fn stripe_buffer_pool_reuses_returned_capacity() { assert_eq!(buffers.points.len(), 16); } +#[test] +fn stripe_buffer_pool_recovers_after_lock_poisoning() { + let pool = StripeBufferPool::default(); + let _ = std::panic::catch_unwind(|| { + let _guard = pool.available.lock().unwrap(); + panic!("poison scratch pool"); + }); + + let mut buffers = pool.take(); + buffers.dots.push(1.0); + pool.put(buffers); + assert_eq!(pool.take().dots, [1.0]); +} + #[test] fn leader_assignment_handles_multiple_stripes() { let points = 2_048; From 83ce7372d86c11d35fe335c10c7fcd6a520c7f11 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:37:22 +0000 Subject: [PATCH 117/229] docs(pipnn): diagram core stage ownership --- diskann/src/graph/pipnn/finalization.rs | 20 ++++++++++++++++++ diskann/src/graph/pipnn/partitioning.rs | 27 +++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 5cae12978a..7388a8ff0b 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -14,6 +14,26 @@ //! The shared kernel owns occlusion and alpha-round semantics; this adapter owns //! only contiguous dataset access and distance specialization for the source //! representation. +//! +//! ```text +//! candidate rows ──> validate row count and every global ID +//! │ +//! ┌────────────┴────────────┐ +//! v v +//! len <= R len > R +//! return row source-distance candidates +//! │ +//! v +//! shared RobustPrune +//! │ +//! v +//! rewrite same row owner +//! ``` +//! +//! | Path | Distance evaluations | Allocation behavior | +//! | --- | --- | --- | +//! | bounded row | none | move row directly to output | +//! | overfull row | source and occlusion distances | reuse Rayon-job workspace | use std::convert::Infallible; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 248b8cb1bb..84c6836ab4 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -9,6 +9,33 @@ //! reuses the partition kernel and dense GEMM. A stage-owned pool leases scratch //! to Rayon chunks and takes it back after each chunk; computation never holds //! the pool lock, and no thread-local cleanup protocol is required. +//! +//! ```text +//! replica root IDs ──> work queue +//! │ +//! v +//! sample leaders +//! │ +//! gather stripes ─> GEMM distances ─> nearest leaders +//! │ +//! v +//! stable scatter by leader +//! │ │ +//! size <= c_max oversized cluster +//! │ │ +//! completed leaf next recursion level +//! └──────────┬────────┘ +//! v +//! global small-leaf merge +//! v +//! coverage/bound validation +//! ``` +//! +//! | Recursion level | Assignment multiplicity | +//! | --- | --- | +//! | `level < fanout.len()` | `fanout[level]` nearest leaders | +//! | later levels | one nearest leader until bounded | +//! | replica boundary | independent deterministic seed | use std::{collections::HashSet, sync::Mutex}; From 749ffd89a45aa2e8b49e3a6e587f60f61ad272b0 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 02:11:33 +0000 Subject: [PATCH 118/229] fix(pipnn): preserve partition quality and scratch reuse --- diskann/src/graph/pipnn/finalization.rs | 7 +- diskann/src/graph/pipnn/mod.rs | 3 +- diskann/src/graph/pipnn/partitioning.rs | 167 +++++++++--------- diskann/src/graph/pipnn/partitioning/tests.rs | 74 +++++--- 4 files changed, 135 insertions(+), 116 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 7388a8ff0b..8acb5fda46 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -132,8 +132,11 @@ where ) .map_err(ANNError::opaque)?; - row.clear(); - row.extend_from_slice(workspace.prune.neighbors()); + // RobustPrune selects distinct candidate positions, so its output IDs + // are unique by construction. `extend_from_slice` would re-derive that + // with an O(degree^2) membership scan per row; the trusted overwrite is + // a copy and still verifies uniqueness under debug assertions. + row.overwrite_trusted(workspace.prune.neighbors()); Ok(row) }) .collect() diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 5f01cbfe7f..e0bbd72bd6 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -390,9 +390,8 @@ where // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); - let partition = partitioning::PartitionConfig::from(&context.config); let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition(data, partition, metric))?; + .in_scope(|| partitioning::partition(data, context.config.clone(), metric))?; // `leaves` is consumed here. Workers borrow individual ID rows during the // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 84c6836ab4..84659095ba 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -37,11 +37,14 @@ //! | later levels | one nearest leader until bounded | //! | replica boundary | independent deterministic seed | -use std::{collections::HashSet, sync::Mutex}; +use std::collections::HashSet; use crate::{utils::VectorRepr, ANNError, ANNResult}; use diskann_linalg::Transpose; -use diskann_utils::views::MatrixView; +use diskann_utils::{ + object_pool::{AsPooled, ObjectPool}, + views::MatrixView, +}; use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; use rand::{prelude::IndexedRandom, SeedableRng}; use rayon::prelude::*; @@ -61,29 +64,6 @@ const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const MAX_PARTITION_ITERATIONS: usize = 30; -/// Policy owned by the partition stage. Leaf-neighbor and merge settings do -/// not cross this boundary. -#[derive(Clone, Debug)] -pub(crate) struct PartitionConfig { - c_max: usize, - c_min: usize, - p_samp: f64, - fanout: Vec, - replicas: usize, -} - -impl From<&PiPNNConfig> for PartitionConfig { - fn from(config: &PiPNNConfig) -> Self { - Self { - c_max: config.c_max, - c_min: config.c_min, - p_samp: config.p_samp, - fanout: config.fanout.clone(), - replicas: config.replicas, - } - } -} - /// A partition failure with enough context to diagnose non-progressing input. #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum PartitionError { @@ -133,35 +113,24 @@ struct StripeBuffers { row_scales: Vec, } -/// Stage-owned high-water scratch storage for partition assignment. -/// -/// Rayon initializes `map_init` state per split job, not per physical worker. -/// Large dimensions would therefore allocate and zero the point/dot buffers many -/// times during recursive partitioning. Chunks instead take ownership of one -/// buffer here and return it afterward. The mutex protects only the short -/// pop/push operations; gather, GEMM, and top-k run without holding it. -#[derive(Default)] -struct StripeBufferPool { - available: Mutex>, -} - -impl StripeBufferPool { - fn take(&self) -> StripeBuffers { - self.available - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .pop() - .unwrap_or_default() +impl AsPooled<()> for StripeBuffers { + fn create(_: ()) -> Self { + Self::default() } - fn put(&self, buffers: StripeBuffers) { - self.available - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .push(buffers); + fn modify(&mut self, _: ()) { + // Scratch retains its high-water allocation across leases; active + // prefixes are established by `assign_stripe` before every read. } } +/// Stage-owned high-water scratch storage for partition assignment. +/// +/// `ObjectPool` owns the short pop/push lock and returns leases through RAII, +/// including error and panic paths. Gather, GEMM, and top-k run while only the +/// leased `StripeBuffers` is held. +type StripeBufferPool = ObjectPool; + /// Partition every configured replica into overlapping bounded leaves. /// /// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped @@ -172,7 +141,7 @@ impl StripeBufferPool { /// covered once per replica. The caller installs the operation in its pool. pub(crate) fn partition( data: MatrixView<'_, T>, - config: PartitionConfig, + config: PiPNNConfig, metric: Metric, ) -> ANNResult>> where @@ -190,7 +159,7 @@ where } let mut leaves = Vec::new(); - let stripe_buffers = StripeBufferPool::default(); + let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = partition_replica(data, &config, metric, seed, &stripe_buffers)?; @@ -205,7 +174,7 @@ where fn partition_replica( data: MatrixView<'_, T>, - config: &PartitionConfig, + config: &PiPNNConfig, metric: Metric, seed: u64, stripe_buffers: &StripeBufferPool, @@ -287,7 +256,7 @@ where fn partition_one_level( data: MatrixView<'_, T>, - config: &PartitionConfig, + config: &PiPNNConfig, metric: Metric, item: WorkItem, stripe_buffers: &StripeBufferPool, @@ -395,7 +364,11 @@ where .iter_mut() .zip(leader_values.chunks_exact(dimensions)) { - *scale = FastL2NormSquared.evaluate(row); + // Leader norms participate in the top-k ordering. Preserve the original + // scalar reduction order: reassociating this short setup pass through a + // SIMD norm changes low bits and can send near-tied points down different + // recursive partition paths. + *scale = row.iter().map(|value| value * value).sum(); if metric == Metric::Cosine { *scale = scale.sqrt(); } @@ -418,27 +391,23 @@ where .par_chunks_mut(worker_assignment) .enumerate() .try_for_each(|(worker, worker_output)| { - let mut buffers = stripe_buffers.take(); - let result = (|| -> ANNResult<()> { - let worker_first = worker * worker_rows; - for (stripe, output) in worker_output.chunks_mut(assignment_stripe).enumerate() { - let first = worker_first + stripe * stripe_rows; - let rows = output.len() / fanout; - assign_stripe( - data, - &points[first..first + rows], - &leader_values, - &leader_scales, - metric, - fanout, - &mut buffers, - output, - )?; - } - Ok(()) - })(); - stripe_buffers.put(buffers); - result + let mut buffers = stripe_buffers.get_ref(()); + let worker_first = worker * worker_rows; + for (stripe, output) in worker_output.chunks_mut(assignment_stripe).enumerate() { + let first = worker_first + stripe * stripe_rows; + let rows = output.len() / fanout; + assign_stripe( + data, + &points[first..first + rows], + &leader_values, + &leader_scales, + metric, + fanout, + &mut buffers, + output, + )?; + } + Ok::<(), ANNError>(()) })?; scatter_assignments(points, &assignments, fanout, leaders.len()) @@ -464,9 +433,23 @@ where let leaders = leader_values.len() / dimensions; let point_values_len = checked_area("point stripe", rows, dimensions)?; let dots_len = checked_area("dot-product stripe", rows, leaders)?; - resize_fallible(&mut buffers.points, point_values_len, 0.0)?; - resize_fallible(&mut buffers.dots, dots_len, 0.0)?; - gather_rows(data, points, &mut buffers.points)?; + // Scratch keeps its high-water length and every consumer receives an + // explicit active prefix. Resizing to the exact stripe shape would be + // correct but re-zeroes the buffer whenever a pooled value moves between + // work items with different leader counts: `stripe_rows` is derived from + // `leaders`, so the point buffer swings between roughly 768 KiB and 6 MiB + // and `Vec::resize` only truncates on the way down, then memsets the whole + // delta on the way back up. + grow_fallible(&mut buffers.points, point_values_len, 0.0)?; + grow_fallible(&mut buffers.dots, dots_len, 0.0)?; + let StripeBuffers { + points: point_buffer, + dots: dot_buffer, + row_scales: row_scale_buffer, + } = buffers; + let point_values = &mut point_buffer[..point_values_len]; + let dots = &mut dot_buffer[..dots_len]; + gather_rows(data, points, point_values)?; diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -474,29 +457,29 @@ where leaders, dimensions, 1.0, - &buffers.points, + point_values, leader_values, None, - &mut buffers.dots, + dots, ) .map_err(ANNError::opaque)?; let row_scales = if metric == Metric::Cosine { - resize_fallible(&mut buffers.row_scales, rows, 0.0)?; - for (scale, row) in buffers - .row_scales + grow_fallible(row_scale_buffer, rows, 0.0)?; + let row_scales = &mut row_scale_buffer[..rows]; + for (scale, row) in row_scales .iter_mut() - .zip(buffers.points.chunks_exact(dimensions)) + .zip(point_values.chunks_exact(dimensions)) { *scale = FastL2NormSquared.evaluate(row); } - buffers.row_scales.as_slice() + &*row_scales } else { &[] }; nearest_leaders( PartitionTopK { - dots: &buffers.dots, + dots, rows, leaders, row_scales, @@ -755,9 +738,17 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { Ok(values) } -fn resize_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { +/// Grow `values` to at least `len` elements, never shrinking it. +/// +/// Callers slice the active prefix themselves. Shrinking would force the next +/// larger stripe to re-zero the reclaimed tail, which is the dominant cost when +/// one pooled buffer serves work items with different stripe shapes. +fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { + if values.len() >= len { + return Ok(()); + } values - .try_reserve(len.saturating_sub(values.len())) + .try_reserve(len - values.len()) .map_err(ANNError::opaque)?; values.resize(len, value); Ok(()) diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index acaffd1c73..31f3e4bcad 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -8,12 +8,13 @@ use diskann_vector::{distance::Metric, Half}; use super::*; -fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PartitionConfig { - PartitionConfig { +fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { + PiPNNConfig { c_max, c_min, p_samp: 0.25, fanout, + k: 1, replicas, } } @@ -229,6 +230,44 @@ fn i8_partition_matches_f32_across_dimension_boundaries() { assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); } +#[test] +fn l2_leader_norms_preserve_scalar_reduction_order() { + fn next(state: &mut u64) -> f32 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 + } + + // This fixed case sits on opposite sides of the top-1 boundary depending + // on whether leader norms use the original scalar reduction or a SIMD + // reassociation. Point/leader dot products still go through the production + // GEMM; only the setup norm calculation is under test. + let dimensions = 129; + let mut state = 0x3a85_f952_c718_6e49; + let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let leader_zero: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let leader_one: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let data: Vec = leader_zero + .into_iter() + .chain(leader_one) + .chain(point) + .collect(); + let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); + + let clusters = assign_to_leaders( + data, + &[2], + &[0, 1], + 1, + Metric::L2, + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); + + assert_eq!(clusters, [vec![], vec![2]]); +} + #[test] fn all_metrics_produce_valid_partitions() { let data = directional_data(64, 8); @@ -267,31 +306,18 @@ fn assignment_stripes_use_power_of_two_row_counts() { #[test] fn stripe_buffer_pool_reuses_returned_capacity() { - let pool = StripeBufferPool::default(); - let mut buffers = pool.take(); - buffers.points.resize(16, 0.0); - let points = buffers.points.as_ptr(); - pool.put(buffers); - - let buffers = pool.take(); + let pool = StripeBufferPool::new((), 0, None); + let points = { + let mut buffers = pool.get_ref(()); + buffers.points.resize(16, 0.0); + buffers.points.as_ptr() + }; + + let buffers = pool.get_ref(()); assert_eq!(buffers.points.as_ptr(), points); assert_eq!(buffers.points.len(), 16); } -#[test] -fn stripe_buffer_pool_recovers_after_lock_poisoning() { - let pool = StripeBufferPool::default(); - let _ = std::panic::catch_unwind(|| { - let _guard = pool.available.lock().unwrap(); - panic!("poison scratch pool"); - }); - - let mut buffers = pool.take(); - buffers.dots.push(1.0); - pool.put(buffers); - assert_eq!(pool.take().dots, [1.0]); -} - #[test] fn leader_assignment_handles_multiple_stripes() { let points = 2_048; @@ -305,7 +331,7 @@ fn leader_assignment_handles_multiple_stripes() { &[0, 2_047], 1, Metric::L2, - &StripeBufferPool::default(), + &StripeBufferPool::new((), 0, None), ) .unwrap(); From ff886327d9b23a8d4800dbe35044e66f9a81b804 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:14:35 +0000 Subject: [PATCH 119/229] refactor(pipnn): reuse prepared kernels --- diskann/src/graph/pipnn/leaf_build.rs | 274 ++++++++++-------- diskann/src/graph/pipnn/leaf_build/tests.rs | 68 +++-- diskann/src/graph/pipnn/partitioning.rs | 192 ++++++------ diskann/src/graph/pipnn/partitioning/tests.rs | 37 ++- 4 files changed, 325 insertions(+), 246 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index f847077ef2..37cf5e3d14 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -7,12 +7,12 @@ //! //! Partitioning supplies leaves as global point IDs. For each leaf this module: //! -//! 1. validates IDs and converts only those rows to reusable `f32` scratch; +//! 1. validates IDs and converts only those point vectors to reusable `f32` scratch; //! 2. computes the lower triangle of `A · Aᵀ`; //! 3. runs the dual-endpoint leaf top-k kernel; and //! 4. translates leaf-local positions back to dataset IDs. //! -//! The final step merges symmetric adjacency rows under per-point locks because +//! The final step merges symmetric adjacency lists under per-point locks because //! overlapping leaves are processed concurrently. Numeric buffers retain their //! high-water length; every consumer therefore receives an explicit active //! prefix rather than treating `Vec::len()` as the current leaf shape. @@ -23,12 +23,13 @@ use std::{ }; use crate::{graph::AdjacencyList, utils::VectorRepr}; -use diskann_utils::views::MatrixView; +use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use rayon::prelude::*; use crate::leaf_kernel::{ - nearest_leaf_neighbors, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, + leaf_neighbor_count, leaf_output_len, LeafInput, LeafKernel, LeafKernelError, + LeafKernelWorkspace, LeafNeighbor, }; /// Failure while converting leaves into direct graph candidates. @@ -80,25 +81,25 @@ pub(crate) enum LeafBuildError { #[source] source: LeafKernelError, }, - #[error("leaf kernel returned local position {position} for a {points}-point leaf")] - InvalidLocalPosition { position: u32, points: usize }, - #[error("candidate row {point} is poisoned")] - PoisonedCandidateRow { point: u32 }, + #[error("leaf kernel returned local target {target} for a {points}-point leaf")] + InvalidLocalTarget { target: u32, points: usize }, + #[error("candidate list for point {point} is poisoned")] + PoisonedCandidateList { point: u32 }, } /// Scratch leased to one Rayon job and reused for successive leaves. /// /// The three numerical vectors retain their largest observed leaf shape. The -/// adjacency rows are prepared separately because zero-k/singleton leaves never +/// adjacency lists are prepared separately because zero-k/singleton leaves never /// write them, and because later candidate-merging modes do not necessarily use /// this representation. #[derive(Default)] struct LeafBuffers { - points: Vec, + point_values: Vec, dots: Vec, - nearest: Vec, - local_graph: Vec>, - top_k: LeafTopKWorkspace, + neighbors: Vec, + local_adjacency: Vec>, + kernel_workspace: LeafKernelWorkspace, seen_ids: HashSet, } @@ -106,51 +107,55 @@ impl LeafBuffers { fn prepare( &mut self, leaf: usize, - points: usize, - dimensions: usize, - k: usize, + point_count: usize, + dimension_count: usize, + requested_k: usize, ) -> Result { - let point_values = points - .checked_mul(dimensions) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: points, - columns: dimensions, - })?; - let dot_values = points - .checked_mul(points) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: points, - columns: points, - })?; - let actual_k = k.min(points.saturating_sub(1)); - let nearest_values = points - .checked_mul(actual_k) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: points, - columns: actual_k, - })?; + let point_value_count = + point_count + .checked_mul(dimension_count) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: point_count, + columns: dimension_count, + })?; + let dot_count = + point_count + .checked_mul(point_count) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: point_count, + columns: point_count, + })?; + let leaf_k = leaf_neighbor_count(point_count, requested_k) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + let neighbor_count = leaf_output_len(point_count, requested_k) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - grow("leaf points", &mut self.points, point_values, 0.0)?; - grow("leaf dot products", &mut self.dots, dot_values, 0.0)?; grow( - "leaf nearest neighbors", - &mut self.nearest, - nearest_values, + "leaf point values", + &mut self.point_values, + point_value_count, + 0.0, + )?; + grow("leaf dot products", &mut self.dots, dot_count, 0.0)?; + grow( + "leaf neighbors", + &mut self.neighbors, + neighbor_count, LeafNeighbor::default(), )?; - Ok(actual_k) + Ok(leaf_k) } - fn prepare_local_graph(&mut self, points: usize) -> Result<(), LeafBuildError> { - let additional = points.saturating_sub(self.local_graph.len()); - self.local_graph + fn prepare_local_adjacency(&mut self, point_count: usize) -> Result<(), LeafBuildError> { + let additional = point_count.saturating_sub(self.local_adjacency.len()); + self.local_adjacency .try_reserve(additional) - .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; - self.local_graph.resize_with(points, AdjacencyList::new); - self.local_graph[..points] + .map_err(|source| allocation_error("leaf adjacency lists", additional, source))?; + self.local_adjacency + .resize_with(point_count, AdjacencyList::new); + self.local_adjacency[..point_count] .iter_mut() .for_each(AdjacencyList::clear); Ok(()) @@ -160,45 +165,50 @@ impl LeafBuffers { /// Concurrent accumulator indexed by global dataset ID. /// /// A point may appear in several overlapping leaves, so workers lock only the -/// destination row long enough to append one leaf's additions. Sorting and +/// destination list long enough to append one leaf's additions. Sorting and /// duplicate removal are deferred until all leaves finish; doing either under /// the lock would lengthen the contended section for no semantic benefit. struct DirectCandidates { - rows: Vec>>, + lists: Vec>>, } impl DirectCandidates { - fn new(points: usize) -> Result { - let mut rows = Vec::new(); - rows.try_reserve_exact(points) - .map_err(|source| allocation_error("candidate rows", points, source))?; - rows.resize_with(points, || Mutex::new(AdjacencyList::new())); - Ok(Self { rows }) + fn new(point_count: usize) -> Result { + let mut lists = Vec::new(); + lists + .try_reserve_exact(point_count) + .map_err(|source| allocation_error("candidate lists", point_count, source))?; + lists.resize_with(point_count, || Mutex::new(AdjacencyList::new())); + Ok(Self { lists }) } fn add_leaf( &self, point_ids: &[u32], - local_graph: &[AdjacencyList], + local_adjacency: &[AdjacencyList], ) -> Result<(), LeafBuildError> { - for (&source, additions) in point_ids.iter().zip(local_graph) { + for (&source, additions) in point_ids.iter().zip(local_adjacency) { // Every point ID is validated before leaf-local work begins. - let row = &self.rows[source as usize]; - let mut row = row.lock().map_err(|_| poisoned_row(source))?; - row.extend_from_slice(additions); + let candidates = &self.lists[source as usize]; + let mut candidates = candidates + .lock() + .map_err(|_| poisoned_candidate_list(source))?; + candidates.extend_from_slice(additions); } Ok(()) } - fn into_rows(self) -> Result>, LeafBuildError> { + fn into_lists(self) -> Result>, LeafBuildError> { let mut output = Vec::new(); output - .try_reserve_exact(self.rows.len()) - .map_err(|source| allocation_error("candidate output", self.rows.len(), source))?; - for (point, row) in self.rows.into_iter().enumerate() { - let mut row = row.into_inner().map_err(|_| poisoned_row(point as u32))?; - row.sort(); - output.push(row); + .try_reserve_exact(self.lists.len()) + .map_err(|source| allocation_error("candidate output", self.lists.len(), source))?; + for (point, candidates) in self.lists.into_iter().enumerate() { + let mut candidates = candidates + .into_inner() + .map_err(|_| poisoned_candidate_list(point as u32))?; + candidates.sort(); + output.push(candidates); } Ok(output) } @@ -209,7 +219,7 @@ impl DirectCandidates { pub(crate) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, - k: usize, + requested_k: usize, metric: Metric, ) -> Result>, LeafBuildError> where @@ -223,16 +233,27 @@ where } let candidates = DirectCandidates::new(data.nrows())?; + // Metric and ISA are selected before Rayon workers start. Workers share + // this Copy handle; each output view supplies its leaf-specific width. + let kernel = LeafKernel::new(metric); leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) + build_leaf( + data, + leaf, + point_ids, + requested_k, + &kernel, + buffers, + &candidates, + ) }, )?; - candidates.into_rows() + candidates.into_lists() } -/// Build and publish one leaf's symmetric nearest-neighbor rows. +/// Build and publish one leaf's symmetric neighbor lists. /// /// Validation precedes all dataset indexing. Sorted partition output takes the /// adjacent-duplicate path, while arbitrary-order callers use `seen_ids`. The @@ -242,8 +263,8 @@ fn build_leaf( data: MatrixView<'_, T>, leaf: usize, point_ids: &[u32], - k: usize, - metric: Metric, + requested_k: usize, + kernel: &LeafKernel, buffers: &mut LeafBuffers, candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> @@ -281,75 +302,92 @@ where } } } - let actual_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), k)?; - if actual_k == 0 { + let leaf_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; + if leaf_k == 0 { return Ok(()); } - let point_values = point_ids.len() * data.ncols(); - let dot_values = point_ids.len() * point_ids.len(); - let nearest_values = point_ids.len() * actual_k; + let point_value_count = point_ids.len() * data.ncols(); + let dot_count = point_ids.len() * point_ids.len(); + let neighbor_value_count = leaf_output_len(point_ids.len(), requested_k) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - for (&point, output) in point_ids + for (&point, point_output) in point_ids .iter() - .zip(buffers.points[..point_values].chunks_exact_mut(data.ncols())) + .zip(buffers.point_values[..point_value_count].chunks_exact_mut(data.ncols())) { - let row = data.row(point as usize); - T::as_f32_into(row, output).map_err(|source| LeafBuildError::Conversion { - leaf, - point, - source: source.into(), + let source_values = data.row(point as usize); + T::as_f32_into(source_values, point_output).map_err(|source| { + LeafBuildError::Conversion { + leaf, + point, + source: source.into(), + } })?; } diskann_linalg::sgemm_aat_lower( - &buffers.points[..point_values], point_ids.len(), data.ncols(), - &mut buffers.dots[..dot_values], + &buffers.point_values[..point_value_count], + &mut buffers.dots[..dot_count], ) .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; - nearest_leaf_neighbors( - LeafTopK { - dots: &buffers.dots[..dot_values], - points: point_ids.len(), - metric, + let dots = MatrixView::try_from(&buffers.dots[..dot_count], point_ids.len(), point_ids.len()) + .map_err(|error| LeafBuildError::Kernel { + leaf, + source: LeafKernelError::InvalidBufferLength { + buffer: "leaf dot-product matrix", + expected: dot_count, + actual: error.into_inner().len(), }, - k, - &mut buffers.nearest[..nearest_values], - &mut buffers.top_k, + })?; + let output = MutMatrixView::try_from( + &mut buffers.neighbors[..neighbor_value_count], + point_ids.len(), + leaf_k, ) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + .map_err(|error| LeafBuildError::Kernel { + leaf, + source: LeafKernelError::InvalidBufferLength { + buffer: "output", + expected: neighbor_value_count, + actual: error.into_inner().len(), + }, + })?; + kernel + .nearest_neighbors(LeafInput { dots }, output, &mut buffers.kernel_workspace) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - buffers.prepare_local_graph(point_ids.len())?; - add_symmetric_edges( + buffers.prepare_local_adjacency(point_ids.len())?; + add_symmetric_neighbors( point_ids, - actual_k, - &buffers.nearest[..nearest_values], - &mut buffers.local_graph[..point_ids.len()], + leaf_k, + &buffers.neighbors[..neighbor_value_count], + &mut buffers.local_adjacency[..point_ids.len()], )?; - candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) + candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]) } -fn add_symmetric_edges( +fn add_symmetric_neighbors( point_ids: &[u32], - k: usize, - nearest: &[LeafNeighbor], - local_graph: &mut [AdjacencyList], + leaf_k: usize, + neighbors: &[LeafNeighbor], + local_adjacency: &mut [AdjacencyList], ) -> Result<(), LeafBuildError> { - for (source, nearest) in nearest.chunks_exact(k).enumerate() { - for neighbor in nearest { - let target = neighbor.position as usize; + for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for neighbor in source_neighbors { + let target = neighbor.target as usize; let Some(&target_id) = point_ids.get(target) else { - return Err(LeafBuildError::InvalidLocalPosition { - position: neighbor.position, + return Err(LeafBuildError::InvalidLocalTarget { + target: neighbor.target, points: point_ids.len(), }); }; let source_id = point_ids[source]; if source_id != target_id { - local_graph[source].push(target_id); - local_graph[target].push(source_id); + local_adjacency[source].push(target_id); + local_adjacency[target].push(source_id); } } } @@ -394,8 +432,8 @@ fn allocation_error( } } -fn poisoned_row(point: u32) -> LeafBuildError { - LeafBuildError::PoisonedCandidateRow { point } +fn poisoned_candidate_list(point: u32) -> LeafBuildError { + LeafBuildError::PoisonedCandidateList { point } } #[cfg(test)] diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index 0236370dee..2d4e66c2e6 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -9,8 +9,8 @@ use half::f16; use std::collections::BTreeSet; use super::{ - add_symmetric_edges, allocation_error, build_leaf_candidates, DirectCandidates, LeafBuffers, - LeafBuildError, + add_symmetric_neighbors, allocation_error, build_leaf_candidates, DirectCandidates, + LeafBuffers, LeafBuildError, }; fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { @@ -36,7 +36,7 @@ where pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) } -fn rows(graph: Vec>) -> Vec> { +fn adjacency_lists(graph: Vec>) -> Vec> { graph.into_iter().map(Vec::from).collect() } @@ -84,7 +84,7 @@ fn leaf_adjacency_matches_an_independent_all_pairs_reference() { ]; let flat: Vec<_> = points.into_iter().flatten().collect(); - let actual = rows( + let actual = adjacency_lists( build( view(&flat, points.len(), 2), &[(0..points.len() as u32).collect()], @@ -105,7 +105,7 @@ fn retains_and_deduplicates_candidates_from_overlapping_leaves() { let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); assert_eq!( - rows(graph), + adjacency_lists(graph), [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] ); } @@ -114,8 +114,8 @@ fn retains_and_deduplicates_candidates_from_overlapping_leaves() { fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { let dimensions = 9; let mut data = vec![0.0_f32; 10 * dimensions]; - for row in 1..10 { - data[row * dimensions + row - 1] = 1.0; + for source in 1..10 { + data[source * dimensions + source - 1] = 1.0; } let graph = build( @@ -143,7 +143,7 @@ fn global_id_translation_is_independent_of_leaf_order() { let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); assert_eq!( - rows(graph), + adjacency_lists(graph), [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] ); } @@ -153,7 +153,7 @@ where T: diskann::utils::VectorRepr + 'static, { let leaves = vec![(0..points as u32).collect()]; - rows(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) + adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) } fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) @@ -167,9 +167,9 @@ where for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let raw: Vec = (0..points * dimensions) .map(|index| { - let row = index / dimensions; - let column = index % dimensions; - ((row * 7 + column * 3 + row * column) % 23) as u8 + let source = index / dimensions; + let dimension = index % dimensions; + ((source * 7 + dimension * 3 + source * dimension) % 23) as u8 }) .collect(); let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); @@ -288,25 +288,28 @@ fn singleton_and_zero_k_leaves_add_no_candidates() { ) .unwrap(); let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); - assert!(singleton.iter().chain(&zero_k).all(|row| row.is_empty())); + assert!(singleton + .iter() + .chain(&zero_k) + .all(|candidates| candidates.is_empty())); } #[test] fn reuses_worker_buffers_for_smaller_leaves() { let mut buffers = LeafBuffers::default(); buffers.prepare(0, 64, 128, 2).unwrap(); - let points = buffers.points.as_ptr(); + let point_values = buffers.point_values.as_ptr(); let dots = buffers.dots.as_ptr(); - let nearest = buffers.nearest.as_ptr(); + let neighbors = buffers.neighbors.as_ptr(); buffers.prepare(1, 8, 128, 2).unwrap(); - assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.point_values.as_ptr(), point_values); assert_eq!(buffers.dots.as_ptr(), dots); - assert_eq!(buffers.nearest.as_ptr(), nearest); - assert_eq!(buffers.points.len(), 64 * 128); + assert_eq!(buffers.neighbors.as_ptr(), neighbors); + assert_eq!(buffers.point_values.len(), 64 * 128); assert_eq!(buffers.dots.len(), 64 * 64); - assert_eq!(buffers.nearest.len(), 64 * 2); + assert_eq!(buffers.neighbors.len(), 64 * 2); } #[test] @@ -321,7 +324,7 @@ fn reports_shape_overflow_before_allocating() { #[test] fn rejects_an_invalid_kernel_position() { let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; - let error = add_symmetric_edges( + let error = add_symmetric_neighbors( &[10, 20], 1, &[ @@ -333,8 +336,8 @@ fn rejects_an_invalid_kernel_position() { .unwrap_err(); assert!(matches!( error, - LeafBuildError::InvalidLocalPosition { - position: 9, + LeafBuildError::InvalidLocalTarget { + target: 9, points: 2 } )); @@ -343,7 +346,7 @@ fn rejects_an_invalid_kernel_position() { #[test] fn skips_duplicate_global_ids_without_self_edges() { let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; - add_symmetric_edges( + add_symmetric_neighbors( &[7, 7], 1, &[ @@ -353,23 +356,23 @@ fn skips_duplicate_global_ids_without_self_edges() { &mut graph, ) .unwrap(); - assert!(graph.iter().all(|row| row.is_empty())); + assert!(graph.iter().all(|neighbors| neighbors.is_empty())); } #[test] -fn poisoned_candidate_rows_return_errors() { +fn poisoned_candidate_lists_return_errors() { let candidates = DirectCandidates::new(1).unwrap(); let _ = std::panic::catch_unwind(|| { - let _guard = candidates.rows[0].lock().unwrap(); - panic!("poison candidate row"); + let _guard = candidates.lists[0].lock().unwrap(); + panic!("poison candidate list"); }); assert!(matches!( candidates.add_leaf(&[0], &[diskann::graph::AdjacencyList::new()]), - Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + Err(LeafBuildError::PoisonedCandidateList { point: 0 }) )); assert!(matches!( - candidates.into_rows(), - Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + candidates.into_lists(), + Err(LeafBuildError::PoisonedCandidateList { point: 0 }) )); } @@ -400,5 +403,8 @@ fn direct_candidate_accumulator_keeps_unique_sorted_rows() { ], ) .unwrap(); - assert_eq!(rows(candidates.into_rows().unwrap()), [vec![1], vec![0]]); + assert_eq!( + adjacency_lists(candidates.into_lists().unwrap()), + [vec![1], vec![0]] + ); } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 84659095ba..88867ce108 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -5,7 +5,7 @@ //! Deterministic overlapping partition construction for PiPNN. //! -//! The stage maps real dataset rows to bounded leaf ID lists. Numerical work +//! The stage maps real dataset points to bounded leaf ID lists. Numerical work //! reuses the partition kernel and dense GEMM. A stage-owned pool leases scratch //! to Rayon chunks and takes it back after each chunk; computation never holds //! the pool lock, and no thread-local cleanup protocol is required. @@ -43,14 +43,14 @@ use crate::{utils::VectorRepr, ANNError, ANNResult}; use diskann_linalg::Transpose; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, - views::MatrixView, + views::{MatrixView, MutMatrixView}, }; use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; use rand::{prelude::IndexedRandom, SeedableRng}; use rayon::prelude::*; use crate::{ - partition_kernel::{nearest_leaders, PartitionTopK}, + partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, PiPNNConfig, }; @@ -59,8 +59,8 @@ const PARTITION_SEED: u64 = 1_000; const REPLICA_SEED_STEP: u64 = 7_919; const LEADER_CAP: usize = 1_000; const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; -const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; -const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; +const MIN_ASSIGNMENT_STRIPE_POINTS: usize = 32; +const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const MAX_PARTITION_ITERATIONS: usize = 30; @@ -71,7 +71,7 @@ pub(crate) enum PartitionError { EmptyDataset, #[error("PiPNN cannot partition vectors with zero dimensions")] EmptyDimensions, - #[error("dataset has {0} rows, which exceeds the u32 ID limit")] + #[error("dataset has {0} points, which exceeds the u32 ID limit")] TooManyPoints(usize), #[error("{buffer} shape {rows} x {cols} overflows usize")] ShapeOverflow { @@ -110,7 +110,7 @@ struct WorkItem { struct StripeBuffers { points: Vec, dots: Vec, - row_scales: Vec, + point_scales: Vec, } impl AsPooled<()> for StripeBuffers { @@ -159,10 +159,15 @@ where } let mut leaves = Vec::new(); + // Prepare metric and ISA dispatch before replicas spawn Rayon work. The + // Copy handle is shared read-only; every stripe calls its direct function + // pointer instead of redispatching in the recursive hot path. + let kernel = PartitionKernel::new(metric); let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas { let seed = replica_seed(replica); - let mut replica_leaves = partition_replica(data, &config, metric, seed, &stripe_buffers)?; + let mut replica_leaves = + partition_replica(data, &config, metric, &kernel, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::opaque)?; @@ -176,6 +181,7 @@ fn partition_replica( data: MatrixView<'_, T>, config: &PiPNNConfig, metric: Metric, + kernel: &PartitionKernel, seed: u64, stripe_buffers: &StripeBufferPool, ) -> ANNResult>> @@ -220,6 +226,7 @@ where data, config, metric, + kernel, item, stripe_buffers, )); @@ -258,6 +265,7 @@ fn partition_one_level( data: MatrixView<'_, T>, config: &PiPNNConfig, metric: Metric, + kernel: &PartitionKernel, item: WorkItem, stripe_buffers: &StripeBufferPool, ) -> ANNResult<(Vec, Vec>)> @@ -277,6 +285,7 @@ where &leaders, fanout, metric, + kernel, stripe_buffers, )?; @@ -334,110 +343,117 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { /// Assign each point to its nearest `fanout` sampled leaders. /// -/// Leader rows are gathered once. Point rows are processed in cache-sized +/// Leader vectors are gathered once. Points are processed in cache-sized /// stripes, while a worker chunk retains one leased scratch buffer across all of /// its stripes. The flat assignment matrix preserves point order and is then /// scattered into per-leader clusters; preserving order is required for fixed /// seed determinism in later recursion levels. fn assign_to_leaders( data: MatrixView<'_, T>, - points: &[u32], - leaders: &[u32], + point_ids: &[u32], + leader_ids: &[u32], fanout: usize, metric: Metric, + kernel: &PartitionKernel, stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where T: VectorRepr + Send + Sync, { - let dimensions = data.ncols(); - let leader_values_len = checked_area("leader data", leaders.len(), dimensions)?; + let dimension_count = data.ncols(); + let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; - gather_rows(data, leaders, &mut leader_values)?; + gather_vectors(data, leader_ids, &mut leader_values)?; let mut leader_scales = if matches!(metric, Metric::L2 | Metric::Cosine) { - filled_vec(leaders.len(), 0.0f32)? + filled_vec(leader_ids.len(), 0.0f32)? } else { Vec::new() }; - for (scale, row) in leader_scales + for (scale, leader_vector) in leader_scales .iter_mut() - .zip(leader_values.chunks_exact(dimensions)) + .zip(leader_values.chunks_exact(dimension_count)) { // Leader norms participate in the top-k ordering. Preserve the original // scalar reduction order: reassociating this short setup pass through a // SIMD norm changes low bits and can send near-tied points down different // recursive partition paths. - *scale = row.iter().map(|value| value * value).sum(); + *scale = leader_vector.iter().map(|value| value * value).sum(); if metric == Metric::Cosine { *scale = scale.sqrt(); } } - let fanout = fanout.min(leaders.len()); - let assignment_len = checked_area("partition assignments", points.len(), fanout)?; + let fanout = fanout.min(leader_ids.len()); + let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; let mut assignments = filled_vec(assignment_len, 0u32)?; - let stripe_rows = assignment_stripe_rows(leaders.len()); - let assignment_stripe = checked_area("assignment stripe", stripe_rows, fanout)?; - let stripes = points.len().div_ceil(stripe_rows); - let worker_stripes = stripes.div_ceil(rayon::current_num_threads().max(1)); - let worker_rows = checked_area("assignment worker", worker_stripes, stripe_rows)?; - let worker_assignment = checked_area("assignment worker", worker_rows, fanout)?; + let stripe_points = assignment_stripe_point_count(leader_ids.len()); + let stripe_assignment_count = checked_area("assignment stripe", stripe_points, fanout)?; + let stripe_count = point_ids.len().div_ceil(stripe_points); + let worker_stripe_count = stripe_count.div_ceil(rayon::current_num_threads().max(1)); + let worker_point_count = checked_area("assignment worker", worker_stripe_count, stripe_points)?; + let worker_assignment_count = checked_area("assignment worker", worker_point_count, fanout)?; // Each worker chunk owns one scratch value and reuses it for its stripes. // build_graph pins this terminal operation to the caller-owned pool. #[allow(clippy::disallowed_methods)] assignments - .par_chunks_mut(worker_assignment) + .par_chunks_mut(worker_assignment_count) .enumerate() - .try_for_each(|(worker, worker_output)| { + .try_for_each(|(worker, worker_assignments)| { let mut buffers = stripe_buffers.get_ref(()); - let worker_first = worker * worker_rows; - for (stripe, output) in worker_output.chunks_mut(assignment_stripe).enumerate() { - let first = worker_first + stripe * stripe_rows; - let rows = output.len() / fanout; + let worker_first = worker * worker_point_count; + for (stripe, stripe_assignments) in worker_assignments + .chunks_mut(stripe_assignment_count) + .enumerate() + { + let first_point = worker_first + stripe * stripe_points; + let stripe_point_count = stripe_assignments.len() / fanout; assign_stripe( data, - &points[first..first + rows], + &point_ids[first_point..first_point + stripe_point_count], &leader_values, &leader_scales, metric, + kernel, fanout, &mut buffers, - output, + stripe_assignments, )?; } Ok::<(), ANNError>(()) })?; - scatter_assignments(points, &assignments, fanout, leaders.len()) + scatter_assignments(point_ids, &assignments, fanout, leader_ids.len()) } #[inline] #[allow(clippy::too_many_arguments)] fn assign_stripe( data: MatrixView<'_, T>, - points: &[u32], + point_ids: &[u32], leader_values: &[f32], leader_scales: &[f32], metric: Metric, + kernel: &PartitionKernel, fanout: usize, buffers: &mut StripeBuffers, - output: &mut [u32], + assignments: &mut [u32], ) -> ANNResult<()> where T: VectorRepr, { - let rows = points.len(); + let point_count = point_ids.len(); let dimensions = data.ncols(); - let leaders = leader_values.len() / dimensions; - let point_values_len = checked_area("point stripe", rows, dimensions)?; - let dots_len = checked_area("dot-product stripe", rows, leaders)?; + let leader_count = leader_values.len() / dimensions; + let point_values_len = checked_area("point stripe", point_count, dimensions)?; + let dots_len = checked_area("dot-product stripe", point_count, leader_count)?; + let output_len = checked_area("partition assignments", point_count, fanout)?; // Scratch keeps its high-water length and every consumer receives an // explicit active prefix. Resizing to the exact stripe shape would be // correct but re-zeroes the buffer whenever a pooled value moves between - // work items with different leader counts: `stripe_rows` is derived from - // `leaders`, so the point buffer swings between roughly 768 KiB and 6 MiB + // work items with different leader counts: `stripe_points` is derived from + // `leader_count`, so the point buffer swings between roughly 768 KiB and 6 MiB // and `Vec::resize` only truncates on the way down, then memsets the whole // delta on the way back up. grow_fallible(&mut buffers.points, point_values_len, 0.0)?; @@ -445,16 +461,16 @@ where let StripeBuffers { points: point_buffer, dots: dot_buffer, - row_scales: row_scale_buffer, + point_scales: point_scale_buffer, } = buffers; let point_values = &mut point_buffer[..point_values_len]; let dots = &mut dot_buffer[..dots_len]; - gather_rows(data, points, point_values)?; + gather_vectors(data, point_ids, point_values)?; diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, - rows, - leaders, + point_count, + leader_count, dimensions, 1.0, point_values, @@ -464,35 +480,49 @@ where ) .map_err(ANNError::opaque)?; - let row_scales = if metric == Metric::Cosine { - grow_fallible(row_scale_buffer, rows, 0.0)?; - let row_scales = &mut row_scale_buffer[..rows]; - for (scale, row) in row_scales + let point_scales = if metric == Metric::Cosine { + grow_fallible(point_scale_buffer, point_count, 0.0)?; + let point_scales = &mut point_scale_buffer[..point_count]; + for (scale, point_values) in point_scales .iter_mut() .zip(point_values.chunks_exact(dimensions)) { - *scale = FastL2NormSquared.evaluate(row); + *scale = FastL2NormSquared.evaluate(point_values); } - &*row_scales + &*point_scales } else { &[] }; - nearest_leaders( - PartitionTopK { - dots, - rows, - leaders, - row_scales, - leader_scales, - metric, + let scales = match metric { + Metric::L2 => PartitionScales::L2 { + leader_squared_norms: leader_scales, }, - fanout, - output, - ) - .map_err(ANNError::opaque) + Metric::Cosine => PartitionScales::Cosine { + point_squared_norms: point_scales, + leader_norms: leader_scales, + }, + Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, + }; + let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { + ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "dot-product stripe", + expected: dots_len, + actual: dots.len(), + }) + })?; + let output = MutMatrixView::try_from(assignments, point_count, fanout).map_err(|error| { + ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "partition assignments", + expected: output_len, + actual: error.into_inner().len(), + }) + })?; + kernel + .nearest_leaders(PartitionInput { dots, scales }, output) + .map_err(ANNError::opaque) } -fn gather_rows(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> +fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> where T: VectorRepr, { @@ -504,8 +534,8 @@ where actual: output.len(), })); } - for (&index, row) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { - T::as_f32_into(data.row(index as usize), row).map_err(Into::::into)?; + for (&index, vector_output) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { + T::as_f32_into(data.row(index as usize), vector_output).map_err(Into::::into)?; } Ok(()) } @@ -526,9 +556,9 @@ fn scatter_assignments( return scatter_serial(points, assignments, fanout, leaders); } - let stripe_rows = points.len().div_ceil(rayon::current_num_threads().max(1)); - let assignment_stripe = checked_area("scatter assignment stripe", stripe_rows, fanout)?; - let stripes = points.len().div_ceil(stripe_rows); + let stripe_points = points.len().div_ceil(rayon::current_num_threads().max(1)); + let stripe_assignment_count = checked_area("scatter assignment stripe", stripe_points, fanout)?; + let stripes = points.len().div_ceil(stripe_points); let mut partials = Vec::new(); partials .try_reserve_exact(stripes) @@ -540,8 +570,8 @@ fn scatter_assignments( .par_iter_mut() .zip( points - .par_chunks(stripe_rows) - .zip(assignments.par_chunks(assignment_stripe)), + .par_chunks(stripe_points) + .zip(assignments.par_chunks(stripe_assignment_count)), ) .for_each(|(slot, (points, assignments))| { *slot = Some(scatter_serial(points, assignments, fanout, leaders)); @@ -608,8 +638,8 @@ fn scatter_serial( })?; } let mut clusters = clusters_with_capacities(&sizes)?; - for (&point, row) in points.iter().zip(assignments.chunks_exact(fanout)) { - for &leader in row { + for (&point, point_assignments) in points.iter().zip(assignments.chunks_exact(fanout)) { + for &leader in point_assignments { clusters[leader as usize].push(point); } } @@ -759,14 +789,14 @@ fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult usize { - let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::()); - let rows = if rows.is_power_of_two() { - rows +fn assignment_stripe_point_count(leader_count: usize) -> usize { + let point_count = ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count.max(1) * size_of::()); + let point_count = if point_count.is_power_of_two() { + point_count } else { - rows.next_power_of_two() / 2 + point_count.next_power_of_two() / 2 }; - rows.clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) + point_count.clamp(MIN_ASSIGNMENT_STRIPE_POINTS, MAX_ASSIGNMENT_STRIPE_POINTS) } #[cfg(test)] diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 31f3e4bcad..7254ed25d9 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -24,10 +24,10 @@ fn clustered_data(points: usize, dimensions: usize) -> Matrix { diskann_utils::views::Init({ let mut position = 0usize; move || { - let row = position / dimensions; - let column = position % dimensions; + let point = position / dimensions; + let dimension = position % dimensions; position += 1; - (row / 8) as f32 * 10.0 + column as f32 * 0.01 + row as f32 * 0.001 + (point / 8) as f32 * 10.0 + dimension as f32 * 0.01 + point as f32 * 0.001 } }), points, @@ -40,11 +40,11 @@ fn directional_data(points: usize, dimensions: usize) -> Matrix { diskann_utils::views::Init({ let mut position = 0usize; move || { - let row = position / dimensions; - let column = position % dimensions; + let point = position / dimensions; + let dimension = position % dimensions; position += 1; - let angle = std::f32::consts::TAU * row as f32 / points as f32; - match column { + let angle = std::f32::consts::TAU * point as f32 / points as f32; + match dimension { 0 => angle.cos(), 1 => angle.sin(), _ => 0.0, @@ -179,14 +179,14 @@ where T: diskann::utils::VectorRepr + Send + Sync, { let points = 64; - // Partition gathering converts source rows before GEMM. Exercise conversion + // Partition gathering converts source vectors before GEMM. Exercise conversion // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let raw: Vec = (0..points * dimensions) .map(|index| { - let row = index / dimensions; - let column = index % dimensions; - ((row * 5 + column * 7 + row * column) % 23) as u8 + let point = index / dimensions; + let dimension = index % dimensions; + ((point * 5 + dimension * 7 + point * dimension) % 23) as u8 }) .collect(); let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); @@ -261,6 +261,7 @@ fn l2_leader_norms_preserve_scalar_reduction_order() { &[0, 1], 1, Metric::L2, + &PartitionKernel::new(Metric::L2), &StripeBufferPool::new((), 0, None), ) .unwrap(); @@ -298,10 +299,13 @@ fn replica_seed_derivation_is_stable_and_distinct() { } #[test] -fn assignment_stripes_use_power_of_two_row_counts() { - assert_eq!(assignment_stripe_rows(1_000), 128); - assert_eq!(assignment_stripe_rows(256), 512); - assert_eq!(assignment_stripe_rows(1), MAX_ASSIGNMENT_STRIPE_ROWS); +fn assignment_stripes_use_power_of_two_point_counts() { + assert_eq!(assignment_stripe_point_count(1_000), 128); + assert_eq!(assignment_stripe_point_count(256), 512); + assert_eq!( + assignment_stripe_point_count(1), + MAX_ASSIGNMENT_STRIPE_POINTS + ); } #[test] @@ -331,6 +335,7 @@ fn leader_assignment_handles_multiple_stripes() { &[0, 2_047], 1, Metric::L2, + &PartitionKernel::new(Metric::L2), &StripeBufferPool::new((), 0, None), ) .unwrap(); @@ -378,7 +383,7 @@ fn rejects_zero_dimensions() { #[test] fn rejects_invalid_gather_output_length() { let data = Matrix::::new(0.0, 2, 2); - let error = gather_rows(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); + let error = gather_vectors(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); assert_eq!( error.downcast::().unwrap(), From 42bc69e99127814f8dacd53027eea665dae95980 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:51:03 +0000 Subject: [PATCH 120/229] refactor(pipnn): name candidate lists --- diskann/src/graph/pipnn/finalization.rs | 134 +++++++++--------- diskann/src/graph/pipnn/finalization/tests.rs | 44 ++++-- diskann/src/graph/pipnn/leaf_build/tests.rs | 4 +- diskann/src/graph/pipnn/mod.rs | 14 +- 4 files changed, 111 insertions(+), 85 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 8acb5fda46..e714d9336d 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -6,9 +6,9 @@ //! Final graph-degree enforcement through the shared Vamana RobustPrune kernel. //! //! Candidate merging may produce more than `R` IDs for a point. This stage first -//! validates every global ID, then processes rows independently in the caller's -//! Rayon pool. Rows already within the degree bound are returned without any -//! distance work. Overfull rows are converted to source-distance candidates, +//! validates every global ID, then processes each point's candidate list in the +//! caller's Rayon pool. Lists already within the degree bound are returned without +//! distance work. Overfull lists are converted to source-distance candidates, //! passed through RobustPrune, and rewritten from the selected output. //! //! The shared kernel owns occlusion and alpha-round semantics; this adapter owns @@ -16,24 +16,24 @@ //! representation. //! //! ```text -//! candidate rows ──> validate row count and every global ID +//! candidate lists ──> validate point count and every global ID //! │ //! ┌────────────┴────────────┐ //! v v //! len <= R len > R -//! return row source-distance candidates +//! return list source-distance candidates //! │ //! v //! shared RobustPrune //! │ //! v -//! rewrite same row owner +//! rewrite same list owner //! ``` //! //! | Path | Distance evaluations | Allocation behavior | //! | --- | --- | --- | -//! | bounded row | none | move row directly to output | -//! | overfull row | source and occlusion distances | reuse Rayon-job workspace | +//! | bounded list | none | move list directly to output | +//! | overfull list | source and occlusion distances | reuse Rayon-job workspace | use std::convert::Infallible; @@ -49,29 +49,31 @@ use rayon::prelude::*; #[derive(Debug, thiserror::Error)] pub(crate) enum FinalizationError { - #[error("candidate row count {rows} does not match the dataset point count {points}")] - RowCountMismatch { rows: usize, points: usize }, - #[error("candidate ID {candidate} in row {row} is outside a {points}-point dataset")] + #[error("candidate list count {lists} does not match the dataset point count {points}")] + CandidateListCountMismatch { lists: usize, points: usize }, + #[error( + "candidate ID {candidate} for source {source_index} is outside a {points}-point dataset" + )] InvalidCandidateId { - row: usize, + source_index: usize, candidate: u32, points: usize, }, } -/// Per-Rayon-job state retained across rows. +/// Per-Rayon-job state retained across source points. /// /// `prune` owns candidate/state/output buffers. `cache` stores provider lookup /// results required by the shared kernel. Reusing both avoids per-node /// allocations, which would otherwise dominate finalization for millions of -/// short rows. +/// short candidate lists. #[derive(Default)] struct Workspace { prune: prune::Scratch, cache: Vec<(f32, Option)>, } -/// Validate candidate IDs and prune only rows whose length exceeds graph degree. +/// Validate candidate IDs and prune only lists whose length exceeds graph degree. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, @@ -81,7 +83,7 @@ pub(crate) fn prune_overfull( where T: VectorRepr + Send + Sync, { - validate_candidates(&candidates, data.nrows()).map_err(ANNError::opaque)?; + validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::opaque)?; let degree = graph.pruned_degree().get(); let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind(), false); @@ -92,70 +94,72 @@ where candidates .into_par_iter() .enumerate() - .map_init(Workspace::default, |workspace, (source, mut row)| { - // Candidate accumulators already enforce uniqueness. A bounded row - // therefore satisfies the graph policy without distance evaluation. - if row.len() <= degree { - return Ok(row); - } + .map_init( + Workspace::default, + |workspace, (source, mut source_candidates)| { + // Candidate accumulators already enforce uniqueness. A bounded list + // therefore satisfies the graph policy without distance evaluation. + if source_candidates.len() <= degree { + return Ok(source_candidates); + } - let source_id = u32::try_from(source).map_err(ANNError::opaque)?; - let source_vector = data.row(source); - let pool = workspace.prune.candidates_mut(); - pool.clear(); - pool.try_reserve(row.len()).map_err(ANNError::opaque)?; - pool.extend(row.iter().copied().map(|candidate| { - Neighbor::new( - candidate, - distance.evaluate_similarity(source_vector, data.row(candidate as usize)), - ) - })); - // as_context sorts the active candidate prefix by source distance. - // The callback below is needed only for selected-to-candidate - // occlusion checks; dimension specialization stays in `distance`. - let candidate_count = pool.len(); - let mut context = workspace.prune.as_context(candidate_count); - prune::robust_prune( - &mut context, - policy, - &mut workspace.cache, - Some, - |left, right| { - Ok::<_, Infallible>( - distance.evaluate_similarity( + let source_id = u32::try_from(source).map_err(ANNError::opaque)?; + let source_vector = data.row(source); + let pool = workspace.prune.candidates_mut(); + pool.clear(); + pool.try_reserve(source_candidates.len()) + .map_err(ANNError::opaque)?; + pool.extend(source_candidates.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + // as_context sorts the active candidate prefix by source distance. + // The callback below is needed only for selected-to-candidate + // occlusion checks; dimension specialization stays in `distance`. + let candidate_count = pool.len(); + let mut context = workspace.prune.as_context(candidate_count); + prune::robust_prune( + &mut context, + policy, + &mut workspace.cache, + Some, + |left, right| { + Ok::<_, Infallible>(distance.evaluate_similarity( data.row(*left as usize), data.row(*right as usize), - ), - ) - }, - |id| id == source_id, - ) - .map_err(ANNError::opaque)?; + )) + }, + |id| id == source_id, + ) + .map_err(ANNError::opaque)?; - // RobustPrune selects distinct candidate positions, so its output IDs - // are unique by construction. `extend_from_slice` would re-derive that - // with an O(degree^2) membership scan per row; the trusted overwrite is - // a copy and still verifies uniqueness under debug assertions. - row.overwrite_trusted(workspace.prune.neighbors()); - Ok(row) - }) + // RobustPrune selects distinct candidate positions, so its output IDs + // are unique by construction. `extend_from_slice` would re-derive that + // with an O(degree^2) membership scan per list; the trusted overwrite + // is a copy and still verifies uniqueness under debug assertions. + source_candidates.overwrite_trusted(workspace.prune.neighbors()); + Ok(source_candidates) + }, + ) .collect() } -fn validate_candidates( +fn validate_candidate_lists( candidates: &[AdjacencyList], points: usize, ) -> Result<(), FinalizationError> { if candidates.len() != points { - return Err(FinalizationError::RowCountMismatch { - rows: candidates.len(), + return Err(FinalizationError::CandidateListCountMismatch { + lists: candidates.len(), points, }); } - for (row_id, row) in candidates.iter().enumerate() { - if let Some(&candidate) = row.iter().find(|&&id| id as usize >= points) { + for (source, source_candidates) in candidates.iter().enumerate() { + if let Some(&candidate) = source_candidates.iter().find(|&&id| id as usize >= points) { return Err(FinalizationError::InvalidCandidateId { - row: row_id, + source_index: source, candidate, points, }); diff --git a/diskann/src/graph/pipnn/finalization/tests.rs b/diskann/src/graph/pipnn/finalization/tests.rs index fde3a6e0ea..42a534bb45 100644 --- a/diskann/src/graph/pipnn/finalization/tests.rs +++ b/diskann/src/graph/pipnn/finalization/tests.rs @@ -25,15 +25,20 @@ fn graph_config(degree: usize) -> Config { .unwrap() } -fn row(ids: impl IntoIterator) -> AdjacencyList { +fn candidate_list(ids: impl IntoIterator) -> AdjacencyList { AdjacencyList::from_iter_untrusted(ids) } #[test] -fn preserves_rows_within_the_degree_bound() { +fn preserves_lists_within_the_degree_bound() { let data = [0.0_f32, 1.0, 2.0, 3.0]; let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let candidates = vec![row([3, 1]), row([]), row([]), row([])]; + let candidates = vec![ + candidate_list([3, 1]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); @@ -41,10 +46,15 @@ fn preserves_rows_within_the_degree_bound() { } #[test] -fn prunes_an_overfull_row_with_the_vamana_kernel() { +fn prunes_an_overfull_list_with_the_vamana_kernel() { let data = [0.0_f32, 1.0, 2.0, -3.0]; let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let candidates = vec![row([3, 2, 1]), row([]), row([]), row([])]; + let candidates = vec![ + candidate_list([3, 2, 1]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); @@ -56,14 +66,18 @@ fn prunes_an_overfull_row_with_the_vamana_kernel() { fn rejects_invalid_candidate_ids_without_panicking() { let data = [0.0_f32, 1.0, 2.0]; let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![row([1, 3]), row([]), row([])]; + let candidates = vec![ + candidate_list([1, 3]), + candidate_list([]), + candidate_list([]), + ]; let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); assert!(matches!( error.downcast_ref::(), Some(FinalizationError::InvalidCandidateId { - row: 0, + source_index: 0, candidate: 3, points: 3, }) @@ -71,16 +85,24 @@ fn rejects_invalid_candidate_ids_without_panicking() { } #[test] -fn rejects_candidate_row_count_mismatch_without_panicking() { +fn rejects_candidate_list_count_mismatch_without_panicking() { let data = [0.0_f32, 1.0, 2.0]; let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![row([]), row([]), row([]), row([])]; + let candidates = vec![ + candidate_list([]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); assert!(matches!( error.downcast_ref::(), - Some(FinalizationError::RowCountMismatch { rows: 4, points: 3 }) + Some(FinalizationError::CandidateListCountMismatch { + lists: 4, + points: 3 + }) )); } @@ -90,7 +112,7 @@ fn rejects_more_candidates_than_the_shared_position_type_can_represent() { let data = vec![0.0_f32; count + 1]; let data = MatrixView::try_from(&data[..], count + 1, 1).unwrap(); let mut candidates = Vec::with_capacity(count + 1); - candidates.push(row(1..=count as u32)); + candidates.push(candidate_list(1..=count as u32)); candidates.resize_with(count + 1, AdjacencyList::new); let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index 2d4e66c2e6..5d6cd92e94 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -322,7 +322,7 @@ fn reports_shape_overflow_before_allocating() { } #[test] -fn rejects_an_invalid_kernel_position() { +fn rejects_an_invalid_kernel_target() { let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; let error = add_symmetric_neighbors( &[10, 20], @@ -392,7 +392,7 @@ fn allocation_errors_preserve_buffer_context() { } #[test] -fn direct_candidate_accumulator_keeps_unique_sorted_rows() { +fn direct_candidate_accumulator_keeps_unique_sorted_lists() { let candidates = DirectCandidates::new(2).unwrap(); candidates .add_leaf( diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index e0bbd72bd6..f328641e12 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -340,11 +340,11 @@ impl<'a> PiPNNBuildContext<'a> { } } -/// Build PiPNN adjacency for real rows in `data`. +/// Build PiPNN adjacency for real points in `data`. /// /// This is the core algorithm boundary. Search entry-point selection, frozen nodes, /// providers, serialization, and index writers belong to the outer build pipelines. -/// For raw `u8` and `i8` rows, `CosineNormalized` is evaluated as `Cosine` because +/// For raw `u8` and `i8` vectors, `CosineNormalized` is evaluated as `Cosine` because /// those representations are converted to f32 scratch but are not unit-normalized. pub fn build_graph( data: MatrixView<'_, T>, @@ -365,7 +365,7 @@ where { if data.nrows() == 0 { return Err(ANNError::log_dimension_mismatch_error( - "PiPNN requires at least one data row".into(), + "PiPNN requires at least one data point".into(), )); } if data.ncols() == 0 { @@ -375,7 +375,7 @@ where } if data.nrows() > u32::MAX as usize { return Err(config_error(format!( - "dataset row count ({}) exceeds the u32 graph ID limit", + "dataset point count ({}) exceeds the u32 graph ID limit", data.nrows() ))); } @@ -386,19 +386,19 @@ where data.ncols() )) })?; - // Integer source rows are not guaranteed unit-normalized after conversion, + // Integer source vectors are not guaranteed unit-normalized after conversion, // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition(data, context.config.clone(), metric))?; - // `leaves` is consumed here. Workers borrow individual ID rows during the + // `leaves` is consumed here. Workers borrow individual ID lists during the // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) .map_err(ANNError::opaque) })?; - // Finalization consumes candidate rows and reuses their allocations for the + // Finalization consumes candidate lists and reuses their allocations for the // resulting adjacency where possible. tracing::info_span!("pipnn.finalization") .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) From c06e2645c29b58e9224124aeda439ad7bc872a6e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 07:27:37 +0000 Subject: [PATCH 121/229] refactor(pipnn): complete graph module migration Move core construction under diskann::graph::pipnn so finalization can reuse private RobustPrune state. Remove standalone PiPNN Cargo benchmarks, group public tests by PiPNN behavior, and delete duplicate or non-discriminating cases. --- diskann-pipnn/README.md | 19 ------- diskann/src/graph/pipnn/finalization.rs | 6 +-- diskann/src/graph/pipnn/finalization/tests.rs | 23 +------- diskann/src/graph/pipnn/leaf_build.rs | 8 +-- diskann/src/graph/pipnn/leaf_build/tests.rs | 52 ++++++++++--------- diskann/src/graph/pipnn/mod.rs | 6 +-- diskann/src/graph/pipnn/partitioning.rs | 42 +++++++-------- diskann/src/graph/pipnn/partitioning/tests.rs | 14 ++--- .../{build_graph.rs => pipnn_build_graph.rs} | 12 +++-- diskann/tests/{config.rs => pipnn_config.rs} | 9 +--- 10 files changed, 77 insertions(+), 114 deletions(-) delete mode 100644 diskann-pipnn/README.md rename diskann/tests/{build_graph.rs => pipnn_build_graph.rs} (96%) rename diskann/tests/{config.rs => pipnn_config.rs} (93%) diff --git a/diskann-pipnn/README.md b/diskann-pipnn/README.md deleted file mode 100644 index 4ac1618f03..0000000000 --- a/diskann-pipnn/README.md +++ /dev/null @@ -1,19 +0,0 @@ -# PiPNN graph construction - -This crate implements the graph-construction stages from [PiPNN: Pick in Partitions for Fast and Accurate ANN Graph Construction](https://arxiv.org/html/2602.21247v1). - -## Boundary - -PiPNN core consumes a dense `MatrixView`, graph policy, and a caller-owned Rayon pool, then returns adjacency lists for the dataset's real point IDs. It does not own start or frozen points, vector or neighbor providers, PQ, disk headers, serialization, or search. Those concerns remain in the outer in-memory and disk pipelines. - -The dense view is intentional: partition assignment and leaf all-pairs kernels operate over the whole source matrix. Materializing provider state inside the algorithm would couple numerical graph construction to storage lifecycle and would require a second dataset copy. Integrations should finish PiPNN scratch before allocating or populating their searchable provider. - -## Policy ownership - -- `PiPNNConfig` owns partition and leaf-selection parameters: leaf bounds, sampling fraction, fanout levels, leaf `k`, and replicas. -- DiskANN graph configuration owns metric, output degree, build-L, alpha, and prune policy. -- Candidate-merging policies are separate validated options; they must not make graph policy fields redundant or silently cap the requested degree. - -## Execution - -A build runs partitioning, leaf construction, candidate merging, then graph finalization. All parallel work executes in the supplied pool. Per-job scratch is initialized through Rayon and is released through normal ownership when its stage completes; the core has no global thread-local buffers or cleanup broadcasts. diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index e714d9336d..2d0451dc73 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -38,13 +38,13 @@ use std::convert::Infallible; use crate::{ - graph::{prune, AdjacencyList, Config}, + ANNError, ANNResult, + graph::{AdjacencyList, Config, prune}, neighbor::Neighbor, utils::VectorRepr, - ANNError, ANNResult, }; use diskann_utils::views::MatrixView; -use diskann_vector::{distance::Metric, DistanceFunction}; +use diskann_vector::{DistanceFunction, distance::Metric}; use rayon::prelude::*; #[derive(Debug, thiserror::Error)] diff --git a/diskann/src/graph/pipnn/finalization/tests.rs b/diskann/src/graph/pipnn/finalization/tests.rs index 42a534bb45..66562049c2 100644 --- a/diskann/src/graph/pipnn/finalization/tests.rs +++ b/diskann/src/graph/pipnn/finalization/tests.rs @@ -4,8 +4,8 @@ */ use crate::graph::{ - config::{self, MaxDegree}, AdjacencyList, + config::{self, MaxDegree}, }; use diskann_utils::views::MatrixView; @@ -58,8 +58,7 @@ fn prunes_an_overfull_list_with_the_vamana_kernel() { let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - assert!(actual[0].len() <= 2); - assert!(actual[0].contains(1)); + assert_eq!(&*actual[0], &[1, 3]); } #[test] @@ -105,21 +104,3 @@ fn rejects_candidate_list_count_mismatch_without_panicking() { }) )); } - -#[test] -fn rejects_more_candidates_than_the_shared_position_type_can_represent() { - let count = u16::MAX as usize + 1; - let data = vec![0.0_f32; count + 1]; - let data = MatrixView::try_from(&data[..], count + 1, 1).unwrap(); - let mut candidates = Vec::with_capacity(count + 1); - candidates.push(candidate_list(1..=count as u32)); - candidates.resize_with(count + 1, AdjacencyList::new); - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::>(), - Some(prune::RobustPruneError::TooManyCandidates { actual, max }) - if *actual == count && *max == u16::MAX as usize - )); -} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 37cf5e3d14..ec8763b190 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -27,9 +27,9 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use rayon::prelude::*; -use crate::leaf_kernel::{ - leaf_neighbor_count, leaf_output_len, LeafInput, LeafKernel, LeafKernelError, - LeafKernelWorkspace, LeafNeighbor, +use super::leaf_kernel::{ + LeafInput, LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, + leaf_output_len, }; /// Failure while converting leaves into direct graph candidates. @@ -67,7 +67,7 @@ pub(crate) enum LeafBuildError { leaf: usize, point: u32, #[source] - source: diskann::ANNError, + source: crate::ANNError, }, #[error("lower-AAT failed for leaf {leaf}")] LowerAat { diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index 5d6cd92e94..6c55e6a577 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -9,8 +9,8 @@ use half::f16; use std::collections::BTreeSet; use super::{ - add_symmetric_neighbors, allocation_error, build_leaf_candidates, DirectCandidates, - LeafBuffers, LeafBuildError, + DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, + build_leaf_candidates, }; fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { @@ -29,14 +29,14 @@ fn build( leaves: &[Vec], k: usize, metric: Metric, -) -> Result>, LeafBuildError> +) -> Result>, LeafBuildError> where - T: diskann::utils::VectorRepr + 'static, + T: crate::utils::VectorRepr + 'static, { pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) } -fn adjacency_lists(graph: Vec>) -> Vec> { +fn adjacency_lists(graph: Vec>) -> Vec> { graph.into_iter().map(Vec::from).collect() } @@ -150,7 +150,7 @@ fn global_id_translation_is_independent_of_leaf_order() { fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> where - T: diskann::utils::VectorRepr + 'static, + T: crate::utils::VectorRepr + 'static, { let leaves = vec![(0..points as u32).collect()]; adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) @@ -158,7 +158,7 @@ where fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) where - T: diskann::utils::VectorRepr + 'static, + T: crate::utils::VectorRepr + 'static, { let points = 8; // Source dimension controls VectorRepr conversion chunking. Cover tails on @@ -213,9 +213,11 @@ fn all_metrics_produce_symmetric_unique_non_self_candidates() { let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); for (source, neighbors) in graph.iter().enumerate() { assert!(neighbors.iter().all(|&target| target as usize != source)); - assert!(neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32))); + assert!( + neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + ); assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); } } @@ -240,7 +242,7 @@ fn parallel_leaf_schedule_does_not_change_candidate_order() { } #[test] -fn rejects_invalid_shape_inputs_without_panicking() { +fn rejects_invalid_dimensions_and_leaf_membership() { let data = [0.0_f32, 1.0]; let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); assert!(matches!( @@ -288,10 +290,12 @@ fn singleton_and_zero_k_leaves_add_no_candidates() { ) .unwrap(); let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); - assert!(singleton - .iter() - .chain(&zero_k) - .all(|candidates| candidates.is_empty())); + assert!( + singleton + .iter() + .chain(&zero_k) + .all(|candidates| candidates.is_empty()) + ); } #[test] @@ -323,13 +327,13 @@ fn reports_shape_overflow_before_allocating() { #[test] fn rejects_an_invalid_kernel_target() { - let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; let error = add_symmetric_neighbors( &[10, 20], 1, &[ - crate::leaf_kernel::LeafNeighbor::new(9, 1.0), - crate::leaf_kernel::LeafNeighbor::new(0, 1.0), + super::super::leaf_kernel::LeafNeighbor::new(9, 1.0), + super::super::leaf_kernel::LeafNeighbor::new(0, 1.0), ], &mut graph, ) @@ -345,13 +349,13 @@ fn rejects_an_invalid_kernel_target() { #[test] fn skips_duplicate_global_ids_without_self_edges() { - let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; add_symmetric_neighbors( &[7, 7], 1, &[ - crate::leaf_kernel::LeafNeighbor::new(1, 0.0), - crate::leaf_kernel::LeafNeighbor::new(0, 0.0), + super::super::leaf_kernel::LeafNeighbor::new(1, 0.0), + super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), ], &mut graph, ) @@ -367,7 +371,7 @@ fn poisoned_candidate_lists_return_errors() { panic!("poison candidate list"); }); assert!(matches!( - candidates.add_leaf(&[0], &[diskann::graph::AdjacencyList::new()]), + candidates.add_leaf(&[0], &[crate::graph::AdjacencyList::new()]), Err(LeafBuildError::PoisonedCandidateList { point: 0 }) )); assert!(matches!( @@ -398,8 +402,8 @@ fn direct_candidate_accumulator_keeps_unique_sorted_lists() { .add_leaf( &[0, 1], &[ - diskann::graph::AdjacencyList::from_iter_untrusted([1, 1]), - diskann::graph::AdjacencyList::from_iter_untrusted([0]), + crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), + crate::graph::AdjacencyList::from_iter_untrusted([0]), ], ) .unwrap(); diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index f328641e12..9a7847aa58 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -//! Provider-independent PiPNN graph construction. +//! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! //! PiPNN means **Pick-in-Partitions Nearest Neighbors**. It 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. This crate constructs that adjacency; it does not execute queries. +//! later query. This module constructs that adjacency; it does not execute queries. //! //! Incremental builders such as Vamana find construction candidates by running //! beam search against a partially built graph: they repeatedly follow graph @@ -153,9 +153,9 @@ mod simd; mod topk; use crate::{ + ANNError, ANNResult, graph::{AdjacencyList, Config}, utils::VectorRepr, - ANNError, ANNResult, }; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 88867ce108..96ed91878b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -39,19 +39,19 @@ use std::collections::HashSet; -use crate::{utils::VectorRepr, ANNError, ANNResult}; +use crate::{ANNError, ANNResult, utils::VectorRepr}; use diskann_linalg::Transpose; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, }; -use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; -use rand::{prelude::IndexedRandom, SeedableRng}; +use diskann_vector::{Norm, distance::Metric, norm::FastL2NormSquared}; +use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; -use crate::{ - partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, +use super::{ PiPNNConfig, + partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, }; // Private algorithm and batching constants live together. None are user policy. @@ -704,22 +704,22 @@ fn global_merge_small( if !small.is_empty() { let mut remainder = drain_sorted(&mut small)?; - if remainder.len() < c_min { - if let Some(last) = merged.last_mut() { - remainder.retain(|id| !last.contains(id)); - let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { - buffer: "small-leaf tail merge", - rows: last.len(), - cols: remainder.len(), - }) - })?; - if combined <= c_max { - last.try_reserve(remainder.len()) - .map_err(ANNError::opaque)?; - last.append(&mut remainder); - last.sort_unstable(); - } + if remainder.len() < c_min + && let Some(last) = merged.last_mut() + { + remainder.retain(|id| !last.contains(id)); + let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "small-leaf tail merge", + rows: last.len(), + cols: remainder.len(), + }) + })?; + if combined <= c_max { + last.try_reserve(remainder.len()) + .map_err(ANNError::opaque)?; + last.append(&mut remainder); + last.sort_unstable(); } } if !remainder.is_empty() { diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 7254ed25d9..8780078ca5 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -4,7 +4,7 @@ */ use diskann_utils::views::{Matrix, MatrixView}; -use diskann_vector::{distance::Metric, Half}; +use diskann_vector::{Half, distance::Metric}; use super::*; @@ -70,9 +70,11 @@ fn sorted_memberships(leaves: &[Vec]) -> Vec> { } fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { - assert!(leaves - .iter() - .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max)); + assert!( + leaves + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) + ); let mut counts = vec![0usize; points]; for leaf in leaves { let mut ids = leaf.clone(); @@ -110,7 +112,7 @@ fn partition_is_fixed_seed_deterministic_and_bounded() { } #[test] -fn recursion_after_fanout_levels_falls_back_to_one() { +fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { let data = clustered_data(80, 4); let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); @@ -176,7 +178,7 @@ fn replicas_cover_every_point_once_or_more_per_replica() { fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) where - T: diskann::utils::VectorRepr + Send + Sync, + T: crate::utils::VectorRepr + Send + Sync, { let points = 64; // Partition gathering converts source vectors before GEMM. Exercise conversion diff --git a/diskann/tests/build_graph.rs b/diskann/tests/pipnn_build_graph.rs similarity index 96% rename from diskann/tests/build_graph.rs rename to diskann/tests/pipnn_build_graph.rs index 0f823778eb..12b8c0ec9a 100644 --- a/diskann/tests/build_graph.rs +++ b/diskann/tests/pipnn_build_graph.rs @@ -5,16 +5,17 @@ #![cfg(feature = "pipnn")] #![allow( + clippy::expect_used, clippy::unwrap_used, reason = "deterministic test fixture construction must abort on invalid setup" )] use diskann::graph::config::{self, MaxDegree}; -use diskann::graph::pipnn::{build_graph, PiPNNBuildContext, PiPNNConfig}; +use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig, build_graph}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; use half::f16; -use rand::{rngs::StdRng, Rng, SeedableRng}; +use rand::{Rng, SeedableRng, rngs::StdRng}; fn pipnn_config() -> PiPNNConfig { PiPNNConfig { @@ -58,9 +59,10 @@ fn assert_graph_invariants( sorted.sort_unstable(); sorted.dedup(); assert_eq!(sorted.len(), row.len()); - assert!(row - .iter() - .all(|&id| (id as usize) < points && id as usize != source)); + assert!( + row.iter() + .all(|&id| (id as usize) < points && id as usize != source) + ); } } diff --git a/diskann/tests/config.rs b/diskann/tests/pipnn_config.rs similarity index 93% rename from diskann/tests/config.rs rename to diskann/tests/pipnn_config.rs index 7d46e9a2a5..23e569c770 100644 --- a/diskann/tests/config.rs +++ b/diskann/tests/pipnn_config.rs @@ -5,6 +5,7 @@ #![cfg(feature = "pipnn")] #![allow( + clippy::expect_used, clippy::unwrap_used, reason = "deterministic test fixture construction must abort on invalid setup" )] @@ -39,14 +40,6 @@ fn pool() -> rayon::ThreadPool { .unwrap() } -#[test] -fn accepts_the_six_algorithm_parameters_with_outer_graph_policy() { - let graph = graph_config(Metric::L2, 1.2); - let pool = pool(); - - PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); -} - #[test] fn rejects_each_invalid_algorithm_parameter() { let graph = graph_config(Metric::L2, 1.2); From da9219fb46988304aae97ba9d5aa4857690c1849 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:46:04 +0000 Subject: [PATCH 122/229] refactor(pipnn): prepare RobustPrune inputs locally Keep sorting, workspace allocation, source exclusion, and adjacency rewriting in PiPNN finalization. The shared internal kernel now sees only prepared candidates and state. --- diskann/src/graph/pipnn/finalization.rs | 97 +++++++++++++++++-------- 1 file changed, 65 insertions(+), 32 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 2d0451dc73..9adf8c92c8 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -39,7 +39,10 @@ use std::convert::Infallible; use crate::{ ANNError, ANNResult, - graph::{AdjacencyList, Config, prune}, + graph::{ + AdjacencyList, Config, + internal::{SortedNeighbors, prune}, + }, neighbor::Neighbor, utils::VectorRepr, }; @@ -61,16 +64,15 @@ pub(crate) enum FinalizationError { }, } -/// Per-Rayon-job state retained across source points. +/// Per-Rayon-job preparation and kernel state retained across source points. /// -/// `prune` owns candidate/state/output buffers. `cache` stores provider lookup -/// results required by the shared kernel. Reusing both avoids per-node -/// allocations, which would otherwise dominate finalization for millions of -/// short candidate lists. +/// PiPNN owns sorting, allocation, and ID translation. The shared internal +/// kernel receives only the prepared candidates and an exactly sized state slice. #[derive(Default)] struct Workspace { - prune: prune::Scratch, - cache: Vec<(f32, Option)>, + pool: Vec>, + prepared: Vec>, + states: Vec, } /// Validate candidate IDs and prune only lists whose length exceeds graph degree. @@ -86,7 +88,7 @@ where validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::opaque)?; let degree = graph.pruned_degree().get(); - let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind(), false); + let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind()); let distance = T::distance(metric, Some(data.ncols())); // build_graph installs the complete call tree in the caller-owned pool. @@ -105,41 +107,72 @@ where let source_id = u32::try_from(source).map_err(ANNError::opaque)?; let source_vector = data.row(source); - let pool = workspace.prune.candidates_mut(); - pool.clear(); - pool.try_reserve(source_candidates.len()) + workspace.pool.clear(); + workspace + .pool + .try_reserve(source_candidates.len()) + .map_err(ANNError::opaque)?; + workspace + .pool + .extend(source_candidates.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance + .evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + + let candidate_count = workspace.pool.len(); + prune::validate_candidate_count::(candidate_count) + .map_err(ANNError::opaque)?; + workspace.prepared.clear(); + workspace + .prepared + .try_reserve(candidate_count) .map_err(ANNError::opaque)?; - pool.extend(source_candidates.iter().copied().map(|candidate| { - Neighbor::new( - candidate, - distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + { + // Sorting/capping precedes source exclusion so filtering cannot + // backfill with farther candidates. + let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); + workspace + .prepared + .extend(sorted.iter().filter_map(|neighbor| { + let id = *neighbor.id(); + (id != source_id) + .then(|| prune::Candidate::new(id, neighbor.distance(), id)) + })); + } + workspace + .states + .try_reserve( + workspace + .prepared + .len() + .saturating_sub(workspace.states.len()), ) - })); - // as_context sorts the active candidate prefix by source distance. - // The callback below is needed only for selected-to-candidate - // occlusion checks; dimension specialization stays in `distance`. - let candidate_count = pool.len(); - let mut context = workspace.prune.as_context(candidate_count); - prune::robust_prune( - &mut context, + .map_err(ANNError::opaque)?; + workspace + .states + .resize(workspace.prepared.len(), prune::State::default()); + + let selected = prune::robust_prune( + &workspace.prepared, + workspace.states.as_mut_slice(), policy, - &mut workspace.cache, - Some, |left, right| { Ok::<_, Infallible>(distance.evaluate_similarity( data.row(*left as usize), data.row(*right as usize), )) }, - |id| id == source_id, ) .map_err(ANNError::opaque)?; - // RobustPrune selects distinct candidate positions, so its output IDs - // are unique by construction. `extend_from_slice` would re-derive that - // with an O(degree^2) membership scan per list; the trusted overwrite - // is a copy and still verifies uniqueness under debug assertions. - source_candidates.overwrite_trusted(workspace.prune.neighbors()); + let mut guard = source_candidates.resize(selected); + for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { + *destination = *workspace.prepared[state.selected_position()].id(); + } + guard.finish(selected); Ok(source_candidates) }, ) From f8be41e87456f70c9fa4f66a2abde4a360038051 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 12:04:57 +0000 Subject: [PATCH 123/229] refactor(pipnn): use shared robust prune core Address the renamed internal modules and preserve graph Config alpha behavior without adding PiPNN-specific validation. --- diskann/src/graph/pipnn/finalization.rs | 7 ++++--- diskann/src/graph/pipnn/mod.rs | 6 ------ diskann/tests/pipnn_config.rs | 7 ++----- 3 files changed, 6 insertions(+), 14 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 9adf8c92c8..d46b8993a5 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -41,7 +41,7 @@ use crate::{ ANNError, ANNResult, graph::{ AdjacencyList, Config, - internal::{SortedNeighbors, prune}, + internal::{SortedNeighbors, robust_prune as prune}, }, neighbor::Neighbor, utils::VectorRepr, @@ -88,7 +88,6 @@ where validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::opaque)?; let degree = graph.pruned_degree().get(); - let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind()); let distance = T::distance(metric, Some(data.ncols())); // build_graph installs the complete call tree in the caller-owned pool. @@ -158,7 +157,9 @@ where let selected = prune::robust_prune( &workspace.prepared, workspace.states.as_mut_slice(), - policy, + degree, + graph.alpha(), + graph.prune_kind(), |left, right| { Ok::<_, Infallible>(distance.evaluate_similarity( data.row(*left as usize), diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 9a7847aa58..9073e2157c 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -318,12 +318,6 @@ impl<'a> PiPNNBuildContext<'a> { pool: &'a ThreadPool, ) -> ANNResult { config.validate()?; - if !graph.alpha().is_finite() || graph.alpha() < 1.0 { - return Err(config_error(format!( - "graph alpha ({}) must be finite and at least 1", - graph.alpha() - ))); - } if graph.prune_kind() != metric.into() { return Err(config_error(format!( "graph prune kind {:?} is incompatible with metric {metric:?}", diff --git a/diskann/tests/pipnn_config.rs b/diskann/tests/pipnn_config.rs index 23e569c770..19a0d1ce8a 100644 --- a/diskann/tests/pipnn_config.rs +++ b/diskann/tests/pipnn_config.rs @@ -114,13 +114,10 @@ fn rejects_graph_policy_for_a_different_metric() { } #[test] -fn rejects_invalid_outer_alpha() { +fn does_not_add_alpha_validation_beyond_graph_config() { let pool = pool(); for alpha in [0.9, f32::NAN, f32::INFINITY] { let graph = graph_config(Metric::L2, alpha); - let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); - - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); - assert!(error.to_string().contains("alpha")); + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); } } From 62a4e9835a5ffc362e3202b1f0943108e7da35cf Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:52:21 +0000 Subject: [PATCH 124/229] test(pipnn): adapt and colocate core tests Use current main error APIs and keep private plus graph/config composition tests beside their owning implementation modules. --- diskann/src/graph/pipnn/finalization.rs | 120 +++- diskann/src/graph/pipnn/finalization/tests.rs | 106 ---- diskann/src/graph/pipnn/leaf_build.rs | 413 +++++++++++++- diskann/src/graph/pipnn/leaf_build/tests.rs | 414 -------------- diskann/src/graph/pipnn/mod.rs | 390 ++++++++++++- diskann/src/graph/pipnn/partitioning.rs | 514 ++++++++++++++++-- diskann/src/graph/pipnn/partitioning/tests.rs | 423 -------------- diskann/src/graph/pipnn/tests.rs | 27 - diskann/tests/pipnn_build_graph.rs | 233 -------- diskann/tests/pipnn_config.rs | 123 ----- 10 files changed, 1365 insertions(+), 1398 deletions(-) delete mode 100644 diskann/src/graph/pipnn/finalization/tests.rs delete mode 100644 diskann/src/graph/pipnn/leaf_build/tests.rs delete mode 100644 diskann/src/graph/pipnn/partitioning/tests.rs delete mode 100644 diskann/src/graph/pipnn/tests.rs delete mode 100644 diskann/tests/pipnn_build_graph.rs delete mode 100644 diskann/tests/pipnn_config.rs diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index d46b8993a5..a10d4df027 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -85,7 +85,7 @@ pub(crate) fn prune_overfull( where T: VectorRepr + Send + Sync, { - validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::opaque)?; + validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::new)?; let degree = graph.pruned_degree().get(); let distance = T::distance(metric, Some(data.ncols())); @@ -104,13 +104,13 @@ where return Ok(source_candidates); } - let source_id = u32::try_from(source).map_err(ANNError::opaque)?; + let source_id = u32::try_from(source).map_err(ANNError::new)?; let source_vector = data.row(source); workspace.pool.clear(); workspace .pool .try_reserve(source_candidates.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; workspace .pool .extend(source_candidates.iter().copied().map(|candidate| { @@ -123,12 +123,12 @@ where let candidate_count = workspace.pool.len(); prune::validate_candidate_count::(candidate_count) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; workspace.prepared.clear(); workspace .prepared .try_reserve(candidate_count) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; { // Sorting/capping precedes source exclusion so filtering cannot // backfill with farther candidates. @@ -138,7 +138,7 @@ where .extend(sorted.iter().filter_map(|neighbor| { let id = *neighbor.id(); (id != source_id) - .then(|| prune::Candidate::new(id, neighbor.distance(), id)) + .then(|| prune::Candidate::new(id, *neighbor.distance(), id)) })); } workspace @@ -149,7 +149,7 @@ where .len() .saturating_sub(workspace.states.len()), ) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; workspace .states .resize(workspace.prepared.len(), prune::State::default()); @@ -167,7 +167,7 @@ where )) }, ) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; let mut guard = source_candidates.resize(selected); for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { @@ -203,4 +203,106 @@ fn validate_candidate_lists( } #[cfg(test)] -mod tests; +mod tests { + use crate::graph::{ + AdjacencyList, + config::{self, MaxDegree}, + }; + use diskann_utils::views::MatrixView; + + use super::*; + + fn graph_config(degree: usize) -> Config { + config::Builder::new_with( + degree, + MaxDegree::same(), + degree, + Metric::L2.into(), + |builder| { + builder.alpha(1.2); + }, + ) + .build() + .unwrap() + } + + fn candidate_list(ids: impl IntoIterator) -> AdjacencyList { + AdjacencyList::from_iter_untrusted(ids) + } + + #[test] + fn preserves_lists_within_the_degree_bound() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![ + candidate_list([3, 1]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert_eq!(&*actual[0], &[1, 3]); + } + + #[test] + fn prunes_an_overfull_list_with_the_vamana_kernel() { + let data = [0.0_f32, 1.0, 2.0, -3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![ + candidate_list([3, 2, 1]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert_eq!(&*actual[0], &[1, 3]); + } + + #[test] + fn rejects_invalid_candidate_ids_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![ + candidate_list([1, 3]), + candidate_list([]), + candidate_list([]), + ]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::InvalidCandidateId { + source_index: 0, + candidate: 3, + points: 3, + }) + )); + } + + #[test] + fn rejects_candidate_list_count_mismatch_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![ + candidate_list([]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::CandidateListCountMismatch { + lists: 4, + points: 3 + }) + )); + } +} diff --git a/diskann/src/graph/pipnn/finalization/tests.rs b/diskann/src/graph/pipnn/finalization/tests.rs deleted file mode 100644 index 66562049c2..0000000000 --- a/diskann/src/graph/pipnn/finalization/tests.rs +++ /dev/null @@ -1,106 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use crate::graph::{ - AdjacencyList, - config::{self, MaxDegree}, -}; -use diskann_utils::views::MatrixView; - -use super::*; - -fn graph_config(degree: usize) -> Config { - config::Builder::new_with( - degree, - MaxDegree::same(), - degree, - Metric::L2.into(), - |builder| { - builder.alpha(1.2); - }, - ) - .build() - .unwrap() -} - -fn candidate_list(ids: impl IntoIterator) -> AdjacencyList { - AdjacencyList::from_iter_untrusted(ids) -} - -#[test] -fn preserves_lists_within_the_degree_bound() { - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let candidates = vec![ - candidate_list([3, 1]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - - let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - - assert_eq!(&*actual[0], &[1, 3]); -} - -#[test] -fn prunes_an_overfull_list_with_the_vamana_kernel() { - let data = [0.0_f32, 1.0, 2.0, -3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let candidates = vec![ - candidate_list([3, 2, 1]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - - let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - - assert_eq!(&*actual[0], &[1, 3]); -} - -#[test] -fn rejects_invalid_candidate_ids_without_panicking() { - let data = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![ - candidate_list([1, 3]), - candidate_list([]), - candidate_list([]), - ]; - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::InvalidCandidateId { - source_index: 0, - candidate: 3, - points: 3, - }) - )); -} - -#[test] -fn rejects_candidate_list_count_mismatch_without_panicking() { - let data = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![ - candidate_list([]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::CandidateListCountMismatch { - lists: 4, - points: 3 - }) - )); -} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index ec8763b190..1c60a40c4e 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -437,4 +437,415 @@ fn poisoned_candidate_list(point: u32) -> LeafBuildError { } #[cfg(test)] -mod tests; +mod tests { + use diskann_utils::views::MatrixView; + use diskann_vector::distance::Metric; + use half::f16; + use std::collections::BTreeSet; + + use super::{ + DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, + build_leaf_candidates, + }; + + fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + MatrixView::try_from(data, rows, columns).unwrap() + } + + fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() + } + + fn build( + data: MatrixView<'_, T>, + leaves: &[Vec], + k: usize, + metric: Metric, + ) -> Result>, LeafBuildError> + where + T: crate::utils::VectorRepr + 'static, + { + pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) + } + + fn adjacency_lists(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() + } + + fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { + let mut graph = vec![BTreeSet::new(); data.len()]; + for (source, left) in data.iter().enumerate() { + let mut nearest: Vec<_> = data + .iter() + .enumerate() + .filter(|(target, _)| *target != source) + .map(|(target, right)| { + let distance = left + .iter() + .zip(right) + .map(|(x, y)| (x - y) * (x - y)) + .sum::(); + (target, distance) + }) + .collect(); + nearest.sort_by(|left, right| { + left.1 + .total_cmp(&right.1) + .then_with(|| left.0.cmp(&right.0)) + }); + for &(target, _) in nearest.iter().take(k) { + graph[source].insert(target as u32); + graph[target].insert(source as u32); + } + } + graph + .into_iter() + .map(|neighbors| neighbors.into_iter().collect()) + .collect() + } + + #[test] + fn leaf_adjacency_matches_an_independent_all_pairs_reference() { + let points = [ + [0.0_f32, 0.0], + [1.0, 0.2], + [3.1, 0.5], + [7.8, 1.4], + [-2.3, 4.1], + [6.7, -3.2], + ]; + let flat: Vec<_> = points.into_iter().flatten().collect(); + + let actual = adjacency_lists( + build( + view(&flat, points.len(), 2), + &[(0..points.len() as u32).collect()], + 2, + Metric::L2, + ) + .unwrap(), + ); + + assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); + } + + #[test] + fn retains_and_deduplicates_candidates_from_overlapping_leaves() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; + + let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + adjacency_lists(graph), + [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] + ); + } + + #[test] + fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { + let dimensions = 9; + let mut data = vec![0.0_f32; 10 * dimensions]; + for source in 1..10 { + data[source * dimensions + source - 1] = 1.0; + } + + let graph = build( + view(&data, 10, dimensions), + &[(0..10).collect()], + 1, + Metric::L2, + ) + .unwrap(); + + assert_eq!(&*graph[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); + assert!(graph.iter().enumerate().all(|(source, neighbors)| { + neighbors.iter().all(|&target| target as usize != source) + && neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + })); + } + + #[test] + fn global_id_translation_is_independent_of_leaf_order() { + let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; + let leaves = vec![vec![4, 1, 3]]; + + let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + adjacency_lists(graph), + [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] + ); + } + + fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> + where + T: crate::utils::VectorRepr + 'static, + { + let leaves = vec![(0..points as u32).collect()]; + adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) + } + + fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) + where + T: crate::utils::VectorRepr + 'static, + { + let points = 8; + // Source dimension controls VectorRepr conversion chunking. Cover tails on + // both sides of 4-, 8-, and 16-element boundaries, then a second 16-lane + // chunk. Input integers remain exact in every tested representation. + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let source = index / dimensions; + let dimension = index % dimensions; + ((source * 7 + dimension * 3 + source * dimension) % 23) as u8 + }) + .collect(); + let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + assert_eq!( + source_graph(&converted, points, dimensions), + source_graph(&f32_data, points, dimensions), + "{label} dimensions={dimensions}" + ); + } + } + + #[test] + fn f16_conversion_matches_f32_across_dimension_boundaries() { + assert_source_conversion_matches_f32("f16", |value| f16::from_f32(value as f32)); + } + + #[test] + fn u8_conversion_matches_f32_across_dimension_boundaries() { + assert_source_conversion_matches_f32("u8", |value| value); + } + + #[test] + fn i8_conversion_matches_f32_across_dimension_boundaries() { + // Applying the same translation to every coordinate preserves L2 pair + // ordering while exercising signed conversion. + assert_source_conversion_matches_f32("i8", |value| value as i8 - 11); + } + + #[test] + fn all_metrics_produce_symmetric_unique_non_self_candidates() { + let data = [1.0_f32, 0.0, 0.8, 0.2, 0.0, 1.0, -1.0, 0.0]; + let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); + for (source, neighbors) in graph.iter().enumerate() { + assert!(neighbors.iter().all(|&target| target as usize != source)); + assert!( + neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + ); + assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); + } + } + } + + #[test] + fn parallel_leaf_schedule_does_not_change_candidate_order() { + let data: Vec = (0..64).map(|value| value as f32).collect(); + let leaves: Vec> = (0..32) + .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) + .collect(); + let pool = pool(); + pool.install(|| { + let expected = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); + for _ in 0..8 { + let actual = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2) + .unwrap(); + assert_eq!(actual, expected); + } + }); + } + + #[test] + fn rejects_invalid_dimensions_and_leaf_membership() { + let data = [0.0_f32, 1.0]; + let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); + assert!(matches!( + build(no_dimensions, &[], 1, Metric::L2), + Err(LeafBuildError::EmptyDimensions) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), + Err(LeafBuildError::EmptyLeaf { leaf: 0 }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { + leaf: 0, + point: 2, + points: 2 + }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![1, 0, 1]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 1 }) + )); + } + + #[test] + fn singleton_and_zero_k_leaves_add_no_candidates() { + let data = [0.0_f32, 1.0, 2.0]; + let singleton = build( + view(&data, 3, 1), + &[vec![0], vec![1], vec![2]], + 1, + Metric::L2, + ) + .unwrap(); + let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); + assert!( + singleton + .iter() + .chain(&zero_k) + .all(|candidates| candidates.is_empty()) + ); + } + + #[test] + fn reuses_worker_buffers_for_smaller_leaves() { + let mut buffers = LeafBuffers::default(); + buffers.prepare(0, 64, 128, 2).unwrap(); + let point_values = buffers.point_values.as_ptr(); + let dots = buffers.dots.as_ptr(); + let neighbors = buffers.neighbors.as_ptr(); + + buffers.prepare(1, 8, 128, 2).unwrap(); + + assert_eq!(buffers.point_values.as_ptr(), point_values); + assert_eq!(buffers.dots.as_ptr(), dots); + assert_eq!(buffers.neighbors.as_ptr(), neighbors); + assert_eq!(buffers.point_values.len(), 64 * 128); + assert_eq!(buffers.dots.len(), 64 * 64); + assert_eq!(buffers.neighbors.len(), 64 * 2); + } + + #[test] + fn reports_shape_overflow_before_allocating() { + let mut buffers = LeafBuffers::default(); + assert!(matches!( + buffers.prepare(7, usize::MAX, 2, 1), + Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) + )); + } + + #[test] + fn rejects_an_invalid_kernel_target() { + let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; + let error = add_symmetric_neighbors( + &[10, 20], + 1, + &[ + super::super::leaf_kernel::LeafNeighbor::new(9, 1.0), + super::super::leaf_kernel::LeafNeighbor::new(0, 1.0), + ], + &mut graph, + ) + .unwrap_err(); + assert!(matches!( + error, + LeafBuildError::InvalidLocalTarget { + target: 9, + points: 2 + } + )); + } + + #[test] + fn skips_duplicate_global_ids_without_self_edges() { + let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; + add_symmetric_neighbors( + &[7, 7], + 1, + &[ + super::super::leaf_kernel::LeafNeighbor::new(1, 0.0), + super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), + ], + &mut graph, + ) + .unwrap(); + assert!(graph.iter().all(|neighbors| neighbors.is_empty())); + } + + #[test] + fn poisoned_candidate_lists_return_errors() { + let candidates = DirectCandidates::new(1).unwrap(); + let _ = std::panic::catch_unwind(|| { + let _guard = candidates.lists[0].lock().unwrap(); + panic!("poison candidate list"); + }); + assert!(matches!( + candidates.add_leaf(&[0], &[crate::graph::AdjacencyList::new()]), + Err(LeafBuildError::PoisonedCandidateList { point: 0 }) + )); + assert!(matches!( + candidates.into_lists(), + Err(LeafBuildError::PoisonedCandidateList { point: 0 }) + )); + } + + #[test] + fn allocation_errors_preserve_buffer_context() { + let mut values = Vec::::new(); + let source = values.try_reserve(usize::MAX).unwrap_err(); + let error = allocation_error("test", 1, source); + assert!(matches!( + error, + LeafBuildError::Allocation { + buffer: "test", + additional: 1, + .. + } + )); + } + + #[test] + fn direct_candidate_accumulator_keeps_unique_sorted_lists() { + let candidates = DirectCandidates::new(2).unwrap(); + candidates + .add_leaf( + &[0, 1], + &[ + crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), + crate::graph::AdjacencyList::from_iter_untrusted([0]), + ], + ) + .unwrap(); + assert_eq!( + adjacency_lists(candidates.into_lists().unwrap()), + [vec![1], vec![0]] + ); + } +} diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs deleted file mode 100644 index 6c55e6a577..0000000000 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ /dev/null @@ -1,414 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use diskann_utils::views::MatrixView; -use diskann_vector::distance::Metric; -use half::f16; -use std::collections::BTreeSet; - -use super::{ - DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, - build_leaf_candidates, -}; - -fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { - MatrixView::try_from(data, rows, columns).unwrap() -} - -fn pool() -> rayon::ThreadPool { - rayon::ThreadPoolBuilder::new() - .num_threads(4) - .build() - .unwrap() -} - -fn build( - data: MatrixView<'_, T>, - leaves: &[Vec], - k: usize, - metric: Metric, -) -> Result>, LeafBuildError> -where - T: crate::utils::VectorRepr + 'static, -{ - pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) -} - -fn adjacency_lists(graph: Vec>) -> Vec> { - graph.into_iter().map(Vec::from).collect() -} - -fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { - let mut graph = vec![BTreeSet::new(); data.len()]; - for (source, left) in data.iter().enumerate() { - let mut nearest: Vec<_> = data - .iter() - .enumerate() - .filter(|(target, _)| *target != source) - .map(|(target, right)| { - let distance = left - .iter() - .zip(right) - .map(|(x, y)| (x - y) * (x - y)) - .sum::(); - (target, distance) - }) - .collect(); - nearest.sort_by(|left, right| { - left.1 - .total_cmp(&right.1) - .then_with(|| left.0.cmp(&right.0)) - }); - for &(target, _) in nearest.iter().take(k) { - graph[source].insert(target as u32); - graph[target].insert(source as u32); - } - } - graph - .into_iter() - .map(|neighbors| neighbors.into_iter().collect()) - .collect() -} - -#[test] -fn leaf_adjacency_matches_an_independent_all_pairs_reference() { - let points = [ - [0.0_f32, 0.0], - [1.0, 0.2], - [3.1, 0.5], - [7.8, 1.4], - [-2.3, 4.1], - [6.7, -3.2], - ]; - let flat: Vec<_> = points.into_iter().flatten().collect(); - - let actual = adjacency_lists( - build( - view(&flat, points.len(), 2), - &[(0..points.len() as u32).collect()], - 2, - Metric::L2, - ) - .unwrap(), - ); - - assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); -} - -#[test] -fn retains_and_deduplicates_candidates_from_overlapping_leaves() { - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; - - let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); - - assert_eq!( - adjacency_lists(graph), - [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] - ); -} - -#[test] -fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { - let dimensions = 9; - let mut data = vec![0.0_f32; 10 * dimensions]; - for source in 1..10 { - data[source * dimensions + source - 1] = 1.0; - } - - let graph = build( - view(&data, 10, dimensions), - &[(0..10).collect()], - 1, - Metric::L2, - ) - .unwrap(); - - assert_eq!(&*graph[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); - assert!(graph.iter().enumerate().all(|(source, neighbors)| { - neighbors.iter().all(|&target| target as usize != source) - && neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32)) - })); -} - -#[test] -fn global_id_translation_is_independent_of_leaf_order() { - let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; - let leaves = vec![vec![4, 1, 3]]; - - let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); - - assert_eq!( - adjacency_lists(graph), - [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] - ); -} - -fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> -where - T: crate::utils::VectorRepr + 'static, -{ - let leaves = vec![(0..points as u32).collect()]; - adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) -} - -fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) -where - T: crate::utils::VectorRepr + 'static, -{ - let points = 8; - // Source dimension controls VectorRepr conversion chunking. Cover tails on - // both sides of 4-, 8-, and 16-element boundaries, then a second 16-lane - // chunk. Input integers remain exact in every tested representation. - for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let raw: Vec = (0..points * dimensions) - .map(|index| { - let source = index / dimensions; - let dimension = index % dimensions; - ((source * 7 + dimension * 3 + source * dimension) % 23) as u8 - }) - .collect(); - let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); - let converted: Vec = raw.iter().copied().map(&convert).collect(); - assert_eq!( - source_graph(&converted, points, dimensions), - source_graph(&f32_data, points, dimensions), - "{label} dimensions={dimensions}" - ); - } -} - -#[test] -fn f16_conversion_matches_f32_across_dimension_boundaries() { - assert_source_conversion_matches_f32("f16", |value| f16::from_f32(value as f32)); -} - -#[test] -fn u8_conversion_matches_f32_across_dimension_boundaries() { - assert_source_conversion_matches_f32("u8", |value| value); -} - -#[test] -fn i8_conversion_matches_f32_across_dimension_boundaries() { - // Applying the same translation to every coordinate preserves L2 pair - // ordering while exercising signed conversion. - assert_source_conversion_matches_f32("i8", |value| value as i8 - 11); -} - -#[test] -fn all_metrics_produce_symmetric_unique_non_self_candidates() { - let data = [1.0_f32, 0.0, 0.8, 0.2, 0.0, 1.0, -1.0, 0.0]; - let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; - - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); - for (source, neighbors) in graph.iter().enumerate() { - assert!(neighbors.iter().all(|&target| target as usize != source)); - assert!( - neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32)) - ); - assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); - } - } -} - -#[test] -fn parallel_leaf_schedule_does_not_change_candidate_order() { - let data: Vec = (0..64).map(|value| value as f32).collect(); - let leaves: Vec> = (0..32) - .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) - .collect(); - let pool = pool(); - pool.install(|| { - let expected = - build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); - for _ in 0..8 { - let actual = - build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); - assert_eq!(actual, expected); - } - }); -} - -#[test] -fn rejects_invalid_dimensions_and_leaf_membership() { - let data = [0.0_f32, 1.0]; - let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); - assert!(matches!( - build(no_dimensions, &[], 1, Metric::L2), - Err(LeafBuildError::EmptyDimensions) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), - Err(LeafBuildError::EmptyLeaf { leaf: 0 }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), - Err(LeafBuildError::InvalidPointId { - leaf: 0, - point: 2, - points: 2 - }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), - Err(LeafBuildError::InvalidPointId { point: 2, .. }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), - Err(LeafBuildError::InvalidPointId { point: 2, .. }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), - Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![1, 0, 1]], 1, Metric::L2), - Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 1 }) - )); -} - -#[test] -fn singleton_and_zero_k_leaves_add_no_candidates() { - let data = [0.0_f32, 1.0, 2.0]; - let singleton = build( - view(&data, 3, 1), - &[vec![0], vec![1], vec![2]], - 1, - Metric::L2, - ) - .unwrap(); - let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); - assert!( - singleton - .iter() - .chain(&zero_k) - .all(|candidates| candidates.is_empty()) - ); -} - -#[test] -fn reuses_worker_buffers_for_smaller_leaves() { - let mut buffers = LeafBuffers::default(); - buffers.prepare(0, 64, 128, 2).unwrap(); - let point_values = buffers.point_values.as_ptr(); - let dots = buffers.dots.as_ptr(); - let neighbors = buffers.neighbors.as_ptr(); - - buffers.prepare(1, 8, 128, 2).unwrap(); - - assert_eq!(buffers.point_values.as_ptr(), point_values); - assert_eq!(buffers.dots.as_ptr(), dots); - assert_eq!(buffers.neighbors.as_ptr(), neighbors); - assert_eq!(buffers.point_values.len(), 64 * 128); - assert_eq!(buffers.dots.len(), 64 * 64); - assert_eq!(buffers.neighbors.len(), 64 * 2); -} - -#[test] -fn reports_shape_overflow_before_allocating() { - let mut buffers = LeafBuffers::default(); - assert!(matches!( - buffers.prepare(7, usize::MAX, 2, 1), - Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) - )); -} - -#[test] -fn rejects_an_invalid_kernel_target() { - let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; - let error = add_symmetric_neighbors( - &[10, 20], - 1, - &[ - super::super::leaf_kernel::LeafNeighbor::new(9, 1.0), - super::super::leaf_kernel::LeafNeighbor::new(0, 1.0), - ], - &mut graph, - ) - .unwrap_err(); - assert!(matches!( - error, - LeafBuildError::InvalidLocalTarget { - target: 9, - points: 2 - } - )); -} - -#[test] -fn skips_duplicate_global_ids_without_self_edges() { - let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; - add_symmetric_neighbors( - &[7, 7], - 1, - &[ - super::super::leaf_kernel::LeafNeighbor::new(1, 0.0), - super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), - ], - &mut graph, - ) - .unwrap(); - assert!(graph.iter().all(|neighbors| neighbors.is_empty())); -} - -#[test] -fn poisoned_candidate_lists_return_errors() { - let candidates = DirectCandidates::new(1).unwrap(); - let _ = std::panic::catch_unwind(|| { - let _guard = candidates.lists[0].lock().unwrap(); - panic!("poison candidate list"); - }); - assert!(matches!( - candidates.add_leaf(&[0], &[crate::graph::AdjacencyList::new()]), - Err(LeafBuildError::PoisonedCandidateList { point: 0 }) - )); - assert!(matches!( - candidates.into_lists(), - Err(LeafBuildError::PoisonedCandidateList { point: 0 }) - )); -} - -#[test] -fn allocation_errors_preserve_buffer_context() { - let mut values = Vec::::new(); - let source = values.try_reserve(usize::MAX).unwrap_err(); - let error = allocation_error("test", 1, source); - assert!(matches!( - error, - LeafBuildError::Allocation { - buffer: "test", - additional: 1, - .. - } - )); -} - -#[test] -fn direct_candidate_accumulator_keeps_unique_sorted_lists() { - let candidates = DirectCandidates::new(2).unwrap(); - candidates - .add_leaf( - &[0, 1], - &[ - crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), - crate::graph::AdjacencyList::from_iter_untrusted([0]), - ], - ) - .unwrap(); - assert_eq!( - adjacency_lists(candidates.into_lists().unwrap()), - [vec![1], vec![0]] - ); -} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 9073e2157c..6112fc0f8e 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -358,13 +358,11 @@ where T: VectorRepr + Send + Sync + 'static, { if data.nrows() == 0 { - return Err(ANNError::log_dimension_mismatch_error( - "PiPNN requires at least one data point".into(), - )); + return Err(ANNError::message("PiPNN requires at least one data point")); } if data.ncols() == 0 { - return Err(ANNError::log_dimension_mismatch_error( - "PiPNN requires at least one data dimension".into(), + return Err(ANNError::message( + "PiPNN requires at least one data dimension", )); } if data.nrows() > u32::MAX as usize { @@ -374,7 +372,7 @@ where ))); } data.nrows().checked_mul(data.ncols()).ok_or_else(|| { - ANNError::log_dimension_mismatch_error(format!( + ANNError::message(format!( "PiPNN dataset shape {} x {} overflows usize", data.nrows(), data.ncols() @@ -390,7 +388,7 @@ where // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) - .map_err(ANNError::opaque) + .map_err(ANNError::new) })?; // Finalization consumes candidate lists and reuses their allocations for the // resulting adjacency where possible. @@ -412,8 +410,382 @@ fn effective_metric(metric: Metric) -> Metric { #[track_caller] fn config_error(message: impl std::fmt::Display) -> ANNError { - ANNError::log_index_config_error("PiPNN".into(), message.to_string()) + ANNError::message(format!("PiPNN configuration: {message}")) } #[cfg(test)] -mod tests; +mod tests { + use super::*; + use half::f16; + + #[test] + fn integer_normalized_cosine_uses_unnormalized_cosine() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let expected = if metric == Metric::CosineNormalized { + Metric::Cosine + } else { + metric + }; + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), metric); + assert_eq!(effective_metric::(metric), metric); + } + } +} +#[cfg(test)] +#[allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] +mod build_graph_tests { + use super::{PiPNNBuildContext, PiPNNConfig, build_graph}; + use crate::graph::config::{self, MaxDegree}; + use diskann_utils::views::MatrixView; + use diskann_vector::distance::Metric; + use half::f16; + use rand::{Rng, SeedableRng, rngs::StdRng}; + + fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 4, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + } + } + + fn graph_config(metric: Metric, degree: usize) -> crate::graph::Config { + config::Builder::new_with(degree, MaxDegree::same(), 8, metric.into(), |builder| { + builder.alpha(1.2); + }) + .build() + .unwrap() + } + + fn pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + } + + fn rows(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() + } + + fn assert_graph_invariants( + graph: &[crate::graph::AdjacencyList], + points: usize, + degree: usize, + ) { + assert_eq!(graph.len(), points); + for (source, row) in graph.iter().enumerate() { + assert!(row.len() <= degree); + let mut sorted = row.to_vec(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!(sorted.len(), row.len()); + assert!( + row.iter() + .all(|&id| (id as usize) < points && id as usize != source) + ); + } + } + + #[test] + fn builds_a_single_leaf_graph_for_real_dataset_ids() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let graph = graph_config(Metric::L2, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); + + let graph = graph_config(Metric::L2, 1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let pruned = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&pruned, 4, 1); + for (source, neighbors) in pruned.iter().enumerate() { + assert_eq!(source.abs_diff(neighbors[0] as usize), 1); + } + } + + #[test] + fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { + let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; + let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); + let graph = graph_config(Metric::L2, 1); + let pool = pool(2); + let config = PiPNNConfig { + c_max: 5, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 4, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&actual, 5, 1); + assert!(actual.iter().all(|row| row.len() == 1)); + } + + #[test] + fn rejects_empty_dataset_dimensions_at_the_public_boundary() { + let graph = graph_config(Metric::L2, 2); + let pool = pool(1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); + let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); + + assert!(build_graph(no_rows, &context).is_err()); + assert!(build_graph(no_columns, &context).is_err()); + } + + #[test] + fn supports_every_source_type_and_metric() { + fn build( + values: &[T], + metric: Metric, + ) { + let data = MatrixView::try_from(values, 6, 2).unwrap(); + let graph = graph_config(metric, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); + let actual = build_graph(data, &context).unwrap(); + assert_graph_invariants(&actual, 6, 2); + } + + let values = [ + 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, + ]; + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + build(&values, metric); + } + build(&values.map(f16::from_f32), Metric::L2); + build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); + build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); + } + + #[test] + fn integer_normalized_cosine_matches_cosine() { + fn assert_match(values: &[T]) { + let data = MatrixView::try_from(values, 8, 2).unwrap(); + let pool = pool(2); + let build = |metric| { + let graph = graph_config(metric, 2); + let config = PiPNNConfig { + c_max: 8, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + rows(build_graph(data, &context).unwrap()) + }; + assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); + } + + assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); + assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); + } + + #[test] + fn is_deterministic_for_a_fixed_pool_size() { + let data: Vec = (0..96 * 4) + .map(|value| ((value * 17 + 3) % 101) as f32) + .collect(); + let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); + let graph = graph_config(Metric::L2, 8); + let pool = pool(4); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.25, + fanout: vec![3, 2], + k: 3, + replicas: 2, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let first = build_graph(data, &context).unwrap(); + let second = build_graph(data, &context).unwrap(); + + assert_eq!(first, second); + assert_graph_invariants(&first, 96, 8); + } + + #[test] + fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { + let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); + for case in 0..24 { + let points = rng.random_range(4..=32); + let dimensions = rng.random_range(1..=8); + let c_max = rng.random_range(4..=points.min(12)); + let c_min = rng.random_range(1..=c_max); + let degree = rng.random_range(1..=points.min(8)); + let values: Vec = (0..points * dimensions) + .map(|_| rng.random_range(-10.0..10.0)) + .collect(); + let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); + let graph = graph_config(Metric::L2, degree); + let pool = pool(2); + let config = PiPNNConfig { + c_max, + c_min, + p_samp: 0.5, + fanout: vec![2], + k: rng.random_range(1..=3), + replicas: rng.random_range(1..=2), + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context) + .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); + assert_graph_invariants(&actual, points, degree); + } + } +} +#[cfg(test)] +#[allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] +mod config_tests { + use super::{PiPNNBuildContext, PiPNNConfig}; + use crate::graph::config::{self, MaxDegree}; + use diskann_vector::distance::Metric; + + fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 512, + c_min: 64, + p_samp: 0.01, + fanout: vec![10, 3], + k: 2, + replicas: 1, + } + } + + fn graph_config(metric: Metric, alpha: f32) -> crate::graph::Config { + config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { + builder.alpha(alpha); + }) + .build() + .unwrap() + } + + fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap() + } + + #[test] + fn rejects_each_invalid_algorithm_parameter() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + let mut cases = [ + PiPNNConfig { + c_max: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 513, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 0.0, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: -0.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 1.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: f64::NAN, + ..pipnn_config() + }, + PiPNNConfig { + fanout: Vec::new(), + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![1, 0], + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![17], + ..pipnn_config() + }, + PiPNNConfig { + k: 0, + ..pipnn_config() + }, + PiPNNConfig { + replicas: 0, + ..pipnn_config() + }, + ]; + + for config in &mut cases { + let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .expect_err("invalid PiPNN config must be rejected"); + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + } + } + + #[test] + fn rejects_graph_policy_for_a_different_metric() { + let graph = graph_config(Metric::InnerProduct, 1.2); + let pool = pool(); + + let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + assert!(error.to_string().contains("prune kind")); + } + + #[test] + fn does_not_add_alpha_validation_beyond_graph_config() { + let pool = pool(); + for alpha in [0.9, f32::NAN, f32::INFINITY] { + let graph = graph_config(Metric::L2, alpha); + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + } + } +} diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 96ed91878b..5de9a68cf5 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -149,13 +149,13 @@ where { let points = data.nrows(); if points == 0 { - return Err(ANNError::opaque(PartitionError::EmptyDataset)); + return Err(ANNError::new(PartitionError::EmptyDataset)); } if data.ncols() == 0 { - return Err(ANNError::opaque(PartitionError::EmptyDimensions)); + return Err(ANNError::new(PartitionError::EmptyDimensions)); } if points > u32::MAX as usize { - return Err(ANNError::opaque(PartitionError::TooManyPoints(points))); + return Err(ANNError::new(PartitionError::TooManyPoints(points))); } let mut leaves = Vec::new(); @@ -170,7 +170,7 @@ where partition_replica(data, &config, metric, &kernel, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; leaves.append(&mut replica_leaves); } validate_leaves(&leaves, config.c_max)?; @@ -191,14 +191,14 @@ where let initial_indices = point_ids(data.nrows())?; if data.nrows() <= config.c_max { let mut leaves = Vec::new(); - leaves.try_reserve_exact(1).map_err(ANNError::opaque)?; + leaves.try_reserve_exact(1).map_err(ANNError::new)?; leaves.push(initial_indices); return Ok(leaves); } let mut leaves = Vec::new(); let mut work = Vec::new(); - work.try_reserve_exact(1).map_err(ANNError::opaque)?; + work.try_reserve_exact(1).map_err(ANNError::new)?; work.push(WorkItem { indices: initial_indices, level: 0, @@ -213,7 +213,7 @@ where let mut results = Vec::new(); results .try_reserve_exact(work.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; results.resize_with(work.len(), || None); // build_graph installs this complete private call tree into the // caller-owned pool; the indexed fill cannot escape that pool. @@ -235,13 +235,11 @@ where let mut next_work = Vec::new(); for result in results { let (mut pending, mut finished) = - result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??; + result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))??; next_work .try_reserve(pending.len()) - .map_err(ANNError::opaque)?; - leaves - .try_reserve(finished.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; + leaves.try_reserve(finished.len()).map_err(ANNError::new)?; next_work.append(&mut pending); leaves.append(&mut finished); } @@ -254,7 +252,7 @@ where let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { return global_merge_small(leaves, config.c_min, config.c_max); }; - Err(ANNError::opaque(PartitionError::IterationLimit { + Err(ANNError::new(PartitionError::IterationLimit { size: largest.indices.len(), level: largest.level, limit: MAX_PARTITION_ITERATIONS, @@ -291,12 +289,10 @@ where let mut pending = Vec::new(); let mut finished = Vec::new(); - pending - .try_reserve(clusters.len()) - .map_err(ANNError::opaque)?; + pending.try_reserve(clusters.len()).map_err(ANNError::new)?; finished .try_reserve(clusters.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; let child_seed = mix_seed(item.seed, points as u64); for cluster in clusters { if cluster.is_empty() { @@ -319,7 +315,7 @@ fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResul let count = sample_num_leaders(points.len(), sampling_fraction); let mut rng = rand::rngs::StdRng::seed_from_u64(seed); let mut leaders = Vec::new(); - leaders.try_reserve_exact(count).map_err(ANNError::opaque)?; + leaders.try_reserve_exact(count).map_err(ANNError::new)?; leaders.extend(points.choose_multiple(&mut rng, count).copied()); Ok(leaders) } @@ -478,7 +474,7 @@ where None, dots, ) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; let point_scales = if metric == Metric::Cosine { grow_fallible(point_scale_buffer, point_count, 0.0)?; @@ -504,14 +500,14 @@ where Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, }; let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { - ANNError::opaque(PartitionError::InvalidBufferLength { + ANNError::new(PartitionError::InvalidBufferLength { buffer: "dot-product stripe", expected: dots_len, actual: dots.len(), }) })?; let output = MutMatrixView::try_from(assignments, point_count, fanout).map_err(|error| { - ANNError::opaque(PartitionError::InvalidBufferLength { + ANNError::new(PartitionError::InvalidBufferLength { buffer: "partition assignments", expected: output_len, actual: error.into_inner().len(), @@ -519,7 +515,7 @@ where })?; kernel .nearest_leaders(PartitionInput { dots, scales }, output) - .map_err(ANNError::opaque) + .map_err(ANNError::new) } fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> @@ -528,7 +524,7 @@ where { let expected = checked_area("gather output", indices.len(), data.ncols())?; if output.len() != expected { - return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + return Err(ANNError::new(PartitionError::InvalidBufferLength { buffer: "gather output", expected, actual: output.len(), @@ -560,9 +556,7 @@ fn scatter_assignments( let stripe_assignment_count = checked_area("scatter assignment stripe", stripe_points, fanout)?; let stripes = points.len().div_ceil(stripe_points); let mut partials = Vec::new(); - partials - .try_reserve_exact(stripes) - .map_err(ANNError::opaque)?; + partials.try_reserve_exact(stripes).map_err(ANNError::new)?; partials.resize_with(stripes, || None); // See the pool invariant at the other partition terminal operations. #[allow(clippy::disallowed_methods)] @@ -578,18 +572,16 @@ fn scatter_assignments( }); let mut locals = Vec::new(); - locals - .try_reserve_exact(stripes) - .map_err(ANNError::opaque)?; + locals.try_reserve_exact(stripes).map_err(ANNError::new)?; for result in partials { - locals.push(result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??); + locals.push(result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))??); } let mut sizes = filled_vec(leaders, 0usize)?; for local in &locals { for (size, cluster) in sizes.iter_mut().zip(local) { *size = size.checked_add(cluster.len()).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { + ANNError::new(PartitionError::ShapeOverflow { buffer: "cluster size", rows: *size, cols: cluster.len(), @@ -605,7 +597,7 @@ fn scatter_assignments( .enumerate() .map(|(leader, size)| { let mut cluster = Vec::new(); - cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + cluster.try_reserve_exact(size).map_err(ANNError::new)?; for local in &locals { cluster.extend_from_slice(&local[leader]); } @@ -623,14 +615,14 @@ fn scatter_serial( let mut sizes = filled_vec(leaders, 0usize)?; for &leader in assignments { let Some(size) = sizes.get_mut(leader as usize) else { - return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + return Err(ANNError::new(PartitionError::InvalidBufferLength { buffer: "leader assignment", expected: leaders, actual: leader as usize + 1, })); }; *size = size.checked_add(1).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { + ANNError::new(PartitionError::ShapeOverflow { buffer: "cluster size", rows: *size, cols: 1, @@ -650,10 +642,10 @@ fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { let mut clusters = Vec::new(); clusters .try_reserve_exact(sizes.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; for &size in sizes { let mut cluster = Vec::new(); - cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + cluster.try_reserve_exact(size).map_err(ANNError::new)?; clusters.push(cluster); } Ok(clusters) @@ -666,10 +658,10 @@ fn global_merge_small( ) -> ANNResult>> { let mut merged = Vec::new(); let mut small_leaves = Vec::new(); - merged.try_reserve(leaves.len()).map_err(ANNError::opaque)?; + merged.try_reserve(leaves.len()).map_err(ANNError::new)?; small_leaves .try_reserve(leaves.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; for leaf in leaves { if leaf.len() >= c_min { merged.push(leaf); @@ -682,11 +674,11 @@ fn global_merge_small( } let mut small = HashSet::new(); - small.try_reserve(c_max).map_err(ANNError::opaque)?; + small.try_reserve(c_max).map_err(ANNError::new)?; for leaf in small_leaves { let combined = small.len().checked_add(leaf.len()).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { + ANNError::new(PartitionError::ShapeOverflow { buffer: "small-leaf merge", rows: small.len(), cols: leaf.len(), @@ -695,7 +687,7 @@ fn global_merge_small( if combined > c_max { merged.push(drain_sorted(&mut small)?); } - small.try_reserve(leaf.len()).map_err(ANNError::opaque)?; + small.try_reserve(leaf.len()).map_err(ANNError::new)?; small.extend(leaf); if small.len() >= c_min { merged.push(drain_sorted(&mut small)?); @@ -709,15 +701,14 @@ fn global_merge_small( { remainder.retain(|id| !last.contains(id)); let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { + ANNError::new(PartitionError::ShapeOverflow { buffer: "small-leaf tail merge", rows: last.len(), cols: remainder.len(), }) })?; if combined <= c_max { - last.try_reserve(remainder.len()) - .map_err(ANNError::opaque)?; + last.try_reserve(remainder.len()).map_err(ANNError::new)?; last.append(&mut remainder); last.sort_unstable(); } @@ -733,9 +724,7 @@ fn global_merge_small( fn drain_sorted(set: &mut HashSet) -> ANNResult> { let mut values = Vec::new(); - values - .try_reserve_exact(set.len()) - .map_err(ANNError::opaque)?; + values.try_reserve_exact(set.len()).map_err(ANNError::new)?; values.extend(set.drain()); values.sort_unstable(); Ok(values) @@ -746,7 +735,7 @@ fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { .iter() .find(|leaf| leaf.is_empty() || leaf.len() > c_max) { - return Err(ANNError::opaque(PartitionError::InvalidLeaf { + return Err(ANNError::new(PartitionError::InvalidLeaf { size: leaf.len(), limit: c_max, })); @@ -756,14 +745,14 @@ fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { fn point_ids(points: usize) -> ANNResult> { let mut ids = Vec::new(); - ids.try_reserve_exact(points).map_err(ANNError::opaque)?; + ids.try_reserve_exact(points).map_err(ANNError::new)?; ids.extend(0..points as u32); Ok(ids) } fn filled_vec(len: usize, value: T) -> ANNResult> { let mut values = Vec::new(); - values.try_reserve_exact(len).map_err(ANNError::opaque)?; + values.try_reserve_exact(len).map_err(ANNError::new)?; values.resize(len, value); Ok(values) } @@ -779,14 +768,14 @@ fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResu } values .try_reserve(len - values.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; values.resize(len, value); Ok(()) } fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { rows.checked_mul(cols) - .ok_or_else(|| ANNError::opaque(PartitionError::ShapeOverflow { buffer, rows, cols })) + .ok_or_else(|| ANNError::new(PartitionError::ShapeOverflow { buffer, rows, cols })) } fn assignment_stripe_point_count(leader_count: usize) -> usize { @@ -800,4 +789,423 @@ fn assignment_stripe_point_count(leader_count: usize) -> usize { } #[cfg(test)] -mod tests; +mod tests { + use diskann_utils::views::{Matrix, MatrixView}; + use diskann_vector::{Half, distance::Metric}; + + use super::*; + + fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { + PiPNNConfig { + c_max, + c_min, + p_samp: 0.25, + fanout, + k: 1, + replicas, + } + } + + fn clustered_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let point = position / dimensions; + let dimension = position % dimensions; + position += 1; + (point / 8) as f32 * 10.0 + dimension as f32 * 0.01 + point as f32 * 0.001 + } + }), + points, + dimensions, + ) + } + + fn directional_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let point = position / dimensions; + let dimension = position % dimensions; + position += 1; + let angle = std::f32::consts::TAU * point as f32 / points as f32; + match dimension { + 0 => angle.cos(), + 1 => angle.sin(), + _ => 0.0, + } + } + }), + points, + dimensions, + ) + } + + fn sorted_memberships(leaves: &[Vec]) -> Vec> { + let mut memberships: Vec> = leaves + .iter() + .map(|leaf| { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids + }) + .collect(); + memberships.sort(); + memberships + } + + fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { + assert!( + leaves + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) + ); + let mut counts = vec![0usize; points]; + for leaf in leaves { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids.len(), leaf.len(), "duplicate ID inside a leaf"); + for &id in leaf { + assert!((id as usize) < points); + counts[id as usize] += 1; + } + } + assert!(counts.iter().all(|&count| count >= replicas)); + } + + #[test] + fn returns_one_leaf_at_and_below_c_max() { + for points in [7, 8] { + let data = clustered_data(points, 3); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + assert_eq!(leaves, vec![(0..points as u32).collect::>()]); + } + } + + #[test] + fn partition_is_fixed_seed_deterministic_and_bounded() { + let data = clustered_data(96, 8); + let config = config(4, 16, vec![3, 2], 2); + + let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); + let second = partition(data.as_view(), config, Metric::L2).unwrap(); + + assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); + assert_valid_partition(&first, 96, 16, 2); + assert!(first.iter().map(Vec::len).sum::() > 96 * 2); + } + + #[test] + fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { + let data = clustered_data(80, 4); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + + assert_valid_partition(&leaves, 80, 8, 1); + } + + #[test] + fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { + let data = Matrix::new(1.0f32, 24, 4); + let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = error.downcast::().unwrap(); + + assert!(matches!( + error, + PartitionError::IterationLimit { + size: 24, + limit: MAX_PARTITION_ITERATIONS, + .. + } + )); + } + + #[test] + fn global_merge_canonicalizes_small_leaf_membership() { + let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; + + let merged = global_merge_small(leaves, 4, 8).unwrap(); + + assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); + } + + #[test] + fn global_merge_never_overfills_before_reaching_c_min() { + let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; + + let merged = global_merge_small(leaves, 11, 11).unwrap(); + + assert_eq!( + merged, + vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] + ); + } + + #[test] + fn global_merge_fills_exact_capacity_before_flushing() { + let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + + assert_eq!(merged, vec![vec![0, 1, 2, 3]]); + } + + #[test] + fn replicas_cover_every_point_once_or_more_per_replica() { + let data = directional_data(72, 5); + let leaves = partition( + data.as_view(), + config(3, 12, vec![3, 2], 3), + Metric::CosineNormalized, + ) + .unwrap(); + + assert_valid_partition(&leaves, 72, 12, 3); + } + + fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) + where + T: crate::utils::VectorRepr + Send + Sync, + { + let points = 64; + // Partition gathering converts source vectors before GEMM. Exercise conversion + // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let point = index / dimensions; + let dimension = index % dimensions; + ((point * 5 + dimension * 7 + point * dimension) % 23) as u8 + }) + .collect(); + let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + let config = config(2, 16, vec![2, 1], 1); + let expected = partition( + MatrixView::try_from(&f32_data, points, dimensions).unwrap(), + config.clone(), + Metric::L2, + ) + .unwrap(); + let actual = partition( + MatrixView::try_from(&converted, points, dimensions).unwrap(), + config, + Metric::L2, + ) + .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); + + assert_valid_partition(&actual, points, 16, 1); + assert_eq!( + sorted_memberships(&actual), + sorted_memberships(&expected), + "{label} dimensions={dimensions}" + ); + } + } + + #[test] + fn f16_partition_matches_f32_across_dimension_boundaries() { + assert_partition_conversion_matches_f32("f16", |value| Half::from_f32(value as f32)); + } + + #[test] + fn u8_partition_matches_f32_across_dimension_boundaries() { + assert_partition_conversion_matches_f32("u8", |value| value); + } + + #[test] + fn i8_partition_matches_f32_across_dimension_boundaries() { + // The same translation in every coordinate preserves L2 ordering. + assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); + } + + #[test] + fn l2_leader_norms_preserve_scalar_reduction_order() { + fn next(state: &mut u64) -> f32 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 + } + + // This fixed case sits on opposite sides of the top-1 boundary depending + // on whether leader norms use the original scalar reduction or a SIMD + // reassociation. Point/leader dot products still go through the production + // GEMM; only the setup norm calculation is under test. + let dimensions = 129; + let mut state = 0x3a85_f952_c718_6e49; + let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let leader_zero: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let leader_one: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let data: Vec = leader_zero + .into_iter() + .chain(leader_one) + .chain(point) + .collect(); + let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); + + let clusters = assign_to_leaders( + data, + &[2], + &[0, 1], + 1, + Metric::L2, + &PartitionKernel::new(Metric::L2), + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); + + assert_eq!(clusters, [vec![], vec![2]]); + } + + #[test] + fn all_metrics_produce_valid_partitions() { + let data = directional_data(64, 8); + let config = config(2, 20, vec![2], 1); + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); + assert_valid_partition(&leaves, 64, 20, 1); + } + } + + #[test] + fn leader_count_is_bounded() { + assert_eq!(sample_num_leaders(1, 1.0), 1); + assert_eq!(sample_num_leaders(10, 0.01), 2); + assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); + } + + #[test] + fn replica_seed_derivation_is_stable_and_distinct() { + assert_eq!(replica_seed(0), 1_000); + assert_eq!(replica_seed(1), 8_919); + } + + #[test] + fn assignment_stripes_use_power_of_two_point_counts() { + assert_eq!(assignment_stripe_point_count(1_000), 128); + assert_eq!(assignment_stripe_point_count(256), 512); + assert_eq!( + assignment_stripe_point_count(1), + MAX_ASSIGNMENT_STRIPE_POINTS + ); + } + + #[test] + fn stripe_buffer_pool_reuses_returned_capacity() { + let pool = StripeBufferPool::new((), 0, None); + let points = { + let mut buffers = pool.get_ref(()); + buffers.points.resize(16, 0.0); + buffers.points.as_ptr() + }; + + let buffers = pool.get_ref(()); + assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.points.len(), 16); + } + + #[test] + fn leader_assignment_handles_multiple_stripes() { + let points = 2_048; + let data: Vec = (0..points).map(|point| point as f32).collect(); + let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); + let point_ids: Vec = (0..points as u32).collect(); + + let clusters = assign_to_leaders( + data, + &point_ids, + &[0, 2_047], + 1, + Metric::L2, + &PartitionKernel::new(Metric::L2), + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); + + assert_eq!(clusters[0], (0..1_024).collect::>()); + assert_eq!(clusters[1], (1_024..2_048).collect::>()); + } + + #[test] + fn parallel_scatter_matches_serial_order() { + let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); + let assignments: Vec = points + .iter() + .flat_map(|point| [point % 7, (point + 3) % 7]) + .collect(); + + let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); + let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); + + assert_eq!(actual, expected); + } + + #[test] + fn rejects_empty_dataset() { + let data = Matrix::::new(0.0, 0, 4); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDataset + ); + } + + #[test] + fn rejects_zero_dimensions() { + let data = Matrix::::new(0.0, 4, 0); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDimensions + ); + } + + #[test] + fn rejects_invalid_gather_output_length() { + let data = Matrix::::new(0.0, 2, 2); + let error = gather_vectors(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "gather output", + expected: 4, + actual: 3, + } + ); + } + + #[test] + fn rejects_assignment_to_an_unknown_leader() { + let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "leader assignment", + expected: 2, + actual: 3, + } + ); + } + + #[test] + fn rejects_empty_and_oversized_leaves() { + for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { + let error = validate_leaves(&leaves, 2).unwrap_err(); + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidLeaf { size, limit: 2 } + ); + } + } +} diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs deleted file mode 100644 index 8780078ca5..0000000000 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ /dev/null @@ -1,423 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use diskann_utils::views::{Matrix, MatrixView}; -use diskann_vector::{Half, distance::Metric}; - -use super::*; - -fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { - PiPNNConfig { - c_max, - c_min, - p_samp: 0.25, - fanout, - k: 1, - replicas, - } -} - -fn clustered_data(points: usize, dimensions: usize) -> Matrix { - Matrix::new( - diskann_utils::views::Init({ - let mut position = 0usize; - move || { - let point = position / dimensions; - let dimension = position % dimensions; - position += 1; - (point / 8) as f32 * 10.0 + dimension as f32 * 0.01 + point as f32 * 0.001 - } - }), - points, - dimensions, - ) -} - -fn directional_data(points: usize, dimensions: usize) -> Matrix { - Matrix::new( - diskann_utils::views::Init({ - let mut position = 0usize; - move || { - let point = position / dimensions; - let dimension = position % dimensions; - position += 1; - let angle = std::f32::consts::TAU * point as f32 / points as f32; - match dimension { - 0 => angle.cos(), - 1 => angle.sin(), - _ => 0.0, - } - } - }), - points, - dimensions, - ) -} - -fn sorted_memberships(leaves: &[Vec]) -> Vec> { - let mut memberships: Vec> = leaves - .iter() - .map(|leaf| { - let mut ids = leaf.clone(); - ids.sort_unstable(); - ids - }) - .collect(); - memberships.sort(); - memberships -} - -fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { - assert!( - leaves - .iter() - .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) - ); - let mut counts = vec![0usize; points]; - for leaf in leaves { - let mut ids = leaf.clone(); - ids.sort_unstable(); - ids.dedup(); - assert_eq!(ids.len(), leaf.len(), "duplicate ID inside a leaf"); - for &id in leaf { - assert!((id as usize) < points); - counts[id as usize] += 1; - } - } - assert!(counts.iter().all(|&count| count >= replicas)); -} - -#[test] -fn returns_one_leaf_at_and_below_c_max() { - for points in [7, 8] { - let data = clustered_data(points, 3); - let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); - assert_eq!(leaves, vec![(0..points as u32).collect::>()]); - } -} - -#[test] -fn partition_is_fixed_seed_deterministic_and_bounded() { - let data = clustered_data(96, 8); - let config = config(4, 16, vec![3, 2], 2); - - let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); - let second = partition(data.as_view(), config, Metric::L2).unwrap(); - - assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); - assert_valid_partition(&first, 96, 16, 2); - assert!(first.iter().map(Vec::len).sum::() > 96 * 2); -} - -#[test] -fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { - let data = clustered_data(80, 4); - let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); - - assert_valid_partition(&leaves, 80, 8, 1); -} - -#[test] -fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { - let data = Matrix::new(1.0f32, 24, 4); - let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); - let error = error.downcast::().unwrap(); - - assert!(matches!( - error, - PartitionError::IterationLimit { - size: 24, - limit: MAX_PARTITION_ITERATIONS, - .. - } - )); -} - -#[test] -fn global_merge_canonicalizes_small_leaf_membership() { - let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; - - let merged = global_merge_small(leaves, 4, 8).unwrap(); - - assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); -} - -#[test] -fn global_merge_never_overfills_before_reaching_c_min() { - let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; - - let merged = global_merge_small(leaves, 11, 11).unwrap(); - - assert_eq!( - merged, - vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] - ); -} - -#[test] -fn global_merge_fills_exact_capacity_before_flushing() { - let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); - - assert_eq!(merged, vec![vec![0, 1, 2, 3]]); -} - -#[test] -fn replicas_cover_every_point_once_or_more_per_replica() { - let data = directional_data(72, 5); - let leaves = partition( - data.as_view(), - config(3, 12, vec![3, 2], 3), - Metric::CosineNormalized, - ) - .unwrap(); - - assert_valid_partition(&leaves, 72, 12, 3); -} - -fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) -where - T: crate::utils::VectorRepr + Send + Sync, -{ - let points = 64; - // Partition gathering converts source vectors before GEMM. Exercise conversion - // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. - for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let raw: Vec = (0..points * dimensions) - .map(|index| { - let point = index / dimensions; - let dimension = index % dimensions; - ((point * 5 + dimension * 7 + point * dimension) % 23) as u8 - }) - .collect(); - let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); - let converted: Vec = raw.iter().copied().map(&convert).collect(); - let config = config(2, 16, vec![2, 1], 1); - let expected = partition( - MatrixView::try_from(&f32_data, points, dimensions).unwrap(), - config.clone(), - Metric::L2, - ) - .unwrap(); - let actual = partition( - MatrixView::try_from(&converted, points, dimensions).unwrap(), - config, - Metric::L2, - ) - .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); - - assert_valid_partition(&actual, points, 16, 1); - assert_eq!( - sorted_memberships(&actual), - sorted_memberships(&expected), - "{label} dimensions={dimensions}" - ); - } -} - -#[test] -fn f16_partition_matches_f32_across_dimension_boundaries() { - assert_partition_conversion_matches_f32("f16", |value| Half::from_f32(value as f32)); -} - -#[test] -fn u8_partition_matches_f32_across_dimension_boundaries() { - assert_partition_conversion_matches_f32("u8", |value| value); -} - -#[test] -fn i8_partition_matches_f32_across_dimension_boundaries() { - // The same translation in every coordinate preserves L2 ordering. - assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); -} - -#[test] -fn l2_leader_norms_preserve_scalar_reduction_order() { - fn next(state: &mut u64) -> f32 { - *state ^= *state << 13; - *state ^= *state >> 7; - *state ^= *state << 17; - (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 - } - - // This fixed case sits on opposite sides of the top-1 boundary depending - // on whether leader norms use the original scalar reduction or a SIMD - // reassociation. Point/leader dot products still go through the production - // GEMM; only the setup norm calculation is under test. - let dimensions = 129; - let mut state = 0x3a85_f952_c718_6e49; - let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); - let leader_zero: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); - let leader_one: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); - let data: Vec = leader_zero - .into_iter() - .chain(leader_one) - .chain(point) - .collect(); - let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); - - let clusters = assign_to_leaders( - data, - &[2], - &[0, 1], - 1, - Metric::L2, - &PartitionKernel::new(Metric::L2), - &StripeBufferPool::new((), 0, None), - ) - .unwrap(); - - assert_eq!(clusters, [vec![], vec![2]]); -} - -#[test] -fn all_metrics_produce_valid_partitions() { - let data = directional_data(64, 8); - let config = config(2, 20, vec![2], 1); - - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); - assert_valid_partition(&leaves, 64, 20, 1); - } -} - -#[test] -fn leader_count_is_bounded() { - assert_eq!(sample_num_leaders(1, 1.0), 1); - assert_eq!(sample_num_leaders(10, 0.01), 2); - assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); -} - -#[test] -fn replica_seed_derivation_is_stable_and_distinct() { - assert_eq!(replica_seed(0), 1_000); - assert_eq!(replica_seed(1), 8_919); -} - -#[test] -fn assignment_stripes_use_power_of_two_point_counts() { - assert_eq!(assignment_stripe_point_count(1_000), 128); - assert_eq!(assignment_stripe_point_count(256), 512); - assert_eq!( - assignment_stripe_point_count(1), - MAX_ASSIGNMENT_STRIPE_POINTS - ); -} - -#[test] -fn stripe_buffer_pool_reuses_returned_capacity() { - let pool = StripeBufferPool::new((), 0, None); - let points = { - let mut buffers = pool.get_ref(()); - buffers.points.resize(16, 0.0); - buffers.points.as_ptr() - }; - - let buffers = pool.get_ref(()); - assert_eq!(buffers.points.as_ptr(), points); - assert_eq!(buffers.points.len(), 16); -} - -#[test] -fn leader_assignment_handles_multiple_stripes() { - let points = 2_048; - let data: Vec = (0..points).map(|point| point as f32).collect(); - let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); - let point_ids: Vec = (0..points as u32).collect(); - - let clusters = assign_to_leaders( - data, - &point_ids, - &[0, 2_047], - 1, - Metric::L2, - &PartitionKernel::new(Metric::L2), - &StripeBufferPool::new((), 0, None), - ) - .unwrap(); - - assert_eq!(clusters[0], (0..1_024).collect::>()); - assert_eq!(clusters[1], (1_024..2_048).collect::>()); -} - -#[test] -fn parallel_scatter_matches_serial_order() { - let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); - let assignments: Vec = points - .iter() - .flat_map(|point| [point % 7, (point + 3) % 7]) - .collect(); - - let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); - let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); - - assert_eq!(actual, expected); -} - -#[test] -fn rejects_empty_dataset() { - let data = Matrix::::new(0.0, 0, 4); - let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::EmptyDataset - ); -} - -#[test] -fn rejects_zero_dimensions() { - let data = Matrix::::new(0.0, 4, 0); - let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::EmptyDimensions - ); -} - -#[test] -fn rejects_invalid_gather_output_length() { - let data = Matrix::::new(0.0, 2, 2); - let error = gather_vectors(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "gather output", - expected: 4, - actual: 3, - } - ); -} - -#[test] -fn rejects_assignment_to_an_unknown_leader() { - let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "leader assignment", - expected: 2, - actual: 3, - } - ); -} - -#[test] -fn rejects_empty_and_oversized_leaves() { - for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { - let error = validate_leaves(&leaves, 2).unwrap_err(); - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidLeaf { size, limit: 2 } - ); - } -} diff --git a/diskann/src/graph/pipnn/tests.rs b/diskann/src/graph/pipnn/tests.rs deleted file mode 100644 index d5689eede6..0000000000 --- a/diskann/src/graph/pipnn/tests.rs +++ /dev/null @@ -1,27 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use super::*; -use half::f16; - -#[test] -fn integer_normalized_cosine_uses_unnormalized_cosine() { - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let expected = if metric == Metric::CosineNormalized { - Metric::Cosine - } else { - metric - }; - assert_eq!(effective_metric::(metric), expected); - assert_eq!(effective_metric::(metric), expected); - assert_eq!(effective_metric::(metric), metric); - assert_eq!(effective_metric::(metric), metric); - } -} diff --git a/diskann/tests/pipnn_build_graph.rs b/diskann/tests/pipnn_build_graph.rs deleted file mode 100644 index 12b8c0ec9a..0000000000 --- a/diskann/tests/pipnn_build_graph.rs +++ /dev/null @@ -1,233 +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 diskann::graph::config::{self, MaxDegree}; -use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig, build_graph}; -use diskann_utils::views::MatrixView; -use diskann_vector::distance::Metric; -use half::f16; -use rand::{Rng, SeedableRng, rngs::StdRng}; - -fn pipnn_config() -> PiPNNConfig { - PiPNNConfig { - c_max: 4, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - k: 1, - replicas: 1, - } -} - -fn graph_config(metric: Metric, degree: usize) -> diskann::graph::Config { - config::Builder::new_with(degree, MaxDegree::same(), 8, metric.into(), |builder| { - builder.alpha(1.2); - }) - .build() - .unwrap() -} - -fn pool(threads: usize) -> rayon::ThreadPool { - rayon::ThreadPoolBuilder::new() - .num_threads(threads) - .build() - .unwrap() -} - -fn rows(graph: Vec>) -> Vec> { - graph.into_iter().map(Vec::from).collect() -} - -fn assert_graph_invariants( - graph: &[diskann::graph::AdjacencyList], - points: usize, - degree: usize, -) { - assert_eq!(graph.len(), points); - for (source, row) in graph.iter().enumerate() { - assert!(row.len() <= degree); - let mut sorted = row.to_vec(); - sorted.sort_unstable(); - sorted.dedup(); - assert_eq!(sorted.len(), row.len()); - assert!( - row.iter() - .all(|&id| (id as usize) < points && id as usize != source) - ); - } -} - -#[test] -fn builds_a_single_leaf_graph_for_real_dataset_ids() { - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let graph = graph_config(Metric::L2, 2); - let pool = pool(2); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); - - assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); - - let graph = graph_config(Metric::L2, 1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - let pruned = build_graph(data, &context).unwrap(); - - assert_graph_invariants(&pruned, 4, 1); - for (source, neighbors) in pruned.iter().enumerate() { - assert_eq!(source.abs_diff(neighbors[0] as usize), 1); - } -} - -#[test] -fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { - let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; - let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); - let graph = graph_config(Metric::L2, 1); - let pool = pool(2); - let config = PiPNNConfig { - c_max: 5, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - k: 4, - replicas: 1, - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); - - assert_graph_invariants(&actual, 5, 1); - assert!(actual.iter().all(|row| row.len() == 1)); -} - -#[test] -fn rejects_empty_dataset_dimensions_at_the_public_boundary() { - let graph = graph_config(Metric::L2, 2); - let pool = pool(1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); - let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); - - assert!(build_graph(no_rows, &context).is_err()); - assert!(build_graph(no_columns, &context).is_err()); -} - -#[test] -fn supports_every_source_type_and_metric() { - fn build(values: &[T], metric: Metric) { - let data = MatrixView::try_from(values, 6, 2).unwrap(); - let graph = graph_config(metric, 2); - let pool = pool(2); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); - let actual = build_graph(data, &context).unwrap(); - assert_graph_invariants(&actual, 6, 2); - } - - let values = [ - 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, - ]; - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - build(&values, metric); - } - build(&values.map(f16::from_f32), Metric::L2); - build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); - build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); -} - -#[test] -fn integer_normalized_cosine_matches_cosine() { - fn assert_match(values: &[T]) { - let data = MatrixView::try_from(values, 8, 2).unwrap(); - let pool = pool(2); - let build = |metric| { - let graph = graph_config(metric, 2); - let config = PiPNNConfig { - c_max: 8, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - k: 1, - replicas: 1, - }; - let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); - rows(build_graph(data, &context).unwrap()) - }; - assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); - } - - assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); - assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); -} - -#[test] -fn is_deterministic_for_a_fixed_pool_size() { - let data: Vec = (0..96 * 4) - .map(|value| ((value * 17 + 3) % 101) as f32) - .collect(); - let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); - let graph = graph_config(Metric::L2, 8); - let pool = pool(4); - let config = PiPNNConfig { - c_max: 16, - c_min: 4, - p_samp: 0.25, - fanout: vec![3, 2], - k: 3, - replicas: 2, - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - let first = build_graph(data, &context).unwrap(); - let second = build_graph(data, &context).unwrap(); - - assert_eq!(first, second); - assert_graph_invariants(&first, 96, 8); -} - -#[test] -fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { - let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); - for case in 0..24 { - let points = rng.random_range(4..=32); - let dimensions = rng.random_range(1..=8); - let c_max = rng.random_range(4..=points.min(12)); - let c_min = rng.random_range(1..=c_max); - let degree = rng.random_range(1..=points.min(8)); - let values: Vec = (0..points * dimensions) - .map(|_| rng.random_range(-10.0..10.0)) - .collect(); - let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); - let graph = graph_config(Metric::L2, degree); - let pool = pool(2); - let config = PiPNNConfig { - c_max, - c_min, - p_samp: 0.5, - fanout: vec![2], - k: rng.random_range(1..=3), - replicas: rng.random_range(1..=2), - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context) - .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); - assert_graph_invariants(&actual, points, degree); - } -} diff --git a/diskann/tests/pipnn_config.rs b/diskann/tests/pipnn_config.rs deleted file mode 100644 index 19a0d1ce8a..0000000000 --- a/diskann/tests/pipnn_config.rs +++ /dev/null @@ -1,123 +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 diskann::graph::config::{self, MaxDegree}; -use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig}; -use diskann_vector::distance::Metric; - -fn pipnn_config() -> PiPNNConfig { - PiPNNConfig { - c_max: 512, - c_min: 64, - p_samp: 0.01, - fanout: vec![10, 3], - k: 2, - replicas: 1, - } -} - -fn graph_config(metric: Metric, alpha: f32) -> diskann::graph::Config { - config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { - builder.alpha(alpha); - }) - .build() - .unwrap() -} - -fn pool() -> rayon::ThreadPool { - rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap() -} - -#[test] -fn rejects_each_invalid_algorithm_parameter() { - let graph = graph_config(Metric::L2, 1.2); - let pool = pool(); - let mut cases = [ - PiPNNConfig { - c_max: 0, - ..pipnn_config() - }, - PiPNNConfig { - c_min: 0, - ..pipnn_config() - }, - PiPNNConfig { - c_min: 513, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: 0.0, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: -0.01, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: 1.01, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: f64::NAN, - ..pipnn_config() - }, - PiPNNConfig { - fanout: Vec::new(), - ..pipnn_config() - }, - PiPNNConfig { - fanout: vec![1, 0], - ..pipnn_config() - }, - PiPNNConfig { - fanout: vec![17], - ..pipnn_config() - }, - PiPNNConfig { - k: 0, - ..pipnn_config() - }, - PiPNNConfig { - replicas: 0, - ..pipnn_config() - }, - ]; - - for config in &mut cases { - let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) - .expect_err("invalid PiPNN config must be rejected"); - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); - } -} - -#[test] -fn rejects_graph_policy_for_a_different_metric() { - let graph = graph_config(Metric::InnerProduct, 1.2); - let pool = pool(); - - let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); - - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); - assert!(error.to_string().contains("prune kind")); -} - -#[test] -fn does_not_add_alpha_validation_beyond_graph_config() { - let pool = pool(); - for alpha in [0.9, f32::NAN, f32::INFINITY] { - let graph = graph_config(Metric::L2, alpha); - PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - } -} From 9c334cebaae3138d4da14bddc47b2f700f31e5fc Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:53:28 +0000 Subject: [PATCH 125/229] test(pipnn): adapt assertions to main errors --- diskann/src/graph/pipnn/mod.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 6112fc0f8e..97ba73fa37 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -763,9 +763,8 @@ mod config_tests { ]; for config in &mut cases { - let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) .expect_err("invalid PiPNN config must be rejected"); - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); } } @@ -776,7 +775,6 @@ mod config_tests { let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); assert!(error.to_string().contains("prune kind")); } From eb843335f1b90b1e82edebf239be77be0088e274 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:14:50 +0000 Subject: [PATCH 126/229] refactor(pipnn): consume positional robust prune Adapt PiPNN-owned preparation and ID translation to the behavior-preserving positional kernel introduced by #1315. --- diskann/src/graph/pipnn/finalization.rs | 36 ++++++++++++++----------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index a10d4df027..e060ad621b 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -35,14 +35,9 @@ //! | bounded list | none | move list directly to output | //! | overfull list | source and occlusion distances | reuse Rayon-job workspace | -use std::convert::Infallible; - use crate::{ ANNError, ANNResult, - graph::{ - AdjacencyList, Config, - internal::{SortedNeighbors, robust_prune as prune}, - }, + graph::{AdjacencyList, Config, internal::SortedNeighbors, robust_prune as prune}, neighbor::Neighbor, utils::VectorRepr, }; @@ -62,6 +57,8 @@ pub(crate) enum FinalizationError { candidate: u32, points: usize, }, + #[error("candidate count {actual} exceeds the u16 position limit {max}")] + TooManyCandidates { actual: usize, max: usize }, } /// Per-Rayon-job preparation and kernel state retained across source points. @@ -71,7 +68,7 @@ pub(crate) enum FinalizationError { #[derive(Default)] struct Workspace { pool: Vec>, - prepared: Vec>, + prepared: Vec<(f32, Option)>, states: Vec, } @@ -122,8 +119,12 @@ where })); let candidate_count = workspace.pool.len(); - prune::validate_candidate_count::(candidate_count) - .map_err(ANNError::new)?; + if candidate_count > u16::MAX as usize { + return Err(ANNError::new(FinalizationError::TooManyCandidates { + actual: candidate_count, + max: u16::MAX as usize, + })); + } workspace.prepared.clear(); workspace .prepared @@ -137,8 +138,7 @@ where .prepared .extend(sorted.iter().filter_map(|neighbor| { let id = *neighbor.id(); - (id != source_id) - .then(|| prune::Candidate::new(id, *neighbor.distance(), id)) + (id != source_id).then(|| (*neighbor.distance(), Some(id))) })); } workspace @@ -161,17 +161,21 @@ where graph.alpha(), graph.prune_kind(), |left, right| { - Ok::<_, Infallible>(distance.evaluate_similarity( + distance.evaluate_similarity( data.row(*left as usize), data.row(*right as usize), - )) + ) }, - ) - .map_err(ANNError::new)?; + ); let mut guard = source_candidates.resize(selected); for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { - *destination = *workspace.prepared[state.selected_position()].id(); + *destination = + workspace.prepared[state.neighbor as usize] + .1 + .ok_or_else(|| { + ANNError::message("RobustPrune selected an unavailable candidate") + })?; } guard.finish(selected); Ok(source_candidates) From 01d14cb8b8ff28c2cf7ead14850bcc155aac615e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 02:57:30 +0000 Subject: [PATCH 127/229] refactor(pipnn): use sorted prune input Pass the existing SortedNeighbors witness into internal prune so source-distance ordering is enforced by type rather than caller documentation. --- diskann/src/graph/pipnn/finalization.rs | 33 +++++++++++-------------- 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index e060ad621b..7d0ccb722d 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -37,7 +37,10 @@ use crate::{ ANNError, ANNResult, - graph::{AdjacencyList, Config, internal::SortedNeighbors, robust_prune as prune}, + graph::{ + AdjacencyList, Config, + internal::{SortedNeighbors, prune}, + }, neighbor::Neighbor, utils::VectorRepr, }; @@ -130,17 +133,15 @@ where .prepared .try_reserve(candidate_count) .map_err(ANNError::new)?; - { - // Sorting/capping precedes source exclusion so filtering cannot - // backfill with farther candidates. - let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); - workspace - .prepared - .extend(sorted.iter().filter_map(|neighbor| { - let id = *neighbor.id(); - (id != source_id).then(|| (*neighbor.distance(), Some(id))) - })); - } + + // Sorting/capping precedes source exclusion so filtering cannot + // backfill with farther candidates. Passing this witness into + // RobustPrune makes source-distance order part of its input type. + let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); + workspace.prepared.extend(sorted.iter().map(|neighbor| { + let id = *neighbor.id(); + (*neighbor.distance(), (id != source_id).then_some(id)) + })); workspace .states .try_reserve( @@ -155,6 +156,7 @@ where .resize(workspace.prepared.len(), prune::State::default()); let selected = prune::robust_prune( + &sorted, &workspace.prepared, workspace.states.as_mut_slice(), degree, @@ -170,12 +172,7 @@ where let mut guard = source_candidates.resize(selected); for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { - *destination = - workspace.prepared[state.neighbor as usize] - .1 - .ok_or_else(|| { - ANNError::message("RobustPrune selected an unavailable candidate") - })?; + *destination = *sorted[state.neighbor as usize].id(); } guard.finish(selected); Ok(source_candidates) From 78bfe0eafa5ca91a1730258c91bf86edebc7bbbc Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:47:35 +0000 Subject: [PATCH 128/229] refactor(pipnn): use direct leaf matrix input --- diskann/src/graph/pipnn/leaf_build.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 1c60a40c4e..3667c2f1e6 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -28,7 +28,7 @@ use diskann_vector::distance::Metric; use rayon::prelude::*; use super::leaf_kernel::{ - LeafInput, LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, + LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, leaf_output_len, }; @@ -356,7 +356,7 @@ where }, })?; kernel - .nearest_neighbors(LeafInput { dots }, output, &mut buffers.kernel_workspace) + .nearest_neighbors(dots, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; buffers.prepare_local_adjacency(point_ids.len())?; From 467c6b4f4b75a6bdc2b732151546dc725c7784c7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 05:56:05 +0000 Subject: [PATCH 129/229] fix(pipnn): validate leaf k capacity Reject k outside 1..=3 at context construction so production never reaches an unsupported leaf-kernel width. --- diskann/src/graph/pipnn/mod.rs | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 97ba73fa37..bfcd2a70f9 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -250,7 +250,7 @@ pub struct PiPNNConfig { pub p_samp: f64, /// Number of nearest leaders retained at each overlapping partition level. pub fanout: Vec, - /// Number of nearest neighbors selected within each leaf. + /// Number of nearest neighbors selected within each leaf (`1..=3`). pub k: usize, /// Number of independent partition passes over the dataset. pub replicas: usize, @@ -290,8 +290,12 @@ impl PiPNNConfig { partition_kernel::MAX_PARTITION_FANOUT ))); } - if self.k == 0 { - return Err(config_error("k must be greater than zero")); + if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.k) { + return Err(config_error(format!( + "k ({}) must be in [1, {}]", + self.k, + leaf_kernel::MAX_LEAF_NEIGHBORS + ))); } if self.replicas == 0 { return Err(config_error("replicas must be greater than zero")); @@ -445,7 +449,7 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod build_graph_tests { - use super::{PiPNNBuildContext, PiPNNConfig, build_graph}; + use super::{PiPNNBuildContext, PiPNNConfig, build_graph, leaf_kernel}; use crate::graph::config::{self, MaxDegree}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; @@ -525,7 +529,7 @@ mod build_graph_tests { } #[test] - fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { + fn prunes_overfull_single_leaf_candidates_to_the_graph_degree() { let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); let graph = graph_config(Metric::L2, 1); @@ -535,7 +539,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - k: 4, + k: leaf_kernel::MAX_LEAF_NEIGHBORS, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -677,7 +681,7 @@ mod build_graph_tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod config_tests { - use super::{PiPNNBuildContext, PiPNNConfig}; + use super::{PiPNNBuildContext, PiPNNConfig, leaf_kernel}; use crate::graph::config::{self, MaxDegree}; use diskann_vector::distance::Metric; @@ -756,6 +760,10 @@ mod config_tests { k: 0, ..pipnn_config() }, + PiPNNConfig { + k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, + ..pipnn_config() + }, PiPNNConfig { replicas: 0, ..pipnn_config() From 372c6e338f12eadccde2f86be7989d046440b46f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:41:17 +0000 Subject: [PATCH 130/229] refactor(pipnn): require sorted leaf IDs Partitioning is the only production source and emits sorted unique IDs. Reject unsorted input linearly and remove the HashSet fallback. --- diskann/src/graph/pipnn/leaf_build.rs | 62 +++++++-------------------- 1 file changed, 16 insertions(+), 46 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 3667c2f1e6..6d54913748 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -5,7 +5,8 @@ //! Leaf-local graph construction and candidate accumulation. //! -//! Partitioning supplies leaves as global point IDs. For each leaf this module: +//! Partitioning supplies strictly increasing, unique global point IDs per leaf. +//! For each leaf this module: //! //! 1. validates IDs and converts only those point vectors to reusable `f32` scratch; //! 2. computes the lower triangle of `A · Aᵀ`; @@ -17,10 +18,7 @@ //! high-water length; every consumer therefore receives an explicit active //! prefix rather than treating `Vec::len()` as the current leaf shape. -use std::{ - collections::{HashSet, TryReserveError}, - sync::Mutex, -}; +use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -49,12 +47,16 @@ pub(crate) enum LeafBuildError { }, #[error("point ID {point} appears more than once in leaf {leaf}")] DuplicatePointId { leaf: usize, point: u32 }, + #[error("point IDs in leaf {leaf} are not strictly increasing")] + UnsortedPointIds { leaf: usize }, #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] ShapeOverflow { leaf: usize, rows: usize, columns: usize, }, + #[error("failed to form {buffer} view for leaf {leaf}")] + InvalidView { leaf: usize, buffer: &'static str }, #[error("failed to reserve {additional} values for {buffer}")] Allocation { buffer: &'static str, @@ -100,7 +102,6 @@ struct LeafBuffers { neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, - seen_ids: HashSet, } impl LeafBuffers { @@ -283,24 +284,14 @@ where }); } } - if point_ids.is_sorted() { - if let Some(pair) = point_ids.windows(2).find(|pair| pair[0] == pair[1]) { + if let Some(pair) = point_ids.windows(2).find(|pair| pair[0] >= pair[1]) { + if pair[0] == pair[1] { return Err(LeafBuildError::DuplicatePointId { leaf, point: pair[0], }); } - } else { - buffers.seen_ids.clear(); - buffers - .seen_ids - .try_reserve(point_ids.len()) - .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; - for &point in point_ids { - if !buffers.seen_ids.insert(point) { - return Err(LeafBuildError::DuplicatePointId { leaf, point }); - } - } + return Err(LeafBuildError::UnsortedPointIds { leaf }); } let leaf_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; if leaf_k == 0 { @@ -334,26 +325,18 @@ where ) .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; let dots = MatrixView::try_from(&buffers.dots[..dot_count], point_ids.len(), point_ids.len()) - .map_err(|error| LeafBuildError::Kernel { + .map_err(|_| LeafBuildError::InvalidView { leaf, - source: LeafKernelError::InvalidBufferLength { - buffer: "leaf dot-product matrix", - expected: dot_count, - actual: error.into_inner().len(), - }, + buffer: "leaf dot-product matrix", })?; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), leaf_k, ) - .map_err(|error| LeafBuildError::Kernel { + .map_err(|_| LeafBuildError::InvalidView { leaf, - source: LeafKernelError::InvalidBufferLength { - buffer: "output", - expected: neighbor_value_count, - actual: error.into_inner().len(), - }, + buffer: "leaf output", })?; kernel .nearest_neighbors(dots, output, &mut buffers.kernel_workspace) @@ -570,19 +553,6 @@ mod tests { })); } - #[test] - fn global_id_translation_is_independent_of_leaf_order() { - let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; - let leaves = vec![vec![4, 1, 3]]; - - let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); - - assert_eq!( - adjacency_lists(graph), - [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] - ); - } - fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> where T: crate::utils::VectorRepr + 'static, @@ -710,8 +680,8 @@ mod tests { Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) )); assert!(matches!( - build(view(&data, 2, 1), &[vec![1, 0, 1]], 1, Metric::L2), - Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 1 }) + build(view(&data, 2, 1), &[vec![1, 0]], 1, Metric::L2), + Err(LeafBuildError::UnsortedPointIds { leaf: 0 }) )); } From 8b69ff60f1ca87bb9314aa464c57813374b4b6dc Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:41:17 +0000 Subject: [PATCH 131/229] refactor(pipnn): borrow partition configuration Rely on validated MatrixView shape and avoid cloning PiPNNConfig and its fanout vector. --- diskann/src/graph/pipnn/mod.rs | 9 +-------- diskann/src/graph/pipnn/partitioning.rs | 26 ++++++++++++------------- 2 files changed, 14 insertions(+), 21 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index bfcd2a70f9..a9297a7789 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -375,19 +375,12 @@ where data.nrows() ))); } - data.nrows().checked_mul(data.ncols()).ok_or_else(|| { - ANNError::message(format!( - "PiPNN dataset shape {} x {} overflows usize", - data.nrows(), - data.ncols() - )) - })?; // Integer source vectors are not guaranteed unit-normalized after conversion, // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition(data, context.config.clone(), metric))?; + .in_scope(|| partitioning::partition(data, &context.config, metric))?; // `leaves` is consumed here. Workers borrow individual ID lists during the // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 5de9a68cf5..5fd341d9cb 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -141,7 +141,7 @@ type StripeBufferPool = ObjectPool; /// covered once per replica. The caller installs the operation in its pool. pub(crate) fn partition( data: MatrixView<'_, T>, - config: PiPNNConfig, + config: &PiPNNConfig, metric: Metric, ) -> ANNResult>> where @@ -167,7 +167,7 @@ where for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = - partition_replica(data, &config, metric, &kernel, seed, &stripe_buffers)?; + partition_replica(data, config, metric, &kernel, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::new)?; @@ -880,7 +880,7 @@ mod tests { fn returns_one_leaf_at_and_below_c_max() { for points in [7, 8] { let data = clustered_data(points, 3); - let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); assert_eq!(leaves, vec![(0..points as u32).collect::>()]); } } @@ -890,8 +890,8 @@ mod tests { let data = clustered_data(96, 8); let config = config(4, 16, vec![3, 2], 2); - let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); - let second = partition(data.as_view(), config, Metric::L2).unwrap(); + let first = partition(data.as_view(), &config, Metric::L2).unwrap(); + let second = partition(data.as_view(), &config, Metric::L2).unwrap(); assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); assert_valid_partition(&first, 96, 16, 2); @@ -901,7 +901,7 @@ mod tests { #[test] fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { let data = clustered_data(80, 4); - let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); assert_valid_partition(&leaves, 80, 8, 1); } @@ -909,7 +909,7 @@ mod tests { #[test] fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { let data = Matrix::new(1.0f32, 24, 4); - let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); let error = error.downcast::().unwrap(); assert!(matches!( @@ -955,7 +955,7 @@ mod tests { let data = directional_data(72, 5); let leaves = partition( data.as_view(), - config(3, 12, vec![3, 2], 3), + &config(3, 12, vec![3, 2], 3), Metric::CosineNormalized, ) .unwrap(); @@ -983,13 +983,13 @@ mod tests { let config = config(2, 16, vec![2, 1], 1); let expected = partition( MatrixView::try_from(&f32_data, points, dimensions).unwrap(), - config.clone(), + &config, Metric::L2, ) .unwrap(); let actual = partition( MatrixView::try_from(&converted, points, dimensions).unwrap(), - config, + &config, Metric::L2, ) .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); @@ -1069,7 +1069,7 @@ mod tests { Metric::CosineNormalized, Metric::InnerProduct, ] { - let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); + let leaves = partition(data.as_view(), &config, metric).unwrap(); assert_valid_partition(&leaves, 64, 20, 1); } } @@ -1150,7 +1150,7 @@ mod tests { #[test] fn rejects_empty_dataset() { let data = Matrix::::new(0.0, 0, 4); - let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); assert_eq!( error.downcast::().unwrap(), @@ -1161,7 +1161,7 @@ mod tests { #[test] fn rejects_zero_dimensions() { let data = Matrix::::new(0.0, 4, 0); - let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); assert_eq!( error.downcast::().unwrap(), From b87230ceccba7a406abbd841714820ceb84f0b67 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:49:02 +0000 Subject: [PATCH 132/229] refactor(pipnn): dispatch partition stages once Run partition orchestration under one architecture/metric specialization and reuse a runtime-sized tracker instead of imposing a fanout cap. --- diskann/src/graph/pipnn/mod.rs | 20 +-- diskann/src/graph/pipnn/partitioning.rs | 170 +++++++++++++++++------- 2 files changed, 128 insertions(+), 62 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index a9297a7789..99337244f7 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -90,9 +90,8 @@ //! [`partition_kernel::PartitionInput`] bundles that tile with typed //! [`partition_kernel::PartitionScales`]. A prepared //! [`partition_kernel::PartitionKernel`] writes sorted leader-local positions to -//! caller-owned output. Fanout is output column count and is bounded by -//! [`partition_kernel::MAX_PARTITION_FANOUT`]. Module documentation describes -//! scale units, validation, `process_points`, and tracker insertion. +//! caller-owned output. Fanout is the output column count and cannot exceed the +//! leaders available for that partition call. //! //! ## [`leaf_kernel`] //! @@ -280,15 +279,8 @@ impl PiPNNConfig { if self.fanout.is_empty() { return Err(config_error("fanout must not be empty")); } - if let Some(&fanout) = self - .fanout - .iter() - .find(|&&fanout| !(1..=partition_kernel::MAX_PARTITION_FANOUT).contains(&fanout)) - { - return Err(config_error(format!( - "fanout ({fanout}) must be in [1, {}]", - partition_kernel::MAX_PARTITION_FANOUT - ))); + if self.fanout.contains(&0) { + return Err(config_error("fanout values must be greater than zero")); } if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.k) { return Err(config_error(format!( @@ -745,10 +737,6 @@ mod config_tests { fanout: vec![1, 0], ..pipnn_config() }, - PiPNNConfig { - fanout: vec![17], - ..pipnn_config() - }, PiPNNConfig { k: 0, ..pipnn_config() diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 5fd341d9cb..4bc7cd164e 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -46,12 +46,19 @@ use diskann_utils::{ views::{MatrixView, MutMatrixView}, }; use diskann_vector::{Norm, distance::Metric, norm::FastL2NormSquared}; +use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, +}; use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; use super::{ PiPNNConfig, - partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, + kernel_metric::{KernelMetric, MetricVisitor, visit_metric}, + partition_kernel::{ + PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders_for, + }, }; // Private algorithm and batching constants live together. None are user policy. @@ -111,6 +118,7 @@ struct StripeBuffers { points: Vec, dots: Vec, point_scales: Vec, + kernel: PartitionKernelWorkspace, } impl AsPooled<()> for StripeBuffers { @@ -146,6 +154,70 @@ pub(crate) fn partition( ) -> ANNResult>> where T: VectorRepr + Send + Sync, +{ + arch::dispatch1_no_features( + RunPartitionStage, + PartitionStageCall { + data, + config, + metric, + }, + ) +} + +struct PartitionStageCall<'a, T> { + data: MatrixView<'a, T>, + config: &'a PiPNNConfig, + metric: Metric, +} + +struct RunPartitionStage; + +impl Target1>>, PartitionStageCall<'_, T>> for RunPartitionStage +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + Send + Sync, +{ + fn run(self, arch: A, call: PartitionStageCall<'_, T>) -> ANNResult>> { + visit_metric(call.metric, ExecutePartitionStage { arch, call }) + } +} + +struct ExecutePartitionStage<'a, A, T> { + arch: A, + call: PartitionStageCall<'a, T>, +} + +impl MetricVisitor for ExecutePartitionStage<'_, A, T> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + Send + Sync, +{ + type Output = ANNResult>>; + + fn visit(self) -> Self::Output { + partition_for::(self.arch, self.call.data, self.call.config) + } +} + +fn partition_for( + arch: A, + data: MatrixView<'_, T>, + config: &PiPNNConfig, +) -> ANNResult>> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, + T: VectorRepr + Send + Sync, { let points = data.nrows(); if points == 0 { @@ -159,15 +231,11 @@ where } let mut leaves = Vec::new(); - // Prepare metric and ISA dispatch before replicas spawn Rayon work. The - // Copy handle is shared read-only; every stripe calls its direct function - // pointer instead of redispatching in the recursive hot path. - let kernel = PartitionKernel::new(metric); let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = - partition_replica(data, config, metric, &kernel, seed, &stripe_buffers)?; + partition_replica::(arch, data, config, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::new)?; @@ -177,15 +245,19 @@ where Ok(leaves) } -fn partition_replica( +fn partition_replica( + arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, - metric: Metric, - kernel: &PartitionKernel, seed: u64, stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr + Send + Sync, { let initial_indices = point_ids(data.nrows())?; @@ -222,11 +294,10 @@ where .par_iter_mut() .zip(work.into_par_iter()) .for_each(|(slot, item)| { - *slot = Some(partition_one_level( + *slot = Some(partition_one_level::( + arch, data, config, - metric, - kernel, item, stripe_buffers, )); @@ -259,15 +330,19 @@ where })) } -fn partition_one_level( +fn partition_one_level( + arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, - metric: Metric, - kernel: &PartitionKernel, item: WorkItem, stripe_buffers: &StripeBufferPool, ) -> ANNResult<(Vec, Vec>)> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr + Send + Sync, { let points = item.indices.len(); @@ -277,15 +352,8 @@ where config.p_samp, mix_seed(item.seed, points as u64), )?; - let clusters = assign_to_leaders( - data, - &item.indices, - &leaders, - fanout, - metric, - kernel, - stripe_buffers, - )?; + let clusters = + assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; let mut pending = Vec::new(); let mut finished = Vec::new(); @@ -344,16 +412,20 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { /// its stripes. The flat assignment matrix preserves point order and is then /// scattered into per-leader clusters; preserving order is required for fixed /// seed determinism in later recursion levels. -fn assign_to_leaders( +fn assign_to_leaders( + arch: A, data: MatrixView<'_, T>, point_ids: &[u32], leader_ids: &[u32], fanout: usize, - metric: Metric, - kernel: &PartitionKernel, stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr + Send + Sync, { let dimension_count = data.ncols(); @@ -361,7 +433,7 @@ where let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; gather_vectors(data, leader_ids, &mut leader_values)?; - let mut leader_scales = if matches!(metric, Metric::L2 | Metric::Cosine) { + let mut leader_scales = if matches!(M::METRIC, Metric::L2 | Metric::Cosine) { filled_vec(leader_ids.len(), 0.0f32)? } else { Vec::new() @@ -375,7 +447,7 @@ where // SIMD norm changes low bits and can send near-tied points down different // recursive partition paths. *scale = leader_vector.iter().map(|value| value * value).sum(); - if metric == Metric::Cosine { + if M::METRIC == Metric::Cosine { *scale = scale.sqrt(); } } @@ -405,13 +477,12 @@ where { let first_point = worker_first + stripe * stripe_points; let stripe_point_count = stripe_assignments.len() / fanout; - assign_stripe( + assign_stripe::( + arch, data, &point_ids[first_point..first_point + stripe_point_count], &leader_values, &leader_scales, - metric, - kernel, fanout, &mut buffers, stripe_assignments, @@ -425,18 +496,22 @@ where #[inline] #[allow(clippy::too_many_arguments)] -fn assign_stripe( +fn assign_stripe( + arch: A, data: MatrixView<'_, T>, point_ids: &[u32], leader_values: &[f32], leader_scales: &[f32], - metric: Metric, - kernel: &PartitionKernel, fanout: usize, buffers: &mut StripeBuffers, assignments: &mut [u32], ) -> ANNResult<()> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr, { let point_count = point_ids.len(); @@ -458,6 +533,7 @@ where points: point_buffer, dots: dot_buffer, point_scales: point_scale_buffer, + kernel: kernel_workspace, } = buffers; let point_values = &mut point_buffer[..point_values_len]; let dots = &mut dot_buffer[..dots_len]; @@ -476,7 +552,7 @@ where ) .map_err(ANNError::new)?; - let point_scales = if metric == Metric::Cosine { + let point_scales = if M::METRIC == Metric::Cosine { grow_fallible(point_scale_buffer, point_count, 0.0)?; let point_scales = &mut point_scale_buffer[..point_count]; for (scale, point_values) in point_scales @@ -489,7 +565,7 @@ where } else { &[] }; - let scales = match metric { + let scales = match M::METRIC { Metric::L2 => PartitionScales::L2 { leader_squared_norms: leader_scales, }, @@ -513,9 +589,13 @@ where actual: error.into_inner().len(), }) })?; - kernel - .nearest_leaders(PartitionInput { dots, scales }, output) - .map_err(ANNError::new) + nearest_leaders_for::( + arch, + PartitionInput { dots, scales }, + output, + kernel_workspace, + ) + .map_err(ANNError::new) } fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> @@ -1044,13 +1124,12 @@ mod tests { .collect(); let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); - let clusters = assign_to_leaders( + let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( + diskann_wide::ARCH, data, &[2], &[0, 1], 1, - Metric::L2, - &PartitionKernel::new(Metric::L2), &StripeBufferPool::new((), 0, None), ) .unwrap(); @@ -1118,13 +1197,12 @@ mod tests { let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); let point_ids: Vec = (0..points as u32).collect(); - let clusters = assign_to_leaders( + let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( + diskann_wide::ARCH, data, &point_ids, &[0, 2_047], 1, - Metric::L2, - &PartitionKernel::new(Metric::L2), &StripeBufferPool::new((), 0, None), ) .unwrap(); From 5d1c14d3b65c6c3b3ca2ee0c97f43e87f9662940 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:49:02 +0000 Subject: [PATCH 133/229] refactor(pipnn): dispatch leaf stages once Keep architecture and metric concrete across the Rayon leaf pass and call the generic kernel directly. --- diskann/src/graph/pipnn/leaf_build.rs | 118 ++++++++++++++++++++++---- 1 file changed, 101 insertions(+), 17 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6d54913748..c5cc541a1b 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -23,11 +23,18 @@ use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; +use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, +}; use rayon::prelude::*; -use super::leaf_kernel::{ - LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, - leaf_output_len, +use super::{ + kernel_metric::{KernelMetric, MetricVisitor, visit_metric}, + leaf_kernel::{ + LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, leaf_output_len, + nearest_neighbors_for, + }, }; /// Failure while converting leaves into direct graph candidates. @@ -216,7 +223,6 @@ impl DirectCandidates { } /// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. -#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. pub(crate) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, @@ -225,6 +231,84 @@ pub(crate) fn build_leaf_candidates( ) -> Result>, LeafBuildError> where T: VectorRepr + 'static, +{ + arch::dispatch1_no_features( + RunLeafStage, + LeafStageCall { + data, + leaves, + requested_k, + metric, + }, + ) +} + +struct LeafStageCall<'a, T> { + data: MatrixView<'a, T>, + leaves: Vec>, + requested_k: usize, + metric: Metric, +} + +struct RunLeafStage; + +impl Target1>, LeafBuildError>, LeafStageCall<'_, T>> + for RunLeafStage +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + 'static, +{ + fn run( + self, + arch: A, + call: LeafStageCall<'_, T>, + ) -> Result>, LeafBuildError> { + visit_metric(call.metric, ExecuteLeafStage { arch, call }) + } +} + +struct ExecuteLeafStage<'a, A, T> { + arch: A, + call: LeafStageCall<'a, T>, +} + +impl MetricVisitor for ExecuteLeafStage<'_, A, T> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + 'static, +{ + type Output = Result>, LeafBuildError>; + + fn visit(self) -> Self::Output { + build_leaf_candidates_for::( + self.arch, + self.call.data, + self.call.leaves, + self.call.requested_k, + ) + } +} + +#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. +fn build_leaf_candidates_for( + arch: A, + data: MatrixView<'_, T>, + leaves: Vec>, + requested_k: usize, +) -> Result>, LeafBuildError> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, + T: VectorRepr + 'static, { if data.ncols() == 0 { return Err(LeafBuildError::EmptyDimensions); @@ -234,18 +318,15 @@ where } let candidates = DirectCandidates::new(data.nrows())?; - // Metric and ISA are selected before Rayon workers start. Workers share - // this Copy handle; each output view supplies its leaf-specific width. - let kernel = LeafKernel::new(metric); leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf( + build_leaf::( + arch, data, leaf, point_ids, requested_k, - &kernel, buffers, &candidates, ) @@ -256,20 +337,24 @@ where /// Build and publish one leaf's symmetric neighbor lists. /// -/// Validation precedes all dataset indexing. Sorted partition output takes the -/// adjacent-duplicate path, while arbitrary-order callers use `seen_ids`. The -/// active lengths computed after `prepare` must be used for every later slice, -/// because the reusable vectors may still be longer than this leaf. -fn build_leaf( +/// Validation precedes all dataset indexing and rejects IDs that are not strictly +/// increasing. Active lengths computed after `prepare` must be used for every +/// later slice because reusable vectors may remain longer than this leaf. +fn build_leaf( + arch: A, data: MatrixView<'_, T>, leaf: usize, point_ids: &[u32], requested_k: usize, - kernel: &LeafKernel, buffers: &mut LeafBuffers, candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr + 'static, { if point_ids.is_empty() { @@ -338,8 +423,7 @@ where leaf, buffer: "leaf output", })?; - kernel - .nearest_neighbors(dots, output, &mut buffers.kernel_workspace) + nearest_neighbors_for::(arch, dots, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; buffers.prepare_local_adjacency(point_ids.len())?; From c74341d946e25117b202af3bcfcefe1ea18af005 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:08:19 +0000 Subject: [PATCH 134/229] refactor(pipnn): dispatch once per graph build --- diskann/src/graph/pipnn/leaf_build.rs | 170 ++++++------- diskann/src/graph/pipnn/mod.rs | 303 ++++++++---------------- diskann/src/graph/pipnn/partitioning.rs | 173 +++++++------- 3 files changed, 263 insertions(+), 383 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index c5cc541a1b..3ff70284d9 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -22,18 +22,14 @@ use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_vector::distance::Metric; -use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, -}; +use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::{KernelMetric, MetricVisitor, visit_metric}, + kernel_metric::KernelMetric, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, leaf_output_len, - nearest_neighbors_for, + nearest_neighbors, }, }; @@ -223,80 +219,10 @@ impl DirectCandidates { } /// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. -pub(crate) fn build_leaf_candidates( - data: MatrixView<'_, T>, - leaves: Vec>, - requested_k: usize, - metric: Metric, -) -> Result>, LeafBuildError> -where - T: VectorRepr + 'static, -{ - arch::dispatch1_no_features( - RunLeafStage, - LeafStageCall { - data, - leaves, - requested_k, - metric, - }, - ) -} - -struct LeafStageCall<'a, T> { - data: MatrixView<'a, T>, - leaves: Vec>, - requested_k: usize, - metric: Metric, -} - -struct RunLeafStage; - -impl Target1>, LeafBuildError>, LeafStageCall<'_, T>> - for RunLeafStage -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - T: VectorRepr + 'static, -{ - fn run( - self, - arch: A, - call: LeafStageCall<'_, T>, - ) -> Result>, LeafBuildError> { - visit_metric(call.metric, ExecuteLeafStage { arch, call }) - } -} - -struct ExecuteLeafStage<'a, A, T> { - arch: A, - call: LeafStageCall<'a, T>, -} - -impl MetricVisitor for ExecuteLeafStage<'_, A, T> -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - T: VectorRepr + 'static, -{ - type Output = Result>, LeafBuildError>; - - fn visit(self) -> Self::Output { - build_leaf_candidates_for::( - self.arch, - self.call.data, - self.call.leaves, - self.call.requested_k, - ) - } -} - +/// +/// The caller has already selected `A` and `M` for the complete graph build. #[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. -fn build_leaf_candidates_for( +pub(super) fn build_leaf_candidates( arch: A, data: MatrixView<'_, T>, leaves: Vec>, @@ -423,7 +349,7 @@ where leaf, buffer: "leaf output", })?; - nearest_neighbors_for::(arch, dots, output, &mut buffers.kernel_workspace) + nearest_neighbors::(arch, dots, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; buffers.prepare_local_adjacency(point_ids.len())?; @@ -507,6 +433,10 @@ fn poisoned_candidate_list(point: u32) -> LeafBuildError { mod tests { use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; + use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, + }; use half::f16; use std::collections::BTreeSet; @@ -526,6 +456,57 @@ mod tests { .unwrap() } + struct LeafBuildCall<'a, T> { + data: MatrixView<'a, T>, + leaves: Vec>, + k: usize, + } + + struct DispatchLeafBuild(Metric); + + impl + Target1< + A, + Result>, LeafBuildError>, + LeafBuildCall<'_, T>, + > for DispatchLeafBuild + where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: crate::utils::VectorRepr + 'static, + { + fn run( + self, + arch: A, + call: LeafBuildCall<'_, T>, + ) -> Result>, LeafBuildError> { + use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + + match self.0 { + Metric::L2 => { + build_leaf_candidates::(arch, call.data, call.leaves, call.k) + } + Metric::Cosine => { + build_leaf_candidates::(arch, call.data, call.leaves, call.k) + } + Metric::CosineNormalized => build_leaf_candidates::( + arch, + call.data, + call.leaves, + call.k, + ), + Metric::InnerProduct => build_leaf_candidates::( + arch, + call.data, + call.leaves, + call.k, + ), + } + } + } + fn build( data: MatrixView<'_, T>, leaves: &[Vec], @@ -535,7 +516,16 @@ mod tests { where T: crate::utils::VectorRepr + 'static, { - pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) + pool().install(|| { + arch::dispatch1_no_features( + DispatchLeafBuild(metric), + LeafBuildCall { + data, + leaves: leaves.to_vec(), + k, + }, + ) + }) } fn adjacency_lists(graph: Vec>) -> Vec> { @@ -718,17 +708,11 @@ mod tests { let leaves: Vec> = (0..32) .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) .collect(); - let pool = pool(); - pool.install(|| { - let expected = - build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); - for _ in 0..8 { - let actual = - build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2) - .unwrap(); - assert_eq!(actual, expected); - } - }); + let expected = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + for _ in 0..8 { + let actual = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + assert_eq!(actual, expected); + } } #[test] diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 99337244f7..8f813009da 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,149 +5,54 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! PiPNN means **Pick-in-Partitions Nearest Neighbors**. It 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. This module constructs that adjacency; it does not execute queries. +//! PiPNN builds graph candidates in bulk instead of searching a partially built +//! graph for every insertion: //! -//! 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 prune.** Candidate edges from overlapping leaves are combined -//! into one unique list per source. Vamana RobustPrune then selects a bounded, -//! directionally diverse adjacency list. +//! 1. `partitioning` recursively samples leaders and assigns each point to its +//! nearest configured fanout, producing overlapping leaves bounded by `c_max`. +//! 2. `leaf_build` gathers each leaf, computes the lower triangle of `A · Aᵀ`, +//! selects up to three local neighbors per point, and merges symmetric global +//! candidate IDs across overlapping leaves. +//! 3. `finalization` applies the shared Vamana RobustPrune policy only to +//! candidate lists that exceed the configured graph degree. //! //! ```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 +//! build_graph +//! └─ caller Rayon pool +//! └─ architecture dispatch + metric match once per build +//! └─ build_graph_for +//! ├─ partitioning::partition +//! │ └─ partition_kernel::nearest_leaders every stripe +//! ├─ leaf_build::build_leaf_candidates +//! │ └─ leaf_kernel::nearest_neighbors every leaf +//! └─ finalization::prune_overfull //! ``` //! -//! 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 crate owns overlapping partition generation, leaf-local nearest-neighbor -//! construction, candidate merging, and graph-degree finalization. The caller -//! supplies a contiguous dataset view, DiskANN graph policy, and the Rayon pool. -//! Providers, start/frozen points, quantization, persistence, and search remain -//! outside this algorithm boundary. -//! -//! 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. -//! -//! Both handles are prepared once per build metric and reused across stripes or -//! leaves. Each output view supplies call-specific fanout or neighbor width. -//! Preparation selects the runtime architecture and returns a direct function -//! pointer; repeated calls do not repeat ISA or metric dispatch. -//! -//! # Main modules and structures -//! -//! ## Public build API -//! -//! - [`PiPNNConfig`] holds Randomized Ball Carving, fanout, leaf size, local `k`, -//! and replication parameters. -//! - [`PiPNNBuildContext`] validates that algorithm parameters, graph pruning -//! policy, metric, and caller-owned Rayon pool agree. -//! - [`build_graph`] runs the full pipeline over a borrowed row-major dataset and -//! returns one dataset-ID adjacency list per input point. -//! -//! ## [`partition_kernel`] -//! -//! Partition callers compute point-by-leader dots with GEMM. -//! [`partition_kernel::PartitionInput`] bundles that tile with typed -//! [`partition_kernel::PartitionScales`]. A prepared -//! [`partition_kernel::PartitionKernel`] writes sorted leader-local positions to -//! caller-owned output. Fanout is the output column count and cannot exceed the -//! leaders available for that partition call. -//! -//! ## [`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 caller-owned output. [`leaf_kernel::leaf_neighbor_count`] derives each -//! leaf's width from point count and requested `k`. Module documentation explains -//! fixed-width selection, `process_pairs`, and stable endpoint insertion. -//! -//! ## Private pipeline stages -//! -//! - `partitioning` recursively samples leaders, invokes the partition kernel, -//! scatters points into overlapping children, and returns bounded leaves. -//! - `leaf_build` gathers each leaf, computes its Gram matrix, invokes the leaf -//! kernel, translates local positions to dataset IDs, and merges candidates. -//! - `finalization` applies shared Vamana RobustPrune to overfull candidate lists. -//! - `kernel_metric` owns norm preparation, exact ranking inputs, and numerical -//! edge cases. The graph build selects one concrete metric for both kernels. -//! -//! # Typical use +//! `diskann-wide` selects concrete architecture `A`; one four-way match selects +//! concrete metric marker `M`. Both types are then carried through every replica, +//! recursive partition, Rayon job, stripe, and leaf. Numerical kernels therefore +//! contain no runtime metric match, visitor, trait object, stored function pointer, +//! or repeated ISA dispatch. //! -//! 1. Construct [`PiPNNConfig`] and DiskANN graph [`Config`]. -//! 2. Create [`PiPNNBuildContext`] with metric and caller-owned Rayon pool. -//! 3. Call [`build_graph`] with one row-major [`MatrixView`] of dataset vectors. -//! 4. The outer index builder chooses start/frozen points and serializes returned -//! adjacency; those policies are intentionally not part of this crate. +//! [`PiPNNConfig`] owns only partition and local-neighbor parameters. +//! [`PiPNNBuildContext`] borrows DiskANN graph policy and the caller-owned Rayon +//! pool. [`build_graph`] borrows a contiguous [`MatrixView`] and returns one +//! dataset-ID adjacency list per real point. Providers, start/frozen points, +//! quantization, persistence, and search remain outside this module. //! -//! Stage outputs are owned values. Leaves move into candidate construction; -//! candidate lists move into finalization. Ownership releases each stage's large -//! scratch before the next outer allocation. -//! -//! # 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 itself never names instruction -//! sets; `diskann-wide` owns architecture selection. - -<<<<<<< HEAD -======= +//! The partition and leaf stages own disjoint reusable scratch. Stage outputs move +//! forward (`leaves → candidates → adjacency`) so large temporary allocations can +//! drop at their consumption boundary. Kernel modules document validation, +//! numerical edge cases, tie order, scalar tails, and unchecked SIMD preconditions. + mod kernel_metric; mod simd; ->>>>>>> a73b9c25c (docs(pipnn): document core stage invariants) mod finalization; mod leaf_build; mod leaf_kernel; -mod leaf_metric; mod partition_kernel; -mod partition_metric; mod partitioning; -mod simd; #[allow(dead_code)] mod topk; @@ -158,82 +63,13 @@ use crate::{ }; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; +use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, +}; use rayon::ThreadPool; -pub(super) struct L2; -pub(super) struct Cosine; -pub(super) struct CosineNormalized; -pub(super) struct InnerProduct; - -/// Convert one dot product and two norms to cosine distance. -/// -/// Treat a zero or subnormal norm as zero similarity. This rule takes precedence -/// over the dot value. Clamp finite similarity to the cosine range. Otherwise, -/// a NaN input produces a NaN distance. -#[inline(always)] -fn cosine_distance(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 / (source_norm * target_norm)).clamp(-1.0, 1.0) - } -} - -#[cfg(test)] -mod cosine_distance_tests { - use super::cosine_distance; - use rstest::rstest; - - #[rstest] - #[case::zero_source(0.0, 0.0, 1.0)] - #[case::zero_target(0.0, 1.0, 0.0)] - #[case::subnormal_source(f32::MIN_POSITIVE.sqrt() / 2.0, f32::MIN_POSITIVE.sqrt() / 2.0, 1.0)] - #[case::subnormal_target(f32::MIN_POSITIVE.sqrt() / 2.0, 1.0, f32::MIN_POSITIVE.sqrt() / 2.0)] - #[case::zero_norm_before_nan_dot(f32::NAN, 0.0, 1.0)] - fn small_norm_produces_unit_distance( - #[case] dot: f32, - #[case] source_norm: f32, - #[case] target_norm: f32, - ) { - assert_eq!(cosine_distance(dot, source_norm, target_norm), 1.0); - } - - #[test] - fn minimum_normal_norm_uses_normalized_similarity() { - let source_norm = f32::MIN_POSITIVE.sqrt(); - let expected_similarity = 0.5; - let dot = expected_similarity * source_norm; - - assert_eq!( - cosine_distance(dot, source_norm, 1.0), - 1.0 - expected_similarity - ); - } - - #[rstest] - #[case::above_one(1.0)] - #[case::below_negative_one(-1.0)] - fn finite_similarity_outside_the_cosine_range_is_clamped(#[case] bounded_similarity: f32) { - let norm = 2.0; - let norm_product = norm * norm; - let rounding_excess = f32::EPSILON * norm_product; - let dot = bounded_similarity * (norm_product + rounding_excess); - - assert_eq!(cosine_distance(dot, norm, norm), 1.0 - bounded_similarity); - } - - #[rstest] - #[case::nan_dot(f32::NAN, 1.0, 1.0)] - #[case::nan_source_norm(0.0, f32::NAN, 1.0)] - #[case::nan_target_norm(0.0, 1.0, f32::NAN)] - fn nan_without_a_zero_norm_produces_nan_distance( - #[case] dot: f32, - #[case] source_norm: f32, - #[case] target_norm: f32, - ) { - assert!(cosine_distance(dot, source_norm, target_norm).is_nan()); - } -} +use self::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; /// Configuration of PiPNN's partitioning and local-neighbor algorithm. /// @@ -370,19 +206,76 @@ where // Integer source vectors are not guaranteed unit-normalized after conversion, // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); + arch::dispatch1_no_features( + RunBuildGraph, + BuildGraphCall { + data, + context, + metric, + }, + ) +} + +struct BuildGraphCall<'data, 'context, 'policy, T> { + data: MatrixView<'data, T>, + context: &'context PiPNNBuildContext<'policy>, + metric: Metric, +} +struct RunBuildGraph; + +impl Target1>>, BuildGraphCall<'_, '_, '_, T>> + for RunBuildGraph +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + Send + Sync + 'static, +{ + fn run( + self, + arch: A, + call: BuildGraphCall<'_, '_, '_, T>, + ) -> ANNResult>> { + match call.metric { + Metric::L2 => build_graph_for::(arch, call.data, call.context), + Metric::Cosine => build_graph_for::(arch, call.data, call.context), + Metric::CosineNormalized => { + build_graph_for::(arch, call.data, call.context) + } + Metric::InnerProduct => { + build_graph_for::(arch, call.data, call.context) + } + } + } +} + +fn build_graph_for( + arch: A, + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, + T: VectorRepr + Send + Sync + 'static, +{ let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition(data, &context.config, metric))?; + .in_scope(|| partitioning::partition::(arch, data, &context.config))?; // `leaves` is consumed here. Workers borrow individual ID lists during the // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) + leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.k) .map_err(ANNError::new) })?; // Finalization consumes candidate lists and reuses their allocations for the // resulting adjacency where possible. tracing::info_span!("pipnn.finalization") - .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, M::METRIC)) } fn effective_metric(metric: Metric) -> Metric { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 4bc7cd164e..4c9fa56842 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -46,18 +46,15 @@ use diskann_utils::{ views::{MatrixView, MutMatrixView}, }; use diskann_vector::{Norm, distance::Metric, norm::FastL2NormSquared}; -use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, -}; +use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::{KernelMetric, MetricVisitor, visit_metric}, + kernel_metric::KernelMetric, partition_kernel::{ - PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders_for, + PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders, }, }; @@ -146,67 +143,9 @@ type StripeBufferPool = ObjectPool; /// leaders for the current level, and recurses only on oversized clusters. /// Levels beyond `fanout.len()` retain one leader assignment. Completed small /// leaves are merged without exceeding `c_max`; every input point must remain -/// covered once per replica. The caller installs the operation in its pool. -pub(crate) fn partition( - data: MatrixView<'_, T>, - config: &PiPNNConfig, - metric: Metric, -) -> ANNResult>> -where - T: VectorRepr + Send + Sync, -{ - arch::dispatch1_no_features( - RunPartitionStage, - PartitionStageCall { - data, - config, - metric, - }, - ) -} - -struct PartitionStageCall<'a, T> { - data: MatrixView<'a, T>, - config: &'a PiPNNConfig, - metric: Metric, -} - -struct RunPartitionStage; - -impl Target1>>, PartitionStageCall<'_, T>> for RunPartitionStage -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - T: VectorRepr + Send + Sync, -{ - fn run(self, arch: A, call: PartitionStageCall<'_, T>) -> ANNResult>> { - visit_metric(call.metric, ExecutePartitionStage { arch, call }) - } -} - -struct ExecutePartitionStage<'a, A, T> { - arch: A, - call: PartitionStageCall<'a, T>, -} - -impl MetricVisitor for ExecutePartitionStage<'_, A, T> -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - T: VectorRepr + Send + Sync, -{ - type Output = ANNResult>>; - - fn visit(self) -> Self::Output { - partition_for::(self.arch, self.call.data, self.call.config) - } -} - -fn partition_for( +/// covered once per replica. The caller has already selected `A` and `M` and +/// installed the operation in its pool. +pub(super) fn partition( arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, @@ -589,7 +528,7 @@ where actual: error.into_inner().len(), }) })?; - nearest_leaders_for::( + nearest_leaders::( arch, PartitionInput { dots, scales }, output, @@ -872,9 +811,55 @@ fn assignment_stripe_point_count(leader_count: usize) -> usize { mod tests { use diskann_utils::views::{Matrix, MatrixView}; use diskann_vector::{Half, distance::Metric}; + use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, + }; use super::*; + struct PartitionCall<'a, T> { + data: MatrixView<'a, T>, + config: &'a PiPNNConfig, + } + + struct DispatchPartition(Metric); + + impl Target1>>, PartitionCall<'_, T>> for DispatchPartition + where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + Send + Sync, + { + fn run(self, arch: A, call: PartitionCall<'_, T>) -> ANNResult>> { + use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + + match self.0 { + Metric::L2 => partition::(arch, call.data, call.config), + Metric::Cosine => partition::(arch, call.data, call.config), + Metric::CosineNormalized => { + partition::(arch, call.data, call.config) + } + Metric::InnerProduct => { + partition::(arch, call.data, call.config) + } + } + } + } + + fn partition_with_runtime_metric( + data: MatrixView<'_, T>, + config: &PiPNNConfig, + metric: Metric, + ) -> ANNResult>> + where + T: VectorRepr + Send + Sync, + { + arch::dispatch1_no_features(DispatchPartition(metric), PartitionCall { data, config }) + } + fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { PiPNNConfig { c_max, @@ -936,7 +921,12 @@ mod tests { memberships } - fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { + fn assert_valid_partition_with_runtime_metric( + leaves: &[Vec], + points: usize, + c_max: usize, + replicas: usize, + ) { assert!( leaves .iter() @@ -960,7 +950,12 @@ mod tests { fn returns_one_leaf_at_and_below_c_max() { for points in [7, 8] { let data = clustered_data(points, 3); - let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition_with_runtime_metric( + data.as_view(), + &config(2, 8, vec![2], 1), + Metric::L2, + ) + .unwrap(); assert_eq!(leaves, vec![(0..points as u32).collect::>()]); } } @@ -970,26 +965,30 @@ mod tests { let data = clustered_data(96, 8); let config = config(4, 16, vec![3, 2], 2); - let first = partition(data.as_view(), &config, Metric::L2).unwrap(); - let second = partition(data.as_view(), &config, Metric::L2).unwrap(); + let first = partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + let second = partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); - assert_valid_partition(&first, 96, 16, 2); + assert_valid_partition_with_runtime_metric(&first, 96, 16, 2); assert!(first.iter().map(Vec::len).sum::() > 96 * 2); } #[test] fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { let data = clustered_data(80, 4); - let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = + partition_with_runtime_metric(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2) + .unwrap(); - assert_valid_partition(&leaves, 80, 8, 1); + assert_valid_partition_with_runtime_metric(&leaves, 80, 8, 1); } #[test] fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { let data = Matrix::new(1.0f32, 24, 4); - let error = partition(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = + partition_with_runtime_metric(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2) + .unwrap_err(); let error = error.downcast::().unwrap(); assert!(matches!( @@ -1033,14 +1032,14 @@ mod tests { #[test] fn replicas_cover_every_point_once_or_more_per_replica() { let data = directional_data(72, 5); - let leaves = partition( + let leaves = partition_with_runtime_metric( data.as_view(), &config(3, 12, vec![3, 2], 3), Metric::CosineNormalized, ) .unwrap(); - assert_valid_partition(&leaves, 72, 12, 3); + assert_valid_partition_with_runtime_metric(&leaves, 72, 12, 3); } fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) @@ -1061,20 +1060,20 @@ mod tests { let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); let converted: Vec = raw.iter().copied().map(&convert).collect(); let config = config(2, 16, vec![2, 1], 1); - let expected = partition( + let expected = partition_with_runtime_metric( MatrixView::try_from(&f32_data, points, dimensions).unwrap(), &config, Metric::L2, ) .unwrap(); - let actual = partition( + let actual = partition_with_runtime_metric( MatrixView::try_from(&converted, points, dimensions).unwrap(), &config, Metric::L2, ) .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); - assert_valid_partition(&actual, points, 16, 1); + assert_valid_partition_with_runtime_metric(&actual, points, 16, 1); assert_eq!( sorted_memberships(&actual), sorted_memberships(&expected), @@ -1148,8 +1147,8 @@ mod tests { Metric::CosineNormalized, Metric::InnerProduct, ] { - let leaves = partition(data.as_view(), &config, metric).unwrap(); - assert_valid_partition(&leaves, 64, 20, 1); + let leaves = partition_with_runtime_metric(data.as_view(), &config, metric).unwrap(); + assert_valid_partition_with_runtime_metric(&leaves, 64, 20, 1); } } @@ -1228,7 +1227,9 @@ mod tests { #[test] fn rejects_empty_dataset() { let data = Matrix::::new(0.0, 0, 4); - let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = + partition_with_runtime_metric(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2) + .unwrap_err(); assert_eq!( error.downcast::().unwrap(), @@ -1239,7 +1240,9 @@ mod tests { #[test] fn rejects_zero_dimensions() { let data = Matrix::::new(0.0, 4, 0); - let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = + partition_with_runtime_metric(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2) + .unwrap_err(); assert_eq!( error.downcast::().unwrap(), From 14dade25b38fbdaaad04f2b282532826397316bb Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:58:45 +0000 Subject: [PATCH 135/229] docs(pipnn): describe the active core flow --- diskann/src/graph/pipnn/finalization.rs | 40 +++--- diskann/src/graph/pipnn/leaf_build.rs | 51 ++++---- diskann/src/graph/pipnn/mod.rs | 82 ++++++------ diskann/src/graph/pipnn/partitioning.rs | 158 ++++++++++++------------ 4 files changed, 162 insertions(+), 169 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 7d0ccb722d..cf1f31d234 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -3,17 +3,18 @@ * Licensed under the MIT license. */ -//! Final graph-degree enforcement through the shared Vamana RobustPrune kernel. +//! Graph-degree enforcement with the Vamana RobustPrune kernel. //! -//! Candidate merging may produce more than `R` IDs for a point. This stage first -//! validates every global ID, then processes each point's candidate list in the -//! caller's Rayon pool. Lists already within the degree bound are returned without -//! distance work. Overfull lists are converted to source-distance candidates, -//! passed through RobustPrune, and rewritten from the selected output. +//! Candidate merging can produce more than `R` IDs for one point. This module +//! checks every global ID before parallel work starts. It returns a list with at +//! most `R` IDs without distance work. //! -//! The shared kernel owns occlusion and alpha-round semantics; this adapter owns -//! only contiguous dataset access and distance specialization for the source -//! representation. +//! For a longer list, the module computes each source distance. It sorts the +//! candidates and calls RobustPrune. The module then writes the selected IDs into +//! the original list allocation. +//! +//! RobustPrune defines occlusion and alpha-round behavior. This module supplies +//! contiguous vector access and a distance function for input type `T`. //! //! ```text //! candidate lists ──> validate point count and every global ID @@ -64,10 +65,11 @@ pub(crate) enum FinalizationError { TooManyCandidates { actual: usize, max: usize }, } -/// Per-Rayon-job preparation and kernel state retained across source points. +/// Reusable buffers for one Rayon job. /// -/// PiPNN owns sorting, allocation, and ID translation. The shared internal -/// kernel receives only the prepared candidates and an exactly sized state slice. +/// `pool` stores candidates with source distances. `prepared` stores each +/// distance and optional non-self ID. `states` stores one RobustPrune state for +/// each sorted candidate. #[derive(Default)] struct Workspace { pool: Vec>, @@ -75,7 +77,7 @@ struct Workspace { states: Vec, } -/// Validate candidate IDs and prune only lists whose length exceeds graph degree. +/// Check candidate IDs and prune each list that exceeds the graph degree. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, @@ -90,7 +92,7 @@ where let degree = graph.pruned_degree().get(); let distance = T::distance(metric, Some(data.ncols())); - // build_graph installs the complete call tree in the caller-owned pool. + // `build_graph` runs this Rayon operation in the pool from the build context. #[allow(clippy::disallowed_methods)] candidates .into_par_iter() @@ -98,8 +100,8 @@ where .map_init( Workspace::default, |workspace, (source, mut source_candidates)| { - // Candidate accumulators already enforce uniqueness. A bounded list - // therefore satisfies the graph policy without distance evaluation. + // Candidate merging already removes duplicate IDs. A list within + // the degree limit needs no distance calculation. if source_candidates.len() <= degree { return Ok(source_candidates); } @@ -134,9 +136,9 @@ where .try_reserve(candidate_count) .map_err(ANNError::new)?; - // Sorting/capping precedes source exclusion so filtering cannot - // backfill with farther candidates. Passing this witness into - // RobustPrune makes source-distance order part of its input type. + // Sort all candidates before the code marks a self-edge as absent. + // Thus, self-edge removal cannot add a farther candidate. The + // `SortedNeighbors` value carries this order into RobustPrune. let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); workspace.prepared.extend(sorted.iter().map(|neighbor| { let id = *neighbor.id(); diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 3ff70284d9..c2143ddb6e 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -5,18 +5,18 @@ //! Leaf-local graph construction and candidate accumulation. //! -//! Partitioning supplies strictly increasing, unique global point IDs per leaf. -//! For each leaf this module: +//! Partitioning supplies sorted, unique global point IDs for each leaf. One leaf +//! job does these steps: //! -//! 1. validates IDs and converts only those point vectors to reusable `f32` scratch; -//! 2. computes the lower triangle of `A · Aᵀ`; -//! 3. runs the dual-endpoint leaf top-k kernel; and -//! 4. translates leaf-local positions back to dataset IDs. +//! 1. Check each ID and convert its vector to reusable `f32` storage. +//! 2. Compute the lower triangle of `A · Aᵀ`. +//! 3. Select local neighbors for both points of each pair. +//! 4. Convert local positions to global point IDs. +//! 5. Add both edge directions to global candidate lists. //! -//! The final step merges symmetric adjacency lists under per-point locks because -//! overlapping leaves are processed concurrently. Numeric buffers retain their -//! high-water length; every consumer therefore receives an explicit active -//! prefix rather than treating `Vec::len()` as the current leaf shape. +//! Overlapping leaves run concurrently. A worker locks one destination list only +//! while it adds one leaf's IDs. Reusable buffers keep their largest allocation. +//! Each operation uses an explicit active prefix. use std::{collections::TryReserveError, sync::Mutex}; @@ -92,12 +92,10 @@ pub(crate) enum LeafBuildError { PoisonedCandidateList { point: u32 }, } -/// Scratch leased to one Rayon job and reused for successive leaves. +/// Reusable buffers for one Rayon leaf job. /// -/// The three numerical vectors retain their largest observed leaf shape. The -/// adjacency lists are prepared separately because zero-k/singleton leaves never -/// write them, and because later candidate-merging modes do not necessarily use -/// this representation. +/// The numerical vectors keep the largest leaf shape that this job observed. +/// The job creates local adjacency lists only when the effective `k` is not zero. #[derive(Default)] struct LeafBuffers { point_values: Vec, @@ -166,12 +164,11 @@ impl LeafBuffers { } } -/// Concurrent accumulator indexed by global dataset ID. +/// Concurrent candidate lists indexed by global point ID. /// -/// A point may appear in several overlapping leaves, so workers lock only the -/// destination list long enough to append one leaf's additions. Sorting and -/// duplicate removal are deferred until all leaves finish; doing either under -/// the lock would lengthen the contended section for no semantic benefit. +/// A point can occur in several overlapping leaves. A worker locks one point's +/// list and adds all IDs from one leaf. `AdjacencyList` removes duplicates during +/// this append. The function sorts each list after all leaf jobs finish. struct DirectCandidates { lists: Vec>>, } @@ -192,7 +189,7 @@ impl DirectCandidates { local_adjacency: &[AdjacencyList], ) -> Result<(), LeafBuildError> { for (&source, additions) in point_ids.iter().zip(local_adjacency) { - // Every point ID is validated before leaf-local work begins. + // `build_leaf` checks every point ID before this append. let candidates = &self.lists[source as usize]; let mut candidates = candidates .lock() @@ -218,9 +215,9 @@ impl DirectCandidates { } } -/// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. +/// Build symmetric leaf-local k-NN graphs and return unique global candidates. /// -/// The caller has already selected `A` and `M` for the complete graph build. +/// The caller supplies concrete architecture `A` and metric `M`. #[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. pub(super) fn build_leaf_candidates( arch: A, @@ -261,11 +258,11 @@ where candidates.into_lists() } -/// Build and publish one leaf's symmetric neighbor lists. +/// Build and publish the symmetric neighbor lists for one leaf. /// -/// Validation precedes all dataset indexing and rejects IDs that are not strictly -/// increasing. Active lengths computed after `prepare` must be used for every -/// later slice because reusable vectors may remain longer than this leaf. +/// The function checks all IDs before it indexes the dataset. IDs must be +/// strictly increasing. Reusable vectors can be longer than this leaf. Every +/// read and write uses the active length from `prepare`. fn build_leaf( arch: A, data: MatrixView<'_, T>, diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 8f813009da..a7d8c9ff27 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,45 +5,41 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! PiPNN builds graph candidates in bulk instead of searching a partially built -//! graph for every insertion: +//! PiPNN builds graph candidates in three steps: //! -//! 1. `partitioning` recursively samples leaders and assigns each point to its -//! nearest configured fanout, producing overlapping leaves bounded by `c_max`. -//! 2. `leaf_build` gathers each leaf, computes the lower triangle of `A · Aᵀ`, -//! selects up to three local neighbors per point, and merges symmetric global -//! candidate IDs across overlapping leaves. -//! 3. `finalization` applies the shared Vamana RobustPrune policy only to -//! candidate lists that exceed the configured graph degree. +//! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has +//! at most `c_max` points. +//! 2. `leaf_build` computes a lower-triangular Gram matrix for each leaf. It +//! selects local neighbors and merges their global point IDs. +//! 3. `finalization` applies Vamana RobustPrune to each candidate list that is +//! longer than the graph degree. //! //! ```text //! build_graph //! └─ caller Rayon pool -//! └─ architecture dispatch + metric match once per build +//! └─ select architecture A and metric M once //! └─ build_graph_for //! ├─ partitioning::partition -//! │ └─ partition_kernel::nearest_leaders every stripe +//! │ └─ partition_kernel::nearest_leaders per stripe //! ├─ leaf_build::build_leaf_candidates -//! │ └─ leaf_kernel::nearest_neighbors every leaf +//! │ └─ leaf_kernel::nearest_neighbors per leaf //! └─ finalization::prune_overfull //! ``` //! -//! `diskann-wide` selects concrete architecture `A`; one four-way match selects -//! concrete metric marker `M`. Both types are then carried through every replica, -//! recursive partition, Rayon job, stripe, and leaf. Numerical kernels therefore -//! contain no runtime metric match, visitor, trait object, stored function pointer, -//! or repeated ISA dispatch. +//! `diskann-wide` selects architecture `A`. One match selects metric marker `M`. +//! The build passes both concrete types through all replicas, recursive +//! partitions, stripes, and leaves. The numerical loops do not dispatch again. //! -//! [`PiPNNConfig`] owns only partition and local-neighbor parameters. -//! [`PiPNNBuildContext`] borrows DiskANN graph policy and the caller-owned Rayon -//! pool. [`build_graph`] borrows a contiguous [`MatrixView`] and returns one -//! dataset-ID adjacency list per real point. Providers, start/frozen points, -//! quantization, persistence, and search remain outside this module. +//! [`PiPNNConfig`] contains partition and local-neighbor parameters. +//! [`PiPNNBuildContext`] borrows graph policy and a Rayon pool. [`build_graph`] +//! borrows one contiguous [`MatrixView`]. It returns one adjacency list for each +//! input point. //! -//! The partition and leaf stages own disjoint reusable scratch. Stage outputs move -//! forward (`leaves → candidates → adjacency`) so large temporary allocations can -//! drop at their consumption boundary. Kernel modules document validation, -//! numerical edge cases, tie order, scalar tails, and unchecked SIMD preconditions. +//! The function does not load providers or select start and frozen points. It +//! also does not quantize, serialize, or search the graph. +//! +//! Partition and leaf work use separate reusable buffers. The build consumes +//! each stage output before it creates the next graph representation. mod kernel_metric; mod simd; @@ -71,10 +67,10 @@ use rayon::ThreadPool; use self::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; -/// Configuration of PiPNN's partitioning and local-neighbor algorithm. +/// Configuration for PiPNN partitioning and local-neighbor selection. /// -/// Graph degree, pruning policy, and alpha belong to DiskANN's graph -/// configuration and are supplied separately through [`PiPNNBuildContext`]. +/// [`PiPNNBuildContext`] supplies graph degree, prune policy, alpha, metric, and +/// the Rayon pool. #[derive(Clone, Debug, PartialEq)] pub struct PiPNNConfig { /// Maximum number of points in a leaf. @@ -132,7 +128,7 @@ impl PiPNNConfig { } } -/// Validated, borrowed policy and execution context for one PiPNN graph build. +/// Checked policy and execution inputs for one PiPNN graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { pub(crate) config: PiPNNConfig, @@ -142,7 +138,7 @@ pub struct PiPNNBuildContext<'a> { } impl<'a> PiPNNBuildContext<'a> { - /// Validate and combine PiPNN configuration with outer graph policy. + /// Check and combine PiPNN configuration with DiskANN graph policy. pub fn new( config: PiPNNConfig, graph: &'a Config, @@ -166,12 +162,14 @@ impl<'a> PiPNNBuildContext<'a> { } } -/// Build PiPNN adjacency for real points in `data`. +/// Build PiPNN adjacency for all points in `data`. +/// +/// The function does not select start or frozen points. It does not load a +/// provider or write an index. /// -/// This is the core algorithm boundary. Search entry-point selection, frozen nodes, -/// providers, serialization, and index writers belong to the outer build pipelines. -/// For raw `u8` and `i8` vectors, `CosineNormalized` is evaluated as `Cosine` because -/// those representations are converted to f32 scratch but are not unit-normalized. +/// Raw `u8` and `i8` vectors are not unit-normalized after conversion to `f32`. +/// Therefore, the function evaluates `CosineNormalized` as `Cosine` for these +/// two input types. pub fn build_graph( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, @@ -203,8 +201,8 @@ where data.nrows() ))); } - // Integer source vectors are not guaranteed unit-normalized after conversion, - // so their normalized-cosine request must use the norm-aware formula. + // Conversion does not make integer vectors unit length. Use the norm-aware + // cosine formula for these vectors. let metric = effective_metric::(context.metric); arch::dispatch1_no_features( RunBuildGraph, @@ -266,14 +264,14 @@ where { let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition::(arch, data, &context.config))?; - // `leaves` is consumed here. Workers borrow individual ID lists during the - // parallel pass, and the complete partition allocation drops on return. + // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, + // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.k) .map_err(ANNError::new) })?; - // Finalization consumes candidate lists and reuses their allocations for the - // resulting adjacency where possible. + // Finalization consumes each candidate list. It reuses that list's allocation + // for the final adjacency when the graph policy permits it. tracing::info_span!("pipnn.finalization") .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, M::METRIC)) } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 4c9fa56842..c208a9685b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -5,37 +5,38 @@ //! Deterministic overlapping partition construction for PiPNN. //! -//! The stage maps real dataset points to bounded leaf ID lists. Numerical work -//! reuses the partition kernel and dense GEMM. A stage-owned pool leases scratch -//! to Rayon chunks and takes it back after each chunk; computation never holds -//! the pool lock, and no thread-local cleanup protocol is required. +//! This module converts dataset point IDs into bounded leaf ID lists. It uses +//! dense GEMM and the partition kernel for leader assignment. An `ObjectPool` +//! supplies one reusable buffer set to each Rayon worker chunk. The pool lock is +//! not held during gather, GEMM, or top-k selection. //! //! ```text -//! replica root IDs ──> work queue -//! │ -//! v -//! sample leaders -//! │ -//! gather stripes ─> GEMM distances ─> nearest leaders -//! │ -//! v -//! stable scatter by leader -//! │ │ -//! size <= c_max oversized cluster -//! │ │ -//! completed leaf next recursion level -//! └──────────┬────────┘ -//! v -//! global small-leaf merge -//! v -//! coverage/bound validation +//! replica point IDs ──> work queue +//! │ +//! v +//! sample leaders +//! │ +//! gather stripes ─> GEMM ─> nearest leaders +//! │ +//! v +//! stable scatter by leader +//! │ │ +//! size <= c_max oversized cluster +//! │ │ +//! completed leaf work queue +//! └────────┬────────┘ +//! v +//! merge undersized leaves +//! │ +//! v +//! check leaf bounds //! ``` //! -//! | Recursion level | Assignment multiplicity | +//! | Condition | Assignments per point | //! | --- | --- | //! | `level < fanout.len()` | `fanout[level]` nearest leaders | -//! | later levels | one nearest leader until bounded | -//! | replica boundary | independent deterministic seed | +//! | `level >= fanout.len()` | one nearest leader | +//! | new replica | independent deterministic seed | use std::collections::HashSet; @@ -58,7 +59,7 @@ use super::{ }, }; -// Private algorithm and batching constants live together. None are user policy. +// These constants control internal batching and deterministic seed generation. const PARTITION_SEED: u64 = 1_000; const REPLICA_SEED_STEP: u64 = 7_919; const LEADER_CAP: usize = 1_000; @@ -68,7 +69,7 @@ const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const MAX_PARTITION_ITERATIONS: usize = 30; -/// A partition failure with enough context to diagnose non-progressing input. +/// Error from partition input checks, allocation, or recursion progress. #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum PartitionError { #[error("PiPNN cannot partition an empty dataset")] @@ -124,27 +125,27 @@ impl AsPooled<()> for StripeBuffers { } fn modify(&mut self, _: ()) { - // Scratch retains its high-water allocation across leases; active - // prefixes are established by `assign_stripe` before every read. + // Keep the largest allocation across leases. `assign_stripe` defines the + // active prefix before each read. } } -/// Stage-owned high-water scratch storage for partition assignment. +/// Reusable partition buffers owned by this partition call. /// -/// `ObjectPool` owns the short pop/push lock and returns leases through RAII, -/// including error and panic paths. Gather, GEMM, and top-k run while only the -/// leased `StripeBuffers` is held. +/// `ObjectPool` locks only when it gives or receives a lease. RAII returns each +/// lease after success, error, or panic. Numerical work holds only the lease. type StripeBufferPool = ObjectPool; -/// Partition every configured replica into overlapping bounded leaves. +/// Partition all configured replicas into overlapping bounded leaves. /// -/// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped -/// to the private leader bound), assigns every point to its nearest `fanout` -/// leaders for the current level, and recurses only on oversized clusters. -/// Levels beyond `fanout.len()` retain one leader assignment. Completed small -/// leaves are merged without exceeding `c_max`; every input point must remain -/// covered once per replica. The caller has already selected `A` and `M` and -/// installed the operation in its pool. +/// An oversized work item samples `ceil(p_samp * points)` leaders, up to +/// `LEADER_CAP`. It assigns each point to the configured number of nearest +/// leaders. It adds each cluster above `c_max` to the work queue. +/// +/// A level without a configured fanout assigns each point to one leader. The +/// final merge does not make a leaf larger than `c_max`. Each replica covers +/// every input point. The caller supplies concrete architecture `A` and metric +/// `M`. pub(super) fn partition( arch: A, data: MatrixView<'_, T>, @@ -226,8 +227,8 @@ where .try_reserve_exact(work.len()) .map_err(ANNError::new)?; results.resize_with(work.len(), || None); - // build_graph installs this complete private call tree into the - // caller-owned pool; the indexed fill cannot escape that pool. + // `build_graph` runs this Rayon operation in the pool from the build + // context. Each worker writes only to its indexed result slot. #[allow(clippy::disallowed_methods)] results .par_iter_mut() @@ -337,20 +338,20 @@ fn replica_seed(replica: usize) -> u64 { PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) } -// A single LCG mixer derives recursive seeds. Wrapping makes the mapping stable -// across debug/release builds and supported platforms. +// This LCG derives child seeds. Wrapping arithmetic gives the same mapping in +// debug and release builds on all supported platforms. fn mix_seed(seed: u64, salt: u64) -> u64 { seed.wrapping_mul(6_364_136_223_846_793_005) .wrapping_add(salt) } -/// Assign each point to its nearest `fanout` sampled leaders. +/// Assign each point to its nearest sampled leaders. +/// +/// The function gathers leader vectors once. It divides points into cache-sized +/// stripes. Each worker chunk reuses one leased buffer set for all its stripes. /// -/// Leader vectors are gathered once. Points are processed in cache-sized -/// stripes, while a worker chunk retains one leased scratch buffer across all of -/// its stripes. The flat assignment matrix preserves point order and is then -/// scattered into per-leader clusters; preserving order is required for fixed -/// seed determinism in later recursion levels. +/// The flat assignment matrix keeps point order. The scatter step keeps the same +/// order in each leader cluster. Recursive sampling depends on this order. fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, @@ -381,10 +382,8 @@ where .iter_mut() .zip(leader_values.chunks_exact(dimension_count)) { - // Leader norms participate in the top-k ordering. Preserve the original - // scalar reduction order: reassociating this short setup pass through a - // SIMD norm changes low bits and can send near-tied points down different - // recursive partition paths. + // Leader norms affect top-k order. Use this scalar reduction order. + // SIMD reassociation changes low bits and can change a near-tie branch. *scale = leader_vector.iter().map(|value| value * value).sum(); if M::METRIC == Metric::Cosine { *scale = scale.sqrt(); @@ -401,8 +400,8 @@ where let worker_point_count = checked_area("assignment worker", worker_stripe_count, stripe_points)?; let worker_assignment_count = checked_area("assignment worker", worker_point_count, fanout)?; - // Each worker chunk owns one scratch value and reuses it for its stripes. - // build_graph pins this terminal operation to the caller-owned pool. + // Each worker chunk reuses one buffer lease for all its stripes. + // `build_graph` runs this operation in the pool from the build context. #[allow(clippy::disallowed_methods)] assignments .par_chunks_mut(worker_assignment_count) @@ -459,13 +458,10 @@ where let point_values_len = checked_area("point stripe", point_count, dimensions)?; let dots_len = checked_area("dot-product stripe", point_count, leader_count)?; let output_len = checked_area("partition assignments", point_count, fanout)?; - // Scratch keeps its high-water length and every consumer receives an - // explicit active prefix. Resizing to the exact stripe shape would be - // correct but re-zeroes the buffer whenever a pooled value moves between - // work items with different leader counts: `stripe_points` is derived from - // `leader_count`, so the point buffer swings between roughly 768 KiB and 6 MiB - // and `Vec::resize` only truncates on the way down, then memsets the whole - // delta on the way back up. + // Keep each buffer at its largest length and use an explicit active prefix. + // Different leader counts change the point buffer from about 768 KiB to + // 6 MiB. Exact resizing truncates the buffer and then zeros the full growth + // when another work item needs the larger shape. grow_fallible(&mut buffers.points, point_values_len, 0.0)?; grow_fallible(&mut buffers.dots, dots_len, 0.0)?; let StripeBuffers { @@ -555,12 +551,11 @@ where Ok(()) } -/// Convert the flat point-major assignment matrix into leader-major clusters. +/// Convert point-major assignments into leader-major clusters. /// -/// Small inputs use one serial exact-capacity pass. Large inputs form at most -/// one partial cluster set per Rayon worker, then merge each leader independently. -/// Concatenating partials in stripe order keeps the same member order as the -/// serial implementation while removing a large serial copy tail. +/// A small input uses one serial pass with exact capacities. A large input makes +/// at most one partial cluster set for each Rayon worker. It then merges each +/// leader independently. Stripe-order concatenation matches serial member order. fn scatter_assignments( points: &[u32], assignments: &[u32], @@ -577,7 +572,8 @@ fn scatter_assignments( let mut partials = Vec::new(); partials.try_reserve_exact(stripes).map_err(ANNError::new)?; partials.resize_with(stripes, || None); - // See the pool invariant at the other partition terminal operations. + // `build_graph` runs this Rayon operation in the pool from the build context. + // Each worker writes only to its indexed partial result. #[allow(clippy::disallowed_methods)] partials .par_iter_mut() @@ -609,7 +605,8 @@ fn scatter_assignments( } } - // See the pool invariant at the other partition terminal operations. + // `build_graph` runs this Rayon operation in the pool from the build context. + // Each worker creates one independent leader cluster. #[allow(clippy::disallowed_methods)] sizes .into_par_iter() @@ -776,11 +773,11 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { Ok(values) } -/// Grow `values` to at least `len` elements, never shrinking it. +/// Grow `values` to at least `len` elements and do not shrink it. /// -/// Callers slice the active prefix themselves. Shrinking would force the next -/// larger stripe to re-zero the reclaimed tail, which is the dominant cost when -/// one pooled buffer serves work items with different stripe shapes. +/// Callers use an explicit active prefix. Shrinking and regrowing the buffer +/// zeros the reclaimed tail. This zeroing dominates buffer reuse across +/// different stripe shapes. fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { if values.len() >= len { return Ok(()); @@ -1047,8 +1044,8 @@ mod tests { T: crate::utils::VectorRepr + Send + Sync, { let points = 64; - // Partition gathering converts source vectors before GEMM. Exercise conversion - // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. + // Partition gather converts source vectors before GEMM. Test conversion + // tails around 4, 8, 16, and 32 elements. for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let raw: Vec = (0..points * dimensions) .map(|index| { @@ -1107,10 +1104,9 @@ mod tests { (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 } - // This fixed case sits on opposite sides of the top-1 boundary depending - // on whether leader norms use the original scalar reduction or a SIMD - // reassociation. Point/leader dot products still go through the production - // GEMM; only the setup norm calculation is under test. + // Scalar and SIMD-reassociated leader norms select different top-1 + // leaders for this case. Dot products still use the production GEMM. The + // test changes only the leader-norm reduction. let dimensions = 129; let mut state = 0x3a85_f952_c718_6e49; let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); From ec55050d4868ad02c3d1b62937ad9cc664ee45e8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:03:46 +0000 Subject: [PATCH 136/229] docs(pipnn): remove core diagrams and tuning notes --- diskann/src/graph/pipnn/finalization.rs | 20 ------------ diskann/src/graph/pipnn/mod.rs | 15 ++------- diskann/src/graph/pipnn/partitioning.rs | 42 +++++-------------------- 3 files changed, 9 insertions(+), 68 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index cf1f31d234..c44c35bebe 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -15,26 +15,6 @@ //! //! RobustPrune defines occlusion and alpha-round behavior. This module supplies //! contiguous vector access and a distance function for input type `T`. -//! -//! ```text -//! candidate lists ──> validate point count and every global ID -//! │ -//! ┌────────────┴────────────┐ -//! v v -//! len <= R len > R -//! return list source-distance candidates -//! │ -//! v -//! shared RobustPrune -//! │ -//! v -//! rewrite same list owner -//! ``` -//! -//! | Path | Distance evaluations | Allocation behavior | -//! | --- | --- | --- | -//! | bounded list | none | move list directly to output | -//! | overfull list | source and occlusion distances | reuse Rayon-job workspace | use crate::{ ANNError, ANNResult, diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index a7d8c9ff27..864b4499f6 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -13,18 +13,7 @@ //! selects local neighbors and merges their global point IDs. //! 3. `finalization` applies Vamana RobustPrune to each candidate list that is //! longer than the graph degree. -//! -//! ```text -//! build_graph -//! └─ caller Rayon pool -//! └─ select architecture A and metric M once -//! └─ build_graph_for -//! ├─ partitioning::partition -//! │ └─ partition_kernel::nearest_leaders per stripe -//! ├─ leaf_build::build_leaf_candidates -//! │ └─ leaf_kernel::nearest_neighbors per leaf -//! └─ finalization::prune_overfull -//! ``` + //! //! `diskann-wide` selects architecture `A`. One match selects metric marker `M`. //! The build passes both concrete types through all replicas, recursive @@ -39,7 +28,7 @@ //! also does not quantize, serialize, or search the graph. //! //! Partition and leaf work use separate reusable buffers. The build consumes -//! each stage output before it creates the next graph representation. +//! each output before it creates another graph representation. mod kernel_metric; mod simd; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index c208a9685b..8fb2c5bdb5 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -10,33 +10,9 @@ //! supplies one reusable buffer set to each Rayon worker chunk. The pool lock is //! not held during gather, GEMM, or top-k selection. //! -//! ```text -//! replica point IDs ──> work queue -//! │ -//! v -//! sample leaders -//! │ -//! gather stripes ─> GEMM ─> nearest leaders -//! │ -//! v -//! stable scatter by leader -//! │ │ -//! size <= c_max oversized cluster -//! │ │ -//! completed leaf work queue -//! └────────┬────────┘ -//! v -//! merge undersized leaves -//! │ -//! v -//! check leaf bounds -//! ``` -//! -//! | Condition | Assignments per point | -//! | --- | --- | -//! | `level < fanout.len()` | `fanout[level]` nearest leaders | -//! | `level >= fanout.len()` | one nearest leader | -//! | new replica | independent deterministic seed | +//! A configured level assigns each point to `fanout[level]` leaders. A deeper +//! level assigns each point to one leader. Each replica uses a different +//! deterministic seed. use std::collections::HashSet; @@ -347,7 +323,7 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { /// Assign each point to its nearest sampled leaders. /// -/// The function gathers leader vectors once. It divides points into cache-sized +/// The function gathers leader vectors once. It divides points into bounded /// stripes. Each worker chunk reuses one leased buffer set for all its stripes. /// /// The flat assignment matrix keeps point order. The scatter step keeps the same @@ -458,10 +434,8 @@ where let point_values_len = checked_area("point stripe", point_count, dimensions)?; let dots_len = checked_area("dot-product stripe", point_count, leader_count)?; let output_len = checked_area("partition assignments", point_count, fanout)?; - // Keep each buffer at its largest length and use an explicit active prefix. - // Different leader counts change the point buffer from about 768 KiB to - // 6 MiB. Exact resizing truncates the buffer and then zeros the full growth - // when another work item needs the larger shape. + // Keep each buffer at its largest length. Every operation uses an explicit + // active prefix. grow_fallible(&mut buffers.points, point_values_len, 0.0)?; grow_fallible(&mut buffers.dots, dots_len, 0.0)?; let StripeBuffers { @@ -775,9 +749,7 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { /// Grow `values` to at least `len` elements and do not shrink it. /// -/// Callers use an explicit active prefix. Shrinking and regrowing the buffer -/// zeros the reclaimed tail. This zeroing dominates buffer reuse across -/// different stripe shapes. +/// Callers use an explicit active prefix. fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { if values.len() >= len { return Ok(()); From e51f2c9795cb41cffe23063cb03906cc6c67bee3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:44:03 +0000 Subject: [PATCH 137/229] refactor(pipnn): remove duplicate partition checks --- diskann/src/graph/pipnn/leaf_build.rs | 25 ++++++++++++++++--------- diskann/src/graph/pipnn/partitioning.rs | 1 - 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index c2143ddb6e..062896d161 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -28,8 +28,7 @@ use rayon::prelude::*; use super::{ kernel_metric::KernelMetric, leaf_kernel::{ - LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, leaf_output_len, - nearest_neighbors, + LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, }; @@ -112,7 +111,7 @@ impl LeafBuffers { point_count: usize, dimension_count: usize, requested_k: usize, - ) -> Result { + ) -> Result<(usize, usize), LeafBuildError> { let point_value_count = point_count .checked_mul(dimension_count) @@ -131,8 +130,17 @@ impl LeafBuffers { })?; let leaf_k = leaf_neighbor_count(point_count, requested_k) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - let neighbor_count = leaf_output_len(point_count, requested_k) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + let neighbor_count = + point_count + .checked_mul(leaf_k) + .ok_or_else(|| LeafBuildError::Kernel { + leaf, + source: LeafKernelError::ShapeOverflow { + buffer: "output", + rows: point_count, + cols: leaf_k, + }, + })?; grow( "leaf point values", @@ -147,7 +155,7 @@ impl LeafBuffers { neighbor_count, LeafNeighbor::default(), )?; - Ok(leaf_k) + Ok((leaf_k, neighbor_count)) } fn prepare_local_adjacency(&mut self, point_count: usize) -> Result<(), LeafBuildError> { @@ -301,15 +309,14 @@ where } return Err(LeafBuildError::UnsortedPointIds { leaf }); } - let leaf_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; + let (leaf_k, neighbor_value_count) = + buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; if leaf_k == 0 { return Ok(()); } let point_value_count = point_ids.len() * data.ncols(); let dot_count = point_ids.len() * point_ids.len(); - let neighbor_value_count = leaf_output_len(point_ids.len(), requested_k) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; for (&point, point_output) in point_ids .iter() diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 8fb2c5bdb5..83568138dc 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -157,7 +157,6 @@ where .map_err(ANNError::new)?; leaves.append(&mut replica_leaves); } - validate_leaves(&leaves, config.c_max)?; Ok(leaves) } From b0f46304b029dbaa56926cf1100f0659deb5544d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:48:34 +0000 Subject: [PATCH 138/229] refactor(pipnn): keep leaf shape validation local --- diskann/src/graph/pipnn/leaf_build.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 062896d161..6e787c02bf 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -133,13 +133,10 @@ impl LeafBuffers { let neighbor_count = point_count .checked_mul(leaf_k) - .ok_or_else(|| LeafBuildError::Kernel { + .ok_or(LeafBuildError::ShapeOverflow { leaf, - source: LeafKernelError::ShapeOverflow { - buffer: "output", - rows: point_count, - cols: leaf_k, - }, + rows: point_count, + columns: leaf_k, })?; grow( From 02b43570400b85912544ba990b71d946510c7e0f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:54:25 +0000 Subject: [PATCH 139/229] refactor(pipnn): propagate partition worker errors --- diskann/src/graph/pipnn/partitioning.rs | 50 ++++++++----------------- 1 file changed, 15 insertions(+), 35 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 83568138dc..d5376c4033 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -69,8 +69,6 @@ pub(crate) enum PartitionError { level: usize, limit: usize, }, - #[error("partition produced an invalid leaf of size {size}; expected 1..={limit}")] - InvalidLeaf { size: usize, limit: usize }, #[error("invalid {buffer} length: expected {expected}, got {actual}")] InvalidBufferLength { buffer: &'static str, @@ -208,20 +206,21 @@ where results .par_iter_mut() .zip(work.into_par_iter()) - .for_each(|(slot, item)| { + .try_for_each(|(slot, item)| { *slot = Some(partition_one_level::( arch, data, config, item, stripe_buffers, - )); - }); + )?); + Ok::<(), ANNError>(()) + })?; let mut next_work = Vec::new(); for result in results { let (mut pending, mut finished) = - result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))??; + result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))?; next_work .try_reserve(pending.len()) .map_err(ANNError::new)?; @@ -555,14 +554,15 @@ fn scatter_assignments( .par_chunks(stripe_points) .zip(assignments.par_chunks(stripe_assignment_count)), ) - .for_each(|(slot, (points, assignments))| { - *slot = Some(scatter_serial(points, assignments, fanout, leaders)); - }); + .try_for_each(|(slot, (points, assignments))| { + *slot = Some(scatter_serial(points, assignments, fanout, leaders)?); + Ok::<(), ANNError>(()) + })?; let mut locals = Vec::new(); locals.try_reserve_exact(stripes).map_err(ANNError::new)?; for result in partials { - locals.push(result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))??); + locals.push(result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))?); } let mut sizes = filled_vec(leaders, 0usize)?; @@ -707,7 +707,11 @@ fn global_merge_small( } } - validate_leaves(&merged, c_max)?; + debug_assert!( + merged + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) + ); Ok(merged) } @@ -719,19 +723,6 @@ fn drain_sorted(set: &mut HashSet) -> ANNResult> { Ok(values) } -fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { - if let Some(leaf) = leaves - .iter() - .find(|leaf| leaf.is_empty() || leaf.len() > c_max) - { - return Err(ANNError::new(PartitionError::InvalidLeaf { - size: leaf.len(), - limit: c_max, - })); - } - Ok(()) -} - fn point_ids(points: usize) -> ANNResult> { let mut ids = Vec::new(); ids.try_reserve_exact(points).map_err(ANNError::new)?; @@ -1245,15 +1236,4 @@ mod tests { } ); } - - #[test] - fn rejects_empty_and_oversized_leaves() { - for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { - let error = validate_leaves(&leaves, 2).unwrap_err(); - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidLeaf { size, limit: 2 } - ); - } - } } From 7177087a8ea5527e79ddc16af781a76d9776d107 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:59:51 +0000 Subject: [PATCH 140/229] refactor(pipnn): remove partition assertions --- diskann/src/graph/pipnn/partitioning.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index d5376c4033..2817c38ca3 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -707,11 +707,6 @@ fn global_merge_small( } } - debug_assert!( - merged - .iter() - .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) - ); Ok(merged) } From 6485e983e5b5b8b61e5304522eeaaf0e85382df5 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:45:44 +0000 Subject: [PATCH 141/229] docs(pipnn): state core function contracts --- diskann/src/graph/pipnn/mod.rs | 4 ++++ diskann/src/graph/pipnn/partitioning.rs | 16 ++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 864b4499f6..9dc6f24b40 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -238,6 +238,10 @@ where } } +/// Build the graph with concrete architecture `A` and metric `M`. +/// +/// The function partitions the dataset, builds direct candidates, and applies +/// final graph-degree pruning. fn build_graph_for( arch: A, data: MatrixView<'_, T>, diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 2817c38ca3..483098bd4b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -158,6 +158,10 @@ where Ok(leaves) } +/// Partition one replica until each leaf has at most `c_max` points. +/// +/// The function processes one work queue per recursion level. It merges leaves +/// smaller than `c_min` after the queue becomes empty. fn partition_replica( arch: A, data: MatrixView<'_, T>, @@ -244,6 +248,10 @@ where })) } +/// Process one partition work item. +/// +/// The function samples leaders, assigns all points, and separates complete +/// leaves from clusters that require another recursion level. fn partition_one_level( arch: A, data: MatrixView<'_, T>, @@ -406,6 +414,10 @@ where scatter_assignments(point_ids, &assignments, fanout, leader_ids.len()) } +/// Assign one point stripe to the selected leaders. +/// +/// The function gathers point vectors, computes point-to-leader dot products, +/// and writes nearest leader-column IDs to `assignments`. #[inline] #[allow(clippy::too_many_arguments)] fn assign_stripe( @@ -640,6 +652,10 @@ fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { Ok(clusters) } +/// Merge leaves smaller than `c_min` without exceeding `c_max`. +/// +/// A `HashSet` removes duplicate point IDs across merged leaves. The function +/// sorts each merged result before it returns. fn global_merge_small( leaves: Vec>, c_min: usize, From 997e832568238d8d3b873792943f50f66813e770 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:05:49 +0000 Subject: [PATCH 142/229] refactor(pipnn): use domain names in core flow --- diskann/src/graph/pipnn/finalization.rs | 64 ++++++++--------- diskann/src/graph/pipnn/leaf_build.rs | 19 ++--- diskann/src/graph/pipnn/mod.rs | 31 +++++---- diskann/src/graph/pipnn/partitioning.rs | 92 +++++++++++++------------ 4 files changed, 106 insertions(+), 100 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index c44c35bebe..55d86103e6 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -6,15 +6,15 @@ //! Graph-degree enforcement with the Vamana RobustPrune kernel. //! //! Candidate merging can produce more than `R` IDs for one point. This module -//! checks every global ID before parallel work starts. It returns a list with at -//! most `R` IDs without distance work. +//! checks every global ID before parallel work starts. A list at or below `R` +//! returns without distance calculations. //! //! For a longer list, the module computes each source distance. It sorts the //! candidates and calls RobustPrune. The module then writes the selected IDs into //! the original list allocation. //! //! RobustPrune defines occlusion and alpha-round behavior. This module supplies -//! contiguous vector access and a distance function for input type `T`. +//! source vectors and metric distances. use crate::{ ANNError, ANNResult, @@ -45,16 +45,15 @@ pub(crate) enum FinalizationError { TooManyCandidates { actual: usize, max: usize }, } -/// Reusable buffers for one Rayon job. +/// RobustPrune state for one Rayon job. /// -/// `pool` stores candidates with source distances. `prepared` stores each -/// distance and optional non-self ID. `states` stores one RobustPrune state for -/// each sorted candidate. +/// `candidate_slots` and `prune_states` stay positionally aligned with +/// `sorted_candidates`. #[derive(Default)] -struct Workspace { - pool: Vec>, - prepared: Vec<(f32, Option)>, - states: Vec, +struct PruneWorkspace { + sorted_candidates: Vec>, + candidate_slots: Vec<(f32, Option)>, + prune_states: Vec, } /// Check candidate IDs and prune each list that exceeds the graph degree. @@ -78,7 +77,7 @@ where .into_par_iter() .enumerate() .map_init( - Workspace::default, + PruneWorkspace::default, |workspace, (source, mut source_candidates)| { // Candidate merging already removes duplicate IDs. A list within // the degree limit needs no distance calculation. @@ -88,13 +87,13 @@ where let source_id = u32::try_from(source).map_err(ANNError::new)?; let source_vector = data.row(source); - workspace.pool.clear(); + workspace.sorted_candidates.clear(); workspace - .pool + .sorted_candidates .try_reserve(source_candidates.len()) .map_err(ANNError::new)?; workspace - .pool + .sorted_candidates .extend(source_candidates.iter().copied().map(|candidate| { Neighbor::new( candidate, @@ -103,44 +102,47 @@ where ) })); - let candidate_count = workspace.pool.len(); + let candidate_count = workspace.sorted_candidates.len(); if candidate_count > u16::MAX as usize { return Err(ANNError::new(FinalizationError::TooManyCandidates { actual: candidate_count, max: u16::MAX as usize, })); } - workspace.prepared.clear(); + workspace.candidate_slots.clear(); workspace - .prepared + .candidate_slots .try_reserve(candidate_count) .map_err(ANNError::new)?; // Sort all candidates before the code marks a self-edge as absent. // Thus, self-edge removal cannot add a farther candidate. The // `SortedNeighbors` value carries this order into RobustPrune. - let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); - workspace.prepared.extend(sorted.iter().map(|neighbor| { - let id = *neighbor.id(); - (*neighbor.distance(), (id != source_id).then_some(id)) - })); + let sorted = + SortedNeighbors::new(&mut workspace.sorted_candidates, candidate_count); workspace - .states + .candidate_slots + .extend(sorted.iter().map(|neighbor| { + let id = *neighbor.id(); + (*neighbor.distance(), (id != source_id).then_some(id)) + })); + workspace + .prune_states .try_reserve( workspace - .prepared + .candidate_slots .len() - .saturating_sub(workspace.states.len()), + .saturating_sub(workspace.prune_states.len()), ) .map_err(ANNError::new)?; workspace - .states - .resize(workspace.prepared.len(), prune::State::default()); + .prune_states + .resize(workspace.candidate_slots.len(), prune::State::default()); let selected = prune::robust_prune( &sorted, - &workspace.prepared, - workspace.states.as_mut_slice(), + &workspace.candidate_slots, + workspace.prune_states.as_mut_slice(), degree, graph.alpha(), graph.prune_kind(), @@ -153,7 +155,7 @@ where ); let mut guard = source_candidates.resize(selected); - for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { + for (destination, state) in guard.iter_mut().zip(workspace.prune_states.iter()) { *destination = *sorted[state.neighbor as usize].id(); } guard.finish(selected); diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6e787c02bf..25d7df8178 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -194,7 +194,7 @@ impl DirectCandidates { local_adjacency: &[AdjacencyList], ) -> Result<(), LeafBuildError> { for (&source, additions) in point_ids.iter().zip(local_adjacency) { - // `build_leaf` checks every point ID before this append. + // `add_direct_leaf_candidates` checks every point ID before this append. let candidates = &self.lists[source as usize]; let mut candidates = candidates .lock() @@ -220,9 +220,10 @@ impl DirectCandidates { } } -/// Build symmetric leaf-local k-NN graphs and return unique global candidates. +/// Build direct graph candidates from all overlapping leaves. /// -/// The caller supplies concrete architecture `A` and metric `M`. +/// Each selected leaf pair contributes both edge directions. Candidate lists use +/// global dataset IDs and contain no duplicate IDs. #[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. pub(super) fn build_leaf_candidates( arch: A, @@ -249,7 +250,7 @@ where leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf::( + add_direct_leaf_candidates::( arch, data, leaf, @@ -263,12 +264,12 @@ where candidates.into_lists() } -/// Build and publish the symmetric neighbor lists for one leaf. +/// Add one leaf's symmetric neighbors to the direct candidate lists. /// -/// The function checks all IDs before it indexes the dataset. IDs must be -/// strictly increasing. Reusable vectors can be longer than this leaf. Every -/// read and write uses the active length from `prepare`. -fn build_leaf( +/// The function rejects empty, duplicate, unsorted, or out-of-range point IDs. +/// Reusable buffers can be longer than this leaf, so all accesses use the current +/// leaf shape. +fn add_direct_leaf_candidates( arch: A, data: MatrixView<'_, T>, leaf: usize, diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 9dc6f24b40..0e216126fa 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,7 +5,8 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! PiPNN builds graph candidates in three steps: +//! PiPNN builds graph candidates in three steps. A leader is a sampled dataset +//! point that acts as the center of one child partition. //! //! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has //! at most `c_max` points. @@ -13,7 +14,6 @@ //! selects local neighbors and merges their global point IDs. //! 3. `finalization` applies Vamana RobustPrune to each candidate list that is //! longer than the graph degree. - //! //! `diskann-wide` selects architecture `A`. One match selects metric marker `M`. //! The build passes both concrete types through all replicas, recursive @@ -56,19 +56,18 @@ use rayon::ThreadPool; use self::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; -/// Configuration for PiPNN partitioning and local-neighbor selection. +/// PiPNN partition and leaf-selection policy. /// -/// [`PiPNNBuildContext`] supplies graph degree, prune policy, alpha, metric, and -/// the Rayon pool. +/// DiskANN graph policy separately supplies degree, alpha, and prune metric. #[derive(Clone, Debug, PartialEq)] pub struct PiPNNConfig { /// Maximum number of points in a leaf. pub c_max: usize, /// Minimum leaf size used by global small-leaf merging. pub c_min: usize, - /// Fraction of a cluster sampled as partition leaders. + /// Fraction of a cluster sampled as child-partition centers. pub p_samp: f64, - /// Number of nearest leaders retained at each overlapping partition level. + /// Number of nearest partition centers assigned to each point at each level. pub fanout: Vec, /// Number of nearest neighbors selected within each leaf (`1..=3`). pub k: usize, @@ -117,7 +116,7 @@ impl PiPNNConfig { } } -/// Checked policy and execution inputs for one PiPNN graph build. +/// PiPNN policy and borrowed execution resources for one graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { pub(crate) config: PiPNNConfig, @@ -151,14 +150,13 @@ impl<'a> PiPNNBuildContext<'a> { } } -/// Build PiPNN adjacency for all points in `data`. +/// Build one PiPNN adjacency list for each point in `data`. /// -/// The function does not select start or frozen points. It does not load a -/// provider or write an index. +/// This graph contains only real dataset points. Start-point selection and index +/// serialization are separate operations. /// /// Raw `u8` and `i8` vectors are not unit-normalized after conversion to `f32`. -/// Therefore, the function evaluates `CosineNormalized` as `Cosine` for these -/// two input types. +/// The build therefore uses norm-aware cosine for these two input types. pub fn build_graph( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, @@ -166,10 +164,13 @@ pub fn build_graph( where T: VectorRepr + Send + Sync + 'static, { - context.pool.install(|| build_graph_inner(data, context)) + context + .pool + .install(|| validate_and_dispatch_build(data, context)) } -fn build_graph_inner( +/// Check dataset bounds and select the architecture and metric implementation. +fn validate_and_dispatch_build( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, ) -> ANNResult>> diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 483098bd4b..3d953ee461 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -5,10 +5,12 @@ //! Deterministic overlapping partition construction for PiPNN. //! -//! This module converts dataset point IDs into bounded leaf ID lists. It uses -//! dense GEMM and the partition kernel for leader assignment. An `ObjectPool` -//! supplies one reusable buffer set to each Rayon worker chunk. The pool lock is -//! not held during gather, GEMM, or top-k selection. +//! A leader is a sampled point that acts as the center of one child partition. +//! A point can join several leaders, so child partitions can overlap. +//! +//! This module recursively splits dataset point IDs into bounded leaves. It uses +//! dense GEMM to compare points with sampled leaders. An `ObjectPool` supplies +//! reusable buffers to Rayon worker chunks. //! //! A configured level assigns each point to `fanout[level]` leaders. A deeper //! level assigns each point to one leader. Each replica uses a different @@ -99,27 +101,23 @@ impl AsPooled<()> for StripeBuffers { } fn modify(&mut self, _: ()) { - // Keep the largest allocation across leases. `assign_stripe` defines the + // Keep the largest allocation across leases. `assign_point_stripe` defines the // active prefix before each read. } } -/// Reusable partition buffers owned by this partition call. +/// Reusable buffers for point-to-leader assignment. /// -/// `ObjectPool` locks only when it gives or receives a lease. RAII returns each -/// lease after success, error, or panic. Numerical work holds only the lease. +/// `ObjectPool` locks only when it gives or receives a lease. Numerical work +/// holds the lease, not the pool lock. type StripeBufferPool = ObjectPool; -/// Partition all configured replicas into overlapping bounded leaves. -/// -/// An oversized work item samples `ceil(p_samp * points)` leaders, up to -/// `LEADER_CAP`. It assigns each point to the configured number of nearest -/// leaders. It adds each cluster above `c_max` to the work queue. +/// Build overlapping bounded leaves for all configured replicas. /// -/// A level without a configured fanout assigns each point to one leader. The -/// final merge does not make a leaf larger than `c_max`. Each replica covers -/// every input point. The caller supplies concrete architecture `A` and metric -/// `M`. +/// Each split samples partition centers and assigns every cluster point to its +/// nearest centers. A cluster above `c_max` is split again. A level without a +/// configured fanout assigns each point to one center. Each replica covers every +/// input point. pub(super) fn partition( arch: A, data: MatrixView<'_, T>, @@ -196,7 +194,7 @@ where for _ in 0..MAX_PARTITION_ITERATIONS { if work.is_empty() { - return global_merge_small(leaves, config.c_min, config.c_max); + return merge_undersized_leaves(leaves, config.c_min, config.c_max); } let mut results = Vec::new(); @@ -211,7 +209,7 @@ where .par_iter_mut() .zip(work.into_par_iter()) .try_for_each(|(slot, item)| { - *slot = Some(partition_one_level::( + *slot = Some(partition_work_item::( arch, data, config, @@ -236,10 +234,10 @@ where } if work.is_empty() { - return global_merge_small(leaves, config.c_min, config.c_max); + return merge_undersized_leaves(leaves, config.c_min, config.c_max); } let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { - return global_merge_small(leaves, config.c_min, config.c_max); + return merge_undersized_leaves(leaves, config.c_min, config.c_max); }; Err(ANNError::new(PartitionError::IterationLimit { size: largest.indices.len(), @@ -248,11 +246,11 @@ where })) } -/// Process one partition work item. +/// Split one oversized cluster into child partitions. /// -/// The function samples leaders, assigns all points, and separates complete -/// leaves from clusters that require another recursion level. -fn partition_one_level( +/// The function samples center points, assigns the cluster points, and returns +/// bounded leaves separately from child clusters that need another split. +fn partition_work_item( arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, @@ -301,8 +299,9 @@ where Ok((pending, finished)) } +/// Sample point IDs that act as centers for one partition split. fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResult> { - let count = sample_num_leaders(points.len(), sampling_fraction); + let count = sampled_leader_count(points.len(), sampling_fraction); let mut rng = rand::rngs::StdRng::seed_from_u64(seed); let mut leaders = Vec::new(); leaders.try_reserve_exact(count).map_err(ANNError::new)?; @@ -310,7 +309,12 @@ fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResul Ok(leaders) } -fn sample_num_leaders(points: usize, sampling_fraction: f64) -> usize { +/// Return the number of centers to sample from one cluster. +/// +/// The count is `ceil(points * sampling_fraction)`, limited by `LEADER_CAP` and +/// the number of available points. A cluster with at least two points uses at +/// least two centers. +fn sampled_leader_count(points: usize, sampling_fraction: f64) -> usize { ((points as f64 * sampling_fraction).ceil() as usize) .clamp(2, LEADER_CAP) .min(points) @@ -327,13 +331,11 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { .wrapping_add(salt) } -/// Assign each point to its nearest sampled leaders. -/// -/// The function gathers leader vectors once. It divides points into bounded -/// stripes. Each worker chunk reuses one leased buffer set for all its stripes. +/// Assign each cluster point to its nearest sampled partition centers. /// -/// The flat assignment matrix keeps point order. The scatter step keeps the same -/// order in each leader cluster. Recursive sampling depends on this order. +/// The function gathers center vectors once and evaluates points in bounded +/// stripes. The assignment matrix keeps point order. Scatter preserves this order +/// inside each child partition, which makes recursive sampling deterministic. fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, @@ -397,7 +399,7 @@ where { let first_point = worker_first + stripe * stripe_points; let stripe_point_count = stripe_assignments.len() / fanout; - assign_stripe::( + assign_point_stripe::( arch, data, &point_ids[first_point..first_point + stripe_point_count], @@ -414,13 +416,13 @@ where scatter_assignments(point_ids, &assignments, fanout, leader_ids.len()) } -/// Assign one point stripe to the selected leaders. +/// Assign one point stripe to sampled partition centers. /// -/// The function gathers point vectors, computes point-to-leader dot products, -/// and writes nearest leader-column IDs to `assignments`. +/// The function gathers point vectors and computes point-to-center dot products. +/// It writes center-column IDs for partition scatter. #[inline] #[allow(clippy::too_many_arguments)] -fn assign_stripe( +fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, point_ids: &[u32], @@ -656,7 +658,7 @@ fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { /// /// A `HashSet` removes duplicate point IDs across merged leaves. The function /// sorts each merged result before it returns. -fn global_merge_small( +fn merge_undersized_leaves( leaves: Vec>, c_min: usize, c_max: usize, @@ -975,7 +977,7 @@ mod tests { fn global_merge_canonicalizes_small_leaf_membership() { let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; - let merged = global_merge_small(leaves, 4, 8).unwrap(); + let merged = merge_undersized_leaves(leaves, 4, 8).unwrap(); assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); } @@ -984,7 +986,7 @@ mod tests { fn global_merge_never_overfills_before_reaching_c_min() { let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; - let merged = global_merge_small(leaves, 11, 11).unwrap(); + let merged = merge_undersized_leaves(leaves, 11, 11).unwrap(); assert_eq!( merged, @@ -994,7 +996,7 @@ mod tests { #[test] fn global_merge_fills_exact_capacity_before_flushing() { - let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + let merged = merge_undersized_leaves(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); assert_eq!(merged, vec![vec![0, 1, 2, 3]]); } @@ -1123,9 +1125,9 @@ mod tests { #[test] fn leader_count_is_bounded() { - assert_eq!(sample_num_leaders(1, 1.0), 1); - assert_eq!(sample_num_leaders(10, 0.01), 2); - assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); + assert_eq!(sampled_leader_count(1, 1.0), 1); + assert_eq!(sampled_leader_count(10, 0.01), 2); + assert_eq!(sampled_leader_count(50_000, 1.0), LEADER_CAP); } #[test] From c1d447c57854b4c94b486790b16d17b37b59b6fa Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:39:55 +0000 Subject: [PATCH 143/229] docs(pipnn): remove layout-restatement comments --- diskann/src/graph/pipnn/mod.rs | 4 ++-- diskann/src/graph/pipnn/partitioning.rs | 10 +++------- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 0e216126fa..773d3e97e0 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -239,9 +239,9 @@ where } } -/// Build the graph with concrete architecture `A` and metric `M`. +/// Run the PiPNN graph pipeline for one selected metric implementation. /// -/// The function partitions the dataset, builds direct candidates, and applies +/// The function builds overlapping leaves, merges direct candidates, and applies /// final graph-degree pruning. fn build_graph_for( arch: A, diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 3d953ee461..c988b38c3b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -537,11 +537,10 @@ where Ok(()) } -/// Convert point-major assignments into leader-major clusters. +/// Group assigned point IDs by child partition. /// -/// A small input uses one serial pass with exact capacities. A large input makes -/// at most one partial cluster set for each Rayon worker. It then merges each -/// leader independently. Stripe-order concatenation matches serial member order. +/// Both the serial and parallel paths preserve point order inside each child. +/// This order is required for deterministic recursive sampling. fn scatter_assignments( points: &[u32], assignments: &[u32], @@ -750,9 +749,6 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { Ok(values) } -/// Grow `values` to at least `len` elements and do not shrink it. -/// -/// Callers use an explicit active prefix. fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { if values.len() >= len { return Ok(()); From 93d35f6c98f7922914110541f5513973542c9de7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:48:39 +0000 Subject: [PATCH 144/229] docs(pipnn): define leaf domain term --- diskann/src/graph/pipnn/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 773d3e97e0..45d74f2bf0 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -6,7 +6,8 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! //! PiPNN builds graph candidates in three steps. A leader is a sampled dataset -//! point that acts as the center of one child partition. +//! point that acts as the center of one child partition. A leaf is a bounded +//! child partition used for local neighbor selection. //! //! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has //! at most `c_max` points. From 55ad5b35d5b1c0189352d1ff1ee5bf05ebecddec Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:05:33 +0000 Subject: [PATCH 145/229] refactor(pipnn): inline poisoned-list errors --- diskann/src/graph/pipnn/leaf_build.rs | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 25d7df8178..6870c05900 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -198,7 +198,7 @@ impl DirectCandidates { let candidates = &self.lists[source as usize]; let mut candidates = candidates .lock() - .map_err(|_| poisoned_candidate_list(source))?; + .map_err(|_| LeafBuildError::PoisonedCandidateList { point: source })?; candidates.extend_from_slice(additions); } Ok(()) @@ -210,9 +210,12 @@ impl DirectCandidates { .try_reserve_exact(self.lists.len()) .map_err(|source| allocation_error("candidate output", self.lists.len(), source))?; for (point, candidates) in self.lists.into_iter().enumerate() { - let mut candidates = candidates - .into_inner() - .map_err(|_| poisoned_candidate_list(point as u32))?; + let mut candidates = + candidates + .into_inner() + .map_err(|_| LeafBuildError::PoisonedCandidateList { + point: point as u32, + })?; candidates.sort(); output.push(candidates); } @@ -427,10 +430,6 @@ fn allocation_error( } } -fn poisoned_candidate_list(point: u32) -> LeafBuildError { - LeafBuildError::PoisonedCandidateList { point } -} - #[cfg(test)] mod tests { use diskann_utils::views::MatrixView; From 708ecc8fe0a7a8ea4b05b9f1f4adcf9b3528994e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:55:56 +0000 Subject: [PATCH 146/229] refactor(pipnn): clarify leaf policy names --- diskann/src/graph/pipnn/mod.rs | 33 +++++++++++++------------ diskann/src/graph/pipnn/partitioning.rs | 2 +- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 45d74f2bf0..785ef254a5 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -68,10 +68,11 @@ pub struct PiPNNConfig { pub c_min: usize, /// Fraction of a cluster sampled as child-partition centers. pub p_samp: f64, - /// Number of nearest partition centers assigned to each point at each level. + /// Number of nearest centers assigned at each recursive partition level. + /// Levels after this schedule assign each point to one center. pub fanout: Vec, /// Number of nearest neighbors selected within each leaf (`1..=3`). - pub k: usize, + pub leaf_k: usize, /// Number of independent partition passes over the dataset. pub replicas: usize, } @@ -91,9 +92,9 @@ impl PiPNNConfig { self.c_min, self.c_max ))); } - if !self.p_samp.is_finite() || !(0.0..=1.0).contains(&self.p_samp) || self.p_samp == 0.0 { + if !(0.0 < self.p_samp && self.p_samp <= 1.0) { return Err(config_error(format!( - "p_samp ({}) must be finite and in (0, 1]", + "p_samp ({}) must be in (0, 1]", self.p_samp ))); } @@ -103,10 +104,10 @@ impl PiPNNConfig { if self.fanout.contains(&0) { return Err(config_error("fanout values must be greater than zero")); } - if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.k) { + if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.leaf_k) { return Err(config_error(format!( - "k ({}) must be in [1, {}]", - self.k, + "leaf_k ({}) must be in [1, {}]", + self.leaf_k, leaf_kernel::MAX_LEAF_NEIGHBORS ))); } @@ -262,7 +263,7 @@ where // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.k) + leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) .map_err(ANNError::new) })?; // Finalization consumes each candidate list. It reuses that list's allocation @@ -333,7 +334,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - k: 1, + leaf_k: 1, replicas: 1, } } @@ -410,7 +411,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - k: leaf_kernel::MAX_LEAF_NEIGHBORS, + leaf_k: leaf_kernel::MAX_LEAF_NEIGHBORS, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -476,7 +477,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - k: 1, + leaf_k: 1, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); @@ -502,7 +503,7 @@ mod build_graph_tests { c_min: 4, p_samp: 0.25, fanout: vec![3, 2], - k: 3, + leaf_k: 3, replicas: 2, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -534,7 +535,7 @@ mod build_graph_tests { c_min, p_samp: 0.5, fanout: vec![2], - k: rng.random_range(1..=3), + leaf_k: rng.random_range(1..=3), replicas: rng.random_range(1..=2), }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -562,7 +563,7 @@ mod config_tests { c_min: 64, p_samp: 0.01, fanout: vec![10, 3], - k: 2, + leaf_k: 2, replicas: 1, } } @@ -624,11 +625,11 @@ mod config_tests { ..pipnn_config() }, PiPNNConfig { - k: 0, + leaf_k: 0, ..pipnn_config() }, PiPNNConfig { - k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, + leaf_k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, ..pipnn_config() }, PiPNNConfig { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index c988b38c3b..b00e5bd02b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -834,7 +834,7 @@ mod tests { c_min, p_samp: 0.25, fanout, - k: 1, + leaf_k: 1, replicas, } } From 1dab27dbdbd7b0c2b2ca440aa7885e6295e1f39d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 07:07:10 +0000 Subject: [PATCH 147/229] fix(pipnn): reject malformed assignments --- diskann/src/graph/pipnn/partitioning.rs | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index b00e5bd02b..431107fff1 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -614,6 +614,15 @@ fn scatter_serial( fanout: usize, leaders: usize, ) -> ANNResult>> { + let expected = checked_area("scatter assignments", points.len(), fanout)?; + if assignments.len() != expected { + return Err(ANNError::new(PartitionError::InvalidBufferLength { + buffer: "scatter assignments", + expected, + actual: assignments.len(), + })); + } + let mut sizes = filled_vec(leaders, 0usize)?; for &leader in assignments { let Some(size) = sizes.get_mut(leader as usize) else { @@ -1232,6 +1241,20 @@ mod tests { ); } + #[test] + fn rejects_invalid_assignment_length() { + let error = scatter_serial(&[7, 8], &[0, 1, 0], 2, 2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "scatter assignments", + expected: 4, + actual: 3, + } + ); + } + #[test] fn rejects_assignment_to_an_unknown_leader() { let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); From dbf540da86026f087c7da9c12c50f21dd6ea9130 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:50:58 +0000 Subject: [PATCH 148/229] refactor(pipnn): use stage metric contracts --- diskann/src/graph/pipnn/leaf_build.rs | 6 +-- diskann/src/graph/pipnn/mod.rs | 11 +++-- diskann/src/graph/pipnn/partitioning.rs | 64 ++++++++++++------------- 3 files changed, 42 insertions(+), 39 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6870c05900..04129b98c4 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -26,7 +26,7 @@ use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::KernelMetric, + kernel_metric::LeafKernelMetric, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, @@ -239,7 +239,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: LeafKernelMetric, T: VectorRepr + 'static, { if data.ncols() == 0 { @@ -286,7 +286,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: LeafKernelMetric, T: VectorRepr + 'static, { if point_ids.is_empty() { diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 785ef254a5..cdf6a820a4 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -55,7 +55,9 @@ use diskann_wide::{ }; use rayon::ThreadPool; -use self::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; +use self::kernel_metric::{ + Cosine, CosineNormalized, InnerProduct, L2, LeafKernelMetric, MetricTag, PartitionKernelMetric, +}; /// PiPNN partition and leaf-selection policy. /// @@ -255,7 +257,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: LeafKernelMetric + PartitionKernelMetric, T: VectorRepr + Send + Sync + 'static, { let leaves = tracing::info_span!("pipnn.partition") @@ -268,8 +270,9 @@ where })?; // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. - tracing::info_span!("pipnn.finalization") - .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, M::METRIC)) + tracing::info_span!("pipnn.finalization").in_scope(|| { + finalization::prune_overfull(data, candidates, context.graph, ::METRIC) + }) } fn effective_metric(metric: Metric) -> Metric { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 431107fff1..811d3d369d 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -31,10 +31,8 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::KernelMetric, - partition_kernel::{ - PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders, - }, + kernel_metric::{MetricTag, PartitionKernelMetric}, + partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, }; // These constants control internal batching and deterministic seed generation. @@ -91,7 +89,7 @@ struct WorkItem { struct StripeBuffers { points: Vec, dots: Vec, - point_scales: Vec, + point_squared_norms: Vec, kernel: PartitionKernelWorkspace, } @@ -128,7 +126,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { let points = data.nrows(); @@ -172,7 +170,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { let initial_indices = point_ids(data.nrows())?; @@ -262,7 +260,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { let points = item.indices.len(); @@ -349,28 +347,29 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { + let metric = ::METRIC; let dimension_count = data.ncols(); let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; gather_vectors(data, leader_ids, &mut leader_values)?; - let mut leader_scales = if matches!(M::METRIC, Metric::L2 | Metric::Cosine) { + let mut leader_norm_values = if matches!(metric, Metric::L2 | Metric::Cosine) { filled_vec(leader_ids.len(), 0.0f32)? } else { Vec::new() }; - for (scale, leader_vector) in leader_scales + for (norm_value, leader_vector) in leader_norm_values .iter_mut() .zip(leader_values.chunks_exact(dimension_count)) { // Leader norms affect top-k order. Use this scalar reduction order. // SIMD reassociation changes low bits and can change a near-tie branch. - *scale = leader_vector.iter().map(|value| value * value).sum(); - if M::METRIC == Metric::Cosine { - *scale = scale.sqrt(); + *norm_value = leader_vector.iter().map(|value| value * value).sum(); + if metric == Metric::Cosine { + *norm_value = norm_value.sqrt(); } } @@ -404,7 +403,7 @@ where data, &point_ids[first_point..first_point + stripe_point_count], &leader_values, - &leader_scales, + &leader_norm_values, fanout, &mut buffers, stripe_assignments, @@ -427,7 +426,7 @@ fn assign_point_stripe( data: MatrixView<'_, T>, point_ids: &[u32], leader_values: &[f32], - leader_scales: &[f32], + leader_norm_values: &[f32], fanout: usize, buffers: &mut StripeBuffers, assignments: &mut [u32], @@ -437,7 +436,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr, { let point_count = point_ids.len(); @@ -453,7 +452,7 @@ where let StripeBuffers { points: point_buffer, dots: dot_buffer, - point_scales: point_scale_buffer, + point_squared_norms: point_squared_norm_buffer, kernel: kernel_workspace, } = buffers; let point_values = &mut point_buffer[..point_values_len]; @@ -473,28 +472,29 @@ where ) .map_err(ANNError::new)?; - let point_scales = if M::METRIC == Metric::Cosine { - grow_fallible(point_scale_buffer, point_count, 0.0)?; - let point_scales = &mut point_scale_buffer[..point_count]; - for (scale, point_values) in point_scales + let metric = ::METRIC; + let point_squared_norms = if metric == Metric::Cosine { + grow_fallible(point_squared_norm_buffer, point_count, 0.0)?; + let point_squared_norms = &mut point_squared_norm_buffer[..point_count]; + for (norm_value, point_values) in point_squared_norms .iter_mut() .zip(point_values.chunks_exact(dimensions)) { - *scale = FastL2NormSquared.evaluate(point_values); + *norm_value = FastL2NormSquared.evaluate(point_values); } - &*point_scales + &*point_squared_norms } else { &[] }; - let scales = match M::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, }; let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { ANNError::new(PartitionError::InvalidBufferLength { @@ -512,7 +512,7 @@ where })?; nearest_leaders::( arch, - PartitionInput { dots, scales }, + PartitionInput { dots, norms }, output, kernel_workspace, ) From fca79819f0c4247f4f57b24dd99aee48b5997ae0 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:43:58 +0000 Subject: [PATCH 149/229] refactor(pipnn): use partition metric identity --- diskann/src/graph/pipnn/mod.rs | 9 +++++++-- diskann/src/graph/pipnn/partitioning.rs | 6 +++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index cdf6a820a4..46a4733eff 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -56,7 +56,7 @@ use diskann_wide::{ use rayon::ThreadPool; use self::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, LeafKernelMetric, MetricTag, PartitionKernelMetric, + Cosine, CosineNormalized, InnerProduct, L2, LeafKernelMetric, PartitionKernelMetric, }; /// PiPNN partition and leaf-selection policy. @@ -271,7 +271,12 @@ where // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. tracing::info_span!("pipnn.finalization").in_scope(|| { - finalization::prune_overfull(data, candidates, context.graph, ::METRIC) + finalization::prune_overfull( + data, + candidates, + context.graph, + ::METRIC, + ) }) } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 811d3d369d..11660adafd 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -31,7 +31,7 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::{MetricTag, PartitionKernelMetric}, + kernel_metric::PartitionKernelMetric, partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, }; @@ -350,7 +350,7 @@ where M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { - let metric = ::METRIC; + let metric = M::METRIC; let dimension_count = data.ncols(); let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; @@ -472,7 +472,7 @@ where ) .map_err(ANNError::new)?; - let metric = ::METRIC; + let metric = M::METRIC; let point_squared_norms = if metric == Metric::Cosine { grow_fallible(point_squared_norm_buffer, point_count, 0.0)?; let point_squared_norms = &mut point_squared_norm_buffer[..point_count]; From 170d27bb24098304292437c743a2835e4fbb3596 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:04:04 +0000 Subject: [PATCH 150/229] refactor(pipnn): prepare norms before kernel calls --- diskann/src/graph/pipnn/leaf_build.rs | 47 +++++++-- diskann/src/graph/pipnn/mod.rs | 45 +++++---- diskann/src/graph/pipnn/partitioning.rs | 122 +++++++++++++++--------- 3 files changed, 141 insertions(+), 73 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 04129b98c4..5bca5b6778 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -22,11 +22,12 @@ use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_vector::distance::Metric; use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::LeafKernelMetric, + kernel_metric::{LeafKernelMetric, norm_from_squared}, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, @@ -99,6 +100,7 @@ pub(crate) enum LeafBuildError { struct LeafBuffers { point_values: Vec, dots: Vec, + norms: Vec, neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, @@ -233,6 +235,7 @@ pub(super) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, requested_k: usize, + metric: Metric, ) -> Result>, LeafBuildError> where A: Architecture, @@ -256,6 +259,7 @@ where add_direct_leaf_candidates::( arch, data, + metric, leaf, point_ids, requested_k, @@ -272,9 +276,11 @@ where /// The function rejects empty, duplicate, unsorted, or out-of-range point IDs. /// Reusable buffers can be longer than this leaf, so all accesses use the current /// leaf shape. +#[allow(clippy::too_many_arguments)] fn add_direct_leaf_candidates( arch: A, data: MatrixView<'_, T>, + metric: Metric, leaf: usize, point_ids: &[u32], requested_k: usize, @@ -345,6 +351,21 @@ where leaf, buffer: "leaf dot-product matrix", })?; + let norm_count = if matches!(metric, Metric::L2 | Metric::Cosine) { + point_ids.len() + } else { + 0 + }; + grow("leaf norms", &mut buffers.norms, norm_count, 0.0)?; + for (point, norm) in buffers.norms[..norm_count].iter_mut().enumerate() { + let squared_norm = dots[(point, point)]; + *norm = if metric == Metric::Cosine { + norm_from_squared(squared_norm) + } else { + squared_norm + }; + } + let norms = &buffers.norms[..norm_count]; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), @@ -354,7 +375,7 @@ where leaf, buffer: "leaf output", })?; - nearest_neighbors::(arch, dots, output, &mut buffers.kernel_workspace) + nearest_neighbors::(arch, dots, norms, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; buffers.prepare_local_adjacency(point_ids.len())?; @@ -486,23 +507,33 @@ mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => { - build_leaf_candidates::(arch, call.data, call.leaves, call.k) - } - Metric::Cosine => { - build_leaf_candidates::(arch, call.data, call.leaves, call.k) - } + Metric::L2 => build_leaf_candidates::( + arch, + call.data, + call.leaves, + call.k, + Metric::L2, + ), + Metric::Cosine => build_leaf_candidates::( + arch, + call.data, + call.leaves, + call.k, + Metric::Cosine, + ), Metric::CosineNormalized => build_leaf_candidates::( arch, call.data, call.leaves, call.k, + Metric::CosineNormalized, ), Metric::InnerProduct => build_leaf_candidates::( arch, call.data, call.leaves, call.k, + Metric::InnerProduct, ), } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 46a4733eff..a96bab6d5f 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -231,14 +231,22 @@ where call: BuildGraphCall<'_, '_, '_, T>, ) -> ANNResult>> { match call.metric { - Metric::L2 => build_graph_for::(arch, call.data, call.context), - Metric::Cosine => build_graph_for::(arch, call.data, call.context), - Metric::CosineNormalized => { - build_graph_for::(arch, call.data, call.context) - } - Metric::InnerProduct => { - build_graph_for::(arch, call.data, call.context) + Metric::L2 => build_graph_for::(arch, call.data, call.context, Metric::L2), + Metric::Cosine => { + build_graph_for::(arch, call.data, call.context, Metric::Cosine) } + Metric::CosineNormalized => build_graph_for::( + arch, + call.data, + call.context, + Metric::CosineNormalized, + ), + Metric::InnerProduct => build_graph_for::( + arch, + call.data, + call.context, + Metric::InnerProduct, + ), } } } @@ -251,6 +259,7 @@ fn build_graph_for( arch: A, data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, + metric: Metric, ) -> ANNResult>> where A: Architecture, @@ -261,23 +270,23 @@ where T: VectorRepr + Send + Sync + 'static, { let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition::(arch, data, &context.config))?; + .in_scope(|| partitioning::partition::(arch, data, metric, &context.config))?; // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) - .map_err(ANNError::new) + leaf_build::build_leaf_candidates::( + arch, + data, + leaves, + context.config.leaf_k, + metric, + ) + .map_err(ANNError::new) })?; // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. - tracing::info_span!("pipnn.finalization").in_scope(|| { - finalization::prune_overfull( - data, - candidates, - context.graph, - ::METRIC, - ) - }) + tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) } fn effective_metric(metric: Metric) -> Metric { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 11660adafd..4254978205 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -31,7 +31,7 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::PartitionKernelMetric, + kernel_metric::{PartitionKernelMetric, norm_from_squared}, partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, }; @@ -87,10 +87,10 @@ struct WorkItem { #[derive(Default)] struct StripeBuffers { - points: Vec, - dots: Vec, - point_squared_norms: Vec, - kernel: PartitionKernelWorkspace, + point_values: Vec, + dot_values: Vec, + point_norms: Vec, + kernel_workspace: PartitionKernelWorkspace, } impl AsPooled<()> for StripeBuffers { @@ -119,6 +119,7 @@ type StripeBufferPool = ObjectPool; pub(super) fn partition( arch: A, data: MatrixView<'_, T>, + metric: Metric, config: &PiPNNConfig, ) -> ANNResult>> where @@ -145,7 +146,7 @@ where for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = - partition_replica::(arch, data, config, seed, &stripe_buffers)?; + partition_replica::(arch, data, metric, config, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::new)?; @@ -161,6 +162,7 @@ where fn partition_replica( arch: A, data: MatrixView<'_, T>, + metric: Metric, config: &PiPNNConfig, seed: u64, stripe_buffers: &StripeBufferPool, @@ -210,6 +212,7 @@ where *slot = Some(partition_work_item::( arch, data, + metric, config, item, stripe_buffers, @@ -251,6 +254,7 @@ where fn partition_work_item( arch: A, data: MatrixView<'_, T>, + metric: Metric, config: &PiPNNConfig, item: WorkItem, stripe_buffers: &StripeBufferPool, @@ -270,8 +274,15 @@ where config.p_samp, mix_seed(item.seed, points as u64), )?; - let clusters = - assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; + let clusters = assign_to_leaders::( + arch, + data, + metric, + &item.indices, + &leaders, + fanout, + stripe_buffers, + )?; let mut pending = Vec::new(); let mut finished = Vec::new(); @@ -337,6 +348,7 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, + metric: Metric, point_ids: &[u32], leader_ids: &[u32], fanout: usize, @@ -350,7 +362,6 @@ where M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { - let metric = M::METRIC; let dimension_count = data.ncols(); let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; @@ -369,7 +380,7 @@ where // SIMD reassociation changes low bits and can change a near-tie branch. *norm_value = leader_vector.iter().map(|value| value * value).sum(); if metric == Metric::Cosine { - *norm_value = norm_value.sqrt(); + *norm_value = norm_from_squared(*norm_value); } } @@ -401,6 +412,7 @@ where assign_point_stripe::( arch, data, + metric, &point_ids[first_point..first_point + stripe_point_count], &leader_values, &leader_norm_values, @@ -424,6 +436,7 @@ where fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, + metric: Metric, point_ids: &[u32], leader_values: &[f32], leader_norm_values: &[f32], @@ -447,17 +460,28 @@ where let output_len = checked_area("partition assignments", point_count, fanout)?; // Keep each buffer at its largest length. Every operation uses an explicit // active prefix. - grow_fallible(&mut buffers.points, point_values_len, 0.0)?; - grow_fallible(&mut buffers.dots, dots_len, 0.0)?; + grow_fallible(&mut buffers.point_values, point_values_len, 0.0)?; + grow_fallible(&mut buffers.dot_values, dots_len, 0.0)?; let StripeBuffers { - points: point_buffer, - dots: dot_buffer, - point_squared_norms: point_squared_norm_buffer, - kernel: kernel_workspace, + point_values: point_buffer, + dot_values: dot_buffer, + point_norms: point_norm_buffer, + kernel_workspace, } = buffers; - let point_values = &mut point_buffer[..point_values_len]; + let mut point_matrix = MutMatrixView::try_from( + &mut point_buffer[..point_values_len], + point_count, + dimensions, + ) + .map_err(|error| { + ANNError::new(PartitionError::InvalidBufferLength { + buffer: "point stripe", + expected: point_values_len, + actual: error.into_inner().len(), + }) + })?; let dots = &mut dot_buffer[..dots_len]; - gather_vectors(data, point_ids, point_values)?; + gather_vectors(data, point_ids, point_matrix.as_mut_slice())?; diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -465,36 +489,29 @@ where leader_count, dimensions, 1.0, - point_values, + point_matrix.as_slice(), leader_values, None, dots, ) .map_err(ANNError::new)?; - let metric = M::METRIC; - let point_squared_norms = if metric == Metric::Cosine { - grow_fallible(point_squared_norm_buffer, point_count, 0.0)?; - let point_squared_norms = &mut point_squared_norm_buffer[..point_count]; - for (norm_value, point_values) in point_squared_norms + let point_norms = if metric == Metric::Cosine { + grow_fallible(point_norm_buffer, point_count, 0.0)?; + let point_norms = &mut point_norm_buffer[..point_count]; + for (norm_value, point_values) in point_norms .iter_mut() - .zip(point_values.chunks_exact(dimensions)) + .zip(point_matrix.as_slice().chunks_exact(dimensions)) { - *norm_value = FastL2NormSquared.evaluate(point_values); + *norm_value = norm_from_squared(FastL2NormSquared.evaluate(point_values)); } - &*point_squared_norms + &*point_norms } else { &[] }; - 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, + let norms = PartitionNorms { + point_norms, + leader_norms: leader_norm_values, }; let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { ANNError::new(PartitionError::InvalidBufferLength { @@ -512,6 +529,7 @@ where })?; nearest_leaders::( arch, + metric, PartitionInput { dots, norms }, output, kernel_workspace, @@ -814,14 +832,22 @@ mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => partition::(arch, call.data, call.config), - Metric::Cosine => partition::(arch, call.data, call.config), - Metric::CosineNormalized => { - partition::(arch, call.data, call.config) - } - Metric::InnerProduct => { - partition::(arch, call.data, call.config) + Metric::L2 => partition::(arch, call.data, Metric::L2, call.config), + Metric::Cosine => { + partition::(arch, call.data, Metric::Cosine, call.config) } + Metric::CosineNormalized => partition::( + arch, + call.data, + Metric::CosineNormalized, + call.config, + ), + Metric::InnerProduct => partition::( + arch, + call.data, + Metric::InnerProduct, + call.config, + ), } } } @@ -1102,6 +1128,7 @@ mod tests { let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, + Metric::L2, &[2], &[0, 1], 1, @@ -1156,13 +1183,13 @@ mod tests { let pool = StripeBufferPool::new((), 0, None); let points = { let mut buffers = pool.get_ref(()); - buffers.points.resize(16, 0.0); - buffers.points.as_ptr() + buffers.point_values.resize(16, 0.0); + buffers.point_values.as_ptr() }; let buffers = pool.get_ref(()); - assert_eq!(buffers.points.as_ptr(), points); - assert_eq!(buffers.points.len(), 16); + assert_eq!(buffers.point_values.as_ptr(), points); + assert_eq!(buffers.point_values.len(), 16); } #[test] @@ -1175,6 +1202,7 @@ mod tests { let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, + Metric::L2, &point_ids, &[0, 2_047], 1, From eee911af09a809f3d6212754b6840b9b778b2d88 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:42:04 +0000 Subject: [PATCH 151/229] refactor(pipnn): dispatch norm preparation through metric types --- diskann/src/graph/pipnn/leaf_build.rs | 50 +++--------- diskann/src/graph/pipnn/mod.rs | 16 ++-- diskann/src/graph/pipnn/partitioning.rs | 101 +++++++----------------- 3 files changed, 47 insertions(+), 120 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 5bca5b6778..273ca54f2c 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -22,12 +22,11 @@ use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_vector::distance::Metric; use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::{LeafKernelMetric, norm_from_squared}, + kernel_metric::LeafMetric, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, @@ -235,14 +234,13 @@ pub(super) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, requested_k: usize, - metric: Metric, ) -> Result>, LeafBuildError> where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: LeafKernelMetric, + M: LeafMetric, T: VectorRepr + 'static, { if data.ncols() == 0 { @@ -259,7 +257,6 @@ where add_direct_leaf_candidates::( arch, data, - metric, leaf, point_ids, requested_k, @@ -280,7 +277,6 @@ where fn add_direct_leaf_candidates( arch: A, data: MatrixView<'_, T>, - metric: Metric, leaf: usize, point_ids: &[u32], requested_k: usize, @@ -292,7 +288,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: LeafKernelMetric, + M: LeafMetric, T: VectorRepr + 'static, { if point_ids.is_empty() { @@ -351,21 +347,9 @@ where leaf, buffer: "leaf dot-product matrix", })?; - let norm_count = if matches!(metric, Metric::L2 | Metric::Cosine) { - point_ids.len() - } else { - 0 - }; - grow("leaf norms", &mut buffers.norms, norm_count, 0.0)?; - for (point, norm) in buffers.norms[..norm_count].iter_mut().enumerate() { - let squared_norm = dots[(point, point)]; - *norm = if metric == Metric::Cosine { - norm_from_squared(squared_norm) - } else { - squared_norm - }; - } - let norms = &buffers.norms[..norm_count]; + M::prepare_norms(dots, &mut buffers.norms) + .map_err(|source| allocation_error("leaf norms", point_ids.len(), source))?; + let norms = &*buffers.norms; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), @@ -507,33 +491,23 @@ mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => build_leaf_candidates::( - arch, - call.data, - call.leaves, - call.k, - Metric::L2, - ), - Metric::Cosine => build_leaf_candidates::( - arch, - call.data, - call.leaves, - call.k, - Metric::Cosine, - ), + Metric::L2 => { + build_leaf_candidates::(arch, call.data, call.leaves, call.k) + } + Metric::Cosine => { + build_leaf_candidates::(arch, call.data, call.leaves, call.k) + } Metric::CosineNormalized => build_leaf_candidates::( arch, call.data, call.leaves, call.k, - Metric::CosineNormalized, ), Metric::InnerProduct => build_leaf_candidates::( arch, call.data, call.leaves, call.k, - Metric::InnerProduct, ), } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index a96bab6d5f..e831741f6b 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -56,7 +56,7 @@ use diskann_wide::{ use rayon::ThreadPool; use self::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, LeafKernelMetric, PartitionKernelMetric, + Cosine, CosineNormalized, InnerProduct, L2, LeafMetric, PartitionMetric, }; /// PiPNN partition and leaf-selection policy. @@ -266,22 +266,16 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: LeafKernelMetric + PartitionKernelMetric, + M: LeafMetric + PartitionMetric, T: VectorRepr + Send + Sync + 'static, { let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition::(arch, data, metric, &context.config))?; + .in_scope(|| partitioning::partition::(arch, data, &context.config))?; // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates::( - arch, - data, - leaves, - context.config.leaf_k, - metric, - ) - .map_err(ANNError::new) + leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) + .map_err(ANNError::new) })?; // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 4254978205..cdd8849919 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -24,14 +24,13 @@ use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, }; -use diskann_vector::{Norm, distance::Metric, norm::FastL2NormSquared}; use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::{PartitionKernelMetric, norm_from_squared}, + kernel_metric::PartitionMetric, partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, }; @@ -119,7 +118,6 @@ type StripeBufferPool = ObjectPool; pub(super) fn partition( arch: A, data: MatrixView<'_, T>, - metric: Metric, config: &PiPNNConfig, ) -> ANNResult>> where @@ -127,7 +125,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr + Send + Sync, { let points = data.nrows(); @@ -146,7 +144,7 @@ where for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = - partition_replica::(arch, data, metric, config, seed, &stripe_buffers)?; + partition_replica::(arch, data, config, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::new)?; @@ -162,7 +160,6 @@ where fn partition_replica( arch: A, data: MatrixView<'_, T>, - metric: Metric, config: &PiPNNConfig, seed: u64, stripe_buffers: &StripeBufferPool, @@ -172,7 +169,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr + Send + Sync, { let initial_indices = point_ids(data.nrows())?; @@ -212,7 +209,6 @@ where *slot = Some(partition_work_item::( arch, data, - metric, config, item, stripe_buffers, @@ -254,7 +250,6 @@ where fn partition_work_item( arch: A, data: MatrixView<'_, T>, - metric: Metric, config: &PiPNNConfig, item: WorkItem, stripe_buffers: &StripeBufferPool, @@ -264,7 +259,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr + Send + Sync, { let points = item.indices.len(); @@ -274,15 +269,8 @@ where config.p_samp, mix_seed(item.seed, points as u64), )?; - let clusters = assign_to_leaders::( - arch, - data, - metric, - &item.indices, - &leaders, - fanout, - stripe_buffers, - )?; + let clusters = + assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; let mut pending = Vec::new(); let mut finished = Vec::new(); @@ -348,7 +336,6 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, - metric: Metric, point_ids: &[u32], leader_ids: &[u32], fanout: usize, @@ -359,7 +346,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr + Send + Sync, { let dimension_count = data.ncols(); @@ -367,22 +354,18 @@ where let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; gather_vectors(data, leader_ids, &mut leader_values)?; - let mut leader_norm_values = if matches!(metric, Metric::L2 | Metric::Cosine) { - filled_vec(leader_ids.len(), 0.0f32)? - } else { - Vec::new() - }; - for (norm_value, leader_vector) in leader_norm_values - .iter_mut() - .zip(leader_values.chunks_exact(dimension_count)) - { - // Leader norms affect top-k order. Use this scalar reduction order. - // SIMD reassociation changes low bits and can change a near-tie branch. - *norm_value = leader_vector.iter().map(|value| value * value).sum(); - if metric == Metric::Cosine { - *norm_value = norm_from_squared(*norm_value); - } - } + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count).map_err( + |error| { + ANNError::new(PartitionError::InvalidBufferLength { + buffer: "leader data", + expected: leader_values_len, + actual: error.into_inner().len(), + }) + }, + )?; + let mut leader_norm_values = Vec::new(); + M::prepare_leader_norms(leader_matrix, &mut leader_norm_values).map_err(ANNError::new)?; let fanout = fanout.min(leader_ids.len()); let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; @@ -412,7 +395,6 @@ where assign_point_stripe::( arch, data, - metric, &point_ids[first_point..first_point + stripe_point_count], &leader_values, &leader_norm_values, @@ -436,7 +418,6 @@ where fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, - metric: Metric, point_ids: &[u32], leader_values: &[f32], leader_norm_values: &[f32], @@ -449,7 +430,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr, { let point_count = point_ids.len(); @@ -496,19 +477,8 @@ where ) .map_err(ANNError::new)?; - let point_norms = if metric == Metric::Cosine { - grow_fallible(point_norm_buffer, point_count, 0.0)?; - let point_norms = &mut point_norm_buffer[..point_count]; - for (norm_value, point_values) in point_norms - .iter_mut() - .zip(point_matrix.as_slice().chunks_exact(dimensions)) - { - *norm_value = norm_from_squared(FastL2NormSquared.evaluate(point_values)); - } - &*point_norms - } else { - &[] - }; + M::prepare_point_norms(point_matrix.as_view(), point_norm_buffer).map_err(ANNError::new)?; + let point_norms = &*point_norm_buffer; let norms = PartitionNorms { point_norms, leader_norms: leader_norm_values, @@ -529,7 +499,6 @@ where })?; nearest_leaders::( arch, - metric, PartitionInput { dots, norms }, output, kernel_workspace, @@ -832,22 +801,14 @@ mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => partition::(arch, call.data, Metric::L2, call.config), - Metric::Cosine => { - partition::(arch, call.data, Metric::Cosine, call.config) + Metric::L2 => partition::(arch, call.data, call.config), + Metric::Cosine => partition::(arch, call.data, call.config), + Metric::CosineNormalized => { + partition::(arch, call.data, call.config) + } + Metric::InnerProduct => { + partition::(arch, call.data, call.config) } - Metric::CosineNormalized => partition::( - arch, - call.data, - Metric::CosineNormalized, - call.config, - ), - Metric::InnerProduct => partition::( - arch, - call.data, - Metric::InnerProduct, - call.config, - ), } } } @@ -1128,7 +1089,6 @@ mod tests { let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, - Metric::L2, &[2], &[0, 1], 1, @@ -1202,7 +1162,6 @@ mod tests { let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, - Metric::L2, &point_ids, &[0, 2_047], 1, From 3a7663f45ac7316f84698ba8f07371b4b4aba5f9 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:22:00 +0000 Subject: [PATCH 152/229] refactor(pipnn): pass metric-owned ranking context --- diskann/src/graph/pipnn/leaf_build.rs | 9 ++++--- diskann/src/graph/pipnn/partitioning.rs | 33 ++++++++++++++----------- 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 273ca54f2c..603ef78ae7 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -26,7 +26,7 @@ use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::LeafMetric, + kernel_metric::{LeafMetric, NormPreparation}, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, @@ -347,8 +347,11 @@ where leaf, buffer: "leaf dot-product matrix", })?; - M::prepare_norms(dots, &mut buffers.norms) - .map_err(|source| allocation_error("leaf norms", point_ids.len(), source))?; + M::prepare_norms(NormPreparation { + values: dots, + norms: &mut buffers.norms, + }) + .map_err(|source| allocation_error("leaf norms", point_ids.len(), source))?; let norms = &*buffers.norms; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index cdd8849919..052ec36608 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -30,8 +30,8 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::PartitionMetric, - partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, + kernel_metric::{NormPreparation, PartitionMetric, PartitionNorms}, + partition_kernel::{PartitionInput, nearest_leaders}, }; // These constants control internal batching and deterministic seed generation. @@ -89,7 +89,7 @@ struct StripeBuffers { point_values: Vec, dot_values: Vec, point_norms: Vec, - kernel_workspace: PartitionKernelWorkspace, + ranked_leaders: Vec<(u32, f32)>, } impl AsPooled<()> for StripeBuffers { @@ -365,7 +365,11 @@ where }, )?; let mut leader_norm_values = Vec::new(); - M::prepare_leader_norms(leader_matrix, &mut leader_norm_values).map_err(ANNError::new)?; + M::prepare_leader_norms(NormPreparation { + values: leader_matrix, + norms: &mut leader_norm_values, + }) + .map_err(ANNError::new)?; let fanout = fanout.min(leader_ids.len()); let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; @@ -447,7 +451,7 @@ where point_values: point_buffer, dot_values: dot_buffer, point_norms: point_norm_buffer, - kernel_workspace, + ranked_leaders, } = buffers; let mut point_matrix = MutMatrixView::try_from( &mut point_buffer[..point_values_len], @@ -477,7 +481,11 @@ where ) .map_err(ANNError::new)?; - M::prepare_point_norms(point_matrix.as_view(), point_norm_buffer).map_err(ANNError::new)?; + M::prepare_point_norms(NormPreparation { + values: point_matrix.as_view(), + norms: point_norm_buffer, + }) + .map_err(ANNError::new)?; let point_norms = &*point_norm_buffer; let norms = PartitionNorms { point_norms, @@ -497,13 +505,8 @@ where actual: error.into_inner().len(), }) })?; - nearest_leaders::( - arch, - PartitionInput { dots, norms }, - output, - kernel_workspace, - ) - .map_err(ANNError::new) + nearest_leaders::(arch, PartitionInput { dots, norms }, output, ranked_leaders) + .map_err(ANNError::new) } fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> @@ -1063,7 +1066,7 @@ mod tests { } #[test] - fn l2_leader_norms_preserve_scalar_reduction_order() { + fn l2_leader_norms_preserve_sequential_reduction_order() { fn next(state: &mut u64) -> f32 { *state ^= *state << 13; *state ^= *state >> 7; @@ -1071,7 +1074,7 @@ mod tests { (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 } - // Scalar and SIMD-reassociated leader norms select different top-1 + // Sequential and SIMD-reassociated leader norms select different top-1 // leaders for this case. Dot products still use the production GEMM. The // test changes only the leader-norm reduction. let dimensions = 129; From 3cd4b2e5941a2ded7e2e2016198ee13b32802929 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:07:09 +0000 Subject: [PATCH 153/229] fix(pipnn): reset reused prune states --- diskann/src/graph/pipnn/finalization.rs | 38 +++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 55d86103e6..62a3bb9fbc 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -138,6 +138,9 @@ where workspace .prune_states .resize(workspace.candidate_slots.len(), prune::State::default()); + // Each candidate list starts a separate RobustPrune state machine. + // Reset retained entries because resize initializes only new entries. + workspace.prune_states.fill(prune::State::default()); let selected = prune::robust_prune( &sorted, @@ -247,6 +250,41 @@ mod tests { assert_eq!(&*actual[0], &[1, 3]); } + #[test] + fn reused_workspace_matches_fresh_pruning() { + let data = [0.0_f32, 1.0, 2.0, -3.0, 4.0]; + let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); + let first = [3, 2, 1]; + let second = [4, 3, 2]; + let candidates = |first: &[u32], second: &[u32]| { + vec![ + candidate_list(first.iter().copied()), + candidate_list(second.iter().copied()), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ] + }; + let graph = graph_config(2); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + let fresh_first = pool + .install(|| prune_overfull(data, candidates(&first, &[]), &graph, Metric::L2)) + .unwrap(); + let fresh_second = pool + .install(|| prune_overfull(data, candidates(&[], &second), &graph, Metric::L2)) + .unwrap(); + + let reused = pool + .install(|| prune_overfull(data, candidates(&first, &second), &graph, Metric::L2)) + .unwrap(); + + assert_eq!(&*reused[0], &*fresh_first[0]); + assert_eq!(&*reused[1], &*fresh_second[1]); + } + #[test] fn rejects_invalid_candidate_ids_without_panicking() { let data = [0.0_f32, 1.0, 2.0]; From e829fec4fa066d84f293969dec43a0e762723cbf Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:05:12 +0000 Subject: [PATCH 154/229] refactor(pipnn): prepare leaf norms in metric policy --- diskann/src/graph/pipnn/leaf_build.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 603ef78ae7..644e05a1c2 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -347,7 +347,7 @@ where leaf, buffer: "leaf dot-product matrix", })?; - M::prepare_norms(NormPreparation { + M::prepare_leaf_norms(NormPreparation { values: dots, norms: &mut buffers.norms, }) From a345bd750e9e20bb68c790dd69dba907c97943aa Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:20:46 +0000 Subject: [PATCH 155/229] refactor(pipnn): simplify graph build stages Delegate GEMM, norm preparation, and ranking to numerical kernels. Keep partition and leaf stages focused on IDs, scheduling, and graph mapping. --- Cargo.lock | 1 + diskann/Cargo.toml | 4 +- diskann/src/graph/pipnn/finalization.rs | 125 +------ diskann/src/graph/pipnn/leaf_build.rs | 373 +++---------------- diskann/src/graph/pipnn/mod.rs | 22 +- diskann/src/graph/pipnn/partitioning.rs | 458 +++++------------------- 6 files changed, 162 insertions(+), 821 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c8a45250fe..5375675666 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -446,6 +446,7 @@ dependencies = [ "half", "hashbrown 0.16.1", "num-traits", + "parking_lot", "pin-project", "rand", "rayon", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 2568ac6669..d8dcf211b8 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -22,6 +22,7 @@ half = { workspace = true, features = ["bytemuck", "num-traits"] } # while other crates use default-features = true. Keeping version 0.16.0 consistent. hashbrown = { version = "0.16.0", default-features = false, features = ["default-hasher"] } num-traits.workspace = true +parking_lot = { version = "0.12.5", optional = true } rand.workspace = true rayon = { workspace = true, optional = true } thiserror.workspace = true @@ -32,7 +33,6 @@ 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 } @@ -61,7 +61,7 @@ panic = "warn" default = ["tracing"] # Enable PiPNN batch graph construction. -pipnn = ["dep:diskann-linalg", "dep:rayon", "tracing"] +pipnn = ["dep:diskann-linalg", "dep:parking_lot", "dep:rayon", "tracing"] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 62a3bb9fbc..6f01c3dd25 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -31,32 +31,24 @@ use rayon::prelude::*; #[derive(Debug, thiserror::Error)] pub(crate) enum FinalizationError { - #[error("candidate list count {lists} does not match the dataset point count {points}")] - CandidateListCountMismatch { lists: usize, points: usize }, - #[error( - "candidate ID {candidate} for source {source_index} is outside a {points}-point dataset" - )] - InvalidCandidateId { - source_index: usize, - candidate: u32, - points: usize, - }, #[error("candidate count {actual} exceeds the u16 position limit {max}")] TooManyCandidates { actual: usize, max: usize }, } /// RobustPrune state for one Rayon job. /// -/// `candidate_slots` and `prune_states` stay positionally aligned with +/// `sorted_cache` and `prune_states` stay positionally aligned with /// `sorted_candidates`. #[derive(Default)] struct PruneWorkspace { sorted_candidates: Vec>, - candidate_slots: Vec<(f32, Option)>, + sorted_cache: Vec<(f32, Option)>, prune_states: Vec, } -/// Check candidate IDs and prune each list that exceeds the graph degree. +/// Prune each candidate list that exceeds the graph degree. +/// +/// Candidate builders supply one list per data row and valid dataset IDs. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, @@ -66,8 +58,6 @@ pub(crate) fn prune_overfull( where T: VectorRepr + Send + Sync, { - validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::new)?; - let degree = graph.pruned_degree().get(); let distance = T::distance(metric, Some(data.ncols())); @@ -88,10 +78,6 @@ where let source_id = u32::try_from(source).map_err(ANNError::new)?; let source_vector = data.row(source); workspace.sorted_candidates.clear(); - workspace - .sorted_candidates - .try_reserve(source_candidates.len()) - .map_err(ANNError::new)?; workspace .sorted_candidates .extend(source_candidates.iter().copied().map(|candidate| { @@ -109,42 +95,25 @@ where max: u16::MAX as usize, })); } - workspace.candidate_slots.clear(); - workspace - .candidate_slots - .try_reserve(candidate_count) - .map_err(ANNError::new)?; - + workspace.sorted_cache.clear(); // Sort all candidates before the code marks a self-edge as absent. - // Thus, self-edge removal cannot add a farther candidate. The - // `SortedNeighbors` value carries this order into RobustPrune. + // Thus, self-edge removal cannot add a farther candidate. Cache + // construction preserves this order for RobustPrune. let sorted = SortedNeighbors::new(&mut workspace.sorted_candidates, candidate_count); - workspace - .candidate_slots - .extend(sorted.iter().map(|neighbor| { - let id = *neighbor.id(); - (*neighbor.distance(), (id != source_id).then_some(id)) - })); + workspace.sorted_cache.extend(sorted.iter().map(|neighbor| { + let id = *neighbor.id(); + (*neighbor.distance(), (id != source_id).then_some(id)) + })); workspace .prune_states - .try_reserve( - workspace - .candidate_slots - .len() - .saturating_sub(workspace.prune_states.len()), - ) - .map_err(ANNError::new)?; - workspace - .prune_states - .resize(workspace.candidate_slots.len(), prune::State::default()); + .resize(workspace.sorted_cache.len(), prune::State::default()); // Each candidate list starts a separate RobustPrune state machine. // Reset retained entries because resize initializes only new entries. workspace.prune_states.fill(prune::State::default()); let selected = prune::robust_prune( - &sorted, - &workspace.candidate_slots, + &workspace.sorted_cache, workspace.prune_states.as_mut_slice(), degree, graph.alpha(), @@ -168,28 +137,6 @@ where .collect() } -fn validate_candidate_lists( - candidates: &[AdjacencyList], - points: usize, -) -> Result<(), FinalizationError> { - if candidates.len() != points { - return Err(FinalizationError::CandidateListCountMismatch { - lists: candidates.len(), - points, - }); - } - for (source, source_candidates) in candidates.iter().enumerate() { - if let Some(&candidate) = source_candidates.iter().find(|&&id| id as usize >= points) { - return Err(FinalizationError::InvalidCandidateId { - source_index: source, - candidate, - points, - }); - } - } - Ok(()) -} - #[cfg(test)] mod tests { use crate::graph::{ @@ -284,48 +231,4 @@ mod tests { assert_eq!(&*reused[0], &*fresh_first[0]); assert_eq!(&*reused[1], &*fresh_second[1]); } - - #[test] - fn rejects_invalid_candidate_ids_without_panicking() { - let data = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![ - candidate_list([1, 3]), - candidate_list([]), - candidate_list([]), - ]; - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::InvalidCandidateId { - source_index: 0, - candidate: 3, - points: 3, - }) - )); - } - - #[test] - fn rejects_candidate_list_count_mismatch_without_panicking() { - let data = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![ - candidate_list([]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::CandidateListCountMismatch { - lists: 4, - points: 3 - }) - )); - } } diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 644e05a1c2..5f02967125 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -8,17 +8,16 @@ //! Partitioning supplies sorted, unique global point IDs for each leaf. One leaf //! job does these steps: //! -//! 1. Check each ID and convert its vector to reusable `f32` storage. -//! 2. Compute the lower triangle of `A · Aᵀ`. -//! 3. Select local neighbors for both points of each pair. -//! 4. Convert local positions to global point IDs. -//! 5. Add both edge directions to global candidate lists. +//! 1. Gather each ID and convert its vector to reusable `f32` storage. +//! 2. Call the leaf kernel for Gram construction, norms, and local ranking. +//! 3. Convert local positions to global point IDs. +//! 4. Add both edge directions to global candidate lists. //! //! Overlapping leaves run concurrently. A worker locks one destination list only //! while it adds one leaf's IDs. Reusable buffers keep their largest allocation. //! Each operation uses an explicit active prefix. -use std::{collections::TryReserveError, sync::Mutex}; +use parking_lot::Mutex; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -26,31 +25,13 @@ use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::{LeafMetric, NormPreparation}, - leaf_kernel::{ - LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, - }, + kernel_metric::LeafMetric, + leaf_kernel::{LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, select_leaf_neighbors}, }; /// Failure while converting leaves into direct graph candidates. #[derive(Debug, thiserror::Error)] pub(crate) enum LeafBuildError { - #[error("leaf build requires at least one dimension")] - EmptyDimensions, - #[error("dataset point count {0} exceeds the u32 ID limit")] - TooManyPoints(usize), - #[error("leaf {leaf} is empty")] - EmptyLeaf { leaf: usize }, - #[error("point ID {point} in leaf {leaf} is outside a {points}-point dataset")] - InvalidPointId { - leaf: usize, - point: u32, - points: usize, - }, - #[error("point ID {point} appears more than once in leaf {leaf}")] - DuplicatePointId { leaf: usize, point: u32 }, - #[error("point IDs in leaf {leaf} are not strictly increasing")] - UnsortedPointIds { leaf: usize }, #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] ShapeOverflow { leaf: usize, @@ -59,13 +40,6 @@ pub(crate) enum LeafBuildError { }, #[error("failed to form {buffer} view for leaf {leaf}")] InvalidView { leaf: usize, buffer: &'static str }, - #[error("failed to reserve {additional} values for {buffer}")] - Allocation { - buffer: &'static str, - additional: usize, - #[source] - source: TryReserveError, - }, #[error("failed to convert point {point} in leaf {leaf}")] Conversion { leaf: usize, @@ -73,22 +47,12 @@ pub(crate) enum LeafBuildError { #[source] source: crate::ANNError, }, - #[error("lower-AAT failed for leaf {leaf}")] - LowerAat { - leaf: usize, - #[source] - source: diskann_linalg::SgemmError, - }, #[error("nearest-neighbor selection failed for leaf {leaf}")] Kernel { leaf: usize, #[source] - source: LeafKernelError, + source: crate::ANNError, }, - #[error("leaf kernel returned local target {target} for a {points}-point leaf")] - InvalidLocalTarget { target: u32, points: usize }, - #[error("candidate list for point {point} is poisoned")] - PoisonedCandidateList { point: u32 }, } /// Reusable buffers for one Rayon leaf job. @@ -98,8 +62,6 @@ pub(crate) enum LeafBuildError { #[derive(Default)] struct LeafBuffers { point_values: Vec, - dots: Vec, - norms: Vec, neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, @@ -121,16 +83,7 @@ impl LeafBuffers { rows: point_count, columns: dimension_count, })?; - let dot_count = - point_count - .checked_mul(point_count) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: point_count, - columns: point_count, - })?; - let leaf_k = leaf_neighbor_count(point_count, requested_k) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + let leaf_k = leaf_neighbor_count(point_count, requested_k); let neighbor_count = point_count .checked_mul(leaf_k) @@ -140,33 +93,17 @@ impl LeafBuffers { columns: leaf_k, })?; - grow( - "leaf point values", - &mut self.point_values, - point_value_count, - 0.0, - )?; - grow("leaf dot products", &mut self.dots, dot_count, 0.0)?; - grow( - "leaf neighbors", - &mut self.neighbors, - neighbor_count, - LeafNeighbor::default(), - )?; + grow(&mut self.point_values, point_value_count, 0.0); + grow(&mut self.neighbors, neighbor_count, LeafNeighbor::default()); Ok((leaf_k, neighbor_count)) } - fn prepare_local_adjacency(&mut self, point_count: usize) -> Result<(), LeafBuildError> { - let additional = point_count.saturating_sub(self.local_adjacency.len()); - self.local_adjacency - .try_reserve(additional) - .map_err(|source| allocation_error("leaf adjacency lists", additional, source))?; + fn prepare_local_adjacency(&mut self, point_count: usize) { self.local_adjacency .resize_with(point_count, AdjacencyList::new); self.local_adjacency[..point_count] .iter_mut() .for_each(AdjacencyList::clear); - Ok(()) } } @@ -180,47 +117,31 @@ struct DirectCandidates { } impl DirectCandidates { - fn new(point_count: usize) -> Result { - let mut lists = Vec::new(); - lists - .try_reserve_exact(point_count) - .map_err(|source| allocation_error("candidate lists", point_count, source))?; - lists.resize_with(point_count, || Mutex::new(AdjacencyList::new())); - Ok(Self { lists }) + fn new(point_count: usize) -> Self { + let lists = (0..point_count) + .map(|_| Mutex::new(AdjacencyList::new())) + .collect(); + Self { lists } } - fn add_leaf( - &self, - point_ids: &[u32], - local_adjacency: &[AdjacencyList], - ) -> Result<(), LeafBuildError> { + fn add_leaf(&self, point_ids: &[u32], local_adjacency: &[AdjacencyList]) { for (&source, additions) in point_ids.iter().zip(local_adjacency) { // `add_direct_leaf_candidates` checks every point ID before this append. - let candidates = &self.lists[source as usize]; - let mut candidates = candidates + self.lists[source as usize] .lock() - .map_err(|_| LeafBuildError::PoisonedCandidateList { point: source })?; - candidates.extend_from_slice(additions); + .extend_from_slice(additions); } - Ok(()) } - fn into_lists(self) -> Result>, LeafBuildError> { - let mut output = Vec::new(); - output - .try_reserve_exact(self.lists.len()) - .map_err(|source| allocation_error("candidate output", self.lists.len(), source))?; - for (point, candidates) in self.lists.into_iter().enumerate() { - let mut candidates = + fn into_lists(self) -> Vec> { + self.lists + .into_iter() + .map(Mutex::into_inner) + .map(|mut candidates| { + candidates.sort(); candidates - .into_inner() - .map_err(|_| LeafBuildError::PoisonedCandidateList { - point: point as u32, - })?; - candidates.sort(); - output.push(candidates); - } - Ok(output) + }) + .collect() } } @@ -243,14 +164,7 @@ where M: LeafMetric, T: VectorRepr + 'static, { - if data.ncols() == 0 { - return Err(LeafBuildError::EmptyDimensions); - } - if data.nrows() > u32::MAX as usize { - return Err(LeafBuildError::TooManyPoints(data.nrows())); - } - - let candidates = DirectCandidates::new(data.nrows())?; + let candidates = DirectCandidates::new(data.nrows()); leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { @@ -265,12 +179,11 @@ where ) }, )?; - candidates.into_lists() + Ok(candidates.into_lists()) } /// Add one leaf's symmetric neighbors to the direct candidate lists. /// -/// The function rejects empty, duplicate, unsorted, or out-of-range point IDs. /// Reusable buffers can be longer than this leaf, so all accesses use the current /// leaf shape. #[allow(clippy::too_many_arguments)] @@ -291,27 +204,6 @@ where M: LeafMetric, T: VectorRepr + 'static, { - if point_ids.is_empty() { - return Err(LeafBuildError::EmptyLeaf { leaf }); - } - for &point in point_ids { - if point as usize >= data.nrows() { - return Err(LeafBuildError::InvalidPointId { - leaf, - point, - points: data.nrows(), - }); - } - } - if let Some(pair) = point_ids.windows(2).find(|pair| pair[0] >= pair[1]) { - if pair[0] == pair[1] { - return Err(LeafBuildError::DuplicatePointId { - leaf, - point: pair[0], - }); - } - return Err(LeafBuildError::UnsortedPointIds { leaf }); - } let (leaf_k, neighbor_value_count) = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; if leaf_k == 0 { @@ -319,11 +211,11 @@ where } let point_value_count = point_ids.len() * data.ncols(); - let dot_count = point_ids.len() * point_ids.len(); + let point_values = &mut buffers.point_values[..point_value_count]; for (&point, point_output) in point_ids .iter() - .zip(buffers.point_values[..point_value_count].chunks_exact_mut(data.ncols())) + .zip(point_values.chunks_exact_mut(data.ncols())) { let source_values = data.row(point as usize); T::as_f32_into(source_values, point_output).map_err(|source| { @@ -335,24 +227,13 @@ where })?; } - diskann_linalg::sgemm_aat_lower( - point_ids.len(), - data.ncols(), - &buffers.point_values[..point_value_count], - &mut buffers.dots[..dot_count], - ) - .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; - let dots = MatrixView::try_from(&buffers.dots[..dot_count], point_ids.len(), point_ids.len()) - .map_err(|_| LeafBuildError::InvalidView { - leaf, - buffer: "leaf dot-product matrix", - })?; - M::prepare_leaf_norms(NormPreparation { - values: dots, - norms: &mut buffers.norms, - }) - .map_err(|source| allocation_error("leaf norms", point_ids.len(), source))?; - let norms = &*buffers.norms; + let points = + MatrixView::try_from(&*point_values, point_ids.len(), data.ncols()).map_err(|_| { + LeafBuildError::InvalidView { + leaf, + buffer: "leaf point matrix", + } + })?; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), @@ -362,17 +243,18 @@ where leaf, buffer: "leaf output", })?; - nearest_neighbors::(arch, dots, norms, output, &mut buffers.kernel_workspace) + select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - buffers.prepare_local_adjacency(point_ids.len())?; + buffers.prepare_local_adjacency(point_ids.len()); add_symmetric_neighbors( point_ids, leaf_k, &buffers.neighbors[..neighbor_value_count], &mut buffers.local_adjacency[..point_ids.len()], - )?; - candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]) + ); + candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]); + Ok(()) } fn add_symmetric_neighbors( @@ -380,61 +262,23 @@ fn add_symmetric_neighbors( leaf_k: usize, neighbors: &[LeafNeighbor], local_adjacency: &mut [AdjacencyList], -) -> Result<(), LeafBuildError> { +) { for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { for neighbor in source_neighbors { let target = neighbor.target as usize; - let Some(&target_id) = point_ids.get(target) else { - return Err(LeafBuildError::InvalidLocalTarget { - target: neighbor.target, - points: point_ids.len(), - }); - }; let source_id = point_ids[source]; + let target_id = point_ids[target]; if source_id != target_id { local_adjacency[source].push(target_id); local_adjacency[target].push(source_id); } } } - Ok(()) } -fn grow( - buffer: &'static str, - values: &mut Vec, - len: usize, - value: T, -) -> Result<(), LeafBuildError> { +fn grow(values: &mut Vec, len: usize, value: T) { if values.len() < len { - resize(buffer, values, len, value)?; - } - Ok(()) -} - -fn resize( - buffer: &'static str, - values: &mut Vec, - len: usize, - value: T, -) -> Result<(), LeafBuildError> { - let additional = len.saturating_sub(values.len()); - values - .try_reserve(additional) - .map_err(|source| allocation_error(buffer, additional, source))?; - values.resize(len, value); - Ok(()) -} - -fn allocation_error( - buffer: &'static str, - additional: usize, - source: TryReserveError, -) -> LeafBuildError { - LeafBuildError::Allocation { - buffer, - additional, - source, + values.resize(len, value); } } @@ -450,7 +294,7 @@ mod tests { use std::collections::BTreeSet; use super::{ - DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, + DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, build_leaf_candidates, }; @@ -724,44 +568,6 @@ mod tests { } } - #[test] - fn rejects_invalid_dimensions_and_leaf_membership() { - let data = [0.0_f32, 1.0]; - let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); - assert!(matches!( - build(no_dimensions, &[], 1, Metric::L2), - Err(LeafBuildError::EmptyDimensions) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), - Err(LeafBuildError::EmptyLeaf { leaf: 0 }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), - Err(LeafBuildError::InvalidPointId { - leaf: 0, - point: 2, - points: 2 - }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), - Err(LeafBuildError::InvalidPointId { point: 2, .. }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), - Err(LeafBuildError::InvalidPointId { point: 2, .. }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), - Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![1, 0]], 1, Metric::L2), - Err(LeafBuildError::UnsortedPointIds { leaf: 0 }) - )); - } - #[test] fn singleton_and_zero_k_leaves_add_no_candidates() { let data = [0.0_f32, 1.0, 2.0]; @@ -786,16 +592,13 @@ mod tests { let mut buffers = LeafBuffers::default(); buffers.prepare(0, 64, 128, 2).unwrap(); let point_values = buffers.point_values.as_ptr(); - let dots = buffers.dots.as_ptr(); let neighbors = buffers.neighbors.as_ptr(); buffers.prepare(1, 8, 128, 2).unwrap(); assert_eq!(buffers.point_values.as_ptr(), point_values); - assert_eq!(buffers.dots.as_ptr(), dots); assert_eq!(buffers.neighbors.as_ptr(), neighbors); assert_eq!(buffers.point_values.len(), 64 * 128); - assert_eq!(buffers.dots.len(), 64 * 64); assert_eq!(buffers.neighbors.len(), 64 * 2); } @@ -808,28 +611,6 @@ mod tests { )); } - #[test] - fn rejects_an_invalid_kernel_target() { - let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; - let error = add_symmetric_neighbors( - &[10, 20], - 1, - &[ - super::super::leaf_kernel::LeafNeighbor::new(9, 1.0), - super::super::leaf_kernel::LeafNeighbor::new(0, 1.0), - ], - &mut graph, - ) - .unwrap_err(); - assert!(matches!( - error, - LeafBuildError::InvalidLocalTarget { - target: 9, - points: 2 - } - )); - } - #[test] fn skips_duplicate_global_ids_without_self_edges() { let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; @@ -841,58 +622,20 @@ mod tests { super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), ], &mut graph, - ) - .unwrap(); + ); assert!(graph.iter().all(|neighbors| neighbors.is_empty())); } - #[test] - fn poisoned_candidate_lists_return_errors() { - let candidates = DirectCandidates::new(1).unwrap(); - let _ = std::panic::catch_unwind(|| { - let _guard = candidates.lists[0].lock().unwrap(); - panic!("poison candidate list"); - }); - assert!(matches!( - candidates.add_leaf(&[0], &[crate::graph::AdjacencyList::new()]), - Err(LeafBuildError::PoisonedCandidateList { point: 0 }) - )); - assert!(matches!( - candidates.into_lists(), - Err(LeafBuildError::PoisonedCandidateList { point: 0 }) - )); - } - - #[test] - fn allocation_errors_preserve_buffer_context() { - let mut values = Vec::::new(); - let source = values.try_reserve(usize::MAX).unwrap_err(); - let error = allocation_error("test", 1, source); - assert!(matches!( - error, - LeafBuildError::Allocation { - buffer: "test", - additional: 1, - .. - } - )); - } - #[test] fn direct_candidate_accumulator_keeps_unique_sorted_lists() { - let candidates = DirectCandidates::new(2).unwrap(); - candidates - .add_leaf( - &[0, 1], - &[ - crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), - crate::graph::AdjacencyList::from_iter_untrusted([0]), - ], - ) - .unwrap(); - assert_eq!( - adjacency_lists(candidates.into_lists().unwrap()), - [vec![1], vec![0]] + let candidates = DirectCandidates::new(2); + candidates.add_leaf( + &[0, 1], + &[ + crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), + crate::graph::AdjacencyList::from_iter_untrusted([0]), + ], ); + assert_eq!(adjacency_lists(candidates.into_lists()), [vec![1], vec![0]]); } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index e831741f6b..bdd51cb510 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -73,7 +73,7 @@ pub struct PiPNNConfig { /// Number of nearest centers assigned at each recursive partition level. /// Levels after this schedule assign each point to one center. pub fanout: Vec, - /// Number of nearest neighbors selected within each leaf (`1..=3`). + /// Number of nearest neighbors selected within each leaf. pub leaf_k: usize, /// Number of independent partition passes over the dataset. pub replicas: usize, @@ -106,12 +106,8 @@ impl PiPNNConfig { if self.fanout.contains(&0) { return Err(config_error("fanout values must be greater than zero")); } - if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.leaf_k) { - return Err(config_error(format!( - "leaf_k ({}) must be in [1, {}]", - self.leaf_k, - leaf_kernel::MAX_LEAF_NEIGHBORS - ))); + if self.leaf_k == 0 { + return Err(config_error("leaf_k must be greater than zero")); } if self.replicas == 0 { return Err(config_error("replicas must be greater than zero")); @@ -332,7 +328,7 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod build_graph_tests { - use super::{PiPNNBuildContext, PiPNNConfig, build_graph, leaf_kernel}; + use super::{PiPNNBuildContext, PiPNNConfig, build_graph}; use crate::graph::config::{self, MaxDegree}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; @@ -422,7 +418,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - leaf_k: leaf_kernel::MAX_LEAF_NEIGHBORS, + leaf_k: 4, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -546,7 +542,7 @@ mod build_graph_tests { c_min, p_samp: 0.5, fanout: vec![2], - leaf_k: rng.random_range(1..=3), + leaf_k: rng.random_range(1..=7), replicas: rng.random_range(1..=2), }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -564,7 +560,7 @@ mod build_graph_tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod config_tests { - use super::{PiPNNBuildContext, PiPNNConfig, leaf_kernel}; + use super::{PiPNNBuildContext, PiPNNConfig}; use crate::graph::config::{self, MaxDegree}; use diskann_vector::distance::Metric; @@ -639,10 +635,6 @@ mod config_tests { leaf_k: 0, ..pipnn_config() }, - PiPNNConfig { - leaf_k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, - ..pipnn_config() - }, PiPNNConfig { replicas: 0, ..pipnn_config() diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 052ec36608..d1a899239c 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -8,9 +8,9 @@ //! A leader is a sampled point that acts as the center of one child partition. //! A point can join several leaders, so child partitions can overlap. //! -//! This module recursively splits dataset point IDs into bounded leaves. It uses -//! dense GEMM to compare points with sampled leaders. An `ObjectPool` supplies -//! reusable buffers to Rayon worker chunks. +//! This module owns recursive splitting, leader sampling, row gathering, and +//! assignment scatter. The partition kernel owns GEMM, norm preparation, and +//! local ranking. An `ObjectPool` supplies reusable scratch to Rayon workers. //! //! A configured level assigns each point to `fanout[level]` leaders. A deeper //! level assigns each point to one leader. Each replica uses a different @@ -19,7 +19,6 @@ use std::collections::HashSet; use crate::{ANNError, ANNResult, utils::VectorRepr}; -use diskann_linalg::Transpose; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, @@ -30,8 +29,8 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::{NormPreparation, PartitionMetric, PartitionNorms}, - partition_kernel::{PartitionInput, nearest_leaders}, + kernel_metric::PartitionMetric, + partition_kernel::{PartitionKernelWorkspace, PreparedLeaders, assign_leaders}, }; // These constants control internal batching and deterministic seed generation. @@ -44,15 +43,9 @@ const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const MAX_PARTITION_ITERATIONS: usize = 30; -/// Error from partition input checks, allocation, or recursion progress. +/// Error from partition shape checks or recursion progress. #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum PartitionError { - #[error("PiPNN cannot partition an empty dataset")] - EmptyDataset, - #[error("PiPNN cannot partition vectors with zero dimensions")] - EmptyDimensions, - #[error("dataset has {0} points, which exceeds the u32 ID limit")] - TooManyPoints(usize), #[error("{buffer} shape {rows} x {cols} overflows usize")] ShapeOverflow { buffer: &'static str, @@ -68,14 +61,6 @@ pub(crate) enum PartitionError { level: usize, limit: usize, }, - #[error("invalid {buffer} length: expected {expected}, got {actual}")] - InvalidBufferLength { - buffer: &'static str, - expected: usize, - actual: usize, - }, - #[error("partition worker did not publish its result")] - MissingWorkerResult, } struct WorkItem { @@ -87,9 +72,7 @@ struct WorkItem { #[derive(Default)] struct StripeBuffers { point_values: Vec, - dot_values: Vec, - point_norms: Vec, - ranked_leaders: Vec<(u32, f32)>, + kernel_workspace: PartitionKernelWorkspace, } impl AsPooled<()> for StripeBuffers { @@ -128,26 +111,12 @@ where M: PartitionMetric, T: VectorRepr + Send + Sync, { - let points = data.nrows(); - if points == 0 { - return Err(ANNError::new(PartitionError::EmptyDataset)); - } - if data.ncols() == 0 { - return Err(ANNError::new(PartitionError::EmptyDimensions)); - } - if points > u32::MAX as usize { - return Err(ANNError::new(PartitionError::TooManyPoints(points))); - } - let mut leaves = Vec::new(); let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = partition_replica::(arch, data, config, seed, &stripe_buffers)?; - leaves - .try_reserve(replica_leaves.len()) - .map_err(ANNError::new)?; leaves.append(&mut replica_leaves); } Ok(leaves) @@ -172,67 +141,38 @@ where M: PartitionMetric, T: VectorRepr + Send + Sync, { - let initial_indices = point_ids(data.nrows())?; + let initial_indices = point_ids(data.nrows()); if data.nrows() <= config.c_max { - let mut leaves = Vec::new(); - leaves.try_reserve_exact(1).map_err(ANNError::new)?; - leaves.push(initial_indices); - return Ok(leaves); + return Ok(vec![initial_indices]); } let mut leaves = Vec::new(); - let mut work = Vec::new(); - work.try_reserve_exact(1).map_err(ANNError::new)?; - work.push(WorkItem { + let mut work = vec![WorkItem { indices: initial_indices, level: 0, seed, - }); + }]; for _ in 0..MAX_PARTITION_ITERATIONS { if work.is_empty() { return merge_undersized_leaves(leaves, config.c_min, config.c_max); } - let mut results = Vec::new(); - results - .try_reserve_exact(work.len()) - .map_err(ANNError::new)?; - results.resize_with(work.len(), || None); - // `build_graph` runs this Rayon operation in the pool from the build - // context. Each worker writes only to its indexed result slot. + // Indexed parallel collection preserves work-item order. #[allow(clippy::disallowed_methods)] - results - .par_iter_mut() - .zip(work.into_par_iter()) - .try_for_each(|(slot, item)| { - *slot = Some(partition_work_item::( - arch, - data, - config, - item, - stripe_buffers, - )?); - Ok::<(), ANNError>(()) - })?; + let results: ANNResult> = work + .into_par_iter() + .map(|item| partition_work_item::(arch, data, config, item, stripe_buffers)) + .collect(); let mut next_work = Vec::new(); - for result in results { - let (mut pending, mut finished) = - result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))?; - next_work - .try_reserve(pending.len()) - .map_err(ANNError::new)?; - leaves.try_reserve(finished.len()).map_err(ANNError::new)?; + for (mut pending, mut finished) in results? { next_work.append(&mut pending); leaves.append(&mut finished); } work = next_work; } - if work.is_empty() { - return merge_undersized_leaves(leaves, config.c_min, config.c_max); - } let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { return merge_undersized_leaves(leaves, config.c_min, config.c_max); }; @@ -268,16 +208,12 @@ where &item.indices, config.p_samp, mix_seed(item.seed, points as u64), - )?; + ); let clusters = assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; let mut pending = Vec::new(); let mut finished = Vec::new(); - pending.try_reserve(clusters.len()).map_err(ANNError::new)?; - finished - .try_reserve(clusters.len()) - .map_err(ANNError::new)?; let child_seed = mix_seed(item.seed, points as u64); for cluster in clusters { if cluster.is_empty() { @@ -297,13 +233,10 @@ where } /// Sample point IDs that act as centers for one partition split. -fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResult> { +fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> Vec { let count = sampled_leader_count(points.len(), sampling_fraction); let mut rng = rand::rngs::StdRng::seed_from_u64(seed); - let mut leaders = Vec::new(); - leaders.try_reserve_exact(count).map_err(ANNError::new)?; - leaders.extend(points.choose_multiple(&mut rng, count).copied()); - Ok(leaders) + points.choose_multiple(&mut rng, count).copied().collect() } /// Return the number of centers to sample from one cluster. @@ -351,30 +284,18 @@ where { let dimension_count = data.ncols(); let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; - let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; + let mut leader_values = vec![0.0f32; leader_values_len]; gather_vectors(data, leader_ids, &mut leader_values)?; let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count).map_err( - |error| { - ANNError::new(PartitionError::InvalidBufferLength { - buffer: "leader data", - expected: leader_values_len, - actual: error.into_inner().len(), - }) - }, - )?; - let mut leader_norm_values = Vec::new(); - M::prepare_leader_norms(NormPreparation { - values: leader_matrix, - norms: &mut leader_norm_values, - }) - .map_err(ANNError::new)?; - - let fanout = fanout.min(leader_ids.len()); + MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count) + .map_err(|error| ANNError::new(error.as_static()))?; + let leaders = PreparedLeaders::::new(leader_matrix); + + let fanout = fanout.min(leaders.len()); let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; - let mut assignments = filled_vec(assignment_len, 0u32)?; - let stripe_points = assignment_stripe_point_count(leader_ids.len()); + let mut assignments = vec![0u32; assignment_len]; + let stripe_points = assignment_stripe_point_count(leaders.len()); let stripe_assignment_count = checked_area("assignment stripe", stripe_points, fanout)?; let stripe_count = point_ids.len().div_ceil(stripe_points); let worker_stripe_count = stripe_count.div_ceil(rayon::current_num_threads().max(1)); @@ -400,8 +321,7 @@ where arch, data, &point_ids[first_point..first_point + stripe_point_count], - &leader_values, - &leader_norm_values, + &leaders, fanout, &mut buffers, stripe_assignments, @@ -415,16 +335,15 @@ where /// Assign one point stripe to sampled partition centers. /// -/// The function gathers point vectors and computes point-to-center dot products. -/// It writes center-column IDs for partition scatter. +/// The function gathers point IDs into a packed `f32` matrix. The partition +/// kernel owns dot products, point norms, and ranking. This function writes the +/// returned leader-column IDs for partition scatter. #[inline] -#[allow(clippy::too_many_arguments)] fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, point_ids: &[u32], - leader_values: &[f32], - leader_norm_values: &[f32], + leaders: &PreparedLeaders<'_, M>, fanout: usize, buffers: &mut StripeBuffers, assignments: &mut [u32], @@ -439,88 +358,30 @@ where { let point_count = point_ids.len(); let dimensions = data.ncols(); - let leader_count = leader_values.len() / dimensions; let point_values_len = checked_area("point stripe", point_count, dimensions)?; - let dots_len = checked_area("dot-product stripe", point_count, leader_count)?; - let output_len = checked_area("partition assignments", point_count, fanout)?; // Keep each buffer at its largest length. Every operation uses an explicit // active prefix. - grow_fallible(&mut buffers.point_values, point_values_len, 0.0)?; - grow_fallible(&mut buffers.dot_values, dots_len, 0.0)?; + grow(&mut buffers.point_values, point_values_len, 0.0); let StripeBuffers { - point_values: point_buffer, - dot_values: dot_buffer, - point_norms: point_norm_buffer, - ranked_leaders, + point_values, + kernel_workspace, } = buffers; - let mut point_matrix = MutMatrixView::try_from( - &mut point_buffer[..point_values_len], - point_count, - dimensions, - ) - .map_err(|error| { - ANNError::new(PartitionError::InvalidBufferLength { - buffer: "point stripe", - expected: point_values_len, - actual: error.into_inner().len(), - }) - })?; - let dots = &mut dot_buffer[..dots_len]; - gather_vectors(data, point_ids, point_matrix.as_mut_slice())?; - diskann_linalg::sgemm( - Transpose::None, - Transpose::Ordinary, + let mut points = MutMatrixView::try_from( + &mut point_values[..point_values_len], point_count, - leader_count, dimensions, - 1.0, - point_matrix.as_slice(), - leader_values, - None, - dots, ) - .map_err(ANNError::new)?; - - M::prepare_point_norms(NormPreparation { - values: point_matrix.as_view(), - norms: point_norm_buffer, - }) - .map_err(ANNError::new)?; - let point_norms = &*point_norm_buffer; - let norms = PartitionNorms { - point_norms, - leader_norms: leader_norm_values, - }; - let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { - ANNError::new(PartitionError::InvalidBufferLength { - buffer: "dot-product stripe", - expected: dots_len, - actual: dots.len(), - }) - })?; - let output = MutMatrixView::try_from(assignments, point_count, fanout).map_err(|error| { - ANNError::new(PartitionError::InvalidBufferLength { - buffer: "partition assignments", - expected: output_len, - actual: error.into_inner().len(), - }) - })?; - nearest_leaders::(arch, PartitionInput { dots, norms }, output, ranked_leaders) - .map_err(ANNError::new) + .map_err(|error| ANNError::new(error.as_static()))?; + gather_vectors(data, point_ids, points.as_mut_slice())?; + let output = MutMatrixView::try_from(assignments, point_count, fanout) + .map_err(|error| ANNError::new(error.as_static()))?; + assign_leaders::(arch, points.as_view(), leaders, output, kernel_workspace) } fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> where T: VectorRepr, { - let expected = checked_area("gather output", indices.len(), data.ncols())?; - if output.len() != expected { - return Err(ANNError::new(PartitionError::InvalidBufferLength { - buffer: "gather output", - expected, - actual: output.len(), - })); - } for (&index, vector_output) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { T::as_f32_into(data.row(index as usize), vector_output).map_err(Into::::into)?; } @@ -538,64 +399,41 @@ fn scatter_assignments( leaders: usize, ) -> ANNResult>> { if points.len() < PARALLEL_SCATTER_MIN_POINTS { - return scatter_serial(points, assignments, fanout, leaders); + return Ok(scatter_serial(points, assignments, fanout, leaders)); } let stripe_points = points.len().div_ceil(rayon::current_num_threads().max(1)); let stripe_assignment_count = checked_area("scatter assignment stripe", stripe_points, fanout)?; - let stripes = points.len().div_ceil(stripe_points); - let mut partials = Vec::new(); - partials.try_reserve_exact(stripes).map_err(ANNError::new)?; - partials.resize_with(stripes, || None); - // `build_graph` runs this Rayon operation in the pool from the build context. - // Each worker writes only to its indexed partial result. + // Indexed parallel collection preserves stripe order. #[allow(clippy::disallowed_methods)] - partials - .par_iter_mut() - .zip( - points - .par_chunks(stripe_points) - .zip(assignments.par_chunks(stripe_assignment_count)), - ) - .try_for_each(|(slot, (points, assignments))| { - *slot = Some(scatter_serial(points, assignments, fanout, leaders)?); - Ok::<(), ANNError>(()) - })?; + let locals: Vec<_> = points + .par_chunks(stripe_points) + .zip(assignments.par_chunks(stripe_assignment_count)) + .map(|(points, assignments)| scatter_serial(points, assignments, fanout, leaders)) + .collect(); - let mut locals = Vec::new(); - locals.try_reserve_exact(stripes).map_err(ANNError::new)?; - for result in partials { - locals.push(result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))?); - } - - let mut sizes = filled_vec(leaders, 0usize)?; + let mut sizes = vec![0usize; leaders]; for local in &locals { for (size, cluster) in sizes.iter_mut().zip(local) { - *size = size.checked_add(cluster.len()).ok_or_else(|| { - ANNError::new(PartitionError::ShapeOverflow { - buffer: "cluster size", - rows: *size, - cols: cluster.len(), - }) - })?; + *size += cluster.len(); } } // `build_graph` runs this Rayon operation in the pool from the build context. // Each worker creates one independent leader cluster. #[allow(clippy::disallowed_methods)] - sizes + let clusters = sizes .into_par_iter() .enumerate() .map(|(leader, size)| { - let mut cluster = Vec::new(); - cluster.try_reserve_exact(size).map_err(ANNError::new)?; + let mut cluster = Vec::with_capacity(size); for local in &locals { cluster.extend_from_slice(&local[leader]); } - Ok(cluster) + cluster }) - .collect() + .collect(); + Ok(clusters) } fn scatter_serial( @@ -603,53 +441,22 @@ fn scatter_serial( assignments: &[u32], fanout: usize, leaders: usize, -) -> ANNResult>> { - let expected = checked_area("scatter assignments", points.len(), fanout)?; - if assignments.len() != expected { - return Err(ANNError::new(PartitionError::InvalidBufferLength { - buffer: "scatter assignments", - expected, - actual: assignments.len(), - })); - } - - let mut sizes = filled_vec(leaders, 0usize)?; +) -> Vec> { + let mut sizes = vec![0usize; leaders]; for &leader in assignments { - let Some(size) = sizes.get_mut(leader as usize) else { - return Err(ANNError::new(PartitionError::InvalidBufferLength { - buffer: "leader assignment", - expected: leaders, - actual: leader as usize + 1, - })); - }; - *size = size.checked_add(1).ok_or_else(|| { - ANNError::new(PartitionError::ShapeOverflow { - buffer: "cluster size", - rows: *size, - cols: 1, - }) - })?; + sizes[leader as usize] += 1; } - let mut clusters = clusters_with_capacities(&sizes)?; + let mut clusters = clusters_with_capacities(&sizes); for (&point, point_assignments) in points.iter().zip(assignments.chunks_exact(fanout)) { for &leader in point_assignments { clusters[leader as usize].push(point); } } - Ok(clusters) + clusters } -fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { - let mut clusters = Vec::new(); - clusters - .try_reserve_exact(sizes.len()) - .map_err(ANNError::new)?; - for &size in sizes { - let mut cluster = Vec::new(); - cluster.try_reserve_exact(size).map_err(ANNError::new)?; - clusters.push(cluster); - } - Ok(clusters) +fn clusters_with_capacities(sizes: &[usize]) -> Vec> { + sizes.iter().map(|&size| Vec::with_capacity(size)).collect() } /// Merge leaves smaller than `c_min` without exceeding `c_max`. @@ -661,12 +468,8 @@ fn merge_undersized_leaves( c_min: usize, c_max: usize, ) -> ANNResult>> { - let mut merged = Vec::new(); + let mut merged = Vec::with_capacity(leaves.len()); let mut small_leaves = Vec::new(); - merged.try_reserve(leaves.len()).map_err(ANNError::new)?; - small_leaves - .try_reserve(leaves.len()) - .map_err(ANNError::new)?; for leaf in leaves { if leaf.len() >= c_min { merged.push(leaf); @@ -678,42 +481,27 @@ fn merge_undersized_leaves( return Ok(merged); } - let mut small = HashSet::new(); - small.try_reserve(c_max).map_err(ANNError::new)?; + let mut small = HashSet::with_capacity(c_max); for leaf in small_leaves { - let combined = small.len().checked_add(leaf.len()).ok_or_else(|| { - ANNError::new(PartitionError::ShapeOverflow { - buffer: "small-leaf merge", - rows: small.len(), - cols: leaf.len(), - }) - })?; + let combined = small.len() + leaf.len(); if combined > c_max { - merged.push(drain_sorted(&mut small)?); + merged.push(drain_sorted(&mut small)); } - small.try_reserve(leaf.len()).map_err(ANNError::new)?; small.extend(leaf); if small.len() >= c_min { - merged.push(drain_sorted(&mut small)?); + merged.push(drain_sorted(&mut small)); } } if !small.is_empty() { - let mut remainder = drain_sorted(&mut small)?; + let mut remainder = drain_sorted(&mut small); if remainder.len() < c_min && let Some(last) = merged.last_mut() { remainder.retain(|id| !last.contains(id)); - let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { - ANNError::new(PartitionError::ShapeOverflow { - buffer: "small-leaf tail merge", - rows: last.len(), - cols: remainder.len(), - }) - })?; + let combined = last.len() + remainder.len(); if combined <= c_max { - last.try_reserve(remainder.len()).map_err(ANNError::new)?; last.append(&mut remainder); last.sort_unstable(); } @@ -726,37 +514,20 @@ fn merge_undersized_leaves( Ok(merged) } -fn drain_sorted(set: &mut HashSet) -> ANNResult> { - let mut values = Vec::new(); - values.try_reserve_exact(set.len()).map_err(ANNError::new)?; - values.extend(set.drain()); +fn drain_sorted(set: &mut HashSet) -> Vec { + let mut values: Vec<_> = set.drain().collect(); values.sort_unstable(); - Ok(values) -} - -fn point_ids(points: usize) -> ANNResult> { - let mut ids = Vec::new(); - ids.try_reserve_exact(points).map_err(ANNError::new)?; - ids.extend(0..points as u32); - Ok(ids) + values } -fn filled_vec(len: usize, value: T) -> ANNResult> { - let mut values = Vec::new(); - values.try_reserve_exact(len).map_err(ANNError::new)?; - values.resize(len, value); - Ok(values) +fn point_ids(points: usize) -> Vec { + (0..points as u32).collect() } -fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { - if values.len() >= len { - return Ok(()); +fn grow(values: &mut Vec, len: usize, value: T) { + if values.len() < len { + values.resize(len, value); } - values - .try_reserve(len - values.len()) - .map_err(ANNError::new)?; - values.resize(len, value); - Ok(()) } fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { @@ -1184,78 +955,9 @@ mod tests { .flat_map(|point| [point % 7, (point + 3) % 7]) .collect(); - let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); + let expected = scatter_serial(&points, &assignments, 2, 7); let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); assert_eq!(actual, expected); } - - #[test] - fn rejects_empty_dataset() { - let data = Matrix::::new(0.0, 0, 4); - let error = - partition_with_runtime_metric(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2) - .unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::EmptyDataset - ); - } - - #[test] - fn rejects_zero_dimensions() { - let data = Matrix::::new(0.0, 4, 0); - let error = - partition_with_runtime_metric(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2) - .unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::EmptyDimensions - ); - } - - #[test] - fn rejects_invalid_gather_output_length() { - let data = Matrix::::new(0.0, 2, 2); - let error = gather_vectors(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "gather output", - expected: 4, - actual: 3, - } - ); - } - - #[test] - fn rejects_invalid_assignment_length() { - let error = scatter_serial(&[7, 8], &[0, 1, 0], 2, 2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "scatter assignments", - expected: 4, - actual: 3, - } - ); - } - - #[test] - fn rejects_assignment_to_an_unknown_leader() { - let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "leader assignment", - expected: 2, - actual: 3, - } - ); - } } From cd69957e74d15483db8ada7ed6c73d7bf1a3e0bc Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:43:50 +0000 Subject: [PATCH 156/229] fix(pipnn): omit unrankable graph candidates --- diskann/src/graph/pipnn/leaf_build.rs | 12 ++++++++++++ diskann/src/graph/pipnn/partitioning.rs | 25 ++++++++++++++++++++++--- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 5f02967125..20d276d180 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -265,6 +265,9 @@ fn add_symmetric_neighbors( ) { for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { for neighbor in source_neighbors { + if !neighbor.is_assigned() { + continue; + } let target = neighbor.target as usize; let source_id = point_ids[source]; let target_id = point_ids[target]; @@ -442,6 +445,15 @@ mod tests { assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); } + #[test] + fn non_rankable_neighbors_are_omitted() { + let data = [0.0_f32, 1.0, f32::NAN]; + + let graph = build(view(&data, 3, 1), &[vec![0, 1, 2]], 2, Metric::InnerProduct).unwrap(); + + assert_eq!(adjacency_lists(graph), [vec![1], vec![0], vec![]]); + } + #[test] fn retains_and_deduplicates_candidates_from_overlapping_leaves() { let data = [0.0_f32, 1.0, 2.0, 3.0]; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index d1a899239c..9d13826d83 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -30,7 +30,9 @@ use rayon::prelude::*; use super::{ PiPNNConfig, kernel_metric::PartitionMetric, - partition_kernel::{PartitionKernelWorkspace, PreparedLeaders, assign_leaders}, + partition_kernel::{ + PartitionKernelWorkspace, PreparedLeaders, UNASSIGNED_LEADER, assign_leaders, + }, }; // These constants control internal batching and deterministic seed generation. @@ -444,12 +446,16 @@ fn scatter_serial( ) -> Vec> { let mut sizes = vec![0usize; leaders]; for &leader in assignments { - sizes[leader as usize] += 1; + if leader != UNASSIGNED_LEADER { + sizes[leader as usize] += 1; + } } let mut clusters = clusters_with_capacities(&sizes); for (&point, point_assignments) in points.iter().zip(assignments.chunks_exact(fanout)) { for &leader in point_assignments { - clusters[leader as usize].push(point); + if leader != UNASSIGNED_LEADER { + clusters[leader as usize].push(point); + } } } clusters @@ -947,6 +953,19 @@ mod tests { assert_eq!(clusters[1], (1_024..2_048).collect::>()); } + #[test] + fn scatter_omits_unassigned_slots() { + assert_eq!( + scatter_serial( + &[10, 11], + &[0, UNASSIGNED_LEADER, UNASSIGNED_LEADER, 1], + 2, + 2 + ), + [vec![10], vec![11]] + ); + } + #[test] fn parallel_scatter_matches_serial_order() { let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); From 4068dad54dc744008c99f1b699bc8340b8260b28 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:50:19 +0000 Subject: [PATCH 157/229] test(pipnn): cover non-rankable build candidates --- diskann/src/graph/pipnn/mod.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index bdd51cb510..9842ff701c 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -407,6 +407,28 @@ mod build_graph_tests { } } + #[test] + fn omits_non_rankable_candidates_without_invalid_ids() { + let values = [0.0_f32, 1.0, f32::NAN]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let graph = graph_config(Metric::InnerProduct, 2); + let pool = pool(1); + let config = PiPNNConfig { + c_max: 2, + c_min: 1, + p_samp: 1.0, + fanout: vec![2], + leaf_k: 1, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::InnerProduct, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&actual, 3, 2); + assert_eq!(rows(actual), [vec![1], vec![0], vec![]]); + } + #[test] fn prunes_overfull_single_leaf_candidates_to_the_graph_degree() { let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; From e009781ad4acd95ee7c10818cfadc5785d83029f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 19 Aug 2026 05:45:11 +0000 Subject: [PATCH 158/229] perf(pipnn): deduplicate direct leaf edges once --- diskann/src/graph/pipnn/leaf_build.rs | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 20d276d180..18b952d329 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -63,7 +63,7 @@ pub(crate) enum LeafBuildError { struct LeafBuffers { point_values: Vec, neighbors: Vec, - local_adjacency: Vec>, + local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, } @@ -99,11 +99,10 @@ impl LeafBuffers { } fn prepare_local_adjacency(&mut self, point_count: usize) { - self.local_adjacency - .resize_with(point_count, AdjacencyList::new); + self.local_adjacency.resize_with(point_count, Vec::new); self.local_adjacency[..point_count] .iter_mut() - .for_each(AdjacencyList::clear); + .for_each(Vec::clear); } } @@ -124,7 +123,7 @@ impl DirectCandidates { Self { lists } } - fn add_leaf(&self, point_ids: &[u32], local_adjacency: &[AdjacencyList]) { + fn add_leaf(&self, point_ids: &[u32], local_adjacency: &[Vec]) { for (&source, additions) in point_ids.iter().zip(local_adjacency) { // `add_direct_leaf_candidates` checks every point ID before this append. self.lists[source as usize] @@ -261,7 +260,7 @@ fn add_symmetric_neighbors( point_ids: &[u32], leaf_k: usize, neighbors: &[LeafNeighbor], - local_adjacency: &mut [AdjacencyList], + local_adjacency: &mut [Vec], ) { for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { for neighbor in source_neighbors { @@ -625,7 +624,7 @@ mod tests { #[test] fn skips_duplicate_global_ids_without_self_edges() { - let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; + let mut graph = vec![Vec::new(); 2]; add_symmetric_neighbors( &[7, 7], 1, @@ -641,13 +640,7 @@ mod tests { #[test] fn direct_candidate_accumulator_keeps_unique_sorted_lists() { let candidates = DirectCandidates::new(2); - candidates.add_leaf( - &[0, 1], - &[ - crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), - crate::graph::AdjacencyList::from_iter_untrusted([0]), - ], - ); + candidates.add_leaf(&[0, 1], &[vec![1, 1], vec![0]]); assert_eq!(adjacency_lists(candidates.into_lists()), [vec![1], vec![0]]); } } From e271b79c3ddc828ecdc2c39c4a8e45ec939eae5f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:09:29 +0000 Subject: [PATCH 159/229] refactor(pipnn): use the SIMD schema in stages --- diskann/src/graph/pipnn/leaf_build.rs | 23 ++++----------- diskann/src/graph/pipnn/mod.rs | 20 ++++--------- diskann/src/graph/pipnn/partitioning.rs | 37 ++++++------------------- 3 files changed, 20 insertions(+), 60 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 18b952d329..bdc4fc06dd 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -21,12 +21,12 @@ use parking_lot::Mutex; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ kernel_metric::LeafMetric, leaf_kernel::{LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, select_leaf_neighbors}, + simd::PiPNNSIMDSchema, }; /// Failure while converting leaves into direct graph candidates. @@ -156,10 +156,7 @@ pub(super) fn build_leaf_candidates( requested_k: usize, ) -> Result>, LeafBuildError> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: LeafMetric, T: VectorRepr + 'static, { @@ -196,10 +193,7 @@ fn add_direct_leaf_candidates( candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: LeafMetric, T: VectorRepr + 'static, { @@ -288,13 +282,11 @@ fn grow(values: &mut Vec, len: usize, value: T) { mod tests { use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; - use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, - }; + use diskann_wide::arch::{self, Target1}; use half::f16; use std::collections::BTreeSet; + use super::super::simd::PiPNNSIMDSchema; use super::{ DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, build_leaf_candidates, @@ -326,10 +318,7 @@ mod tests { LeafBuildCall<'_, T>, > for DispatchLeafBuild where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, T: crate::utils::VectorRepr + 'static, { fn run( diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 9842ff701c..d55f90eb9d 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -49,14 +49,12 @@ use crate::{ }; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; -use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, -}; +use diskann_wide::arch::{self, Target1}; use rayon::ThreadPool; -use self::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, LeafMetric, PartitionMetric, +use self::{ + kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2, LeafMetric, PartitionMetric}, + simd::PiPNNSIMDSchema, }; /// PiPNN partition and leaf-selection policy. @@ -215,10 +213,7 @@ struct RunBuildGraph; impl Target1>>, BuildGraphCall<'_, '_, '_, T>> for RunBuildGraph where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, T: VectorRepr + Send + Sync + 'static, { fn run( @@ -258,10 +253,7 @@ fn build_graph_for( metric: Metric, ) -> ANNResult>> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: LeafMetric + PartitionMetric, T: VectorRepr + Send + Sync + 'static, { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 9d13826d83..7325630f67 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -23,7 +23,6 @@ use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, }; -use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; @@ -33,6 +32,7 @@ use super::{ partition_kernel::{ PartitionKernelWorkspace, PreparedLeaders, UNASSIGNED_LEADER, assign_leaders, }, + simd::PiPNNSIMDSchema, }; // These constants control internal batching and deterministic seed generation. @@ -106,10 +106,7 @@ pub(super) fn partition( config: &PiPNNConfig, ) -> ANNResult>> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { @@ -136,10 +133,7 @@ fn partition_replica( stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { @@ -197,10 +191,7 @@ fn partition_work_item( stripe_buffers: &StripeBufferPool, ) -> ANNResult<(Vec, Vec>)> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { @@ -277,10 +268,7 @@ fn assign_to_leaders( stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { @@ -351,10 +339,7 @@ fn assign_point_stripe( assignments: &mut [u32], ) -> ANNResult<()> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr, { @@ -555,10 +540,7 @@ fn assignment_stripe_point_count(leader_count: usize) -> usize { mod tests { use diskann_utils::views::{Matrix, MatrixView}; use diskann_vector::{Half, distance::Metric}; - use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, - }; + use diskann_wide::arch::{self, Target1}; use super::*; @@ -571,10 +553,7 @@ mod tests { impl Target1>>, PartitionCall<'_, T>> for DispatchPartition where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, T: VectorRepr + Send + Sync, { fn run(self, arch: A, call: PartitionCall<'_, T>) -> ANNResult>> { From c0b978d0de158cbb78ac6ed292e8e4212c821571 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:08:41 +0000 Subject: [PATCH 160/229] test(pipnn): clarify core graph contracts Keep exact graph-output, stale-state, boundary, and deterministic behavior checks while dropping allocation-only assertions. --- diskann/src/graph/pipnn/finalization.rs | 39 ++++++--- diskann/src/graph/pipnn/leaf_build.rs | 54 ++++++------- diskann/src/graph/pipnn/partitioning.rs | 100 +++++++++++------------- 3 files changed, 99 insertions(+), 94 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 6f01c3dd25..754865bc9c 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -166,7 +166,8 @@ mod tests { } #[test] - fn preserves_lists_within_the_degree_bound() { + fn row_within_the_degree_bound_is_only_canonicalized() { + // Given let data = [0.0_f32, 1.0, 2.0, 3.0]; let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); let candidates = vec![ @@ -175,14 +176,18 @@ mod tests { candidate_list([]), candidate_list([]), ]; + let expected_canonical_row = [1, 3]; - let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + // When + let actual_rows = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - assert_eq!(&*actual[0], &[1, 3]); + // Then + assert_eq!(&*actual_rows[0], &expected_canonical_row); } #[test] - fn prunes_an_overfull_list_with_the_vamana_kernel() { + fn overfull_row_keeps_the_nearest_unoccluded_neighbors() { + // Given let data = [0.0_f32, 1.0, 2.0, -3.0]; let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); let candidates = vec![ @@ -191,14 +196,18 @@ mod tests { candidate_list([]), candidate_list([]), ]; + let expected_pruned_row = [1, 3]; - let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + // When + let actual_rows = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - assert_eq!(&*actual[0], &[1, 3]); + // Then + assert_eq!(&*actual_rows[0], &expected_pruned_row); } #[test] fn reused_workspace_matches_fresh_pruning() { + // Given let data = [0.0_f32, 1.0, 2.0, -3.0, 4.0]; let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); let first = [3, 2, 1]; @@ -217,18 +226,26 @@ mod tests { .num_threads(1) .build() .unwrap(); - let fresh_first = pool + // When + let expected_first_row_from_fresh_workspace = pool .install(|| prune_overfull(data, candidates(&first, &[]), &graph, Metric::L2)) .unwrap(); - let fresh_second = pool + let expected_second_row_from_fresh_workspace = pool .install(|| prune_overfull(data, candidates(&[], &second), &graph, Metric::L2)) .unwrap(); - let reused = pool + let actual_rows_from_reused_workspace = pool .install(|| prune_overfull(data, candidates(&first, &second), &graph, Metric::L2)) .unwrap(); - assert_eq!(&*reused[0], &*fresh_first[0]); - assert_eq!(&*reused[1], &*fresh_second[1]); + // Then + assert_eq!( + &*actual_rows_from_reused_workspace[0], + &*expected_first_row_from_fresh_workspace[0] + ); + assert_eq!( + &*actual_rows_from_reused_workspace[1], + &*expected_second_row_from_fresh_workspace[1] + ); } } diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index bdc4fc06dd..2390e0713c 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -410,6 +410,7 @@ mod tests { #[test] fn leaf_adjacency_matches_an_independent_all_pairs_reference() { + // Given let points = [ [0.0_f32, 0.0], [1.0, 0.2], @@ -419,8 +420,10 @@ mod tests { [6.7, -3.2], ]; let flat: Vec<_> = points.into_iter().flatten().collect(); + let expected_adjacency = brute_force_symmetric_l2(&points, 2); - let actual = adjacency_lists( + // When + let actual_adjacency = adjacency_lists( build( view(&flat, points.len(), 2), &[(0..points.len() as u32).collect()], @@ -430,33 +433,41 @@ mod tests { .unwrap(), ); - assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); + // Then + assert_eq!(actual_adjacency, expected_adjacency); } #[test] fn non_rankable_neighbors_are_omitted() { + // Given let data = [0.0_f32, 1.0, f32::NAN]; + let expected_adjacency = [vec![1], vec![0], vec![]]; + // When let graph = build(view(&data, 3, 1), &[vec![0, 1, 2]], 2, Metric::InnerProduct).unwrap(); + let actual_adjacency = adjacency_lists(graph); - assert_eq!(adjacency_lists(graph), [vec![1], vec![0], vec![]]); + // Then + assert_eq!(actual_adjacency, expected_adjacency); } #[test] - fn retains_and_deduplicates_candidates_from_overlapping_leaves() { + fn overlapping_leaves_contribute_each_candidate_once() { + // Given let data = [0.0_f32, 1.0, 2.0, 3.0]; let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; + let expected_adjacency = [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]]; + // When let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); + let actual_adjacency = adjacency_lists(graph); - assert_eq!( - adjacency_lists(graph), - [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] - ); + // Then + assert_eq!(actual_adjacency, expected_adjacency); } #[test] - fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { + fn symmetric_edges_can_give_a_center_more_than_two_k_neighbors() { let dimensions = 9; let mut data = vec![0.0_f32; 10 * dimensions]; for source in 1..10 { @@ -561,10 +572,10 @@ mod tests { let leaves: Vec> = (0..32) .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) .collect(); - let expected = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + let expected_candidate_order = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); for _ in 0..8 { - let actual = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); - assert_eq!(actual, expected); + let actual_candidate_order = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + assert_eq!(actual_candidate_order, expected_candidate_order); } } @@ -588,22 +599,7 @@ mod tests { } #[test] - fn reuses_worker_buffers_for_smaller_leaves() { - let mut buffers = LeafBuffers::default(); - buffers.prepare(0, 64, 128, 2).unwrap(); - let point_values = buffers.point_values.as_ptr(); - let neighbors = buffers.neighbors.as_ptr(); - - buffers.prepare(1, 8, 128, 2).unwrap(); - - assert_eq!(buffers.point_values.as_ptr(), point_values); - assert_eq!(buffers.neighbors.as_ptr(), neighbors); - assert_eq!(buffers.point_values.len(), 64 * 128); - assert_eq!(buffers.neighbors.len(), 64 * 2); - } - - #[test] - fn reports_shape_overflow_before_allocating() { + fn leaf_buffer_preparation_reports_shape_overflow_before_allocating() { let mut buffers = LeafBuffers::default(); assert!(matches!( buffers.prepare(7, usize::MAX, 2, 1), @@ -612,7 +608,7 @@ mod tests { } #[test] - fn skips_duplicate_global_ids_without_self_edges() { + fn symmetric_edge_mapping_skips_duplicate_ids_instead_of_adding_self_edges() { let mut graph = vec![Vec::new(); 2]; add_symmetric_neighbors( &[7, 7], diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 7325630f67..395f71ede1 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -670,7 +670,7 @@ mod tests { } #[test] - fn returns_one_leaf_at_and_below_c_max() { + fn partition_returns_one_leaf_when_point_count_does_not_exceed_c_max() { for points in [7, 8] { let data = clustered_data(points, 3); let leaves = partition_with_runtime_metric( @@ -726,30 +726,42 @@ mod tests { #[test] fn global_merge_canonicalizes_small_leaf_membership() { + // Given let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; + let expected_canonical_membership = vec![vec![1, 2, 3, 8, 9]]; - let merged = merge_undersized_leaves(leaves, 4, 8).unwrap(); + // When + let actual_leaves = merge_undersized_leaves(leaves, 4, 8).unwrap(); - assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); + // Then + assert_eq!(actual_leaves, expected_canonical_membership); } #[test] fn global_merge_never_overfills_before_reaching_c_min() { + // Given let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; + let expected_capacity_bounded_leaves = + vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]]; - let merged = merge_undersized_leaves(leaves, 11, 11).unwrap(); + // When + let actual_leaves = merge_undersized_leaves(leaves, 11, 11).unwrap(); - assert_eq!( - merged, - vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] - ); + // Then + assert_eq!(actual_leaves, expected_capacity_bounded_leaves); } #[test] fn global_merge_fills_exact_capacity_before_flushing() { - let merged = merge_undersized_leaves(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + // Given + let leaves = vec![vec![0, 1], vec![2, 3]]; + let expected_exact_capacity_leaf = vec![vec![0, 1, 2, 3]]; + + // When + let actual_leaves = merge_undersized_leaves(leaves, 4, 4).unwrap(); - assert_eq!(merged, vec![vec![0, 1, 2, 3]]); + // Then + assert_eq!(actual_leaves, expected_exact_capacity_leaf); } #[test] @@ -783,23 +795,23 @@ mod tests { let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); let converted: Vec = raw.iter().copied().map(&convert).collect(); let config = config(2, 16, vec![2, 1], 1); - let expected = partition_with_runtime_metric( + let expected_f32_partition = partition_with_runtime_metric( MatrixView::try_from(&f32_data, points, dimensions).unwrap(), &config, Metric::L2, ) .unwrap(); - let actual = partition_with_runtime_metric( + let actual_converted_partition = partition_with_runtime_metric( MatrixView::try_from(&converted, points, dimensions).unwrap(), &config, Metric::L2, ) .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); - assert_valid_partition_with_runtime_metric(&actual, points, 16, 1); + assert_valid_partition_with_runtime_metric(&actual_converted_partition, points, 16, 1); assert_eq!( - sorted_memberships(&actual), - sorted_memberships(&expected), + sorted_memberships(&actual_converted_partition), + sorted_memberships(&expected_f32_partition), "{label} dimensions={dimensions}" ); } @@ -859,7 +871,7 @@ mod tests { } #[test] - fn all_metrics_produce_valid_partitions() { + fn every_metric_covers_all_points_without_exceeding_c_max() { let data = directional_data(64, 8); let config = config(2, 20, vec![2], 1); @@ -875,7 +887,7 @@ mod tests { } #[test] - fn leader_count_is_bounded() { + fn sampled_leader_count_is_at_least_one_and_never_exceeds_the_cap() { assert_eq!(sampled_leader_count(1, 1.0), 1); assert_eq!(sampled_leader_count(10, 0.01), 2); assert_eq!(sampled_leader_count(50_000, 1.0), LEADER_CAP); @@ -888,31 +900,7 @@ mod tests { } #[test] - fn assignment_stripes_use_power_of_two_point_counts() { - assert_eq!(assignment_stripe_point_count(1_000), 128); - assert_eq!(assignment_stripe_point_count(256), 512); - assert_eq!( - assignment_stripe_point_count(1), - MAX_ASSIGNMENT_STRIPE_POINTS - ); - } - - #[test] - fn stripe_buffer_pool_reuses_returned_capacity() { - let pool = StripeBufferPool::new((), 0, None); - let points = { - let mut buffers = pool.get_ref(()); - buffers.point_values.resize(16, 0.0); - buffers.point_values.as_ptr() - }; - - let buffers = pool.get_ref(()); - assert_eq!(buffers.point_values.as_ptr(), points); - assert_eq!(buffers.point_values.len(), 16); - } - - #[test] - fn leader_assignment_handles_multiple_stripes() { + fn leader_assignment_preserves_clusters_across_multiple_stripes() { let points = 2_048; let data: Vec = (0..points).map(|point| point as f32).collect(); let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); @@ -934,28 +922,32 @@ mod tests { #[test] fn scatter_omits_unassigned_slots() { - assert_eq!( - scatter_serial( - &[10, 11], - &[0, UNASSIGNED_LEADER, UNASSIGNED_LEADER, 1], - 2, - 2 - ), - [vec![10], vec![11]] - ); + // Given + let points = [10, 11]; + let assignments = [0, UNASSIGNED_LEADER, UNASSIGNED_LEADER, 1]; + let expected_clusters = [vec![10], vec![11]]; + + // When + let actual_clusters = scatter_serial(&points, &assignments, 2, 2); + + // Then + assert_eq!(actual_clusters, expected_clusters); } #[test] fn parallel_scatter_matches_serial_order() { + // Given let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); let assignments: Vec = points .iter() .flat_map(|point| [point % 7, (point + 3) % 7]) .collect(); - let expected = scatter_serial(&points, &assignments, 2, 7); - let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); + // When + let expected_serial_clusters = scatter_serial(&points, &assignments, 2, 7); + let actual_parallel_clusters = scatter_assignments(&points, &assignments, 2, 7).unwrap(); - assert_eq!(actual, expected); + // Then + assert_eq!(actual_parallel_clusters, expected_serial_clusters); } } From 6251b0337738ecf59960c57132a5008f75c9ce89 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 24 Aug 2026 08:50:37 +0000 Subject: [PATCH 161/229] test(pipnn): use semantic graph fixtures Derive partition and candidate inputs from named point geometry, split distinct behaviors, and parameterize equivalent boundaries. --- diskann/src/graph/pipnn/finalization.rs | 9 +- diskann/src/graph/pipnn/leaf_build.rs | 163 +++++++++++------ diskann/src/graph/pipnn/partitioning.rs | 230 ++++++++++++++++-------- 3 files changed, 276 insertions(+), 126 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 754865bc9c..d67fb17a3c 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -188,8 +188,13 @@ mod tests { #[test] fn overfull_row_keeps_the_nearest_unoccluded_neighbors() { // Given - let data = [0.0_f32, 1.0, 2.0, -3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let point_values = [ + 0.0_f32, // source + 1.0, // nearest point on the positive ray + 2.0, // farther point occluded by point 1 + -3.0, // point on the opposite ray + ]; + let data = MatrixView::try_from(&point_values[..], 4, 1).unwrap(); let candidates = vec![ candidate_list([3, 2, 1]), candidate_list([]), diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 2390e0713c..d680deb7e2 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -284,6 +284,7 @@ mod tests { use diskann_vector::distance::Metric; use diskann_wide::arch::{self, Target1}; use half::f16; + use rstest::rstest; use std::collections::BTreeSet; use super::super::simd::PiPNNSIMDSchema; @@ -292,11 +293,11 @@ mod tests { build_leaf_candidates, }; - fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + fn matrix_view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { MatrixView::try_from(data, rows, columns).unwrap() } - fn pool() -> rayon::ThreadPool { + fn leaf_build_pool() -> rayon::ThreadPool { rayon::ThreadPoolBuilder::new() .num_threads(4) .build() @@ -351,7 +352,7 @@ mod tests { } } - fn build( + fn build_candidate_graph( data: MatrixView<'_, T>, leaves: &[Vec], k: usize, @@ -360,7 +361,7 @@ mod tests { where T: crate::utils::VectorRepr + 'static, { - pool().install(|| { + leaf_build_pool().install(|| { arch::dispatch1_no_features( DispatchLeafBuild(metric), LeafBuildCall { @@ -376,6 +377,20 @@ mod tests { graph.into_iter().map(Vec::from).collect() } + fn circular_overlapping_leaves( + point_count: usize, + leaf_count: usize, + leaf_size: usize, + ) -> Vec> { + (0..leaf_count) + .map(|offset| { + (0..leaf_size) + .map(|point| ((point + offset) % point_count) as u32) + .collect() + }) + .collect() + } + fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { let mut graph = vec![BTreeSet::new(); data.len()]; for (source, left) in data.iter().enumerate() { @@ -424,8 +439,8 @@ mod tests { // When let actual_adjacency = adjacency_lists( - build( - view(&flat, points.len(), 2), + build_candidate_graph( + matrix_view(&flat, points.len(), 2), &[(0..points.len() as u32).collect()], 2, Metric::L2, @@ -444,7 +459,13 @@ mod tests { let expected_adjacency = [vec![1], vec![0], vec![]]; // When - let graph = build(view(&data, 3, 1), &[vec![0, 1, 2]], 2, Metric::InnerProduct).unwrap(); + let graph = build_candidate_graph( + matrix_view(&data, 3, 1), + &[vec![0, 1, 2]], + 2, + Metric::InnerProduct, + ) + .unwrap(); let actual_adjacency = adjacency_lists(graph); // Then @@ -459,7 +480,8 @@ mod tests { let expected_adjacency = [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]]; // When - let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); + let graph = + build_candidate_graph(matrix_view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); let actual_adjacency = adjacency_lists(graph); // Then @@ -474,8 +496,8 @@ mod tests { data[source * dimensions + source - 1] = 1.0; } - let graph = build( - view(&data, 10, dimensions), + let graph = build_candidate_graph( + matrix_view(&data, 10, dimensions), &[(0..10).collect()], 1, Metric::L2, @@ -491,12 +513,20 @@ mod tests { })); } - fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> + fn build_l2_candidate_graph(data: &[T], points: usize, dimensions: usize) -> Vec> where T: crate::utils::VectorRepr + 'static, { let leaves = vec![(0..points as u32).collect()]; - adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) + adjacency_lists( + build_candidate_graph( + matrix_view(data, points, dimensions), + &leaves, + 2, + Metric::L2, + ) + .unwrap(), + ) } fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) @@ -512,14 +542,14 @@ mod tests { .map(|index| { let source = index / dimensions; let dimension = index % dimensions; - ((source * 7 + dimension * 3 + source * dimension) % 23) as u8 + (source + dimension) as u8 }) .collect(); let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); let converted: Vec = raw.iter().copied().map(&convert).collect(); assert_eq!( - source_graph(&converted, points, dimensions), - source_graph(&f32_data, points, dimensions), + build_l2_candidate_graph(&converted, points, dimensions), + build_l2_candidate_graph(&f32_data, points, dimensions), "{label} dimensions={dimensions}" ); } @@ -542,60 +572,87 @@ mod tests { assert_source_conversion_matches_f32("i8", |value| value as i8 - 11); } - #[test] - fn all_metrics_produce_symmetric_unique_non_self_candidates() { - let data = [1.0_f32, 0.0, 0.8, 0.2, 0.0, 1.0, -1.0, 0.0]; - let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; - - for metric in [ + #[rstest] + fn candidate_graph_is_symmetric_unique_and_non_self( + #[values( Metric::L2, Metric::Cosine, Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); - for (source, neighbors) in graph.iter().enumerate() { - assert!(neighbors.iter().all(|&target| target as usize != source)); - assert!( - neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32)) - ); - assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); - } + Metric::InnerProduct + )] + metric: Metric, + ) { + // Given + let diagonal = std::f32::consts::FRAC_1_SQRT_2; + let unit_vectors = [1.0_f32, 0.0, diagonal, diagonal, 0.0, 1.0, -1.0, 0.0]; + let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; + + // When + let graph = + build_candidate_graph(matrix_view(&unit_vectors, 4, 2), &leaves, 2, metric).unwrap(); + + // Then + for (source, neighbors) in graph.iter().enumerate() { + assert!(neighbors.iter().all(|&target| target as usize != source)); + assert!( + neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + ); + assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); } } #[test] fn parallel_leaf_schedule_does_not_change_candidate_order() { let data: Vec = (0..64).map(|value| value as f32).collect(); - let leaves: Vec> = (0..32) - .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) - .collect(); - let expected_candidate_order = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + let leaves = circular_overlapping_leaves(64, 32, 16); + let expected_candidate_order = + build_candidate_graph(matrix_view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); for _ in 0..8 { - let actual_candidate_order = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + let actual_candidate_order = + build_candidate_graph(matrix_view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); assert_eq!(actual_candidate_order, expected_candidate_order); } } #[test] - fn singleton_and_zero_k_leaves_add_no_candidates() { - let data = [0.0_f32, 1.0, 2.0]; - let singleton = build( - view(&data, 3, 1), - &[vec![0], vec![1], vec![2]], - 1, - Metric::L2, - ) - .unwrap(); - let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); - assert!( - singleton - .iter() - .chain(&zero_k) - .all(|candidates| candidates.is_empty()) + fn singleton_leaves_add_no_candidates() { + // Given + let point_values = [0.0_f32, 1.0, 2.0]; + let singleton_leaves = [vec![0], vec![1], vec![2]]; + let expected_adjacency: [Vec; 3] = [vec![], vec![], vec![]]; + + // When + let actual_adjacency = adjacency_lists( + build_candidate_graph( + matrix_view(&point_values, 3, 1), + &singleton_leaves, + 1, + Metric::L2, + ) + .unwrap(), ); + + // Then + assert_eq!(actual_adjacency, expected_adjacency); + } + + #[test] + fn zero_k_adds_no_candidates() { + // Given + let point_values = [0.0_f32, 1.0, 2.0]; + let leaves = [vec![0, 1, 2]]; + let expected_adjacency: [Vec; 3] = [vec![], vec![], vec![]]; + + // When + let actual_adjacency = adjacency_lists( + build_candidate_graph(matrix_view(&point_values, 3, 1), &leaves, 0, Metric::L2) + .unwrap(), + ); + + // Then + assert_eq!(actual_adjacency, expected_adjacency); } #[test] diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 395f71ede1..4d07f5cede 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -541,6 +541,7 @@ mod tests { use diskann_utils::views::{Matrix, MatrixView}; use diskann_vector::{Half, distance::Metric}; use diskann_wide::arch::{self, Target1}; + use rstest::rstest; use super::*; @@ -583,7 +584,12 @@ mod tests { arch::dispatch1_no_features(DispatchPartition(metric), PartitionCall { data, config }) } - fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { + fn partition_config( + c_min: usize, + c_max: usize, + fanout: Vec, + replicas: usize, + ) -> PiPNNConfig { PiPNNConfig { c_max, c_min, @@ -594,7 +600,12 @@ mod tests { } } - fn clustered_data(points: usize, dimensions: usize) -> Matrix { + fn separated_point_clusters(points: usize, dimensions: usize) -> Matrix { + const POINTS_PER_CLUSTER: usize = 8; + const CLUSTER_SEPARATION: f32 = 10.0; + const POINT_OFFSET: f32 = 0.001; + const DIMENSION_OFFSET: f32 = 0.01; + Matrix::new( diskann_utils::views::Init({ let mut position = 0usize; @@ -602,7 +613,10 @@ mod tests { let point = position / dimensions; let dimension = position % dimensions; position += 1; - (point / 8) as f32 * 10.0 + dimension as f32 * 0.01 + point as f32 * 0.001 + let cluster = point / POINTS_PER_CLUSTER; + cluster as f32 * CLUSTER_SEPARATION + + point as f32 * POINT_OFFSET + + dimension as f32 * DIMENSION_OFFSET } }), points, @@ -610,7 +624,7 @@ mod tests { ) } - fn directional_data(points: usize, dimensions: usize) -> Matrix { + fn unit_circle_points(points: usize, dimensions: usize) -> Matrix { Matrix::new( diskann_utils::views::Init({ let mut position = 0usize; @@ -644,7 +658,7 @@ mod tests { memberships } - fn assert_valid_partition_with_runtime_metric( + fn assert_partition_invariants( leaves: &[Vec], points: usize, c_max: usize, @@ -669,49 +683,82 @@ mod tests { assert!(counts.iter().all(|&count| count >= replicas)); } + #[rstest] + #[case::below_c_max(7)] + #[case::at_c_max(8)] + fn partition_returns_one_leaf_when_point_count_does_not_exceed_c_max( + #[case] point_count: usize, + ) { + // Given + let data = separated_point_clusters(point_count, 3); + let expected_leaf = vec![(0..point_count as u32).collect::>()]; + + // When + let actual_leaves = partition_with_runtime_metric( + data.as_view(), + &partition_config(2, 8, vec![2], 1), + Metric::L2, + ) + .unwrap(); + + // Then + assert_eq!(actual_leaves, expected_leaf); + } + #[test] - fn partition_returns_one_leaf_when_point_count_does_not_exceed_c_max() { - for points in [7, 8] { - let data = clustered_data(points, 3); - let leaves = partition_with_runtime_metric( - data.as_view(), - &config(2, 8, vec![2], 1), - Metric::L2, - ) - .unwrap(); - assert_eq!(leaves, vec![(0..points as u32).collect::>()]); - } + fn partition_membership_is_deterministic_for_a_fixed_seed() { + // Given + let data = separated_point_clusters(96, 8); + let config = partition_config(4, 16, vec![3, 2], 2); + + // When + let first_partition = + partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + let second_partition = + partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + + // Then + assert_eq!( + sorted_memberships(&first_partition), + sorted_memberships(&second_partition) + ); } #[test] - fn partition_is_fixed_seed_deterministic_and_bounded() { - let data = clustered_data(96, 8); - let config = config(4, 16, vec![3, 2], 2); + fn partition_respects_capacity_and_replica_coverage() { + // Given + let data = separated_point_clusters(96, 8); + let config = partition_config(4, 16, vec![3, 2], 2); - let first = partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); - let second = partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + // When + let partition = partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); - assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); - assert_valid_partition_with_runtime_metric(&first, 96, 16, 2); - assert!(first.iter().map(Vec::len).sum::() > 96 * 2); + // Then + assert_partition_invariants(&partition, 96, 16, 2); } #[test] fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { - let data = clustered_data(80, 4); - let leaves = - partition_with_runtime_metric(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2) - .unwrap(); + let data = separated_point_clusters(80, 4); + let leaves = partition_with_runtime_metric( + data.as_view(), + &partition_config(2, 8, vec![2], 1), + Metric::L2, + ) + .unwrap(); - assert_valid_partition_with_runtime_metric(&leaves, 80, 8, 1); + assert_partition_invariants(&leaves, 80, 8, 1); } #[test] fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { let data = Matrix::new(1.0f32, 24, 4); - let error = - partition_with_runtime_metric(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2) - .unwrap_err(); + let error = partition_with_runtime_metric( + data.as_view(), + &partition_config(2, 4, vec![1], 1), + Metric::L2, + ) + .unwrap_err(); let error = error.downcast::().unwrap(); assert!(matches!( @@ -766,15 +813,15 @@ mod tests { #[test] fn replicas_cover_every_point_once_or_more_per_replica() { - let data = directional_data(72, 5); + let data = unit_circle_points(72, 5); let leaves = partition_with_runtime_metric( data.as_view(), - &config(3, 12, vec![3, 2], 3), + &partition_config(3, 12, vec![3, 2], 3), Metric::CosineNormalized, ) .unwrap(); - assert_valid_partition_with_runtime_metric(&leaves, 72, 12, 3); + assert_partition_invariants(&leaves, 72, 12, 3); } fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) @@ -789,12 +836,12 @@ mod tests { .map(|index| { let point = index / dimensions; let dimension = index % dimensions; - ((point * 5 + dimension * 7 + point * dimension) % 23) as u8 + (point + dimension) as u8 }) .collect(); let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); let converted: Vec = raw.iter().copied().map(&convert).collect(); - let config = config(2, 16, vec![2, 1], 1); + let config = partition_config(2, 16, vec![2, 1], 1); let expected_f32_partition = partition_with_runtime_metric( MatrixView::try_from(&f32_data, points, dimensions).unwrap(), &config, @@ -808,7 +855,7 @@ mod tests { ) .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); - assert_valid_partition_with_runtime_metric(&actual_converted_partition, points, 16, 1); + assert_partition_invariants(&actual_converted_partition, points, 16, 1); assert_eq!( sorted_memberships(&actual_converted_partition), sorted_memberships(&expected_f32_partition), @@ -834,8 +881,8 @@ mod tests { } #[test] - fn l2_leader_norms_preserve_sequential_reduction_order() { - fn next(state: &mut u64) -> f32 { + fn leader_assignment_preserves_sequential_norm_reduction_with_l2() { + fn next_reassociation_regression_value(state: &mut u64) -> f32 { *state ^= *state << 13; *state ^= *state >> 7; *state ^= *state << 17; @@ -845,19 +892,27 @@ mod tests { // Sequential and SIMD-reassociated leader norms select different top-1 // leaders for this case. Dot products still use the production GEMM. The // test changes only the leader-norm reduction. - let dimensions = 129; - let mut state = 0x3a85_f952_c718_6e49; - let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); - let leader_zero: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); - let leader_one: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + const REDUCTION_BOUNDARY_DIMENSIONS: usize = 129; + const REASSOCIATION_REGRESSION_SEED: u64 = 0x3a85_f952_c718_6e49; + let mut state = REASSOCIATION_REGRESSION_SEED; + let point: Vec = (0..REDUCTION_BOUNDARY_DIMENSIONS) + .map(|_| next_reassociation_regression_value(&mut state)) + .collect(); + let leader_zero: Vec = (0..REDUCTION_BOUNDARY_DIMENSIONS) + .map(|_| next_reassociation_regression_value(&mut state)) + .collect(); + let leader_one: Vec = (0..REDUCTION_BOUNDARY_DIMENSIONS) + .map(|_| next_reassociation_regression_value(&mut state)) + .collect(); let data: Vec = leader_zero .into_iter() .chain(leader_one) .chain(point) .collect(); - let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); + let data = MatrixView::try_from(data.as_slice(), 3, REDUCTION_BOUNDARY_DIMENSIONS).unwrap(); + let expected_clusters = [vec![], vec![2]]; - let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( + let actual_clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, &[2], @@ -867,36 +922,50 @@ mod tests { ) .unwrap(); - assert_eq!(clusters, [vec![], vec![2]]); + assert_eq!(actual_clusters, expected_clusters); } - #[test] - fn every_metric_covers_all_points_without_exceeding_c_max() { - let data = directional_data(64, 8); - let config = config(2, 20, vec![2], 1); - - for metric in [ + #[rstest] + fn partition_covers_all_points_without_exceeding_c_max( + #[values( Metric::L2, Metric::Cosine, Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let leaves = partition_with_runtime_metric(data.as_view(), &config, metric).unwrap(); - assert_valid_partition_with_runtime_metric(&leaves, 64, 20, 1); - } + Metric::InnerProduct + )] + metric: Metric, + ) { + // Given + let data = unit_circle_points(64, 8); + let config = partition_config(2, 20, vec![2], 1); + + // When + let leaves = partition_with_runtime_metric(data.as_view(), &config, metric).unwrap(); + + // Then + assert_partition_invariants(&leaves, 64, 20, 1); } - #[test] - fn sampled_leader_count_is_at_least_one_and_never_exceeds_the_cap() { - assert_eq!(sampled_leader_count(1, 1.0), 1); - assert_eq!(sampled_leader_count(10, 0.01), 2); - assert_eq!(sampled_leader_count(50_000, 1.0), LEADER_CAP); + #[rstest] + #[case::single_point_minimum(1, 1.0, 1)] + #[case::sampled_count(10, 0.01, 2)] + #[case::leader_cap(50_000, 1.0, LEADER_CAP)] + fn sampled_leader_count_respects_data_size_sampling_and_cap( + #[case] point_count: usize, + #[case] sampling_probability: f64, + #[case] expected_leader_count: usize, + ) { + assert_eq!( + sampled_leader_count(point_count, sampling_probability), + expected_leader_count + ); } - #[test] - fn replica_seed_derivation_is_stable_and_distinct() { - assert_eq!(replica_seed(0), 1_000); - assert_eq!(replica_seed(1), 8_919); + #[rstest] + #[case::first_replica(0, 1_000)] + #[case::second_replica(1, 8_919)] + fn replica_seed_is_stable(#[case] replica: usize, #[case] expected_seed: u64) { + assert_eq!(replica_seed(replica), expected_seed); } #[test] @@ -937,15 +1006,34 @@ mod tests { #[test] fn parallel_scatter_matches_serial_order() { // Given + let leader_count = 7; + let leaders_per_point = 2; + let second_leader_offset = 3; let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); let assignments: Vec = points .iter() - .flat_map(|point| [point % 7, (point + 3) % 7]) + .flat_map(|point| { + [ + point % leader_count, + (point + second_leader_offset) % leader_count, + ] + }) .collect(); // When - let expected_serial_clusters = scatter_serial(&points, &assignments, 2, 7); - let actual_parallel_clusters = scatter_assignments(&points, &assignments, 2, 7).unwrap(); + let expected_serial_clusters = scatter_serial( + &points, + &assignments, + leaders_per_point, + leader_count as usize, + ); + let actual_parallel_clusters = scatter_assignments( + &points, + &assignments, + leaders_per_point, + leader_count as usize, + ) + .unwrap(); // Then assert_eq!(actual_parallel_clusters, expected_serial_clusters); From 49717983d4c44df4277a32cb1ba0c99d49212cd5 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 24 Aug 2026 11:54:59 +0000 Subject: [PATCH 162/229] test(pipnn): name zero-k fixture inputs --- diskann/src/graph/pipnn/leaf_build.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index d680deb7e2..6c1f2c822e 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -642,13 +642,21 @@ mod tests { fn zero_k_adds_no_candidates() { // Given let point_values = [0.0_f32, 1.0, 2.0]; + let point_count = 3; + let dimensions = 1; + let zero_k = 0; let leaves = [vec![0, 1, 2]]; let expected_adjacency: [Vec; 3] = [vec![], vec![], vec![]]; // When let actual_adjacency = adjacency_lists( - build_candidate_graph(matrix_view(&point_values, 3, 1), &leaves, 0, Metric::L2) - .unwrap(), + build_candidate_graph( + matrix_view(&point_values, point_count, dimensions), + &leaves, + zero_k, + Metric::L2, + ) + .unwrap(), ); // Then From a101b56ae29adf529ace1fff0d9e2b84f6be4518 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 4 Sep 2026 06:04:26 +0000 Subject: [PATCH 163/229] refactor(pipnn): adopt metric-owned buffers --- diskann/src/graph/pipnn/finalization.rs | 18 +++-- diskann/src/graph/pipnn/leaf_build.rs | 46 +++--------- diskann/src/graph/pipnn/mod.rs | 95 +++++++++++++++++++++---- diskann/src/graph/pipnn/partitioning.rs | 25 ++++--- 4 files changed, 114 insertions(+), 70 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index d67fb17a3c..d4ec402774 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -42,7 +42,7 @@ pub(crate) enum FinalizationError { #[derive(Default)] struct PruneWorkspace { sorted_candidates: Vec>, - sorted_cache: Vec<(f32, Option)>, + sorted_cache: Vec>>, prune_states: Vec, } @@ -95,25 +95,23 @@ where max: u16::MAX as usize, })); } - workspace.sorted_cache.clear(); // Sort all candidates before the code marks a self-edge as absent. - // Thus, self-edge removal cannot add a farther candidate. Cache - // construction preserves this order for RobustPrune. + // Thus, self-edge removal cannot add a farther candidate. The + // mapped list keeps the same positions for RobustPrune states. let sorted = SortedNeighbors::new(&mut workspace.sorted_candidates, candidate_count); - workspace.sorted_cache.extend(sorted.iter().map(|neighbor| { - let id = *neighbor.id(); - (*neighbor.distance(), (id != source_id).then_some(id)) - })); + let sorted_cache = sorted.map_in(&mut workspace.sorted_cache, |id| { + (*id != source_id).then_some(*id) + }); workspace .prune_states - .resize(workspace.sorted_cache.len(), prune::State::default()); + .resize(sorted_cache.len(), prune::State::default()); // Each candidate list starts a separate RobustPrune state machine. // Reset retained entries because resize initializes only new entries. workspace.prune_states.fill(prune::State::default()); let selected = prune::robust_prune( - &workspace.sorted_cache, + sorted_cache, workspace.prune_states.as_mut_slice(), degree, graph.alpha(), diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6c1f2c822e..eb6f7690bc 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -9,7 +9,7 @@ //! job does these steps: //! //! 1. Gather each ID and convert its vector to reusable `f32` storage. -//! 2. Call the leaf kernel for Gram construction, norms, and local ranking. +//! 2. Call the leaf kernel for ranking-distance construction and local ranking. //! 3. Convert local positions to global point IDs. //! 4. Add both edge directions to global candidate lists. //! @@ -24,9 +24,10 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use rayon::prelude::*; use super::{ - kernel_metric::LeafMetric, - leaf_kernel::{LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, select_leaf_neighbors}, + leaf_kernel::{LeafKernelWorkspace, leaf_neighbor_count, select_leaf_neighbors}, + leaf_metric::LeafMetric, simd::PiPNNSIMDSchema, + topk::Candidate, }; /// Failure while converting leaves into direct graph candidates. @@ -62,7 +63,7 @@ pub(crate) enum LeafBuildError { #[derive(Default)] struct LeafBuffers { point_values: Vec, - neighbors: Vec, + neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, } @@ -94,7 +95,7 @@ impl LeafBuffers { })?; grow(&mut self.point_values, point_value_count, 0.0); - grow(&mut self.neighbors, neighbor_count, LeafNeighbor::default()); + grow(&mut self.neighbors, neighbor_count, Candidate::default()); Ok((leaf_k, neighbor_count)) } @@ -253,7 +254,7 @@ where fn add_symmetric_neighbors( point_ids: &[u32], leaf_k: usize, - neighbors: &[LeafNeighbor], + neighbors: &[Candidate], local_adjacency: &mut [Vec], ) { for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { @@ -261,7 +262,7 @@ fn add_symmetric_neighbors( if !neighbor.is_assigned() { continue; } - let target = neighbor.target as usize; + let target = neighbor.local_idx as usize; let source_id = point_ids[source]; let target_id = point_ids[target]; if source_id != target_id { @@ -327,7 +328,7 @@ mod tests { arch: A, call: LeafBuildCall<'_, T>, ) -> Result>, LeafBuildError> { - use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use super::super::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { Metric::L2 => { @@ -638,31 +639,6 @@ mod tests { assert_eq!(actual_adjacency, expected_adjacency); } - #[test] - fn zero_k_adds_no_candidates() { - // Given - let point_values = [0.0_f32, 1.0, 2.0]; - let point_count = 3; - let dimensions = 1; - let zero_k = 0; - let leaves = [vec![0, 1, 2]]; - let expected_adjacency: [Vec; 3] = [vec![], vec![], vec![]]; - - // When - let actual_adjacency = adjacency_lists( - build_candidate_graph( - matrix_view(&point_values, point_count, dimensions), - &leaves, - zero_k, - Metric::L2, - ) - .unwrap(), - ); - - // Then - assert_eq!(actual_adjacency, expected_adjacency); - } - #[test] fn leaf_buffer_preparation_reports_shape_overflow_before_allocating() { let mut buffers = LeafBuffers::default(); @@ -679,8 +655,8 @@ mod tests { &[7, 7], 1, &[ - super::super::leaf_kernel::LeafNeighbor::new(1, 0.0), - super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), + super::super::topk::Candidate::new(1, 0.0), + super::super::topk::Candidate::new(0, 0.0), ], &mut graph, ); diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index d55f90eb9d..73324c286f 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -11,14 +11,14 @@ //! //! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has //! at most `c_max` points. -//! 2. `leaf_build` computes a lower-triangular Gram matrix for each leaf. It +//! 2. `leaf_build` computes one ranking-distance buffer for each leaf. It //! selects local neighbors and merges their global point IDs. //! 3. `finalization` applies Vamana RobustPrune to each candidate list that is //! longer than the graph degree. //! -//! `diskann-wide` selects architecture `A`. One match selects metric marker `M`. -//! The build passes both concrete types through all replicas, recursive -//! partitions, stripes, and leaves. The numerical loops do not dispatch again. +//! `diskann-wide` selects architecture `A` for the ranking kernels. One match +//! selects metric marker `M`. Metric computation stays architecture-neutral. +//! The numerical loops do not dispatch again. //! //! [`PiPNNConfig`] contains partition and local-neighbor parameters. //! [`PiPNNBuildContext`] borrows graph policy and a Rayon pool. [`build_graph`] @@ -31,15 +31,14 @@ //! Partition and leaf work use separate reusable buffers. The build consumes //! each output before it creates another graph representation. -mod kernel_metric; -mod simd; - mod finalization; mod leaf_build; mod leaf_kernel; +mod leaf_metric; mod partition_kernel; +mod partition_metric; mod partitioning; -#[allow(dead_code)] +mod simd; mod topk; use crate::{ @@ -52,10 +51,82 @@ use diskann_vector::distance::Metric; use diskann_wide::arch::{self, Target1}; use rayon::ThreadPool; -use self::{ - kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2, LeafMetric, PartitionMetric}, - simd::PiPNNSIMDSchema, -}; +use self::{leaf_metric::LeafMetric, partition_metric::PartitionMetric, simd::PiPNNSIMDSchema}; + +pub(super) struct L2; +pub(super) struct Cosine; +pub(super) struct CosineNormalized; +pub(super) struct InnerProduct; + +/// Convert one dot product and two norms to cosine distance. +/// +/// Treat a zero or subnormal norm as zero similarity. This rule takes precedence +/// over the dot value. Clamp finite similarity to the cosine range. Otherwise, +/// a NaN input produces a NaN distance. +#[inline(always)] +fn cosine_distance(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 / (source_norm * target_norm)).clamp(-1.0, 1.0) + } +} + +#[cfg(test)] +mod cosine_distance_tests { + use super::cosine_distance; + use rstest::rstest; + + #[rstest] + #[case::zero_source(0.0, 0.0, 1.0)] + #[case::zero_target(0.0, 1.0, 0.0)] + #[case::subnormal_source(f32::MIN_POSITIVE.sqrt() / 2.0, f32::MIN_POSITIVE.sqrt() / 2.0, 1.0)] + #[case::subnormal_target(f32::MIN_POSITIVE.sqrt() / 2.0, 1.0, f32::MIN_POSITIVE.sqrt() / 2.0)] + #[case::zero_norm_before_nan_dot(f32::NAN, 0.0, 1.0)] + fn small_norm_produces_unit_distance( + #[case] dot: f32, + #[case] source_norm: f32, + #[case] target_norm: f32, + ) { + assert_eq!(cosine_distance(dot, source_norm, target_norm), 1.0); + } + + #[test] + fn minimum_normal_norm_uses_normalized_similarity() { + let source_norm = f32::MIN_POSITIVE.sqrt(); + let expected_similarity = 0.5; + let dot = expected_similarity * source_norm; + + assert_eq!( + cosine_distance(dot, source_norm, 1.0), + 1.0 - expected_similarity + ); + } + + #[rstest] + #[case::above_one(1.0)] + #[case::below_negative_one(-1.0)] + fn finite_similarity_outside_the_cosine_range_is_clamped(#[case] bounded_similarity: f32) { + let norm = 2.0; + let norm_product = norm * norm; + let rounding_excess = f32::EPSILON * norm_product; + let dot = bounded_similarity * (norm_product + rounding_excess); + + assert_eq!(cosine_distance(dot, norm, norm), 1.0 - bounded_similarity); + } + + #[rstest] + #[case::nan_dot(f32::NAN, 1.0, 1.0)] + #[case::nan_source_norm(0.0, f32::NAN, 1.0)] + #[case::nan_target_norm(0.0, 1.0, f32::NAN)] + fn nan_without_a_zero_norm_produces_nan_distance( + #[case] dot: f32, + #[case] source_norm: f32, + #[case] target_norm: f32, + ) { + assert!(cosine_distance(dot, source_norm, target_norm).is_nan()); + } +} /// PiPNN partition and leaf-selection policy. /// diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 4d07f5cede..86a885baad 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -28,10 +28,8 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::PartitionMetric, - partition_kernel::{ - PartitionKernelWorkspace, PreparedLeaders, UNASSIGNED_LEADER, assign_leaders, - }, + partition_kernel::{PartitionKernelWorkspace, UNASSIGNED_LEADER, assign_leaders}, + partition_metric::PartitionMetric, simd::PiPNNSIMDSchema, }; @@ -280,12 +278,13 @@ where let leader_matrix = MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count) .map_err(|error| ANNError::new(error.as_static()))?; - let leaders = PreparedLeaders::::new(leader_matrix); + let leaders = M::create_leaders(leader_matrix); + let leader_count = M::leader_count(&leaders); - let fanout = fanout.min(leaders.len()); + let fanout = fanout.min(leader_count); let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; let mut assignments = vec![0u32; assignment_len]; - let stripe_points = assignment_stripe_point_count(leaders.len()); + let stripe_points = assignment_stripe_point_count(leader_count); let stripe_assignment_count = checked_area("assignment stripe", stripe_points, fanout)?; let stripe_count = point_ids.len().div_ceil(stripe_points); let worker_stripe_count = stripe_count.div_ceil(rayon::current_num_threads().max(1)); @@ -320,20 +319,20 @@ where Ok::<(), ANNError>(()) })?; - scatter_assignments(point_ids, &assignments, fanout, leader_ids.len()) + scatter_assignments(point_ids, &assignments, fanout, leader_count) } /// Assign one point stripe to sampled partition centers. /// /// The function gathers point IDs into a packed `f32` matrix. The partition -/// kernel owns dot products, point norms, and ranking. This function writes the +/// kernel owns ranking-distance construction and ranking. This function writes the /// returned leader-column IDs for partition scatter. #[inline] fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, point_ids: &[u32], - leaders: &PreparedLeaders<'_, M>, + leaders: &M::Leaders<'_>, fanout: usize, buffers: &mut StripeBuffers, assignments: &mut [u32], @@ -558,7 +557,7 @@ mod tests { T: VectorRepr + Send + Sync, { fn run(self, arch: A, call: PartitionCall<'_, T>) -> ANNResult>> { - use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use super::super::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { Metric::L2 => partition::(arch, call.data, call.config), @@ -912,7 +911,7 @@ mod tests { let data = MatrixView::try_from(data.as_slice(), 3, REDUCTION_BOUNDARY_DIMENSIONS).unwrap(); let expected_clusters = [vec![], vec![2]]; - let actual_clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( + let actual_clusters = assign_to_leaders::<_, super::super::L2, _>( diskann_wide::ARCH, data, &[2], @@ -975,7 +974,7 @@ mod tests { let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); let point_ids: Vec = (0..points as u32).collect(); - let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( + let clusters = assign_to_leaders::<_, super::super::L2, _>( diskann_wide::ARCH, data, &point_ids, From 2c26d8b6ea7efa8f62a6df7601bffdc86855e030 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 8 Sep 2026 07:50:35 +0000 Subject: [PATCH 164/229] docs(pipnn): clarify ranking contracts --- diskann/src/graph/pipnn/leaf_kernel.rs | 19 +++++++------------ diskann/src/graph/pipnn/leaf_metric.rs | 5 +++++ diskann/src/graph/pipnn/partition_kernel.rs | 14 +++++++------- diskann/src/graph/pipnn/partition_metric.rs | 4 ++++ 4 files changed, 23 insertions(+), 19 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 662a37cfdc..dd2333a506 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -5,18 +5,8 @@ //! Leaf-local top-k selection from packed `f32` point vectors. //! -//! The metric fills the lower triangle of a distance matrix. The kernel -//! reads each strict-lower point pair once and updates both points. -//! -//! The output is an `n × k` matrix of sorted [`Candidate`] values. Each target -//! is a position in the leaf. -//! -//! Equal distances can select either candidate. The kernel does not rank NaN. -//! An unfilled output slot contains [`Candidate::default`]. All supported -//! metrics use the same SIMD-group and single-value traversal. -//! -//! The caller supplies concrete architecture `A` and metric `M`. -//! [`LeafKernelWorkspace`] stores reusable numerical scratch. +//! [`select_leaf_neighbors`] asks the metric to fill a lower-triangle ranking +//! buffer. It scans each point pair once and updates both points' neighbor lists. use crate::{ANNError, ANNResult}; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -57,6 +47,11 @@ pub(super) fn leaf_neighbor_count(points: usize, requested_k: usize) -> usize { /// Compute local nearest neighbors for one packed leaf matrix. /// +/// `output` contains one row per point, ordered by increasing ranking distance. +/// Candidate IDs are positions in the leaf. Equal distances can select either +/// candidate. NaN and positive infinity are not retained; unfilled slots contain +/// [`Candidate::default`]. +/// /// # Errors /// /// Returns an error for invalid linear-algebra input or output width. diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index 9a3d73bcef..280fec7f0c 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -21,6 +21,11 @@ use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; pub(super) trait LeafMetric: Send + Sync + 'static { /// Compute ranking distances for all unordered point pairs. /// + /// Values preserve nearest-first order, but need not equal metric distances. + /// L2 returns squared distances. Normalized cosine and inner product return + /// `-dot`, omitting the constant in normalized cosine's `1 - dot`. + /// Cosine returns `1 - similarity`. + /// /// `storage` has `points.nrows() * points.nrows()` elements. fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()>; } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 00d62cbfac..75656a1071 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -5,13 +5,9 @@ //! Select partition centers for PiPNN point assignment. //! -//! A leader is a sampled dataset point that represents one child partition. -//! The metric builds final point-to-leader distances. The kernel returns nearest -//! leader-column IDs for partition scatter. -//! -//! L2 omits the assigned point's norm because it is constant across all sampled -//! leaders. Equal distances can select either leader. NaN is not rankable. An -//! unfilled output slot contains [`UNASSIGNED_LEADER`]. +//! [`assign_leaders`] computes point-to-leader ranking values and selects the +//! nearest leader-column IDs for partition scatter. Each sampled leader +//! represents one child partition. use crate::{ANNError, ANNResult}; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -34,6 +30,10 @@ pub(super) struct PartitionKernelWorkspace { /// Assign one packed point stripe to metric-owned partition leaders. /// +/// Each output row contains leader-column IDs ordered by increasing ranking +/// value. Equal values can select either leader. NaN and positive infinity are +/// not retained. +/// /// A point can have fewer assignments than the output width. Each remaining /// slot contains [`UNASSIGNED_LEADER`]. /// diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 70b2552013..fd898f0548 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -42,6 +42,10 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { /// Compute one row-major point-to-leader ranking buffer. /// + /// Values preserve nearest-first order. L2 omits the point's squared norm, + /// which is constant across its row. Normalized cosine and inner product + /// return `-dot`; cosine returns `1 - similarity`. + /// /// `storage` has `points.nrows() * leader_count` elements. /// A zero distance can have either sign. Equal distances can select either leader. fn compute_distances( From 00b8a5cdd84f48838682d521a22dcbcde9a18b68 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 8 Sep 2026 08:06:47 +0000 Subject: [PATCH 165/229] fix(pipnn): address kernel review feedback --- .github/workflows/ci.yml | 1 + diskann/src/graph/pipnn/leaf_kernel.rs | 53 ++++++- diskann/src/graph/pipnn/leaf_metric.rs | 13 +- diskann/src/graph/pipnn/mod.rs | 26 ++-- diskann/src/graph/pipnn/partition_kernel.rs | 52 ++++++- diskann/src/graph/pipnn/partition_metric.rs | 164 ++++++++++++-------- 6 files changed, 202 insertions(+), 107 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f857d27a17..f8b585fcc4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -520,6 +520,7 @@ jobs: cargo llvm-cov nextest --locked \ --config "$RUST_CONFIG" \ --workspace \ + --features diskann/pipnn \ --lcov --output-path lcov.info - name: Generate miri code coverage diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index dd2333a506..6f4a8c5f11 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -24,9 +24,12 @@ pub(super) struct LeafKernelWorkspace { worst: Vec, } -/// Invalid output width for leaf-neighbor selection. +/// Invalid output shape for leaf-neighbor selection. #[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] pub(super) enum LeafKernelError { + /// The output must contain one row per input point. + #[error("invalid leaf output row count {rows} for {points} points")] + InvalidOutputRows { points: usize, rows: usize }, /// A source requests more neighbors than the leaf has other points. #[error("invalid leaf neighbor count {neighbors} for {points} points; maximum is {maximum}")] InvalidNeighborCount { @@ -54,8 +57,8 @@ pub(super) fn leaf_neighbor_count(points: usize, requested_k: usize) -> usize { /// /// # Errors /// -/// Returns an error for invalid linear-algebra input or output width. -/// Invalid output widths leave the output and workspace unchanged. +/// Returns an error for invalid linear-algebra input or output shape. +/// Invalid output shapes leave the output and workspace unchanged. pub(super) fn select_leaf_neighbors( arch: A, points: MatrixView<'_, f32>, @@ -67,7 +70,7 @@ where M: LeafMetric, { let point_count = points.nrows(); - validate_neighbor_count(point_count, &output).map_err(ANNError::new)?; + validate_output(point_count, &output).map_err(ANNError::new)?; let distance_count = point_count * point_count; let LeafKernelWorkspace { distance_scratch, @@ -80,8 +83,7 @@ where &mut distance_scratch[..distance_count], point_count, point_count, - ) - .map_err(|error| ANNError::new(error.as_static()))?; + )?; M::compute_distances(points, distances.as_mut_slice())?; rank_leaf_distances(arch, distances.as_view(), output, worst); Ok(()) @@ -114,11 +116,17 @@ fn rank_leaf_distances( }); } -/// Check the output width against the number of non-self points. -fn validate_neighbor_count( +/// Check for one output row per point and a valid non-self neighbor count. +fn validate_output( point_count: usize, output: &MutMatrixView<'_, Candidate>, ) -> Result<(), LeafKernelError> { + if output.nrows() != point_count { + return Err(LeafKernelError::InvalidOutputRows { + points: point_count, + rows: output.nrows(), + }); + } let maximum_neighbors = point_count.saturating_sub(1); let neighbor_count = output.ncols(); if neighbor_count > maximum_neighbors { @@ -244,6 +252,35 @@ mod tests { mod select_leaf_neighbors_tests { use super::*; + #[rstest::rstest] + #[case::missing_row(2)] + #[case::extra_row(4)] + fn invalid_output_rows_leave_buffers_unchanged(#[case] rows: usize) { + let values = [0.0_f32, 1.0, 3.0]; + let expected_output = vec![Candidate::new(42, -1.0); rows]; + let mut output = expected_output.clone(); + let mut workspace = LeafKernelWorkspace { + distance_scratch: vec![99.0], + worst: vec![7.0], + }; + + let error = select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + MatrixView::try_from(&values[..], 3, 1).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), rows, 1).unwrap(), + &mut workspace, + ) + .unwrap_err(); + + assert_eq!( + error.downcast_ref::(), + Some(&LeafKernelError::InvalidOutputRows { points: 3, rows }) + ); + assert_eq!(output, expected_output); + assert_eq!(workspace.distance_scratch, [99.0]); + assert_eq!(workspace.worst, [7.0]); + } + #[test] fn invalid_neighbor_width_leaves_buffers_unchanged() { // Given diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index 280fec7f0c..f6b2c9cdf6 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -86,12 +86,10 @@ impl LeafMetric for Cosine { } } -impl LeafMetric for CosineNormalized { +impl LeafMetric for InnerProduct { fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { - let point_count = points.nrows(); - // The constant in `1 - dot` does not change nearest-first order. diskann_linalg::sgemm_aat_lower( - point_count, + points.nrows(), points.ncols(), -1.0, points.as_slice(), @@ -102,15 +100,14 @@ impl LeafMetric for CosineNormalized { } } -impl LeafMetric for InnerProduct { +impl LeafMetric for CosineNormalized { fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { - // Both metrics rank with `-dot`. Their graph-pruning policies stay separate. - CosineNormalized::compute_distances(points, storage) + // The constant in `1 - dot` does not change nearest-first order. + InnerProduct::compute_distances(points, storage) } } #[cfg(test)] -#[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] mod tests { use super::*; use rstest::rstest; diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index e91ed5fff8..8892c4b333 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -61,26 +61,18 @@ mod cosine_distance_tests { #[test] fn minimum_normal_norm_uses_normalized_similarity() { - let source_norm = f32::MIN_POSITIVE.sqrt(); - let expected_similarity = 0.5; - let dot = expected_similarity * source_norm; - - assert_eq!( - cosine_distance(dot, source_norm, 1.0), - 1.0 - expected_similarity - ); + let norm = f32::MIN_POSITIVE.sqrt(); + assert_eq!(cosine_distance(0.5 * norm, norm, 1.0), 0.5); } #[rstest] - #[case::above_one(1.0)] - #[case::below_negative_one(-1.0)] - fn finite_similarity_outside_the_cosine_range_is_clamped(#[case] bounded_similarity: f32) { - let norm = 2.0; - let norm_product = norm * norm; - let rounding_excess = f32::EPSILON * norm_product; - let dot = bounded_similarity * (norm_product + rounding_excess); - - assert_eq!(cosine_distance(dot, norm, norm), 1.0 - bounded_similarity); + #[case::above_one(1.0 + f32::EPSILON, 0.0)] + #[case::below_negative_one(-1.0 - f32::EPSILON, 2.0)] + fn finite_similarity_outside_the_cosine_range_is_clamped( + #[case] similarity: f32, + #[case] expected: f32, + ) { + assert_eq!(cosine_distance(similarity, 1.0, 1.0), expected); } #[rstest] diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 75656a1071..d0a15f4cdd 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -39,7 +39,8 @@ pub(super) struct PartitionKernelWorkspace { /// /// # Errors /// -/// Returns an error for invalid GEMM input. +/// Returns an error for invalid GEMM input or an output row count different from +/// the point count. A row-count mismatch leaves output and workspace unchanged. pub(super) fn assign_leaders( arch: A, points: MatrixView<'_, f32>, @@ -52,6 +53,12 @@ where M: PartitionMetric, { let point_count = points.nrows(); + if output.nrows() != point_count { + return Err(ANNError::message(format!( + "invalid partition output row count {} for {point_count} points", + output.nrows() + ))); + } let leader_count = M::leader_count(leaders); let distance_count = point_count * leader_count; let PartitionKernelWorkspace { @@ -61,14 +68,13 @@ where if distance_scratch.len() < distance_count { distance_scratch.resize(distance_count, 0.0); } - M::compute_distances(points, leaders, &mut distance_scratch[..distance_count])?; - let distances = MatrixView::try_from( - &distance_scratch[..distance_count], + let mut distances = MutMatrixView::try_from( + &mut distance_scratch[..distance_count], point_count, leader_count, - ) - .map_err(|error| ANNError::new(error.as_static()))?; - rank_leader_distances(arch, distances, output, ranked_leader_scratch); + )?; + M::compute_distances(points, leaders, distances.as_mut_view())?; + rank_leader_distances(arch, distances.as_view(), output, ranked_leader_scratch); Ok(()) } @@ -157,6 +163,38 @@ mod tests { mod assign_leaders_tests { use super::*; + #[rstest::rstest] + #[case::missing_row(2)] + #[case::extra_row(4)] + fn invalid_output_rows_leave_buffers_unchanged(#[case] rows: usize) { + let values = [1.0_f32, 2.0, 3.0]; + let points = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let leaders = Cosine::create_leaders(points); + let expected_output = vec![42; rows]; + let mut output = expected_output.clone(); + let expected_ranked = [Candidate::new(7, -1.0)]; + let mut workspace = PartitionKernelWorkspace { + distance_scratch: vec![99.0], + ranked_leader_scratch: expected_ranked.to_vec(), + }; + + let error = assign_leaders::<_, Cosine>( + diskann_wide::ARCH, + points, + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), rows, 1).unwrap(), + &mut workspace, + ) + .unwrap_err(); + + assert!(error.to_string().contains(&format!( + "invalid partition output row count {rows} for 3 points" + ))); + assert_eq!(output, expected_output); + assert_eq!(workspace.distance_scratch, [99.0]); + assert_eq!(workspace.ranked_leader_scratch, expected_ranked); + } + #[test] fn cosine_assignment_reuses_workspace_and_output_for_a_different_point() { // Given diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index fd898f0548..0933d22220 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -10,7 +10,7 @@ use crate::{ANNError, ANNResult}; use diskann_linalg::Transpose; -use diskann_utils::views::MatrixView; +use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::{Norm, norm::FastL2NormSquared}; use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; @@ -24,7 +24,7 @@ pub(super) struct PartitionLeaders<'a, Norms> { norms: Norms, } -/// Fill one flattened point-to-leader ranking buffer. +/// Fill one point-to-leader ranking matrix. /// /// The associated leader type hides metric data from the caller. The caller /// creates one value and shares it across all point stripes. @@ -46,12 +46,12 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { /// which is constant across its row. Normalized cosine and inner product /// return `-dot`; cosine returns `1 - similarity`. /// - /// `storage` has `points.nrows() * leader_count` elements. + /// `storage` has `points.nrows()` rows and `leader_count` columns. /// A zero distance can have either sign. Equal distances can select either leader. fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, - storage: &mut [f32], + storage: MutMatrixView<'_, f32>, ) -> ANNResult<()>; } @@ -90,18 +90,15 @@ impl PartitionMetric for L2 { fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, - storage: &mut [f32], + mut storage: MutMatrixView<'_, f32>, ) -> ANNResult<()> { + if storage.nrows() != points.nrows() || storage.ncols() != leaders.values.nrows() { + return Err(ANNError::message("point-to-leader output shape mismatch")); + } // The point norm is constant across a point row. It cannot change ranking. - let leader_norms = &leaders.norms; // Initialize each point row before GEMM adds the dot-product term. - let leader_count = leader_norms.len(); - for row in storage.chunks_exact_mut(leader_count) { - let mut leader = 0; - while leader < leader_count { - row[leader] = leader_norms[leader]; - leader += 1; - } + for row in storage.row_iter_mut() { + row.copy_from_slice(&leaders.norms); } diskann_linalg::sgemm( Transpose::None, @@ -113,7 +110,7 @@ impl PartitionMetric for L2 { points.as_slice(), leaders.values.as_slice(), Some(1.0), - storage, + storage.as_mut_slice(), ) .map_err(ANNError::new)?; Ok(()) @@ -137,8 +134,11 @@ impl PartitionMetric for Cosine { fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, - storage: &mut [f32], + mut storage: MutMatrixView<'_, f32>, ) -> ANNResult<()> { + if storage.nrows() != points.nrows() || storage.ncols() != leaders.values.nrows() { + return Err(ANNError::message("point-to-leader output shape mismatch")); + } diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -149,17 +149,13 @@ impl PartitionMetric for Cosine { points.as_slice(), leaders.values.as_slice(), None, - storage, + storage.as_mut_slice(), ) .map_err(ANNError::new)?; let point_norms = cosine_norms(points); let leader_norms = &leaders.norms; - let leader_count = leaders.values.nrows(); // Convert each dot to cosine distance. Reuse leader norms across stripes. - for (row, &point_norm) in storage - .chunks_exact_mut(leader_count) - .zip(point_norms.iter()) - { + for (row, &point_norm) in storage.row_iter_mut().zip(point_norms.iter()) { for (distance, &leader_norm) in row.iter_mut().zip(leader_norms.iter()) { *distance = cosine_distance(*distance, point_norm, leader_norm); } @@ -168,9 +164,7 @@ impl PartitionMetric for Cosine { } } -// Normalized cosine and inner product have the same ranking expression. -// Both metrics rank candidates with `-dot`. -impl PartitionMetric for CosineNormalized { +impl PartitionMetric for InnerProduct { type Leaders<'a> = PartitionLeaders<'a, ()>; fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { @@ -184,8 +178,11 @@ impl PartitionMetric for CosineNormalized { fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, - storage: &mut [f32], + mut storage: MutMatrixView<'_, f32>, ) -> ANNResult<()> { + if storage.nrows() != points.nrows() || storage.ncols() != leaders.values.nrows() { + return Err(ANNError::message("point-to-leader output shape mismatch")); + } diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -196,35 +193,35 @@ impl PartitionMetric for CosineNormalized { points.as_slice(), leaders.values.as_slice(), None, - storage, + storage.as_mut_slice(), ) .map_err(ANNError::new)?; Ok(()) } } -impl PartitionMetric for InnerProduct { - type Leaders<'a> = ::Leaders<'a>; +impl PartitionMetric for CosineNormalized { + type Leaders<'a> = ::Leaders<'a>; fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { - CosineNormalized::create_leaders(values) + InnerProduct::create_leaders(values) } fn leader_count(leaders: &Self::Leaders<'_>) -> usize { - CosineNormalized::leader_count(leaders) + InnerProduct::leader_count(leaders) } fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, - storage: &mut [f32], + storage: MutMatrixView<'_, f32>, ) -> ANNResult<()> { - CosineNormalized::compute_distances(points, leaders, storage) + // The constant in `1 - dot` does not change nearest-first order. + InnerProduct::compute_distances(points, leaders, storage) } } #[cfg(test)] -#[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] mod tests { use super::*; @@ -243,7 +240,12 @@ mod tests { let leaders = M::create_leaders(matrix(&leader, 1)); let mut storage = [STALE_DISTANCE]; - M::compute_distances(matrix(&point, 1), &leaders, &mut storage).unwrap(); + M::compute_distances( + matrix(&point, 1), + &leaders, + MutMatrixView::row_vector(&mut storage[..]), + ) + .unwrap(); storage[0] } @@ -271,20 +273,57 @@ mod tests { mod compute_distances_tests { use super::*; + #[rstest::rstest] + #[case::wrong_rows_with_equal_area(1, 4)] + #[case::wrong_columns(2, 1)] + fn output_shape_mismatch_leaves_distances_unchanged( + #[values(L2, Cosine, CosineNormalized, InnerProduct)] _metric: M, + #[case] rows: usize, + #[case] columns: usize, + ) { + // Given: two points and two leaders require a 2-by-2 output. + let leader_values = [1.0, 0.0, 0.0, 2.0]; + let leaders = M::create_leaders(matrix(&leader_values, 2)); + let points = matrix(&[0.0, 4.0, -2.0, 0.0], 2); + let expected = vec![STALE_DISTANCE; rows * columns]; + let mut storage = expected.clone(); + + // When + let error = M::compute_distances( + points, + &leaders, + MutMatrixView::try_from(storage.as_mut_slice(), rows, columns).unwrap(), + ) + .unwrap_err(); + + // Then: a shape error is returned before any output is written. + assert!( + error + .to_string() + .contains("point-to-leader output shape mismatch") + ); + assert_eq!(storage, expected); + } + #[test] fn squared_l2_ranking_omits_the_point_norm() { - // Given - let point = [0.0_f32, 4.0]; - let leader = [3.0_f32, 4.0]; - let leader_squared_norm = leader[0].mul_add(leader[0], leader[1] * leader[1]); - let dot = point[0].mul_add(leader[0], point[1] * leader[1]); - let expected = (-2.0_f32).mul_add(dot, leader_squared_norm); + // Given: leader norms are 25 and 4; each entry is norm - 2 * dot. + let leader_values = [3.0_f32, 4.0, 0.0, 2.0]; + let leaders = L2::create_leaders(matrix(&leader_values, 2)); + let points = matrix(&[2.0, 1.0, -1.0, 2.0], 2); + let expected = [5.0, 0.0, 15.0, -4.0]; + let mut storage = [STALE_DISTANCE; 4]; // When - let actual = compute_one_ranking::(point, leader); - - // Then - assert_eq!(actual, expected); + L2::compute_distances( + points, + &leaders, + MutMatrixView::try_from(&mut storage[..], 2, 2).unwrap(), + ) + .unwrap(); + + // Then: every row replaces stale distances before GEMM accumulates. + assert_eq!(storage, expected); } #[test] @@ -376,33 +415,24 @@ mod tests { } #[test] - fn l2_reinitializes_every_output_row_for_a_new_point_stripe() { - let leader_values = [3.0_f32, 4.0, 0.0, 2.0]; - let leaders = L2::create_leaders(matrix(&leader_values, 2)); - let mut output = [STALE_DISTANCE; 4]; - - // Leader norms are 25 and 4; each entry is norm - 2 * dot. - L2::compute_distances(matrix(&[1.0, 0.0, 0.0, 1.0], 2), &leaders, &mut output).unwrap(); - assert_eq!(output, [19.0, 4.0, 17.0, 0.0]); - - L2::compute_distances(matrix(&[2.0, 1.0, -1.0, 2.0], 2), &leaders, &mut output) - .unwrap(); - assert_eq!(output, [5.0, 0.0, 15.0, -4.0]); - } - - #[test] - fn cosine_overwrites_output_for_a_new_point_stripe() { + fn cosine_overwrites_every_point_row_with_normalized_distances() { + // Given: points and leaders lie on signed coordinate axes. let leader_values = [1.0_f32, 0.0, 0.0, 2.0]; let leaders = Cosine::create_leaders(matrix(&leader_values, 2)); - let mut output = [STALE_DISTANCE; 4]; + let points = matrix(&[0.0, 4.0, -2.0, 0.0], 2); + let expected = [1.0, 0.0, 2.0, 1.0]; + let mut storage = [STALE_DISTANCE; 4]; - Cosine::compute_distances(matrix(&[1.0, 0.0, 0.0, 3.0], 2), &leaders, &mut output) - .unwrap(); - assert_eq!(output, [0.0, 1.0, 1.0, 0.0]); + // When + Cosine::compute_distances( + points, + &leaders, + MutMatrixView::try_from(&mut storage[..], 2, 2).unwrap(), + ) + .unwrap(); - Cosine::compute_distances(matrix(&[0.0, 4.0, -2.0, 0.0], 2), &leaders, &mut output) - .unwrap(); - assert_eq!(output, [1.0, 0.0, 2.0, 1.0]); + // Then + assert_eq!(storage, expected); } } } From 927ceb6c10b8b60d8cf431e26a51dc34a5e87fe9 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 8 Sep 2026 09:16:58 +0000 Subject: [PATCH 166/229] refactor(pipnn): move distance blocks to topk --- diskann/src/graph/pipnn/leaf_kernel.rs | 4 +- diskann/src/graph/pipnn/partition_kernel.rs | 4 +- diskann/src/graph/pipnn/simd.rs | 45 ------------------- diskann/src/graph/pipnn/topk.rs | 50 ++++++++++++++++++++- 4 files changed, 52 insertions(+), 51 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 6f4a8c5f11..9c4becdd90 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -13,8 +13,8 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ leaf_metric::LeafMetric, - simd::{PiPNNSIMDSchema, distance_blocks}, - topk::{Candidate, with_topk_rows}, + simd::PiPNNSIMDSchema, + topk::{Candidate, distance_blocks, with_topk_rows}, }; /// Reusable storage for one leaf numerical pipeline. diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index d0a15f4cdd..e603e89f06 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -14,8 +14,8 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ partition_metric::PartitionMetric, - simd::{PiPNNSIMDSchema, distance_blocks}, - topk::{Candidate, UNASSIGNED, with_topk_rows}, + simd::PiPNNSIMDSchema, + topk::{Candidate, UNASSIGNED, distance_blocks, with_topk_rows}, }; /// No sampled partition center was rankable for this output slot. diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 964b65dbd8..b8fad0d746 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -43,48 +43,3 @@ where { type Vector = DefaultVector; } - -/// One group of distances and their column indexes in the supplied slice. -#[derive(Clone, Copy)] -pub(super) enum DistanceBlock<'a, F: PiPNNSIMDVector> { - Simd { - first_idx: usize, - values: F, - lanes: &'a [f32], - }, - Scalar { - idx: usize, - distance: f32, - }, -} - -/// Iterate over distance groups without a callback on the ranking hot path. -/// -/// Scalar lanes borrow the input slice, so sharing a block across updates needs -/// no temporary array. Tail entries follow in column order. The caller keeps the -/// selected architecture's execution scope active. -#[inline(always)] -pub(super) fn distance_blocks( - arch: A, - distances: &[f32], -) -> impl Iterator> { - let simd_end = distances.len() - distances.len() % A::Vector::LANES; - distances[..simd_end] - .chunks_exact(A::Vector::LANES) - .enumerate() - .map(move |(group, lanes)| DistanceBlock::Simd { - first_idx: group * A::Vector::LANES, - // SAFETY: chunks_exact yields one complete SIMD group. - values: unsafe { A::Vector::load_simd(arch, lanes.as_ptr()) }, - lanes, - }) - .chain( - distances[simd_end..] - .iter() - .enumerate() - .map(move |(tail, &distance)| DistanceBlock::Scalar { - idx: simd_end + tail, - distance, - }), - ) -} diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 51f437739f..447f927f06 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -7,8 +7,54 @@ //! Updates allocate no storage. NaN and positive infinity are not retained. use diskann_utils::views::MutMatrixView; +use diskann_wide::SIMDVector; + +use super::simd::{PiPNNSIMDSchema, PiPNNSIMDVector}; + +/// One group of distances and their column indexes in the supplied slice. +#[derive(Clone, Copy)] +pub(super) enum DistanceBlock<'a, F: PiPNNSIMDVector> { + Simd { + first_idx: usize, + values: F, + lanes: &'a [f32], + }, + Scalar { + idx: usize, + distance: f32, + }, +} -use super::simd::{DistanceBlock, PiPNNSIMDVector}; +/// Iterate over distance groups without a callback on the ranking hot path. +/// +/// Scalar lanes borrow the input slice, so sharing a block across updates needs +/// no temporary array. Tail entries follow in column order. The caller keeps the +/// selected architecture's execution scope active. +#[inline(always)] +pub(super) fn distance_blocks( + arch: A, + distances: &[f32], +) -> impl Iterator> { + let simd_end = distances.len() - distances.len() % A::Vector::LANES; + distances[..simd_end] + .chunks_exact(A::Vector::LANES) + .enumerate() + .map(move |(group, lanes)| DistanceBlock::Simd { + first_idx: group * A::Vector::LANES, + // SAFETY: chunks_exact yields one complete SIMD group. + values: unsafe { A::Vector::load_simd(arch, lanes.as_ptr()) }, + lanes, + }) + .chain( + distances[simd_end..] + .iter() + .enumerate() + .map(move |(tail, &distance)| DistanceBlock::Scalar { + idx: simd_end + tail, + distance, + }), + ) +} /// An output slot with no assigned candidate. pub(super) const UNASSIGNED: u32 = u32::MAX; @@ -289,7 +335,7 @@ mod tests { mod update_tests { use super::*; - use crate::graph::pipnn::simd::distance_blocks; + use crate::graph::pipnn::topk::distance_blocks; use rstest::rstest; #[test] From 5ef1097e1f2791c30ee74cf3ab69bbda4986f6df Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 8 Sep 2026 09:23:47 +0000 Subject: [PATCH 167/229] test(pipnn): simplify partition row rejection --- diskann/src/graph/pipnn/partition_kernel.rs | 27 ++++++--------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index e603e89f06..2ba2c86d27 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -166,33 +166,22 @@ mod tests { #[rstest::rstest] #[case::missing_row(2)] #[case::extra_row(4)] - fn invalid_output_rows_leave_buffers_unchanged(#[case] rows: usize) { + fn invalid_output_rows_return_error(#[case] rows: usize) { let values = [1.0_f32, 2.0, 3.0]; let points = MatrixView::try_from(&values[..], 3, 1).unwrap(); let leaders = Cosine::create_leaders(points); - let expected_output = vec![42; rows]; - let mut output = expected_output.clone(); - let expected_ranked = [Candidate::new(7, -1.0)]; - let mut workspace = PartitionKernelWorkspace { - distance_scratch: vec![99.0], - ranked_leader_scratch: expected_ranked.to_vec(), - }; - - let error = assign_leaders::<_, Cosine>( + let mut output = vec![UNASSIGNED_LEADER; rows]; + let mut workspace = PartitionKernelWorkspace::default(); + + let result = assign_leaders::<_, Cosine>( diskann_wide::ARCH, points, &leaders, MutMatrixView::try_from(output.as_mut_slice(), rows, 1).unwrap(), &mut workspace, - ) - .unwrap_err(); - - assert!(error.to_string().contains(&format!( - "invalid partition output row count {rows} for 3 points" - ))); - assert_eq!(output, expected_output); - assert_eq!(workspace.distance_scratch, [99.0]); - assert_eq!(workspace.ranked_leader_scratch, expected_ranked); + ); + + assert!(result.is_err()); } #[test] From f803283ed229ed16b6864e94d1148608df08257b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 8 Sep 2026 09:29:23 +0000 Subject: [PATCH 168/229] refactor(pipnn): merge SIMD schema traits --- diskann/src/graph/pipnn/simd.rs | 35 +++++++++++++-------------------- diskann/src/graph/pipnn/topk.rs | 29 ++++++++++++++------------- 2 files changed, 29 insertions(+), 35 deletions(-) diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index b8fad0d746..0cc6b9678f 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -5,41 +5,34 @@ //! SIMD schema for PiPNN numerical kernels. -use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDVector, SupportedLaneCount}; +use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDVector}; /// Default SIMD representation used by both PiPNN ranking kernels. /// /// This alias is the single build-time width selection. type DefaultVector = ::f32x16; -/// Operations required by PiPNN SIMD vectors. -pub(super) trait PiPNNSIMDVector: SIMDVector + SIMDFloat { - /// Return one bit for each selected lane. - fn active_lanes(mask: Self::Mask) -> u64; -} - -impl PiPNNSIMDVector for F -where - F: SIMDVector> + SIMDFloat, - Const: SupportedLaneCount, - u64: From<<::BitMask as SIMDMask>::Underlying>, -{ - #[inline(always)] - fn active_lanes(mask: Self::Mask) -> u64 { - u64::from(mask.bitmask().to_underlying()) - } -} - /// PiPNN SIMD representation for one architecture. pub(super) trait PiPNNSIMDSchema: Architecture { /// SIMD vector used by both ranking kernels. - type Vector: PiPNNSIMDVector; + type Vector: SIMDVector + SIMDFloat; + + /// Return one bit for each selected lane. + fn active_lanes(mask: ::Mask) -> u64; } impl PiPNNSIMDSchema for A where A: Architecture, - DefaultVector: PiPNNSIMDVector, + DefaultVector: SIMDVector + SIMDFloat, + u64: From< + << as SIMDVector>::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying, + >, { type Vector = DefaultVector; + + #[inline(always)] + fn active_lanes(mask: ::Mask) -> u64 { + u64::from(mask.bitmask().to_underlying()) + } } diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 447f927f06..6ac712522d 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -7,16 +7,16 @@ //! Updates allocate no storage. NaN and positive infinity are not retained. use diskann_utils::views::MutMatrixView; -use diskann_wide::SIMDVector; +use diskann_wide::{SIMDPartialOrd, SIMDVector}; -use super::simd::{PiPNNSIMDSchema, PiPNNSIMDVector}; +use super::simd::PiPNNSIMDSchema; /// One group of distances and their column indexes in the supplied slice. #[derive(Clone, Copy)] -pub(super) enum DistanceBlock<'a, F: PiPNNSIMDVector> { +pub(super) enum DistanceBlock<'a, A: PiPNNSIMDSchema> { Simd { first_idx: usize, - values: F, + values: A::Vector, lanes: &'a [f32], }, Scalar { @@ -34,12 +34,12 @@ pub(super) enum DistanceBlock<'a, F: PiPNNSIMDVector> { pub(super) fn distance_blocks( arch: A, distances: &[f32], -) -> impl Iterator> { +) -> impl Iterator> { let simd_end = distances.len() - distances.len() % A::Vector::LANES; distances[..simd_end] .chunks_exact(A::Vector::LANES) .enumerate() - .map(move |(group, lanes)| DistanceBlock::Simd { + .map(move |(group, lanes)| DistanceBlock::::Simd { first_idx: group * A::Vector::LANES, // SAFETY: chunks_exact yields one complete SIMD group. values: unsafe { A::Vector::load_simd(arch, lanes.as_ptr()) }, @@ -173,10 +173,10 @@ impl<'a, const K: usize> TopKRows<'a, K> { /// Offer the block's candidates to one row using their slice-column indexes. #[inline(always)] - pub(super) fn update_one( + pub(super) fn update_one( &mut self, row_idx: usize, - block: &DistanceBlock<'_, F>, + block: &DistanceBlock<'_, A>, ) { match *block { DistanceBlock::Scalar { idx, distance } => { @@ -188,7 +188,8 @@ impl<'a, const K: usize> TopKRows<'a, K> { lanes, } => { let mut worst = self.worst[row_idx]; - let mut eligible = F::active_lanes(values.lt_simd(F::splat(values.arch(), worst))); + let mut eligible = + A::active_lanes(values.lt_simd(A::Vector::splat(values.arch(), worst))); if eligible == 0 { return; } @@ -211,10 +212,10 @@ impl<'a, const K: usize> TopKRows<'a, K> { /// Offer one candidate to each row identified by the block's slice-column indexes. #[inline(always)] - pub(super) fn update_many( + pub(super) fn update_many( &mut self, candidate_idx: u32, - block: &DistanceBlock<'_, F>, + block: &DistanceBlock<'_, A>, ) { match *block { DistanceBlock::Scalar { idx, distance } => { @@ -225,10 +226,10 @@ impl<'a, const K: usize> TopKRows<'a, K> { values, lanes, } => { - let thresholds = &mut self.worst[first_idx..][..F::LANES]; + let thresholds = &mut self.worst[first_idx..][..A::Vector::LANES]; // SAFETY: the slice above has one threshold for every SIMD lane. - let worst = unsafe { F::load_simd(values.arch(), thresholds.as_ptr()) }; - let mut eligible = F::active_lanes(values.lt_simd(worst)); + let worst = unsafe { A::Vector::load_simd(values.arch(), thresholds.as_ptr()) }; + let mut eligible = A::active_lanes(values.lt_simd(worst)); while eligible != 0 { let lane = eligible.trailing_zeros() as usize; eligible &= eligible - 1; From 303ff6d7eaa9376c33f93e9c50be7e0682359dc8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:47:32 +0000 Subject: [PATCH 169/229] refactor(pipnn): encapsulate top-k updates Keep width dispatch and SIMD distance blocks inside TopK. Reuse the same row storage for replacement and dual updates, and retain the measured unfilled-result path. Consolidate tests around reachable kernel behavior, independent output oracles, and workspace reuse. --- diskann/src/graph/pipnn/leaf_kernel.rs | 356 ++------ diskann/src/graph/pipnn/partition_kernel.rs | 277 +++--- diskann/src/graph/pipnn/topk.rs | 941 +++++++++++++------- 3 files changed, 782 insertions(+), 792 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 9c4becdd90..e50d0da729 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -14,7 +14,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ leaf_metric::LeafMetric, simd::PiPNNSIMDSchema, - topk::{Candidate, distance_blocks, with_topk_rows}, + topk::{Candidate, with_topk}, }; /// Reusable storage for one leaf numerical pipeline. @@ -89,28 +89,18 @@ where Ok(()) } -/// Traverse the strict lower triangle and update each pair's two candidate rows. -fn rank_leaf_distances( +/// Offer each lower-triangle row to the two endpoint neighbor lists. +fn rank_leaf_distances( arch: A, distances: MatrixView<'_, f32>, output: MutMatrixView<'_, Candidate>, worst: &mut Vec, -) where - A: PiPNNSIMDSchema, -{ - if output.ncols() == 0 { - return; - } - worst.resize(distances.nrows(), f32::INFINITY); - with_topk_rows!(output, worst, |topks| { - // Rayon outlines leaf workers. Reapply target features before the SIMD scan. - arch.run(move || { - for source_idx in 1..distances.nrows() { - let source_distances = &distances.row(source_idx)[..source_idx]; - for block in distance_blocks(arch, source_distances) { - topks.update_one(source_idx, &block); - topks.update_many(source_idx as u32, &block); - } +) { + worst.resize(output.nrows(), f32::INFINITY); + arch.run(move || { + with_topk!(output, worst.as_mut_slice(), |topk| { + for point_idx in 1..distances.nrows() { + topk.update_dual_topk(arch, point_idx, &distances.row(point_idx)[..point_idx]); } }); }); @@ -143,93 +133,39 @@ fn validate_output( mod tests { use super::*; use crate::graph::pipnn::L2; - use diskann_utils::views::{MatrixView, MutMatrixView}; mod test_support { - use std::cmp::Ordering; - use super::*; use diskann_wide::arch::{self, Target1}; struct KernelCall<'a> { distances: MatrixView<'a, f32>, output: MutMatrixView<'a, Candidate>, - worst: &'a mut Vec, } struct RankDistances; - impl Target1> for RankDistances - where - A: PiPNNSIMDSchema, - { + impl Target1> for RankDistances { fn run(self, arch: A, call: KernelCall<'_>) { - rank_leaf_distances(arch, call.distances, call.output, call.worst); + rank_leaf_distances(arch, call.distances, call.output, &mut Vec::new()); } } pub(super) fn rank_distance_fixture( distances: &[f32], points: usize, - output_width: usize, + width: usize, ) -> Vec { - let mut output = vec![Candidate::default(); points * output_width]; + let mut output = vec![Candidate::default(); points * width]; arch::dispatch1_no_features( RankDistances, KernelCall { distances: MatrixView::try_from(distances, points, points).unwrap(), - output: MutMatrixView::try_from(output.as_mut_slice(), points, output_width) - .unwrap(), - worst: &mut Vec::new(), + output: MutMatrixView::try_from(output.as_mut_slice(), points, width).unwrap(), }, ); output } - - pub(super) fn reference_neighbors( - distances: &[f32], - points: usize, - width: usize, - ) -> Vec { - let mut output = vec![Candidate::default(); points * width]; - for source in 0..points { - let mut candidates = Vec::with_capacity(points.saturating_sub(1)); - for target in 0..points { - if source == target { - continue; - } - let row = source.max(target); - let column = source.min(target); - let distance = distances[row * points + column]; - if distance < f32::INFINITY { - candidates.push(Candidate::new(target as u32, distance)); - } - } - candidates.sort_unstable_by(|left, right| { - left.distance - .partial_cmp(&right.distance) - .unwrap_or(Ordering::Equal) - }); - let retained = candidates.len().min(width); - output[source * width..source * width + retained] - .copy_from_slice(&candidates[..retained]); - } - output - } - - /// Give each pair a unique distance, with alternating signs in each row. - pub(super) fn lower_triangle_distances(points: usize) -> Vec { - // NaN in unused entries detects reads of the diagonal or upper triangle. - let mut distances = vec![f32::NAN; points * points]; - for source in 1..points { - for target in 0..source { - let distance = (source * points + target + 1) as f32; - distances[source * points + target] = - if target % 2 == 0 { -distance } else { distance }; - } - } - distances - } } mod leaf_neighbor_count_tests { @@ -255,14 +191,10 @@ mod tests { #[rstest::rstest] #[case::missing_row(2)] #[case::extra_row(4)] - fn invalid_output_rows_leave_buffers_unchanged(#[case] rows: usize) { + fn invalid_output_rows_return_an_error(#[case] rows: usize) { let values = [0.0_f32, 1.0, 3.0]; - let expected_output = vec![Candidate::new(42, -1.0); rows]; - let mut output = expected_output.clone(); - let mut workspace = LeafKernelWorkspace { - distance_scratch: vec![99.0], - worst: vec![7.0], - }; + let mut output = vec![Candidate::default(); rows]; + let mut workspace = LeafKernelWorkspace::default(); let error = select_leaf_neighbors::<_, L2>( diskann_wide::ARCH, @@ -276,58 +208,35 @@ mod tests { error.downcast_ref::(), Some(&LeafKernelError::InvalidOutputRows { points: 3, rows }) ); - assert_eq!(output, expected_output); - assert_eq!(workspace.distance_scratch, [99.0]); - assert_eq!(workspace.worst, [7.0]); } #[test] - fn invalid_neighbor_width_leaves_buffers_unchanged() { - // Given + fn requesting_self_as_an_extra_neighbor_returns_an_error() { let values = [0.0_f32, 1.0, 3.0]; - let point_count = values.len(); - let invalid_width = point_count; - let points = MatrixView::try_from(&values[..], point_count, 1).unwrap(); - let expected_output = [Candidate::default(); 9]; - // A smaller prior leaf forces scratch growth if validation runs too late. - let expected_distances = [99.0; 4]; - let expected_thresholds = [7.0; 2]; - let expected_error = LeafKernelError::InvalidNeighborCount { - points: point_count, - neighbors: invalid_width, - maximum: point_count - 1, - }; - let mut output = expected_output; - let mut workspace = LeafKernelWorkspace { - distance_scratch: expected_distances.to_vec(), - worst: expected_thresholds.to_vec(), + let mut output = [Candidate::default(); 9]; + let mut workspace = LeafKernelWorkspace::default(); + let expected = LeafKernelError::InvalidNeighborCount { + points: 3, + neighbors: 3, + maximum: 2, }; - // When let error = select_leaf_neighbors::<_, L2>( diskann_wide::ARCH, - points, - MutMatrixView::try_from(&mut output[..], point_count, invalid_width).unwrap(), + MatrixView::try_from(&values[..], 3, 1).unwrap(), + MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), &mut workspace, ) .unwrap_err(); - // Then - assert_eq!( - error.downcast_ref::(), - Some(&expected_error) - ); - assert_eq!(output, expected_output); - assert_eq!(workspace.distance_scratch, expected_distances); - assert_eq!(workspace.worst, expected_thresholds); + assert_eq!(error.downcast_ref::(), Some(&expected)); } #[test] - fn l2_selection_reuses_output_and_scratch_for_a_smaller_leaf() { - // Given + fn l2_neighbors_follow_new_geometry_when_workspace_is_reused() { + // Given: squared distances on the line determine each point's two neighbors. let values = [0.0_f32, 1.0, 3.0, 10.0]; - let points = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let expected_neighbors = [ + let expected = [ Candidate::new(1, (values[0] - values[1]).powi(2)), Candidate::new(2, (values[0] - values[2]).powi(2)), Candidate::new(0, (values[1] - values[0]).powi(2)), @@ -337,30 +246,30 @@ mod tests { Candidate::new(2, (values[3] - values[2]).powi(2)), Candidate::new(1, (values[3] - values[1]).powi(2)), ]; - let mut actual_neighbors = [Candidate::default(); 8]; + let mut output = [Candidate::default(); 8]; let mut workspace = LeafKernelWorkspace::default(); - // When + // When: first populate output and scratch with a valid larger leaf. select_leaf_neighbors::<_, L2>( diskann_wide::ARCH, - points, - MutMatrixView::try_from(&mut actual_neighbors[..], 4, 2).unwrap(), + MatrixView::try_from(&values[..], 4, 1).unwrap(), + MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), &mut workspace, ) .unwrap(); - assert_eq!(actual_neighbors, expected_neighbors); + assert_eq!(output, expected); - // Reuse both buffers with a smaller, different dataset and unfilled slots. + // A smaller leaf has one finite pair and leaves the other slots unassigned. let smaller = [0.0, 2.0, f32::NAN]; select_leaf_neighbors::<_, L2>( diskann_wide::ARCH, MatrixView::try_from(&smaller[..], 3, 1).unwrap(), - MutMatrixView::try_from(&mut actual_neighbors[..6], 3, 2).unwrap(), + MutMatrixView::try_from(&mut output[..6], 3, 2).unwrap(), &mut workspace, ) .unwrap(); assert_eq!( - actual_neighbors[..6], + output[..6], [ Candidate::new(1, 4.0), Candidate::default(), @@ -370,169 +279,64 @@ mod tests { Candidate::default(), ] ); + + // Growing again must replace the NaN scratch and stale rows. Doubling + // every coordinate preserves IDs and multiplies squared distances by four. + let scaled = values.map(|value| 2.0 * value); + let expected = expected + .map(|candidate| Candidate::new(candidate.local_idx, 4.0 * candidate.distance)); + select_leaf_neighbors::<_, L2>( + diskann_wide::ARCH, + MatrixView::try_from(&scaled[..], 4, 1).unwrap(), + MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), + &mut workspace, + ) + .unwrap(); + assert_eq!(output, expected); } } mod rank_leaf_distances_tests { - use super::test_support::*; + use super::test_support::rank_distance_fixture; 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(16, 3)] - #[case::one_complete_lane_fixed_one(17, 1)] - #[case::one_complete_lane_fixed_two(17, 2)] - #[case::one_complete_lane_fixed_three(17, 3)] - #[case::lane_plus_one_runtime_width(18, 4)] - #[case::all_non_self_neighbors(17, 16)] - #[case::two_lanes_minus_one_runtime_width(32, 7)] - #[case::two_complete_lanes_runtime_width(33, 7)] - #[case::two_lanes_plus_one_runtime_width(34, 7)] - #[case::maximum_leaf_size_runtime_width(512, 7)] - #[trace] - fn dispatched_leaf_ranking_matches_scalar_reference_across_lane_boundaries( - #[case] point_count: usize, - #[case] requested_k: usize, - ) { - // Miri covers the pointer boundaries in the smaller lane cases. - if cfg!(miri) && point_count > 64 { - return; - } - - // Given: the last source row has `point_count - 1` distances. - let distances = lower_triangle_distances(point_count); - let expected_neighbors = reference_neighbors(&distances, point_count, requested_k); - - // When - let actual_neighbors = rank_distance_fixture(&distances, point_count, requested_k); - - // Then - assert_eq!(actual_neighbors, expected_neighbors); - } - - #[rstest] - #[case::finite_negative(-f32::EPSILON)] - #[case::negative_infinity(f32::NEG_INFINITY)] - fn negative_distances_remain_rankable(#[case] distance: f32) { - // Given - let distances = [f32::NAN, f32::NAN, distance, f32::NAN]; - let expected = [Candidate::new(1, distance), Candidate::new(0, distance)]; - - // When - let actual = rank_distance_fixture(&distances, 2, 1); - - // Then - assert_eq!(actual, expected); - } #[test] - fn tied_distances_fill_capacity_with_distinct_non_self_neighbors() { - // Given: each point has three equally distant candidates for two slots. - let point_count = 4; - let width = 2; - let distances = [1.0; 16]; - - // When - let actual = rank_distance_fixture(&distances, point_count, width); - - // Then: any two candidates are valid; no tie order is required. - for (source, neighbors) in actual.chunks_exact(width).enumerate() { - assert_eq!(neighbors[0].distance, 1.0); - assert_eq!(neighbors[1].distance, 1.0); - assert_ne!(neighbors[0].local_idx, neighbors[1].local_idx); - for neighbor in neighbors { - assert!(neighbor.local_idx < point_count as u32); - assert_ne!(neighbor.local_idx, source as u32); - } + fn lower_triangle_pairs_update_both_endpoints_across_a_simd_group_and_tail() { + // Given: only three lower-triangle pairs rank. The diagonal and upper + // triangle contain better distances, so reading either changes the answer. + let mut distances = [f32::NEG_INFINITY; 18 * 18]; + let mut matrix = MutMatrixView::try_from(&mut distances[..], 18, 18).unwrap(); + for source in 1..18 { + matrix.row_mut(source)[..source].fill(f32::INFINITY); } + matrix.row_mut(16)[0] = 3.0; + matrix.row_mut(17)[1] = 2.0; + matrix.row_mut(17)[16] = 1.0; + let mut expected = [[Candidate::default(); 2]; 18]; + expected[0][0] = Candidate::new(16, 3.0); + expected[1][0] = Candidate::new(17, 2.0); + expected[16] = [Candidate::new(17, 1.0), Candidate::new(0, 3.0)]; + expected[17] = [Candidate::new(16, 1.0), Candidate::new(1, 2.0)]; + + // When: row 16 supplies a full SIMD group; row 17 also supplies a tail. + let actual = rank_distance_fixture(matrix.as_slice(), 18, 2); + + // Then: every reciprocal ID is local to its output row; other rows stay empty. + assert_eq!(actual.as_slice(), expected.as_flattened()); } #[test] - fn f32_max_distance_is_still_a_rankable_neighbor() { - // Given - let point_count = 4; - let width = 3; - let source = 3; - let mut distances = [1.0; 16]; - distances[source * point_count] = f32::MAX; - let expected_last_neighbor = Candidate::new(0, f32::MAX); - - // When - let actual = rank_distance_fixture(&distances, point_count, width); - - // Then - assert_eq!(actual[source * width + width - 1], expected_last_neighbor); - } - - #[rstest] - #[case::nan(f32::NAN)] - #[case::positive_infinity(f32::INFINITY)] - fn non_rankable_distances_leave_neighbor_slots_unassigned(#[case] distance: f32) { - // One leaf exercises scalar rows, a full SIMD group, and its tail. - let point_count = 18; - let distances = vec![distance; point_count * point_count]; - let expected = vec![Candidate::default(); point_count]; - - // When - let actual = rank_distance_fixture(&distances, point_count, 1); - - // Then - assert_eq!(actual, expected); - } - - #[rstest] - #[case::nan(f32::NAN)] - #[case::positive_infinity(f32::INFINITY)] - fn non_rankable_candidates_cannot_replace_finite_neighbors( - #[case] distance: f32, - #[values(17, 18)] point_count: usize, - ) { - // Given: the last candidate is in a full SIMD group or its scalar tail. - let source = point_count - 1; - let invalid_target = source - 1; - let mut distances = vec![1.0; point_count * point_count]; - distances[source * point_count + invalid_target] = distance; - - // When - let actual = rank_distance_fixture(&distances, point_count, 1); - - // Then - assert!(actual[source].local_idx < invalid_target as u32); - assert_eq!(actual[source].distance, 1.0); - } + fn singleton_leaf_has_no_neighbors() { + let actual = rank_distance_fixture(&[f32::NAN], 1, 0); - #[test] - fn scalar_tail_updates_both_endpoints_after_a_rejected_simd_group() { - // Given: the final pair is the only rankable pair in the leaf. - let point_count = 18; - let source = 17; - let target = 16; - let mut distances = vec![f32::INFINITY; point_count * point_count]; - distances[source * point_count + target] = -1.0; - let mut expected = vec![Candidate::default(); point_count]; - expected[source] = Candidate::new(target as u32, -1.0); - expected[target] = Candidate::new(source as u32, -1.0); - - // When - let actual = rank_distance_fixture(&distances, point_count, 1); - - // Then - assert_eq!(actual, expected); + assert!(actual.is_empty()); } #[test] - fn singleton_leaf_has_no_neighbors() { - // Given - let distances = [f32::NAN]; - let expected: [Candidate; 0] = []; - - // When - let actual = rank_distance_fixture(&distances, 1, 0); + fn zero_neighbor_capacity_returns_empty_rows() { + let actual = rank_distance_fixture(&[1.0; 9], 3, 0); - // Then - assert_eq!(actual, expected); + assert!(actual.is_empty()); } } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 2ba2c86d27..e67229d5d7 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -15,7 +15,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ partition_metric::PartitionMetric, simd::PiPNNSIMDSchema, - topk::{Candidate, UNASSIGNED, distance_blocks, with_topk_rows}, + topk::{Candidate, UNASSIGNED, with_topk}, }; /// No sampled partition center was rankable for this output slot. @@ -25,7 +25,7 @@ pub(super) const UNASSIGNED_LEADER: u32 = UNASSIGNED; #[derive(Default)] pub(super) struct PartitionKernelWorkspace { distance_scratch: Vec, - ranked_leader_scratch: Vec, + ranked_leaders: Vec, } /// Assign one packed point stripe to metric-owned partition leaders. @@ -63,7 +63,7 @@ where let distance_count = point_count * leader_count; let PartitionKernelWorkspace { distance_scratch, - ranked_leader_scratch, + ranked_leaders, } = workspace; if distance_scratch.len() < distance_count { distance_scratch.resize(distance_count, 0.0); @@ -74,40 +74,28 @@ where leader_count, )?; M::compute_distances(points, leaders, distances.as_mut_view())?; - rank_leader_distances(arch, distances.as_view(), output, ranked_leader_scratch); + rank_leader_distances(arch, distances.as_view(), output, ranked_leaders); Ok(()) } -/// Rank final point-to-leader distances. -/// -/// Each distance row stores all leader distances for one point. The caller -/// supplies at least one point, one leader, and one output column. -fn rank_leader_distances( +/// Select leader columns from each point's distance row. +fn rank_leader_distances( arch: A, distances: MatrixView<'_, f32>, mut output: MutMatrixView<'_, u32>, - ranked_leaders: &mut Vec, -) where - A: PiPNNSIMDSchema, -{ - let fanout = output.ncols(); - ranked_leaders.resize(fanout, Candidate::default()); - // Rayon outlines stripe workers. Reapply target features before ranking leaders. + candidates: &mut Vec, +) { + candidates.resize(output.ncols(), Candidate::default()); arch.run(move || { let mut worst = [f32::INFINITY]; - with_topk_rows!( - MutMatrixView::row_vector(ranked_leaders.as_mut_slice()), - &mut worst[..], - |topks| { - for (point_distances, point_output) in - distances.row_iter().zip(output.row_iter_mut()) - { - topks.reset(); - for block in distance_blocks(arch, point_distances) { - topks.update_one(0, &block); - } - for (destination, leader) in point_output.iter_mut().zip(topks.row(0)) { - *destination = leader.local_idx; + with_topk!( + MutMatrixView::row_vector(candidates.as_mut_slice()), + &mut worst, + |topk| { + for (distances, output) in distances.row_iter().zip(output.row_iter_mut()) { + topk.replace_topk(arch, 0, distances); + for (destination, candidate) in output.iter_mut().zip(topk.candidates(0)) { + *destination = candidate.local_idx; } } } @@ -118,45 +106,30 @@ fn rank_leader_distances( #[cfg(test)] mod tests { use super::*; - use crate::graph::pipnn::Cosine; - use diskann_utils::views::{Matrix, MatrixView, MutMatrixView}; - - mod test_support { - use super::*; - use diskann_wide::arch::{self, Target1}; - - struct KernelCall<'a> { - distances: MatrixView<'a, f32>, - output: MutMatrixView<'a, u32>, - ranked_leaders: &'a mut Vec, - } - - struct RankDistances; - - impl Target1> for RankDistances - where - A: PiPNNSIMDSchema, - { - fn run(self, arch: A, call: KernelCall<'_>) { - rank_leader_distances(arch, call.distances, call.output, call.ranked_leaders); - } - } - - pub(super) fn rank_distance_fixture( - distances: MatrixView<'_, f32>, - nearest_leader_count: usize, - ) -> Vec { - // A previous valid leader must be overwritten even when no candidate is rankable. - let mut output = Matrix::new(0, distances.nrows(), nearest_leader_count); - arch::dispatch1_no_features( - RankDistances, - KernelCall { - distances, - output: output.as_mut_view(), - ranked_leaders: &mut Vec::new(), - }, - ); - output.into_inner().into_vec() + use crate::graph::pipnn::{Cosine, L2}; + use diskann_utils::views::{MatrixView, MutMatrixView}; + use diskann_wide::arch::{self, Target3}; + + struct AssignL2; + + impl + Target3, MatrixView<'_, f32>, MatrixView<'_, f32>, MutMatrixView<'_, u32>> + for AssignL2 + { + fn run( + self, + arch: A, + points: MatrixView<'_, f32>, + leader_values: MatrixView<'_, f32>, + output: MutMatrixView<'_, u32>, + ) -> ANNResult<()> { + assign_leaders::<_, L2>( + arch, + points, + &L2::create_leaders(leader_values), + output, + &mut PartitionKernelWorkspace::default(), + ) } } @@ -185,16 +158,43 @@ mod tests { } #[test] - fn cosine_assignment_reuses_workspace_and_output_for_a_different_point() { - // Given + fn assignment_projects_each_points_l2_leader_columns_with_simd_and_tail() { + // Given: leader j lies at x=j. Seventeen leaders enter one SIMD block and a tail. + // The nearest two integer positions to 15.75, -0.25, and 7.25 are respectively + // (16, 15), (0, 1), and (7, 8); their squared distances are (0.0625, 0.5625), + // (0.0625, 1.5625), and (0.0625, 0.5625). The last point is unrankable. + let leader_values: [f32; 17] = std::array::from_fn(|leader| leader as f32); + let point_values = [15.75, -0.25, 7.25, f32::NAN]; + let expected = [16, 15, 0, 1, 7, 8, UNASSIGNED_LEADER, UNASSIGNED_LEADER]; + let mut output = [0; 8]; + + // When: dispatch reaches the real metric and kernel on the available architecture. + arch::dispatch3_no_features( + AssignL2, + MatrixView::try_from(&point_values[..], 4, 1).unwrap(), + MatrixView::try_from(&leader_values[..], 17, 1).unwrap(), + MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), + ) + .unwrap(); + + // Then: each point replaces the previous row, including an unrankable final row. + assert_eq!(output, expected); + } + + #[test] + fn assignment_reuses_workspace_after_shrinking_and_regrowing_a_cosine_stripe() { + // Given: unit leaders point right, up, and left. The first point's cosine + // similarities have order right > up > left; the second reverses right and left. let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; let leaders = Cosine::create_leaders(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 expected_reassigned_leaders = [2, 1]; - let mut actual_assignments = [UNASSIGNED_LEADER; 4]; + let expected = [0, 1, 2, 1]; + let expected_smaller_stripe = [2, 1]; + let regrown_point_values = [-0.8, 0.2, 0.9, 0.1]; + let expected_regrown_stripe = [2, 1, 0, 1]; + let mut output = [UNASSIGNED_LEADER; 4]; let mut workspace = PartitionKernelWorkspace::default(); // When @@ -202,133 +202,56 @@ mod tests { diskann_wide::ARCH, points, &leaders, - MutMatrixView::try_from(&mut actual_assignments[..], 2, 2).unwrap(), + MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), &mut workspace, ) .unwrap(); // Then - assert_eq!( - actual_assignments, - expected_leaders_by_descending_cosine_similarity - ); + assert_eq!(output, expected); // When: shrink to the second point, whose nearest leader differs from the old prefix. assign_leaders::<_, Cosine>( diskann_wide::ARCH, MatrixView::try_from(&point_values[2..], 1, 2).unwrap(), &leaders, - MutMatrixView::try_from(&mut actual_assignments[..2], 1, 2).unwrap(), + MutMatrixView::try_from(&mut output[..2], 1, 2).unwrap(), &mut workspace, ) .unwrap(); // Then - assert_eq!(actual_assignments[..2], expected_reassigned_leaders); - } - } - - mod rank_leader_distances_tests { - use super::test_support::*; - use super::*; - use rstest::rstest; - - #[rstest] - #[case::two_leaders_select_one(2, 1)] - #[case::scalar_select_two(7, 2)] - #[case::lane_minus_one(15, 3)] - #[case::one_complete_lane(16, 3)] - #[case::lane_plus_one(17, 4)] - #[case::all_leaders(17, 17)] - #[case::two_lanes_minus_one(31, 7)] - #[case::two_complete_lanes(32, 7)] - #[case::two_lanes_plus_one(33, 7)] - #[case::root_partition(1000, 64)] - #[trace] - fn dispatched_partition_ranking_selects_nearest_columns_across_lane_boundaries( - #[case] leader_count: usize, - #[case] nearest_leader_count: usize, - ) { - // Given: unique distances increase in one row and decrease in the other. - let values: Vec<_> = (0..leader_count) - .chain((0..leader_count).rev()) - .map(|leader| leader as f32 - leader_count as f32 / 2.0) - .collect(); - let distances = MatrixView::try_from(values.as_slice(), 2, leader_count).unwrap(); - let expected: Vec<_> = (0..nearest_leader_count) - .chain((leader_count - nearest_leader_count..leader_count).rev()) - .map(|leader| leader as u32) - .collect(); - - // When - let actual = rank_distance_fixture(distances, nearest_leader_count); - - // Then - assert_eq!(actual, expected); - } - - #[rstest] - #[case::scalar(4)] - #[case::simd_with_tail(17)] - fn nearer_leader_displaces_one_of_the_tied_candidates(#[case] leader_count: usize) { - // Given - let nearest_leader = leader_count - 1; - let mut values = vec![1.0; leader_count]; - values[nearest_leader] = 0.0; - let distances = MatrixView::try_from(values.as_slice(), 1, leader_count).unwrap(); + assert_eq!(output[..2], expected_smaller_stripe); - // When - let actual = rank_distance_fixture(distances, 3); - - // Then: the remaining two slots can contain any distinct tied leaders. - assert_eq!(actual[0], nearest_leader as u32); - assert!(actual[1] < nearest_leader as u32); - assert!(actual[2] < nearest_leader as u32); - assert_ne!(actual[1], actual[2]); - } - - #[rstest] - #[case::nan(f32::NAN)] - #[case::positive_infinity(f32::INFINITY)] - fn non_rankable_distances_leave_only_unfilled_ranks_unassigned(#[case] unrankable: f32) { - // Given: only the first row has rankable candidates, spanning SIMD and scalar tail. - let leader_count = 17; - let mut values = vec![unrankable; 2 * leader_count]; - values[leader_count - 2] = -1.0; - values[leader_count - 1] = 2.0; - let distances = MatrixView::try_from(values.as_slice(), 2, leader_count).unwrap(); - let expected = [ - (leader_count - 2) as u32, - (leader_count - 1) as u32, - UNASSIGNED_LEADER, - UNASSIGNED_LEADER, - UNASSIGNED_LEADER, - UNASSIGNED_LEADER, - ]; - - // When - let actual = rank_distance_fixture(distances, 3); + // When: regrow with reversed rows, so the old tail cannot pass unchanged. + assign_leaders::<_, Cosine>( + diskann_wide::ARCH, + MatrixView::try_from(®rown_point_values[..], 2, 2).unwrap(), + &leaders, + MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), + &mut workspace, + ) + .unwrap(); // Then - assert_eq!(actual, expected); + assert_eq!(output, expected_regrown_stripe); } #[test] - fn negative_infinity_and_f32_max_remain_rankable() { - // Given: the smallest distance is in the scalar tail. - let mut values = [f32::NAN; 17]; - values[0] = f32::MAX; - values[1] = 0.0; - values[2] = -f32::EPSILON; - values[16] = f32::NEG_INFINITY; - let distances = MatrixView::try_from(&values[..], 1, values.len()).unwrap(); - let expected = [16, 2, 1, 0]; - - // When - let actual = rank_distance_fixture(distances, 4); + fn empty_point_stripe_leaves_output_empty() { + let leader_values = [1.0_f32]; + let leaders = L2::create_leaders(MatrixView::row_vector(&leader_values[..])); + let points = MatrixView::try_from(&[][..], 0, 1).unwrap(); + let mut output = []; - // Then - assert_eq!(actual, expected); + assign_leaders::<_, L2>( + diskann_wide::ARCH, + points, + &leaders, + MutMatrixView::try_from(&mut output[..], 0, 1).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap(); } } } diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 6ac712522d..730c516305 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -3,59 +3,14 @@ * Licensed under the MIT license. */ -//! Fixed-capacity candidate rows with cached farthest distances. -//! Updates allocate no storage. NaN and positive infinity are not retained. +//! Top-k updates from caller-selected distance slices. +//! NaN and positive infinity are not retained. use diskann_utils::views::MutMatrixView; use diskann_wide::{SIMDPartialOrd, SIMDVector}; use super::simd::PiPNNSIMDSchema; -/// One group of distances and their column indexes in the supplied slice. -#[derive(Clone, Copy)] -pub(super) enum DistanceBlock<'a, A: PiPNNSIMDSchema> { - Simd { - first_idx: usize, - values: A::Vector, - lanes: &'a [f32], - }, - Scalar { - idx: usize, - distance: f32, - }, -} - -/// Iterate over distance groups without a callback on the ranking hot path. -/// -/// Scalar lanes borrow the input slice, so sharing a block across updates needs -/// no temporary array. Tail entries follow in column order. The caller keeps the -/// selected architecture's execution scope active. -#[inline(always)] -pub(super) fn distance_blocks( - arch: A, - distances: &[f32], -) -> impl Iterator> { - let simd_end = distances.len() - distances.len() % A::Vector::LANES; - distances[..simd_end] - .chunks_exact(A::Vector::LANES) - .enumerate() - .map(move |(group, lanes)| DistanceBlock::::Simd { - first_idx: group * A::Vector::LANES, - // SAFETY: chunks_exact yields one complete SIMD group. - values: unsafe { A::Vector::load_simd(arch, lanes.as_ptr()) }, - lanes, - }) - .chain( - distances[simd_end..] - .iter() - .enumerate() - .map(move |(tail, &distance)| DistanceBlock::Scalar { - idx: simd_end + tail, - distance, - }), - ) -} - /// An output slot with no assigned candidate. pub(super) const UNASSIGNED: u32 = u32::MAX; @@ -85,158 +40,275 @@ impl Default for Candidate { } } -// Private dispatch tag; callers supply only the result shape. +// Fixed widths preserve row lengths for inlined insertion. Zero selects runtime width. pub(super) const RUNTIME_WIDTH: usize = 0; -/// Select the top-k implementation once from the row width. -/// -/// For widths one to three, `row_mut` exposes a row with a compile-time length. -/// After inlining, the compiler can unroll the insertion loop and remove bounds checks. -/// Other widths use the same insertion loop with a runtime row length. -/// -/// Each branch selects a different `TopKRows` type. -/// Use a code block to keep that type through the candidate loops without a callback. -/// Like an ordinary block, `return` and `?` affect the caller. -macro_rules! with_topk_rows { - ($rows:expr, $worst:expr, |$topks:ident| $body:block) => {{ - let rows = $rows; - let worst = $worst; - match rows.ncols() { +/// Select the width specialization once for a batch. +/// Storage expressions are evaluated once. The body has ordinary block control flow. +macro_rules! with_topk { + ($candidates:expr, $thresholds:expr, |$topk:ident| $body:block) => {{ + let candidates = $candidates; + let thresholds = $thresholds; + match candidates.ncols() { 1 => { - let mut $topks = $crate::graph::pipnn::topk::TopKRows::<1>::new(rows, worst); + let mut $topk = $crate::graph::pipnn::topk::TopK::<1>::new(candidates, thresholds); $body } 2 => { - let mut $topks = $crate::graph::pipnn::topk::TopKRows::<2>::new(rows, worst); + let mut $topk = $crate::graph::pipnn::topk::TopK::<2>::new(candidates, thresholds); $body } 3 => { - let mut $topks = $crate::graph::pipnn::topk::TopKRows::<3>::new(rows, worst); + let mut $topk = $crate::graph::pipnn::topk::TopK::<3>::new(candidates, thresholds); + $body + } + 10 => { + let mut $topk = $crate::graph::pipnn::topk::TopK::<10>::new(candidates, thresholds); $body } _ => { - let mut $topks = $crate::graph::pipnn::topk::TopKRows::< + let mut $topk = $crate::graph::pipnn::topk::TopK::< { $crate::graph::pipnn::topk::RUNTIME_WIDTH }, - >::new(rows, worst); + >::new(candidates, thresholds); $body } } }}; } -pub(super) use with_topk_rows; +pub(super) use with_topk; -/// Borrowed top-k results and one contiguous cached threshold per row. +/// Nearest candidates and cached distance limits for one or more result sets. /// -/// Each row has the same positive capacity. Only this view mutates the rows and -/// their thresholds while borrowed. Callers supply each candidate ID at most once -/// per row; updates do not deduplicate. Equal distances need no fixed tie order. -/// The with_topk_rows macro selects the implementation from the result shape. -pub(super) struct TopKRows<'a, const K: usize> { - rows: MutMatrixView<'a, Candidate>, - worst: &'a mut [f32], +/// Each result set has the same capacity. This view borrows the candidate matrix +/// and one threshold per set; kernels provide their storage. Candidate IDs are +/// local indexes defined by the update operation, not dataset IDs. Updates do +/// not deduplicate: callers supply each candidate ID at most once per result set. +/// Equal distances need no fixed tie order. +/// `K` is the result capacity; `RUNTIME_WIDTH` uses the matrix's column count. +/// The construction macro selects the width once per batch. Kernels keep the +/// architecture scope active across their updates. +pub(super) struct TopK<'a, const K: usize> { + candidates: MutMatrixView<'a, Candidate>, + thresholds: &'a mut [f32], } -impl<'a, const K: usize> TopKRows<'a, K> { +impl<'a, const K: usize> TopK<'a, K> { /// Bind and clear caller-owned result and threshold buffers. - pub(super) fn new(rows: MutMatrixView<'a, Candidate>, worst: &'a mut [f32]) -> Self { - debug_assert!(rows.ncols() > 0, "top-k row capacity must be positive"); - debug_assert!( - K == RUNTIME_WIDTH || K == rows.ncols(), - "top-k compile-time width must match its rows" + /// A fixed `K` must equal the output width. + pub(super) fn new(candidates: MutMatrixView<'a, Candidate>, thresholds: &'a mut [f32]) -> Self { + assert!( + K == RUNTIME_WIDTH || K == candidates.ncols(), + "top-k width must match its output" ); debug_assert_eq!( - rows.nrows(), - worst.len(), - "top-k rows and thresholds must match" + candidates.nrows(), + thresholds.len(), + "top-k candidates and thresholds must match" ); - let mut topk = Self { rows, worst }; + let mut topk = Self { + candidates, + thresholds, + }; topk.reset(); topk } - pub(super) fn reset(&mut self) { - self.rows.as_mut_slice().fill(Candidate::default()); - self.worst.fill(f32::INFINITY); + fn reset(&mut self) { + self.candidates.as_mut_slice().fill(Candidate::default()); + self.thresholds.fill(f32::INFINITY); } - pub(super) fn row(&self, row_idx: usize) -> &[Candidate] { - self.rows.row(row_idx) + /// Return one result set in nearest-first order, including unassigned slots. + pub(super) fn candidates(&self, result_idx: usize) -> &[Candidate] { + self.candidates.row(result_idx) } - /// Offer one candidate and synchronize the row's cached threshold. + /// Replace one result set with the nearest candidates in the distance slice. + /// Slice positions become candidate IDs; `result_idx` only selects output storage. #[inline(always)] - pub(super) fn insert(&mut self, row_idx: usize, candidate: Candidate) { - if candidate.distance < self.worst[row_idx] { - self.worst[row_idx] = insert_eligible(row_mut::(&mut self.rows, row_idx), candidate); + pub(super) fn replace_topk( + &mut self, + arch: A, + result_idx: usize, + distances: &[f32], + ) { + let k = self.candidates.ncols(); + if k == 0 { + return; } + let nearest = &mut self.candidates.as_mut_slice()[result_idx * k..][..k]; + nearest.fill(Candidate::default()); + self.thresholds[result_idx] = distance_blocks(arch, distances).fold( + f32::INFINITY, + #[inline(always)] + |limit, block| block.update_one(nearest, limit), + ); } - /// Offer the block's candidates to one row using their slice-column indexes. - #[inline(always)] - pub(super) fn update_one( + /// Add each pair `(point_idx, j)` independently to both endpoints' top-k results. + /// Preserve existing candidates. All slice indexes must be less than `point_idx`. + #[inline] + pub(super) fn update_dual_topk( &mut self, - row_idx: usize, - block: &DistanceBlock<'_, A>, + arch: A, + point_idx: usize, + distances: &[f32], ) { - match *block { - DistanceBlock::Scalar { idx, distance } => { - self.insert(row_idx, Candidate::new(idx as u32, distance)); + let k = self.candidates.ncols(); + if k == 0 { + return; + } + debug_assert!( + distances.len() <= point_idx, + "candidate indexes must exclude the updated point" + ); + // Keep this point's result and local limit independent of reciprocal updates. + let (others, remaining) = self.candidates.as_mut_slice().split_at_mut(point_idx * k); + let nearest = &mut remaining[..k]; + let limit = distance_blocks(arch, distances).fold( + self.thresholds[point_idx], + #[inline(always)] + |limit, block| { + let limit = block.update_one(nearest, limit); + block.update_many(others, self.thresholds, point_idx as u32, k); + limit + }, + ); + self.thresholds[point_idx] = limit; + } +} + +/// A complete SIMD group or one tail distance. This never leaves the TopK module. +/// Both update directions share the loaded vector; scalar reads borrow the input. +#[derive(Clone, Copy)] +enum DistanceBlock<'a, A: PiPNNSIMDSchema> { + Simd { + first_candidate: usize, + values: A::Vector, + distances: &'a [f32], + }, + Scalar { + candidate_idx: usize, + distance: f32, + }, +} + +/// Keep all distance loads and tail indexing inside TopK's architecture scope. +#[inline(always)] +fn distance_blocks( + arch: A, + distances: &[f32], +) -> impl Iterator> { + let simd_end = distances.len() - distances.len() % A::Vector::LANES; + distances[..simd_end] + .chunks_exact(A::Vector::LANES) + .enumerate() + .map(move |(group, distances)| DistanceBlock::Simd { + first_candidate: group * A::Vector::LANES, + // SAFETY: chunks_exact supplies one distance for every SIMD lane. + values: unsafe { A::Vector::load_simd(arch, distances.as_ptr()) }, + distances, + }) + .chain( + distances[simd_end..] + .iter() + .enumerate() + .map(move |(tail, &distance)| DistanceBlock::Scalar { + candidate_idx: simd_end + tail, + distance, + }), + ) +} + +impl DistanceBlock<'_, A> { + /// Offer the block's candidates, preserving and returning the updated limit. + #[inline(always)] + fn update_one(self, nearest: &mut [Candidate], mut max_distance: f32) -> f32 { + match self { + Self::Scalar { + candidate_idx, + distance, + } => { + if distance < max_distance { + max_distance = + insert_sorted(nearest, Candidate::new(candidate_idx as u32, distance)); + } } - DistanceBlock::Simd { - first_idx, + Self::Simd { + first_candidate, values, - lanes, + distances, } => { - let mut worst = self.worst[row_idx]; - let mut eligible = - A::active_lanes(values.lt_simd(A::Vector::splat(values.arch(), worst))); - if eligible == 0 { - return; - } - let row = row_mut::(&mut self.rows, row_idx); - while eligible != 0 { - let lane = eligible.trailing_zeros() as usize; - eligible &= eligible - 1; - // Earlier candidates in this same block can lower the threshold. - if lanes[lane] < worst { - worst = insert_eligible( - row, - Candidate::new((first_idx + lane) as u32, lanes[lane]), - ); + // An unfilled result admits all rankable values. Scan them in + // order without constructing and traversing an all-eligible mask. + if max_distance == f32::INFINITY { + for (lane, &distance) in distances.iter().enumerate() { + if distance < max_distance { + max_distance = insert_sorted( + nearest, + Candidate::new((first_candidate + lane) as u32, distance), + ); + } + } + } else { + let mut eligible = A::active_lanes( + values.lt_simd(A::Vector::splat(values.arch(), max_distance)), + ); + while eligible != 0 { + let lane = eligible.trailing_zeros() as usize; + eligible &= eligible - 1; + // Earlier insertions in this block can lower the limit. + if distances[lane] < max_distance { + max_distance = insert_sorted( + nearest, + Candidate::new((first_candidate + lane) as u32, distances[lane]), + ); + } } } - self.worst[row_idx] = worst; } } + max_distance } - /// Offer one candidate to each row identified by the block's slice-column indexes. + /// Offer one point to each result set identified by this block's candidate IDs. #[inline(always)] - pub(super) fn update_many( - &mut self, - candidate_idx: u32, - block: &DistanceBlock<'_, A>, + fn update_many( + self, + candidates: &mut [Candidate], + thresholds: &mut [f32], + point_idx: u32, + k: usize, ) { - match *block { - DistanceBlock::Scalar { idx, distance } => { - self.insert(idx, Candidate::new(candidate_idx, distance)); + match self { + Self::Scalar { + candidate_idx, + distance, + } => { + let limit = &mut thresholds[candidate_idx]; + if distance < *limit { + *limit = insert_sorted( + &mut candidates[candidate_idx * k..][..k], + Candidate::new(point_idx, distance), + ); + } } - DistanceBlock::Simd { - first_idx, + Self::Simd { + first_candidate, values, - lanes, + distances, } => { - let thresholds = &mut self.worst[first_idx..][..A::Vector::LANES]; - // SAFETY: the slice above has one threshold for every SIMD lane. - let worst = unsafe { A::Vector::load_simd(values.arch(), thresholds.as_ptr()) }; - let mut eligible = A::active_lanes(values.lt_simd(worst)); + let limits = &mut thresholds[first_candidate..][..A::Vector::LANES]; + // SAFETY: the slice above has one distance limit per SIMD lane. + let limit_values = unsafe { A::Vector::load_simd(values.arch(), limits.as_ptr()) }; + let mut eligible = A::active_lanes(values.lt_simd(limit_values)); while eligible != 0 { let lane = eligible.trailing_zeros() as usize; eligible &= eligible - 1; - // Each lane updates a different row once; its threshold is still current. - thresholds[lane] = insert_eligible( - row_mut::(&mut self.rows, first_idx + lane), - Candidate::new(candidate_idx, lanes[lane]), + // Each lane updates a different result; its limit is current. + limits[lane] = insert_sorted( + &mut candidates[(first_candidate + lane) * k..][..k], + Candidate::new(point_idx, distances[lane]), ); } } @@ -244,258 +316,449 @@ impl<'a, const K: usize> TopKRows<'a, K> { } } -/// Preserve fixed row lengths for inlined candidate insertion. -#[inline(always)] -fn row_mut<'a, const K: usize>( - rows: &'a mut MutMatrixView<'_, Candidate>, - row_idx: usize, -) -> &'a mut [Candidate] { - if K == RUNTIME_WIDTH { - rows.row_mut(row_idx) - } else { - let (rows, _) = rows.as_mut_slice().as_chunks_mut::(); - &mut rows[row_idx] - } -} - -/// Insert after the threshold check; keep candidates in nearest-first order. +/// Insert an eligible candidate in nearest-first order and return the new distance limit. +/// The caller must check the candidate against the current limit before insertion. #[inline(always)] -fn insert_eligible(row: &mut [Candidate], candidate: Candidate) -> f32 { - let last = row.len() - 1; +fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate) -> f32 { + let last = nearest.len() - 1; let mut slot = last; - while slot > 0 && candidate.distance < row[slot - 1].distance { - row[slot] = row[slot - 1]; + while slot > 0 && candidate.distance < nearest[slot - 1].distance { + nearest[slot] = nearest[slot - 1]; slot -= 1; } - row[slot] = candidate; - row[last].distance + nearest[slot] = candidate; + nearest[last].distance } #[cfg(test)] mod tests { use super::*; + use diskann_wide::arch::{self, Target1}; + use rstest::rstest; + + // These adapters only select the real CPU backend; all updates use production TopK. + struct ReplaceRow; + impl Target1, usize, &[f32])> + for ReplaceRow + { + fn run(self, arch: A, (topk, row, distances): (&mut TopK<'_, K>, usize, &[f32])) { + topk.replace_topk(arch, row, distances); + } + } + + struct UpdatePair; + impl Target1, usize, &[f32])> + for UpdatePair + { + fn run(self, arch: A, (topk, point, distances): (&mut TopK<'_, K>, usize, &[f32])) { + topk.update_dual_topk(arch, point, distances); + } + } + + fn sorted_candidates(distances: &[f32], width: usize) -> Vec { + let mut candidates: Vec<_> = distances + .iter() + .enumerate() + .filter(|(_, distance)| **distance < f32::INFINITY) + .map(|(id, &distance)| Candidate::new(id as u32, distance)) + .collect(); + candidates.sort_by(|left, right| left.distance.partial_cmp(&right.distance).unwrap()); + candidates.truncate(width); + candidates.resize(width, Candidate::default()); + candidates + } - mod insert_eligible_tests { + mod new_tests { use super::*; - use rstest::rstest; - #[rstest] - #[case::singleton(&[(1, 4.0)], (2, 2.0), &[(2, 2.0)])] - #[case::no_shift(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 3.0), &[(1, 1.0), (2, 2.0), (4, 3.0)])] - #[case::one_shift(&[(1, 1.0), (2, 2.0), (3, 4.0)], (4, 1.5), &[(1, 1.0), (4, 1.5), (2, 2.0)])] - #[case::multiple_shifts(&[(1, 1.0), (2, 2.0), (3, 3.0), (4, 5.0)], (5, 0.5), &[(5, 0.5), (1, 1.0), (2, 2.0), (3, 3.0)])] - #[case::negative_infinity(&[(1, -1.0), (2, 1.0), (UNASSIGNED, f32::INFINITY)], (3, f32::NEG_INFINITY), &[(3, f32::NEG_INFINITY), (1, -1.0), (2, 1.0)])] - #[case::positive_zero(&[(1, 1.0)], (2, 0.0), &[(2, 0.0)])] - #[case::negative_zero(&[(1, 1.0)], (2, -0.0), &[(2, -0.0)])] - fn insertion_keeps_nearest_candidates_and_returns_threshold( - #[case] initial: &[(u32, f32)], - #[case] candidate: (u32, f32), - #[case] expected: &[(u32, f32)], - ) { - // Given - let expected_threshold = expected.last().unwrap().1.to_bits(); - let expected: Vec<_> = expected - .iter() - .map(|&(idx, distance)| (idx, distance.to_bits())) - .collect(); - let mut row: Vec<_> = initial - .iter() - .map(|&(idx, distance)| Candidate::new(idx, distance)) - .collect(); + #[test] + fn reused_storage_is_cleared_before_dual_updates() { + let mut output = [ + Candidate::new(1, -10.0), + Candidate::new(2, -9.0), + Candidate::new(0, -10.0), + Candidate::new(2, -8.0), + Candidate::new(0, -9.0), + Candidate::new(1, -8.0), + ]; + let mut thresholds = [-9.0, -8.0, -8.0]; + let expected = [Candidate::new(1, 2.0), Candidate::default()]; + let mut topk = TopK::<2>::new( + MutMatrixView::try_from(&mut output[..], 3, 2).unwrap(), + &mut thresholds, + ); - // When - let threshold = insert_eligible(&mut row, Candidate::new(candidate.0, candidate.1)); + arch::dispatch1_no_features(UpdatePair, (&mut topk, 1, &[2.0][..])); - // Then - let actual: Vec<_> = row - .iter() - .map(|candidate| (candidate.local_idx, candidate.distance.to_bits())) - .collect(); - assert_eq!(actual, expected); - assert_eq!(threshold.to_bits(), expected_threshold); + assert_eq!(topk.candidates(0), expected); + assert_eq!( + topk.candidates(1), + [Candidate::new(0, 2.0), Candidate::default()] + ); + assert_eq!(topk.candidates(2), [Candidate::default(); 2]); } } #[test] - fn reset_clears_rows_and_thresholds() { - // Given - let mut output = [Candidate::new(10, -1.0), Candidate::new(11, 2.0)]; - let mut worst = [-1.0, 2.0]; - let mut topk = TopKRows::<1> { - rows: MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), - worst: &mut worst, - }; + fn zero_capacity_accepts_both_update_operations() { + let mut output = []; + let mut thresholds = [f32::INFINITY; 2]; + let mut topk = TopK::::new( + MutMatrixView::try_from(&mut output[..], 2, 0).unwrap(), + &mut thresholds, + ); - // When - topk.reset(); + arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &[1.0, 2.0][..])); + arch::dispatch1_no_features(UpdatePair, (&mut topk, 1, &[0.0][..])); - // Then - assert_eq!(topk.rows.as_slice(), [Candidate::default(); 2]); - assert_eq!(topk.worst, [f32::INFINITY; 2]); + assert!(topk.candidates(0).is_empty()); + assert!(topk.candidates(1).is_empty()); } - mod update_tests { + mod replace_topk_tests { use super::*; - use crate::graph::pipnn::topk::distance_blocks; - use rstest::rstest; - #[test] - fn one_row_rechecks_the_threshold_after_each_insert() { - let mut distances = [f32::INFINITY; 16]; - distances[0] = 1.0; - distances[1] = 9.0; - let mut output = [Candidate::default(); 1]; - let mut worst = [f32::INFINITY]; - let mut topk = TopKRows::::new( - MutMatrixView::row_vector(&mut output[..]), - &mut worst, + #[rstest] + #[case::lane_minus_one(15, 3)] + #[case::complete_lane(16, 2)] + #[case::lane_and_tail(17, 10)] + #[case::two_lanes_minus_one(31, 1)] + #[case::two_complete_lanes(32, 4)] + #[case::two_lanes_and_tail(33, 17)] + fn dispatched_selection_matches_independent_sort( + #[case] count: usize, + #[case] width: usize, + ) { + // Alternating signs force insertions at different positions, not just append. + let distances: Vec<_> = (0..count) + .map(|i| if i % 2 == 0 { i as f32 } else { -(i as f32) }) + .collect(); + let expected = sorted_candidates(&distances, width); + let mut output = vec![Candidate::default(); width]; + let mut thresholds = [f32::INFINITY]; + + with_topk!( + MutMatrixView::row_vector(output.as_mut_slice()), + &mut thresholds, + |topk| { + arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, distances.as_slice())); + assert_eq!(topk.candidates(0), expected); + } ); + } - for block in distance_blocks(diskann_wide::ARCH, &distances) { - topk.update_one(0, &block); - } + #[rstest] + #[case::unrankable(&[f32::NAN, f32::INFINITY], [Candidate::default(); 2])] + #[case::partly_filled( + &[f32::NAN, f32::INFINITY, 9.0], + [Candidate::new(2, 9.0), Candidate::default()] + )] + fn each_point_replaces_the_previous_points_candidates( + #[case] distances: &[f32], + #[case] expected: [Candidate; 2], + ) { + let mut output = [Candidate::default(); 2]; + let mut thresholds = [f32::INFINITY]; + let mut topk = + TopK::<2>::new(MutMatrixView::row_vector(&mut output[..]), &mut thresholds); + topk.replace_topk(diskann_wide::ARCH, 0, &[3.0, 1.0]); + + arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, distances)); - assert_eq!(topk.row(0), [Candidate::new(0, 1.0)]); - assert_eq!(topk.worst, [1.0]); + assert_eq!(topk.candidates(0), expected); } #[test] - fn many_rows_use_their_own_threshold_and_distance() { - let mut distances = [2.0; 16]; - distances[1] = 4.0; - distances[2] = 6.0; - let mut output = [Candidate::default(); 16]; - let mut worst = [f32::INFINITY; 16]; - let mut topk = TopKRows::::new( - MutMatrixView::try_from(&mut output[..], 16, 1).unwrap(), - &mut worst, + fn finite_simd_mask_rechecks_the_threshold_after_insertion() { + // Group one fills K=1 at 5. Both 1 and 4 pass group two's initial + // mask, but inserting 1 must prevent the later 4 from replacing it. + let mut distances = [f32::INFINITY; 32]; + distances[0] = 5.0; + distances[16] = 1.0; + distances[17] = 4.0; + let mut output = [Candidate::default()]; + let mut thresholds = [f32::INFINITY]; + let mut topk = TopK::<1>::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + + arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &distances[..])); + + assert_eq!(topk.candidates(0), [Candidate::new(16, 1.0)]); + } + + #[test] + fn closer_candidates_shift_the_middle_of_a_full_result() { + let mut output = [Candidate::default(); 3]; + let mut thresholds = [f32::INFINITY]; + let mut topk = TopK::<3>::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + + // After retaining [0, 4, 6], distances 2 and 3 must enter between + // existing candidates instead of replacing the last slot directly. + arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &[0.0, 4.0, 6.0, 2.0, 3.0][..])); + + assert_eq!( + topk.candidates(0), + [ + Candidate::new(0, 0.0), + Candidate::new(3, 2.0), + Candidate::new(4, 3.0) + ] ); - topk.insert(0, Candidate::new(80, 1.0)); - topk.insert(1, Candidate::new(81, 5.0)); - topk.insert(2, Candidate::new(82, 5.0)); + } - for block in distance_blocks(diskann_wide::ARCH, &distances) { - topk.update_many(42, &block); - } + #[test] + fn non_rankable_first_block_preserves_later_candidate_ids() { + let mut distances = [f32::INFINITY; 33]; + distances[..8].fill(f32::NAN); + distances[16] = 4.0; + distances[17] = 2.0; + distances[31] = 3.0; + distances[32] = 1.0; + let mut output = [Candidate::default(); 3]; + let mut thresholds = [f32::INFINITY]; + let mut topk = TopK::<3>::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + + arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &distances[..])); + + assert_eq!( + topk.candidates(0), + [ + Candidate::new(32, 1.0), + Candidate::new(17, 2.0), + Candidate::new(31, 3.0), + ] + ); + } + + #[test] + fn nearer_tail_displaces_a_tied_candidate_without_duplicate_ids() { + let mut distances = [2.0; 17]; + distances[16] = 1.0; + let mut output = [Candidate::default(); 3]; + let mut thresholds = [f32::INFINITY]; + let mut topk = TopK::<3>::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + + arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &distances[..])); - assert_eq!(topk.row(0), [Candidate::new(80, 1.0)]); - assert_eq!(topk.row(1), [Candidate::new(42, 4.0)]); - assert_eq!(topk.row(2), [Candidate::new(82, 5.0)]); + assert_eq!(topk.candidates(0)[0], Candidate::new(16, 1.0)); + let tied = &topk.candidates(0)[1..]; assert!( - topk.rows - .row_iter() - .skip(3) - .all(|row| row == [Candidate::new(42, 2.0)]) + tied.iter() + .all(|candidate| candidate.local_idx < 16 && candidate.distance == 2.0) ); - assert_eq!(topk.worst[..3], [1.0, 4.0, 5.0]); - assert!(topk.worst[3..].iter().all(|&worst| worst == 2.0)); + assert_ne!(tied[0].local_idx, tied[1].local_idx); } #[test] - fn one_and_many_updates_have_independent_acceptance() { - let distances = [2.0; 16]; - let mut output = [Candidate::default(); 17]; - let mut worst = [f32::INFINITY; 17]; - let mut topk = TopKRows::::new( - MutMatrixView::try_from(&mut output[..], 17, 1).unwrap(), - &mut worst, + fn special_rankable_distances_keep_distinct_ids_and_original_bits() { + let mut distances = [f32::INFINITY; 33]; + distances[0] = -0.0; + distances[16] = 0.0; + distances[17] = f32::MAX; + distances[32] = f32::NEG_INFINITY; + let mut output = [Candidate::default(); 4]; + let mut thresholds = [f32::INFINITY]; + let mut topk = + TopK::::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + + arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &distances[..])); + + assert_eq!(topk.candidates(0)[0], Candidate::new(32, f32::NEG_INFINITY)); + // Tie order is unspecified; each retained value must match its own ID. + let mut actual: Vec<_> = topk + .candidates(0) + .iter() + .map(|candidate| (candidate.local_idx, candidate.distance.to_bits())) + .collect(); + actual.sort_unstable_by_key(|&(id, _)| id); + assert_eq!( + actual, + [ + (0, (-0.0_f32).to_bits()), + (16, 0.0_f32.to_bits()), + (17, f32::MAX.to_bits()), + (32, f32::NEG_INFINITY.to_bits()) + ] ); - topk.insert(16, Candidate::new(70, 1.0)); + } + } + + mod update_dual_topk_tests { + use super::*; - for block in distance_blocks(diskann_wide::ARCH, &distances) { - topk.update_one(16, &block); - topk.update_many(16, &block); + #[test] + fn endpoints_accept_independently_in_both_simd_groups_and_tail() { + let mut distances = [f32::INFINITY; 33]; + distances[0] = 0.5; + distances[1] = 4.0; + distances[2] = 6.0; + distances[16] = 2.0; + distances[32] = 0.75; + let mut output = [Candidate::default(); 34]; + let mut thresholds = [f32::INFINITY; 34]; + let mut topk = TopK::<1>::new( + MutMatrixView::try_from(&mut output[..], 34, 1).unwrap(), + &mut thresholds, + ); + // Earlier leaf rows establish each destination's threshold. + let mut previous = [f32::INFINITY; 32]; + for point in 1..33 { + if point == 3 { + previous[..3].copy_from_slice(&[0.2, 5.0, 5.0]); + } + arch::dispatch1_no_features(UpdatePair, (&mut topk, point, &previous[..point])); + previous.fill(f32::INFINITY); } + let mut expected = [Candidate::default(); 34]; + expected[0] = Candidate::new(3, 0.2); + expected[1] = Candidate::new(33, 4.0); + expected[2] = Candidate::new(3, 5.0); + expected[3] = Candidate::new(0, 0.2); + expected[16] = Candidate::new(33, 2.0); + expected[32] = Candidate::new(33, 0.75); + expected[33] = Candidate::new(0, 0.5); + + arch::dispatch1_no_features(UpdatePair, (&mut topk, 33, &distances[..])); + + assert_eq!(topk.candidates.as_slice(), expected); + } - assert_eq!(topk.row(16), [Candidate::new(70, 1.0)]); - assert!( - topk.rows - .row_iter() - .take(16) - .all(|row| row == [Candidate::new(16, 2.0)]) + #[test] + fn a_source_fills_across_blocks_with_few_rankable_distances() { + let mut output = [Candidate::default(); 34 * 3]; + let mut thresholds = [f32::INFINITY; 34]; + let mut topk = TopK::<3>::new( + MutMatrixView::try_from(&mut output[..], 34, 3).unwrap(), + &mut thresholds, + ); + let mut distances = [f32::INFINITY; 33]; + distances[..8].fill(f32::NAN); + distances[15] = 4.0; + distances[16] = 2.0; + distances[17] = 9.0; + distances[31] = 3.0; + distances[32] = f32::NEG_INFINITY; + + // A leaf source starts empty. It remains partially filled after + // group one and acquires its finite threshold within group two. + arch::dispatch1_no_features(UpdatePair, (&mut topk, 33, &distances[..])); + + assert_eq!( + topk.candidates(33), + [ + Candidate::new(32, f32::NEG_INFINITY), + Candidate::new(16, 2.0), + Candidate::new(31, 3.0) + ] + ); + assert_eq!( + topk.candidates(17), + [ + Candidate::new(33, 9.0), + Candidate::default(), + Candidate::default() + ] ); } - #[rstest] - #[case::nan(f32::NAN)] - #[case::infinity(f32::INFINITY)] - fn non_rankable_candidates_leave_reset_rows_unassigned(#[case] distance: f32) { - let distances = [distance; 16]; - let mut output = [Candidate::new(70, 1.0); 16]; - let mut worst = [1.0; 16]; - let mut topk = TopKRows::::new( - MutMatrixView::try_from(&mut output[..], 16, 1).unwrap(), - &mut worst, + #[test] + fn later_pairs_use_the_sources_updated_threshold() { + let mut output = [Candidate::default(); 3]; + let mut thresholds = [f32::INFINITY; 3]; + let mut topk = TopK::<1>::new( + MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), + &mut thresholds, ); - for block in distance_blocks(diskann_wide::ARCH, &distances) { - topk.update_one(0, &block); - topk.update_many(42, &block); - } - topk.insert(0, Candidate::new(43, distance)); + arch::dispatch1_no_features(UpdatePair, (&mut topk, 1, &[1.0][..])); + arch::dispatch1_no_features(UpdatePair, (&mut topk, 2, &[3.0, 2.0][..])); - assert!( - topk.rows - .as_slice() - .iter() - .all(|&candidate| candidate == Candidate::default()) + assert_eq!( + topk.candidates.as_slice(), + [ + Candidate::new(1, 1.0), + Candidate::new(0, 1.0), + Candidate::new(1, 2.0) + ] + ); + } + + #[test] + fn runtime_capacity_updates_rows_filled_by_earlier_leaf_pairs() { + let mut output = [Candidate::default(); 18 * 4]; + let mut thresholds = [f32::INFINITY; 18]; + let mut previous = [f32::INFINITY; 17]; + let mut distances = [f32::INFINITY; 17]; + distances[0] = 4.0; + distances[2] = 3.0; + distances[15] = 2.0; + distances[16] = 1.0; + + with_topk!( + MutMatrixView::try_from(&mut output[..], 18, 4).unwrap(), + &mut thresholds, + |topk| { + for point in 1..17 { + if point == 15 { + previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); + } + arch::dispatch1_no_features( + UpdatePair, + (&mut topk, point, &previous[..point]), + ); + previous.fill(f32::INFINITY); + } + arch::dispatch1_no_features(UpdatePair, (&mut topk, 17, &distances[..])); + + assert_eq!( + topk.candidates(17), + [ + Candidate::new(16, 1.0), + Candidate::new(15, 2.0), + Candidate::new(2, 3.0), + Candidate::new(0, 4.0) + ] + ); + assert_eq!( + topk.candidates(15), + [ + Candidate::new(0, -1.0), + Candidate::new(17, 2.0), + Candidate::new(1, 3.0), + Candidate::new(2, 5.0) + ] + ); + assert_eq!( + topk.candidates(16), + [ + Candidate::new(17, 1.0), + Candidate::default(), + Candidate::default(), + Candidate::default() + ] + ); + } ); - assert!(topk.worst.iter().all(|&worst| worst == f32::INFINITY)); } #[rstest] - #[case::empty(0)] - #[case::lane_minus_one(15)] - #[case::one_lane(16)] - #[case::lane_with_tail(17)] - #[case::two_lanes_with_tail(33)] - #[trace] - fn shared_block_updates_match_runtime_scalar_insertion( - #[case] count: usize, - #[values(1, 2, 3, 4)] width: usize, - ) { - // Given - let rows = count + 1; - let source_idx = count; - let distances: Vec<_> = (0..count) - .map(|idx| { - let distance = idx as f32 + 0.25; - if idx % 2 == 0 { -distance } else { distance } - }) - .collect(); - let mut actual_rows = vec![Candidate::default(); rows * width]; - let mut expected_rows = actual_rows.clone(); - let mut actual_worst = vec![f32::INFINITY; rows]; - let mut expected_worst = actual_worst.clone(); - let mut expected = TopKRows::::new( - MutMatrixView::try_from(expected_rows.as_mut_slice(), rows, width).unwrap(), - &mut expected_worst, + #[case::simd(15)] + #[case::tail(16)] + fn negative_infinity_reaches_both_endpoints(#[case] candidate: usize) { + let mut distances = [f32::NAN; 17]; + distances[candidate] = f32::NEG_INFINITY; + let mut output = [Candidate::default(); 18 * 2]; + let mut thresholds = [f32::INFINITY; 18]; + let mut topk = TopK::<2>::new( + MutMatrixView::try_from(&mut output[..], 18, 2).unwrap(), + &mut thresholds, ); - expected.insert(source_idx, Candidate::new(90, 0.0)); - expected.insert(0, Candidate::new(91, -4.0)); - for (idx, &distance) in distances.iter().enumerate() { - expected.insert(source_idx, Candidate::new(idx as u32, distance)); - expected.insert(idx, Candidate::new(42, distance)); - } + let mut expected = [[Candidate::default(); 2]; 18]; + expected[17][0] = Candidate::new(candidate as u32, f32::NEG_INFINITY); + expected[candidate][0] = Candidate::new(17, f32::NEG_INFINITY); - with_topk_rows!( - MutMatrixView::try_from(actual_rows.as_mut_slice(), rows, width).unwrap(), - &mut actual_worst[..], - |actual| { - actual.insert(source_idx, Candidate::new(90, 0.0)); - actual.insert(0, Candidate::new(91, -4.0)); - - // When - for block in distance_blocks(diskann_wide::ARCH, &distances) { - actual.update_one(source_idx, &block); - actual.update_many(42, &block); - } + arch::dispatch1_no_features(UpdatePair, (&mut topk, 17, &distances[..])); - // Then - assert_eq!(actual.rows.as_slice(), expected.rows.as_slice()); - assert_eq!(actual.worst, expected.worst); - } - ); + assert_eq!(topk.candidates.as_slice(), expected.as_flattened()); } } } From 4a3aed73efe0c49341e590376c8a3d0bc76b564d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:31:33 +0000 Subject: [PATCH 170/229] test(pipnn): focus coverage on live contracts Remove unreachable empty-kernel smoke cases and duplicate delegated checks. Keep point replacement fixtures at a fixed leader count, combine cosine formula and multi-row checks, and exercise rectangular AAT inputs. --- diskann-linalg/src/lib.rs | 53 +++------------- diskann/src/graph/pipnn/leaf_kernel.rs | 14 ----- diskann/src/graph/pipnn/leaf_metric.rs | 1 - diskann/src/graph/pipnn/partition_kernel.rs | 17 ----- diskann/src/graph/pipnn/partition_metric.rs | 70 +++++++++------------ diskann/src/graph/pipnn/topk.rs | 4 +- 6 files changed, 40 insertions(+), 119 deletions(-) diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index 94582f555a..2c88e5a359 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -725,27 +725,14 @@ mod sgemm_aat_lower_tests { #[test] fn replace_writes_the_scaled_lower_triangle_and_preserves_the_upper_triangle() { - // Given - let first_row = [1.0_f32, 2.0]; - let second_row = [3.0_f32, 4.0]; - let row_count = 2; - let dimension_count = first_row.len(); - let matrix = [first_row[0], first_row[1], second_row[0], second_row[1]]; - let scale = -1.0_f32; - let first_self_dot = first_row[0].mul_add(first_row[0], first_row[1] * first_row[1]); - let cross_dot = first_row[0].mul_add(second_row[0], first_row[1] * second_row[1]); - let second_self_dot = second_row[0].mul_add(second_row[0], second_row[1] * second_row[1]); + // Given: rows [1, 2, 3] and [4, 5, 6] have Gram entries 14, 32, and 77. + let matrix = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; let upper_triangle_sentinel = -123.0_f32; - let expected = [ - scale * first_self_dot, - upper_triangle_sentinel, - scale * cross_dot, - scale * second_self_dot, - ]; + let expected = [-14.0, upper_triangle_sentinel, -32.0, -77.0]; let mut actual = [upper_triangle_sentinel; 4]; // When - sgemm_aat_lower(row_count, dimension_count, scale, &matrix, &mut actual).unwrap(); + sgemm_aat_lower(2, 3, -1.0, &matrix, &mut actual).unwrap(); // Then assert_eq!(actual, expected); @@ -753,35 +740,15 @@ mod sgemm_aat_lower_tests { #[test] fn add_accumulates_the_scaled_lower_triangle_and_preserves_the_upper_triangle() { - // Given - let first_row = [1.0_f32, 2.0]; - let second_row = [3.0_f32, 4.0]; - let row_count = 2; - let dimension_count = first_row.len(); - let matrix = [first_row[0], first_row[1], second_row[0], second_row[1]]; - let scale = -2.0_f32; - let first_self_dot = first_row[0].mul_add(first_row[0], first_row[1] * first_row[1]); - let cross_dot = first_row[0].mul_add(second_row[0], first_row[1] * second_row[1]); - let second_self_dot = second_row[0].mul_add(second_row[0], second_row[1] * second_row[1]); - let initial_first = 10.0_f32; - let initial_cross = 20.0_f32; - let initial_second = 30.0_f32; + // Given: adding -2 times the Gram entries (14, 32, 77) to (10, 20, 30) + // produces (-18, -44, -124), while the upper triangle remains unchanged. + let matrix = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; let upper_triangle_sentinel = -123.0_f32; - let expected = [ - scale.mul_add(first_self_dot, initial_first), - upper_triangle_sentinel, - scale.mul_add(cross_dot, initial_cross), - scale.mul_add(second_self_dot, initial_second), - ]; - let mut actual = [ - initial_first, - upper_triangle_sentinel, - initial_cross, - initial_second, - ]; + let expected = [-18.0, upper_triangle_sentinel, -44.0, -124.0]; + let mut actual = [10.0, upper_triangle_sentinel, 20.0, 30.0]; // When - sgemm_aat_lower_add(row_count, dimension_count, scale, &matrix, &mut actual).unwrap(); + sgemm_aat_lower_add(2, 3, -2.0, &matrix, &mut actual).unwrap(); // Then assert_eq!(actual, expected); diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index e50d0da729..92d7b6a229 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -324,19 +324,5 @@ mod tests { // Then: every reciprocal ID is local to its output row; other rows stay empty. assert_eq!(actual.as_slice(), expected.as_flattened()); } - - #[test] - fn singleton_leaf_has_no_neighbors() { - let actual = rank_distance_fixture(&[f32::NAN], 1, 0); - - assert!(actual.is_empty()); - } - - #[test] - fn zero_neighbor_capacity_returns_empty_rows() { - let actual = rank_distance_fixture(&[1.0; 9], 3, 0); - - assert!(actual.is_empty()); - } } } diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index f6b2c9cdf6..bbd6dd2436 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -224,7 +224,6 @@ mod tests { #[rstest] #[case::l2(compute_pair_ranking::)] #[case::cosine(compute_pair_ranking::)] - #[case::normalized_cosine(compute_pair_ranking::)] #[case::inner_product(compute_pair_ranking::)] fn nan_coordinate_produces_nan_ranking( #[case] compute: fn([f32; DIMENSION_COUNT], [f32; DIMENSION_COUNT]) -> f32, diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index e67229d5d7..57118d2671 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -236,22 +236,5 @@ mod tests { // Then assert_eq!(output, expected_regrown_stripe); } - - #[test] - fn empty_point_stripe_leaves_output_empty() { - let leader_values = [1.0_f32]; - let leaders = L2::create_leaders(MatrixView::row_vector(&leader_values[..])); - let points = MatrixView::try_from(&[][..], 0, 1).unwrap(); - let mut output = []; - - assign_leaders::<_, L2>( - diskann_wide::ARCH, - points, - &leaders, - MutMatrixView::try_from(&mut output[..], 0, 1).unwrap(), - &mut PartitionKernelWorkspace::default(), - ) - .unwrap(); - } } } diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 0933d22220..4d3282d1f1 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -276,8 +276,8 @@ mod tests { #[rstest::rstest] #[case::wrong_rows_with_equal_area(1, 4)] #[case::wrong_columns(2, 1)] - fn output_shape_mismatch_leaves_distances_unchanged( - #[values(L2, Cosine, CosineNormalized, InnerProduct)] _metric: M, + fn output_shape_mismatch_returns_error( + #[values(L2, Cosine, InnerProduct)] _metric: M, #[case] rows: usize, #[case] columns: usize, ) { @@ -285,24 +285,17 @@ mod tests { let leader_values = [1.0, 0.0, 0.0, 2.0]; let leaders = M::create_leaders(matrix(&leader_values, 2)); let points = matrix(&[0.0, 4.0, -2.0, 0.0], 2); - let expected = vec![STALE_DISTANCE; rows * columns]; - let mut storage = expected.clone(); + let mut storage = vec![STALE_DISTANCE; rows * columns]; // When - let error = M::compute_distances( + let result = M::compute_distances( points, &leaders, MutMatrixView::try_from(storage.as_mut_slice(), rows, columns).unwrap(), - ) - .unwrap_err(); - - // Then: a shape error is returned before any output is written. - assert!( - error - .to_string() - .contains("point-to-leader output shape mismatch") ); - assert_eq!(storage, expected); + + // Then: invalid output shapes stop the build. + assert!(result.is_err()); } #[test] @@ -326,26 +319,6 @@ mod tests { assert_eq!(storage, expected); } - #[test] - fn cosine_ranking_equals_one_minus_normalized_similarity() { - // Given - let point = [2.0_f32, 0.0]; - let leader = [1.0_f32, 1.0]; - let dot = point[0].mul_add(leader[0], point[1] * leader[1]); - let point_norm = point[0].hypot(point[1]); - let leader_norm = leader[0].hypot(leader[1]); - let expected = 1.0 - dot / (point_norm * leader_norm); - - // When - let actual = compute_one_ranking::(point, leader); - - // Then - assert!( - (actual - expected).abs() <= FLOAT_TOLERANCE, - "actual {actual} differs from expected {expected}" - ); - } - #[rstest::rstest] #[case::zero_point([0.0, 0.0], [1.0, 0.0])] #[case::zero_leader([1.0, 0.0], [0.0, 0.0])] @@ -368,7 +341,6 @@ mod tests { #[rstest::rstest] #[case::l2(compute_one_ranking::)] #[case::cosine(compute_one_ranking::)] - #[case::normalized_cosine(compute_one_ranking::)] #[case::inner_product(compute_one_ranking::)] fn nan_coordinate_produces_nan_ranking( #[case] compute: fn([f32; DIMENSION_COUNT], [f32; DIMENSION_COUNT]) -> f32, @@ -416,23 +388,37 @@ mod tests { #[test] fn cosine_overwrites_every_point_row_with_normalized_distances() { - // Given: points and leaders lie on signed coordinate axes. - let leader_values = [1.0_f32, 0.0, 0.0, 2.0]; + // Given: non-unit points point right, up, and left. Leaders point + // diagonally up-right and up, so similarities include ±1/sqrt(2). + let leader_values = [1.0_f32, 1.0, 0.0, 2.0]; let leaders = Cosine::create_leaders(matrix(&leader_values, 2)); - let points = matrix(&[0.0, 4.0, -2.0, 0.0], 2); - let expected = [1.0, 0.0, 2.0, 1.0]; - let mut storage = [STALE_DISTANCE; 4]; + let points = matrix(&[2.0, 0.0, 0.0, 4.0, -2.0, 0.0], 3); + let diagonal_similarity = std::f32::consts::FRAC_1_SQRT_2; + let expected = [ + 1.0 - diagonal_similarity, + 1.0, + 1.0 - diagonal_similarity, + 0.0, + 1.0 + diagonal_similarity, + 1.0, + ]; + let mut storage = [STALE_DISTANCE; 6]; // When Cosine::compute_distances( points, &leaders, - MutMatrixView::try_from(&mut storage[..], 2, 2).unwrap(), + MutMatrixView::try_from(&mut storage[..], 3, 2).unwrap(), ) .unwrap(); // Then - assert_eq!(storage, expected); + for (actual, expected) in storage.into_iter().zip(expected) { + assert!( + (actual - expected).abs() <= FLOAT_TOLERANCE, + "actual {actual} differs from expected {expected}" + ); + } } } } diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 730c516305..d11270f6cd 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -448,7 +448,7 @@ mod tests { } #[rstest] - #[case::unrankable(&[f32::NAN, f32::INFINITY], [Candidate::default(); 2])] + #[case::unrankable(&[f32::NAN, f32::INFINITY, f32::INFINITY], [Candidate::default(); 2])] #[case::partly_filled( &[f32::NAN, f32::INFINITY, 9.0], [Candidate::new(2, 9.0), Candidate::default()] @@ -461,7 +461,7 @@ mod tests { let mut thresholds = [f32::INFINITY]; let mut topk = TopK::<2>::new(MutMatrixView::row_vector(&mut output[..]), &mut thresholds); - topk.replace_topk(diskann_wide::ARCH, 0, &[3.0, 1.0]); + topk.replace_topk(diskann_wide::ARCH, 0, &[3.0, 1.0, 2.0]); arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, distances)); From dc10909160fc32073ef7d223c0ceaae39a225951 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:44:09 +0000 Subject: [PATCH 171/229] refactor(pipnn): borrow top-k buffers per call Keep partition selection thresholds local and initialize reusable leaf state explicitly. Kernels own candidate storage across updates. --- diskann/src/graph/pipnn/leaf_kernel.rs | 14 +- diskann/src/graph/pipnn/partition_kernel.rs | 17 +- diskann/src/graph/pipnn/topk.rs | 522 +++++++++++--------- 3 files changed, 316 insertions(+), 237 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 92d7b6a229..9c4c3c5c7f 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -93,14 +93,20 @@ where fn rank_leaf_distances( arch: A, distances: MatrixView<'_, f32>, - output: MutMatrixView<'_, Candidate>, + mut output: MutMatrixView<'_, Candidate>, worst: &mut Vec, ) { - worst.resize(output.nrows(), f32::INFINITY); arch.run(move || { - with_topk!(output, worst.as_mut_slice(), |topk| { + with_topk!(output.ncols(), |topk| { + topk.initialize(output.as_mut_view(), worst); for point_idx in 1..distances.nrows() { - topk.update_dual_topk(arch, point_idx, &distances.row(point_idx)[..point_idx]); + topk.update_dual_topk( + arch, + point_idx, + &distances.row(point_idx)[..point_idx], + output.as_mut_view(), + worst.as_mut_slice(), + ); } }); }); diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 57118d2671..d6047cd261 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -87,19 +87,14 @@ fn rank_leader_distances( ) { candidates.resize(output.ncols(), Candidate::default()); arch.run(move || { - let mut worst = [f32::INFINITY]; - with_topk!( - MutMatrixView::row_vector(candidates.as_mut_slice()), - &mut worst, - |topk| { - for (distances, output) in distances.row_iter().zip(output.row_iter_mut()) { - topk.replace_topk(arch, 0, distances); - for (destination, candidate) in output.iter_mut().zip(topk.candidates(0)) { - *destination = candidate.local_idx; - } + with_topk!(output.ncols(), |topk| { + for (distances, output) in distances.row_iter().zip(output.row_iter_mut()) { + topk.select_topk(arch, distances, candidates.as_mut_slice()); + for (destination, candidate) in output.iter_mut().zip(candidates.iter()) { + *destination = candidate.local_idx; } } - ); + }); }); } diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index d11270f6cd..7c06fad69b 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -44,32 +44,31 @@ impl Default for Candidate { pub(super) const RUNTIME_WIDTH: usize = 0; /// Select the width specialization once for a batch. -/// Storage expressions are evaluated once. The body has ordinary block control flow. +/// The width is evaluated once. The body has ordinary block control flow. macro_rules! with_topk { - ($candidates:expr, $thresholds:expr, |$topk:ident| $body:block) => {{ - let candidates = $candidates; - let thresholds = $thresholds; - match candidates.ncols() { + ($width:expr, |$topk:ident| $body:block) => {{ + let width = $width; + match width { 1 => { - let mut $topk = $crate::graph::pipnn::topk::TopK::<1>::new(candidates, thresholds); + let $topk = $crate::graph::pipnn::topk::TopK::<1>::new(width); $body } 2 => { - let mut $topk = $crate::graph::pipnn::topk::TopK::<2>::new(candidates, thresholds); + let $topk = $crate::graph::pipnn::topk::TopK::<2>::new(width); $body } 3 => { - let mut $topk = $crate::graph::pipnn::topk::TopK::<3>::new(candidates, thresholds); + let $topk = $crate::graph::pipnn::topk::TopK::<3>::new(width); $body } 10 => { - let mut $topk = $crate::graph::pipnn::topk::TopK::<10>::new(candidates, thresholds); + let $topk = $crate::graph::pipnn::topk::TopK::<10>::new(width); $body } _ => { - let mut $topk = $crate::graph::pipnn::topk::TopK::< + let $topk = $crate::graph::pipnn::topk::TopK::< { $crate::graph::pipnn::topk::RUNTIME_WIDTH }, - >::new(candidates, thresholds); + >::new(width); $body } } @@ -77,84 +76,85 @@ macro_rules! with_topk { } pub(super) use with_topk; -/// Nearest candidates and cached distance limits for one or more result sets. +/// Top-k selection with a fixed result capacity and caller-owned storage. /// -/// Each result set has the same capacity. This view borrows the candidate matrix -/// and one threshold per set; kernels provide their storage. Candidate IDs are -/// local indexes defined by the update operation, not dataset IDs. Updates do -/// not deduplicate: callers supply each candidate ID at most once per result set. -/// Equal distances need no fixed tie order. -/// `K` is the result capacity; `RUNTIME_WIDTH` uses the matrix's column count. -/// The construction macro selects the width once per batch. Kernels keep the -/// architecture scope active across their updates. -pub(super) struct TopK<'a, const K: usize> { - candidates: MutMatrixView<'a, Candidate>, - thresholds: &'a mut [f32], +/// This object stores no candidates or thresholds. Each operation borrows only +/// the buffers it updates. Candidate IDs are slice positions, not dataset IDs. +/// Callers supply each candidate at most once per result; equal distances need +/// no fixed tie order. The construction macro specializes the width per batch. +pub(super) struct TopK { + k: usize, } -impl<'a, const K: usize> TopK<'a, K> { - /// Bind and clear caller-owned result and threshold buffers. - /// A fixed `K` must equal the output width. - pub(super) fn new(candidates: MutMatrixView<'a, Candidate>, thresholds: &'a mut [f32]) -> Self { +impl TopK { + /// Configure the result capacity; a fixed K must match it. + pub(super) fn new(k: usize) -> Self { assert!( - K == RUNTIME_WIDTH || K == candidates.ncols(), - "top-k width must match its output" + K == RUNTIME_WIDTH || K == k, + "top-k width must match its capacity" ); - debug_assert_eq!( - candidates.nrows(), - thresholds.len(), - "top-k candidates and thresholds must match" - ); - let mut topk = Self { - candidates, - thresholds, - }; - topk.reset(); - topk + Self { k } } - fn reset(&mut self) { - self.candidates.as_mut_slice().fill(Candidate::default()); - self.thresholds.fill(f32::INFINITY); - } - - /// Return one result set in nearest-first order, including unassigned slots. - pub(super) fn candidates(&self, result_idx: usize) -> &[Candidate] { - self.candidates.row(result_idx) + /// Clear results and prepare reusable thresholds before a sequence of dual updates. + pub(super) fn initialize( + &self, + mut output: MutMatrixView<'_, Candidate>, + thresholds: &mut Vec, + ) { + debug_assert_eq!( + output.ncols(), + self.k, + "top-k output width must match its capacity" + ); + output.as_mut_slice().fill(Candidate::default()); + thresholds.resize(output.nrows(), f32::INFINITY); + thresholds.fill(f32::INFINITY); } - /// Replace one result set with the nearest candidates in the distance slice. - /// Slice positions become candidate IDs; `result_idx` only selects output storage. + /// Select the nearest candidates, replacing the previous contents of output. + /// Slice positions become candidate IDs. The threshold is local to this call. #[inline(always)] - pub(super) fn replace_topk( - &mut self, + pub(super) fn select_topk( + &self, arch: A, - result_idx: usize, distances: &[f32], + output: &mut [Candidate], ) { - let k = self.candidates.ncols(); - if k == 0 { + debug_assert_eq!( + output.len(), + self.k, + "top-k output width must match its capacity" + ); + if self.k == 0 { return; } - let nearest = &mut self.candidates.as_mut_slice()[result_idx * k..][..k]; - nearest.fill(Candidate::default()); - self.thresholds[result_idx] = distance_blocks(arch, distances).fold( + output.fill(Candidate::default()); + distance_blocks(arch, distances).fold( f32::INFINITY, #[inline(always)] - |limit, block| block.update_one(nearest, limit), + |limit, block| block.update_one(output, limit), ); } - /// Add each pair `(point_idx, j)` independently to both endpoints' top-k results. - /// Preserve existing candidates. All slice indexes must be less than `point_idx`. + /// Add each pair (point_idx, j) independently to both endpoints' top-k results. + /// Call [`Self::initialize`] once before the pair scan; subsequent updates preserve its state. + /// Preserve existing results; all slice indexes must be less than point_idx. #[inline] pub(super) fn update_dual_topk( - &mut self, + &self, arch: A, point_idx: usize, distances: &[f32], + mut output: MutMatrixView<'_, Candidate>, + thresholds: &mut [f32], ) { - let k = self.candidates.ncols(); + let k = self.k; + debug_assert_eq!( + output.ncols(), + k, + "top-k output width must match its capacity" + ); if k == 0 { return; } @@ -162,19 +162,24 @@ impl<'a, const K: usize> TopK<'a, K> { distances.len() <= point_idx, "candidate indexes must exclude the updated point" ); + debug_assert_eq!( + output.nrows(), + thresholds.len(), + "top-k candidates and thresholds must match" + ); // Keep this point's result and local limit independent of reciprocal updates. - let (others, remaining) = self.candidates.as_mut_slice().split_at_mut(point_idx * k); + let (others, remaining) = output.as_mut_slice().split_at_mut(point_idx * k); let nearest = &mut remaining[..k]; let limit = distance_blocks(arch, distances).fold( - self.thresholds[point_idx], + thresholds[point_idx], #[inline(always)] |limit, block| { let limit = block.update_one(nearest, limit); - block.update_many(others, self.thresholds, point_idx as u32, k); + block.update_many(others, thresholds, point_idx as u32, k); limit }, ); - self.thresholds[point_idx] = limit; + thresholds[point_idx] = limit; } } @@ -337,21 +342,41 @@ mod tests { use rstest::rstest; // These adapters only select the real CPU backend; all updates use production TopK. - struct ReplaceRow; - impl Target1, usize, &[f32])> - for ReplaceRow + struct SelectTopK; + impl Target1, &[f32], &mut [Candidate])> + for SelectTopK { - fn run(self, arch: A, (topk, row, distances): (&mut TopK<'_, K>, usize, &[f32])) { - topk.replace_topk(arch, row, distances); + fn run(self, arch: A, (topk, distances, output): (&TopK, &[f32], &mut [Candidate])) { + topk.select_topk(arch, distances, output); } } struct UpdatePair; - impl Target1, usize, &[f32])> - for UpdatePair + impl + Target1< + A, + (), + ( + &TopK, + usize, + &[f32], + MutMatrixView<'_, Candidate>, + &mut [f32], + ), + > for UpdatePair { - fn run(self, arch: A, (topk, point, distances): (&mut TopK<'_, K>, usize, &[f32])) { - topk.update_dual_topk(arch, point, distances); + fn run( + self, + arch: A, + (topk, point, distances, output, thresholds): ( + &TopK, + usize, + &[f32], + MutMatrixView<'_, Candidate>, + &mut [f32], + ), + ) { + topk.update_dual_topk(arch, point, distances, output, thresholds); } } @@ -368,54 +393,57 @@ mod tests { candidates } - mod new_tests { - use super::*; - - #[test] - fn reused_storage_is_cleared_before_dual_updates() { - let mut output = [ - Candidate::new(1, -10.0), - Candidate::new(2, -9.0), - Candidate::new(0, -10.0), - Candidate::new(2, -8.0), - Candidate::new(0, -9.0), - Candidate::new(1, -8.0), - ]; - let mut thresholds = [-9.0, -8.0, -8.0]; - let expected = [Candidate::new(1, 2.0), Candidate::default()]; - let mut topk = TopK::<2>::new( - MutMatrixView::try_from(&mut output[..], 3, 2).unwrap(), - &mut thresholds, - ); - - arch::dispatch1_no_features(UpdatePair, (&mut topk, 1, &[2.0][..])); + #[test] + fn initialize_clears_candidates_and_resizes_thresholds() { + // Given: candidate storage contains old results; the threshold buffer + // still has the smaller previous leaf's length. + let mut output = [ + Candidate::new(1, 1.0), + Candidate::new(2, 4.0), + Candidate::new(0, 1.0), + Candidate::new(2, 3.0), + Candidate::new(1, 3.0), + Candidate::new(0, 4.0), + ]; + let mut thresholds = vec![4.0, 3.0]; + let expected_output = [Candidate::default(); 6]; + let expected_thresholds = [f32::INFINITY; 3]; + let topk = TopK::<2>::new(2); + + // When + topk.initialize( + MutMatrixView::try_from(&mut output[..], 3, 2).unwrap(), + &mut thresholds, + ); - assert_eq!(topk.candidates(0), expected); - assert_eq!( - topk.candidates(1), - [Candidate::new(0, 2.0), Candidate::default()] - ); - assert_eq!(topk.candidates(2), [Candidate::default(); 2]); - } + // Then: clear both old thresholds and the newly added slot. + assert_eq!(output, expected_output); + assert_eq!(thresholds, expected_thresholds); } #[test] fn zero_capacity_accepts_both_update_operations() { let mut output = []; let mut thresholds = [f32::INFINITY; 2]; - let mut topk = TopK::::new( - MutMatrixView::try_from(&mut output[..], 2, 0).unwrap(), - &mut thresholds, + let topk = TopK::::new(0); + + arch::dispatch1_no_features(SelectTopK, (&topk, &[1.0, 2.0][..], &mut output[..])); + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 1, + &[0.0][..], + MutMatrixView::try_from(&mut output[..], 2, 0).unwrap(), + &mut thresholds[..], + ), ); - arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &[1.0, 2.0][..])); - arch::dispatch1_no_features(UpdatePair, (&mut topk, 1, &[0.0][..])); - - assert!(topk.candidates(0).is_empty()); - assert!(topk.candidates(1).is_empty()); + assert!(output.is_empty()); + assert_eq!(thresholds, [f32::INFINITY; 2]); } - mod replace_topk_tests { + mod select_topk_tests { use super::*; #[rstest] @@ -435,16 +463,14 @@ mod tests { .collect(); let expected = sorted_candidates(&distances, width); let mut output = vec![Candidate::default(); width]; - let mut thresholds = [f32::INFINITY]; - - with_topk!( - MutMatrixView::row_vector(output.as_mut_slice()), - &mut thresholds, - |topk| { - arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, distances.as_slice())); - assert_eq!(topk.candidates(0), expected); - } - ); + + with_topk!(width, |topk| { + arch::dispatch1_no_features( + SelectTopK, + (&topk, distances.as_slice(), output.as_mut_slice()), + ); + assert_eq!(output, expected); + }); } #[rstest] @@ -458,14 +484,12 @@ mod tests { #[case] expected: [Candidate; 2], ) { let mut output = [Candidate::default(); 2]; - let mut thresholds = [f32::INFINITY]; - let mut topk = - TopK::<2>::new(MutMatrixView::row_vector(&mut output[..]), &mut thresholds); - topk.replace_topk(diskann_wide::ARCH, 0, &[3.0, 1.0, 2.0]); + let topk = TopK::<2>::new(2); + topk.select_topk(diskann_wide::ARCH, &[3.0, 1.0, 2.0], &mut output); - arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, distances)); + arch::dispatch1_no_features(SelectTopK, (&topk, distances, &mut output[..])); - assert_eq!(topk.candidates(0), expected); + assert_eq!(output, expected); } #[test] @@ -477,26 +501,27 @@ mod tests { distances[16] = 1.0; distances[17] = 4.0; let mut output = [Candidate::default()]; - let mut thresholds = [f32::INFINITY]; - let mut topk = TopK::<1>::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + let topk = TopK::<1>::new(1); - arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &distances[..])); + arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); - assert_eq!(topk.candidates(0), [Candidate::new(16, 1.0)]); + assert_eq!(output, [Candidate::new(16, 1.0)]); } #[test] fn closer_candidates_shift_the_middle_of_a_full_result() { let mut output = [Candidate::default(); 3]; - let mut thresholds = [f32::INFINITY]; - let mut topk = TopK::<3>::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + let topk = TopK::<3>::new(3); // After retaining [0, 4, 6], distances 2 and 3 must enter between // existing candidates instead of replacing the last slot directly. - arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &[0.0, 4.0, 6.0, 2.0, 3.0][..])); + arch::dispatch1_no_features( + SelectTopK, + (&topk, &[0.0, 4.0, 6.0, 2.0, 3.0][..], &mut output[..]), + ); assert_eq!( - topk.candidates(0), + output, [ Candidate::new(0, 0.0), Candidate::new(3, 2.0), @@ -514,13 +539,12 @@ mod tests { distances[31] = 3.0; distances[32] = 1.0; let mut output = [Candidate::default(); 3]; - let mut thresholds = [f32::INFINITY]; - let mut topk = TopK::<3>::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + let topk = TopK::<3>::new(3); - arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &distances[..])); + arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); assert_eq!( - topk.candidates(0), + output, [ Candidate::new(32, 1.0), Candidate::new(17, 2.0), @@ -534,13 +558,12 @@ mod tests { let mut distances = [2.0; 17]; distances[16] = 1.0; let mut output = [Candidate::default(); 3]; - let mut thresholds = [f32::INFINITY]; - let mut topk = TopK::<3>::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + let topk = TopK::<3>::new(3); - arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &distances[..])); + arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); - assert_eq!(topk.candidates(0)[0], Candidate::new(16, 1.0)); - let tied = &topk.candidates(0)[1..]; + assert_eq!(output[0], Candidate::new(16, 1.0)); + let tied = &output[1..]; assert!( tied.iter() .all(|candidate| candidate.local_idx < 16 && candidate.distance == 2.0) @@ -556,16 +579,13 @@ mod tests { distances[17] = f32::MAX; distances[32] = f32::NEG_INFINITY; let mut output = [Candidate::default(); 4]; - let mut thresholds = [f32::INFINITY]; - let mut topk = - TopK::::new(MutMatrixView::row_vector(&mut output), &mut thresholds); + let topk = TopK::::new(4); - arch::dispatch1_no_features(ReplaceRow, (&mut topk, 0, &distances[..])); + arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); - assert_eq!(topk.candidates(0)[0], Candidate::new(32, f32::NEG_INFINITY)); + assert_eq!(output[0], Candidate::new(32, f32::NEG_INFINITY)); // Tie order is unspecified; each retained value must match its own ID. - let mut actual: Vec<_> = topk - .candidates(0) + let mut actual: Vec<_> = output .iter() .map(|candidate| (candidate.local_idx, candidate.distance.to_bits())) .collect(); @@ -595,17 +615,24 @@ mod tests { distances[32] = 0.75; let mut output = [Candidate::default(); 34]; let mut thresholds = [f32::INFINITY; 34]; - let mut topk = TopK::<1>::new( - MutMatrixView::try_from(&mut output[..], 34, 1).unwrap(), - &mut thresholds, - ); + let topk = TopK::<1>::new(1); + let mut rows = MutMatrixView::try_from(&mut output[..], 34, 1).unwrap(); // Earlier leaf rows establish each destination's threshold. let mut previous = [f32::INFINITY; 32]; for point in 1..33 { if point == 3 { previous[..3].copy_from_slice(&[0.2, 5.0, 5.0]); } - arch::dispatch1_no_features(UpdatePair, (&mut topk, point, &previous[..point])); + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + point, + &previous[..point], + rows.as_mut_view(), + &mut thresholds[..], + ), + ); previous.fill(f32::INFINITY); } let mut expected = [Candidate::default(); 34]; @@ -617,19 +644,26 @@ mod tests { expected[32] = Candidate::new(33, 0.75); expected[33] = Candidate::new(0, 0.5); - arch::dispatch1_no_features(UpdatePair, (&mut topk, 33, &distances[..])); + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 33, + &distances[..], + rows.as_mut_view(), + &mut thresholds[..], + ), + ); - assert_eq!(topk.candidates.as_slice(), expected); + assert_eq!(rows.as_slice(), expected); } #[test] fn a_source_fills_across_blocks_with_few_rankable_distances() { let mut output = [Candidate::default(); 34 * 3]; let mut thresholds = [f32::INFINITY; 34]; - let mut topk = TopK::<3>::new( - MutMatrixView::try_from(&mut output[..], 34, 3).unwrap(), - &mut thresholds, - ); + let topk = TopK::<3>::new(3); + let mut rows = MutMatrixView::try_from(&mut output[..], 34, 3).unwrap(); let mut distances = [f32::INFINITY; 33]; distances[..8].fill(f32::NAN); distances[15] = 4.0; @@ -640,10 +674,19 @@ mod tests { // A leaf source starts empty. It remains partially filled after // group one and acquires its finite threshold within group two. - arch::dispatch1_no_features(UpdatePair, (&mut topk, 33, &distances[..])); + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 33, + &distances[..], + rows.as_mut_view(), + &mut thresholds[..], + ), + ); assert_eq!( - topk.candidates(33), + rows.row(33), [ Candidate::new(32, f32::NEG_INFINITY), Candidate::new(16, 2.0), @@ -651,7 +694,7 @@ mod tests { ] ); assert_eq!( - topk.candidates(17), + rows.row(17), [ Candidate::new(33, 9.0), Candidate::default(), @@ -664,16 +707,32 @@ mod tests { fn later_pairs_use_the_sources_updated_threshold() { let mut output = [Candidate::default(); 3]; let mut thresholds = [f32::INFINITY; 3]; - let mut topk = TopK::<1>::new( - MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), - &mut thresholds, + let topk = TopK::<1>::new(1); + let mut rows = MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(); + + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 1, + &[1.0][..], + rows.as_mut_view(), + &mut thresholds[..], + ), + ); + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 2, + &[3.0, 2.0][..], + rows.as_mut_view(), + &mut thresholds[..], + ), ); - - arch::dispatch1_no_features(UpdatePair, (&mut topk, 1, &[1.0][..])); - arch::dispatch1_no_features(UpdatePair, (&mut topk, 2, &[3.0, 2.0][..])); assert_eq!( - topk.candidates.as_slice(), + rows.as_slice(), [ Candidate::new(1, 1.0), Candidate::new(0, 1.0), @@ -693,51 +752,63 @@ mod tests { distances[15] = 2.0; distances[16] = 1.0; - with_topk!( - MutMatrixView::try_from(&mut output[..], 18, 4).unwrap(), - &mut thresholds, - |topk| { - for point in 1..17 { - if point == 15 { - previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); - } - arch::dispatch1_no_features( - UpdatePair, - (&mut topk, point, &previous[..point]), - ); - previous.fill(f32::INFINITY); + let mut rows = MutMatrixView::try_from(&mut output[..], 18, 4).unwrap(); + with_topk!(4, |topk| { + for point in 1..17 { + if point == 15 { + previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); } - arch::dispatch1_no_features(UpdatePair, (&mut topk, 17, &distances[..])); - - assert_eq!( - topk.candidates(17), - [ - Candidate::new(16, 1.0), - Candidate::new(15, 2.0), - Candidate::new(2, 3.0), - Candidate::new(0, 4.0) - ] - ); - assert_eq!( - topk.candidates(15), - [ - Candidate::new(0, -1.0), - Candidate::new(17, 2.0), - Candidate::new(1, 3.0), - Candidate::new(2, 5.0) - ] - ); - assert_eq!( - topk.candidates(16), - [ - Candidate::new(17, 1.0), - Candidate::default(), - Candidate::default(), - Candidate::default() - ] + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + point, + &previous[..point], + rows.as_mut_view(), + &mut thresholds[..], + ), ); + previous.fill(f32::INFINITY); } - ); + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 17, + &distances[..], + rows.as_mut_view(), + &mut thresholds[..], + ), + ); + + assert_eq!( + rows.row(17), + [ + Candidate::new(16, 1.0), + Candidate::new(15, 2.0), + Candidate::new(2, 3.0), + Candidate::new(0, 4.0) + ] + ); + assert_eq!( + rows.row(15), + [ + Candidate::new(0, -1.0), + Candidate::new(17, 2.0), + Candidate::new(1, 3.0), + Candidate::new(2, 5.0) + ] + ); + assert_eq!( + rows.row(16), + [ + Candidate::new(17, 1.0), + Candidate::default(), + Candidate::default(), + Candidate::default() + ] + ); + }); } #[rstest] @@ -748,17 +819,24 @@ mod tests { distances[candidate] = f32::NEG_INFINITY; let mut output = [Candidate::default(); 18 * 2]; let mut thresholds = [f32::INFINITY; 18]; - let mut topk = TopK::<2>::new( - MutMatrixView::try_from(&mut output[..], 18, 2).unwrap(), - &mut thresholds, - ); + let topk = TopK::<2>::new(2); + let mut rows = MutMatrixView::try_from(&mut output[..], 18, 2).unwrap(); let mut expected = [[Candidate::default(); 2]; 18]; expected[17][0] = Candidate::new(candidate as u32, f32::NEG_INFINITY); expected[candidate][0] = Candidate::new(17, f32::NEG_INFINITY); - arch::dispatch1_no_features(UpdatePair, (&mut topk, 17, &distances[..])); + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 17, + &distances[..], + rows.as_mut_view(), + &mut thresholds[..], + ), + ); - assert_eq!(topk.candidates.as_slice(), expected.as_flattened()); + assert_eq!(rows.as_slice(), expected.as_flattened()); } } } From 78f5e95e15588c6c96b47e2da883e3736d52277c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 10 Sep 2026 06:32:33 +0000 Subject: [PATCH 172/229] fix(pipnn): keep SIMD scope inside top-k Enter the architecture scope per TopK operation with explicit buffer arguments so kernels only traverse rows. Preserve const capacity inside the feature boundary to keep sorted insertion specialized in release. --- diskann/src/graph/pipnn/leaf_kernel.rs | 24 ++--- diskann/src/graph/pipnn/partition_kernel.rs | 14 ++- diskann/src/graph/pipnn/topk.rs | 109 ++++++++++++-------- 3 files changed, 85 insertions(+), 62 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 9c4c3c5c7f..6827b6b9dd 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -96,19 +96,17 @@ fn rank_leaf_distances( mut output: MutMatrixView<'_, Candidate>, worst: &mut Vec, ) { - arch.run(move || { - with_topk!(output.ncols(), |topk| { - topk.initialize(output.as_mut_view(), worst); - for point_idx in 1..distances.nrows() { - topk.update_dual_topk( - arch, - point_idx, - &distances.row(point_idx)[..point_idx], - output.as_mut_view(), - worst.as_mut_slice(), - ); - } - }); + with_topk!(output.ncols(), |topk| { + topk.initialize(output.as_mut_view(), worst); + for point_idx in 1..distances.nrows() { + topk.update_dual_topk( + arch, + point_idx, + &distances.row(point_idx)[..point_idx], + output.as_mut_view(), + worst.as_mut_slice(), + ); + } }); } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index d6047cd261..78bbeb88f4 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -86,15 +86,13 @@ fn rank_leader_distances( candidates: &mut Vec, ) { candidates.resize(output.ncols(), Candidate::default()); - arch.run(move || { - with_topk!(output.ncols(), |topk| { - for (distances, output) in distances.row_iter().zip(output.row_iter_mut()) { - topk.select_topk(arch, distances, candidates.as_mut_slice()); - for (destination, candidate) in output.iter_mut().zip(candidates.iter()) { - *destination = candidate.local_idx; - } + with_topk!(output.ncols(), |topk| { + for (distances, output) in distances.row_iter().zip(output.row_iter_mut()) { + topk.select_topk(arch, distances, candidates.as_mut_slice()); + for (destination, candidate) in output.iter_mut().zip(candidates.iter()) { + *destination = candidate.local_idx; } - }); + } }); } diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 7c06fad69b..e2f1515b8b 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -82,6 +82,7 @@ pub(super) use with_topk; /// the buffers it updates. Candidate IDs are slice positions, not dataset IDs. /// Callers supply each candidate at most once per result; equal distances need /// no fixed tie order. The construction macro specializes the width per batch. +/// Selection and dual updates run inside their own architecture scope. pub(super) struct TopK { k: usize, } @@ -96,6 +97,12 @@ impl TopK { Self { k } } + // Keep the const width visible inside the target-feature boundary. + #[inline(always)] + fn capacity(&self) -> usize { + if K == RUNTIME_WIDTH { self.k } else { K } + } + /// Clear results and prepare reusable thresholds before a sequence of dual updates. pub(super) fn initialize( &self, @@ -121,19 +128,29 @@ impl TopK { distances: &[f32], output: &mut [Candidate], ) { - debug_assert_eq!( - output.len(), - self.k, - "top-k output width must match its capacity" - ); - if self.k == 0 { - return; - } - output.fill(Candidate::default()); - distance_blocks(arch, distances).fold( - f32::INFINITY, + arch.run2( #[inline(always)] - |limit, block| block.update_one(output, limit), + move |distances: &[f32], output: &mut [Candidate]| { + let k = self.capacity(); + debug_assert_eq!( + output.len(), + k, + "top-k output width must match its capacity" + ); + if k == 0 { + return; + } + // Preserve the specialized capacity through the inlined insertion. + let output = &mut output[..k]; + output.fill(Candidate::default()); + distance_blocks(arch, distances).fold( + f32::INFINITY, + #[inline(always)] + |limit, block| block.update_one(output, limit), + ); + }, + distances, + output, ); } @@ -146,40 +163,50 @@ impl TopK { arch: A, point_idx: usize, distances: &[f32], - mut output: MutMatrixView<'_, Candidate>, + output: MutMatrixView<'_, Candidate>, thresholds: &mut [f32], ) { - let k = self.k; - debug_assert_eq!( - output.ncols(), - k, - "top-k output width must match its capacity" - ); - if k == 0 { - return; - } - debug_assert!( - distances.len() <= point_idx, - "candidate indexes must exclude the updated point" - ); - debug_assert_eq!( - output.nrows(), - thresholds.len(), - "top-k candidates and thresholds must match" - ); - // Keep this point's result and local limit independent of reciprocal updates. - let (others, remaining) = output.as_mut_slice().split_at_mut(point_idx * k); - let nearest = &mut remaining[..k]; - let limit = distance_blocks(arch, distances).fold( - thresholds[point_idx], + arch.run3( #[inline(always)] - |limit, block| { - let limit = block.update_one(nearest, limit); - block.update_many(others, thresholds, point_idx as u32, k); - limit + move |distances: &[f32], + mut output: MutMatrixView<'_, Candidate>, + thresholds: &mut [f32]| { + let k = self.capacity(); + debug_assert_eq!( + output.ncols(), + k, + "top-k output width must match its capacity" + ); + if k == 0 { + return; + } + debug_assert!( + distances.len() <= point_idx, + "candidate indexes must exclude the updated point" + ); + debug_assert_eq!( + output.nrows(), + thresholds.len(), + "top-k candidates and thresholds must match" + ); + // Keep this point's result and local limit independent of reciprocal updates. + let (others, remaining) = output.as_mut_slice().split_at_mut(point_idx * k); + let nearest = &mut remaining[..k]; + let limit = distance_blocks(arch, distances).fold( + thresholds[point_idx], + #[inline(always)] + |limit, block| { + let limit = block.update_one(nearest, limit); + block.update_many(others, thresholds, point_idx as u32, k); + limit + }, + ); + thresholds[point_idx] = limit; }, + distances, + output, + thresholds, ); - thresholds[point_idx] = limit; } } From c627be81c7265602fbdb6a4f7bc062e306e533fe Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 10 Sep 2026 06:55:02 +0000 Subject: [PATCH 173/229] fix(ci): run relocated top-k tests under Miri --- .github/workflows/nightly.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index c29d60275f..58505d0028 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -147,4 +147,4 @@ jobs: cargo +nightly miri test --locked -p diskann --features pipnn --lib \ graph::pipnn::leaf_kernel::tests::rank_leaf_distances_tests cargo +nightly miri test --locked -p diskann --features pipnn --lib \ - graph::pipnn::partition_kernel::tests::rank_leader_distances_tests + graph::pipnn::topk From 6a8d13e53b83c73ec1f8eee8702616a0902af450 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 10 Sep 2026 09:07:46 +0000 Subject: [PATCH 174/229] perf(pipnn): specialize top-k widths through ten --- diskann/src/graph/pipnn/topk.rs | 122 ++++++++++++++++++++------------ 1 file changed, 76 insertions(+), 46 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index e2f1515b8b..fdc4626d8c 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -44,6 +44,8 @@ impl Default for Candidate { pub(super) const RUNTIME_WIDTH: usize = 0; /// Select the width specialization once for a batch. +/// Specialize small capacities 1..=10; larger capacities use the runtime path. +/// Typical partition K is 10 or 3, and leaf K is 3 or 2. /// The width is evaluated once. The body has ordinary block control flow. macro_rules! with_topk { ($width:expr, |$topk:ident| $body:block) => {{ @@ -61,6 +63,30 @@ macro_rules! with_topk { let $topk = $crate::graph::pipnn::topk::TopK::<3>::new(width); $body } + 4 => { + let $topk = $crate::graph::pipnn::topk::TopK::<4>::new(width); + $body + } + 5 => { + let $topk = $crate::graph::pipnn::topk::TopK::<5>::new(width); + $body + } + 6 => { + let $topk = $crate::graph::pipnn::topk::TopK::<6>::new(width); + $body + } + 7 => { + let $topk = $crate::graph::pipnn::topk::TopK::<7>::new(width); + $body + } + 8 => { + let $topk = $crate::graph::pipnn::topk::TopK::<8>::new(width); + $body + } + 9 => { + let $topk = $crate::graph::pipnn::topk::TopK::<9>::new(width); + $body + } 10 => { let $topk = $crate::graph::pipnn::topk::TopK::<10>::new(width); $body @@ -480,6 +506,11 @@ mod tests { #[case::two_lanes_minus_one(31, 1)] #[case::two_complete_lanes(32, 4)] #[case::two_lanes_and_tail(33, 17)] + #[case::capacity_five(33, 5)] + #[case::capacity_six(33, 6)] + #[case::capacity_seven(33, 7)] + #[case::capacity_eight(33, 8)] + #[case::capacity_nine(33, 9)] fn dispatched_selection_matches_independent_sort( #[case] count: usize, #[case] width: usize, @@ -780,62 +811,61 @@ mod tests { distances[16] = 1.0; let mut rows = MutMatrixView::try_from(&mut output[..], 18, 4).unwrap(); - with_topk!(4, |topk| { - for point in 1..17 { - if point == 15 { - previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); - } - arch::dispatch1_no_features( - UpdatePair, - ( - &topk, - point, - &previous[..point], - rows.as_mut_view(), - &mut thresholds[..], - ), - ); - previous.fill(f32::INFINITY); + let topk = TopK::::new(4); + for point in 1..17 { + if point == 15 { + previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); } arch::dispatch1_no_features( UpdatePair, ( &topk, - 17, - &distances[..], + point, + &previous[..point], rows.as_mut_view(), &mut thresholds[..], ), ); + previous.fill(f32::INFINITY); + } + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 17, + &distances[..], + rows.as_mut_view(), + &mut thresholds[..], + ), + ); - assert_eq!( - rows.row(17), - [ - Candidate::new(16, 1.0), - Candidate::new(15, 2.0), - Candidate::new(2, 3.0), - Candidate::new(0, 4.0) - ] - ); - assert_eq!( - rows.row(15), - [ - Candidate::new(0, -1.0), - Candidate::new(17, 2.0), - Candidate::new(1, 3.0), - Candidate::new(2, 5.0) - ] - ); - assert_eq!( - rows.row(16), - [ - Candidate::new(17, 1.0), - Candidate::default(), - Candidate::default(), - Candidate::default() - ] - ); - }); + assert_eq!( + rows.row(17), + [ + Candidate::new(16, 1.0), + Candidate::new(15, 2.0), + Candidate::new(2, 3.0), + Candidate::new(0, 4.0) + ] + ); + assert_eq!( + rows.row(15), + [ + Candidate::new(0, -1.0), + Candidate::new(17, 2.0), + Candidate::new(1, 3.0), + Candidate::new(2, 5.0) + ] + ); + assert_eq!( + rows.row(16), + [ + Candidate::new(17, 1.0), + Candidate::default(), + Candidate::default(), + Candidate::default() + ] + ); } #[rstest] From acda5846c5f84a7e805a265c85d25a7527f8387b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:03:20 +0000 Subject: [PATCH 175/229] refactor(pipnn): make insertion capacity explicit --- diskann/src/graph/pipnn/topk.rs | 54 +++++++++++++-------------------- 1 file changed, 21 insertions(+), 33 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index fdc4626d8c..f9af24493e 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -116,14 +116,14 @@ pub(super) struct TopK { impl TopK { /// Configure the result capacity; a fixed K must match it. pub(super) fn new(k: usize) -> Self { - assert!( + debug_assert!( K == RUNTIME_WIDTH || K == k, "top-k width must match its capacity" ); Self { k } } - // Keep the const width visible inside the target-feature boundary. + // Keep initialization and row addressing specialized across the architecture boundary. #[inline(always)] fn capacity(&self) -> usize { if K == RUNTIME_WIDTH { self.k } else { K } @@ -135,11 +135,6 @@ impl TopK { mut output: MutMatrixView<'_, Candidate>, thresholds: &mut Vec, ) { - debug_assert_eq!( - output.ncols(), - self.k, - "top-k output width must match its capacity" - ); output.as_mut_slice().fill(Candidate::default()); thresholds.resize(output.nrows(), f32::INFINITY); thresholds.fill(f32::INFINITY); @@ -157,22 +152,15 @@ impl TopK { arch.run2( #[inline(always)] move |distances: &[f32], output: &mut [Candidate]| { - let k = self.capacity(); - debug_assert_eq!( - output.len(), - k, - "top-k output width must match its capacity" - ); - if k == 0 { + let output = &mut output[..self.capacity()]; + if output.is_empty() { return; } - // Preserve the specialized capacity through the inlined insertion. - let output = &mut output[..k]; output.fill(Candidate::default()); distance_blocks(arch, distances).fold( f32::INFINITY, #[inline(always)] - |limit, block| block.update_one(output, limit), + |limit, block| block.update_one::(output, limit), ); }, distances, @@ -210,11 +198,6 @@ impl TopK { distances.len() <= point_idx, "candidate indexes must exclude the updated point" ); - debug_assert_eq!( - output.nrows(), - thresholds.len(), - "top-k candidates and thresholds must match" - ); // Keep this point's result and local limit independent of reciprocal updates. let (others, remaining) = output.as_mut_slice().split_at_mut(point_idx * k); let nearest = &mut remaining[..k]; @@ -222,8 +205,8 @@ impl TopK { thresholds[point_idx], #[inline(always)] |limit, block| { - let limit = block.update_one(nearest, limit); - block.update_many(others, thresholds, point_idx as u32, k); + let limit = block.update_one::(nearest, limit); + block.update_many::(others, thresholds, point_idx as u32, k); limit }, ); @@ -281,7 +264,7 @@ fn distance_blocks( impl DistanceBlock<'_, A> { /// Offer the block's candidates, preserving and returning the updated limit. #[inline(always)] - fn update_one(self, nearest: &mut [Candidate], mut max_distance: f32) -> f32 { + fn update_one(self, nearest: &mut [Candidate], mut max_distance: f32) -> f32 { match self { Self::Scalar { candidate_idx, @@ -289,7 +272,7 @@ impl DistanceBlock<'_, A> { } => { if distance < max_distance { max_distance = - insert_sorted(nearest, Candidate::new(candidate_idx as u32, distance)); + insert_sorted::(nearest, Candidate::new(candidate_idx as u32, distance)); } } Self::Simd { @@ -302,7 +285,7 @@ impl DistanceBlock<'_, A> { if max_distance == f32::INFINITY { for (lane, &distance) in distances.iter().enumerate() { if distance < max_distance { - max_distance = insert_sorted( + max_distance = insert_sorted::( nearest, Candidate::new((first_candidate + lane) as u32, distance), ); @@ -317,7 +300,7 @@ impl DistanceBlock<'_, A> { eligible &= eligible - 1; // Earlier insertions in this block can lower the limit. if distances[lane] < max_distance { - max_distance = insert_sorted( + max_distance = insert_sorted::( nearest, Candidate::new((first_candidate + lane) as u32, distances[lane]), ); @@ -331,7 +314,7 @@ impl DistanceBlock<'_, A> { /// Offer one point to each result set identified by this block's candidate IDs. #[inline(always)] - fn update_many( + fn update_many( self, candidates: &mut [Candidate], thresholds: &mut [f32], @@ -345,7 +328,7 @@ impl DistanceBlock<'_, A> { } => { let limit = &mut thresholds[candidate_idx]; if distance < *limit { - *limit = insert_sorted( + *limit = insert_sorted::( &mut candidates[candidate_idx * k..][..k], Candidate::new(point_idx, distance), ); @@ -364,7 +347,7 @@ impl DistanceBlock<'_, A> { let lane = eligible.trailing_zeros() as usize; eligible &= eligible - 1; // Each lane updates a different result; its limit is current. - limits[lane] = insert_sorted( + limits[lane] = insert_sorted::( &mut candidates[(first_candidate + lane) * k..][..k], Candidate::new(point_idx, distances[lane]), ); @@ -376,9 +359,14 @@ impl DistanceBlock<'_, A> { /// Insert an eligible candidate in nearest-first order and return the new distance limit. /// The caller must check the candidate against the current limit before insertion. +/// Fixed K exposes the insertion capacity to loop unrolling; runtime K uses the slice length. #[inline(always)] -fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate) -> f32 { - let last = nearest.len() - 1; +fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate) -> f32 { + let last = if K == RUNTIME_WIDTH { + nearest.len() - 1 + } else { + K - 1 + }; let mut slot = last; while slot > 0 && candidate.distance < nearest[slot - 1].distance { nearest[slot] = nearest[slot - 1]; From 5e4c08a83147e42770fdb8340bed5734bb7562c7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 16 Sep 2026 07:45:43 +0000 Subject: [PATCH 176/229] refactor(pipnn): model top-k width with types --- diskann/src/graph/pipnn/topk.rs | 179 +++++++++++++++++++------------- 1 file changed, 107 insertions(+), 72 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index f9af24493e..9012caff27 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -40,8 +40,32 @@ impl Default for Candidate { } } -// Fixed widths preserve row lengths for inlined insertion. Zero selects runtime width. -pub(super) const RUNTIME_WIDTH: usize = 0; +/// A result capacity known at compile time. +#[derive(Clone, Copy, Debug)] +pub(super) struct Fixed; + +/// A result capacity selected at runtime. +#[derive(Clone, Copy, Debug)] +pub(super) struct Runtime(pub(super) usize); + +/// Keep fixed capacities available to inlined insertion across architecture scopes. +pub(super) trait Width: Copy { + fn capacity(self) -> usize; +} + +impl Width for Fixed { + #[inline(always)] + fn capacity(self) -> usize { + K + } +} + +impl Width for Runtime { + #[inline(always)] + fn capacity(self) -> usize { + self.0 + } +} /// Select the width specialization once for a batch. /// Specialize small capacities 1..=10; larger capacities use the runtime path. @@ -52,49 +76,59 @@ macro_rules! with_topk { let width = $width; match width { 1 => { - let $topk = $crate::graph::pipnn::topk::TopK::<1>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<1>); $body } 2 => { - let $topk = $crate::graph::pipnn::topk::TopK::<2>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<2>); $body } 3 => { - let $topk = $crate::graph::pipnn::topk::TopK::<3>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<3>); $body } 4 => { - let $topk = $crate::graph::pipnn::topk::TopK::<4>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<4>); $body } 5 => { - let $topk = $crate::graph::pipnn::topk::TopK::<5>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<5>); $body } 6 => { - let $topk = $crate::graph::pipnn::topk::TopK::<6>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<6>); $body } 7 => { - let $topk = $crate::graph::pipnn::topk::TopK::<7>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<7>); $body } 8 => { - let $topk = $crate::graph::pipnn::topk::TopK::<8>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<8>); $body } 9 => { - let $topk = $crate::graph::pipnn::topk::TopK::<9>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<9>); $body } 10 => { - let $topk = $crate::graph::pipnn::topk::TopK::<10>::new(width); + let $topk = + $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<10>); $body } _ => { - let $topk = $crate::graph::pipnn::topk::TopK::< - { $crate::graph::pipnn::topk::RUNTIME_WIDTH }, - >::new(width); + let $topk = $crate::graph::pipnn::topk::TopK::new( + $crate::graph::pipnn::topk::Runtime(width), + ); $body } } @@ -109,24 +143,14 @@ pub(super) use with_topk; /// Callers supply each candidate at most once per result; equal distances need /// no fixed tie order. The construction macro specializes the width per batch. /// Selection and dual updates run inside their own architecture scope. -pub(super) struct TopK { - k: usize, +pub(super) struct TopK { + width: W, } -impl TopK { - /// Configure the result capacity; a fixed K must match it. - pub(super) fn new(k: usize) -> Self { - debug_assert!( - K == RUNTIME_WIDTH || K == k, - "top-k width must match its capacity" - ); - Self { k } - } - - // Keep initialization and row addressing specialized across the architecture boundary. - #[inline(always)] - fn capacity(&self) -> usize { - if K == RUNTIME_WIDTH { self.k } else { K } +impl TopK { + /// Configure the result capacity. + pub(super) fn new(width: W) -> Self { + Self { width } } /// Clear results and prepare reusable thresholds before a sequence of dual updates. @@ -152,7 +176,7 @@ impl TopK { arch.run2( #[inline(always)] move |distances: &[f32], output: &mut [Candidate]| { - let output = &mut output[..self.capacity()]; + let output = &mut output[..self.width.capacity()]; if output.is_empty() { return; } @@ -160,7 +184,7 @@ impl TopK { distance_blocks(arch, distances).fold( f32::INFINITY, #[inline(always)] - |limit, block| block.update_one::(output, limit), + |limit, block| block.update_one(output, limit, self.width), ); }, distances, @@ -185,7 +209,7 @@ impl TopK { move |distances: &[f32], mut output: MutMatrixView<'_, Candidate>, thresholds: &mut [f32]| { - let k = self.capacity(); + let k = self.width.capacity(); debug_assert_eq!( output.ncols(), k, @@ -205,8 +229,8 @@ impl TopK { thresholds[point_idx], #[inline(always)] |limit, block| { - let limit = block.update_one::(nearest, limit); - block.update_many::(others, thresholds, point_idx as u32, k); + let limit = block.update_one(nearest, limit, self.width); + block.update_many(others, thresholds, point_idx as u32, self.width); limit }, ); @@ -264,15 +288,23 @@ fn distance_blocks( impl DistanceBlock<'_, A> { /// Offer the block's candidates, preserving and returning the updated limit. #[inline(always)] - fn update_one(self, nearest: &mut [Candidate], mut max_distance: f32) -> f32 { + fn update_one( + self, + nearest: &mut [Candidate], + mut max_distance: f32, + width: W, + ) -> f32 { match self { Self::Scalar { candidate_idx, distance, } => { if distance < max_distance { - max_distance = - insert_sorted::(nearest, Candidate::new(candidate_idx as u32, distance)); + max_distance = insert_sorted( + nearest, + Candidate::new(candidate_idx as u32, distance), + width, + ); } } Self::Simd { @@ -285,9 +317,10 @@ impl DistanceBlock<'_, A> { if max_distance == f32::INFINITY { for (lane, &distance) in distances.iter().enumerate() { if distance < max_distance { - max_distance = insert_sorted::( + max_distance = insert_sorted( nearest, Candidate::new((first_candidate + lane) as u32, distance), + width, ); } } @@ -300,9 +333,10 @@ impl DistanceBlock<'_, A> { eligible &= eligible - 1; // Earlier insertions in this block can lower the limit. if distances[lane] < max_distance { - max_distance = insert_sorted::( + max_distance = insert_sorted( nearest, Candidate::new((first_candidate + lane) as u32, distances[lane]), + width, ); } } @@ -314,13 +348,14 @@ impl DistanceBlock<'_, A> { /// Offer one point to each result set identified by this block's candidate IDs. #[inline(always)] - fn update_many( + fn update_many( self, candidates: &mut [Candidate], thresholds: &mut [f32], point_idx: u32, - k: usize, + width: W, ) { + let k = width.capacity(); match self { Self::Scalar { candidate_idx, @@ -328,9 +363,10 @@ impl DistanceBlock<'_, A> { } => { let limit = &mut thresholds[candidate_idx]; if distance < *limit { - *limit = insert_sorted::( + *limit = insert_sorted( &mut candidates[candidate_idx * k..][..k], Candidate::new(point_idx, distance), + width, ); } } @@ -347,9 +383,10 @@ impl DistanceBlock<'_, A> { let lane = eligible.trailing_zeros() as usize; eligible &= eligible - 1; // Each lane updates a different result; its limit is current. - limits[lane] = insert_sorted::( + limits[lane] = insert_sorted( &mut candidates[(first_candidate + lane) * k..][..k], Candidate::new(point_idx, distances[lane]), + width, ); } } @@ -359,14 +396,10 @@ impl DistanceBlock<'_, A> { /// Insert an eligible candidate in nearest-first order and return the new distance limit. /// The caller must check the candidate against the current limit before insertion. -/// Fixed K exposes the insertion capacity to loop unrolling; runtime K uses the slice length. +/// The width exposes fixed capacities to loop unrolling. #[inline(always)] -fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate) -> f32 { - let last = if K == RUNTIME_WIDTH { - nearest.len() - 1 - } else { - K - 1 - }; +fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate, width: W) -> f32 { + let last = width.capacity() - 1; let mut slot = last; while slot > 0 && candidate.distance < nearest[slot - 1].distance { nearest[slot] = nearest[slot - 1]; @@ -384,21 +417,21 @@ mod tests { // These adapters only select the real CPU backend; all updates use production TopK. struct SelectTopK; - impl Target1, &[f32], &mut [Candidate])> + impl Target1, &[f32], &mut [Candidate])> for SelectTopK { - fn run(self, arch: A, (topk, distances, output): (&TopK, &[f32], &mut [Candidate])) { + fn run(self, arch: A, (topk, distances, output): (&TopK, &[f32], &mut [Candidate])) { topk.select_topk(arch, distances, output); } } struct UpdatePair; - impl + impl Target1< A, (), ( - &TopK, + &TopK, usize, &[f32], MutMatrixView<'_, Candidate>, @@ -410,7 +443,7 @@ mod tests { self, arch: A, (topk, point, distances, output, thresholds): ( - &TopK, + &TopK, usize, &[f32], MutMatrixView<'_, Candidate>, @@ -449,7 +482,7 @@ mod tests { let mut thresholds = vec![4.0, 3.0]; let expected_output = [Candidate::default(); 6]; let expected_thresholds = [f32::INFINITY; 3]; - let topk = TopK::<2>::new(2); + let topk = TopK::new(Fixed::<2>); // When topk.initialize( @@ -462,11 +495,13 @@ mod tests { assert_eq!(thresholds, expected_thresholds); } - #[test] - fn zero_capacity_accepts_both_update_operations() { + #[rstest] + #[case::fixed(Fixed::<0>)] + #[case::runtime(Runtime(0))] + fn zero_capacity_accepts_both_update_operations(#[case] width: impl Width) { let mut output = []; let mut thresholds = [f32::INFINITY; 2]; - let topk = TopK::::new(0); + let topk = TopK::new(width); arch::dispatch1_no_features(SelectTopK, (&topk, &[1.0, 2.0][..], &mut output[..])); arch::dispatch1_no_features( @@ -530,7 +565,7 @@ mod tests { #[case] expected: [Candidate; 2], ) { let mut output = [Candidate::default(); 2]; - let topk = TopK::<2>::new(2); + let topk = TopK::new(Fixed::<2>); topk.select_topk(diskann_wide::ARCH, &[3.0, 1.0, 2.0], &mut output); arch::dispatch1_no_features(SelectTopK, (&topk, distances, &mut output[..])); @@ -547,7 +582,7 @@ mod tests { distances[16] = 1.0; distances[17] = 4.0; let mut output = [Candidate::default()]; - let topk = TopK::<1>::new(1); + let topk = TopK::new(Fixed::<1>); arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); @@ -557,7 +592,7 @@ mod tests { #[test] fn closer_candidates_shift_the_middle_of_a_full_result() { let mut output = [Candidate::default(); 3]; - let topk = TopK::<3>::new(3); + let topk = TopK::new(Fixed::<3>); // After retaining [0, 4, 6], distances 2 and 3 must enter between // existing candidates instead of replacing the last slot directly. @@ -585,7 +620,7 @@ mod tests { distances[31] = 3.0; distances[32] = 1.0; let mut output = [Candidate::default(); 3]; - let topk = TopK::<3>::new(3); + let topk = TopK::new(Fixed::<3>); arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); @@ -604,7 +639,7 @@ mod tests { let mut distances = [2.0; 17]; distances[16] = 1.0; let mut output = [Candidate::default(); 3]; - let topk = TopK::<3>::new(3); + let topk = TopK::new(Fixed::<3>); arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); @@ -625,7 +660,7 @@ mod tests { distances[17] = f32::MAX; distances[32] = f32::NEG_INFINITY; let mut output = [Candidate::default(); 4]; - let topk = TopK::::new(4); + let topk = TopK::new(Runtime(4)); arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); @@ -661,7 +696,7 @@ mod tests { distances[32] = 0.75; let mut output = [Candidate::default(); 34]; let mut thresholds = [f32::INFINITY; 34]; - let topk = TopK::<1>::new(1); + let topk = TopK::new(Fixed::<1>); let mut rows = MutMatrixView::try_from(&mut output[..], 34, 1).unwrap(); // Earlier leaf rows establish each destination's threshold. let mut previous = [f32::INFINITY; 32]; @@ -708,7 +743,7 @@ mod tests { fn a_source_fills_across_blocks_with_few_rankable_distances() { let mut output = [Candidate::default(); 34 * 3]; let mut thresholds = [f32::INFINITY; 34]; - let topk = TopK::<3>::new(3); + let topk = TopK::new(Fixed::<3>); let mut rows = MutMatrixView::try_from(&mut output[..], 34, 3).unwrap(); let mut distances = [f32::INFINITY; 33]; distances[..8].fill(f32::NAN); @@ -753,7 +788,7 @@ mod tests { fn later_pairs_use_the_sources_updated_threshold() { let mut output = [Candidate::default(); 3]; let mut thresholds = [f32::INFINITY; 3]; - let topk = TopK::<1>::new(1); + let topk = TopK::new(Fixed::<1>); let mut rows = MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(); arch::dispatch1_no_features( @@ -799,7 +834,7 @@ mod tests { distances[16] = 1.0; let mut rows = MutMatrixView::try_from(&mut output[..], 18, 4).unwrap(); - let topk = TopK::::new(4); + let topk = TopK::new(Runtime(4)); for point in 1..17 { if point == 15 { previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); @@ -864,7 +899,7 @@ mod tests { distances[candidate] = f32::NEG_INFINITY; let mut output = [Candidate::default(); 18 * 2]; let mut thresholds = [f32::INFINITY; 18]; - let topk = TopK::<2>::new(2); + let topk = TopK::new(Fixed::<2>); let mut rows = MutMatrixView::try_from(&mut output[..], 18, 2).unwrap(); let mut expected = [[Candidate::default(); 2]; 18]; expected[17][0] = Candidate::new(candidate as u32, f32::NEG_INFINITY); From 31ef4a34201165620d6b369ad2b64b5ef127d095 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:58:56 +0000 Subject: [PATCH 177/229] perf(pipnn): simplify core and dispatch FP16 Use runtime-dispatched SliceCast for leaf and partition FP16 gathers. Simplify core ownership and validation, retain leaf scratch allocations, and replace overlapping tests with focused graph cases. --- diskann/src/graph/pipnn/conversion.rs | 90 +++ diskann/src/graph/pipnn/finalization.rs | 348 +++++---- diskann/src/graph/pipnn/leaf_build.rs | 535 ++++---------- diskann/src/graph/pipnn/mod.rs | 519 +++++++------ diskann/src/graph/pipnn/partitioning.rs | 928 +++++++++++++----------- 5 files changed, 1166 insertions(+), 1254 deletions(-) create mode 100644 diskann/src/graph/pipnn/conversion.rs diff --git a/diskann/src/graph/pipnn/conversion.rs b/diskann/src/graph/pipnn/conversion.rs new file mode 100644 index 0000000000..9c93a88223 --- /dev/null +++ b/diskann/src/graph/pipnn/conversion.rs @@ -0,0 +1,90 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Convert gathered vectors with the available CPU's FP16 instructions. + +use std::any::TypeId; + +use crate::utils::VectorRepr; +use diskann_vector::conversion::SliceCast; +use diskann_wide::arch; +use half::f16; + +/// Preserve VectorRepr conversion errors while accelerating FP16 slices. +#[inline] +pub(super) fn as_f32_into(source: &[T], output: &mut [f32]) -> Result<(), T::Error> { + if TypeId::of::() == TypeId::of::() && source.len() == output.len() { + // VectorRepr's default conversion uses the compile-time architecture. + // Dispatch one whole slice so portable builds can use F16C or NEON too. + let source: &[f16] = bytemuck::cast_slice(source); + arch::dispatch2(SliceCast::::new(), output, source); + Ok(()) + } else { + T::as_f32_into(source, output) + } +} + +#[cfg(test)] +mod as_f32_into_tests { + use super::*; + use rstest::rstest; + + #[test] + fn fp16_bits_are_preserved_through_a_vector_and_tail() { + // Given: 13 values leave a tail for both four- and eight-lane converters. + // Normal values gain 112 in exponent bias and shift the fraction by 13; + // the smallest FP16 subnormal becomes 2^-24. NaN keeps its payload. + let half_bits = [ + 0x0000, 0x8000, 0x0001, 0x03ff, 0x0400, 0x3c00, 0xbc00, 0x3555, 0x7bff, 0x7c00, 0xfc00, + 0x7e01, 0x8400, + ]; + let expected = [ + 0x00000000, 0x80000000, 0x33800000, 0x387fc000, 0x38800000, 0x3f800000, 0xbf800000, + 0x3eaaa000, 0x477fe000, 0x7f800000, 0xff800000, 0x7fc02000, 0xb8800000, + ]; + let source = half_bits.map(f16::from_bits); + let mut output = [123.0_f32; 13]; + + // When + as_f32_into(&source, &mut output).unwrap(); + + // Then + assert_eq!(output.map(f32::to_bits), expected); + } + + #[rstest] + #[case::short_output(7)] + #[case::long_output(9)] + fn mismatched_fp16_lengths_return_the_vector_repr_error(#[case] output_len: usize) { + // Given + let source = [f16::ONE; 8]; + let mut output = vec![0.0_f32; output_len]; + let expected = f16::as_f32_into(&source, &mut output).unwrap_err(); + + // When + let actual = as_f32_into(&source, &mut output).unwrap_err(); + + // Then + assert_eq!(actual, expected); + } + + #[rstest] + #[case::f32([-2.5_f32, 1.25], [-2.5, 1.25])] + #[case::i8([-2_i8, 127], [-2.0, 127.0])] + #[case::u8([0_u8, 255], [0.0, 255.0])] + fn other_representations_keep_their_conversion( + #[case] source: [T; 2], + #[case] expected: [f32; 2], + ) { + // Given + let mut output = [f32::NAN; 2]; + + // When + as_f32_into(&source, &mut output).unwrap(); + + // Then + assert_eq!(output, expected); + } +} diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index d4ec402774..4c5ef60fc1 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -5,15 +5,12 @@ //! Graph-degree enforcement with the Vamana RobustPrune kernel. //! -//! Candidate merging can produce more than `R` IDs for one point. This module -//! checks every global ID before parallel work starts. A list at or below `R` -//! returns without distance calculations. +//! Candidate merging can produce more than the graph degree for one point. +//! Lists within the degree limit pass through unchanged. Longer lists are sorted +//! by source distance and passed to RobustPrune, then overwritten with the +//! selected IDs. Reusable per-job scratch stores candidates and pruning state. //! -//! For a longer list, the module computes each source distance. It sorts the -//! candidates and calls RobustPrune. The module then writes the selected IDs into -//! the original list allocation. -//! -//! RobustPrune defines occlusion and alpha-round behavior. This module supplies +//! RobustPrune owns occlusion and alpha-round behavior. This module supplies //! source vectors and metric distances. use crate::{ @@ -36,112 +33,100 @@ pub(crate) enum FinalizationError { } /// RobustPrune state for one Rayon job. -/// -/// `sorted_cache` and `prune_states` stay positionally aligned with -/// `sorted_candidates`. #[derive(Default)] struct PruneWorkspace { - sorted_candidates: Vec>, - sorted_cache: Vec>>, - prune_states: Vec, + candidates: Vec>, + mapped_candidates: Vec>>, + states: Vec, +} + +impl PruneWorkspace { + fn prune( + &mut self, + data: MatrixView<'_, T>, + source: usize, + mut source_candidates: AdjacencyList, + graph: &Config, + distance: &T::Distance, + ) -> ANNResult> { + let degree = graph.pruned_degree().get(); + let candidate_count = source_candidates.len(); + if candidate_count <= degree { + return Ok(source_candidates); + } + if candidate_count > u16::MAX as usize { + return Err(ANNError::new(FinalizationError::TooManyCandidates { + actual: candidate_count, + max: u16::MAX as usize, + })); + } + + let source_vector = data.row(source); + self.candidates.clear(); + self.candidates + .extend(source_candidates.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + // Sort before mapping: changing the element type can change the unstable + // sort's order for equal distances, and therefore RobustPrune's choices. + let sorted = SortedNeighbors::new(&mut self.candidates, candidate_count); + let mapped = sorted.map_in(&mut self.mapped_candidates, |&candidate| { + (candidate as usize != source).then_some(candidate) + }); + self.states.clear(); + self.states.resize(candidate_count, prune::State::default()); + let selected = prune::robust_prune( + mapped, + &mut self.states, + degree, + graph.alpha(), + graph.prune_kind(), + |left, right| { + distance.evaluate_similarity(data.row(*left as usize), data.row(*right as usize)) + }, + ); + + let mut output = source_candidates.resize(selected); + for (destination, state) in output.iter_mut().zip(&self.states) { + *destination = *sorted[state.neighbor as usize].id(); + } + output.finish(selected); + Ok(source_candidates) + } } /// Prune each candidate list that exceeds the graph degree. /// /// Candidate builders supply one list per data row and valid dataset IDs. -pub(crate) fn prune_overfull( +pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, graph: &Config, metric: Metric, -) -> ANNResult>> -where - T: VectorRepr + Send + Sync, -{ - let degree = graph.pruned_degree().get(); +) -> ANNResult>> { let distance = T::distance(metric, Some(data.ncols())); - // `build_graph` runs this Rayon operation in the pool from the build context. + // The build context supplies the Rayon pool for this terminal operation. #[allow(clippy::disallowed_methods)] candidates .into_par_iter() .enumerate() .map_init( PruneWorkspace::default, - |workspace, (source, mut source_candidates)| { - // Candidate merging already removes duplicate IDs. A list within - // the degree limit needs no distance calculation. - if source_candidates.len() <= degree { - return Ok(source_candidates); - } - - let source_id = u32::try_from(source).map_err(ANNError::new)?; - let source_vector = data.row(source); - workspace.sorted_candidates.clear(); - workspace - .sorted_candidates - .extend(source_candidates.iter().copied().map(|candidate| { - Neighbor::new( - candidate, - distance - .evaluate_similarity(source_vector, data.row(candidate as usize)), - ) - })); - - let candidate_count = workspace.sorted_candidates.len(); - if candidate_count > u16::MAX as usize { - return Err(ANNError::new(FinalizationError::TooManyCandidates { - actual: candidate_count, - max: u16::MAX as usize, - })); - } - // Sort all candidates before the code marks a self-edge as absent. - // Thus, self-edge removal cannot add a farther candidate. The - // mapped list keeps the same positions for RobustPrune states. - let sorted = - SortedNeighbors::new(&mut workspace.sorted_candidates, candidate_count); - let sorted_cache = sorted.map_in(&mut workspace.sorted_cache, |id| { - (*id != source_id).then_some(*id) - }); - workspace - .prune_states - .resize(sorted_cache.len(), prune::State::default()); - // Each candidate list starts a separate RobustPrune state machine. - // Reset retained entries because resize initializes only new entries. - workspace.prune_states.fill(prune::State::default()); - - let selected = prune::robust_prune( - sorted_cache, - workspace.prune_states.as_mut_slice(), - degree, - graph.alpha(), - graph.prune_kind(), - |left, right| { - distance.evaluate_similarity( - data.row(*left as usize), - data.row(*right as usize), - ) - }, - ); - - let mut guard = source_candidates.resize(selected); - for (destination, state) in guard.iter_mut().zip(workspace.prune_states.iter()) { - *destination = *sorted[state.neighbor as usize].id(); - } - guard.finish(selected); - Ok(source_candidates) + |workspace, (source, candidates)| { + workspace.prune(data, source, candidates, graph, &distance) }, ) .collect() } #[cfg(test)] -mod tests { - use crate::graph::{ - AdjacencyList, - config::{self, MaxDegree}, - }; - use diskann_utils::views::MatrixView; +mod prune_tests { + use crate::graph::config::{self, MaxDegree}; + use rstest::rstest; use super::*; @@ -159,96 +144,153 @@ mod tests { .unwrap() } - fn candidate_list(ids: impl IntoIterator) -> AdjacencyList { - AdjacencyList::from_iter_untrusted(ids) + #[rstest] + #[case::empty(&[])] + #[case::below_degree(&[3])] + #[case::at_degree(&[3, 1])] + fn rows_within_degree_are_unchanged(#[case] ids: &[u32]) { + // Given: preserve the caller's order, including a non-sorted row. + let values = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::column_vector(&values[..]); + let mut candidates = AdjacencyList::new(); + candidates.overwrite_trusted(ids); + let distance = f32::distance(Metric::L2, Some(1)); + let mut workspace = PruneWorkspace::default(); + + // When + let actual = workspace + .prune(data, 0, candidates, &graph_config(2), &distance) + .unwrap(); + + // Then + assert_eq!(&*actual, ids); } #[test] - fn row_within_the_degree_bound_is_only_canonicalized() { - // Given - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + fn overfull_rows_keep_nearest_unoccluded_neighbors_at_their_source_id() { + // Given: source 1 is at 0. Its squared distances to IDs 2, 0, 3 are + // 1, 4, 9. ID 2 occludes ID 0 because 4 / (2 - 1)^2 = 4 > alpha. + // ID 3 survives because 9 / (-3 - 1)^2 = 9/16 <= 1. + let values = [2.0_f32, 0.0, 1.0, -3.0]; + let data = MatrixView::column_vector(&values[..]); let candidates = vec![ - candidate_list([3, 1]), - candidate_list([]), - candidate_list([]), - candidate_list([]), + AdjacencyList::new(), + AdjacencyList::from_iter_untrusted([0, 2, 3]), + AdjacencyList::new(), + AdjacencyList::new(), ]; - let expected_canonical_row = [1, 3]; + let expected = [vec![], vec![2, 3], vec![], vec![]]; // When - let actual_rows = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + let actual: Vec> = prune_overfull(data, candidates, &graph_config(2), Metric::L2) + .unwrap() + .into_iter() + .map(Vec::from) + .collect(); // Then - assert_eq!(&*actual_rows[0], &expected_canonical_row); + assert_eq!(actual, expected); } #[test] - fn overfull_row_keeps_the_nearest_unoccluded_neighbors() { - // Given - let point_values = [ - 0.0_f32, // source - 1.0, // nearest point on the positive ray - 2.0, // farther point occluded by point 1 - -3.0, // point on the opposite ray - ]; - let data = MatrixView::try_from(&point_values[..], 4, 1).unwrap(); - let candidates = vec![ - candidate_list([3, 2, 1]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - let expected_pruned_row = [1, 3]; + fn self_candidates_do_not_shift_selected_neighbor_ids() { + // Given: self occupies the first sorted position but cannot be selected. + // Of the other points, 1 occludes 2, while -3 remains on the opposite ray. + let values = [0.0_f32, 1.0, 2.0, -3.0]; + let data = MatrixView::column_vector(&values[..]); + let candidates = AdjacencyList::from_iter_untrusted([0, 1, 2, 3]); + let distance = f32::distance(Metric::L2, Some(1)); + let expected = [1, 3]; + let mut workspace = PruneWorkspace::default(); // When - let actual_rows = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + let actual = workspace + .prune(data, 0, candidates, &graph_config(2), &distance) + .unwrap(); // Then - assert_eq!(&*actual_rows[0], &expected_pruned_row); + assert_eq!(&*actual, &expected); } #[test] - fn reused_workspace_matches_fresh_pruning() { - // Given - let data = [0.0_f32, 1.0, 2.0, -3.0, 4.0]; - let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); - let first = [3, 2, 1]; - let second = [4, 3, 2]; - let candidates = |first: &[u32], second: &[u32]| { - vec![ - candidate_list(first.iter().copied()), - candidate_list(second.iter().copied()), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ] - }; + fn reused_workspace_clears_previous_candidates_and_occlusion_states() { + // Given: the same workspace prunes four candidates, then three, then four. + let values = [0.0_f32, 1.0, 2.0, -3.0, 4.0]; + let data = MatrixView::column_vector(&values[..]); let graph = graph_config(2); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() + let distance = f32::distance(Metric::L2, Some(1)); + let mut workspace = PruneWorkspace::default(); + workspace + .prune( + data, + 0, + AdjacencyList::from_iter_untrusted([1, 2, 3, 4]), + &graph, + &distance, + ) .unwrap(); - // When - let expected_first_row_from_fresh_workspace = pool - .install(|| prune_overfull(data, candidates(&first, &[]), &graph, Metric::L2)) + workspace + .prune( + data, + 4, + AdjacencyList::from_iter_untrusted([0, 1, 3]), + &graph, + &distance, + ) .unwrap(); - let expected_second_row_from_fresh_workspace = pool - .install(|| prune_overfull(data, candidates(&[], &second), &graph, Metric::L2)) + // For source 2 at coordinate 2, ID 1 at 1 is nearest. It occludes ID 0 + // (4/1 > alpha); ID 4 at 4 survives (4/9 <= 1). + let expected = [1, 4]; + let candidates = AdjacencyList::from_iter_untrusted([0, 1, 3, 4]); + + // When + let actual = workspace + .prune(data, 2, candidates, &graph, &distance) .unwrap(); - let actual_rows_from_reused_workspace = pool - .install(|| prune_overfull(data, candidates(&first, &second), &graph, Metric::L2)) + // Then + assert_eq!(&*actual, &expected); + } + + #[test] + fn maximum_u16_candidate_count_can_be_pruned() { + // Given: all candidates lie on the positive ray, so ID 1 is nearest. + let count = u16::MAX as usize; + let values: Vec = (0..=count).map(|id| id as f32).collect(); + let data = MatrixView::column_vector(&values[..]); + let candidates = AdjacencyList::from_iter_untrusted(1..=count as u32); + let distance = f32::distance(Metric::L2, Some(1)); + let mut workspace = PruneWorkspace::default(); + + // When + let actual = workspace + .prune(data, 0, candidates, &graph_config(1), &distance) .unwrap(); // Then - assert_eq!( - &*actual_rows_from_reused_workspace[0], - &*expected_first_row_from_fresh_workspace[0] - ); - assert_eq!( - &*actual_rows_from_reused_workspace[1], - &*expected_second_row_from_fresh_workspace[1] - ); + assert_eq!(&*actual, &[1]); + } + + #[test] + fn candidate_count_above_u16_limit_is_rejected() { + // Given: the list is one entry longer than RobustPrune's position limit. + let count = u16::MAX as usize + 1; + let values = vec![0.0_f32; count + 1]; + let data = MatrixView::column_vector(&values[..]); + let candidates = AdjacencyList::from_iter_untrusted(1..=count as u32); + let distance = f32::distance(Metric::L2, Some(1)); + let mut workspace = PruneWorkspace::default(); + + // When + let error = workspace + .prune(data, 0, candidates, &graph_config(1), &distance) + .unwrap_err(); + + // Then + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::TooManyCandidates { actual, max }) + if *actual == count && *max == u16::MAX as usize + )); } } diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index eb6f7690bc..8683db3daa 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -39,8 +39,6 @@ pub(crate) enum LeafBuildError { rows: usize, columns: usize, }, - #[error("failed to form {buffer} view for leaf {leaf}")] - InvalidView { leaf: usize, buffer: &'static str }, #[error("failed to convert point {point} in leaf {leaf}")] Conversion { leaf: usize, @@ -100,51 +98,15 @@ impl LeafBuffers { } fn prepare_local_adjacency(&mut self, point_count: usize) { - self.local_adjacency.resize_with(point_count, Vec::new); + if self.local_adjacency.len() < point_count { + self.local_adjacency.resize_with(point_count, Vec::new); + } self.local_adjacency[..point_count] .iter_mut() .for_each(Vec::clear); } } -/// Concurrent candidate lists indexed by global point ID. -/// -/// A point can occur in several overlapping leaves. A worker locks one point's -/// list and adds all IDs from one leaf. `AdjacencyList` removes duplicates during -/// this append. The function sorts each list after all leaf jobs finish. -struct DirectCandidates { - lists: Vec>>, -} - -impl DirectCandidates { - fn new(point_count: usize) -> Self { - let lists = (0..point_count) - .map(|_| Mutex::new(AdjacencyList::new())) - .collect(); - Self { lists } - } - - fn add_leaf(&self, point_ids: &[u32], local_adjacency: &[Vec]) { - for (&source, additions) in point_ids.iter().zip(local_adjacency) { - // `add_direct_leaf_candidates` checks every point ID before this append. - self.lists[source as usize] - .lock() - .extend_from_slice(additions); - } - } - - fn into_lists(self) -> Vec> { - self.lists - .into_iter() - .map(Mutex::into_inner) - .map(|mut candidates| { - candidates.sort(); - candidates - }) - .collect() - } -} - /// Build direct graph candidates from all overlapping leaves. /// /// Each selected leaf pair contributes both edge directions. Candidate lists use @@ -159,9 +121,11 @@ pub(super) fn build_leaf_candidates( where A: PiPNNSIMDSchema, M: LeafMetric, - T: VectorRepr + 'static, + T: VectorRepr, { - let candidates = DirectCandidates::new(data.nrows()); + let candidates: Vec<_> = (0..data.nrows()) + .map(|_| Mutex::new(AdjacencyList::new())) + .collect(); leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { @@ -176,7 +140,14 @@ where ) }, )?; - Ok(candidates.into_lists()) + Ok(candidates + .into_iter() + .map(Mutex::into_inner) + .map(|mut neighbors| { + neighbors.sort(); + neighbors + }) + .collect()) } /// Add one leaf's symmetric neighbors to the direct candidate lists. @@ -184,6 +155,10 @@ where /// Reusable buffers can be longer than this leaf, so all accesses use the current /// leaf shape. #[allow(clippy::too_many_arguments)] +#[expect( + clippy::expect_used, + reason = "buffer prefixes have the checked leaf shape" +)] fn add_direct_leaf_candidates( arch: A, data: MatrixView<'_, T>, @@ -191,12 +166,12 @@ fn add_direct_leaf_candidates( point_ids: &[u32], requested_k: usize, buffers: &mut LeafBuffers, - candidates: &DirectCandidates, + candidates: &[Mutex>], ) -> Result<(), LeafBuildError> where A: PiPNNSIMDSchema, M: LeafMetric, - T: VectorRepr + 'static, + T: VectorRepr, { let (leaf_k, neighbor_value_count) = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; @@ -212,7 +187,7 @@ where .zip(point_values.chunks_exact_mut(data.ncols())) { let source_values = data.row(point as usize); - T::as_f32_into(source_values, point_output).map_err(|source| { + super::conversion::as_f32_into(source_values, point_output).map_err(|source| { LeafBuildError::Conversion { leaf, point, @@ -221,22 +196,14 @@ where })?; } - let points = - MatrixView::try_from(&*point_values, point_ids.len(), data.ncols()).map_err(|_| { - LeafBuildError::InvalidView { - leaf, - buffer: "leaf point matrix", - } - })?; + let points = MatrixView::try_from(&*point_values, point_ids.len(), data.ncols()) + .expect("point buffer prefix has the checked leaf shape"); let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), leaf_k, ) - .map_err(|_| LeafBuildError::InvalidView { - leaf, - buffer: "leaf output", - })?; + .expect("neighbor buffer prefix has the checked leaf shape"); select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; @@ -247,7 +214,11 @@ where &buffers.neighbors[..neighbor_value_count], &mut buffers.local_adjacency[..point_ids.len()], ); - candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]); + for (&point_id, additions) in point_ids.iter().zip(&buffers.local_adjacency) { + candidates[point_id as usize] + .lock() + .extend_from_slice(additions); + } Ok(()) } @@ -265,10 +236,9 @@ fn add_symmetric_neighbors( let target = neighbor.local_idx as usize; let source_id = point_ids[source]; let target_id = point_ids[target]; - if source_id != target_id { - local_adjacency[source].push(target_id); - local_adjacency[target].push(source_id); - } + // Leaves contain unique IDs, and the kernel excludes each point itself. + local_adjacency[source].push(target_id); + local_adjacency[target].push(source_id); } } } @@ -280,393 +250,170 @@ fn grow(values: &mut Vec, len: usize, value: T) { } #[cfg(test)] -mod tests { - use diskann_utils::views::MatrixView; - use diskann_vector::distance::Metric; - use diskann_wide::arch::{self, Target1}; +mod build_leaf_candidates_tests { + use diskann_wide::arch::Scalar; use half::f16; use rstest::rstest; - use std::collections::BTreeSet; - use super::super::simd::PiPNNSIMDSchema; - use super::{ - DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, - build_leaf_candidates, - }; + use super::*; + use crate::graph::pipnn::{InnerProduct, L2}; - fn matrix_view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { - MatrixView::try_from(data, rows, columns).unwrap() - } - - fn leaf_build_pool() -> rayon::ThreadPool { + fn build_candidates( + data: MatrixView<'_, T>, + leaves: Vec>, + k: usize, + threads: usize, + ) -> Vec> { rayon::ThreadPoolBuilder::new() - .num_threads(4) + .num_threads(threads) .build() .unwrap() + .install(|| build_leaf_candidates::<_, M, _>(Scalar, data, leaves, k)) + .unwrap() + .into_iter() + .map(Vec::from) + .collect() } - struct LeafBuildCall<'a, T> { - data: MatrixView<'a, T>, - leaves: Vec>, - k: usize, - } + #[rstest] + #[case::f32([5.5_f32, 100.0, 1.25, 50.0, 0.25])] + #[case::f16([5.5, 100.0, 1.25, 50.0, 0.25].map(f16::from_f32))] + #[case::u8([5_u8, 100, 1, 50, 0])] + #[case::i8([-6_i8, 100, -10, 50, -11])] + fn selected_local_positions_become_global_ids(#[case] values: [T; 5]) { + // Given: only IDs 0, 2, 4 belong to the leaf. Their coordinates decrease, + // so 0 selects 2, while 2 and 4 select each other. Reverse edges add 2 -> 0. + let data = MatrixView::column_vector(&values[..]); + let leaves = vec![vec![0, 2, 4]]; + let expected = [vec![2], vec![], vec![0, 4], vec![], vec![2]]; - struct DispatchLeafBuild(Metric); - - impl - Target1< - A, - Result>, LeafBuildError>, - LeafBuildCall<'_, T>, - > for DispatchLeafBuild - where - A: PiPNNSIMDSchema, - T: crate::utils::VectorRepr + 'static, - { - fn run( - self, - arch: A, - call: LeafBuildCall<'_, T>, - ) -> Result>, LeafBuildError> { - use super::super::{Cosine, CosineNormalized, InnerProduct, L2}; - - match self.0 { - Metric::L2 => { - build_leaf_candidates::(arch, call.data, call.leaves, call.k) - } - Metric::Cosine => { - build_leaf_candidates::(arch, call.data, call.leaves, call.k) - } - Metric::CosineNormalized => build_leaf_candidates::( - arch, - call.data, - call.leaves, - call.k, - ), - Metric::InnerProduct => build_leaf_candidates::( - arch, - call.data, - call.leaves, - call.k, - ), - } - } - } + // When + let actual = build_candidates::<_, L2>(data, leaves, 1, 1); - fn build_candidate_graph( - data: MatrixView<'_, T>, - leaves: &[Vec], - k: usize, - metric: Metric, - ) -> Result>, LeafBuildError> - where - T: crate::utils::VectorRepr + 'static, - { - leaf_build_pool().install(|| { - arch::dispatch1_no_features( - DispatchLeafBuild(metric), - LeafBuildCall { - data, - leaves: leaves.to_vec(), - k, - }, - ) - }) + // Then + assert_eq!(actual, expected); } - fn adjacency_lists(graph: Vec>) -> Vec> { - graph.into_iter().map(Vec::from).collect() - } + #[test] + fn unassigned_kernel_results_add_no_edges() { + // Given: the NaN point has no rankable distances. The two finite points + // select each other, leaving their second output slot unassigned. + let values = [0.0_f32, 1.0, f32::NAN]; + let data = MatrixView::column_vector(&values[..]); + let expected = [vec![1], vec![0], vec![]]; - fn circular_overlapping_leaves( - point_count: usize, - leaf_count: usize, - leaf_size: usize, - ) -> Vec> { - (0..leaf_count) - .map(|offset| { - (0..leaf_size) - .map(|point| ((point + offset) % point_count) as u32) - .collect() - }) - .collect() - } + // When + let actual = build_candidates::<_, InnerProduct>(data, vec![vec![0, 1, 2]], 2, 1); - fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { - let mut graph = vec![BTreeSet::new(); data.len()]; - for (source, left) in data.iter().enumerate() { - let mut nearest: Vec<_> = data - .iter() - .enumerate() - .filter(|(target, _)| *target != source) - .map(|(target, right)| { - let distance = left - .iter() - .zip(right) - .map(|(x, y)| (x - y) * (x - y)) - .sum::(); - (target, distance) - }) - .collect(); - nearest.sort_by(|left, right| { - left.1 - .total_cmp(&right.1) - .then_with(|| left.0.cmp(&right.0)) - }); - for &(target, _) in nearest.iter().take(k) { - graph[source].insert(target as u32); - graph[target].insert(source as u32); - } - } - graph - .into_iter() - .map(|neighbors| neighbors.into_iter().collect()) - .collect() + // Then + assert_eq!(actual, expected); } - #[test] - fn leaf_adjacency_matches_an_independent_all_pairs_reference() { - // Given - let points = [ - [0.0_f32, 0.0], - [1.0, 0.2], - [3.1, 0.5], - [7.8, 1.4], - [-2.3, 4.1], - [6.7, -3.2], - ]; - let flat: Vec<_> = points.into_iter().flatten().collect(); - let expected_adjacency = brute_force_symmetric_l2(&points, 2); + #[rstest] + #[case::one_worker(1)] + #[case::four_workers(4)] + fn overlapping_leaves_merge_into_sorted_unique_neighbors(#[case] threads: usize) { + // Given: k = 2 connects every pair in each three-point leaf. The union + // contains all pairs except 1 <-> 3, regardless of repeated leaf jobs. + let values = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::column_vector(&values[..]); + let leaves = [vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]] + .into_iter() + .cycle() + .take(48) + .collect(); + let expected = [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]]; // When - let actual_adjacency = adjacency_lists( - build_candidate_graph( - matrix_view(&flat, points.len(), 2), - &[(0..points.len() as u32).collect()], - 2, - Metric::L2, - ) - .unwrap(), - ); + let actual = build_candidates::<_, L2>(data, leaves, 2, threads); // Then - assert_eq!(actual_adjacency, expected_adjacency); + assert_eq!(actual, expected); } #[test] - fn non_rankable_neighbors_are_omitted() { - // Given - let data = [0.0_f32, 1.0, f32::NAN]; - let expected_adjacency = [vec![1], vec![0], vec![]]; + fn reverse_edges_can_give_a_point_more_than_twice_k_neighbors() { + // Given: each outer point is distance 1 from the origin and at least + // sqrt(2) from another outer point. Their three reverse edges exceed 2k. + let points = [[0.0_f32, 0.0], [1.0, 0.0], [-1.0, 0.0], [0.0, 1.0]]; + let data = MatrixView::try_from(points.as_flattened(), 4, 2).unwrap(); + let expected = [vec![1, 2, 3], vec![0], vec![0], vec![0]]; // When - let graph = build_candidate_graph( - matrix_view(&data, 3, 1), - &[vec![0, 1, 2]], - 2, - Metric::InnerProduct, - ) - .unwrap(); - let actual_adjacency = adjacency_lists(graph); + let actual = build_candidates::<_, L2>(data, vec![vec![0, 1, 2, 3]], 1, 1); // Then - assert_eq!(actual_adjacency, expected_adjacency); + assert_eq!(actual, expected); } #[test] - fn overlapping_leaves_contribute_each_candidate_once() { + fn singleton_leaves_add_no_edges() { // Given - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; - let expected_adjacency = [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]]; + let values = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::column_vector(&values[..]); + let leaves = vec![vec![0], vec![1], vec![2]]; + let expected: [Vec; 3] = [vec![], vec![], vec![]]; // When - let graph = - build_candidate_graph(matrix_view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); - let actual_adjacency = adjacency_lists(graph); + let actual = build_candidates::<_, L2>(data, leaves, 1, 1); // Then - assert_eq!(actual_adjacency, expected_adjacency); + assert_eq!(actual, expected); } #[test] - fn symmetric_edges_can_give_a_center_more_than_two_k_neighbors() { - let dimensions = 9; - let mut data = vec![0.0_f32; 10 * dimensions]; - for source in 1..10 { - data[source * dimensions + source - 1] = 1.0; - } - - let graph = build_candidate_graph( - matrix_view(&data, 10, dimensions), - &[(0..10).collect()], - 1, - Metric::L2, - ) - .unwrap(); - - assert_eq!(&*graph[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); - assert!(graph.iter().enumerate().all(|(source, neighbors)| { - neighbors.iter().all(|&target| target as usize != source) - && neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32)) - })); - } - - fn build_l2_candidate_graph(data: &[T], points: usize, dimensions: usize) -> Vec> - where - T: crate::utils::VectorRepr + 'static, - { - let leaves = vec![(0..points as u32).collect()]; - adjacency_lists( - build_candidate_graph( - matrix_view(data, points, dimensions), - &leaves, + fn reused_leaf_buffers_do_not_add_edges_from_previous_leaves() { + // Given: use the same buffers for four points, two points, then three. + // The final leaf excludes ID 1 and connects all pairs among 0, 2, 3. + let values = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::column_vector(&values[..]); + let mut buffers = LeafBuffers::default(); + let mut build_leaf = |point_ids: &[u32]| { + let candidates: Vec<_> = (0..data.nrows()) + .map(|_| Mutex::new(AdjacencyList::new())) + .collect(); + add_direct_leaf_candidates::<_, L2, _>( + Scalar, + data, + 0, + point_ids, 2, - Metric::L2, + &mut buffers, + &candidates, ) - .unwrap(), - ) - } - - fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) - where - T: crate::utils::VectorRepr + 'static, - { - let points = 8; - // Source dimension controls VectorRepr conversion chunking. Cover tails on - // both sides of 4-, 8-, and 16-element boundaries, then a second 16-lane - // chunk. Input integers remain exact in every tested representation. - for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let raw: Vec = (0..points * dimensions) - .map(|index| { - let source = index / dimensions; - let dimension = index % dimensions; - (source + dimension) as u8 + .unwrap(); + candidates + .into_iter() + .map(|list| { + let mut ids = Vec::from(list.into_inner()); + ids.sort_unstable(); + ids }) - .collect(); - let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); - let converted: Vec = raw.iter().copied().map(&convert).collect(); - assert_eq!( - build_l2_candidate_graph(&converted, points, dimensions), - build_l2_candidate_graph(&f32_data, points, dimensions), - "{label} dimensions={dimensions}" - ); - } - } - - #[test] - fn f16_conversion_matches_f32_across_dimension_boundaries() { - assert_source_conversion_matches_f32("f16", |value| f16::from_f32(value as f32)); - } - - #[test] - fn u8_conversion_matches_f32_across_dimension_boundaries() { - assert_source_conversion_matches_f32("u8", |value| value); - } - - #[test] - fn i8_conversion_matches_f32_across_dimension_boundaries() { - // Applying the same translation to every coordinate preserves L2 pair - // ordering while exercising signed conversion. - assert_source_conversion_matches_f32("i8", |value| value as i8 - 11); - } - - #[rstest] - fn candidate_graph_is_symmetric_unique_and_non_self( - #[values( - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct - )] - metric: Metric, - ) { - // Given - let diagonal = std::f32::consts::FRAC_1_SQRT_2; - let unit_vectors = [1.0_f32, 0.0, diagonal, diagonal, 0.0, 1.0, -1.0, 0.0]; - let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; + .collect::>() + }; + build_leaf(&[0, 1, 2, 3]); + build_leaf(&[1, 3]); + let expected = [vec![2, 3], vec![], vec![0, 3], vec![0, 2]]; // When - let graph = - build_candidate_graph(matrix_view(&unit_vectors, 4, 2), &leaves, 2, metric).unwrap(); + let actual = build_leaf(&[0, 2, 3]); // Then - for (source, neighbors) in graph.iter().enumerate() { - assert!(neighbors.iter().all(|&target| target as usize != source)); - assert!( - neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32)) - ); - assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); - } - } - - #[test] - fn parallel_leaf_schedule_does_not_change_candidate_order() { - let data: Vec = (0..64).map(|value| value as f32).collect(); - let leaves = circular_overlapping_leaves(64, 32, 16); - let expected_candidate_order = - build_candidate_graph(matrix_view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); - for _ in 0..8 { - let actual_candidate_order = - build_candidate_graph(matrix_view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); - assert_eq!(actual_candidate_order, expected_candidate_order); - } + assert_eq!(actual, expected); } #[test] - fn singleton_leaves_add_no_candidates() { + fn overflowing_leaf_shape_is_rejected() { // Given - let point_values = [0.0_f32, 1.0, 2.0]; - let singleton_leaves = [vec![0], vec![1], vec![2]]; - let expected_adjacency: [Vec; 3] = [vec![], vec![], vec![]]; + let mut buffers = LeafBuffers::default(); // When - let actual_adjacency = adjacency_lists( - build_candidate_graph( - matrix_view(&point_values, 3, 1), - &singleton_leaves, - 1, - Metric::L2, - ) - .unwrap(), - ); + let result = buffers.prepare(7, usize::MAX, 2, 1); // Then - assert_eq!(actual_adjacency, expected_adjacency); - } - - #[test] - fn leaf_buffer_preparation_reports_shape_overflow_before_allocating() { - let mut buffers = LeafBuffers::default(); assert!(matches!( - buffers.prepare(7, usize::MAX, 2, 1), + result, Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) )); } - - #[test] - fn symmetric_edge_mapping_skips_duplicate_ids_instead_of_adding_self_edges() { - let mut graph = vec![Vec::new(); 2]; - add_symmetric_neighbors( - &[7, 7], - 1, - &[ - super::super::topk::Candidate::new(1, 0.0), - super::super::topk::Candidate::new(0, 0.0), - ], - &mut graph, - ); - assert!(graph.iter().all(|neighbors| neighbors.is_empty())); - } - - #[test] - fn direct_candidate_accumulator_keeps_unique_sorted_lists() { - let candidates = DirectCandidates::new(2); - candidates.add_leaf(&[0, 1], &[vec![1, 1], vec![0]]); - assert_eq!(adjacency_lists(candidates.into_lists()), [vec![1], vec![0]]); - } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 653c088754..a22aae559c 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -9,8 +9,8 @@ //! point that acts as the center of one child partition. A leaf is a bounded //! child partition used for local neighbor selection. //! -//! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has -//! at most `c_max` points. +//! 1. `partitioning` samples leaders and makes overlapping leaves of at most +//! `c_max` points. Points without rankable leaders contribute no assignments. //! 2. `leaf_build` computes one ranking-distance buffer for each leaf. It //! selects local neighbors and merges their global point IDs. //! 3. `finalization` applies Vamana RobustPrune to each candidate list that is @@ -18,7 +18,7 @@ //! //! `diskann-wide` selects architecture `A` for the ranking kernels. One match //! selects metric marker `M`. Metric computation stays architecture-neutral. -//! The numerical loops do not dispatch again. +//! Ranking kernels reapply architecture features inside Rayon jobs. //! //! [`PiPNNConfig`] contains partition and local-neighbor parameters. //! [`PiPNNBuildContext`] borrows graph policy and a Rayon pool. [`build_graph`] @@ -31,6 +31,7 @@ //! Partition and leaf work use separate reusable buffers. The build consumes //! each output before it creates another graph representation. +mod conversion; mod finalization; mod leaf_build; mod leaf_kernel; @@ -48,7 +49,7 @@ use crate::{ }; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; -use diskann_wide::arch::{self, Target1}; +use diskann_wide::arch::{self, Target2}; use rayon::ThreadPool; use self::{leaf_metric::LeafMetric, partition_metric::PartitionMetric, simd::PiPNNSIMDSchema}; @@ -180,10 +181,10 @@ impl PiPNNConfig { /// PiPNN policy and borrowed execution resources for one graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { - pub(crate) config: PiPNNConfig, - pub(crate) graph: &'a Config, - pub(crate) metric: Metric, - pub(crate) pool: &'a ThreadPool, + config: PiPNNConfig, + graph: &'a Config, + metric: Metric, + pool: &'a ThreadPool, } impl<'a> PiPNNBuildContext<'a> { @@ -223,7 +224,7 @@ pub fn build_graph( context: &PiPNNBuildContext<'_>, ) -> ANNResult>> where - T: VectorRepr + Send + Sync + 'static, + T: VectorRepr, { context .pool @@ -236,7 +237,7 @@ fn validate_and_dispatch_build( context: &PiPNNBuildContext<'_>, ) -> ANNResult>> where - T: VectorRepr + Send + Sync + 'static, + T: VectorRepr, { if data.nrows() == 0 { return Err(ANNError::message("PiPNN requires at least one data point")); @@ -252,55 +253,34 @@ where data.nrows() ))); } - // Conversion does not make integer vectors unit length. Use the norm-aware - // cosine formula for these vectors. - let metric = effective_metric::(context.metric); - arch::dispatch1_no_features( - RunBuildGraph, - BuildGraphCall { - data, - context, - metric, - }, - ) -} - -struct BuildGraphCall<'data, 'context, 'policy, T> { - data: MatrixView<'data, T>, - context: &'context PiPNNBuildContext<'policy>, - metric: Metric, + arch::dispatch2_no_features(BuildGraph, data, context) } -struct RunBuildGraph; +struct BuildGraph; -impl Target1>>, BuildGraphCall<'_, '_, '_, T>> - for RunBuildGraph +impl Target2>>, MatrixView<'_, T>, &PiPNNBuildContext<'_>> + for BuildGraph where A: PiPNNSIMDSchema, - T: VectorRepr + Send + Sync + 'static, + T: VectorRepr, { fn run( self, arch: A, - call: BuildGraphCall<'_, '_, '_, T>, + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, ) -> ANNResult>> { - match call.metric { - Metric::L2 => build_graph_for::(arch, call.data, call.context, Metric::L2), - Metric::Cosine => { - build_graph_for::(arch, call.data, call.context, Metric::Cosine) + // Converting raw integer coordinates does not normalize their vectors. + let metric = effective_metric::(context.metric); + match metric { + Metric::L2 => build_graph_for::(arch, data, context, metric), + Metric::Cosine => build_graph_for::(arch, data, context, metric), + Metric::CosineNormalized => { + build_graph_for::(arch, data, context, metric) + } + Metric::InnerProduct => { + build_graph_for::(arch, data, context, metric) } - Metric::CosineNormalized => build_graph_for::( - arch, - call.data, - call.context, - Metric::CosineNormalized, - ), - Metric::InnerProduct => build_graph_for::( - arch, - call.data, - call.context, - Metric::InnerProduct, - ), } } } @@ -318,7 +298,7 @@ fn build_graph_for( where A: PiPNNSIMDSchema, M: LeafMetric + PartitionMetric, - T: VectorRepr + Send + Sync + 'static, + T: VectorRepr, { let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition::(arch, data, &context.config))?; @@ -334,7 +314,7 @@ where .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) } -fn effective_metric(metric: Metric) -> Metric { +fn effective_metric(metric: Metric) -> Metric { use std::any::TypeId; if metric == Metric::CosineNormalized @@ -352,28 +332,40 @@ fn config_error(message: impl std::fmt::Display) -> ANNError { } #[cfg(test)] -mod tests { +mod effective_metric_tests { use super::*; use half::f16; + use rstest::rstest; - #[test] - fn integer_normalized_cosine_uses_unnormalized_cosine() { - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let expected = if metric == Metric::CosineNormalized { - Metric::Cosine - } else { - metric - }; - assert_eq!(effective_metric::(metric), expected); - assert_eq!(effective_metric::(metric), expected); - assert_eq!(effective_metric::(metric), metric); - assert_eq!(effective_metric::(metric), metric); - } + #[rstest] + #[case::l2(Metric::L2)] + #[case::cosine(Metric::Cosine)] + #[case::normalized_cosine(Metric::CosineNormalized)] + #[case::inner_product(Metric::InnerProduct)] + fn floating_point_formats_keep_the_requested_metric(#[case] metric: Metric) { + // Given / When: floating-point callers define whether their data is normalized. + let actual = ( + effective_metric::(metric), + effective_metric::(metric), + ); + + // Then: both supported floating-point formats preserve the request. + assert_eq!(actual, (metric, metric)); + } + + #[rstest] + #[case::l2(Metric::L2)] + #[case::cosine(Metric::Cosine)] + #[case::inner_product(Metric::InnerProduct)] + fn integer_formats_keep_metrics_that_do_not_assume_normalization(#[case] metric: Metric) { + // Given / When + let actual = ( + effective_metric::(metric), + effective_metric::(metric), + ); + + // Then: only normalized cosine needs the integer-data fallback. + assert_eq!(actual, (metric, metric)); } } #[cfg(test)] @@ -387,8 +379,7 @@ mod build_graph_tests { use crate::graph::config::{self, MaxDegree}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; - use half::f16; - use rand::{Rng, SeedableRng, rngs::StdRng}; + use rstest::rstest; fn pipnn_config() -> PiPNNConfig { PiPNNConfig { @@ -420,50 +411,28 @@ mod build_graph_tests { graph.into_iter().map(Vec::from).collect() } - fn assert_graph_invariants( - graph: &[crate::graph::AdjacencyList], - points: usize, - degree: usize, - ) { - assert_eq!(graph.len(), points); - for (source, row) in graph.iter().enumerate() { - assert!(row.len() <= degree); - let mut sorted = row.to_vec(); - sorted.sort_unstable(); - sorted.dedup(); - assert_eq!(sorted.len(), row.len()); - assert!( - row.iter() - .all(|&id| (id as usize) < points && id as usize != source) - ); - } - } - #[test] - fn builds_a_single_leaf_graph_for_real_dataset_ids() { - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + fn single_leaf_merges_neighbors_in_both_directions() { + // Given: unequal gaps give nearest choices 0->1, 1->0, 2->1, 3->2. + // Symmetric insertion adds the reverse edge for each choice. + let values = [0.0_f32, 1.0, 3.0, 7.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let expected = [vec![1], vec![0, 2], vec![1, 3], vec![2]]; let graph = graph_config(Metric::L2, 2); let pool = pool(2); let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + // When let actual = build_graph(data, &context).unwrap(); - assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); - - let graph = graph_config(Metric::L2, 1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - let pruned = build_graph(data, &context).unwrap(); - - assert_graph_invariants(&pruned, 4, 1); - for (source, neighbors) in pruned.iter().enumerate() { - assert_eq!(source.abs_diff(neighbors[0] as usize), 1); - } + // Then + assert_eq!(rows(actual), expected); } #[test] fn omits_non_rankable_candidates_without_invalid_ids() { + // Given: three points force partitioning at c_max=2. The NaN point + // has no rankable leader or neighbor, so only the finite points connect. let values = [0.0_f32, 1.0, f32::NAN]; let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); let graph = graph_config(Metric::InnerProduct, 2); @@ -477,111 +446,161 @@ mod build_graph_tests { replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::InnerProduct, &pool).unwrap(); + let expected = [vec![1], vec![0], vec![]]; + // When let actual = build_graph(data, &context).unwrap(); - assert_graph_invariants(&actual, 3, 2); - assert_eq!(rows(actual), [vec![1], vec![0], vec![]]); + // Then + assert_eq!(rows(actual), expected); } #[test] - fn prunes_overfull_single_leaf_candidates_to_the_graph_degree() { - let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; - let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); + fn finalization_keeps_the_nearest_neighbor_when_degree_is_one() { + // Given: selecting both peers makes each row overfull. With degree one, + // RobustPrune keeps the nearest: 0->1, 1->0, 2->1. + let values = [0.0_f32, 1.0, 3.0]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let expected = [vec![1], vec![0], vec![1]]; let graph = graph_config(Metric::L2, 1); let pool = pool(2); let config = PiPNNConfig { - c_max: 5, + c_max: 3, c_min: 1, p_samp: 0.5, fanout: vec![2], - leaf_k: 4, + leaf_k: 2, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + // When let actual = build_graph(data, &context).unwrap(); - assert_graph_invariants(&actual, 5, 1); - assert!(actual.iter().all(|row| row.len() == 1)); + // Then + assert_eq!(rows(actual), expected); } - #[test] - fn rejects_empty_dataset_dimensions_at_the_public_boundary() { + #[rstest] + #[case::no_points(0, 4, "PiPNN requires at least one data point")] + #[case::no_dimensions(4, 0, "PiPNN requires at least one data dimension")] + fn empty_input_shape_returns_an_error( + #[case] points: usize, + #[case] dimensions: usize, + #[case] message: &str, + ) { + // Given + let data = MatrixView::try_from(&[] as &[f32], points, dimensions).unwrap(); let graph = graph_config(Metric::L2, 2); let pool = pool(1); let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); - let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); + // When + let error = build_graph(data, &context).unwrap_err(); - assert!(build_graph(no_rows, &context).is_err()); - assert!(build_graph(no_columns, &context).is_err()); + // Then + assert!(error.to_string().contains(message), "{error}"); } - #[test] - fn supports_every_source_type_and_metric() { - fn build( - values: &[T], - metric: Metric, - ) { - let data = MatrixView::try_from(values, 6, 2).unwrap(); - let graph = graph_config(metric, 2); - let pool = pool(2); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); - let actual = build_graph(data, &context).unwrap(); - assert_graph_invariants(&actual, 6, 2); - } + #[rstest] + #[case::l2(Metric::L2, [vec![1, 2], vec![0], vec![0]])] + #[case::inner_product(Metric::InnerProduct, [vec![1], vec![0, 2], vec![1]])] + fn selected_metric_changes_the_neighbor_graph( + #[case] metric: Metric, + #[case] expected: [Vec; 3], + ) { + // Given: squared L2 scores for pairs (0,1),(0,2),(1,2) are 4,0.8,6.4; + // negative-dot scores are -3,-0.6,-1.8. The nearest choices differ. + let values = [1.0_f32, 0.0, 3.0, 0.0, 0.6, 0.8]; + let data = MatrixView::try_from(&values[..], 3, 2).unwrap(); + let graph = graph_config(metric, 2); + let pool = pool(1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); - let values = [ - 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, - ]; - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - build(&values, metric); - } - build(&values.map(f16::from_f32), Metric::L2); - build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); - build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); + // When + let actual = build_graph(data, &context).unwrap(); + + // Then + assert_eq!(rows(actual), expected); + } + + #[rstest] + #[case::cosine(Metric::Cosine)] + #[case::normalized_cosine(Metric::CosineNormalized)] + fn cosine_metrics_rank_unit_vectors_by_direction(#[case] metric: Metric) { + // Given: pairwise dots are 0.6,-1,-0.6, so 0 chooses 1, 1 chooses 0, + // and 2 chooses 1. Symmetrization gives the expected three-node chain. + let values = [1.0_f32, 0.0, 0.6, 0.8, -1.0, 0.0]; + let data = MatrixView::try_from(&values[..], 3, 2).unwrap(); + let expected = [vec![1], vec![0, 2], vec![1]]; + let graph = graph_config(metric, 2); + let pool = pool(1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); + + // When + let actual = build_graph(data, &context).unwrap(); + + // Then + assert_eq!(rows(actual), expected); } #[test] - fn integer_normalized_cosine_matches_cosine() { - fn assert_match(values: &[T]) { - let data = MatrixView::try_from(values, 8, 2).unwrap(); - let pool = pool(2); - let build = |metric| { - let graph = graph_config(metric, 2); - let config = PiPNNConfig { - c_max: 8, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - leaf_k: 1, - replicas: 1, - }; - let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); - rows(build_graph(data, &context).unwrap()) + fn singleton_has_no_self_neighbor() { + // Given: a singleton has zero non-self candidates even though leaf_k=1. + let values = [3.0_f32, 4.0]; + let data = MatrixView::try_from(&values[..], 1, 2).unwrap(); + let graph = graph_config(Metric::L2, 2); + let pool = pool(1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + // When + let actual = build_graph(data, &context).unwrap(); + + // Then + assert_eq!(rows(actual), [Vec::::new()]); + } + + #[rstest] + #[case::unsigned([1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2])] + #[case::signed([1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2])] + fn integer_normalized_cosine_matches_cosine( + #[case] values: [T; 16], + ) { + // Given: unequal norms distinguish cosine from raw dot ranking. + // Keep one leaf so partition randomness cannot hide metric selection. + let data = MatrixView::try_from(&values[..], 8, 2).unwrap(); + let pool = pool(2); + let build = |metric| { + let graph = graph_config(metric, 2); + let config = PiPNNConfig { + c_max: 8, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + leaf_k: 1, + replicas: 1, }; - assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); - } + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + rows(build_graph(data, &context).unwrap()) + }; + let expected = build(Metric::Cosine); - assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); - assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); + // When + let actual = build(Metric::CosineNormalized); + + // Then + assert_eq!(actual, expected); } - #[test] - fn is_deterministic_for_a_fixed_pool_size() { - let data: Vec = (0..96 * 4) - .map(|value| ((value * 17 + 3) % 101) as f32) - .collect(); - let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); + #[rstest] + #[case::one_worker(1)] + #[case::four_workers(4)] + fn recursive_build_is_deterministic_for_a_fixed_pool(#[case] threads: usize) { + // Given: distinct points require multiple partition levels and two replicas. + let values: Vec = (0..96).flat_map(|i| [i as f32, (i * i) as f32]).collect(); + let data = MatrixView::try_from(values.as_slice(), 96, 2).unwrap(); let graph = graph_config(Metric::L2, 8); - let pool = pool(4); + let pool = pool(threads); let config = PiPNNConfig { c_max: 16, c_min: 4, @@ -592,42 +611,13 @@ mod build_graph_tests { }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - let first = build_graph(data, &context).unwrap(); - let second = build_graph(data, &context).unwrap(); + let expected = build_graph(data, &context).unwrap(); - assert_eq!(first, second); - assert_graph_invariants(&first, 96, 8); - } - - #[test] - fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { - let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); - for case in 0..24 { - let points = rng.random_range(4..=32); - let dimensions = rng.random_range(1..=8); - let c_max = rng.random_range(4..=points.min(12)); - let c_min = rng.random_range(1..=c_max); - let degree = rng.random_range(1..=points.min(8)); - let values: Vec = (0..points * dimensions) - .map(|_| rng.random_range(-10.0..10.0)) - .collect(); - let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); - let graph = graph_config(Metric::L2, degree); - let pool = pool(2); - let config = PiPNNConfig { - c_max, - c_min, - p_samp: 0.5, - fanout: vec![2], - leaf_k: rng.random_range(1..=7), - replicas: rng.random_range(1..=2), - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + // When + let actual = build_graph(data, &context).unwrap(); - let actual = build_graph(data, &context) - .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); - assert_graph_invariants(&actual, points, degree); - } + // Then + assert_eq!(actual, expected); } } #[cfg(test)] @@ -640,6 +630,7 @@ mod config_tests { use super::{PiPNNBuildContext, PiPNNConfig}; use crate::graph::config::{self, MaxDegree}; use diskann_vector::distance::Metric; + use rstest::rstest; fn pipnn_config() -> PiPNNConfig { PiPNNConfig { @@ -667,79 +658,69 @@ mod config_tests { .unwrap() } + #[rstest] + #[case::zero_maximum(PiPNNConfig { c_max: 0, ..pipnn_config() }, "c_max must be greater than zero")] + #[case::zero_minimum(PiPNNConfig { c_min: 0, ..pipnn_config() }, "c_min must be greater than zero")] + #[case::minimum_exceeds_maximum(PiPNNConfig { c_min: 513, ..pipnn_config() }, "c_min (513) must not exceed c_max (512)")] + #[case::zero_sampling(PiPNNConfig { p_samp: 0.0, ..pipnn_config() }, "p_samp (0) must be in (0, 1]")] + #[case::negative_sampling(PiPNNConfig { p_samp: -0.5, ..pipnn_config() }, "p_samp (-0.5) must be in (0, 1]")] + #[case::sampling_above_one(PiPNNConfig { p_samp: 1.5, ..pipnn_config() }, "p_samp (1.5) must be in (0, 1]")] + #[case::nan_sampling(PiPNNConfig { p_samp: f64::NAN, ..pipnn_config() }, "p_samp (NaN) must be in (0, 1]")] + #[case::infinite_sampling(PiPNNConfig { p_samp: f64::INFINITY, ..pipnn_config() }, "p_samp (inf) must be in (0, 1]")] + #[case::empty_fanout(PiPNNConfig { fanout: vec![], ..pipnn_config() }, "fanout must not be empty")] + #[case::zero_fanout(PiPNNConfig { fanout: vec![1, 0], ..pipnn_config() }, "fanout values must be greater than zero")] + #[case::zero_neighbors(PiPNNConfig { leaf_k: 0, ..pipnn_config() }, "leaf_k must be greater than zero")] + #[case::zero_replicas(PiPNNConfig { replicas: 0, ..pipnn_config() }, "replicas must be greater than zero")] + fn invalid_policy_reports_the_rejected_parameter( + #[case] config: PiPNNConfig, + #[case] message: &str, + ) { + // Given: one invalid field in an otherwise valid policy. + let expected = format!("PiPNN configuration: {message}"); + + // When + let error = config.validate().unwrap_err(); + + // Then + assert!(error.to_string().contains(&expected), "{error}"); + } + #[test] - fn rejects_each_invalid_algorithm_parameter() { + fn context_validates_the_partition_policy() { + // Given: verify the public constructor invokes the policy validation. + let config = PiPNNConfig { + replicas: 0, + ..pipnn_config() + }; let graph = graph_config(Metric::L2, 1.2); let pool = pool(); - let mut cases = [ - PiPNNConfig { - c_max: 0, - ..pipnn_config() - }, - PiPNNConfig { - c_min: 0, - ..pipnn_config() - }, - PiPNNConfig { - c_min: 513, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: 0.0, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: -0.01, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: 1.01, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: f64::NAN, - ..pipnn_config() - }, - PiPNNConfig { - fanout: Vec::new(), - ..pipnn_config() - }, - PiPNNConfig { - fanout: vec![1, 0], - ..pipnn_config() - }, - PiPNNConfig { - leaf_k: 0, - ..pipnn_config() - }, - PiPNNConfig { - replicas: 0, - ..pipnn_config() - }, - ]; - - for config in &mut cases { - PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) - .expect_err("invalid PiPNN config must be rejected"); - } + + // When + let error = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap_err(); + + // Then + assert!( + error + .to_string() + .contains("PiPNN configuration: replicas must be greater than zero"), + "{error}", + ); } #[test] fn rejects_graph_policy_for_a_different_metric() { + // Given: inner-product occlusion cannot prune an L2 build. let graph = graph_config(Metric::InnerProduct, 1.2); let pool = pool(); + let expected = format!( + "PiPNN configuration: graph prune kind {:?} is incompatible with metric L2", + graph.prune_kind(), + ); + // When let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); - assert!(error.to_string().contains("prune kind")); - } - - #[test] - fn does_not_add_alpha_validation_beyond_graph_config() { - let pool = pool(); - for alpha in [0.9, f32::NAN, f32::INFINITY] { - let graph = graph_config(Metric::L2, alpha); - PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - } + // Then + assert!(error.to_string().contains(&expected), "{error}"); } } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 86a885baad..f187a20f3c 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -63,12 +63,17 @@ pub(crate) enum PartitionError { }, } -struct WorkItem { - indices: Vec, +struct PendingPartition { + point_ids: Vec, level: usize, seed: u64, } +struct PartitionSplit { + pending: Vec, + leaves: Vec>, +} + #[derive(Default)] struct StripeBuffers { point_values: Vec, @@ -96,8 +101,8 @@ type StripeBufferPool = ObjectPool; /// /// Each split samples partition centers and assigns every cluster point to its /// nearest centers. A cluster above `c_max` is split again. A level without a -/// configured fanout assigns each point to one center. Each replica covers every -/// input point. +/// configured fanout assigns each point to one center. Non-rankable distances +/// do not produce assignments. pub(super) fn partition( arch: A, data: MatrixView<'_, T>, @@ -135,43 +140,48 @@ where M: PartitionMetric, T: VectorRepr + Send + Sync, { - let initial_indices = point_ids(data.nrows()); + let initial_point_ids = (0..data.nrows() as u32).collect(); if data.nrows() <= config.c_max { - return Ok(vec![initial_indices]); + return Ok(vec![initial_point_ids]); } let mut leaves = Vec::new(); - let mut work = vec![WorkItem { - indices: initial_indices, + let mut pending = vec![PendingPartition { + point_ids: initial_point_ids, level: 0, seed, }]; for _ in 0..MAX_PARTITION_ITERATIONS { - if work.is_empty() { - return merge_undersized_leaves(leaves, config.c_min, config.c_max); + if pending.is_empty() { + return Ok(merge_undersized_leaves(leaves, config.c_min, config.c_max)); } - // Indexed parallel collection preserves work-item order. + // Indexed parallel collection preserves parent-partition order. #[allow(clippy::disallowed_methods)] - let results: ANNResult> = work + let splits: ANNResult> = pending .into_par_iter() - .map(|item| partition_work_item::(arch, data, config, item, stripe_buffers)) + .map(|partition| { + split_partition::(arch, data, config, partition, stripe_buffers) + }) .collect(); - let mut next_work = Vec::new(); - for (mut pending, mut finished) in results? { - next_work.append(&mut pending); - leaves.append(&mut finished); + let mut next_level = Vec::new(); + for mut split in splits? { + next_level.append(&mut split.pending); + leaves.append(&mut split.leaves); } - work = next_work; + pending = next_level; } - let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { - return merge_undersized_leaves(leaves, config.c_min, config.c_max); + let Some(largest) = pending + .iter() + .max_by_key(|partition| partition.point_ids.len()) + else { + return Ok(merge_undersized_leaves(leaves, config.c_min, config.c_max)); }; Err(ANNError::new(PartitionError::IterationLimit { - size: largest.indices.len(), + size: largest.point_ids.len(), level: largest.level, limit: MAX_PARTITION_ITERATIONS, })) @@ -181,46 +191,47 @@ where /// /// The function samples center points, assigns the cluster points, and returns /// bounded leaves separately from child clusters that need another split. -fn partition_work_item( +fn split_partition( arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, - item: WorkItem, + partition: PendingPartition, stripe_buffers: &StripeBufferPool, -) -> ANNResult<(Vec, Vec>)> +) -> ANNResult where A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { - let points = item.indices.len(); - let fanout = config.fanout.get(item.level).copied().unwrap_or(1); - let leaders = sample_leaders( - &item.indices, - config.p_samp, - mix_seed(item.seed, points as u64), - ); - let clusters = - assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; + let split_seed = mix_seed(partition.seed, partition.point_ids.len() as u64); + let fanout = config.fanout.get(partition.level).copied().unwrap_or(1); + let leaders = sample_leaders(&partition.point_ids, config.p_samp, split_seed); + let clusters = assign_to_leaders::( + arch, + data, + &partition.point_ids, + &leaders, + fanout, + stripe_buffers, + )?; let mut pending = Vec::new(); - let mut finished = Vec::new(); - let child_seed = mix_seed(item.seed, points as u64); + let mut leaves = Vec::new(); for cluster in clusters { if cluster.is_empty() { continue; } if cluster.len() <= config.c_max { - finished.push(cluster); + leaves.push(cluster); } else { - pending.push(WorkItem { - indices: cluster, - level: item.level + 1, - seed: child_seed, + pending.push(PendingPartition { + point_ids: cluster, + level: partition.level + 1, + seed: split_seed, }); } } - Ok((pending, finished)) + Ok(PartitionSplit { pending, leaves }) } /// Sample point IDs that act as centers for one partition split. @@ -232,13 +243,10 @@ fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> Vec /// Return the number of centers to sample from one cluster. /// -/// The count is `ceil(points * sampling_fraction)`, limited by `LEADER_CAP` and -/// the number of available points. A cluster with at least two points uses at -/// least two centers. +/// Splitting requires at least two points. The validated sampling fraction is in +/// `(0, 1]`; round up its point count and keep between two and `LEADER_CAP` centers. fn sampled_leader_count(points: usize, sampling_fraction: f64) -> usize { - ((points as f64 * sampling_fraction).ceil() as usize) - .clamp(2, LEADER_CAP) - .min(points) + ((points as f64 * sampling_fraction).ceil() as usize).clamp(2, LEADER_CAP) } fn replica_seed(replica: usize) -> u64 { @@ -285,9 +293,9 @@ where let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; let mut assignments = vec![0u32; assignment_len]; let stripe_points = assignment_stripe_point_count(leader_count); - let stripe_assignment_count = checked_area("assignment stripe", stripe_points, fanout)?; + let stripe_assignment_count = stripe_points * fanout; let stripe_count = point_ids.len().div_ceil(stripe_points); - let worker_stripe_count = stripe_count.div_ceil(rayon::current_num_threads().max(1)); + let worker_stripe_count = stripe_count.div_ceil(rayon::current_num_threads()); let worker_point_count = checked_area("assignment worker", worker_stripe_count, stripe_points)?; let worker_assignment_count = checked_area("assignment worker", worker_point_count, fanout)?; @@ -319,7 +327,12 @@ where Ok::<(), ANNError>(()) })?; - scatter_assignments(point_ids, &assignments, fanout, leader_count) + Ok(scatter_assignments( + point_ids, + &assignments, + fanout, + leader_count, + )) } /// Assign one point stripe to sampled partition centers. @@ -347,7 +360,9 @@ where let point_values_len = checked_area("point stripe", point_count, dimensions)?; // Keep each buffer at its largest length. Every operation uses an explicit // active prefix. - grow(&mut buffers.point_values, point_values_len, 0.0); + if buffers.point_values.len() < point_values_len { + buffers.point_values.resize(point_values_len, 0.0); + } let StripeBuffers { point_values, kernel_workspace, @@ -369,7 +384,8 @@ where T: VectorRepr, { for (&index, vector_output) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { - T::as_f32_into(data.row(index as usize), vector_output).map_err(Into::::into)?; + super::conversion::as_f32_into(data.row(index as usize), vector_output) + .map_err(Into::::into)?; } Ok(()) } @@ -383,13 +399,13 @@ fn scatter_assignments( assignments: &[u32], fanout: usize, leaders: usize, -) -> ANNResult>> { +) -> Vec> { if points.len() < PARALLEL_SCATTER_MIN_POINTS { - return Ok(scatter_serial(points, assignments, fanout, leaders)); + return scatter_serial(points, assignments, fanout, leaders); } - let stripe_points = points.len().div_ceil(rayon::current_num_threads().max(1)); - let stripe_assignment_count = checked_area("scatter assignment stripe", stripe_points, fanout)?; + let stripe_points = points.len().div_ceil(rayon::current_num_threads()); + let stripe_assignment_count = stripe_points * fanout; // Indexed parallel collection preserves stripe order. #[allow(clippy::disallowed_methods)] let locals: Vec<_> = points @@ -408,7 +424,7 @@ fn scatter_assignments( // `build_graph` runs this Rayon operation in the pool from the build context. // Each worker creates one independent leader cluster. #[allow(clippy::disallowed_methods)] - let clusters = sizes + sizes .into_par_iter() .enumerate() .map(|(leader, size)| { @@ -418,8 +434,7 @@ fn scatter_assignments( } cluster }) - .collect(); - Ok(clusters) + .collect() } fn scatter_serial( @@ -434,7 +449,7 @@ fn scatter_serial( sizes[leader as usize] += 1; } } - let mut clusters = clusters_with_capacities(&sizes); + let mut clusters: Vec> = sizes.into_iter().map(Vec::with_capacity).collect(); for (&point, point_assignments) in points.iter().zip(assignments.chunks_exact(fanout)) { for &leader in point_assignments { if leader != UNASSIGNED_LEADER { @@ -445,19 +460,11 @@ fn scatter_serial( clusters } -fn clusters_with_capacities(sizes: &[usize]) -> Vec> { - sizes.iter().map(|&size| Vec::with_capacity(size)).collect() -} - /// Merge leaves smaller than `c_min` without exceeding `c_max`. /// /// A `HashSet` removes duplicate point IDs across merged leaves. The function /// sorts each merged result before it returns. -fn merge_undersized_leaves( - leaves: Vec>, - c_min: usize, - c_max: usize, -) -> ANNResult>> { +fn merge_undersized_leaves(leaves: Vec>, c_min: usize, c_max: usize) -> Vec> { let mut merged = Vec::with_capacity(leaves.len()); let mut small_leaves = Vec::new(); for leaf in leaves { @@ -468,7 +475,7 @@ fn merge_undersized_leaves( } } if small_leaves.is_empty() { - return Ok(merged); + return merged; } let mut small = HashSet::with_capacity(c_max); @@ -501,7 +508,7 @@ fn merge_undersized_leaves( } } - Ok(merged) + merged } fn drain_sorted(set: &mut HashSet) -> Vec { @@ -510,23 +517,13 @@ fn drain_sorted(set: &mut HashSet) -> Vec { values } -fn point_ids(points: usize) -> Vec { - (0..points as u32).collect() -} - -fn grow(values: &mut Vec, len: usize, value: T) { - if values.len() < len { - values.resize(len, value); - } -} - fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { rows.checked_mul(cols) .ok_or_else(|| ANNError::new(PartitionError::ShapeOverflow { buffer, rows, cols })) } fn assignment_stripe_point_count(leader_count: usize) -> usize { - let point_count = ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count.max(1) * size_of::()); + let point_count = ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count * size_of::()); let point_count = if point_count.is_power_of_two() { point_count } else { @@ -542,6 +539,7 @@ mod tests { use diskann_wide::arch::{self, Target1}; use rstest::rstest; + use super::super::{Cosine, CosineNormalized, InnerProduct, L2}; use super::*; struct PartitionCall<'a, T> { @@ -557,8 +555,6 @@ mod tests { T: VectorRepr + Send + Sync, { fn run(self, arch: A, call: PartitionCall<'_, T>) -> ANNResult>> { - use super::super::{Cosine, CosineNormalized, InnerProduct, L2}; - match self.0 { Metric::L2 => partition::(arch, call.data, call.config), Metric::Cosine => partition::(arch, call.data, call.config), @@ -599,62 +595,14 @@ mod tests { } } - fn separated_point_clusters(points: usize, dimensions: usize) -> Matrix { - const POINTS_PER_CLUSTER: usize = 8; - const CLUSTER_SEPARATION: f32 = 10.0; - const POINT_OFFSET: f32 = 0.001; - const DIMENSION_OFFSET: f32 = 0.01; - - Matrix::new( - diskann_utils::views::Init({ - let mut position = 0usize; - move || { - let point = position / dimensions; - let dimension = position % dimensions; - position += 1; - let cluster = point / POINTS_PER_CLUSTER; - cluster as f32 * CLUSTER_SEPARATION - + point as f32 * POINT_OFFSET - + dimension as f32 * DIMENSION_OFFSET - } - }), - points, - dimensions, - ) - } - - fn unit_circle_points(points: usize, dimensions: usize) -> Matrix { - Matrix::new( - diskann_utils::views::Init({ - let mut position = 0usize; - move || { - let point = position / dimensions; - let dimension = position % dimensions; - position += 1; - let angle = std::f32::consts::TAU * point as f32 / points as f32; - match dimension { - 0 => angle.cos(), - 1 => angle.sin(), - _ => 0.0, - } - } - }), - points, - dimensions, - ) - } - - fn sorted_memberships(leaves: &[Vec]) -> Vec> { - let mut memberships: Vec> = leaves - .iter() - .map(|leaf| { - let mut ids = leaf.clone(); - ids.sort_unstable(); - ids + fn unit_circle_points(points: usize) -> Matrix { + let coordinates: Vec = (0..points) + .flat_map(|point| { + let angle = std::f32::consts::TAU * point as f32 / points as f32; + [angle.cos(), angle.sin()] }) .collect(); - memberships.sort(); - memberships + Matrix::try_from(coordinates.into_boxed_slice(), points, 2).unwrap() } fn assert_partition_invariants( @@ -663,378 +611,482 @@ mod tests { c_max: usize, replicas: usize, ) { - assert!( - leaves - .iter() - .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) - ); - let mut counts = vec![0usize; points]; + let mut occurrences = vec![0usize; points]; for leaf in leaves { - let mut ids = leaf.clone(); - ids.sort_unstable(); - ids.dedup(); - assert_eq!(ids.len(), leaf.len(), "duplicate ID inside a leaf"); + assert!(!leaf.is_empty() && leaf.len() <= c_max); + let unique_ids: HashSet<_> = leaf.iter().copied().collect(); + assert_eq!(unique_ids.len(), leaf.len(), "duplicate ID inside a leaf"); for &id in leaf { assert!((id as usize) < points); - counts[id as usize] += 1; + occurrences[id as usize] += 1; } } - assert!(counts.iter().all(|&count| count >= replicas)); + assert!(occurrences.iter().all(|&count| count >= replicas)); } - #[rstest] - #[case::below_c_max(7)] - #[case::at_c_max(8)] - fn partition_returns_one_leaf_when_point_count_does_not_exceed_c_max( - #[case] point_count: usize, - ) { - // Given - let data = separated_point_clusters(point_count, 3); - let expected_leaf = vec![(0..point_count as u32).collect::>()]; + mod partition_tests { + use super::*; + + #[rstest] + #[case::single_point(1)] + #[case::below_capacity(7)] + #[case::at_capacity(8)] + fn one_leaf_contains_all_points_when_no_split_is_needed(#[case] point_count: usize) { + // Given + let data = Matrix::new(0.0f32, point_count, 1); + let config = partition_config(2, 8, vec![2], 1); + let expected_leaves = vec![(0..point_count as u32).collect::>()]; + + // When + let leaves = + partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + + // Then + assert_eq!(leaves, expected_leaves); + } - // When - let actual_leaves = partition_with_runtime_metric( - data.as_view(), - &partition_config(2, 8, vec![2], 1), - Metric::L2, - ) - .unwrap(); + #[test] + fn each_replica_emits_its_leaf_when_no_split_is_needed() { + // Given + let data = MatrixView::column_vector(&[0.0f32, 1.0]); + let config = partition_config(1, 2, vec![1], 2); + let expected_leaves = vec![vec![0, 1], vec![0, 1]]; - // Then - assert_eq!(actual_leaves, expected_leaf); - } + // When + let leaves = partition_with_runtime_metric(data, &config, Metric::L2).unwrap(); - #[test] - fn partition_membership_is_deterministic_for_a_fixed_seed() { - // Given - let data = separated_point_clusters(96, 8); - let config = partition_config(4, 16, vec![3, 2], 2); + // Then + assert_eq!(leaves, expected_leaves); + } - // When - let first_partition = - partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); - let second_partition = - partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + #[test] + fn leaf_order_and_point_order_are_independent_of_worker_count() { + // Given: both runs split the same replicas through several levels. + let data = unit_circle_points(64); + let config = partition_config(2, 8, vec![3, 2], 2); + let single_worker = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + let three_workers = rayon::ThreadPoolBuilder::new() + .num_threads(3) + .build() + .unwrap(); + + // When + let serial = single_worker.install(|| { + partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap() + }); + let parallel = three_workers.install(|| { + partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap() + }); - // Then - assert_eq!( - sorted_memberships(&first_partition), - sorted_memberships(&second_partition) - ); - } + // Then: compare the full order; sorting would hide changes to recursive sampling. + assert_eq!(parallel, serial); + } - #[test] - fn partition_respects_capacity_and_replica_coverage() { - // Given - let data = separated_point_clusters(96, 8); - let config = partition_config(4, 16, vec![3, 2], 2); + #[rstest] + #[case::l2(Metric::L2)] + #[case::cosine(Metric::Cosine)] + #[case::normalized_cosine(Metric::CosineNormalized)] + #[case::inner_product(Metric::InnerProduct)] + fn bounded_leaves_preserve_point_coverage(#[case] metric: Metric) { + // Given: distinct unit vectors allow each metric to make progress. + let data = unit_circle_points(64); + let config = partition_config(2, 16, vec![3, 2], 3); + + // When + let leaves = partition_with_runtime_metric(data.as_view(), &config, metric).unwrap(); + + // Then + assert_partition_invariants(&leaves, 64, 16, 3); + } - // When - let partition = partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + #[test] + fn identical_points_return_an_iteration_limit_error() { + // Given: every point ties for the same first leader, so splitting cannot shrink it. + let data = Matrix::new(1.0f32, 5, 1); + let config = partition_config(1, 2, vec![1], 1); + + // When + let error = partition_with_runtime_metric(data.as_view(), &config, Metric::L2) + .unwrap_err() + .downcast::() + .unwrap(); + + // Then + assert!(matches!( + error, + PartitionError::IterationLimit { size: 5, .. } + )); + } - // Then - assert_partition_invariants(&partition, 96, 16, 2); + #[test] + fn no_leaves_are_emitted_when_every_distance_is_non_rankable() { + // Given: a split is required, but no point can be assigned to any leader. + let data = Matrix::new(f32::NAN, 3, 1); + let config = partition_config(1, 2, vec![1], 1); + + // When + let leaves = + partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + + // Then + assert!(leaves.is_empty()); + } } #[test] - fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { - let data = separated_point_clusters(80, 4); - let leaves = partition_with_runtime_metric( - data.as_view(), - &partition_config(2, 8, vec![2], 1), - Metric::L2, + fn split_assigns_one_leader_after_the_fanout_schedule_ends() { + // Given: all points are leaders; at level one the only configured fanout is exhausted. + let data = MatrixView::column_vector(&[0.0f32, 10.0, 20.0, 30.0]); + let config = PiPNNConfig { + p_samp: 1.0, + ..partition_config(1, 2, vec![2], 1) + }; + let parent = PendingPartition { + point_ids: vec![0, 1, 2, 3], + level: 1, + seed: PARTITION_SEED, + }; + let expected_leaves = vec![vec![0], vec![1], vec![2], vec![3]]; + + // When: each point's unique nearest leader is itself. + let mut split = split_partition::<_, L2, _>( + diskann_wide::ARCH, + data, + &config, + parent, + &StripeBufferPool::new((), 0, None), ) .unwrap(); + split.leaves.sort(); - assert_partition_invariants(&leaves, 80, 8, 1); + // Then + assert!(split.pending.is_empty()); + assert_eq!(split.leaves, expected_leaves); } - #[test] - fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { - let data = Matrix::new(1.0f32, 24, 4); - let error = partition_with_runtime_metric( - data.as_view(), - &partition_config(2, 4, vec![1], 1), - Metric::L2, - ) - .unwrap_err(); - let error = error.downcast::().unwrap(); + mod assign_to_leaders_tests { + use super::*; + + #[test] + fn clusters_contain_global_point_ids_in_input_order() { + // Given: the leader-column order differs from global ID order. + let data = MatrixView::column_vector(&[0.0f32, 100.0, 4.0, 96.0]); + let points = [3, 2, 1]; + let leaders = [1, 0]; + // Points 3 and 1 are nearest coordinate 100; point 2 is nearest coordinate 0. + let expected_clusters = vec![vec![3, 1], vec![2]]; + + // When + let clusters = assign_to_leaders::<_, L2, _>( + diskann_wide::ARCH, + data, + &points, + &leaders, + 1, + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); - assert!(matches!( - error, - PartitionError::IterationLimit { - size: 24, - limit: MAX_PARTITION_ITERATIONS, - .. - } - )); - } + // Then + assert_eq!(clusters, expected_clusters); + } - #[test] - fn global_merge_canonicalizes_small_leaf_membership() { - // Given - let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; - let expected_canonical_membership = vec![vec![1, 2, 3, 8, 9]]; + #[test] + fn fanout_larger_than_the_leader_set_assigns_every_leader_once() { + // Given + let data = MatrixView::column_vector(&[0.0f32, 10.0, 4.0]); + let expected_clusters = vec![vec![2], vec![2]]; + + // When + let clusters = assign_to_leaders::<_, L2, _>( + diskann_wide::ARCH, + data, + &[2], + &[0, 1], + 3, + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); - // When - let actual_leaves = merge_undersized_leaves(leaves, 4, 8).unwrap(); + // Then + assert_eq!(clusters, expected_clusters); + } - // Then - assert_eq!(actual_leaves, expected_canonical_membership); - } + #[test] + fn points_without_rankable_distances_are_omitted() { + // Given: the NaN point has no leader assignments; the finite point has two. + let data = MatrixView::column_vector(&[0.0f32, 10.0, f32::NAN, 1.0]); + let expected_clusters = vec![vec![3], vec![3]]; + + // When + let clusters = assign_to_leaders::<_, L2, _>( + diskann_wide::ARCH, + data, + &[2, 3], + &[0, 1], + 2, + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); - #[test] - fn global_merge_never_overfills_before_reaching_c_min() { - // Given - let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; - let expected_capacity_bounded_leaves = - vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]]; + // Then + assert_eq!(clusters, expected_clusters); + } - // When - let actual_leaves = merge_undersized_leaves(leaves, 11, 11).unwrap(); + #[rstest] + fn partial_final_stripe_preserves_point_assignments(#[values(1, 2)] worker_count: usize) { + // Given: three stripes, the last containing one point. One worker reuses its lease. + let point_count = 2 * MAX_ASSIGNMENT_STRIPE_POINTS + 1; + let data: Vec = (0..point_count).map(|point| point as f32).collect(); + let points: Vec = (0..point_count as u32).collect(); + let last_point = point_count as u32 - 1; + let leaders = [0, last_point]; + // On the line [0, last_point], the midpoint ties and selects leader column 0. + let midpoint = last_point / 2; + let expected_clusters = vec![ + (0..=midpoint).collect::>(), + (midpoint + 1..=last_point).collect(), + ]; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(worker_count) + .build() + .unwrap(); + + // When + let clusters = pool.install(|| { + assign_to_leaders::<_, L2, _>( + diskann_wide::ARCH, + MatrixView::column_vector(&data), + &points, + &leaders, + 1, + &StripeBufferPool::new((), 0, None), + ) + .unwrap() + }); - // Then - assert_eq!(actual_leaves, expected_capacity_bounded_leaves); + // Then + assert_eq!(clusters, expected_clusters); + } } - #[test] - fn global_merge_fills_exact_capacity_before_flushing() { - // Given - let leaves = vec![vec![0, 1], vec![2, 3]]; - let expected_exact_capacity_leaf = vec![vec![0, 1, 2, 3]]; + #[rstest] + #[case::f32([-2.0f32, 1.0, 3.0, -4.0, 5.0, 6.0], [5.0, 6.0, -2.0, 1.0, 5.0, 6.0])] + #[case::f16([-2.0f32, 1.0, 3.0, -4.0, 5.0, 6.0].map(Half::from_f32), [5.0, 6.0, -2.0, 1.0, 5.0, 6.0])] + #[case::i8([-2i8, 1, 3, -4, 5, 6], [5.0, 6.0, -2.0, 1.0, 5.0, 6.0])] + #[case::u8([2u8, 1, 3, 4, 5, 6], [5.0, 6.0, 2.0, 1.0, 5.0, 6.0])] + fn gather_converts_rows_in_requested_id_order( + #[case] values: [T; 6], + #[case] expected: [f32; 6], + ) { + // Given: request the final row, then the first row, then repeat the final row. + let data = MatrixView::try_from(values.as_slice(), 3, 2).unwrap(); + let mut gathered = [f32::NAN; 6]; // When - let actual_leaves = merge_undersized_leaves(leaves, 4, 4).unwrap(); + gather_vectors(data, &[2, 0, 2], &mut gathered).unwrap(); // Then - assert_eq!(actual_leaves, expected_exact_capacity_leaf); + assert_eq!(gathered, expected); } - #[test] - fn replicas_cover_every_point_once_or_more_per_replica() { - let data = unit_circle_points(72, 5); - let leaves = partition_with_runtime_metric( - data.as_view(), - &partition_config(3, 12, vec![3, 2], 3), - Metric::CosineNormalized, - ) - .unwrap(); + mod merge_undersized_leaves_tests { + use super::*; - assert_partition_invariants(&leaves, 72, 12, 3); - } + #[test] + fn total_below_minimum_is_kept_as_one_leaf() { + // Given: all available points together still fall below the minimum. + let leaves = vec![vec![0], vec![1]]; + let expected_leaves = vec![vec![0, 1]]; - fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) - where - T: crate::utils::VectorRepr + Send + Sync, - { - let points = 64; - // Partition gather converts source vectors before GEMM. Test conversion - // tails around 4, 8, 16, and 32 elements. - for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let raw: Vec = (0..points * dimensions) - .map(|index| { - let point = index / dimensions; - let dimension = index % dimensions; - (point + dimension) as u8 - }) - .collect(); - let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); - let converted: Vec = raw.iter().copied().map(&convert).collect(); - let config = partition_config(2, 16, vec![2, 1], 1); - let expected_f32_partition = partition_with_runtime_metric( - MatrixView::try_from(&f32_data, points, dimensions).unwrap(), - &config, - Metric::L2, - ) - .unwrap(); - let actual_converted_partition = partition_with_runtime_metric( - MatrixView::try_from(&converted, points, dimensions).unwrap(), - &config, - Metric::L2, - ) - .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); - - assert_partition_invariants(&actual_converted_partition, points, 16, 1); - assert_eq!( - sorted_memberships(&actual_converted_partition), - sorted_memberships(&expected_f32_partition), - "{label} dimensions={dimensions}" - ); + // When + let merged = merge_undersized_leaves(leaves, 3, 4); + + // Then + assert_eq!(merged, expected_leaves); } - } - #[test] - fn f16_partition_matches_f32_across_dimension_boundaries() { - assert_partition_conversion_matches_f32("f16", |value| Half::from_f32(value as f32)); - } + #[test] + fn leaves_already_at_minimum_keep_their_order() { + // Given + let leaves = vec![vec![3, 1, 2], vec![7, 5, 6]]; + let expected_leaves = leaves.clone(); - #[test] - fn u8_partition_matches_f32_across_dimension_boundaries() { - assert_partition_conversion_matches_f32("u8", |value| value); - } + // When + let merged = merge_undersized_leaves(leaves, 3, 4); - #[test] - fn i8_partition_matches_f32_across_dimension_boundaries() { - // The same translation in every coordinate preserves L2 ordering. - assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); - } + // Then + assert_eq!(merged, expected_leaves); + } - #[test] - fn leader_assignment_preserves_sequential_norm_reduction_with_l2() { - fn next_reassociation_regression_value(state: &mut u64) -> f32 { - *state ^= *state << 13; - *state ^= *state >> 7; - *state ^= *state << 17; - (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 + #[test] + fn merged_leaves_contain_sorted_unique_ids() { + // Given: the shared ID 1 counts only once toward the minimum of three. + let leaves = vec![vec![3, 1], vec![1, 2]]; + let expected_leaves = vec![vec![1, 2, 3]]; + + // When + let merged = merge_undersized_leaves(leaves, 3, 4); + + // Then + assert_eq!(merged, expected_leaves); } - // Sequential and SIMD-reassociated leader norms select different top-1 - // leaders for this case. Dot products still use the production GEMM. The - // test changes only the leader-norm reduction. - const REDUCTION_BOUNDARY_DIMENSIONS: usize = 129; - const REASSOCIATION_REGRESSION_SEED: u64 = 0x3a85_f952_c718_6e49; - let mut state = REASSOCIATION_REGRESSION_SEED; - let point: Vec = (0..REDUCTION_BOUNDARY_DIMENSIONS) - .map(|_| next_reassociation_regression_value(&mut state)) - .collect(); - let leader_zero: Vec = (0..REDUCTION_BOUNDARY_DIMENSIONS) - .map(|_| next_reassociation_regression_value(&mut state)) - .collect(); - let leader_one: Vec = (0..REDUCTION_BOUNDARY_DIMENSIONS) - .map(|_| next_reassociation_regression_value(&mut state)) - .collect(); - let data: Vec = leader_zero - .into_iter() - .chain(leader_one) - .chain(point) - .collect(); - let data = MatrixView::try_from(data.as_slice(), 3, REDUCTION_BOUNDARY_DIMENSIONS).unwrap(); - let expected_clusters = [vec![], vec![2]]; + #[test] + fn full_batch_is_emitted_before_another_leaf_would_exceed_capacity() { + // Given: three pairs cannot fit in a five-point leaf; the first two can. + let leaves = vec![vec![0, 1], vec![2, 3], vec![4, 5]]; + let expected_leaves = vec![vec![0, 1, 2, 3], vec![4, 5]]; - let actual_clusters = assign_to_leaders::<_, super::super::L2, _>( - diskann_wide::ARCH, - data, - &[2], - &[0, 1], - 1, - &StripeBufferPool::new((), 0, None), - ) - .unwrap(); + // When + let merged = merge_undersized_leaves(leaves, 5, 5); - assert_eq!(actual_clusters, expected_clusters); - } + // Then + assert_eq!(merged, expected_leaves); + } - #[rstest] - fn partition_covers_all_points_without_exceeding_c_max( - #[values( - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct - )] - metric: Metric, - ) { - // Given - let data = unit_circle_points(64, 8); - let config = partition_config(2, 20, vec![2], 1); + #[test] + fn leaves_that_exactly_fill_capacity_are_combined() { + // Given + let leaves = vec![vec![0, 1], vec![2, 3]]; + let expected_leaves = vec![vec![0, 1, 2, 3]]; - // When - let leaves = partition_with_runtime_metric(data.as_view(), &config, metric).unwrap(); + // When + let merged = merge_undersized_leaves(leaves, 4, 4); - // Then - assert_partition_invariants(&leaves, 64, 20, 1); + // Then + assert_eq!(merged, expected_leaves); + } + + #[test] + fn overlapping_remainder_fits_after_duplicate_ids_are_removed() { + // Given: raw lengths total five, but the union has four IDs and fits. + let leaves = vec![vec![0, 1, 2], vec![2, 3]]; + let expected_leaves = vec![vec![0, 1, 2, 3]]; + + // When + let merged = merge_undersized_leaves(leaves, 3, 4); + + // Then + assert_eq!(merged, expected_leaves); + } + + #[test] + fn new_remainder_ids_form_a_separate_leaf_when_the_last_leaf_is_full() { + // Given: ID 3 already exists, but the new ID 4 cannot fit in the full leaf. + let leaves = vec![vec![0, 1, 2, 3], vec![3, 4]]; + let expected_leaves = vec![vec![0, 1, 2, 3], vec![4]]; + + // When + let merged = merge_undersized_leaves(leaves, 3, 4); + + // Then + assert_eq!(merged, expected_leaves); + } + + #[test] + fn entirely_overlapping_remainder_adds_no_leaf() { + // Given: every remainder ID is already present in the last leaf. + let leaves = vec![vec![0, 1, 2], vec![1, 2]]; + let expected_leaves = vec![vec![0, 1, 2]]; + + // When + let merged = merge_undersized_leaves(leaves, 3, 4); + + // Then + assert_eq!(merged, expected_leaves); + } } #[rstest] - #[case::single_point_minimum(1, 1.0, 1)] - #[case::sampled_count(10, 0.01, 2)] + #[case::minimum_two_leaders(10, 0.01, 2)] + #[case::fraction_rounds_up(10, 0.25, 3)] #[case::leader_cap(50_000, 1.0, LEADER_CAP)] - fn sampled_leader_count_respects_data_size_sampling_and_cap( + fn sampled_leader_count_respects_sampling_bounds( #[case] point_count: usize, - #[case] sampling_probability: f64, + #[case] sampling_fraction: f64, #[case] expected_leader_count: usize, ) { - assert_eq!( - sampled_leader_count(point_count, sampling_probability), - expected_leader_count - ); - } + // When + let count = sampled_leader_count(point_count, sampling_fraction); - #[rstest] - #[case::first_replica(0, 1_000)] - #[case::second_replica(1, 8_919)] - fn replica_seed_is_stable(#[case] replica: usize, #[case] expected_seed: u64) { - assert_eq!(replica_seed(replica), expected_seed); + // Then + assert_eq!(count, expected_leader_count); } - #[test] - fn leader_assignment_preserves_clusters_across_multiple_stripes() { - let points = 2_048; - let data: Vec = (0..points).map(|point| point as f32).collect(); - let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); - let point_ids: Vec = (0..points as u32).collect(); + mod scatter_assignments_tests { + use super::*; - let clusters = assign_to_leaders::<_, super::super::L2, _>( - diskann_wide::ARCH, - data, - &point_ids, - &[0, 2_047], - 1, - &StripeBufferPool::new((), 0, None), - ) - .unwrap(); - - assert_eq!(clusters[0], (0..1_024).collect::>()); - assert_eq!(clusters[1], (1_024..2_048).collect::>()); - } + #[test] + fn unassigned_slots_are_skipped_without_reordering_points() { + // Given: each assignment pair belongs to one point; leader 2 has no points. + let points = [30, 10, 20]; + let assignments = [1, UNASSIGNED_LEADER, 0, 1, UNASSIGNED_LEADER, 0]; + let expected_clusters = vec![vec![10, 20], vec![30, 10], vec![]]; - #[test] - fn scatter_omits_unassigned_slots() { - // Given - let points = [10, 11]; - let assignments = [0, UNASSIGNED_LEADER, UNASSIGNED_LEADER, 1]; - let expected_clusters = [vec![10], vec![11]]; + // When + let clusters = scatter_assignments(&points, &assignments, 2, 3); - // When - let actual_clusters = scatter_serial(&points, &assignments, 2, 2); + // Then + assert_eq!(clusters, expected_clusters); + } - // Then - assert_eq!(actual_clusters, expected_clusters); + #[rstest] + fn parallel_stripes_preserve_order_and_skip_unassigned_slots( + #[values(1, 3)] worker_count: usize, + ) { + // Given: three repeating assignments, with a partial final worker chunk. + let points: Vec = (0..(PARALLEL_SCATTER_MIN_POINTS as u32 + 1)) + .rev() + .collect(); + let pattern = [0, UNASSIGNED_LEADER, 1, 0, UNASSIGNED_LEADER, 1]; + let assignments: Vec = + pattern.into_iter().cycle().take(2 * points.len()).collect(); + // Leader 0 receives rows 0 and 1 of each triple; leader 1 receives rows 1 and 2. + let expected_clusters = vec![ + points + .iter() + .enumerate() + .filter(|(row, _)| row % 3 != 2) + .map(|(_, &id)| id) + .collect::>(), + points + .iter() + .enumerate() + .filter(|(row, _)| row % 3 != 0) + .map(|(_, &id)| id) + .collect(), + vec![], + ]; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(worker_count) + .build() + .unwrap(); + + // When + let clusters = pool.install(|| scatter_assignments(&points, &assignments, 2, 3)); + + // Then + assert_eq!(clusters, expected_clusters); + } } #[test] - fn parallel_scatter_matches_serial_order() { - // Given - let leader_count = 7; - let leaders_per_point = 2; - let second_leader_offset = 3; - let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); - let assignments: Vec = points - .iter() - .flat_map(|point| { - [ - point % leader_count, - (point + second_leader_offset) % leader_count, - ] - }) - .collect(); + fn overflowing_buffer_shape_returns_an_error() { + // Given: this shape cannot be represented by an allocation length. + let rows = usize::MAX; + let columns = 2; // When - let expected_serial_clusters = scatter_serial( - &points, - &assignments, - leaders_per_point, - leader_count as usize, - ); - let actual_parallel_clusters = scatter_assignments( - &points, - &assignments, - leaders_per_point, - leader_count as usize, - ) - .unwrap(); + let error = checked_area("partition assignments", rows, columns).unwrap_err(); // Then - assert_eq!(actual_parallel_clusters, expected_serial_clusters); + assert!(matches!( + error.downcast::().unwrap(), + PartitionError::ShapeOverflow { .. } + )); } } From fca1df9a46046421bb0ca7d46635964434806f01 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 14 Sep 2026 05:28:42 +0000 Subject: [PATCH 178/229] perf(pipnn): dispatch vector gathers per batch --- diskann/src/graph/pipnn/conversion.rs | 120 ++++++++++++------------ diskann/src/graph/pipnn/leaf_build.rs | 20 ++-- diskann/src/graph/pipnn/partitioning.rs | 6 +- 3 files changed, 66 insertions(+), 80 deletions(-) diff --git a/diskann/src/graph/pipnn/conversion.rs b/diskann/src/graph/pipnn/conversion.rs index 9c93a88223..f8482ba0c6 100644 --- a/diskann/src/graph/pipnn/conversion.rs +++ b/diskann/src/graph/pipnn/conversion.rs @@ -8,83 +8,79 @@ use std::any::TypeId; use crate::utils::VectorRepr; +use diskann_utils::views::MatrixView; use diskann_vector::conversion::SliceCast; -use diskann_wide::arch; +use diskann_wide::{ + Architecture, + arch::{self, Target2, Target3}, +}; use half::f16; -/// Preserve VectorRepr conversion errors while accelerating FP16 slices. -#[inline] -pub(super) fn as_f32_into(source: &[T], output: &mut [f32]) -> Result<(), T::Error> { - if TypeId::of::() == TypeId::of::() && source.len() == output.len() { - // VectorRepr's default conversion uses the compile-time architecture. - // Dispatch one whole slice so portable builds can use F16C or NEON too. - let source: &[f16] = bytemuck::cast_slice(source); - arch::dispatch2(SliceCast::::new(), output, source); - Ok(()) +/// Gather rows into packed f32 storage, dispatching once for the whole FP16 batch. +pub(super) fn gather_as_f32( + data: MatrixView<'_, T>, + ids: &[u32], + output: &mut [f32], +) -> Result<(), (u32, T::Error)> { + if TypeId::of::() == TypeId::of::() { + arch::dispatch3(GatherFp16, data, ids, output); } else { - T::as_f32_into(source, output) + for (&id, destination) in ids.iter().zip(output.chunks_exact_mut(data.ncols())) { + T::as_f32_into(data.row(id as usize), destination).map_err(|error| (id, error))?; + } + } + Ok(()) +} + +struct GatherFp16; + +impl Target3, &[u32], &mut [f32]> for GatherFp16 +where + A: Architecture, + T: VectorRepr, + for<'a, 'b> SliceCast: Target2, +{ + #[inline(always)] + fn run(self, arch: A, data: MatrixView<'_, T>, ids: &[u32], output: &mut [f32]) { + for (&id, destination) in ids.iter().zip(output.chunks_exact_mut(data.ncols())) { + let source = bytemuck::cast_slice(data.row(id as usize)); + SliceCast::::new().run(arch, destination, source); + } } } #[cfg(test)] -mod as_f32_into_tests { +mod gather_as_f32_tests { use super::*; - use rstest::rstest; #[test] - fn fp16_bits_are_preserved_through_a_vector_and_tail() { - // Given: 13 values leave a tail for both four- and eight-lane converters. - // Normal values gain 112 in exponent bias and shift the fraction by 13; - // the smallest FP16 subnormal becomes 2^-24. NaN keeps its payload. - let half_bits = [ - 0x0000, 0x8000, 0x0001, 0x03ff, 0x0400, 0x3c00, 0xbc00, 0x3555, 0x7bff, 0x7c00, 0xfc00, - 0x7e01, 0x8400, - ]; - let expected = [ - 0x00000000, 0x80000000, 0x33800000, 0x387fc000, 0x38800000, 0x3f800000, 0xbf800000, - 0x3eaaa000, 0x477fe000, 0x7f800000, 0xff800000, 0x7fc02000, 0xb8800000, + fn fp16_gather_preserves_bits_and_requested_row_order_with_a_tail() { + // Given: repeated, noncontiguous rows of nine elements exercise tails + // and the row stride while overwriting stale destination values. + let points = [ + [-4.0_f32, -3.5, -2.0, -1.5, -0.0, 0.5, 1.0, 2.5, 3.0], + [1.0, 2.5, 3.0, 0.5, -1.0, 4.0, -2.5, 0.0, -3.5], + [8.0, 2.0, -7.0, 0.25, 0.75, 3.5, -4.5, 0.0, -2.0], ]; - let source = half_bits.map(f16::from_bits); - let mut output = [123.0_f32; 13]; - - // When - as_f32_into(&source, &mut output).unwrap(); - - // Then - assert_eq!(output.map(f32::to_bits), expected); - } - - #[rstest] - #[case::short_output(7)] - #[case::long_output(9)] - fn mismatched_fp16_lengths_return_the_vector_repr_error(#[case] output_len: usize) { - // Given - let source = [f16::ONE; 8]; - let mut output = vec![0.0_f32; output_len]; - let expected = f16::as_f32_into(&source, &mut output).unwrap_err(); - - // When - let actual = as_f32_into(&source, &mut output).unwrap_err(); - - // Then - assert_eq!(actual, expected); - } - - #[rstest] - #[case::f32([-2.5_f32, 1.25], [-2.5, 1.25])] - #[case::i8([-2_i8, 127], [-2.0, 127.0])] - #[case::u8([0_u8, 255], [0.0, 255.0])] - fn other_representations_keep_their_conversion( - #[case] source: [T; 2], - #[case] expected: [f32; 2], - ) { - // Given - let mut output = [f32::NAN; 2]; + let half_points = points.map(|row| row.map(f16::from_f32_const)); + let data = MatrixView::try_from(half_points.as_flattened(), 3, 9).unwrap(); + let ids = [2, 0, 2]; + // All coordinates are exactly representable in FP16, including -0.0. + let expected = [points[2], points[0], points[2]]; + let mut actual = [123.0_f32; 27]; // When - as_f32_into(&source, &mut output).unwrap(); + gather_as_f32(data, &ids, &mut actual).unwrap(); // Then - assert_eq!(output, expected); + assert_eq!( + actual.map(f32::to_bits).as_slice(), + expected + .as_flattened() + .iter() + .copied() + .map(f32::to_bits) + .collect::>() + ); } } diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 8683db3daa..2f09f38064 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -182,19 +182,13 @@ where let point_value_count = point_ids.len() * data.ncols(); let point_values = &mut buffers.point_values[..point_value_count]; - for (&point, point_output) in point_ids - .iter() - .zip(point_values.chunks_exact_mut(data.ncols())) - { - let source_values = data.row(point as usize); - super::conversion::as_f32_into(source_values, point_output).map_err(|source| { - LeafBuildError::Conversion { - leaf, - point, - source: source.into(), - } - })?; - } + super::conversion::gather_as_f32(data, point_ids, point_values).map_err( + |(point, source)| LeafBuildError::Conversion { + leaf, + point, + source: source.into(), + }, + )?; let points = MatrixView::try_from(&*point_values, point_ids.len(), data.ncols()) .expect("point buffer prefix has the checked leaf shape"); diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index f187a20f3c..9ef05983bc 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -383,11 +383,7 @@ fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32 where T: VectorRepr, { - for (&index, vector_output) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { - super::conversion::as_f32_into(data.row(index as usize), vector_output) - .map_err(Into::::into)?; - } - Ok(()) + super::conversion::gather_as_f32(data, indices, output).map_err(|(_, error)| error.into()) } /// Group assigned point IDs by child partition. From 08c2fba93b816986687c5739e98c7e953a411b19 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 17 Sep 2026 02:14:23 +0000 Subject: [PATCH 179/229] refactor(pipnn): localize insertion slicing --- diskann/src/graph/pipnn/topk.rs | 58 ++++++++++++++++++++++++++++----- 1 file changed, 49 insertions(+), 9 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 9012caff27..596166b341 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -176,11 +176,11 @@ impl TopK { arch.run2( #[inline(always)] move |distances: &[f32], output: &mut [Candidate]| { - let output = &mut output[..self.width.capacity()]; - if output.is_empty() { + let k = self.width.capacity(); + if k == 0 { return; } - output.fill(Candidate::default()); + output[..k].fill(Candidate::default()); distance_blocks(arch, distances).fold( f32::INFINITY, #[inline(always)] @@ -224,12 +224,11 @@ impl TopK { ); // Keep this point's result and local limit independent of reciprocal updates. let (others, remaining) = output.as_mut_slice().split_at_mut(point_idx * k); - let nearest = &mut remaining[..k]; let limit = distance_blocks(arch, distances).fold( thresholds[point_idx], #[inline(always)] |limit, block| { - let limit = block.update_one(nearest, limit, self.width); + let limit = block.update_one(remaining, limit, self.width); block.update_many(others, thresholds, point_idx as u32, self.width); limit }, @@ -364,7 +363,7 @@ impl DistanceBlock<'_, A> { let limit = &mut thresholds[candidate_idx]; if distance < *limit { *limit = insert_sorted( - &mut candidates[candidate_idx * k..][..k], + &mut candidates[candidate_idx * k..], Candidate::new(point_idx, distance), width, ); @@ -384,7 +383,7 @@ impl DistanceBlock<'_, A> { eligible &= eligible - 1; // Each lane updates a different result; its limit is current. limits[lane] = insert_sorted( - &mut candidates[(first_candidate + lane) * k..][..k], + &mut candidates[(first_candidate + lane) * k..], Candidate::new(point_idx, distances[lane]), width, ); @@ -395,11 +394,16 @@ impl DistanceBlock<'_, A> { } /// Insert an eligible candidate in nearest-first order and return the new distance limit. +/// Only the first `width.capacity()` slots belong to this result. /// The caller must check the candidate against the current limit before insertion. -/// The width exposes fixed capacities to loop unrolling. #[inline(always)] fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate, width: W) -> f32 { - let last = width.capacity() - 1; + let k = width.capacity(); + // Keep this slice here for performance. + // Fixed gives the slice a compile-time constant length. + // After inlining, LLVM can unroll the loop and remove bounds checks. + let nearest = &mut nearest[..k]; + let last = nearest.len() - 1; let mut slot = last; while slot > 0 && candidate.distance < nearest[slot - 1].distance { nearest[slot] = nearest[slot - 1]; @@ -519,6 +523,42 @@ mod tests { assert_eq!(thresholds, [f32::INFINITY; 2]); } + mod insert_sorted_tests { + use super::*; + + #[rstest] + #[case::fixed(Fixed::<3>)] + #[case::runtime(Runtime(3))] + fn insertion_preserves_later_result_rows(#[case] width: impl Width) { + // Given: the new candidate displaces the first result's farthest neighbor. + let first_result = [ + Candidate::new(0, 1.0), + Candidate::new(1, 4.0), + Candidate::new(2, 6.0), + ]; + let next_result = [ + Candidate::new(3, 8.0), + Candidate::new(4, 10.0), + Candidate::new(5, 12.0), + ]; + let candidate = Candidate::new(6, 3.0); + let expected_first_result = [ + Candidate::new(0, 1.0), + Candidate::new(6, 3.0), + Candidate::new(1, 4.0), + ]; + let expected_limit = 4.0; + let mut results = [first_result, next_result]; + + // When + let limit = insert_sorted(results.as_flattened_mut(), candidate, width); + + // Then + assert_eq!(results, [expected_first_result, next_result]); + assert_eq!(limit, expected_limit); + } + } + mod select_topk_tests { use super::*; From 1bc033d1112c06b17a9860bfb264c01f373d64d3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 17 Sep 2026 02:30:24 +0000 Subject: [PATCH 180/229] refactor(pipnn): name top-k update directions --- diskann/src/graph/pipnn/topk.rs | 40 +++++++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 596166b341..00637d1a9b 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -184,7 +184,7 @@ impl TopK { distance_blocks(arch, distances).fold( f32::INFINITY, #[inline(always)] - |limit, block| block.update_one(output, limit, self.width), + |limit, block| block.update_topk_from_candidates(output, limit, self.width), ); }, distances, @@ -228,8 +228,13 @@ impl TopK { thresholds[point_idx], #[inline(always)] |limit, block| { - let limit = block.update_one(remaining, limit, self.width); - block.update_many(others, thresholds, point_idx as u32, self.width); + let limit = block.update_topk_from_candidates(remaining, limit, self.width); + block.update_topks_with_point( + others, + thresholds, + point_idx as u32, + self.width, + ); limit }, ); @@ -285,9 +290,17 @@ fn distance_blocks( } impl DistanceBlock<'_, A> { - /// Offer the block's candidates, preserving and returning the updated limit. + /// Update one top-k result from this block's candidates. + /// + /// Each candidate uses its index in the full distance row as its local ID. + /// Partition ranking uses leader IDs. Leaf ranking uses IDs of earlier points in the leaf. + /// + /// Compare candidates against the existing result without clearing its slots. + /// The first `width.capacity()` slots of `nearest` hold this result. + /// The distance limit is the last slot's distance, or positive infinity while slots remain empty. + /// Each insertion can lower this limit. Return the updated limit for the next block. #[inline(always)] - fn update_one( + fn update_topk_from_candidates( self, nearest: &mut [Candidate], mut max_distance: f32, @@ -345,9 +358,22 @@ impl DistanceBlock<'_, A> { max_distance } - /// Offer one point to each result set identified by this block's candidate IDs. + /// Offer one leaf point to the top-k results of earlier leaf points. + /// + /// This block stores distances from `point_idx` to those earlier points. + /// Each candidate's local ID selects its result row in the full leaf, not its SIMD lane. + /// If that row accepts the pair, insert `point_idx` with the pair's distance. + /// + /// `candidates` starts at leaf row zero, with `width.capacity()` slots per row. + /// `thresholds` holds one distance limit per row. + /// The caller initializes both buffers once before the leaf scan. + /// These updates retain earlier candidates unless a nearer point displaces them. + /// + /// Each row compares the pair distance against its own limit. + /// Its decision does not depend on the current point's top-k result. + /// Update the row's limit after an insertion. #[inline(always)] - fn update_many( + fn update_topks_with_point( self, candidates: &mut [Candidate], thresholds: &mut [f32], From fa560c33789f69cc3379d3ade6d728ff233e2f6d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 17 Sep 2026 08:20:47 +0000 Subject: [PATCH 181/229] test(pipnn): cover numerical kernel matrices Check complete metric matrices and count/width ranking grids. Keep ranking and distance tolerances separate. Clarify test names and reusable-state fixtures. Run GEMM-backed tests natively and keep them out of Miri. --- diskann/src/graph/pipnn/leaf_kernel.rs | 116 +++++- diskann/src/graph/pipnn/leaf_metric.rs | 71 +++- diskann/src/graph/pipnn/mod.rs | 153 +++++++- diskann/src/graph/pipnn/partition_kernel.rs | 77 +++- diskann/src/graph/pipnn/partition_metric.rs | 142 ++++++-- diskann/src/graph/pipnn/topk.rs | 378 ++++++++++++++------ 6 files changed, 768 insertions(+), 169 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 6827b6b9dd..b667d17628 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -136,12 +136,36 @@ fn validate_output( #[cfg(test)] mod tests { use super::*; - use crate::graph::pipnn::L2; mod test_support { use super::*; + #[cfg(not(miri))] + use diskann_wide::arch::Target2; use diskann_wide::arch::{self, Target1}; + #[cfg(not(miri))] + pub(super) struct SelectLeaf(pub(super) M); + + #[cfg(not(miri))] + impl + Target2, MatrixView<'_, f32>, MutMatrixView<'_, Candidate>> + for SelectLeaf + { + fn run( + self, + arch: A, + points: MatrixView<'_, f32>, + output: MutMatrixView<'_, Candidate>, + ) -> ANNResult<()> { + select_leaf_neighbors::<_, M>( + arch, + points, + output, + &mut LeafKernelWorkspace::default(), + ) + } + } + struct KernelCall<'a> { distances: MatrixView<'a, f32>, output: MutMatrixView<'a, Candidate>, @@ -189,8 +213,92 @@ mod tests { } } + // These entry-point tests call GEMM. Miri checks the distance scan below. + #[cfg(not(miri))] mod select_leaf_neighbors_tests { + use super::test_support::SelectLeaf; use super::*; + use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2, scalar_ranking}; + use diskann_vector::distance::Metric; + use diskann_wide::arch; + use rstest::rstest; + + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn leaf_neighbors_match_scalar_ranking( + #[case] metric: M, + #[case] scalar_metric: Metric, + #[values(2, 7, 8, 9, 15, 16, 17, 127, 128, 129, 384, 768)] dimensions: usize, + #[values(1, 3, 10, 11)] width: usize, + ) { + // Given: 34 points supply two SIMD groups and a tail in the last row. + // All coordinates contribute. Only normalized cosine receives unit vectors. + let values = scalar_ranking::arc_vectors( + (0..34).map(f64::from), + dimensions, + scalar_metric == Metric::CosineNormalized, + ); + let points = MatrixView::try_from(values.as_slice(), 34, dimensions).unwrap(); + let expected: Vec> = points + .row_iter() + .enumerate() + .map(|(source, point)| { + let mut row: Vec<_> = points + .row_iter() + .enumerate() + .filter(|&(target, _)| source != target) + .map(|(target, vector)| { + ( + target as u32, + scalar_ranking::distance(scalar_metric, point, vector), + ) + }) + .collect(); + row.sort_by(|left, right| left.1.total_cmp(&right.1)); + row + }) + .collect(); + let ranking_tolerance = scalar_ranking::ranking_tolerance(scalar_metric, dimensions); + let mut output = vec![Candidate::default(); 34 * width]; + + // When: dispatch the complete points-to-metric-to-neighbors pipeline. + arch::dispatch2_no_features( + SelectLeaf(metric), + points, + MutMatrixView::try_from(output.as_mut_slice(), 34, width).unwrap(), + ) + .unwrap(); + + // Then: compare in returned order, including each ID's own scalar score. + for (source, (row, expected)) in output.chunks_exact(width).zip(&expected).enumerate() { + let ids: Vec<_> = row.iter().map(|candidate| candidate.local_idx).collect(); + scalar_ranking::assert_ranked_ids(&ids, expected, ranking_tolerance); + assert!( + row.windows(2) + .all(|pair| pair[0].distance <= pair[1].distance), + "unordered row {source}" + ); + for candidate in row { + let score = expected + .iter() + .find(|&&(id, _)| id == candidate.local_idx) + .unwrap() + .1; + let distance_tolerance = scalar_ranking::distance_tolerance( + scalar_metric, + points.row(source), + points.row(candidate.local_idx as usize), + ); + assert!( + (f64::from(candidate.distance) - score).abs() <= distance_tolerance, + "row {source}, candidate {candidate:?}, scalar {score}, tolerance {distance_tolerance}" + ); + } + } + } #[rstest::rstest] #[case::missing_row(2)] @@ -215,7 +323,7 @@ mod tests { } #[test] - fn requesting_self_as_an_extra_neighbor_returns_an_error() { + fn excess_neighbor_count_returns_error() { let values = [0.0_f32, 1.0, 3.0]; let mut output = [Candidate::default(); 9]; let mut workspace = LeafKernelWorkspace::default(); @@ -237,7 +345,7 @@ mod tests { } #[test] - fn l2_neighbors_follow_new_geometry_when_workspace_is_reused() { + fn l2_neighbors_update_with_reused_workspace() { // Given: squared distances on the line determine each point's two neighbors. let values = [0.0_f32, 1.0, 3.0, 10.0]; let expected = [ @@ -305,7 +413,7 @@ mod tests { use super::*; #[test] - fn lower_triangle_pairs_update_both_endpoints_across_a_simd_group_and_tail() { + fn lower_triangle_pairs_update_both_endpoints() { // Given: only three lower-triangle pairs rank. The diagonal and upper // triangle contain better distances, so reading either changes the answer. let mut distances = [f32::NEG_INFINITY; 18 * 18]; diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index bbd6dd2436..f1ae207b8f 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -107,9 +107,13 @@ impl LeafMetric for CosineNormalized { } } +// GEMM-backed tests run natively, not under Miri. #[cfg(test)] +#[cfg(not(miri))] mod tests { use super::*; + use crate::graph::pipnn::scalar_ranking; + use diskann_vector::distance::Metric; use rstest::rstest; const POINT_COUNT: usize = 2; @@ -140,8 +144,63 @@ mod tests { mod compute_distances_tests { use super::*; + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn lower_triangle_matches_scalar_ranking( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[values(64, 128, 256)] point_count: usize, + #[values(128, 384, 768)] dimensions: usize, + ) { + // Given: spread all points over the same arc as the matrix grows. + // NaN marks unwritten entries. The scalar oracle includes self pairs. + let values = scalar_ranking::arc_vectors( + (0..point_count).map(|point| 32.0 * point as f64 / point_count as f64), + dimensions, + scalar_metric == Metric::CosineNormalized, + ); + let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let expected: Vec> = points + .row_iter() + .enumerate() + .map(|(source, point)| { + points + .row_iter() + .take(source + 1) + .map(|target| { + ( + scalar_ranking::distance(scalar_metric, point, target), + scalar_ranking::distance_tolerance(scalar_metric, point, target), + ) + }) + .collect() + }) + .collect(); + let mut distances = vec![f32::NAN; point_count * point_count]; + + // When: call the metric directly, without TopK selection. + M::compute_distances(points, &mut distances).unwrap(); + + // Then: check every required entry, including the diagonal. + // The upper triangle is outside the metric contract. + let actual = + MatrixView::try_from(distances.as_slice(), point_count, point_count).unwrap(); + for (source, row) in expected.iter().enumerate() { + for (target, &(expected, tolerance)) in row.iter().enumerate() { + let actual = f64::from(actual.row(source)[target]); + assert!( + (actual - expected).abs() <= tolerance, + "{scalar_metric:?}, {point_count} points, D={dimensions}, pair ({source}, {target}): actual {actual}, expected {expected}, tolerance {tolerance}" + ); + } + } + } + #[test] - fn squared_l2_ranking_equals_the_sum_of_squared_coordinate_differences() { + fn l2_writes_squared_distances() { // Given let first_point = [3.0_f32, 4.0]; let second_point = [0.0_f32, 4.0]; @@ -157,7 +216,7 @@ mod tests { } #[test] - fn cosine_ranking_equals_one_minus_normalized_similarity() { + fn cosine_writes_one_minus_similarity() { // Given let first_point = [2.0_f32, 0.0]; let second_point = [1.0_f32, 1.0]; @@ -179,7 +238,7 @@ mod tests { #[rstest] #[case::zero([0.0, 0.0])] #[case::subnormal([f32::MIN_POSITIVE.sqrt() / 2.0, 0.0])] - fn small_norm_produces_unit_cosine_ranking(#[case] small_point: [f32; DIMENSION_COUNT]) { + fn cosine_writes_one_for_small_norms(#[case] small_point: [f32; DIMENSION_COUNT]) { // Given: a zero or subnormal norm represents zero similarity. let unit_point = [1.0_f32, 0.0]; let expected = 1.0; @@ -192,7 +251,7 @@ mod tests { } #[test] - fn normalized_cosine_ranking_equals_the_negative_dot_product() { + fn normalized_cosine_writes_negative_dot_products() { // Given let first_point = [1.0_f32, 0.0]; let second_point = [0.6_f32, 0.8]; @@ -207,7 +266,7 @@ mod tests { } #[test] - fn inner_product_ranking_equals_the_negative_dot_product() { + fn inner_product_writes_negative_dot_products() { // Given let first_point = [2.0_f32, -1.0]; let second_point = [3.0_f32, 4.0]; @@ -250,7 +309,7 @@ mod tests { } #[test] - fn inner_product_uses_the_last_coordinate_beyond_a_complete_dimension_block() { + fn inner_product_includes_last_coordinate() { let dimensions = 129; let mut values = vec![0.0; 2 * dimensions]; values[dimensions - 1] = 2.0; diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 8892c4b333..a9479c2418 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -40,6 +40,149 @@ fn cosine_distance(dot: f32, source_norm: f32, target_norm: f32) -> f32 { } } +#[cfg(test)] +#[cfg(not(miri))] +mod scalar_ranking { + use diskann_vector::distance::Metric; + + // Sample an irregular arc with increasing radius. Repeat its x/y coordinates + // through every dimension, so an odd final coordinate changes the geometry. + // Unequal angle steps avoid systematic nearest-neighbor ties. + pub(super) fn arc_vectors( + parameters: impl IntoIterator, + dimensions: usize, + normalize: bool, + ) -> Vec { + parameters + .into_iter() + .flat_map(|t| { + let angle = 0.03 * t + 0.001 * t * t; + let radius = 1.0 + t / 64.0; + let coordinates = [radius * angle.cos(), radius * angle.sin()]; + let mut vector: Vec<_> = (0..dimensions) + .map(|dimension| coordinates[dimension % 2] as f32) + .collect(); + if normalize { + let norm = vector + .iter() + .map(|&x| f64::from(x).powi(2)) + .sum::() + .sqrt(); + for x in &mut vector { + *x = (f64::from(*x) / norm) as f32; + } + } + vector + }) + .collect() + } + + // Use f64 scalar arithmetic on the actual f32 inputs, not GEMM or a kernel + // metric. These fixtures have finite coordinates and nonzero normal norms. + // L2 includes the point norm; partition ranking can omit this row constant. + pub(super) fn distance(metric: Metric, point: &[f32], target: &[f32]) -> f64 { + let dot: f64 = point + .iter() + .zip(target) + .map(|(&x, &y)| f64::from(x) * f64::from(y)) + .sum(); + match metric { + Metric::L2 => point + .iter() + .zip(target) + .map(|(&x, &y)| (f64::from(x) - f64::from(y)).powi(2)) + .sum(), + Metric::Cosine => { + let point_norm = point + .iter() + .map(|&x| f64::from(x).powi(2)) + .sum::() + .sqrt(); + let target_norm = target + .iter() + .map(|&x| f64::from(x).powi(2)) + .sum::() + .sqrt(); + 1.0 - (dot / (point_norm * target_norm)).clamp(-1.0, 1.0) + } + Metric::CosineNormalized | Metric::InnerProduct => -dot, + } + } + + // Keep rank comparisons tight enough to distinguish nearby candidates. + // L2 and inner-product scores grow with dimension. Cosine scores stay unit-scale. + pub(super) fn ranking_tolerance(metric: Metric, dimensions: usize) -> f64 { + let scale = match metric { + Metric::L2 | Metric::InnerProduct => dimensions as f64, + Metric::Cosine | Metric::CosineNormalized => 1.0, + }; + 32.0 * f64::from(f32::EPSILON) * scale + } + + // Bound distance error from the accumulated terms, not the final distance. + // L2 can subtract large dot products from large squared norms. + // This bound applies to the finite, nonzero-norm geometry above. + pub(super) fn distance_tolerance(metric: Metric, point: &[f32], target: &[f32]) -> f64 { + // gamma = n*u/(1-n*u) bounds rounding across n operations. + // Allow norm reduction, Gram accumulation, and the final arithmetic. + let unit_roundoff = f64::from(f32::EPSILON) / 2.0; + let steps = 2.0 * point.len() as f64 + 3.0; + let accumulated_roundoff = steps * unit_roundoff; + let gamma = accumulated_roundoff / (1.0 - accumulated_roundoff); + match metric { + Metric::L2 => { + let terms: f64 = point + .iter() + .zip(target) + .map(|(&x, &y)| (f64::from(x).abs() + f64::from(y).abs()).powi(2)) + .sum(); + gamma * terms + } + Metric::InnerProduct | Metric::CosineNormalized => { + let terms: f64 = point + .iter() + .zip(target) + .map(|(&x, &y)| (f64::from(x) * f64::from(y)).abs()) + .sum(); + gamma * terms + } + Metric::Cosine => { + // Both the dot product and the norm product can carry error. + // Division adds rounding. The final subtraction adds at most 2*u. + (2.0 * gamma + unit_roundoff * (1.0 + gamma)) / (1.0 - gamma) + 2.0 * unit_roundoff + } + } + } + + // Compare in output order. An isolated score admits only its expected ID. + // Scores within the rounding bound can exchange IDs, but cannot duplicate one. + #[track_caller] + pub(super) fn assert_ranked_ids(actual: &[u32], sorted: &[(u32, f64)], tolerance: f64) { + let mut previous = f64::NEG_INFINITY; + for (rank, &id) in actual.iter().enumerate() { + assert!( + !actual[..rank].contains(&id), + "duplicate ID {id} at rank {rank}" + ); + let score = sorted + .iter() + .find(|&&(candidate, _)| candidate == id) + .unwrap_or_else(|| panic!("unexpected ID {id} at rank {rank}")) + .1; + let (expected_id, expected_score) = sorted[rank]; + assert!( + (score - expected_score).abs() <= tolerance, + "rank {rank}: ID {id} score {score}, expected ID {expected_id} score {expected_score}, tolerance {tolerance}" + ); + assert!( + previous <= score + tolerance, + "rank {rank}: score {score} follows farther score {previous}" + ); + previous = score; + } + } +} + #[cfg(test)] mod cosine_distance_tests { use super::cosine_distance; @@ -60,7 +203,8 @@ mod cosine_distance_tests { } #[test] - fn minimum_normal_norm_uses_normalized_similarity() { + fn norm_at_cutoff_preserves_positive_similarity() { + // At the cutoff, similarity remains 0.5 instead of zero. let norm = f32::MIN_POSITIVE.sqrt(); assert_eq!(cosine_distance(0.5 * norm, norm, 1.0), 0.5); } @@ -68,10 +212,7 @@ mod cosine_distance_tests { #[rstest] #[case::above_one(1.0 + f32::EPSILON, 0.0)] #[case::below_negative_one(-1.0 - f32::EPSILON, 2.0)] - fn finite_similarity_outside_the_cosine_range_is_clamped( - #[case] similarity: f32, - #[case] expected: f32, - ) { + fn cosine_clamps_finite_similarity(#[case] similarity: f32, #[case] expected: f32) { assert_eq!(cosine_distance(similarity, 1.0, 1.0), expected); } @@ -79,7 +220,7 @@ mod cosine_distance_tests { #[case::nan_dot(f32::NAN, 1.0, 1.0)] #[case::nan_source_norm(0.0, f32::NAN, 1.0)] #[case::nan_target_norm(0.0, 1.0, f32::NAN)] - fn nan_without_a_zero_norm_produces_nan_distance( + fn nan_propagates_without_small_norms( #[case] dot: f32, #[case] source_norm: f32, #[case] target_norm: f32, diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 78bbeb88f4..f945229a8c 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -96,18 +96,21 @@ fn rank_leader_distances( }); } +// These entry-point tests call GEMM. Miri checks ranking separately. #[cfg(test)] +#[cfg(not(miri))] mod tests { use super::*; - use crate::graph::pipnn::{Cosine, L2}; + use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2, scalar_ranking}; use diskann_utils::views::{MatrixView, MutMatrixView}; + use diskann_vector::distance::Metric; use diskann_wide::arch::{self, Target3}; - struct AssignL2; + struct AssignLeaders(M); - impl + impl Target3, MatrixView<'_, f32>, MatrixView<'_, f32>, MutMatrixView<'_, u32>> - for AssignL2 + for AssignLeaders { fn run( self, @@ -116,10 +119,10 @@ mod tests { leader_values: MatrixView<'_, f32>, output: MutMatrixView<'_, u32>, ) -> ANNResult<()> { - assign_leaders::<_, L2>( + assign_leaders::<_, M>( arch, points, - &L2::create_leaders(leader_values), + &M::create_leaders(leader_values), output, &mut PartitionKernelWorkspace::default(), ) @@ -128,6 +131,62 @@ mod tests { mod assign_leaders_tests { use super::*; + use rstest::rstest; + + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn leader_ids_match_scalar_ranking( + #[case] metric: M, + #[case] scalar_metric: Metric, + #[values(2, 7, 8, 9, 15, 16, 17, 127, 128, 129, 384, 768)] dimensions: usize, + #[values(1, 3, 10, 11)] width: usize, + ) { + // Given: 33 leaders supply two SIMD groups and a scalar tail. + // Four points lie between different leaders on the same irregular arc. + let normalize = scalar_metric == Metric::CosineNormalized; + let leader_values = + scalar_ranking::arc_vectors((0..33).map(f64::from), dimensions, normalize); + let point_values = + scalar_ranking::arc_vectors([1.25, 10.5, 21.75, 31.25], dimensions, normalize); + let leaders = MatrixView::try_from(leader_values.as_slice(), 33, dimensions).unwrap(); + let points = MatrixView::try_from(point_values.as_slice(), 4, dimensions).unwrap(); + let expected: Vec> = points + .row_iter() + .map(|point| { + let mut row: Vec<_> = leaders + .row_iter() + .enumerate() + .map(|(leader, vector)| { + ( + leader as u32, + scalar_ranking::distance(scalar_metric, point, vector), + ) + }) + .collect(); + row.sort_by(|left, right| left.1.total_cmp(&right.1)); + row + }) + .collect(); + let tolerance = scalar_ranking::ranking_tolerance(scalar_metric, dimensions); + let mut output = vec![UNASSIGNED_LEADER; 4 * width]; + + // When: create metric-owned leaders and dispatch the complete assignment. + arch::dispatch3_no_features( + AssignLeaders(metric), + points, + leaders, + MutMatrixView::try_from(output.as_mut_slice(), 4, width).unwrap(), + ) + .unwrap(); + + // Then: the full L2 distance has the same order as the omitted-norm score. + for (row, expected) in output.chunks_exact(width).zip(&expected) { + scalar_ranking::assert_ranked_ids(row, expected, tolerance); + } + } #[rstest::rstest] #[case::missing_row(2)] @@ -151,7 +210,7 @@ mod tests { } #[test] - fn assignment_projects_each_points_l2_leader_columns_with_simd_and_tail() { + fn l2_points_select_nearest_leaders() { // Given: leader j lies at x=j. Seventeen leaders enter one SIMD block and a tail. // The nearest two integer positions to 15.75, -0.25, and 7.25 are respectively // (16, 15), (0, 1), and (7, 8); their squared distances are (0.0625, 0.5625), @@ -163,7 +222,7 @@ mod tests { // When: dispatch reaches the real metric and kernel on the available architecture. arch::dispatch3_no_features( - AssignL2, + AssignLeaders(L2), MatrixView::try_from(&point_values[..], 4, 1).unwrap(), MatrixView::try_from(&leader_values[..], 17, 1).unwrap(), MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), @@ -175,7 +234,7 @@ mod tests { } #[test] - fn assignment_reuses_workspace_after_shrinking_and_regrowing_a_cosine_stripe() { + fn cosine_assignments_update_after_workspace_reuse() { // Given: unit leaders point right, up, and left. The first point's cosine // similarities have order right > up > left; the second reverses right and left. let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 4d3282d1f1..96b51b082d 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -225,36 +225,11 @@ impl PartitionMetric for CosineNormalized { mod tests { use super::*; - const DIMENSION_COUNT: usize = 2; - const STALE_DISTANCE: f32 = 99.0; - const FLOAT_TOLERANCE: f32 = 1.0e-6; - - fn matrix(values: &[f32], rows: usize) -> MatrixView<'_, f32> { - MatrixView::try_from(values, rows, DIMENSION_COUNT).unwrap() - } - - fn compute_one_ranking( - point: [f32; DIMENSION_COUNT], - leader: [f32; DIMENSION_COUNT], - ) -> f32 { - let leaders = M::create_leaders(matrix(&leader, 1)); - let mut storage = [STALE_DISTANCE]; - - M::compute_distances( - matrix(&point, 1), - &leaders, - MutMatrixView::row_vector(&mut storage[..]), - ) - .unwrap(); - - storage[0] - } - mod create_leaders_tests { use super::*; #[test] - fn l2_leader_norms_preserve_the_sequential_reduction_order() { + fn l2_norms_follow_sequential_sum_order() { // Given: each small square is half an ULP at the first squared norm. // Each sequential addition rounds back to that norm. let mut values = [1.0_f32; 129]; @@ -270,8 +245,113 @@ mod tests { } } + // GEMM-backed tests run natively, not under Miri. + #[cfg(not(miri))] mod compute_distances_tests { use super::*; + use crate::graph::pipnn::scalar_ranking; + use diskann_vector::distance::Metric; + use rstest::rstest; + + const DIMENSION_COUNT: usize = 2; + const STALE_DISTANCE: f32 = 99.0; + const FLOAT_TOLERANCE: f32 = 1.0e-6; + + fn matrix(values: &[f32], rows: usize) -> MatrixView<'_, f32> { + MatrixView::try_from(values, rows, DIMENSION_COUNT).unwrap() + } + + fn compute_one_ranking( + point: [f32; DIMENSION_COUNT], + leader: [f32; DIMENSION_COUNT], + ) -> f32 { + let leaders = M::create_leaders(matrix(&leader, 1)); + let mut storage = [STALE_DISTANCE]; + + M::compute_distances( + matrix(&point, 1), + &leaders, + MutMatrixView::row_vector(&mut storage[..]), + ) + .unwrap(); + + storage[0] + } + + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn point_leader_matrix_matches_scalar_ranking( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[values(32, 128)] point_count: usize, + #[values(64, 256)] leader_count: usize, + #[values(128, 384, 768)] dimensions: usize, + ) { + // Given: points and leaders sample the same bounded arc on different grids. + // NaN marks unwritten entries in every point-to-leader row. + let normalize = scalar_metric == Metric::CosineNormalized; + let point_values = scalar_ranking::arc_vectors( + (0..point_count).map(|point| 32.0 * (point as f64 + 0.5) / point_count as f64), + dimensions, + normalize, + ); + let leader_values = scalar_ranking::arc_vectors( + (0..leader_count).map(|leader| 32.0 * leader as f64 / leader_count as f64), + dimensions, + normalize, + ); + let points = + MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); + let leader_vectors = + MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); + let expected: Vec> = points + .row_iter() + .map(|point| { + // Derive the omitted L2 term from input coordinates, not cached norms. + let point_term = if scalar_metric == Metric::L2 { + point.iter().map(|&value| f64::from(value).powi(2)).sum() + } else { + 0.0 + }; + leader_vectors + .row_iter() + .map(|leader| { + ( + scalar_ranking::distance(scalar_metric, point, leader) - point_term, + scalar_ranking::distance_tolerance(scalar_metric, point, leader), + ) + }) + .collect() + }) + .collect(); + let mut distances = vec![f32::NAN; point_count * leader_count]; + + // When: prepare leaders and call the metric without partition selection. + let leaders = M::create_leaders(leader_vectors); + M::compute_distances( + points, + &leaders, + MutMatrixView::try_from(distances.as_mut_slice(), point_count, leader_count) + .unwrap(), + ) + .unwrap(); + + // Then: check every leader column for every point, not only selected leaders. + let actual = + MatrixView::try_from(distances.as_slice(), point_count, leader_count).unwrap(); + for (point, row) in expected.iter().enumerate() { + for (leader, &(expected, tolerance)) in row.iter().enumerate() { + let actual = f64::from(actual.row(point)[leader]); + assert!( + (actual - expected).abs() <= tolerance, + "{scalar_metric:?}, {point_count} points, {leader_count} leaders, D={dimensions}, pair ({point}, {leader}): actual {actual}, expected {expected}, tolerance {tolerance}" + ); + } + } + } #[rstest::rstest] #[case::wrong_rows_with_equal_area(1, 4)] @@ -299,7 +379,7 @@ mod tests { } #[test] - fn squared_l2_ranking_omits_the_point_norm() { + fn l2_omits_squared_point_norm() { // Given: leader norms are 25 and 4; each entry is norm - 2 * dot. let leader_values = [3.0_f32, 4.0, 0.0, 2.0]; let leaders = L2::create_leaders(matrix(&leader_values, 2)); @@ -324,7 +404,7 @@ mod tests { #[case::zero_leader([1.0, 0.0], [0.0, 0.0])] #[case::subnormal_point([f32::MIN_POSITIVE.sqrt() / 2.0, 0.0], [1.0, 0.0])] #[case::subnormal_leader([1.0, 0.0], [f32::MIN_POSITIVE.sqrt() / 2.0, 0.0])] - fn small_norm_produces_unit_cosine_ranking( + fn cosine_writes_one_for_small_norms( #[case] point: [f32; DIMENSION_COUNT], #[case] leader: [f32; DIMENSION_COUNT], ) { @@ -357,7 +437,7 @@ mod tests { } #[test] - fn normalized_cosine_ranking_equals_the_negative_dot_product() { + fn normalized_cosine_writes_negative_dot_products() { // Given let point = [1.0_f32, 0.0]; let leader = [0.6_f32, 0.8]; @@ -372,7 +452,7 @@ mod tests { } #[test] - fn inner_product_ranking_equals_the_negative_dot_product() { + fn inner_product_writes_negative_dot_products() { // Given let point = [2.0_f32, -1.0]; let leader = [3.0_f32, 4.0]; @@ -387,7 +467,7 @@ mod tests { } #[test] - fn cosine_overwrites_every_point_row_with_normalized_distances() { + fn cosine_replaces_all_distance_rows() { // Given: non-unit points point right, up, and left. Leaders point // diagonally up-right and up, so similarities include ±1/sqrt(2). let leader_values = [1.0_f32, 1.0, 0.0, 2.0]; diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 00637d1a9b..c73fadead5 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -498,7 +498,7 @@ mod tests { } #[test] - fn initialize_clears_candidates_and_resizes_thresholds() { + fn initialize_resets_results_and_limits() { // Given: candidate storage contains old results; the threshold buffer // still has the smaller previous leaf's length. let mut output = [ @@ -528,7 +528,7 @@ mod tests { #[rstest] #[case::fixed(Fixed::<0>)] #[case::runtime(Runtime(0))] - fn zero_capacity_accepts_both_update_operations(#[case] width: impl Width) { + fn zero_capacity_preserves_results_and_limits(#[case] width: impl Width) { let mut output = []; let mut thresholds = [f32::INFINITY; 2]; let topk = TopK::new(width); @@ -555,7 +555,7 @@ mod tests { #[rstest] #[case::fixed(Fixed::<3>)] #[case::runtime(Runtime(3))] - fn insertion_preserves_later_result_rows(#[case] width: impl Width) { + fn insertion_preserves_other_rows(#[case] width: impl Width) { // Given: the new candidate displaces the first result's farthest neighbor. let first_result = [ Candidate::new(0, 1.0), @@ -589,35 +589,73 @@ mod tests { use super::*; #[rstest] - #[case::lane_minus_one(15, 3)] - #[case::complete_lane(16, 2)] - #[case::lane_and_tail(17, 10)] - #[case::two_lanes_minus_one(31, 1)] - #[case::two_complete_lanes(32, 4)] - #[case::two_lanes_and_tail(33, 17)] - #[case::capacity_five(33, 5)] - #[case::capacity_six(33, 6)] - #[case::capacity_seven(33, 7)] - #[case::capacity_eight(33, 8)] - #[case::capacity_nine(33, 9)] + #[case::before_first_simd_block(15)] + #[case::one_simd_block(16)] + #[case::one_simd_block_and_tail(17)] + #[case::before_second_simd_block(31)] + #[case::two_simd_blocks(32)] + #[case::two_simd_blocks_and_tail(33)] + #[case::three_simd_blocks_and_tail(49)] fn dispatched_selection_matches_independent_sort( #[case] count: usize, - #[case] width: usize, + #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)] width: usize, ) { - // Alternating signs force insertions at different positions, not just append. + // Given: alternating signs force replacement and interior insertion. let distances: Vec<_> = (0..count) .map(|i| if i % 2 == 0 { i as f32 } else { -(i as f32) }) .collect(); let expected = sorted_candidates(&distances, width); let mut output = vec![Candidate::default(); width]; + let mut runtime_output = output.clone(); + // When with_topk!(width, |topk| { arch::dispatch1_no_features( SelectTopK, (&topk, distances.as_slice(), output.as_mut_slice()), ); - assert_eq!(output, expected); }); + arch::dispatch1_no_features( + SelectTopK, + ( + &TopK::new(Runtime(width)), + distances.as_slice(), + runtime_output.as_mut_slice(), + ), + ); + + // Then: both width representations must satisfy the independent oracle. + assert_eq!(output, expected); + assert_eq!(runtime_output, expected); + } + + #[rstest] + #[case::before_second_simd_block(31)] + #[case::two_simd_blocks(32)] + #[case::two_simd_blocks_and_tail(33)] + #[case::three_simd_blocks_and_tail(49)] + fn runtime_selection_matches_independent_sort( + #[case] count: usize, + #[values(11, 17)] width: usize, + ) { + // Given: K=17 fills during the second block. The third block starts + // with a finite limit. Each input has more candidates than result slots. + let distances: Vec<_> = (0..count) + .map(|i| if i % 2 == 0 { i as f32 } else { -(i as f32) }) + .collect(); + let expected = sorted_candidates(&distances, width); + let mut output = vec![Candidate::default(); width]; + + // When: these capacities select the runtime branch of the batch dispatch. + with_topk!(width, |topk| { + arch::dispatch1_no_features( + SelectTopK, + (&topk, distances.as_slice(), output.as_mut_slice()), + ); + }); + + // Then + assert_eq!(output, expected); } #[rstest] @@ -626,21 +664,23 @@ mod tests { &[f32::NAN, f32::INFINITY, 9.0], [Candidate::new(2, 9.0), Candidate::default()] )] - fn each_point_replaces_the_previous_points_candidates( + fn selection_replaces_previous_results( #[case] distances: &[f32], #[case] expected: [Candidate; 2], ) { - let mut output = [Candidate::default(); 2]; + // Given: both slots contain candidates from the previous selection. + let mut output = [Candidate::new(1, 1.0), Candidate::new(2, 2.0)]; let topk = TopK::new(Fixed::<2>); - topk.select_topk(diskann_wide::ARCH, &[3.0, 1.0, 2.0], &mut output); + // When: select again with the same output buffer. arch::dispatch1_no_features(SelectTopK, (&topk, distances, &mut output[..])); + // Then: unused slots contain no candidates from the previous selection. assert_eq!(output, expected); } #[test] - fn finite_simd_mask_rechecks_the_threshold_after_insertion() { + fn selection_rechecks_candidates_against_updated_limit() { // Group one fills K=1 at 5. Both 1 and 4 pass group two's initial // mask, but inserting 1 must prevent the later 4 from replacing it. let mut distances = [f32::INFINITY; 32]; @@ -656,7 +696,7 @@ mod tests { } #[test] - fn closer_candidates_shift_the_middle_of_a_full_result() { + fn middle_insertions_keep_nearest_first_order() { let mut output = [Candidate::default(); 3]; let topk = TopK::new(Fixed::<3>); @@ -678,7 +718,9 @@ mod tests { } #[test] - fn non_rankable_first_block_preserves_later_candidate_ids() { + fn non_rankable_block_preserves_input_indices() { + // The first block contributes no candidates. Later IDs still refer + // to positions in the full input, not positions among rankable values. let mut distances = [f32::INFINITY; 33]; distances[..8].fill(f32::NAN); distances[16] = 4.0; @@ -701,7 +743,7 @@ mod tests { } #[test] - fn nearer_tail_displaces_a_tied_candidate_without_duplicate_ids() { + fn tail_candidate_replaces_a_tied_neighbor() { let mut distances = [2.0; 17]; distances[16] = 1.0; let mut output = [Candidate::default(); 3]; @@ -719,7 +761,7 @@ mod tests { } #[test] - fn special_rankable_distances_keep_distinct_ids_and_original_bits() { + fn selection_preserves_ids_and_float_bits() { let mut distances = [f32::INFINITY; 33]; distances[0] = -0.0; distances[16] = 0.0; @@ -752,61 +794,173 @@ mod tests { mod update_dual_topk_tests { use super::*; - #[test] - fn endpoints_accept_independently_in_both_simd_groups_and_tail() { - let mut distances = [f32::INFINITY; 33]; - distances[0] = 0.5; - distances[1] = 4.0; - distances[2] = 6.0; - distances[16] = 2.0; - distances[32] = 0.75; - let mut output = [Candidate::default(); 34]; - let mut thresholds = [f32::INFINITY; 34]; - let topk = TopK::new(Fixed::<1>); - let mut rows = MutMatrixView::try_from(&mut output[..], 34, 1).unwrap(); - // Earlier leaf rows establish each destination's threshold. - let mut previous = [f32::INFINITY; 32]; - for point in 1..33 { - if point == 3 { - previous[..3].copy_from_slice(&[0.2, 5.0, 5.0]); + // Assign one unique score to each unordered pair. Shuffling prevents the + // leaf scan order from also being distance order. The diagonal cannot rank. + fn shuffled_symmetric_distances(count: usize) -> Vec { + use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; + + let mut scores: Vec<_> = (1..=count * (count - 1) / 2).map(|i| i as f32).collect(); + scores.shuffle(&mut StdRng::seed_from_u64(1287)); + let mut distances = vec![f32::INFINITY; count * count]; + for source in 1..count { + for target in 0..source { + let score = scores.pop().unwrap(); + distances[source * count + target] = score; + distances[target * count + source] = score; } - arch::dispatch1_no_features( - UpdatePair, - ( - &topk, - point, - &previous[..point], - rows.as_mut_view(), - &mut thresholds[..], - ), + } + distances + } + + #[rstest] + #[case::one_simd_block_and_tail(18)] + #[case::two_simd_blocks(33)] + #[case::two_simd_blocks_and_tail(34)] + #[case::three_simd_blocks_and_tail(50)] + fn each_point_retains_its_nearest_neighbors( + #[case] count: usize, + #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17)] width: usize, + ) { + // Given: the oracle sorts complete rows, independent of the pair scan. + // Every row has enough non-self candidates to fill its result. + let distances = shuffled_symmetric_distances(count); + let expected: Vec<_> = distances + .chunks_exact(count) + .map(|row| sorted_candidates(row, width)) + .collect(); + let mut output = vec![Candidate::default(); count * width]; + let mut thresholds = Vec::new(); + let mut rows = MutMatrixView::try_from(output.as_mut_slice(), count, width).unwrap(); + + // When: each lower-triangle pair reaches both endpoints exactly once. + with_topk!(width, |topk| { + topk.initialize(rows.as_mut_view(), &mut thresholds); + for source in 1..count { + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + source, + &distances[source * count..source * count + source], + rows.as_mut_view(), + thresholds.as_mut_slice(), + ), + ); + } + }); + + // Then: check all source and reciprocal rows with their persisted limits. + for (source, expected) in expected.iter().enumerate() { + assert_eq!( + rows.row(source), + expected, + "source={source}, N={count}, K={width}" + ); + assert_eq!( + thresholds[source], + expected[width - 1].distance, + "threshold for source={source}, N={count}, K={width}" ); - previous.fill(f32::INFINITY); } - let mut expected = [Candidate::default(); 34]; - expected[0] = Candidate::new(3, 0.2); - expected[1] = Candidate::new(33, 4.0); - expected[2] = Candidate::new(3, 5.0); - expected[3] = Candidate::new(0, 0.2); - expected[16] = Candidate::new(33, 2.0); - expected[32] = Candidate::new(33, 0.75); - expected[33] = Candidate::new(0, 0.5); + } + #[test] + fn only_the_source_accepts_the_pair() { + // Given: point 2 is the source. Point 0 is the target. + // Points 0 and 1 already retain each other at distance 1. + // The source has no neighbors. Distance 2 improves only its result. + let mut output = [ + Candidate::new(1, 1.0), + Candidate::new(0, 1.0), + Candidate::default(), + ]; + let mut thresholds = output.map(|candidate| candidate.distance); + let expected = [output[0], output[1], Candidate::new(0, 2.0)]; + + // When arch::dispatch1_no_features( UpdatePair, ( - &topk, - 33, - &distances[..], - rows.as_mut_view(), + &TopK::new(Fixed::<1>), + 2, + &[2.0, 3.0][..], + MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), + &mut thresholds[..], + ), + ); + + // Then + assert_eq!(output, expected); + assert_eq!(thresholds, expected.map(|candidate| candidate.distance)); + } + + #[test] + fn only_the_target_accepts_the_pair() { + // Given: point 2 is the source. Point 1 is the target. + // Points 0 and 1 retain each other at distance 3. + // The source first selects point 0 at distance 1. + // Its later distance 2 to the target improves only the target's result. + let mut output = [ + Candidate::new(1, 3.0), + Candidate::new(0, 3.0), + Candidate::default(), + ]; + let mut thresholds = output.map(|candidate| candidate.distance); + let expected = [ + Candidate::new(2, 1.0), + Candidate::new(2, 2.0), + Candidate::new(0, 1.0), + ]; + + // When + arch::dispatch1_no_features( + UpdatePair, + ( + &TopK::new(Fixed::<1>), + 2, + &[1.0, 2.0][..], + MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), &mut thresholds[..], ), ); - assert_eq!(rows.as_slice(), expected); + // Then + assert_eq!(output, expected); + assert_eq!(thresholds, expected.map(|candidate| candidate.distance)); } #[test] - fn a_source_fills_across_blocks_with_few_rankable_distances() { + fn both_points_reject_the_farther_pair() { + // Given: points 0 and 1 retain each other at distance 1. + // Point 2 first selects point 0 at distance 0.5. The later pair + // (2, 1) at distance 2 improves neither endpoint. + let mut output = [ + Candidate::new(1, 1.0), + Candidate::new(0, 1.0), + Candidate::default(), + ]; + let mut thresholds = output.map(|candidate| candidate.distance); + let expected = [Candidate::new(2, 0.5), output[1], Candidate::new(0, 0.5)]; + + // When + arch::dispatch1_no_features( + UpdatePair, + ( + &TopK::new(Fixed::<1>), + 2, + &[0.5, 2.0][..], + MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), + &mut thresholds[..], + ), + ); + + // Then + assert_eq!(output, expected); + assert_eq!(thresholds, expected.map(|candidate| candidate.distance)); + } + + #[test] + fn later_blocks_fill_source_neighbor_slots() { let mut output = [Candidate::default(); 34 * 3]; let mut thresholds = [f32::INFINITY; 34]; let topk = TopK::new(Fixed::<3>); @@ -851,7 +1005,7 @@ mod tests { } #[test] - fn later_pairs_use_the_sources_updated_threshold() { + fn later_pairs_respect_saved_distance_limits() { let mut output = [Candidate::default(); 3]; let mut thresholds = [f32::INFINITY; 3]; let topk = TopK::new(Fixed::<1>); @@ -889,34 +1043,52 @@ mod tests { } #[test] - fn runtime_capacity_updates_rows_filled_by_earlier_leaf_pairs() { + fn target_replaces_its_farthest_neighbor() { + // Given: point 15 retains points 0, 1, 2, and 3. It is the last + // lane in group one. Point 16 is the tail target and has no neighbors. + // Point 17 enters the middle of point 15's full row. let mut output = [Candidate::default(); 18 * 4]; let mut thresholds = [f32::INFINITY; 18]; - let mut previous = [f32::INFINITY; 17]; + let mut previous = [f32::INFINITY; 15]; + previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); let mut distances = [f32::INFINITY; 17]; distances[0] = 4.0; distances[2] = 3.0; distances[15] = 2.0; distances[16] = 1.0; + let expected_source = [ + Candidate::new(16, distances[16]), + Candidate::new(15, distances[15]), + Candidate::new(2, distances[2]), + Candidate::new(0, distances[0]), + ]; + let expected_prior_row = [ + Candidate::new(0, previous[0]), + Candidate::new(17, distances[15]), + Candidate::new(1, previous[1]), + Candidate::new(2, previous[2]), + ]; + let expected_tail_row = [ + Candidate::new(17, distances[16]), + Candidate::default(), + Candidate::default(), + Candidate::default(), + ]; let mut rows = MutMatrixView::try_from(&mut output[..], 18, 4).unwrap(); let topk = TopK::new(Runtime(4)); - for point in 1..17 { - if point == 15 { - previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); - } - arch::dispatch1_no_features( - UpdatePair, - ( - &topk, - point, - &previous[..point], - rows.as_mut_view(), - &mut thresholds[..], - ), - ); - previous.fill(f32::INFINITY); - } + arch::dispatch1_no_features( + UpdatePair, + ( + &topk, + 15, + &previous[..], + rows.as_mut_view(), + &mut thresholds[..], + ), + ); + + // When arch::dispatch1_no_features( UpdatePair, ( @@ -928,39 +1100,19 @@ mod tests { ), ); - assert_eq!( - rows.row(17), - [ - Candidate::new(16, 1.0), - Candidate::new(15, 2.0), - Candidate::new(2, 3.0), - Candidate::new(0, 4.0) - ] - ); - assert_eq!( - rows.row(15), - [ - Candidate::new(0, -1.0), - Candidate::new(17, 2.0), - Candidate::new(1, 3.0), - Candidate::new(2, 5.0) - ] - ); - assert_eq!( - rows.row(16), - [ - Candidate::new(17, 1.0), - Candidate::default(), - Candidate::default(), - Candidate::default() - ] - ); + // Then: the source fills, the prior row evicts, and the tail stays underfilled. + assert_eq!(rows.row(17), expected_source); + assert_eq!(rows.row(15), expected_prior_row); + assert_eq!(rows.row(16), expected_tail_row); + assert_eq!(thresholds[17], expected_source[3].distance); + assert_eq!(thresholds[15], expected_prior_row[3].distance); + assert_eq!(thresholds[16], f32::INFINITY); } #[rstest] #[case::simd(15)] #[case::tail(16)] - fn negative_infinity_reaches_both_endpoints(#[case] candidate: usize) { + fn both_points_accept_negative_infinity(#[case] candidate: usize) { let mut distances = [f32::NAN; 17]; distances[candidate] = f32::NEG_INFINITY; let mut output = [Candidate::default(); 18 * 2]; From d0f048dd4fbdee110959cca6600d2e29d8469640 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:07:44 +0000 Subject: [PATCH 182/229] refactor(pipnn): simplify workspace access Name the dual-update test and dispatcher after their entry point. Record pending Neon width validation. --- diskann/src/graph/pipnn/leaf_kernel.rs | 12 +++------ diskann/src/graph/pipnn/partition_kernel.rs | 17 ++++++------ diskann/src/graph/pipnn/simd.rs | 3 +++ diskann/src/graph/pipnn/topk.rs | 29 +++++++++++---------- 4 files changed, 31 insertions(+), 30 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index b667d17628..7288b3ea0c 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -72,20 +72,16 @@ where let point_count = points.nrows(); validate_output(point_count, &output).map_err(ANNError::new)?; let distance_count = point_count * point_count; - let LeafKernelWorkspace { - distance_scratch, - worst, - } = workspace; - if distance_scratch.len() < distance_count { - distance_scratch.resize(distance_count, 0.0); + if workspace.distance_scratch.len() < distance_count { + workspace.distance_scratch.resize(distance_count, 0.0); } let mut distances = MutMatrixView::try_from( - &mut distance_scratch[..distance_count], + &mut workspace.distance_scratch[..distance_count], point_count, point_count, )?; M::compute_distances(points, distances.as_mut_slice())?; - rank_leaf_distances(arch, distances.as_view(), output, worst); + rank_leaf_distances(arch, distances.as_view(), output, &mut workspace.worst); Ok(()) } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index f945229a8c..00507edcca 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -61,20 +61,21 @@ where } let leader_count = M::leader_count(leaders); let distance_count = point_count * leader_count; - let PartitionKernelWorkspace { - distance_scratch, - ranked_leaders, - } = workspace; - if distance_scratch.len() < distance_count { - distance_scratch.resize(distance_count, 0.0); + if workspace.distance_scratch.len() < distance_count { + workspace.distance_scratch.resize(distance_count, 0.0); } let mut distances = MutMatrixView::try_from( - &mut distance_scratch[..distance_count], + &mut workspace.distance_scratch[..distance_count], point_count, leader_count, )?; M::compute_distances(points, leaders, distances.as_mut_view())?; - rank_leader_distances(arch, distances.as_view(), output, ranked_leaders); + rank_leader_distances( + arch, + distances.as_view(), + output, + &mut workspace.ranked_leaders, + ); Ok(()) } diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 0cc6b9678f..687aa19ad7 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -10,6 +10,9 @@ use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDVector}; /// Default SIMD representation used by both PiPNN ranking kernels. /// /// This alias is the single build-time width selection. +/// +/// TODO: This 16-lane choice lacks Neon-specific validation. +/// Add Neon coverage in a follow-up. type DefaultVector = ::f32x16; /// PiPNN SIMD representation for one architecture. diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index c73fadead5..0dad92b2cf 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -455,7 +455,7 @@ mod tests { } } - struct UpdatePair; + struct UpdateDualTopK; impl Target1< A, @@ -467,7 +467,7 @@ mod tests { MutMatrixView<'_, Candidate>, &mut [f32], ), - > for UpdatePair + > for UpdateDualTopK { fn run( self, @@ -535,7 +535,7 @@ mod tests { arch::dispatch1_no_features(SelectTopK, (&topk, &[1.0, 2.0][..], &mut output[..])); arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &topk, 1, @@ -812,12 +812,13 @@ mod tests { distances } + // Happy path: every point retains K nearest neighbors from finite pair distances. #[rstest] #[case::one_simd_block_and_tail(18)] #[case::two_simd_blocks(33)] #[case::two_simd_blocks_and_tail(34)] #[case::three_simd_blocks_and_tail(50)] - fn each_point_retains_its_nearest_neighbors( + fn update_dual_topk_selects_nearest_neighbors( #[case] count: usize, #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17)] width: usize, ) { @@ -837,7 +838,7 @@ mod tests { topk.initialize(rows.as_mut_view(), &mut thresholds); for source in 1..count { arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &topk, source, @@ -879,7 +880,7 @@ mod tests { // When arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &TopK::new(Fixed::<1>), 2, @@ -914,7 +915,7 @@ mod tests { // When arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &TopK::new(Fixed::<1>), 2, @@ -944,7 +945,7 @@ mod tests { // When arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &TopK::new(Fixed::<1>), 2, @@ -976,7 +977,7 @@ mod tests { // A leaf source starts empty. It remains partially filled after // group one and acquires its finite threshold within group two. arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &topk, 33, @@ -1012,7 +1013,7 @@ mod tests { let mut rows = MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(); arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &topk, 1, @@ -1022,7 +1023,7 @@ mod tests { ), ); arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &topk, 2, @@ -1078,7 +1079,7 @@ mod tests { let mut rows = MutMatrixView::try_from(&mut output[..], 18, 4).unwrap(); let topk = TopK::new(Runtime(4)); arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &topk, 15, @@ -1090,7 +1091,7 @@ mod tests { // When arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &topk, 17, @@ -1124,7 +1125,7 @@ mod tests { expected[candidate][0] = Candidate::new(17, f32::NEG_INFINITY); arch::dispatch1_no_features( - UpdatePair, + UpdateDualTopK, ( &topk, 17, From 4e75562f96fcc01bb1d78b2cb2e091153362dc33 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:21:49 +0000 Subject: [PATCH 183/229] test(pipnn): clarify neighbor selection cases --- diskann/src/graph/pipnn/topk.rs | 37 ++++++++++++++++----------------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 0dad92b2cf..7a1eba9a29 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -794,8 +794,9 @@ mod tests { mod update_dual_topk_tests { use super::*; - // Assign one unique score to each unordered pair. Shuffling prevents the - // leaf scan order from also being distance order. The diagonal cannot rank. + // The lower triangle contains unique scores. The upper triangle mirrors them. + // Shuffle the scores so matrix order differs from distance order. + // Diagonal values cannot rank, so no point selects itself. fn shuffled_symmetric_distances(count: usize) -> Vec { use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; @@ -812,7 +813,7 @@ mod tests { distances } - // Happy path: every point retains K nearest neighbors from finite pair distances. + // Happy path: every point retains K nearest neighbors from finite distances. #[rstest] #[case::one_simd_block_and_tail(18)] #[case::two_simd_blocks(33)] @@ -822,7 +823,7 @@ mod tests { #[case] count: usize, #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17)] width: usize, ) { - // Given: the oracle sorts complete rows, independent of the pair scan. + // Given: the oracle sorts complete rows, independent of the update order. // Every row has enough non-self candidates to fill its result. let distances = shuffled_symmetric_distances(count); let expected: Vec<_> = distances @@ -833,7 +834,7 @@ mod tests { let mut thresholds = Vec::new(); let mut rows = MutMatrixView::try_from(output.as_mut_slice(), count, width).unwrap(); - // When: each lower-triangle pair reaches both endpoints exactly once. + // When: each distance updates the two points' neighbor lists once. with_topk!(width, |topk| { topk.initialize(rows.as_mut_view(), &mut thresholds); for source in 1..count { @@ -866,10 +867,9 @@ mod tests { } #[test] - fn only_the_source_accepts_the_pair() { - // Given: point 2 is the source. Point 0 is the target. - // Points 0 and 1 already retain each other at distance 1. - // The source has no neighbors. Distance 2 improves only its result. + fn farther_points_do_not_replace_closer_neighbors() { + // Given: points 0 and 1 select each other at distance 1. + // Point 2 has no neighbors. Its distances to points 0 and 1 are 2 and 3. let mut output = [ Candidate::new(1, 1.0), Candidate::new(0, 1.0), @@ -896,11 +896,10 @@ mod tests { } #[test] - fn only_the_target_accepts_the_pair() { - // Given: point 2 is the source. Point 1 is the target. - // Points 0 and 1 retain each other at distance 3. - // The source first selects point 0 at distance 1. - // Its later distance 2 to the target improves only the target's result. + fn closer_points_replace_existing_neighbors() { + // Given: points 0 and 1 select each other at distance 3. + // Point 2 is closer to both: distance 1 from point 0 and 2 from point 1. + // Points 0 and 1 must select point 2. Point 2 must select point 0. let mut output = [ Candidate::new(1, 3.0), Candidate::new(0, 3.0), @@ -931,10 +930,10 @@ mod tests { } #[test] - fn both_points_reject_the_farther_pair() { - // Given: points 0 and 1 retain each other at distance 1. - // Point 2 first selects point 0 at distance 0.5. The later pair - // (2, 1) at distance 2 improves neither endpoint. + fn both_points_prefer_the_closer_neighbor() { + // Given: points 0 and 1 select each other at distance 1. + // Point 2 is 0.5 from point 0 and 2 from point 1. + // Points 1 and 2 must each select point 0, not each other. let mut output = [ Candidate::new(1, 1.0), Candidate::new(0, 1.0), @@ -1006,7 +1005,7 @@ mod tests { } #[test] - fn later_pairs_respect_saved_distance_limits() { + fn later_updates_respect_saved_distance_limits() { let mut output = [Candidate::default(); 3]; let mut thresholds = [f32::INFINITY; 3]; let topk = TopK::new(Fixed::<1>); From d3ae858326676d73bd721710fc79dc4f88ce1921 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:02:59 +0000 Subject: [PATCH 184/229] test(pipnn): rewrite numerical kernel contracts --- Cargo.lock | 1 + diskann-linalg/Cargo.toml | 1 + diskann-linalg/src/lib.rs | 281 ++++-- diskann-utils/src/views.rs | 49 +- diskann/src/graph/pipnn/leaf_kernel.rs | 586 ++++++------ diskann/src/graph/pipnn/leaf_metric.rs | 369 ++++---- diskann/src/graph/pipnn/mod.rs | 310 +++--- diskann/src/graph/pipnn/partition_kernel.rs | 447 +++++---- diskann/src/graph/pipnn/partition_metric.rs | 492 +++++----- diskann/src/graph/pipnn/topk.rs | 984 +++++++------------- 10 files changed, 1739 insertions(+), 1781 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 85a46b7e09..5cc9e8288b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -677,6 +677,7 @@ dependencies = [ "faer", "rand", "rand_distr", + "rstest", "serde", "serde_json", "thiserror 2.0.17", diff --git a/diskann-linalg/Cargo.toml b/diskann-linalg/Cargo.toml index d60cbf3031..c798185b03 100644 --- a/diskann-linalg/Cargo.toml +++ b/diskann-linalg/Cargo.toml @@ -19,6 +19,7 @@ clippy.undocumented_unsafe_blocks = "warn" clippy.uninlined_format_args = "allow" [dev-dependencies] +rstest.workspace = true serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } approx = { workspace = true } diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index 2c88e5a359..8f98ef5b99 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -718,121 +718,214 @@ mod tests { } } } + #[cfg(test)] -#[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] mod sgemm_aat_lower_tests { - use super::{sgemm_aat_lower, sgemm_aat_lower_add, MatrixName, SgemmError}; + use super::*; + use rstest::rstest; - #[test] - fn replace_writes_the_scaled_lower_triangle_and_preserves_the_upper_triangle() { - // Given: rows [1, 2, 3] and [4, 5, 6] have Gram entries 14, 32, and 77. - let matrix = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; - let upper_triangle_sentinel = -123.0_f32; - let expected = [-14.0, upper_triangle_sentinel, -32.0, -77.0]; - let mut actual = [upper_triangle_sentinel; 4]; - - // When - sgemm_aat_lower(2, 3, -1.0, &matrix, &mut actual).unwrap(); - - // Then - assert_eq!(actual, expected); + type LowerProduct = fn(usize, usize, f32, &[f32], &mut [f32]) -> Result<(), SgemmError>; + + #[rstest] + #[case::replace(sgemm_aat_lower, [-8.0, 101.0, 102.0, 8.0, -26.0, 103.0, 12.0, -36.0, -50.0])] + #[case::add(sgemm_aat_lower_add, [2.0, 101.0, 102.0, 28.0, 4.0, 103.0, 52.0, 14.0, 10.0])] + fn scaled_products_change_only_the_diagonal_and_lower_triangle( + #[case] operation: LowerProduct, + #[case] expected: [f32; 9], + ) { + // Rows [2,0], [-2,3], [-3,4] have lower Gram entries 4,-4,13,-6,18,25. + let input = [2.0, 0.0, -2.0, 3.0, -3.0, 4.0]; + let mut output = [10.0, 101.0, 102.0, 20.0, 30.0, 103.0, 40.0, 50.0, 60.0]; + + operation(3, 2, -2.0, &input, &mut output).unwrap(); + + assert_eq!(output, expected); } - #[test] - fn add_accumulates_the_scaled_lower_triangle_and_preserves_the_upper_triangle() { - // Given: adding -2 times the Gram entries (14, 32, 77) to (10, 20, 30) - // produces (-18, -44, -124), while the upper triangle remains unchanged. - let matrix = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; - let upper_triangle_sentinel = -123.0_f32; - let expected = [-18.0, upper_triangle_sentinel, -44.0, -124.0]; - let mut actual = [10.0, upper_triangle_sentinel, 20.0, 30.0]; - - // When - sgemm_aat_lower_add(2, 3, -2.0, &matrix, &mut actual).unwrap(); - - // Then - assert_eq!(actual, expected); + #[rstest] + #[case::replace(sgemm_aat_lower, false)] + #[case::add(sgemm_aat_lower_add, true)] + fn rectangular_inputs_match_scalar_dot_products( + #[case] operation: LowerProduct, + #[case] add: bool, + #[values(1, 3, 17)] rows: usize, + #[values(0, 1, 7, 8, 9, 16, 17, 33)] dimensions: usize, + ) { + let input: Vec<_> = (0..rows * dimensions) + .map(|i| (i % 9) as f32 - 4.0) + .collect(); + let previous: Vec<_> = (0..rows * rows).map(|i| 100.0 + i as f32).collect(); + let mut output = previous.clone(); + + operation(rows, dimensions, -2.0, &input, &mut output).unwrap(); + + for row in 0..rows { + for column in 0..rows { + let index = row * rows + column; + if column > row { + assert_eq!( + output[index].to_bits(), + previous[index].to_bits(), + "upper ({row},{column})" + ); + } else { + let dot: f64 = (0..dimensions) + .map(|d| { + f64::from(input[row * dimensions + d]) + * f64::from(input[column * dimensions + d]) + }) + .sum(); + let expected = -2.0 * dot + if add { f64::from(previous[index]) } else { 0.0 }; + assert_eq!( + f64::from(output[index]), + expected, + "lower ({row},{column}), dimensions={dimensions}" + ); + } + } + } + } + + #[rstest] + #[case::replace(sgemm_aat_lower, false)] + #[case::add(sgemm_aat_lower_add, true)] + fn large_dense_inputs_match_scalar_dot_products( + #[case] operation: LowerProduct, + #[case] add: bool, + #[values((33, 384), (65, 768), (129, 1536), (17, 1537), (513, 129), (17, 4097))] shape: ( + usize, + usize, + ), + ) { + use rand::{rngs::StdRng, SeedableRng}; + + let (rows, dimensions) = shape; + let mut rng = StdRng::seed_from_u64(1287); + // Bounded multiples of 1/8 keep products and sums exact in f32. + let mut input: Vec<_> = (0..rows * dimensions) + .map(|_| rng.random_range(-16..=16) as f32 / 8.0) + .collect(); + for (row, values) in input.chunks_exact_mut(dimensions).enumerate() { + values[dimensions - 1] = 8.0 + (row % 7) as f32; + } + let previous: Vec<_> = (0..rows * rows).map(|i| (i % 97) as f32 - 48.0).collect(); + let mut output = previous.clone(); + + operation(rows, dimensions, -2.0, &input, &mut output).unwrap(); + + // Check every entry, including the untouched upper triangle. + for row in 0..rows { + for column in 0..rows { + let index = row * rows + column; + if column > row { + assert_eq!( + output[index].to_bits(), + previous[index].to_bits(), + "shape={shape:?}, upper ({row},{column})" + ); + } else { + let dot: f64 = (0..dimensions) + .map(|dimension| { + f64::from(input[row * dimensions + dimension]) + * f64::from(input[column * dimensions + dimension]) + }) + .sum(); + let expected = -2.0 * dot + if add { f64::from(previous[index]) } else { 0.0 }; + assert_eq!( + f64::from(output[index]), + expected, + "shape={shape:?}, lower ({row},{column})" + ); + } + } + } } #[test] - fn invalid_input_length_returns_the_a_matrix_error() { - // Given - let row_count = 2; - let dimension_count = 2; - let invalid_input = [0.0_f32; 3]; - let mut output = [0.0_f32; 4]; - let expected = SgemmError::InvalidMatrixDimensions { - matrix_name: MatrixName::A, - expected_rows: row_count, - expected_cols: dimension_count, - actual_len: invalid_input.len(), - }; + fn replacement_does_not_read_old_lower_triangle_values() { + let input = [1.0, 2.0, -3.0, 4.0]; + let mut output = [f32::NAN, 17.0, f32::NAN, f32::NAN]; - // When - let actual = sgemm_aat_lower(row_count, dimension_count, 1.0, &invalid_input, &mut output) - .unwrap_err(); + sgemm_aat_lower(2, 2, 1.0, &input, &mut output).unwrap(); - // Then - assert_eq!(actual, expected); + assert_eq!(output, [5.0, 17.0, 5.0, 25.0]); } - #[test] - fn invalid_output_length_returns_the_c_matrix_error() { - // Given - let row_count = 2; - let dimension_count = 2; - let input = [0.0_f32; 4]; - let mut invalid_output = [0.0_f32; 3]; - let expected = SgemmError::InvalidMatrixDimensions { - matrix_name: MatrixName::C, - expected_rows: row_count, - expected_cols: row_count, - actual_len: invalid_output.len(), - }; + #[rstest] + #[case::replace(sgemm_aat_lower, [0.0, 17.0, 0.0, 0.0])] + #[case::add(sgemm_aat_lower_add, [3.0, 17.0, 5.0, 7.0])] + fn zero_scale_respects_the_accumulation_mode( + #[case] operation: LowerProduct, + #[case] expected: [f32; 4], + ) { + let mut output = [3.0, 17.0, 5.0, 7.0]; - // When - let actual = - sgemm_aat_lower_add(row_count, dimension_count, 1.0, &input, &mut invalid_output) - .unwrap_err(); + operation(2, 1, 0.0, &[2.0, 4.0], &mut output).unwrap(); - // Then - assert_eq!(actual, expected); + assert_eq!(output, expected); } - #[test] - fn input_size_overflow_returns_the_a_matrix_error() { - // Given - let row_count = usize::MAX; - let dimension_count = 2; - let expected = SgemmError::DimensionOverflow { - matrix_name: MatrixName::A, - rows: row_count, - cols: dimension_count, - }; + #[rstest] + #[case::replace(sgemm_aat_lower)] + #[case::add(sgemm_aat_lower_add)] + fn empty_matrices_require_no_storage(#[case] operation: LowerProduct) { + assert_eq!(operation(0, 3, 2.0, &[], &mut []), Ok(())); + } + + #[rstest] + #[case::short_input(3, 4, MatrixName::A, 2, 2, 3)] + #[case::long_input(5, 4, MatrixName::A, 2, 2, 5)] + #[case::short_output(4, 3, MatrixName::C, 2, 2, 3)] + #[case::long_output(4, 5, MatrixName::C, 2, 2, 5)] + fn length_mismatch_reports_the_matrix_and_preserves_output( + #[case] input_len: usize, + #[case] output_len: usize, + #[case] matrix_name: MatrixName, + #[case] expected_rows: usize, + #[case] expected_cols: usize, + #[case] actual_len: usize, + #[values(sgemm_aat_lower, sgemm_aat_lower_add)] operation: LowerProduct, + ) { + let input = vec![1.0; input_len]; + let mut output = vec![19.0; output_len]; - // When - let actual = sgemm_aat_lower(row_count, dimension_count, 1.0, &[], &mut []).unwrap_err(); + let error = operation(2, 2, 1.0, &input, &mut output).unwrap_err(); - // Then - assert_eq!(actual, expected); + assert_eq!( + error, + SgemmError::InvalidMatrixDimensions { + matrix_name, + expected_rows, + expected_cols, + actual_len + } + ); + assert_eq!(output, vec![19.0; output_len]); } - #[test] - fn output_size_overflow_returns_the_c_matrix_error() { - // Given - let row_count = usize::MAX; - let dimension_count = 0; - let expected = SgemmError::DimensionOverflow { - matrix_name: MatrixName::C, - rows: row_count, - cols: row_count, - }; + #[rstest] + #[case::input(usize::MAX, 2, MatrixName::A, usize::MAX, 2)] + #[case::output(usize::MAX, 0, MatrixName::C, usize::MAX, usize::MAX)] + fn shape_overflow_reports_the_matrix_and_preserves_output( + #[case] rows: usize, + #[case] dimensions: usize, + #[case] matrix_name: MatrixName, + #[case] error_rows: usize, + #[case] error_cols: usize, + #[values(sgemm_aat_lower, sgemm_aat_lower_add)] operation: LowerProduct, + ) { + let mut output = [23.0, 29.0]; - // When - let actual = - sgemm_aat_lower_add(row_count, dimension_count, 1.0, &[], &mut []).unwrap_err(); + let error = operation(rows, dimensions, 1.0, &[], &mut output).unwrap_err(); - // Then - assert_eq!(actual, expected); + assert_eq!( + error, + SgemmError::DimensionOverflow { + matrix_name, + rows: error_rows, + cols: error_cols + } + ); + assert_eq!(output, [23.0, 29.0]); } } diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs index e38b5b6c58..eb7cfe5d44 100644 --- a/diskann-utils/src/views.rs +++ b/diskann-utils/src/views.rs @@ -731,6 +731,53 @@ where // Tests // /////////// +#[cfg(test)] +mod matrix_shape_tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::wraps_to_empty(usize::MAX / 2 + 1, 2, &[])] + #[case::wraps_to_one(usize::MAX, usize::MAX, &[7])] + fn overflowing_shape_is_rejected_even_when_wrapping_matches_the_buffer( + #[case] rows: usize, + #[case] columns: usize, + #[case] values: &[u32], + ) { + let data = values.to_vec().into_boxed_slice(); + + let error = Matrix::try_from(data, rows, columns).unwrap_err(); + + assert_eq!((error.nrows, error.ncols), (rows, columns)); + assert_eq!(error.into_inner().as_ref(), values); + } + + #[test] + fn rejected_mutable_view_returns_the_original_buffer() { + let mut values = [11_u32]; + + let error = MutMatrixView::try_from(&mut values[..], usize::MAX, usize::MAX).unwrap_err(); + + let returned = error.into_inner(); + assert_eq!(returned, [11]); + returned[0] = 13; + assert_eq!(values, [13]); + } + + #[rstest] + #[case::zero_rows(0, usize::MAX)] + #[case::zero_columns(usize::MAX, 0)] + fn zero_sized_shapes_remain_valid_with_a_large_other_dimension( + #[case] rows: usize, + #[case] columns: usize, + ) { + let matrix = MatrixView::try_from(&[] as &[u32], rows, columns).unwrap(); + + assert_eq!((matrix.nrows(), matrix.ncols()), (rows, columns)); + assert!(matrix.as_slice().is_empty()); + } +} + #[cfg(test)] mod tests { use super::*; @@ -1053,8 +1100,6 @@ 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 + 1, 2).is_err()); } #[test] diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 7288b3ea0c..2bc32454fa 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -132,305 +132,365 @@ fn validate_output( #[cfg(test)] mod tests { use super::*; + #[cfg(not(miri))] + use crate::graph::pipnn::test_support; + use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2}; + #[cfg(not(miri))] + use diskann_vector::distance::Metric; + use diskann_wide::ARCH; + use rstest::rstest; + + const EMPTY: Candidate = Candidate::new(super::super::topk::UNASSIGNED, f32::INFINITY); + + #[rstest] + #[case::empty(0, 8, 0)] + #[case::singleton(1, 8, 0)] + #[case::zero_requested(5, 0, 0)] + #[case::below_available(5, 2, 2)] + #[case::all_available(5, 4, 4)] + #[case::above_available(5, 9, 4)] + fn neighbor_count_is_limited_to_other_points( + #[case] points: usize, + #[case] requested: usize, + #[case] expected: usize, + ) { + assert_eq!(leaf_neighbor_count(points, requested), expected); + } - mod test_support { - use super::*; - #[cfg(not(miri))] - use diskann_wide::arch::Target2; - use diskann_wide::arch::{self, Target1}; - - #[cfg(not(miri))] - pub(super) struct SelectLeaf(pub(super) M); - - #[cfg(not(miri))] - impl - Target2, MatrixView<'_, f32>, MutMatrixView<'_, Candidate>> - for SelectLeaf - { - fn run( - self, - arch: A, - points: MatrixView<'_, f32>, - output: MutMatrixView<'_, Candidate>, - ) -> ANNResult<()> { - select_leaf_neighbors::<_, M>( - arch, - points, - output, - &mut LeafKernelWorkspace::default(), - ) + #[cfg(not(miri))] + #[rstest] + #[case::l2(L2, Metric::L2, [[1,3,2,4], [0,2,3,4], [4,0,1,3], [0,2,1,4], [2,1,0,3]])] + #[case::cosine(Cosine, Metric::Cosine, [[3,1,2,4], [0,3,4,2], [4,3,0,1], [0,2,1,4], [2,1,3,0]])] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized, [[3,1,2,4], [0,3,4,2], [4,3,0,1], [0,2,1,4], [2,1,3,0]])] + #[case::inner_product(InnerProduct, Metric::InnerProduct, [[3,1,2,4], [0,3,2,4], [4,3,0,1], [0,2,1,4], [2,1,0,3]])] + fn neighbors_are_local_non_self_ids_in_metric_order( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[case] expected_ids: [[u32; 4]; 5], + #[values(0, 1, 2, 3, 4)] neighbors: usize, + #[values(2, 7, 8, 9, 15, 16, 17, 128, 129)] dimensions: usize, + ) { + // These five points have distinct scores in each row for every metric. + let coordinates = [ + [2.0, 2.0], + [-1.0, 3.0], + [0.0, -2.0], + [5.0, 0.0], + [-3.0, -4.0], + ]; + let values = test_support::packed_points( + &coordinates, + dimensions, + scalar_metric == Metric::CosineNormalized, + ); + let points = MatrixView::try_from(values.as_slice(), 5, dimensions).unwrap(); + let mut output = vec![Candidate::new(0, -100.0); 5 * neighbors]; + let mut workspace = LeafKernelWorkspace::default(); + + select_leaf_neighbors::<_, M>( + ARCH, + points, + MutMatrixView::try_from(output.as_mut_slice(), 5, neighbors).unwrap(), + &mut workspace, + ) + .unwrap(); + + for point in 0..5 { + let actual = &output[point * neighbors..(point + 1) * neighbors]; + assert_eq!( + actual.iter().map(|c| c.local_idx).collect::>(), + expected_ids[point][..neighbors], + "point={point}" + ); + for candidate in actual { + let expected = test_support::distance( + scalar_metric, + points.row(point), + points.row(candidate.local_idx as usize), + ); + // Only two coordinates are nonzero, so eight f32 ulps at the score scale + // cover the dot/norm rounding without admitting another neighbor. + let tolerance = 8.0 * f64::from(f32::EPSILON) * expected.abs().max(1.0); + assert!( + (f64::from(candidate.distance) - expected).abs() <= tolerance, + "point={point}, candidate={candidate:?}, expected={expected}" + ); } } + } - struct KernelCall<'a> { - distances: MatrixView<'a, f32>, - output: MutMatrixView<'a, Candidate>, - } - - struct RankDistances; - - impl Target1> for RankDistances { - fn run(self, arch: A, call: KernelCall<'_>) { - rank_leaf_distances(arch, call.distances, call.output, &mut Vec::new()); - } + #[cfg(not(miri))] + #[rstest] + fn neighbors_match_scalar_ranking_across_leaf_sizes_and_counts( + #[values(1, 2, 16, 17, 18, 33, 34)] point_count: usize, + #[values(1, 2, 3, 10, 11, 17)] requested: usize, + ) { + // Slightly increasing gaps avoid ties between the two sides of each point. + // Center the coordinates so intermediate norm sums stay exactly representable. + let mut values: Vec<_> = (0..point_count).map(|i| (64 * i + i * i) as f32).collect(); + let center = values[point_count - 1] / 2.0; + for value in &mut values { + *value -= center; } - - pub(super) fn rank_distance_fixture( - distances: &[f32], - points: usize, - width: usize, - ) -> Vec { - let mut output = vec![Candidate::default(); points * width]; - arch::dispatch1_no_features( - RankDistances, - KernelCall { - distances: MatrixView::try_from(distances, points, points).unwrap(), - output: MutMatrixView::try_from(output.as_mut_slice(), points, width).unwrap(), - }, + let neighbors = requested.min(point_count - 1); + let mut output = vec![EMPTY; point_count * neighbors]; + + select_leaf_neighbors::<_, L2>( + ARCH, + MatrixView::try_from(values.as_slice(), point_count, 1).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap(); + + for point in 0..point_count { + let mut expected: Vec<_> = values + .iter() + .enumerate() + .filter(|&(other, _)| other != point) + .map(|(other, &value)| { + Candidate::new(other as u32, (values[point] - value).powi(2)) + }) + .collect(); + expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); + expected.truncate(neighbors); + assert_eq!( + &output[point * neighbors..(point + 1) * neighbors], + expected, + "point={point}, point_count={point_count}, neighbors={neighbors}" ); - output - } - } - - mod leaf_neighbor_count_tests { - use super::leaf_neighbor_count; - use rstest::rstest; - - #[rstest] - #[case::singleton(1, 3, 0)] - #[case::clamped(4, 4, 3)] - #[case::requested(8, 5, 5)] - fn count_is_bounded_by_non_self_points( - #[case] points: usize, - #[case] requested: usize, - #[case] expected: usize, - ) { - assert_eq!(leaf_neighbor_count(points, requested), expected); } } - // These entry-point tests call GEMM. Miri checks the distance scan below. #[cfg(not(miri))] - mod select_leaf_neighbors_tests { - use super::test_support::SelectLeaf; - use super::*; - use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2, scalar_ranking}; - use diskann_vector::distance::Metric; - use diskann_wide::arch; - use rstest::rstest; - - #[rstest] - #[case::l2(L2, Metric::L2)] - #[case::cosine(Cosine, Metric::Cosine)] - #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] - #[case::inner_product(InnerProduct, Metric::InnerProduct)] - fn leaf_neighbors_match_scalar_ranking( - #[case] metric: M, - #[case] scalar_metric: Metric, - #[values(2, 7, 8, 9, 15, 16, 17, 127, 128, 129, 384, 768)] dimensions: usize, - #[values(1, 3, 10, 11)] width: usize, - ) { - // Given: 34 points supply two SIMD groups and a tail in the last row. - // All coordinates contribute. Only normalized cosine receives unit vectors. - let values = scalar_ranking::arc_vectors( - (0..34).map(f64::from), - dimensions, - scalar_metric == Metric::CosineNormalized, - ); - let points = MatrixView::try_from(values.as_slice(), 34, dimensions).unwrap(); - let expected: Vec> = points - .row_iter() - .enumerate() - .map(|(source, point)| { - let mut row: Vec<_> = points - .row_iter() - .enumerate() - .filter(|&(target, _)| source != target) - .map(|(target, vector)| { - ( - target as u32, - scalar_ranking::distance(scalar_metric, point, vector), - ) - }) - .collect(); - row.sort_by(|left, right| left.1.total_cmp(&right.1)); - row + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn large_dense_leaves_select_nearest_non_self_neighbors( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[values((33, 384), (65, 768), (129, 1536), (35, 1537))] shape: (usize, usize), + #[values(3, 11)] neighbors: usize, + ) { + let (point_count, dimensions) = shape; + let mut values = test_support::dense_points(point_count, dimensions, 1287); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut values, dimensions); + } + let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let mut output = vec![EMPTY; point_count * neighbors]; + + select_leaf_neighbors::<_, M>( + ARCH, + points, + MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap(); + + // Unnormalized dyadic sums are exact; cosine rounds sqrt/division, while + // normalized dot products also accumulate coordinate rounding. + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + }; + for point in 0..point_count { + let mut expected: Vec<_> = (0..point_count) + .filter(|&other| other != point) + .map(|other| { + ( + other as u32, + test_support::distance(scalar_metric, points.row(point), points.row(other)), + ) }) .collect(); - let ranking_tolerance = scalar_ranking::ranking_tolerance(scalar_metric, dimensions); - let mut output = vec![Candidate::default(); 34 * width]; + expected.sort_by(|left, right| left.1.total_cmp(&right.1)); + let actual = &output[point * neighbors..(point + 1) * neighbors]; + for (rank, candidate) in actual.iter().enumerate() { + assert!( + !actual[..rank] + .iter() + .any(|previous| previous.local_idx == candidate.local_idx), + "duplicate neighbor for point={point}: {candidate:?}" + ); + let own_score = expected + .iter() + .find(|&&(id, _)| id == candidate.local_idx) + .unwrap_or_else(|| { + panic!("invalid or self neighbor for point={point}: {candidate:?}") + }) + .1; + assert!( + (f64::from(candidate.distance) - own_score).abs() <= tolerance, + "shape={shape:?}, point={point}, candidate={candidate:?}, score={own_score}" + ); + // Equal scores may use either ID, but every rank must be nearest-first. + assert!( + (own_score - expected[rank].1).abs() <= tolerance, + "shape={shape:?}, point={point}, rank={rank}, score={own_score}, expected={:?}", + expected[rank] + ); + } + } + } - // When: dispatch the complete points-to-metric-to-neighbors pipeline. - arch::dispatch2_no_features( - SelectLeaf(metric), + #[cfg(not(miri))] + #[test] + fn workspace_reuse_does_not_mix_results_from_different_leaves() { + let values = [0.0, 1.0, 4.0, 10.0, 21.0]; + let mut workspace = LeafKernelWorkspace::default(); + let mut output = Vec::new(); + + // Grow, shrink, change K, and finish with a singleton using the same buffers. + for (count, neighbors) in [(4, 1), (2, 1), (5, 3), (1, 0)] { + output.resize(count * neighbors, Candidate::new(4, -100.0)); + let points = MatrixView::try_from(&values[..count], count, 1).unwrap(); + select_leaf_neighbors::<_, L2>( + ARCH, points, - MutMatrixView::try_from(output.as_mut_slice(), 34, width).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), count, neighbors).unwrap(), + &mut workspace, ) .unwrap(); - // Then: compare in returned order, including each ID's own scalar score. - for (source, (row, expected)) in output.chunks_exact(width).zip(&expected).enumerate() { - let ids: Vec<_> = row.iter().map(|candidate| candidate.local_idx).collect(); - scalar_ranking::assert_ranked_ids(&ids, expected, ranking_tolerance); - assert!( - row.windows(2) - .all(|pair| pair[0].distance <= pair[1].distance), - "unordered row {source}" + for point in 0..count { + let mut expected: Vec<_> = (0..count) + .filter(|&other| other != point) + .map(|other| { + Candidate::new(other as u32, (values[point] - values[other]).powi(2)) + }) + .collect(); + expected.sort_by(|a, b| a.distance.total_cmp(&b.distance)); + expected.truncate(neighbors); + assert_eq!( + &output[point * neighbors..(point + 1) * neighbors], + expected, + "count={count}, neighbors={neighbors}, point={point}" ); - for candidate in row { - let score = expected - .iter() - .find(|&&(id, _)| id == candidate.local_idx) - .unwrap() - .1; - let distance_tolerance = scalar_ranking::distance_tolerance( - scalar_metric, - points.row(source), - points.row(candidate.local_idx as usize), - ); - assert!( - (f64::from(candidate.distance) - score).abs() <= distance_tolerance, - "row {source}, candidate {candidate:?}, scalar {score}, tolerance {distance_tolerance}" - ); - } } } + } - #[rstest::rstest] - #[case::missing_row(2)] - #[case::extra_row(4)] - fn invalid_output_rows_return_an_error(#[case] rows: usize) { - let values = [0.0_f32, 1.0, 3.0]; - let mut output = vec![Candidate::default(); rows]; - let mut workspace = LeafKernelWorkspace::default(); - - let error = select_leaf_neighbors::<_, L2>( - diskann_wide::ARCH, - MatrixView::try_from(&values[..], 3, 1).unwrap(), - MutMatrixView::try_from(output.as_mut_slice(), rows, 1).unwrap(), - &mut workspace, - ) - .unwrap_err(); + #[rstest] + #[case::wrong_rows((2, 1), LeafKernelError::InvalidOutputRows { points: 3, rows: 2 })] + #[case::too_many_neighbors((3, 3), LeafKernelError::InvalidNeighborCount { points: 3, neighbors: 3, maximum: 2 })] + fn invalid_output_shape_preserves_output_and_workspace( + #[case] shape: (usize, usize), + #[case] expected: LeafKernelError, + ) { + let values = [1.0, 2.0, 4.0]; + let mut output = vec![Candidate::new(2, 9.0); shape.0 * shape.1]; + let previous_output = output.clone(); + let mut workspace = LeafKernelWorkspace { + distance_scratch: vec![11.0, 13.0], + worst: vec![3.0, 7.0], + }; + + let error = select_leaf_neighbors::<_, L2>( + ARCH, + MatrixView::try_from(&values[..], 3, 1).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), shape.0, shape.1).unwrap(), + &mut workspace, + ) + .unwrap_err(); + + assert_eq!(error.downcast_ref::(), Some(&expected)); + assert_eq!(output, previous_output); + assert_eq!(workspace.distance_scratch, [11.0, 13.0]); + assert_eq!(workspace.worst, [3.0, 7.0]); + } - assert_eq!( - error.downcast_ref::(), - Some(&LeafKernelError::InvalidOutputRows { points: 3, rows }) - ); + #[test] + fn metric_failure_is_returned_without_publishing_neighbors() { + #[derive(Debug, thiserror::Error)] + #[error("distance computation failed")] + struct MetricFailure; + struct FailingMetric; + impl LeafMetric for FailingMetric { + fn compute_distances(_: MatrixView<'_, f32>, _: &mut [f32]) -> ANNResult<()> { + Err(ANNError::new(MetricFailure)) + } } + let values = [1.0, 2.0]; + let mut output = [Candidate::new(1, 7.0), Candidate::new(0, 7.0)]; + + let error = select_leaf_neighbors::<_, FailingMetric>( + ARCH, + MatrixView::try_from(&values[..], 2, 1).unwrap(), + MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap_err(); + + assert!(error.is::()); + assert_eq!(output, [Candidate::new(1, 7.0), Candidate::new(0, 7.0)]); + } + + // Keep this module path available to the nightly Miri selector. + mod rank_leaf_distances_tests { + use super::*; #[test] - fn excess_neighbor_count_returns_error() { - let values = [0.0_f32, 1.0, 3.0]; - let mut output = [Candidate::default(); 9]; - let mut workspace = LeafKernelWorkspace::default(); - let expected = LeafKernelError::InvalidNeighborCount { - points: 3, - neighbors: 3, - maximum: 2, - }; - - let error = select_leaf_neighbors::<_, L2>( - diskann_wide::ARCH, - MatrixView::try_from(&values[..], 3, 1).unwrap(), - MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), - &mut workspace, - ) - .unwrap_err(); + fn ranking_reads_only_pairs_in_the_strict_lower_triangle() { + // Diagonal and upper entries are deliberately better than every real pair. + let distances = [ + -100.0, -100.0, -100.0, -100.0, 7.0, -100.0, -100.0, -100.0, 3.0, 8.0, -100.0, + -100.0, 5.0, 2.0, 6.0, -100.0, + ]; + let mut output = [Candidate::new(0, -200.0); 8]; + let mut limits = vec![-200.0; 4]; - assert_eq!(error.downcast_ref::(), Some(&expected)); - } + rank_leaf_distances( + ARCH, + MatrixView::try_from(&distances[..], 4, 4).unwrap(), + MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), + &mut limits, + ); - #[test] - fn l2_neighbors_update_with_reused_workspace() { - // Given: squared distances on the line determine each point's two neighbors. - let values = [0.0_f32, 1.0, 3.0, 10.0]; let expected = [ - Candidate::new(1, (values[0] - values[1]).powi(2)), - Candidate::new(2, (values[0] - values[2]).powi(2)), - Candidate::new(0, (values[1] - values[0]).powi(2)), - Candidate::new(2, (values[1] - values[2]).powi(2)), - Candidate::new(1, (values[2] - values[1]).powi(2)), - Candidate::new(0, (values[2] - values[0]).powi(2)), - Candidate::new(2, (values[3] - values[2]).powi(2)), - Candidate::new(1, (values[3] - values[1]).powi(2)), + [Candidate::new(2, 3.0), Candidate::new(3, 5.0)], + [Candidate::new(3, 2.0), Candidate::new(0, 7.0)], + [Candidate::new(0, 3.0), Candidate::new(3, 6.0)], + [Candidate::new(1, 2.0), Candidate::new(0, 5.0)], ]; - let mut output = [Candidate::default(); 8]; - let mut workspace = LeafKernelWorkspace::default(); + assert_eq!(output, expected.as_flattened()); + assert_eq!(limits, [5.0, 7.0, 6.0, 5.0]); + } - // When: first populate output and scratch with a valid larger leaf. - select_leaf_neighbors::<_, L2>( - diskann_wide::ARCH, - MatrixView::try_from(&values[..], 4, 1).unwrap(), - MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), - &mut workspace, - ) - .unwrap(); - assert_eq!(output, expected); + #[rstest] + #[case::nan(f32::NAN)] + #[case::infinity(f32::INFINITY)] + fn unrankable_pairs_leave_unassigned_slots(#[case] unrankable: f32) { + let mut distances = [unrankable; 9]; + distances[6] = 4.0; + let mut output = [Candidate::new(1, -100.0); 6]; + let mut limits = vec![-100.0; 5]; + + rank_leaf_distances( + ARCH, + MatrixView::try_from(&distances[..], 3, 3).unwrap(), + MutMatrixView::try_from(&mut output[..], 3, 2).unwrap(), + &mut limits, + ); - // A smaller leaf has one finite pair and leaves the other slots unassigned. - let smaller = [0.0, 2.0, f32::NAN]; - select_leaf_neighbors::<_, L2>( - diskann_wide::ARCH, - MatrixView::try_from(&smaller[..], 3, 1).unwrap(), - MutMatrixView::try_from(&mut output[..6], 3, 2).unwrap(), - &mut workspace, - ) - .unwrap(); assert_eq!( - output[..6], + output, [ - Candidate::new(1, 4.0), - Candidate::default(), + Candidate::new(2, 4.0), + EMPTY, + EMPTY, + EMPTY, Candidate::new(0, 4.0), - Candidate::default(), - Candidate::default(), - Candidate::default(), + EMPTY ] ); - - // Growing again must replace the NaN scratch and stale rows. Doubling - // every coordinate preserves IDs and multiplies squared distances by four. - let scaled = values.map(|value| 2.0 * value); - let expected = expected - .map(|candidate| Candidate::new(candidate.local_idx, 4.0 * candidate.distance)); - select_leaf_neighbors::<_, L2>( - diskann_wide::ARCH, - MatrixView::try_from(&scaled[..], 4, 1).unwrap(), - MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), - &mut workspace, - ) - .unwrap(); - assert_eq!(output, expected); - } - } - - mod rank_leaf_distances_tests { - use super::test_support::rank_distance_fixture; - use super::*; - - #[test] - fn lower_triangle_pairs_update_both_endpoints() { - // Given: only three lower-triangle pairs rank. The diagonal and upper - // triangle contain better distances, so reading either changes the answer. - let mut distances = [f32::NEG_INFINITY; 18 * 18]; - let mut matrix = MutMatrixView::try_from(&mut distances[..], 18, 18).unwrap(); - for source in 1..18 { - matrix.row_mut(source)[..source].fill(f32::INFINITY); - } - matrix.row_mut(16)[0] = 3.0; - matrix.row_mut(17)[1] = 2.0; - matrix.row_mut(17)[16] = 1.0; - let mut expected = [[Candidate::default(); 2]; 18]; - expected[0][0] = Candidate::new(16, 3.0); - expected[1][0] = Candidate::new(17, 2.0); - expected[16] = [Candidate::new(17, 1.0), Candidate::new(0, 3.0)]; - expected[17] = [Candidate::new(16, 1.0), Candidate::new(1, 2.0)]; - - // When: row 16 supplies a full SIMD group; row 17 also supplies a tail. - let actual = rank_distance_fixture(matrix.as_slice(), 18, 2); - - // Then: every reciprocal ID is local to its output row; other rows stay empty. - assert_eq!(actual.as_slice(), expected.as_flattened()); + assert_eq!(limits, [f32::INFINITY; 3]); } } } diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index f1ae207b8f..275382a40d 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -107,224 +107,197 @@ impl LeafMetric for CosineNormalized { } } -// GEMM-backed tests run natively, not under Miri. #[cfg(test)] -#[cfg(not(miri))] mod tests { use super::*; - use crate::graph::pipnn::scalar_ranking; + #[cfg(not(miri))] + use crate::graph::pipnn::test_support; + #[cfg(not(miri))] use diskann_vector::distance::Metric; use rstest::rstest; - const POINT_COUNT: usize = 2; - const DIMENSION_COUNT: usize = 2; - const FIRST_POINT: usize = 0; - const SECOND_POINT: usize = 1; - const STALE_DISTANCE: f32 = 99.0; - const FLOAT_TOLERANCE: f32 = 1.0e-6; - - fn compute_pair_ranking( - first_point: [f32; DIMENSION_COUNT], - second_point: [f32; DIMENSION_COUNT], - ) -> f32 { - let point_values = [ - first_point[0], - first_point[1], - second_point[0], - second_point[1], - ]; - let points = MatrixView::try_from(&point_values[..], POINT_COUNT, DIMENSION_COUNT).unwrap(); - let mut storage = [STALE_DISTANCE; POINT_COUNT * POINT_COUNT]; - - M::compute_distances(points, &mut storage).unwrap(); - - storage[SECOND_POINT * POINT_COUNT + FIRST_POINT] - } - - mod compute_distances_tests { - use super::*; - - #[rstest] - #[case::l2(L2, Metric::L2)] - #[case::cosine(Cosine, Metric::Cosine)] - #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] - #[case::inner_product(InnerProduct, Metric::InnerProduct)] - fn lower_triangle_matches_scalar_ranking( - #[case] _metric: M, - #[case] scalar_metric: Metric, - #[values(64, 128, 256)] point_count: usize, - #[values(128, 384, 768)] dimensions: usize, - ) { - // Given: spread all points over the same arc as the matrix grows. - // NaN marks unwritten entries. The scalar oracle includes self pairs. - let values = scalar_ranking::arc_vectors( - (0..point_count).map(|point| 32.0 * point as f64 / point_count as f64), - dimensions, - scalar_metric == Metric::CosineNormalized, - ); - let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); - let expected: Vec> = points - .row_iter() - .enumerate() - .map(|(source, point)| { - points - .row_iter() - .take(source + 1) - .map(|target| { - ( - scalar_ranking::distance(scalar_metric, point, target), - scalar_ranking::distance_tolerance(scalar_metric, point, target), - ) - }) - .collect() - }) - .collect(); - let mut distances = vec![f32::NAN; point_count * point_count]; - - // When: call the metric directly, without TopK selection. - M::compute_distances(points, &mut distances).unwrap(); - - // Then: check every required entry, including the diagonal. - // The upper triangle is outside the metric contract. - let actual = - MatrixView::try_from(distances.as_slice(), point_count, point_count).unwrap(); - for (source, row) in expected.iter().enumerate() { - for (target, &(expected, tolerance)) in row.iter().enumerate() { - let actual = f64::from(actual.row(source)[target]); - assert!( - (actual - expected).abs() <= tolerance, - "{scalar_metric:?}, {point_count} points, D={dimensions}, pair ({source}, {target}): actual {actual}, expected {expected}, tolerance {tolerance}" - ); - } - } - } - - #[test] - fn l2_writes_squared_distances() { - // Given - let first_point = [3.0_f32, 4.0]; - let second_point = [0.0_f32, 4.0]; - let x_difference = first_point[0] - second_point[0]; - let y_difference = first_point[1] - second_point[1]; - let expected = x_difference.mul_add(x_difference, y_difference * y_difference); - - // When - let actual = compute_pair_ranking::(first_point, second_point); - - // Then - assert_eq!(actual, expected); + #[cfg(not(miri))] + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn lower_triangle_matches_scalar_distances( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[values(1, 4, 17)] point_count: usize, + #[values(1, 2, 7, 8, 9, 15, 16, 17, 127, 128, 129)] dimensions: usize, + ) { + // Small integers keep L2 and dot products exact; the first coordinate gives + // every vector a nonzero norm. Every dimension contributes to the result. + let mut values: Vec<_> = (0..point_count * dimensions) + .map(|index| (index % 11) as f32 - 5.0) + .collect(); + for (point, row) in values.chunks_exact_mut(dimensions).enumerate() { + row[0] = point as f32 + 1.0; } - - #[test] - fn cosine_writes_one_minus_similarity() { - // Given - let first_point = [2.0_f32, 0.0]; - let second_point = [1.0_f32, 1.0]; - let dot = first_point[0].mul_add(second_point[0], first_point[1] * second_point[1]); - let first_norm = first_point[0].hypot(first_point[1]); - let second_norm = second_point[0].hypot(second_point[1]); - let expected = 1.0 - dot / (first_norm * second_norm); - - // When - let actual = compute_pair_ranking::(first_point, second_point); - - // Then - assert!( - (actual - expected).abs() <= FLOAT_TOLERANCE, - "actual {actual} differs from expected {expected}" - ); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut values, dimensions); } + let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let mut output = vec![f32::NAN; point_count * point_count]; - #[rstest] - #[case::zero([0.0, 0.0])] - #[case::subnormal([f32::MIN_POSITIVE.sqrt() / 2.0, 0.0])] - fn cosine_writes_one_for_small_norms(#[case] small_point: [f32; DIMENSION_COUNT]) { - // Given: a zero or subnormal norm represents zero similarity. - let unit_point = [1.0_f32, 0.0]; - let expected = 1.0; - - // When - let actual = compute_pair_ranking::(small_point, unit_point); + M::compute_distances(points, &mut output).unwrap(); - // Then - assert_eq!(actual, expected); + for source in 0..point_count { + for target in 0..=source { + let expected = + test_support::distance(scalar_metric, points.row(source), points.row(target)); + let actual = f64::from(output[source * point_count + target]); + // Cosine reductions and normalization round in f32; allow eight ulps + // per dimension at unit scale. Integer L2 and dot products are exact. + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + Metric::Cosine | Metric::CosineNormalized => { + 8.0 * f64::from(f32::EPSILON) * dimensions as f64 + } + }; + assert!( + (actual - expected).abs() <= tolerance, + "pair=({source},{target}), dimensions={dimensions}: {actual} != {expected}" + ); + } } + } - #[test] - fn normalized_cosine_writes_negative_dot_products() { - // Given - let first_point = [1.0_f32, 0.0]; - let second_point = [0.6_f32, 0.8]; - let dot = first_point[0].mul_add(second_point[0], first_point[1] * second_point[1]); - let expected = -dot; - - // When - let actual = compute_pair_ranking::(first_point, second_point); - - // Then - assert_eq!(actual, expected); + #[cfg(not(miri))] + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn large_dense_inputs_match_scalar_distances( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[values((33, 384), (65, 768), (129, 1536), (17, 1537), (513, 129), (17, 4097))] shape: ( + usize, + usize, + ), + ) { + let (point_count, dimensions) = shape; + let mut values = test_support::dense_points(point_count, dimensions, 1287); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut values, dimensions); } - - #[test] - fn inner_product_writes_negative_dot_products() { - // Given - let first_point = [2.0_f32, -1.0]; - let second_point = [3.0_f32, 4.0]; - let dot = first_point[0].mul_add(second_point[0], first_point[1] * second_point[1]); - let expected = -dot; - - // When - let actual = compute_pair_ranking::(first_point, second_point); - - // Then - assert_eq!(actual, expected); + let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let mut output = vec![f32::NAN; point_count * point_count]; + + M::compute_distances(points, &mut output).unwrap(); + + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + // Dyadic dot/norm sums are exact; only square roots and division round. + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + // Normalized f32 coordinates need a dot-product rounding bound. + // Sum |x*y| is at most approximately one; EPSILON allows both + // product and reduction rounding in gamma = n*u/(1-n*u). + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + }; + for source in 0..point_count { + for target in 0..=source { + let expected = + test_support::distance(scalar_metric, points.row(source), points.row(target)); + let actual = f64::from(output[source * point_count + target]); + assert!( + (actual - expected).abs() <= tolerance, + "shape={shape:?}, pair=({source},{target}): {actual} != {expected}, tolerance={tolerance}" + ); + } } + } - #[rstest] - #[case::l2(compute_pair_ranking::)] - #[case::cosine(compute_pair_ranking::)] - #[case::inner_product(compute_pair_ranking::)] - fn nan_coordinate_produces_nan_ranking( - #[case] compute: fn([f32; DIMENSION_COUNT], [f32; DIMENSION_COUNT]) -> f32, - ) { - // Given: neither vector has zero norm. - let first_point = [f32::NAN, 1.0]; - let second_point = [1.0, 0.0]; + #[cfg(not(miri))] + #[rstest] + #[case::squared_l2(L2, &[2.0, 0.0, 0.0, 3.0, -4.0, 0.0], [0.0, 13.0, 0.0, 36.0, 25.0, 0.0])] + #[case::negative_dot(InnerProduct, &[2.0, 0.0, 0.0, 3.0, -4.0, 0.0], [-4.0, 0.0, -9.0, 8.0, 0.0, -16.0])] + #[case::cosine(Cosine, &[2.0, 0.0, 0.0, 3.0, -4.0, 0.0], [0.0, 1.0, 0.0, 2.0, 1.0, 0.0])] + #[case::normalized_cosine(CosineNormalized, &[1.0, 0.0, 0.0, 1.0, -1.0, 0.0], [-1.0, 0.0, -1.0, 1.0, 0.0, -1.0])] + fn ranking_values_follow_the_metric_definition( + #[case] _metric: M, + #[case] values: &[f32], + #[case] expected: [f32; 6], + ) { + let points = MatrixView::try_from(values, 3, 2).unwrap(); + let mut output = [42.0; 9]; + + M::compute_distances(points, &mut output).unwrap(); + + assert_eq!( + [ + output[0], output[3], output[4], output[6], output[7], output[8] + ], + expected + ); + } - // When - let actual = compute(first_point, second_point); + #[cfg(not(miri))] + #[rstest] + #[case::zero(0.0)] + #[case::squared_norm_underflows(f32::MIN_POSITIVE)] + fn cosine_gives_unit_distance_to_points_with_small_norms(#[case] coordinate: f32) { + let values = [coordinate, 0.0, 0.0, 2.0]; + let points = MatrixView::try_from(&values[..], 2, 2).unwrap(); + let mut output = [42.0; 4]; - // Then - assert!(actual.is_nan()); - } + Cosine::compute_distances(points, &mut output).unwrap(); - #[rstest] - #[case::inner_product(compute_pair_ranking::)] - #[case::normalized_cosine(compute_pair_ranking::)] - fn orthogonal_vectors_have_zero_ranking( - #[case] compute: fn([f32; DIMENSION_COUNT], [f32; DIMENSION_COUNT]) -> f32, - ) { - // The contract permits either zero sign at both metric entry points. - assert_eq!(compute([1.0, 0.0], [0.0, 1.0]), 0.0); - } + assert_eq!([output[0], output[2], output[3]], [1.0, 1.0, 0.0]); + } - #[test] - fn inner_product_includes_last_coordinate() { - let dimensions = 129; - let mut values = vec![0.0; 2 * dimensions]; - values[dimensions - 1] = 2.0; - values[dimensions] = -1.0; - values[2 * dimensions - 1] = 3.0; - let points = MatrixView::try_from(values.as_slice(), 2, dimensions).unwrap(); - let mut distances = [STALE_DISTANCE; 4]; + #[cfg(not(miri))] + #[rstest] + #[case::l2(L2, 2.0)] + #[case::cosine(Cosine, 1.0)] + #[case::normalized_cosine(CosineNormalized, 0.0)] + #[case::inner_product(InnerProduct, 0.0)] + fn a_nan_point_does_not_change_other_pair_distances( + #[case] _metric: M, + #[case] expected_finite_pair: f32, + ) { + let values = [1.0, 0.0, 0.0, -1.0, f32::NAN, f32::NAN]; + let mut output = [42.0; 9]; + + M::compute_distances( + MatrixView::try_from(&values[..], 3, 2).unwrap(), + &mut output, + ) + .unwrap(); - InnerProduct::compute_distances(points, &mut distances).unwrap(); + assert_eq!(output[3], expected_finite_pair); + assert!(output[6..=8].iter().all(|distance| distance.is_nan())); + } - // Dot products are 4, 6 and 10; all nonzero pair contribution is in the tail. - assert_eq!( - [distances[0], distances[2], distances[3]], - [-4.0, -6.0, -10.0] - ); - } + #[rstest] + #[case::l2(L2)] + #[case::cosine(Cosine)] + #[case::normalized_cosine(CosineNormalized)] + #[case::inner_product(InnerProduct)] + fn storage_length_mismatch_reports_the_output_matrix_error(#[case] _metric: M) { + let values = [1.0, 0.0, 0.0, 2.0]; + let mut output = [42.0; 5]; + + let error = M::compute_distances( + MatrixView::try_from(&values[..], 2, 2).unwrap(), + &mut output, + ) + .unwrap_err(); + + assert_eq!( + error.downcast_ref::(), + Some(&diskann_linalg::SgemmError::InvalidMatrixDimensions { + matrix_name: diskann_linalg::MatrixName::C, + expected_rows: 2, + expected_cols: 2, + actual_len: 5, + }) + ); } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index a9479c2418..0dff853d77 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -41,190 +41,190 @@ fn cosine_distance(dot: f32, source_norm: f32, target_norm: f32) -> f32 { } #[cfg(test)] -#[cfg(not(miri))] -mod scalar_ranking { +mod cosine_distance_tests { + use super::cosine_distance; + use rstest::rstest; + + #[rstest] + #[case::same_direction(6.0, 2.0, 3.0, 0.0)] + #[case::opposite_directions(-6.0, 2.0, 3.0, 2.0)] + #[case::orthogonal(0.0, 2.0, 3.0, 1.0)] + #[case::positive_similarity(3.0, 2.0, 3.0, 0.5)] + #[case::negative_similarity(-3.0, 2.0, 3.0, 1.5)] + fn distance_uses_both_vector_norms( + #[case] dot: f32, + #[case] source_norm: f32, + #[case] target_norm: f32, + #[case] expected: f32, + ) { + assert_eq!(cosine_distance(dot, source_norm, target_norm), expected); + } + + #[rstest] + #[case::above_one(1.0 + f32::EPSILON, 0.0)] + #[case::below_minus_one(-1.0 - f32::EPSILON, 2.0)] + fn rounding_outside_the_similarity_range_is_clamped( + #[case] similarity: f32, + #[case] expected: f32, + ) { + assert_eq!(cosine_distance(similarity, 1.0, 1.0), expected); + } + + #[rstest] + #[case::zero(0.0)] + #[case::below_cutoff(f32::from_bits(f32::MIN_POSITIVE.sqrt().to_bits() - 1))] + fn a_small_norm_takes_precedence_over_the_dot_product( + #[case] small_norm: f32, + #[values(0.75, f32::NAN)] dot: f32, + #[values(false, true)] small_source: bool, + ) { + let (source_norm, target_norm) = if small_source { + (small_norm, 1.0) + } else { + (1.0, small_norm) + }; + + assert_eq!(cosine_distance(dot, source_norm, target_norm), 1.0); + } + + #[rstest] + #[case::source(true)] + #[case::target(false)] + fn a_norm_at_the_cutoff_still_contributes_similarity(#[case] source_at_cutoff: bool) { + let norm = f32::MIN_POSITIVE.sqrt(); + let (source_norm, target_norm) = if source_at_cutoff { + (norm, 1.0) + } else { + (1.0, norm) + }; + + assert_eq!(cosine_distance(norm / 4.0, source_norm, target_norm), 0.75); + } + + #[rstest] + #[case::dot(f32::NAN, 2.0, 3.0)] + #[case::source_norm(1.0, f32::NAN, 3.0)] + #[case::target_norm(1.0, 2.0, f32::NAN)] + fn nan_propagates_when_neither_norm_is_small( + #[case] dot: f32, + #[case] source_norm: f32, + #[case] target_norm: f32, + ) { + assert!(cosine_distance(dot, source_norm, target_norm).is_nan()); + } +} + +#[cfg(all(test, not(miri)))] +mod test_support { use diskann_vector::distance::Metric; - // Sample an irregular arc with increasing radius. Repeat its x/y coordinates - // through every dimension, so an odd final coordinate changes the geometry. - // Unequal angle steps avoid systematic nearest-neighbor ties. - pub(super) fn arc_vectors( - parameters: impl IntoIterator, - dimensions: usize, - normalize: bool, - ) -> Vec { - parameters - .into_iter() - .flat_map(|t| { - let angle = 0.03 * t + 0.001 * t * t; - let radius = 1.0 + t / 64.0; - let coordinates = [radius * angle.cos(), radius * angle.sin()]; - let mut vector: Vec<_> = (0..dimensions) - .map(|dimension| coordinates[dimension % 2] as f32) - .collect(); - if normalize { - let norm = vector - .iter() - .map(|&x| f64::from(x).powi(2)) - .sum::() - .sqrt(); - for x in &mut vector { - *x = (f64::from(*x) / norm) as f32; - } - } - vector - }) - .collect() + pub(super) fn dense_points(rows: usize, dimensions: usize, seed: u64) -> Vec { + use rand::{Rng, SeedableRng, rngs::StdRng}; + + // Multiples of 1/8 keep unnormalized products exact at the tested sizes. + let mut rng = StdRng::seed_from_u64(seed); + let mut values: Vec<_> = (0..rows * dimensions) + .map(|_| rng.random_range(-16..=16) as f32 / 8.0) + .collect(); + for (point, row) in values.chunks_exact_mut(dimensions).enumerate() { + // A substantial final coordinate makes an omitted dimension visible + // even after normalization. It also guarantees nonzero norms. + row[dimensions - 1] = 8.0 + (point % 7) as f32; + } + values } - // Use f64 scalar arithmetic on the actual f32 inputs, not GEMM or a kernel - // metric. These fixtures have finite coordinates and nonzero normal norms. - // L2 includes the point norm; partition ranking can omit this row constant. + // These scalar definitions use the actual f32 inputs, with f64 arithmetic. + // Callers supply finite vectors with nonzero norms. L2 includes the point norm. pub(super) fn distance(metric: Metric, point: &[f32], target: &[f32]) -> f64 { - let dot: f64 = point - .iter() - .zip(target) - .map(|(&x, &y)| f64::from(x) * f64::from(y)) - .sum(); match metric { Metric::L2 => point .iter() .zip(target) .map(|(&x, &y)| (f64::from(x) - f64::from(y)).powi(2)) .sum(), + Metric::InnerProduct | Metric::CosineNormalized => -point + .iter() + .zip(target) + .map(|(&x, &y)| f64::from(x) * f64::from(y)) + .sum::(), Metric::Cosine => { - let point_norm = point - .iter() - .map(|&x| f64::from(x).powi(2)) - .sum::() - .sqrt(); - let target_norm = target - .iter() - .map(|&x| f64::from(x).powi(2)) - .sum::() - .sqrt(); - 1.0 - (dot / (point_norm * target_norm)).clamp(-1.0, 1.0) - } - Metric::CosineNormalized | Metric::InnerProduct => -dot, - } - } - - // Keep rank comparisons tight enough to distinguish nearby candidates. - // L2 and inner-product scores grow with dimension. Cosine scores stay unit-scale. - pub(super) fn ranking_tolerance(metric: Metric, dimensions: usize) -> f64 { - let scale = match metric { - Metric::L2 | Metric::InnerProduct => dimensions as f64, - Metric::Cosine | Metric::CosineNormalized => 1.0, - }; - 32.0 * f64::from(f32::EPSILON) * scale - } - - // Bound distance error from the accumulated terms, not the final distance. - // L2 can subtract large dot products from large squared norms. - // This bound applies to the finite, nonzero-norm geometry above. - pub(super) fn distance_tolerance(metric: Metric, point: &[f32], target: &[f32]) -> f64 { - // gamma = n*u/(1-n*u) bounds rounding across n operations. - // Allow norm reduction, Gram accumulation, and the final arithmetic. - let unit_roundoff = f64::from(f32::EPSILON) / 2.0; - let steps = 2.0 * point.len() as f64 + 3.0; - let accumulated_roundoff = steps * unit_roundoff; - let gamma = accumulated_roundoff / (1.0 - accumulated_roundoff); - match metric { - Metric::L2 => { - let terms: f64 = point - .iter() - .zip(target) - .map(|(&x, &y)| (f64::from(x).abs() + f64::from(y).abs()).powi(2)) - .sum(); - gamma * terms - } - Metric::InnerProduct | Metric::CosineNormalized => { - let terms: f64 = point + let dot: f64 = point .iter() .zip(target) - .map(|(&x, &y)| (f64::from(x) * f64::from(y)).abs()) + .map(|(&x, &y)| f64::from(x) * f64::from(y)) .sum(); - gamma * terms - } - Metric::Cosine => { - // Both the dot product and the norm product can carry error. - // Division adds rounding. The final subtraction adds at most 2*u. - (2.0 * gamma + unit_roundoff * (1.0 + gamma)) / (1.0 - gamma) + 2.0 * unit_roundoff + let norm = |row: &[f32]| { + row.iter() + .map(|&x| f64::from(x).powi(2)) + .sum::() + .sqrt() + }; + 1.0 - (dot / (norm(point) * norm(target))).clamp(-1.0, 1.0) } } } - // Compare in output order. An isolated score admits only its expected ID. - // Scores within the rounding bound can exchange IDs, but cannot duplicate one. - #[track_caller] - pub(super) fn assert_ranked_ids(actual: &[u32], sorted: &[(u32, f64)], tolerance: f64) { - let mut previous = f64::NEG_INFINITY; - for (rank, &id) in actual.iter().enumerate() { - assert!( - !actual[..rank].contains(&id), - "duplicate ID {id} at rank {rank}" - ); - let score = sorted + pub(super) fn normalize(values: &mut [f32], dimensions: usize) { + for row in values.chunks_exact_mut(dimensions) { + let norm = row .iter() - .find(|&&(candidate, _)| candidate == id) - .unwrap_or_else(|| panic!("unexpected ID {id} at rank {rank}")) - .1; - let (expected_id, expected_score) = sorted[rank]; - assert!( - (score - expected_score).abs() <= tolerance, - "rank {rank}: ID {id} score {score}, expected ID {expected_id} score {expected_score}, tolerance {tolerance}" - ); - assert!( - previous <= score + tolerance, - "rank {rank}: score {score} follows farther score {previous}" - ); - previous = score; + .map(|&x| f64::from(x).powi(2)) + .sum::() + .sqrt(); + for value in row { + *value = (f64::from(*value) / norm) as f32; + } } } -} - -#[cfg(test)] -mod cosine_distance_tests { - use super::cosine_distance; - use rstest::rstest; - #[rstest] - #[case::zero_source(0.0, 0.0, 1.0)] - #[case::zero_target(0.0, 1.0, 0.0)] - #[case::subnormal_source(f32::MIN_POSITIVE.sqrt() / 2.0, f32::MIN_POSITIVE.sqrt() / 2.0, 1.0)] - #[case::subnormal_target(f32::MIN_POSITIVE.sqrt() / 2.0, 1.0, f32::MIN_POSITIVE.sqrt() / 2.0)] - #[case::zero_norm_before_nan_dot(f32::NAN, 0.0, 1.0)] - fn small_norm_produces_unit_distance( - #[case] dot: f32, - #[case] source_norm: f32, - #[case] target_norm: f32, - ) { - assert_eq!(cosine_distance(dot, source_norm, target_norm), 1.0); + // Put the second coordinate in the last dimension so omitting a tail changes ranking. + pub(super) fn packed_points( + coordinates: &[[f32; 2]], + dimensions: usize, + unit_norm: bool, + ) -> Vec { + let mut values = vec![0.0; coordinates.len() * dimensions]; + for (row, &[x, y]) in values.chunks_exact_mut(dimensions).zip(coordinates) { + row[0] = x; + row[dimensions - 1] = y; + } + if unit_norm { + normalize(&mut values, dimensions); + } + values } #[test] - fn norm_at_cutoff_preserves_positive_similarity() { - // At the cutoff, similarity remains 0.5 instead of zero. - let norm = f32::MIN_POSITIVE.sqrt(); - assert_eq!(cosine_distance(0.5 * norm, norm, 1.0), 0.5); - } + fn dense_fixtures_are_replayable_and_exercise_the_last_dimension() { + let values = dense_points(3, 17, 1287); - #[rstest] - #[case::above_one(1.0 + f32::EPSILON, 0.0)] - #[case::below_negative_one(-1.0 - f32::EPSILON, 2.0)] - fn cosine_clamps_finite_similarity(#[case] similarity: f32, #[case] expected: f32) { - assert_eq!(cosine_distance(similarity, 1.0, 1.0), expected); + assert_eq!(values, dense_points(3, 17, 1287)); + assert_eq!(values.len(), 3 * 17); + assert_ne!(&values[..17], &values[17..34]); + assert_eq!([values[16], values[33], values[50]], [8.0, 9.0, 10.0]); + assert!(values[..16].iter().filter(|&&x| x != 0.0).count() > 8); } - #[rstest] - #[case::nan_dot(f32::NAN, 1.0, 1.0)] - #[case::nan_source_norm(0.0, f32::NAN, 1.0)] - #[case::nan_target_norm(0.0, 1.0, f32::NAN)] - fn nan_propagates_without_small_norms( - #[case] dot: f32, - #[case] source_norm: f32, - #[case] target_norm: f32, + #[rstest::rstest] + #[case::squared_l2(Metric::L2, 18.0)] + #[case::negative_dot(Metric::InnerProduct, -2.0)] + #[case::cosine(Metric::Cosine, 0.7830695421813438)] + fn scalar_reference_matches_hand_calculated_distances( + #[case] metric: Metric, + #[case] expected: f64, ) { - assert!(cosine_distance(dot, source_norm, target_norm).is_nan()); + // [1, 2] and [4, -1]: squared difference 18, dot 2, squared norms 5 and 17. + assert!((distance(metric, &[1.0, 2.0], &[4.0, -1.0]) - expected).abs() < 1.0e-14); + } + + #[test] + fn normalized_points_keep_their_direction_and_last_coordinate() { + let actual = packed_points(&[[3.0, 4.0], [0.0, -2.0]], 3, true); + assert_eq!(actual, [0.6, 0.0, 0.8, 0.0, 0.0, -1.0]); + assert_eq!( + distance(Metric::CosineNormalized, &actual[..3], &actual[3..]), + f64::from(0.8_f32) + ); } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 00507edcca..a7bae7956d 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -97,197 +97,320 @@ fn rank_leader_distances( }); } -// These entry-point tests call GEMM. Miri checks ranking separately. #[cfg(test)] -#[cfg(not(miri))] mod tests { use super::*; - use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2, scalar_ranking}; - use diskann_utils::views::{MatrixView, MutMatrixView}; + #[cfg(not(miri))] + use crate::graph::pipnn::test_support; + use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2}; + #[cfg(not(miri))] use diskann_vector::distance::Metric; - use diskann_wide::arch::{self, Target3}; + use diskann_wide::ARCH; + use rstest::rstest; - struct AssignLeaders(M); + #[cfg(not(miri))] + #[rstest] + #[case::l2(L2, false, [[0,1,3,2], [1,2,3,0], [3,0,2,1]])] + #[case::cosine(Cosine, false, [[0,1,3,2], [1,2,0,3], [3,0,2,1]])] + #[case::normalized_cosine(CosineNormalized, true, [[0,1,3,2], [1,2,0,3], [3,0,2,1]])] + #[case::inner_product(InnerProduct, false, [[0,1,3,2], [1,2,3,0], [0,3,2,1]])] + fn assignments_are_leader_ids_in_metric_order( + #[case] _metric: M, + #[case] unit_norm: bool, + #[case] expected_ids: [[u32; 4]; 3], + #[values(1, 2, 3, 4, 6)] assignments: usize, + #[values(2, 7, 8, 9, 15, 16, 17, 128, 129)] dimensions: usize, + ) { + let point_values = test_support::packed_points( + &[[3.0, 2.0], [-1.0, 3.0], [2.0, -4.0]], + dimensions, + unit_norm, + ); + let leader_values = test_support::packed_points( + &[[4.0, 0.0], [0.0, 3.0], [-2.0, 0.0], [0.0, -1.0]], + dimensions, + unit_norm, + ); + let leaders = M::create_leaders( + MatrixView::try_from(leader_values.as_slice(), 4, dimensions).unwrap(), + ); + let mut output = vec![0; 3 * assignments]; - impl - Target3, MatrixView<'_, f32>, MatrixView<'_, f32>, MutMatrixView<'_, u32>> - for AssignLeaders - { - fn run( - self, - arch: A, - points: MatrixView<'_, f32>, - leader_values: MatrixView<'_, f32>, - output: MutMatrixView<'_, u32>, - ) -> ANNResult<()> { - assign_leaders::<_, M>( - arch, - points, - &M::create_leaders(leader_values), - output, - &mut PartitionKernelWorkspace::default(), - ) + assign_leaders::<_, M>( + ARCH, + MatrixView::try_from(point_values.as_slice(), 3, dimensions).unwrap(), + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap(); + + for point in 0..3 { + let mut expected = expected_ids[point].to_vec(); + expected.truncate(assignments); + expected.resize(assignments, UNASSIGNED_LEADER); + assert_eq!( + &output[point * assignments..(point + 1) * assignments], + expected, + "point={point}" + ); } } - mod assign_leaders_tests { - use super::*; - use rstest::rstest; + #[cfg(not(miri))] + #[rstest] + fn assignments_match_nearest_leaders_across_counts_and_widths( + #[values(1, 5, 16, 17, 33)] leader_count: usize, + #[values(1, 3, 10, 11, 17)] assignments: usize, + ) { + // Queries avoid midpoints, so each leader has a distinct distance. + let leader_values: Vec<_> = (0..leader_count).map(|i| 2.0 * i as f32).collect(); + let point_values = [-1.0, 1.5, 2.0 * (leader_count - 1) as f32 + 0.25]; + let leaders = L2::create_leaders( + MatrixView::try_from(leader_values.as_slice(), leader_count, 1).unwrap(), + ); + let mut output = vec![0; 3 * assignments]; - #[rstest] - #[case::l2(L2, Metric::L2)] - #[case::cosine(Cosine, Metric::Cosine)] - #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] - #[case::inner_product(InnerProduct, Metric::InnerProduct)] - fn leader_ids_match_scalar_ranking( - #[case] metric: M, - #[case] scalar_metric: Metric, - #[values(2, 7, 8, 9, 15, 16, 17, 127, 128, 129, 384, 768)] dimensions: usize, - #[values(1, 3, 10, 11)] width: usize, - ) { - // Given: 33 leaders supply two SIMD groups and a scalar tail. - // Four points lie between different leaders on the same irregular arc. - let normalize = scalar_metric == Metric::CosineNormalized; - let leader_values = - scalar_ranking::arc_vectors((0..33).map(f64::from), dimensions, normalize); - let point_values = - scalar_ranking::arc_vectors([1.25, 10.5, 21.75, 31.25], dimensions, normalize); - let leaders = MatrixView::try_from(leader_values.as_slice(), 33, dimensions).unwrap(); - let points = MatrixView::try_from(point_values.as_slice(), 4, dimensions).unwrap(); - let expected: Vec> = points - .row_iter() - .map(|point| { - let mut row: Vec<_> = leaders - .row_iter() - .enumerate() - .map(|(leader, vector)| { - ( - leader as u32, - scalar_ranking::distance(scalar_metric, point, vector), - ) - }) - .collect(); - row.sort_by(|left, right| left.1.total_cmp(&right.1)); - row - }) + assign_leaders::<_, L2>( + ARCH, + MatrixView::try_from(&point_values[..], 3, 1).unwrap(), + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap(); + + for (point, &coordinate) in point_values.iter().enumerate() { + let mut ordered: Vec<_> = leader_values + .iter() + .enumerate() + .map(|(id, &value)| (id as u32, (coordinate - value).powi(2))) + .collect(); + ordered.sort_by(|left, right| left.1.total_cmp(&right.1)); + let mut expected: Vec<_> = ordered + .iter() + .take(assignments) + .map(|&(id, _)| id) .collect(); - let tolerance = scalar_ranking::ranking_tolerance(scalar_metric, dimensions); - let mut output = vec![UNASSIGNED_LEADER; 4 * width]; + expected.resize(assignments, UNASSIGNED_LEADER); + assert_eq!( + &output[point * assignments..(point + 1) * assignments], + expected, + "point={point}, leader_count={leader_count}, assignments={assignments}" + ); + } + } - // When: create metric-owned leaders and dispatch the complete assignment. - arch::dispatch3_no_features( - AssignLeaders(metric), - points, - leaders, - MutMatrixView::try_from(output.as_mut_slice(), 4, width).unwrap(), - ) - .unwrap(); + #[cfg(not(miri))] + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn large_dense_stripes_select_nearest_leaders( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[values((9, 33, 384), (17, 65, 768), (33, 129, 1536), (5, 35, 1537))] shape: ( + usize, + usize, + usize, + ), + #[values(3, 11)] assignments: usize, + ) { + let (point_count, leader_count, dimensions) = shape; + let mut point_values = test_support::dense_points(point_count, dimensions, 1287); + let mut leader_values = test_support::dense_points(leader_count, dimensions, 2026); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut point_values, dimensions); + test_support::normalize(&mut leader_values, dimensions); + } + let points = + MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); + let leaders = M::create_leaders(leader_matrix); + let mut output = vec![UNASSIGNED_LEADER; point_count * assignments]; + + assign_leaders::<_, M>( + ARCH, + points, + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), point_count, assignments).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap(); - // Then: the full L2 distance has the same order as the omitted-norm score. - for (row, expected) in output.chunks_exact(width).zip(&expected) { - scalar_ranking::assert_ranked_ids(row, expected, tolerance); + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + // Dyadic inputs make the dot/norm sums exact before cosine conversion. + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + }; + for point in 0..point_count { + let mut expected: Vec<_> = (0..leader_count) + .map(|leader| { + ( + leader as u32, + test_support::distance( + scalar_metric, + points.row(point), + leader_matrix.row(leader), + ), + ) + }) + .collect(); + expected.sort_by(|left, right| left.1.total_cmp(&right.1)); + let actual = &output[point * assignments..(point + 1) * assignments]; + for (rank, &id) in actual.iter().enumerate() { + assert!( + !actual[..rank].contains(&id), + "duplicate leader {id} for point={point}" + ); + let score = expected + .iter() + .find(|&&(leader, _)| leader == id) + .unwrap_or_else(|| panic!("invalid leader {id} for point={point}")) + .1; + assert!( + (score - expected[rank].1).abs() <= tolerance, + "shape={shape:?}, point={point}, rank={rank}, score={score}, expected={:?}", + expected[rank] + ); } } + } - #[rstest::rstest] - #[case::missing_row(2)] - #[case::extra_row(4)] - fn invalid_output_rows_return_error(#[case] rows: usize) { - let values = [1.0_f32, 2.0, 3.0]; - let points = MatrixView::try_from(&values[..], 3, 1).unwrap(); - let leaders = Cosine::create_leaders(points); - let mut output = vec![UNASSIGNED_LEADER; rows]; - let mut workspace = PartitionKernelWorkspace::default(); + #[cfg(not(miri))] + #[test] + fn workspace_reuse_replaces_assignments_for_each_stripe() { + let leader_values = [0.0, 5.0, 12.0]; + let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); + let point_values = [1.0, 7.0, 11.0]; + let expected_ids = [[0, 1, 2], [1, 2, 0], [2, 1, 0]]; + let mut output = Vec::new(); + let mut workspace = PartitionKernelWorkspace::default(); - let result = assign_leaders::<_, Cosine>( - diskann_wide::ARCH, - points, + for (point_count, assignments) in [(3, 1), (1, 3), (2, 2), (3, 1)] { + output.resize(point_count * assignments, 0); + assign_leaders::<_, L2>( + ARCH, + MatrixView::try_from(&point_values[..point_count], point_count, 1).unwrap(), &leaders, - MutMatrixView::try_from(output.as_mut_slice(), rows, 1).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), point_count, assignments).unwrap(), &mut workspace, - ); - - assert!(result.is_err()); - } - - #[test] - fn l2_points_select_nearest_leaders() { - // Given: leader j lies at x=j. Seventeen leaders enter one SIMD block and a tail. - // The nearest two integer positions to 15.75, -0.25, and 7.25 are respectively - // (16, 15), (0, 1), and (7, 8); their squared distances are (0.0625, 0.5625), - // (0.0625, 1.5625), and (0.0625, 0.5625). The last point is unrankable. - let leader_values: [f32; 17] = std::array::from_fn(|leader| leader as f32); - let point_values = [15.75, -0.25, 7.25, f32::NAN]; - let expected = [16, 15, 0, 1, 7, 8, UNASSIGNED_LEADER, UNASSIGNED_LEADER]; - let mut output = [0; 8]; - - // When: dispatch reaches the real metric and kernel on the available architecture. - arch::dispatch3_no_features( - AssignLeaders(L2), - MatrixView::try_from(&point_values[..], 4, 1).unwrap(), - MatrixView::try_from(&leader_values[..], 17, 1).unwrap(), - MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), ) .unwrap(); - // Then: each point replaces the previous row, including an unrankable final row. - assert_eq!(output, expected); + for point in 0..point_count { + assert_eq!( + &output[point * assignments..(point + 1) * assignments], + &expected_ids[point][..assignments] + ); + } } + } - #[test] - fn cosine_assignments_update_after_workspace_reuse() { - // Given: unit leaders point right, up, and left. The first point's cosine - // similarities have order right > up > left; the second reverses right and left. - let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; - let leaders = - Cosine::create_leaders(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 = [0, 1, 2, 1]; - let expected_smaller_stripe = [2, 1]; - let regrown_point_values = [-0.8, 0.2, 0.9, 0.1]; - let expected_regrown_stripe = [2, 1, 0, 1]; - let mut output = [UNASSIGNED_LEADER; 4]; - let mut workspace = PartitionKernelWorkspace::default(); + #[test] + fn each_distance_row_replaces_previous_assignments_and_unused_slots() { + let distances = [ + 3.0, + 1.0, + 2.0, + 4.0, + f32::NAN, + f32::INFINITY, + f32::NAN, + f32::INFINITY, + f32::NAN, + 5.0, + f32::INFINITY, + f32::NAN, + ]; + let mut output = [2; 9]; + let mut candidates = vec![Candidate::new(3, -10.0); 4]; - // When - assign_leaders::<_, Cosine>( - diskann_wide::ARCH, - points, - &leaders, - MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), - &mut workspace, - ) - .unwrap(); + rank_leader_distances( + ARCH, + MatrixView::try_from(&distances[..], 3, 4).unwrap(), + MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), + &mut candidates, + ); - // Then - assert_eq!(output, expected); + assert_eq!( + output, + [ + 1, + 2, + 0, + UNASSIGNED_LEADER, + UNASSIGNED_LEADER, + UNASSIGNED_LEADER, + 1, + UNASSIGNED_LEADER, + UNASSIGNED_LEADER + ] + ); + } - // When: shrink to the second point, whose nearest leader differs from the old prefix. - assign_leaders::<_, Cosine>( - diskann_wide::ARCH, - MatrixView::try_from(&point_values[2..], 1, 2).unwrap(), - &leaders, - MutMatrixView::try_from(&mut output[..2], 1, 2).unwrap(), - &mut workspace, - ) - .unwrap(); + #[test] + fn invalid_output_rows_preserve_assignments_and_workspace() { + let point_values = [1.0, 6.0]; + let leader_values = [0.0, 5.0, 12.0]; + let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); + let mut output = [2]; + let previous_candidates = vec![Candidate::new(1, 7.0)]; + let mut workspace = PartitionKernelWorkspace { + distance_scratch: vec![9.0, 11.0], + ranked_leaders: previous_candidates.clone(), + }; - // Then - assert_eq!(output[..2], expected_smaller_stripe); + let error = assign_leaders::<_, L2>( + ARCH, + MatrixView::try_from(&point_values[..], 2, 1).unwrap(), + &leaders, + MutMatrixView::try_from(&mut output[..], 1, 1).unwrap(), + &mut workspace, + ) + .unwrap_err(); - // When: regrow with reversed rows, so the old tail cannot pass unchanged. - assign_leaders::<_, Cosine>( - diskann_wide::ARCH, - MatrixView::try_from(®rown_point_values[..], 2, 2).unwrap(), - &leaders, - MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), - &mut workspace, - ) - .unwrap(); + assert!( + error + .to_string() + .contains("invalid partition output row count 1 for 2 points") + ); + assert_eq!(output, [2]); + assert_eq!(workspace.distance_scratch, [9.0, 11.0]); + assert_eq!(workspace.ranked_leaders, previous_candidates); + } - // Then - assert_eq!(output, expected_regrown_stripe); - } + #[test] + fn metric_dimension_errors_are_returned_without_publishing_assignments() { + let point_values = [1.0, 2.0, 3.0]; + let leader_values = [0.0, 1.0, 2.0, 3.0]; + let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 2, 2).unwrap()); + let mut output = [1]; + + let error = assign_leaders::<_, L2>( + ARCH, + MatrixView::try_from(&point_values[..], 1, 3).unwrap(), + &leaders, + MutMatrixView::try_from(&mut output[..], 1, 1).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap_err(); + + assert_eq!( + error.downcast_ref::(), + Some(&diskann_linalg::SgemmError::InvalidMatrixDimensions { + matrix_name: diskann_linalg::MatrixName::B, + expected_rows: 3, + expected_cols: 2, + actual_len: 4, + }) + ); + assert_eq!(output, [1]); } } diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 96b51b082d..8ee0e67d97 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -224,281 +224,259 @@ impl PartitionMetric for CosineNormalized { #[cfg(test)] mod tests { use super::*; + #[cfg(not(miri))] + use crate::graph::pipnn::test_support; + #[cfg(not(miri))] + use diskann_vector::distance::Metric; + use rstest::rstest; - mod create_leaders_tests { - use super::*; - - #[test] - fn l2_norms_follow_sequential_sum_order() { - // Given: each small square is half an ULP at the first squared norm. - // Each sequential addition rounds back to that norm. - let mut values = [1.0_f32; 129]; - values[0] = 4096.0; - let expected_squared_norm = 16_777_216.0_f32; - let matrix = MatrixView::try_from(&values[..], 1, values.len()).unwrap(); - - // When - let leaders = L2::create_leaders(matrix); - - // Then - assert_eq!(leaders.norms, [expected_squared_norm]); - } - } - - // GEMM-backed tests run natively, not under Miri. #[cfg(not(miri))] - mod compute_distances_tests { - use super::*; - use crate::graph::pipnn::scalar_ranking; - use diskann_vector::distance::Metric; - use rstest::rstest; - - const DIMENSION_COUNT: usize = 2; - const STALE_DISTANCE: f32 = 99.0; - const FLOAT_TOLERANCE: f32 = 1.0e-6; - - fn matrix(values: &[f32], rows: usize) -> MatrixView<'_, f32> { - MatrixView::try_from(values, rows, DIMENSION_COUNT).unwrap() + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn point_to_leader_scores_match_scalar_distances( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[values(1, 3)] point_count: usize, + #[values(1, 4, 17)] leader_count: usize, + #[values(1, 2, 7, 8, 9, 15, 16, 17, 127, 128, 129)] dimensions: usize, + ) { + // Small integer coordinates make unnormalized dot products exact. Points + // and leaders use different values, so swapped rows or columns change scores. + let mut point_values: Vec<_> = (0..point_count * dimensions) + .map(|index| (index % 7) as f32 - 3.0) + .collect(); + let mut leader_values: Vec<_> = (0..leader_count * dimensions) + .map(|index| (index % 11) as f32 - 5.0) + .collect(); + for (point, row) in point_values.chunks_exact_mut(dimensions).enumerate() { + row[0] = point as f32 + 1.0; } - - fn compute_one_ranking( - point: [f32; DIMENSION_COUNT], - leader: [f32; DIMENSION_COUNT], - ) -> f32 { - let leaders = M::create_leaders(matrix(&leader, 1)); - let mut storage = [STALE_DISTANCE]; - - M::compute_distances( - matrix(&point, 1), - &leaders, - MutMatrixView::row_vector(&mut storage[..]), - ) - .unwrap(); - - storage[0] + for (leader, row) in leader_values.chunks_exact_mut(dimensions).enumerate() { + row[0] = leader as f32 + 2.0; } - - #[rstest] - #[case::l2(L2, Metric::L2)] - #[case::cosine(Cosine, Metric::Cosine)] - #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] - #[case::inner_product(InnerProduct, Metric::InnerProduct)] - fn point_leader_matrix_matches_scalar_ranking( - #[case] _metric: M, - #[case] scalar_metric: Metric, - #[values(32, 128)] point_count: usize, - #[values(64, 256)] leader_count: usize, - #[values(128, 384, 768)] dimensions: usize, - ) { - // Given: points and leaders sample the same bounded arc on different grids. - // NaN marks unwritten entries in every point-to-leader row. - let normalize = scalar_metric == Metric::CosineNormalized; - let point_values = scalar_ranking::arc_vectors( - (0..point_count).map(|point| 32.0 * (point as f64 + 0.5) / point_count as f64), - dimensions, - normalize, - ); - let leader_values = scalar_ranking::arc_vectors( - (0..leader_count).map(|leader| 32.0 * leader as f64 / leader_count as f64), - dimensions, - normalize, - ); - let points = - MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); - let leader_vectors = - MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); - let expected: Vec> = points - .row_iter() - .map(|point| { - // Derive the omitted L2 term from input coordinates, not cached norms. - let point_term = if scalar_metric == Metric::L2 { - point.iter().map(|&value| f64::from(value).powi(2)).sum() - } else { - 0.0 - }; - leader_vectors - .row_iter() - .map(|leader| { - ( - scalar_ranking::distance(scalar_metric, point, leader) - point_term, - scalar_ranking::distance_tolerance(scalar_metric, point, leader), - ) - }) - .collect() - }) - .collect(); - let mut distances = vec![f32::NAN; point_count * leader_count]; - - // When: prepare leaders and call the metric without partition selection. - let leaders = M::create_leaders(leader_vectors); - M::compute_distances( - points, - &leaders, - MutMatrixView::try_from(distances.as_mut_slice(), point_count, leader_count) - .unwrap(), - ) - .unwrap(); - - // Then: check every leader column for every point, not only selected leaders. - let actual = - MatrixView::try_from(distances.as_slice(), point_count, leader_count).unwrap(); - for (point, row) in expected.iter().enumerate() { - for (leader, &(expected, tolerance)) in row.iter().enumerate() { - let actual = f64::from(actual.row(point)[leader]); - assert!( - (actual - expected).abs() <= tolerance, - "{scalar_metric:?}, {point_count} points, {leader_count} leaders, D={dimensions}, pair ({point}, {leader}): actual {actual}, expected {expected}, tolerance {tolerance}" - ); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut point_values, dimensions); + test_support::normalize(&mut leader_values, dimensions); + } + let points = + MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); + let leaders = M::create_leaders(leader_matrix); + let mut output = vec![f32::NAN; point_count * leader_count]; + + M::compute_distances( + points, + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), point_count, leader_count).unwrap(), + ) + .unwrap(); + + assert_eq!(M::leader_count(&leaders), leader_count); + for point in 0..point_count { + for leader in 0..leader_count { + let mut expected = test_support::distance( + scalar_metric, + points.row(point), + leader_matrix.row(leader), + ); + if scalar_metric == Metric::L2 { + // Partition L2 omits exactly this constant from every column of the row. + expected -= points + .row(point) + .iter() + .map(|&x| f64::from(x).powi(2)) + .sum::(); } + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + // Bound f32 norm/dot reductions at unit scale, as in the leaf metric. + Metric::Cosine | Metric::CosineNormalized => { + 8.0 * f64::from(f32::EPSILON) * dimensions as f64 + } + }; + let actual = f64::from(output[point * leader_count + leader]); + assert!( + (actual - expected).abs() <= tolerance, + "point={point}, leader={leader}, dimensions={dimensions}: {actual} != {expected}" + ); } } + } - #[rstest::rstest] - #[case::wrong_rows_with_equal_area(1, 4)] - #[case::wrong_columns(2, 1)] - fn output_shape_mismatch_returns_error( - #[values(L2, Cosine, InnerProduct)] _metric: M, - #[case] rows: usize, - #[case] columns: usize, - ) { - // Given: two points and two leaders require a 2-by-2 output. - let leader_values = [1.0, 0.0, 0.0, 2.0]; - let leaders = M::create_leaders(matrix(&leader_values, 2)); - let points = matrix(&[0.0, 4.0, -2.0, 0.0], 2); - let mut storage = vec![STALE_DISTANCE; rows * columns]; - - // When - let result = M::compute_distances( - points, - &leaders, - MutMatrixView::try_from(storage.as_mut_slice(), rows, columns).unwrap(), - ); - - // Then: invalid output shapes stop the build. - assert!(result.is_err()); + #[cfg(not(miri))] + #[rstest] + #[case::l2(L2, Metric::L2)] + #[case::cosine(Cosine, Metric::Cosine)] + #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] + #[case::inner_product(InnerProduct, Metric::InnerProduct)] + fn large_dense_inputs_match_scalar_distances( + #[case] _metric: M, + #[case] scalar_metric: Metric, + #[values((17, 33, 384), (33, 65, 768), (65, 129, 1536), (9, 17, 1537), (257, 513, 129), (5, 17, 4097))] + shape: (usize, usize, usize), + ) { + let (point_count, leader_count, dimensions) = shape; + let mut point_values = test_support::dense_points(point_count, dimensions, 1287); + let mut leader_values = test_support::dense_points(leader_count, dimensions, 2026); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut point_values, dimensions); + test_support::normalize(&mut leader_values, dimensions); } + let points = + MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); + let leaders = M::create_leaders(leader_matrix); + let mut output = vec![f32::NAN; point_count * leader_count]; + + M::compute_distances( + points, + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), point_count, leader_count).unwrap(), + ) + .unwrap(); + + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + // Dyadic dot/norm sums are exact; square roots and division round. + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + // Bound product and sum rounding for normalized f32 coordinates. + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + }; + for point in 0..point_count { + let point_norm: f64 = points + .row(point) + .iter() + .map(|&x| f64::from(x).powi(2)) + .sum(); + for leader in 0..leader_count { + let mut expected = test_support::distance( + scalar_metric, + points.row(point), + leader_matrix.row(leader), + ); + if scalar_metric == Metric::L2 { + expected -= point_norm; + } + let actual = f64::from(output[point * leader_count + leader]); + assert!( + (actual - expected).abs() <= tolerance, + "shape={shape:?}, point={point}, leader={leader}: {actual} != {expected}, tolerance={tolerance}" + ); + } + } + } - #[test] - fn l2_omits_squared_point_norm() { - // Given: leader norms are 25 and 4; each entry is norm - 2 * dot. - let leader_values = [3.0_f32, 4.0, 0.0, 2.0]; - let leaders = L2::create_leaders(matrix(&leader_values, 2)); - let points = matrix(&[2.0, 1.0, -1.0, 2.0], 2); - let expected = [5.0, 0.0, 15.0, -4.0]; - let mut storage = [STALE_DISTANCE; 4]; - - // When - L2::compute_distances( - points, + #[cfg(not(miri))] + #[rstest] + #[case::l2(L2, &[2.0, 0.0, 0.0, 3.0], &[1.0, 0.0, 0.0, -2.0, -3.0, 0.0], [-3.0, 4.0, 21.0, 1.0, 16.0, 9.0])] + #[case::cosine(Cosine, &[2.0, 0.0, 0.0, 3.0], &[1.0, 0.0, 0.0, -2.0, -3.0, 0.0], [0.0, 1.0, 2.0, 1.0, 2.0, 1.0])] + #[case::inner_product(InnerProduct, &[2.0, 0.0, 0.0, 3.0], &[1.0, 0.0, 0.0, -2.0, -3.0, 0.0], [-2.0, 0.0, 6.0, 0.0, 6.0, 0.0])] + #[case::normalized_cosine(CosineNormalized, &[1.0, 0.0, 0.0, 1.0], &[1.0, 0.0, 0.0, -1.0, -1.0, 0.0], [-1.0, 0.0, 1.0, 0.0, 1.0, 0.0])] + fn leader_sets_can_be_reused_across_point_stripes( + #[case] _metric: M, + #[case] point_values: &[f32], + #[case] leader_values: &[f32], + #[case] expected: [f32; 6], + ) { + let leaders = M::create_leaders(MatrixView::try_from(leader_values, 3, 2).unwrap()); + let mut output = [-100.0; 3]; + + for (point, expected) in point_values.chunks_exact(2).zip(expected.chunks_exact(3)) { + M::compute_distances( + MatrixView::try_from(point, 1, 2).unwrap(), &leaders, - MutMatrixView::try_from(&mut storage[..], 2, 2).unwrap(), + MutMatrixView::try_from(&mut output[..], 1, 3).unwrap(), ) .unwrap(); - // Then: every row replaces stale distances before GEMM accumulates. - assert_eq!(storage, expected); - } - - #[rstest::rstest] - #[case::zero_point([0.0, 0.0], [1.0, 0.0])] - #[case::zero_leader([1.0, 0.0], [0.0, 0.0])] - #[case::subnormal_point([f32::MIN_POSITIVE.sqrt() / 2.0, 0.0], [1.0, 0.0])] - #[case::subnormal_leader([1.0, 0.0], [f32::MIN_POSITIVE.sqrt() / 2.0, 0.0])] - fn cosine_writes_one_for_small_norms( - #[case] point: [f32; DIMENSION_COUNT], - #[case] leader: [f32; DIMENSION_COUNT], - ) { - // Given: a zero or subnormal norm represents zero similarity. - let expected = 1.0; - - // When - let actual = compute_one_ranking::(point, leader); - - // Then - assert_eq!(actual, expected); - } - - #[rstest::rstest] - #[case::l2(compute_one_ranking::)] - #[case::cosine(compute_one_ranking::)] - #[case::inner_product(compute_one_ranking::)] - fn nan_coordinate_produces_nan_ranking( - #[case] compute: fn([f32; DIMENSION_COUNT], [f32; DIMENSION_COUNT]) -> f32, - ) { - // Given: neither vector has zero norm. - let point = [f32::NAN, 1.0]; - let leader = [1.0, 0.0]; - - // When - let actual = compute(point, leader); - - // Then - assert!(actual.is_nan()); - } - - #[test] - fn normalized_cosine_writes_negative_dot_products() { - // Given - let point = [1.0_f32, 0.0]; - let leader = [0.6_f32, 0.8]; - let dot = point[0].mul_add(leader[0], point[1] * leader[1]); - let expected = -dot; - - // When - let actual = compute_one_ranking::(point, leader); - - // Then - assert_eq!(actual, expected); + assert_eq!(output, expected); } + } - #[test] - fn inner_product_writes_negative_dot_products() { - // Given - let point = [2.0_f32, -1.0]; - let leader = [3.0_f32, 4.0]; - let dot = point[0].mul_add(leader[0], point[1] * leader[1]); - let expected = -dot; - - // When - let actual = compute_one_ranking::(point, leader); + #[cfg(not(miri))] + #[test] + fn cosine_assigns_unit_distance_when_either_vector_has_zero_norm() { + let point_values = [0.0, 0.0, 0.0, 2.0]; + let leader_values = [3.0, 0.0, 0.0, 0.0, 0.0, -4.0]; + let leaders = + Cosine::create_leaders(MatrixView::try_from(&leader_values[..], 3, 2).unwrap()); + let mut output = [42.0; 6]; + + Cosine::compute_distances( + MatrixView::try_from(&point_values[..], 2, 2).unwrap(), + &leaders, + MutMatrixView::try_from(&mut output[..], 2, 3).unwrap(), + ) + .unwrap(); - // Then - assert_eq!(actual, expected); - } + assert_eq!(output, [1.0, 1.0, 1.0, 1.0, 1.0, 2.0]); + } - #[test] - fn cosine_replaces_all_distance_rows() { - // Given: non-unit points point right, up, and left. Leaders point - // diagonally up-right and up, so similarities include ±1/sqrt(2). - let leader_values = [1.0_f32, 1.0, 0.0, 2.0]; - let leaders = Cosine::create_leaders(matrix(&leader_values, 2)); - let points = matrix(&[2.0, 0.0, 0.0, 4.0, -2.0, 0.0], 3); - let diagonal_similarity = std::f32::consts::FRAC_1_SQRT_2; - let expected = [ - 1.0 - diagonal_similarity, - 1.0, - 1.0 - diagonal_similarity, - 0.0, - 1.0 + diagonal_similarity, - 1.0, - ]; - let mut storage = [STALE_DISTANCE; 6]; - - // When - Cosine::compute_distances( - points, - &leaders, - MutMatrixView::try_from(&mut storage[..], 3, 2).unwrap(), - ) - .unwrap(); + #[rstest] + #[case::l2(L2)] + #[case::cosine(Cosine)] + #[case::normalized_cosine(CosineNormalized)] + #[case::inner_product(InnerProduct)] + fn output_shape_mismatch_is_rejected_before_writing( + #[case] _metric: M, + #[values((1, 3), (2, 2), (3, 3), (2, 4))] shape: (usize, usize), + ) { + let point_values = [1.0, 2.0, 3.0, 4.0]; + let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; + let leaders = M::create_leaders(MatrixView::try_from(&leader_values[..], 3, 2).unwrap()); + let mut output = vec![17.0; shape.0 * shape.1]; + + let error = M::compute_distances( + MatrixView::try_from(&point_values[..], 2, 2).unwrap(), + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), shape.0, shape.1).unwrap(), + ) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("point-to-leader output shape mismatch") + ); + assert_eq!(output, vec![17.0; shape.0 * shape.1]); + } - // Then - for (actual, expected) in storage.into_iter().zip(expected) { - assert!( - (actual - expected).abs() <= FLOAT_TOLERANCE, - "actual {actual} differs from expected {expected}" - ); - } - } + #[rstest] + #[case::l2(L2)] + #[case::cosine(Cosine)] + #[case::normalized_cosine(CosineNormalized)] + #[case::inner_product(InnerProduct)] + fn mismatched_vector_dimensions_report_the_leader_matrix_error( + #[case] _metric: M, + ) { + let point_values = [1.0, 2.0, 3.0]; + let leader_values = [1.0, 2.0, 3.0, 4.0]; + let leaders = M::create_leaders(MatrixView::try_from(&leader_values[..], 2, 2).unwrap()); + let mut output = [17.0; 2]; + + let error = M::compute_distances( + MatrixView::try_from(&point_values[..], 1, 3).unwrap(), + &leaders, + MutMatrixView::try_from(&mut output[..], 1, 2).unwrap(), + ) + .unwrap_err(); + + assert_eq!( + error.downcast_ref::(), + Some(&diskann_linalg::SgemmError::InvalidMatrixDimensions { + matrix_name: diskann_linalg::MatrixName::B, + expected_rows: 3, + expected_cols: 2, + actual_len: 4, + }) + ); } } diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 7a1eba9a29..d50611d4df 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -442,699 +442,383 @@ fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate, widt #[cfg(test)] mod tests { use super::*; - use diskann_wide::arch::{self, Target1}; + use diskann_wide::{ARCH, arch::Current}; use rstest::rstest; - // These adapters only select the real CPU backend; all updates use production TopK. - struct SelectTopK; - impl Target1, &[f32], &mut [Candidate])> - for SelectTopK - { - fn run(self, arch: A, (topk, distances, output): (&TopK, &[f32], &mut [Candidate])) { - topk.select_topk(arch, distances, output); - } - } - - struct UpdateDualTopK; - impl - Target1< - A, - (), - ( - &TopK, - usize, - &[f32], - MutMatrixView<'_, Candidate>, - &mut [f32], - ), - > for UpdateDualTopK - { - fn run( - self, - arch: A, - (topk, point, distances, output, thresholds): ( - &TopK, - usize, - &[f32], - MutMatrixView<'_, Candidate>, - &mut [f32], - ), - ) { - topk.update_dual_topk(arch, point, distances, output, thresholds); - } - } + const EMPTY: Candidate = Candidate::new(UNASSIGNED, f32::INFINITY); + const LANES: usize = <::Vector as SIMDVector>::LANES; - fn sorted_candidates(distances: &[f32], width: usize) -> Vec { - let mut candidates: Vec<_> = distances - .iter() - .enumerate() - .filter(|(_, distance)| **distance < f32::INFINITY) - .map(|(id, &distance)| Candidate::new(id as u32, distance)) - .collect(); - candidates.sort_by(|left, right| left.distance.partial_cmp(&right.distance).unwrap()); - candidates.truncate(width); - candidates.resize(width, Candidate::default()); - candidates - } - - #[test] - fn initialize_resets_results_and_limits() { - // Given: candidate storage contains old results; the threshold buffer - // still has the smaller previous leaf's length. - let mut output = [ - Candidate::new(1, 1.0), - Candidate::new(2, 4.0), - Candidate::new(0, 1.0), - Candidate::new(2, 3.0), - Candidate::new(1, 3.0), - Candidate::new(0, 4.0), - ]; - let mut thresholds = vec![4.0, 3.0]; - let expected_output = [Candidate::default(); 6]; - let expected_thresholds = [f32::INFINITY; 3]; + #[rstest] + #[case::first_batch(0, 3)] + #[case::larger_batch(2, 5)] + #[case::smaller_batch(5, 2)] + #[case::same_size_batch(3, 3)] + #[case::empty_batch(3, 0)] + fn starting_a_batch_clears_results_and_sizes_its_limits( + #[case] previous_rows: usize, + #[case] rows: usize, + ) { + // Given: reusable storage contains assigned candidates and finite limits. let topk = TopK::new(Fixed::<2>); + let mut output = vec![[Candidate::new(0, -12.0), Candidate::new(1, -6.0)]; rows]; + let mut limits = vec![-6.0; previous_rows]; - // When + // When: prepare storage for a new batch. topk.initialize( - MutMatrixView::try_from(&mut output[..], 3, 2).unwrap(), - &mut thresholds, + MutMatrixView::try_from(output.as_flattened_mut(), rows, 2).unwrap(), + &mut limits, ); - // Then: clear both old thresholds and the newly added slot. - assert_eq!(output, expected_output); - assert_eq!(thresholds, expected_thresholds); + // Then: no old candidate or limit belongs to the new batch. + assert_eq!(output, vec![[EMPTY; 2]; rows]); + assert_eq!(limits, vec![f32::INFINITY; rows]); } #[rstest] - #[case::fixed(Fixed::<0>)] - #[case::runtime(Runtime(0))] - fn zero_capacity_preserves_results_and_limits(#[case] width: impl Width) { - let mut output = []; - let mut thresholds = [f32::INFINITY; 2]; - let topk = TopK::new(width); - - arch::dispatch1_no_features(SelectTopK, (&topk, &[1.0, 2.0][..], &mut output[..])); - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &topk, - 1, - &[0.0][..], - MutMatrixView::try_from(&mut output[..], 2, 0).unwrap(), - &mut thresholds[..], - ), - ); + #[case::empty_input(&[], [EMPTY; 3])] + #[case::no_rankable_values(&[f32::INFINITY, f32::NAN], [EMPTY; 3])] + #[case::fewer_candidates_than_slots( + &[11.0, -4.0], + [Candidate::new(1, -4.0), Candidate::new(0, 11.0), EMPTY], + )] + #[case::already_nearest_first( + &[-7.0, 2.0, 8.0, 13.0], + [Candidate::new(0, -7.0), Candidate::new(1, 2.0), Candidate::new(2, 8.0)], + )] + #[case::replacements_and_reordering( + &[8.0, -7.0, 13.0, 2.0], + [Candidate::new(1, -7.0), Candidate::new(3, 2.0), Candidate::new(0, 8.0)], + )] + #[case::unrankable_values_leave_gaps_in_ids( + &[f32::NAN, 6.0, f32::INFINITY, -2.0, 1.0], + [Candidate::new(3, -2.0), Candidate::new(4, 1.0), Candidate::new(1, 6.0)], + )] + fn selection_returns_nearest_input_positions_and_clears_unused_slots( + #[case] distances: &[f32], + #[case] expected: [Candidate; 3], + ) { + // Given: output still holds results from a different input. + let topk = TopK::new(Fixed::<3>); + let mut output = [ + Candidate::new(2, -20.0), + Candidate::new(4, -10.0), + Candidate::new(9, -1.0), + ]; - assert!(output.is_empty()); - assert_eq!(thresholds, [f32::INFINITY; 2]); - } + topk.select_topk(ARCH, distances, &mut output); - mod insert_sorted_tests { - use super::*; - - #[rstest] - #[case::fixed(Fixed::<3>)] - #[case::runtime(Runtime(3))] - fn insertion_preserves_other_rows(#[case] width: impl Width) { - // Given: the new candidate displaces the first result's farthest neighbor. - let first_result = [ - Candidate::new(0, 1.0), - Candidate::new(1, 4.0), - Candidate::new(2, 6.0), - ]; - let next_result = [ - Candidate::new(3, 8.0), - Candidate::new(4, 10.0), - Candidate::new(5, 12.0), - ]; - let candidate = Candidate::new(6, 3.0); - let expected_first_result = [ - Candidate::new(0, 1.0), - Candidate::new(6, 3.0), - Candidate::new(1, 4.0), - ]; - let expected_limit = 4.0; - let mut results = [first_result, next_result]; - - // When - let limit = insert_sorted(results.as_flattened_mut(), candidate, width); - - // Then - assert_eq!(results, [expected_first_result, next_result]); - assert_eq!(limit, expected_limit); - } + assert_eq!(output, expected); } - mod select_topk_tests { - use super::*; - - #[rstest] - #[case::before_first_simd_block(15)] - #[case::one_simd_block(16)] - #[case::one_simd_block_and_tail(17)] - #[case::before_second_simd_block(31)] - #[case::two_simd_blocks(32)] - #[case::two_simd_blocks_and_tail(33)] - #[case::three_simd_blocks_and_tail(49)] - fn dispatched_selection_matches_independent_sort( - #[case] count: usize, - #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)] width: usize, - ) { - // Given: alternating signs force replacement and interior insertion. - let distances: Vec<_> = (0..count) - .map(|i| if i % 2 == 0 { i as f32 } else { -(i as f32) }) - .collect(); - let expected = sorted_candidates(&distances, width); - let mut output = vec![Candidate::default(); width]; - let mut runtime_output = output.clone(); - - // When - with_topk!(width, |topk| { - arch::dispatch1_no_features( - SelectTopK, - (&topk, distances.as_slice(), output.as_mut_slice()), - ); - }); - arch::dispatch1_no_features( - SelectTopK, - ( - &TopK::new(Runtime(width)), - distances.as_slice(), - runtime_output.as_mut_slice(), - ), - ); - - // Then: both width representations must satisfy the independent oracle. - assert_eq!(output, expected); - assert_eq!(runtime_output, expected); - } - - #[rstest] - #[case::before_second_simd_block(31)] - #[case::two_simd_blocks(32)] - #[case::two_simd_blocks_and_tail(33)] - #[case::three_simd_blocks_and_tail(49)] - fn runtime_selection_matches_independent_sort( - #[case] count: usize, - #[values(11, 17)] width: usize, - ) { - // Given: K=17 fills during the second block. The third block starts - // with a finite limit. Each input has more candidates than result slots. - let distances: Vec<_> = (0..count) - .map(|i| if i % 2 == 0 { i as f32 } else { -(i as f32) }) - .collect(); - let expected = sorted_candidates(&distances, width); - let mut output = vec![Candidate::default(); width]; - - // When: these capacities select the runtime branch of the batch dispatch. - with_topk!(width, |topk| { - arch::dispatch1_no_features( - SelectTopK, - (&topk, distances.as_slice(), output.as_mut_slice()), - ); - }); - - // Then - assert_eq!(output, expected); - } - - #[rstest] - #[case::unrankable(&[f32::NAN, f32::INFINITY, f32::INFINITY], [Candidate::default(); 2])] - #[case::partly_filled( - &[f32::NAN, f32::INFINITY, 9.0], - [Candidate::new(2, 9.0), Candidate::default()] - )] - fn selection_replaces_previous_results( - #[case] distances: &[f32], - #[case] expected: [Candidate; 2], - ) { - // Given: both slots contain candidates from the previous selection. - let mut output = [Candidate::new(1, 1.0), Candidate::new(2, 2.0)]; - let topk = TopK::new(Fixed::<2>); - - // When: select again with the same output buffer. - arch::dispatch1_no_features(SelectTopK, (&topk, distances, &mut output[..])); - - // Then: unused slots contain no candidates from the previous selection. - assert_eq!(output, expected); - } + #[test] + fn reusing_output_starts_a_selection_with_a_fresh_distance_limit() { + let topk = TopK::new(Fixed::<2>); + let mut output = [EMPTY; 2]; + topk.select_topk(ARCH, &[-20.0, -10.0], &mut output); + assert_eq!(output, [Candidate::new(0, -20.0), Candidate::new(1, -10.0)]); - #[test] - fn selection_rechecks_candidates_against_updated_limit() { - // Group one fills K=1 at 5. Both 1 and 4 pass group two's initial - // mask, but inserting 1 must prevent the later 4 from replacing it. - let mut distances = [f32::INFINITY; 32]; - distances[0] = 5.0; - distances[16] = 1.0; - distances[17] = 4.0; - let mut output = [Candidate::default()]; - let topk = TopK::new(Fixed::<1>); + // Every new distance exceeds the previous selection's limit. + topk.select_topk(ARCH, &[16.0, 3.0, 9.0], &mut output); - arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); + assert_eq!(output, [Candidate::new(1, 3.0), Candidate::new(2, 9.0)]); + } - assert_eq!(output, [Candidate::new(16, 1.0)]); + #[rstest] + fn batch_selection_matches_a_full_sort( + #[values(0, 1, LANES - 1, LANES, LANES + 1, 2 * LANES, 2 * LANES + 3, 3 * LANES + 2)] + count: usize, + #[values(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, LANES + 1)] capacity: usize, + ) { + // Each pair offers a nearer score before a farther one. Later pairs improve + // on earlier pairs, so a full result must keep lowering its cutoff mid-vector. + let mut distances: Vec<_> = (0..count).map(|i| -(i as f32) - 1.0).collect(); + for pair in distances.chunks_exact_mut(2) { + pair.swap(0, 1); } - #[test] - fn middle_insertions_keep_nearest_first_order() { - let mut output = [Candidate::default(); 3]; - let topk = TopK::new(Fixed::<3>); - - // After retaining [0, 4, 6], distances 2 and 3 must enter between - // existing candidates instead of replacing the last slot directly. - arch::dispatch1_no_features( - SelectTopK, - (&topk, &[0.0, 4.0, 6.0, 2.0, 3.0][..], &mut output[..]), - ); - - assert_eq!( - output, - [ - Candidate::new(0, 0.0), - Candidate::new(3, 2.0), - Candidate::new(4, 3.0) - ] - ); - } + // Sorting the whole input is independent of TopK's insertion and cutoff logic. + let mut expected: Vec<_> = distances + .iter() + .enumerate() + .map(|(index, &distance)| Candidate::new(index as u32, distance)) + .collect(); + expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); + expected.truncate(capacity); + expected.resize(capacity, EMPTY); + let mut output = vec![EMPTY; capacity]; + let mut capacity_reads = 0; + + with_topk!( + { + capacity_reads += 1; + capacity + }, + |topk| { + topk.select_topk(ARCH, &distances, &mut output); + } + ); - #[test] - fn non_rankable_block_preserves_input_indices() { - // The first block contributes no candidates. Later IDs still refer - // to positions in the full input, not positions among rankable values. - let mut distances = [f32::INFINITY; 33]; - distances[..8].fill(f32::NAN); - distances[16] = 4.0; - distances[17] = 2.0; - distances[31] = 3.0; - distances[32] = 1.0; - let mut output = [Candidate::default(); 3]; - let topk = TopK::new(Fixed::<3>); - - arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); + assert_eq!(output, expected, "count={count}, capacity={capacity}"); + assert_eq!(capacity_reads, 1); + } - assert_eq!( - output, - [ - Candidate::new(32, 1.0), - Candidate::new(17, 2.0), - Candidate::new(31, 3.0), - ] - ); - } + #[test] + fn selection_returns_distinct_candidates_when_distances_tie() { + let best_index = 2 * LANES + 2; + let mut distances = vec![8.0; best_index + 1]; + distances[best_index] = -2.0; + let mut output = [EMPTY; 4]; + + TopK::new(Fixed::<4>).select_topk(ARCH, &distances, &mut output); + + assert_eq!(output[0], Candidate::new(best_index as u32, -2.0)); + // Any three distinct input positions with distance 8 are valid; tie order is free. + let ties = &output[1..]; + assert!( + ties.iter() + .all(|c| c.local_idx < best_index as u32 && c.distance == 8.0), + "invalid tied candidates: {ties:?}" + ); + let mut ids: Vec<_> = ties.iter().map(|c| c.local_idx).collect(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids.len(), 3, "a candidate must not occupy multiple slots"); + } - #[test] - fn tail_candidate_replaces_a_tied_neighbor() { - let mut distances = [2.0; 17]; - distances[16] = 1.0; - let mut output = [Candidate::default(); 3]; - let topk = TopK::new(Fixed::<3>); + #[rstest] + #[case::negative_zero(-0.0)] + #[case::positive_zero(0.0)] + #[case::lowest_finite(f32::MIN)] + #[case::highest_finite(f32::MAX)] + #[case::negative_infinity(f32::NEG_INFINITY)] + fn a_rankable_value_keeps_its_input_id_and_float_bits( + #[case] distance: f32, + #[values(0, LANES - 1, LANES, 2 * LANES + 1)] index: usize, + ) { + let mut distances = vec![f32::NAN; 2 * LANES + 2]; + distances[index] = distance; + let mut output = [EMPTY; 2]; - arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); + TopK::new(Runtime(2)).select_topk(ARCH, &distances, &mut output); - assert_eq!(output[0], Candidate::new(16, 1.0)); - let tied = &output[1..]; - assert!( - tied.iter() - .all(|candidate| candidate.local_idx < 16 && candidate.distance == 2.0) - ); - assert_ne!(tied[0].local_idx, tied[1].local_idx); - } + assert_eq!(output[0].local_idx, index as u32); + assert_eq!(output[0].distance.to_bits(), distance.to_bits()); + assert_eq!(output[1], EMPTY); + } - #[test] - fn selection_preserves_ids_and_float_bits() { - let mut distances = [f32::INFINITY; 33]; - distances[0] = -0.0; - distances[16] = 0.0; - distances[17] = f32::MAX; - distances[32] = f32::NEG_INFINITY; - let mut output = [Candidate::default(); 4]; - let topk = TopK::new(Runtime(4)); - - arch::dispatch1_no_features(SelectTopK, (&topk, &distances[..], &mut output[..])); - - assert_eq!(output[0], Candidate::new(32, f32::NEG_INFINITY)); - // Tie order is unspecified; each retained value must match its own ID. - let mut actual: Vec<_> = output - .iter() - .map(|candidate| (candidate.local_idx, candidate.distance.to_bits())) - .collect(); - actual.sort_unstable_by_key(|&(id, _)| id); - assert_eq!( - actual, - [ - (0, (-0.0_f32).to_bits()), - (16, 0.0_f32.to_bits()), - (17, f32::MAX.to_bits()), - (32, f32::NEG_INFINITY.to_bits()) - ] - ); - } + #[rstest] + #[case::fixed(Fixed::<2>)] + #[case::runtime(Runtime(2))] + fn successive_pair_rows_accumulate_each_points_nearest_neighbors(#[case] width: impl Width) { + let topk = TopK::new(width); + let mut output = [[EMPTY; 2]; 4]; + let mut limits = Vec::new(); + let mut rows = MutMatrixView::try_from(output.as_flattened_mut(), 4, 2).unwrap(); + topk.initialize(rows.as_mut_view(), &mut limits); + + // Given: pairs 0-1=9, 0-2=2, and 1-2=7; point 3 has no offered pairs. + topk.update_dual_topk(ARCH, 1, &[9.0], rows.as_mut_view(), &mut limits); + topk.update_dual_topk(ARCH, 2, &[2.0, 7.0], rows.as_mut_view(), &mut limits); + let before_last_point = [ + [Candidate::new(2, 2.0), Candidate::new(1, 9.0)], + [Candidate::new(2, 7.0), Candidate::new(0, 9.0)], + [Candidate::new(0, 2.0), Candidate::new(1, 7.0)], + [EMPTY; 2], + ]; + assert_eq!(rows.as_slice(), before_last_point.as_flattened()); + assert_eq!(limits, [9.0, 9.0, 7.0, f32::INFINITY]); + + // When: point 3 is at distances 6, 1, and 4 from points 0, 1, and 2. + topk.update_dual_topk(ARCH, 3, &[6.0, 1.0, 4.0], rows.as_mut_view(), &mut limits); + + // Then: each endpoint keeps its own nearest two, with its farthest retained limit. + let expected = [ + [Candidate::new(2, 2.0), Candidate::new(3, 6.0)], + [Candidate::new(3, 1.0), Candidate::new(2, 7.0)], + [Candidate::new(0, 2.0), Candidate::new(3, 4.0)], + [Candidate::new(1, 1.0), Candidate::new(2, 4.0)], + ]; + assert_eq!(rows.as_slice(), expected.as_flattened()); + assert_eq!(limits, [6.0, 7.0, 4.0, 4.0]); } - mod update_dual_topk_tests { - use super::*; - - // The lower triangle contains unique scores. The upper triangle mirrors them. - // Shuffle the scores so matrix order differs from distance order. - // Diagonal values cannot rank, so no point selects itself. - fn shuffled_symmetric_distances(count: usize) -> Vec { - use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; - - let mut scores: Vec<_> = (1..=count * (count - 1) / 2).map(|i| i as f32).collect(); - scores.shuffle(&mut StdRng::seed_from_u64(1287)); - let mut distances = vec![f32::INFINITY; count * count]; - for source in 1..count { - for target in 0..source { - let score = scores.pop().unwrap(); - distances[source * count + target] = score; - distances[target * count + source] = score; - } + #[rstest] + fn pair_scans_match_full_row_sort_across_capacities_and_lengths( + #[values(LANES, LANES + 1, LANES + 2, 2 * LANES, 2 * LANES + 1, 2 * LANES + 2)] + point_count: usize, + #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17)] capacity: usize, + ) { + // Given: XOR gives symmetric distances with distinct, exact scores in each row. + // The final pair row has point_count - 1 distances, straddling vector boundaries. + let distances: Vec> = (0..point_count) + .map(|point| { + (0..point_count) + .map(|other| (point ^ other) as f32) + .collect() + }) + .collect(); + let mut output = vec![EMPTY; point_count * capacity]; + let mut limits = Vec::new(); + let mut rows = + MutMatrixView::try_from(output.as_mut_slice(), point_count, capacity).unwrap(); + + // When: offer every non-self pair once, through the production width dispatch. + with_topk!(capacity, |topk| { + topk.initialize(rows.as_mut_view(), &mut limits); + for (point, row) in distances.iter().enumerate().skip(1) { + topk.update_dual_topk(ARCH, point, &row[..point], rows.as_mut_view(), &mut limits); } - distances - } + }); - // Happy path: every point retains K nearest neighbors from finite distances. - #[rstest] - #[case::one_simd_block_and_tail(18)] - #[case::two_simd_blocks(33)] - #[case::two_simd_blocks_and_tail(34)] - #[case::three_simd_blocks_and_tail(50)] - fn update_dual_topk_selects_nearest_neighbors( - #[case] count: usize, - #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17)] width: usize, - ) { - // Given: the oracle sorts complete rows, independent of the update order. - // Every row has enough non-self candidates to fill its result. - let distances = shuffled_symmetric_distances(count); - let expected: Vec<_> = distances - .chunks_exact(count) - .map(|row| sorted_candidates(row, width)) + // Then: independently sort each complete row, excluding only the point itself. + for (point, row) in distances.iter().enumerate() { + let mut expected: Vec<_> = row + .iter() + .enumerate() + .filter(|(other, _)| *other != point) + .map(|(other, &distance)| Candidate::new(other as u32, distance)) .collect(); - let mut output = vec![Candidate::default(); count * width]; - let mut thresholds = Vec::new(); - let mut rows = MutMatrixView::try_from(output.as_mut_slice(), count, width).unwrap(); - - // When: each distance updates the two points' neighbor lists once. - with_topk!(width, |topk| { - topk.initialize(rows.as_mut_view(), &mut thresholds); - for source in 1..count { - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &topk, - source, - &distances[source * count..source * count + source], - rows.as_mut_view(), - thresholds.as_mut_slice(), - ), - ); - } - }); - - // Then: check all source and reciprocal rows with their persisted limits. - for (source, expected) in expected.iter().enumerate() { - assert_eq!( - rows.row(source), - expected, - "source={source}, N={count}, K={width}" - ); - assert_eq!( - thresholds[source], - expected[width - 1].distance, - "threshold for source={source}, N={count}, K={width}" - ); - } - } - - #[test] - fn farther_points_do_not_replace_closer_neighbors() { - // Given: points 0 and 1 select each other at distance 1. - // Point 2 has no neighbors. Its distances to points 0 and 1 are 2 and 3. - let mut output = [ - Candidate::new(1, 1.0), - Candidate::new(0, 1.0), - Candidate::default(), - ]; - let mut thresholds = output.map(|candidate| candidate.distance); - let expected = [output[0], output[1], Candidate::new(0, 2.0)]; - - // When - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &TopK::new(Fixed::<1>), - 2, - &[2.0, 3.0][..], - MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), - &mut thresholds[..], - ), - ); - - // Then - assert_eq!(output, expected); - assert_eq!(thresholds, expected.map(|candidate| candidate.distance)); - } - - #[test] - fn closer_points_replace_existing_neighbors() { - // Given: points 0 and 1 select each other at distance 3. - // Point 2 is closer to both: distance 1 from point 0 and 2 from point 1. - // Points 0 and 1 must select point 2. Point 2 must select point 0. - let mut output = [ - Candidate::new(1, 3.0), - Candidate::new(0, 3.0), - Candidate::default(), - ]; - let mut thresholds = output.map(|candidate| candidate.distance); - let expected = [ - Candidate::new(2, 1.0), - Candidate::new(2, 2.0), - Candidate::new(0, 1.0), - ]; - - // When - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &TopK::new(Fixed::<1>), - 2, - &[1.0, 2.0][..], - MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), - &mut thresholds[..], - ), - ); - - // Then - assert_eq!(output, expected); - assert_eq!(thresholds, expected.map(|candidate| candidate.distance)); - } - - #[test] - fn both_points_prefer_the_closer_neighbor() { - // Given: points 0 and 1 select each other at distance 1. - // Point 2 is 0.5 from point 0 and 2 from point 1. - // Points 1 and 2 must each select point 0, not each other. - let mut output = [ - Candidate::new(1, 1.0), - Candidate::new(0, 1.0), - Candidate::default(), - ]; - let mut thresholds = output.map(|candidate| candidate.distance); - let expected = [Candidate::new(2, 0.5), output[1], Candidate::new(0, 0.5)]; - - // When - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &TopK::new(Fixed::<1>), - 2, - &[0.5, 2.0][..], - MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), - &mut thresholds[..], - ), - ); - - // Then - assert_eq!(output, expected); - assert_eq!(thresholds, expected.map(|candidate| candidate.distance)); - } - - #[test] - fn later_blocks_fill_source_neighbor_slots() { - let mut output = [Candidate::default(); 34 * 3]; - let mut thresholds = [f32::INFINITY; 34]; - let topk = TopK::new(Fixed::<3>); - let mut rows = MutMatrixView::try_from(&mut output[..], 34, 3).unwrap(); - let mut distances = [f32::INFINITY; 33]; - distances[..8].fill(f32::NAN); - distances[15] = 4.0; - distances[16] = 2.0; - distances[17] = 9.0; - distances[31] = 3.0; - distances[32] = f32::NEG_INFINITY; - - // A leaf source starts empty. It remains partially filled after - // group one and acquires its finite threshold within group two. - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &topk, - 33, - &distances[..], - rows.as_mut_view(), - &mut thresholds[..], - ), - ); + expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); + expected.truncate(capacity); + expected.resize(capacity, EMPTY); assert_eq!( - rows.row(33), - [ - Candidate::new(32, f32::NEG_INFINITY), - Candidate::new(16, 2.0), - Candidate::new(31, 3.0) - ] + rows.row(point), + expected, + "point={point}, point_count={point_count}, capacity={capacity}" ); assert_eq!( - rows.row(17), - [ - Candidate::new(33, 9.0), - Candidate::default(), - Candidate::default() - ] + limits[point], + expected[capacity - 1].distance, + "limit for point={point}, point_count={point_count}, capacity={capacity}" ); } + } - #[test] - fn later_updates_respect_saved_distance_limits() { - let mut output = [Candidate::default(); 3]; - let mut thresholds = [f32::INFINITY; 3]; - let topk = TopK::new(Fixed::<1>); - let mut rows = MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(); - - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &topk, - 1, - &[1.0][..], - rows.as_mut_view(), - &mut thresholds[..], - ), - ); - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &topk, - 2, - &[3.0, 2.0][..], - rows.as_mut_view(), - &mut thresholds[..], - ), - ); + #[rstest] + fn a_source_rejection_does_not_prevent_the_target_from_accepting( + #[values(1, LANES - 1, LANES, 2 * LANES + 1)] target: usize, + ) { + let source = 2 * LANES + 2; + let topk = TopK::new(Fixed::<1>); + let mut output = vec![EMPTY; source + 1]; + let mut limits = Vec::new(); + let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); + topk.initialize(rows.as_mut_view(), &mut limits); + let mut distances = vec![f32::INFINITY; source]; + distances[0] = -6.0; + distances[target] = 4.0; + + // The source fills its only slot with point 0 before reaching this target. + topk.update_dual_topk(ARCH, source, &distances, rows.as_mut_view(), &mut limits); + + let mut expected = vec![EMPTY; source + 1]; + expected[0] = Candidate::new(source as u32, -6.0); + expected[target] = Candidate::new(source as u32, 4.0); + expected[source] = Candidate::new(0, -6.0); + let expected_limits: Vec<_> = expected.iter().map(|c| c.distance).collect(); + assert_eq!(rows.as_slice(), expected, "target={target}"); + assert_eq!(limits, expected_limits); + } - assert_eq!( - rows.as_slice(), - [ - Candidate::new(1, 1.0), - Candidate::new(0, 1.0), - Candidate::new(1, 2.0) - ] - ); - } + #[rstest] + fn a_target_rejection_does_not_prevent_the_source_from_accepting( + #[values(1, LANES - 1, LANES, 2 * LANES + 1)] target: usize, + ) { + let source = 2 * LANES + 2; + let topk = TopK::new(Fixed::<1>); + let mut output = vec![EMPTY; source + 1]; + let mut limits = Vec::new(); + let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); + topk.initialize(rows.as_mut_view(), &mut limits); + let mut previous_distances = vec![f32::INFINITY; target]; + previous_distances[0] = 1.0; + topk.update_dual_topk( + ARCH, + target, + &previous_distances, + rows.as_mut_view(), + &mut limits, + ); + let mut distances = vec![f32::INFINITY; source]; + distances[target] = 5.0; + + // The target already has a closer neighbor; the source still needs one. + topk.update_dual_topk(ARCH, source, &distances, rows.as_mut_view(), &mut limits); + + let mut expected = vec![EMPTY; source + 1]; + expected[0] = Candidate::new(target as u32, 1.0); + expected[target] = Candidate::new(0, 1.0); + expected[source] = Candidate::new(target as u32, 5.0); + let expected_limits: Vec<_> = expected.iter().map(|c| c.distance).collect(); + assert_eq!(rows.as_slice(), expected, "target={target}"); + assert_eq!(limits, expected_limits); + } - #[test] - fn target_replaces_its_farthest_neighbor() { - // Given: point 15 retains points 0, 1, 2, and 3. It is the last - // lane in group one. Point 16 is the tail target and has no neighbors. - // Point 17 enters the middle of point 15's full row. - let mut output = [Candidate::default(); 18 * 4]; - let mut thresholds = [f32::INFINITY; 18]; - let mut previous = [f32::INFINITY; 15]; - previous[..4].copy_from_slice(&[-1.0, 3.0, 5.0, 7.0]); - let mut distances = [f32::INFINITY; 17]; - distances[0] = 4.0; - distances[2] = 3.0; - distances[15] = 2.0; - distances[16] = 1.0; - let expected_source = [ - Candidate::new(16, distances[16]), - Candidate::new(15, distances[15]), - Candidate::new(2, distances[2]), - Candidate::new(0, distances[0]), - ]; - let expected_prior_row = [ - Candidate::new(0, previous[0]), - Candidate::new(17, distances[15]), - Candidate::new(1, previous[1]), - Candidate::new(2, previous[2]), - ]; - let expected_tail_row = [ - Candidate::new(17, distances[16]), - Candidate::default(), - Candidate::default(), - Candidate::default(), - ]; - - let mut rows = MutMatrixView::try_from(&mut output[..], 18, 4).unwrap(); - let topk = TopK::new(Runtime(4)); - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &topk, - 15, - &previous[..], - rows.as_mut_view(), - &mut thresholds[..], - ), - ); + #[rstest] + #[case::finite(12.0)] + #[case::negative_infinity(f32::NEG_INFINITY)] + #[case::negative_zero(-0.0)] + fn a_single_pair_updates_exactly_its_two_endpoints( + #[case] distance: f32, + #[values(LANES - 1, LANES, 2 * LANES + 1)] target: usize, + ) { + let source = 2 * LANES + 2; + let topk = TopK::new(Fixed::<2>); + let mut output = vec![EMPTY; (source + 1) * 2]; + let mut limits = Vec::new(); + let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 2).unwrap(); + topk.initialize(rows.as_mut_view(), &mut limits); + let mut distances = vec![f32::INFINITY; source]; + distances[target] = distance; + + topk.update_dual_topk(ARCH, source, &distances, rows.as_mut_view(), &mut limits); + + let mut expected = vec![EMPTY; (source + 1) * 2]; + expected[source * 2] = Candidate::new(target as u32, distance); + expected[target * 2] = Candidate::new(source as u32, distance); + assert_eq!(rows.as_slice(), expected); + assert_eq!(rows.row(source)[0].distance.to_bits(), distance.to_bits()); + assert_eq!(rows.row(target)[0].distance.to_bits(), distance.to_bits()); + // One neighbor leaves each two-slot result open to another candidate. + assert_eq!(limits, vec![f32::INFINITY; source + 1]); + } - // When - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &topk, - 17, - &distances[..], - rows.as_mut_view(), - &mut thresholds[..], - ), - ); + #[rstest] + #[case::no_pairs(&[])] + #[case::nan_pairs(&[f32::NAN; 2 * LANES + 1])] + #[case::infinite_pairs(&[f32::INFINITY; 2 * LANES + 1])] + fn an_update_without_rankable_pairs_preserves_existing_results(#[case] distances: &[f32]) { + let source = 2 * LANES + 1; + let topk = TopK::new(Fixed::<1>); + let mut output = vec![EMPTY; source + 1]; + let mut limits = Vec::new(); + let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); + topk.initialize(rows.as_mut_view(), &mut limits); + topk.update_dual_topk(ARCH, 1, &[3.0], rows.as_mut_view(), &mut limits); + let previous_output = rows.as_slice().to_vec(); + let previous_limits = limits.clone(); + + topk.update_dual_topk(ARCH, source, distances, rows.as_mut_view(), &mut limits); + + assert_eq!(rows.as_slice(), previous_output); + assert_eq!(limits, previous_limits); + } - // Then: the source fills, the prior row evicts, and the tail stays underfilled. - assert_eq!(rows.row(17), expected_source); - assert_eq!(rows.row(15), expected_prior_row); - assert_eq!(rows.row(16), expected_tail_row); - assert_eq!(thresholds[17], expected_source[3].distance); - assert_eq!(thresholds[15], expected_prior_row[3].distance); - assert_eq!(thresholds[16], f32::INFINITY); - } + #[rstest] + #[case::fixed(Fixed::<0>)] + #[case::runtime(Runtime(0))] + fn zero_capacity_pair_scans_leave_no_neighbors_or_finite_limits(#[case] width: impl Width) { + let topk = TopK::new(width); + let mut output = []; + let mut limits = Vec::new(); + let mut rows = MutMatrixView::try_from(&mut output[..], 3, 0).unwrap(); + topk.initialize(rows.as_mut_view(), &mut limits); - #[rstest] - #[case::simd(15)] - #[case::tail(16)] - fn both_points_accept_negative_infinity(#[case] candidate: usize) { - let mut distances = [f32::NAN; 17]; - distances[candidate] = f32::NEG_INFINITY; - let mut output = [Candidate::default(); 18 * 2]; - let mut thresholds = [f32::INFINITY; 18]; - let topk = TopK::new(Fixed::<2>); - let mut rows = MutMatrixView::try_from(&mut output[..], 18, 2).unwrap(); - let mut expected = [[Candidate::default(); 2]; 18]; - expected[17][0] = Candidate::new(candidate as u32, f32::NEG_INFINITY); - expected[candidate][0] = Candidate::new(17, f32::NEG_INFINITY); - - arch::dispatch1_no_features( - UpdateDualTopK, - ( - &topk, - 17, - &distances[..], - rows.as_mut_view(), - &mut thresholds[..], - ), - ); + topk.update_dual_topk(ARCH, 1, &[4.0], rows.as_mut_view(), &mut limits); + topk.update_dual_topk(ARCH, 2, &[2.0, 7.0], rows.as_mut_view(), &mut limits); - assert_eq!(rows.as_slice(), expected.as_flattened()); - } + assert!(rows.as_slice().is_empty()); + assert_eq!(limits, [f32::INFINITY; 3]); } } From 4690241b17ededf08f42e133ab6ac0f6de5ca76f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 18 Sep 2026 07:06:08 +0000 Subject: [PATCH 185/229] refactor(pipnn): dispatch TopK with a visitor --- diskann/src/graph/pipnn/leaf_kernel.rs | 40 ++++-- diskann/src/graph/pipnn/partition_kernel.rs | 34 ++++- diskann/src/graph/pipnn/topk.rs | 147 ++++++++++---------- 3 files changed, 127 insertions(+), 94 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 2bc32454fa..c9b37230d1 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -14,7 +14,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ leaf_metric::LeafMetric, simd::PiPNNSIMDSchema, - topk::{Candidate, with_topk}, + topk::{Candidate, TopK, TopKVisitor, Width, with_topk}, }; /// Reusable storage for one leaf numerical pipeline. @@ -89,21 +89,41 @@ where fn rank_leaf_distances( arch: A, distances: MatrixView<'_, f32>, - mut output: MutMatrixView<'_, Candidate>, + output: MutMatrixView<'_, Candidate>, worst: &mut Vec, ) { - with_topk!(output.ncols(), |topk| { - topk.initialize(output.as_mut_view(), worst); - for point_idx in 1..distances.nrows() { + with_topk( + output.ncols(), + RankLeaf { + arch, + distances, + output, + thresholds: worst, + }, + ); +} + +struct RankLeaf<'a, A> { + arch: A, + distances: MatrixView<'a, f32>, + output: MutMatrixView<'a, Candidate>, + thresholds: &'a mut Vec, +} + +impl TopKVisitor for RankLeaf<'_, A> { + #[inline] + fn visit(mut self, topk: TopK) { + topk.initialize(self.output.as_mut_view(), self.thresholds); + for point_idx in 1..self.distances.nrows() { topk.update_dual_topk( - arch, + self.arch, point_idx, - &distances.row(point_idx)[..point_idx], - output.as_mut_view(), - worst.as_mut_slice(), + &self.distances.row(point_idx)[..point_idx], + self.output.as_mut_view(), + self.thresholds.as_mut_slice(), ); } - }); + } } /// Check for one output row per point and a valid non-self neighbor count. diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index a7bae7956d..e16a56d976 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -15,7 +15,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ partition_metric::PartitionMetric, simd::PiPNNSIMDSchema, - topk::{Candidate, UNASSIGNED, with_topk}, + topk::{Candidate, TopK, TopKVisitor, UNASSIGNED, Width, with_topk}, }; /// No sampled partition center was rankable for this output slot. @@ -83,18 +83,38 @@ where fn rank_leader_distances( arch: A, distances: MatrixView<'_, f32>, - mut output: MutMatrixView<'_, u32>, + output: MutMatrixView<'_, u32>, candidates: &mut Vec, ) { candidates.resize(output.ncols(), Candidate::default()); - with_topk!(output.ncols(), |topk| { - for (distances, output) in distances.row_iter().zip(output.row_iter_mut()) { - topk.select_topk(arch, distances, candidates.as_mut_slice()); - for (destination, candidate) in output.iter_mut().zip(candidates.iter()) { + with_topk( + output.ncols(), + RankLeaders { + arch, + distances, + output, + candidates, + }, + ); +} + +struct RankLeaders<'a, A> { + arch: A, + distances: MatrixView<'a, f32>, + output: MutMatrixView<'a, u32>, + candidates: &'a mut Vec, +} + +impl TopKVisitor for RankLeaders<'_, A> { + #[inline] + fn visit(mut self, topk: TopK) { + for (distances, output) in self.distances.row_iter().zip(self.output.row_iter_mut()) { + topk.select_topk(self.arch, distances, self.candidates.as_mut_slice()); + for (destination, candidate) in output.iter_mut().zip(self.candidates.iter()) { *destination = candidate.local_idx; } } - }); + } } #[cfg(test)] diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index d50611d4df..abb8d9fc14 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -67,81 +67,39 @@ impl Width for Runtime { } } -/// Select the width specialization once for a batch. -/// Specialize small capacities 1..=10; larger capacities use the runtime path. +/// A batch operation that accepts any TopK width specialization. +pub(super) trait TopKVisitor { + /// Consume the operation and run it with the selected width. + fn visit(self, topk: TopK); +} + +/// Select the width specialization and run the visitor once for the batch. +/// +/// Specialize capacities 1..=10; zero and larger capacities use the runtime path. /// Typical partition K is 10 or 3, and leaf K is 3 or 2. -/// The width is evaluated once. The body has ordinary block control flow. -macro_rules! with_topk { - ($width:expr, |$topk:ident| $body:block) => {{ - let width = $width; - match width { - 1 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<1>); - $body - } - 2 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<2>); - $body - } - 3 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<3>); - $body - } - 4 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<4>); - $body - } - 5 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<5>); - $body - } - 6 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<6>); - $body - } - 7 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<7>); - $body - } - 8 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<8>); - $body - } - 9 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<9>); - $body - } - 10 => { - let $topk = - $crate::graph::pipnn::topk::TopK::new($crate::graph::pipnn::topk::Fixed::<10>); - $body - } - _ => { - let $topk = $crate::graph::pipnn::topk::TopK::new( - $crate::graph::pipnn::topk::Runtime(width), - ); - $body - } - } - }}; +#[inline] +pub(super) fn with_topk(width: usize, visitor: V) { + match width { + 1 => visitor.visit(TopK::new(Fixed::<1>)), + 2 => visitor.visit(TopK::new(Fixed::<2>)), + 3 => visitor.visit(TopK::new(Fixed::<3>)), + 4 => visitor.visit(TopK::new(Fixed::<4>)), + 5 => visitor.visit(TopK::new(Fixed::<5>)), + 6 => visitor.visit(TopK::new(Fixed::<6>)), + 7 => visitor.visit(TopK::new(Fixed::<7>)), + 8 => visitor.visit(TopK::new(Fixed::<8>)), + 9 => visitor.visit(TopK::new(Fixed::<9>)), + 10 => visitor.visit(TopK::new(Fixed::<10>)), + _ => visitor.visit(TopK::new(Runtime(width))), + } } -pub(super) use with_topk; /// Top-k selection with a fixed result capacity and caller-owned storage. /// /// This object stores no candidates or thresholds. Each operation borrows only /// the buffers it updates. Candidate IDs are slice positions, not dataset IDs. /// Callers supply each candidate at most once per result; equal distances need -/// no fixed tie order. The construction macro specializes the width per batch. +/// no fixed tie order. [`with_topk`] specializes the width per batch. /// Selection and dual updates run inside their own architecture scope. pub(super) struct TopK { width: W, @@ -529,6 +487,17 @@ mod tests { count: usize, #[values(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, LANES + 1)] capacity: usize, ) { + struct SelectRow<'a> { + distances: &'a [f32], + output: &'a mut [Candidate], + } + + impl TopKVisitor for SelectRow<'_> { + fn visit(self, topk: TopK) { + topk.select_topk(ARCH, self.distances, self.output); + } + } + // Each pair offers a nearer score before a farther one. Later pairs improve // on earlier pairs, so a full result must keep lowering its cutoff mid-vector. let mut distances: Vec<_> = (0..count).map(|i| -(i as f32) - 1.0).collect(); @@ -548,14 +517,15 @@ mod tests { let mut output = vec![EMPTY; capacity]; let mut capacity_reads = 0; - with_topk!( + with_topk( { capacity_reads += 1; capacity }, - |topk| { - topk.select_topk(ARCH, &distances, &mut output); - } + SelectRow { + distances: &distances, + output: &mut output, + }, ); assert_eq!(output, expected, "count={count}, capacity={capacity}"); @@ -648,6 +618,27 @@ mod tests { point_count: usize, #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17)] capacity: usize, ) { + struct ScanPairs<'a> { + distances: &'a [Vec], + output: MutMatrixView<'a, Candidate>, + limits: &'a mut Vec, + } + + impl TopKVisitor for ScanPairs<'_> { + fn visit(mut self, topk: TopK) { + topk.initialize(self.output.as_mut_view(), self.limits); + for (point, row) in self.distances.iter().enumerate().skip(1) { + topk.update_dual_topk( + ARCH, + point, + &row[..point], + self.output.as_mut_view(), + self.limits, + ); + } + } + } + // Given: XOR gives symmetric distances with distinct, exact scores in each row. // The final pair row has point_count - 1 distances, straddling vector boundaries. let distances: Vec> = (0..point_count) @@ -663,12 +654,14 @@ mod tests { MutMatrixView::try_from(output.as_mut_slice(), point_count, capacity).unwrap(); // When: offer every non-self pair once, through the production width dispatch. - with_topk!(capacity, |topk| { - topk.initialize(rows.as_mut_view(), &mut limits); - for (point, row) in distances.iter().enumerate().skip(1) { - topk.update_dual_topk(ARCH, point, &row[..point], rows.as_mut_view(), &mut limits); - } - }); + with_topk( + capacity, + ScanPairs { + distances: &distances, + output: rows.as_mut_view(), + limits: &mut limits, + }, + ); // Then: independently sort each complete row, excluding only the point itself. for (point, row) in distances.iter().enumerate() { From 34722b263b7e46445dc5ec3c35f04b31d50a7d77 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 18 Sep 2026 07:20:28 +0000 Subject: [PATCH 186/229] refactor(pipnn): reuse fast vector norms --- diskann/src/graph/pipnn/leaf_metric.rs | 37 +++++++++++++++-- diskann/src/graph/pipnn/partition_metric.rs | 46 ++++++++++++++++++--- 2 files changed, 73 insertions(+), 10 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index 275382a40d..bcbbc374d2 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -10,6 +10,10 @@ use crate::{ANNError, ANNResult}; use diskann_utils::views::MatrixView; +use diskann_vector::{ + Norm, + norm::{FastL2Norm, FastL2NormSquared}, +}; use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; @@ -36,7 +40,7 @@ impl LeafMetric for L2 { // The expanded L2 formula is `||x||² + ||y||² - 2(x·y)`. let squared_norms: Vec = points .row_iter() - .map(|point| point.iter().map(|value| value * value).sum()) + .map(|point| FastL2NormSquared.evaluate(point)) .collect(); // Initialize the first two terms before GEMM adds the dot-product term. for source in 0..point_count { @@ -63,7 +67,6 @@ impl LeafMetric for L2 { impl LeafMetric for Cosine { fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { let point_count = points.nrows(); - // The diagonal supplies each point norm after GEMM computes all dots. diskann_linalg::sgemm_aat_lower( point_count, points.ncols(), @@ -72,8 +75,9 @@ impl LeafMetric for Cosine { storage, ) .map_err(ANNError::new)?; - let norms: Vec = (0..point_count) - .map(|point| storage[point * point_count + point].sqrt()) + let norms: Vec = points + .row_iter() + .map(|point| FastL2Norm.evaluate(point)) .collect(); // Convert each lower-triangle dot to the bounded cosine distance. for source in 0..point_count { @@ -116,6 +120,31 @@ mod tests { use diskann_vector::distance::Metric; use rstest::rstest; + #[cfg(not(miri))] + #[test] + fn l2_distance_retains_small_coordinate_contributions() { + // Given: 4096^2 + 127 unit coordinates, orthogonal to a unit vector. + // Scalar left-to-right f32 summation loses every unit after 4096^2. + let mut values = vec![1.0; 2 * 129]; + values[0] = 4096.0; + values[128] = 0.0; + values[129..].fill(0.0); + values[257] = 1.0; + let points = MatrixView::try_from(values.as_slice(), 2, 129).unwrap(); + let mut output = [f32::NAN; 4]; + let expected = 16_777_344.0; // 4096^2 + 128. + + L2::compute_distances(points, &mut output).unwrap(); + + // Allow eight f32 ULPs at this scale, far less than the lost 128. + let tolerance = 16.0; + assert!( + (output[2] - expected).abs() <= tolerance, + "{} != {expected}", + output[2] + ); + } + #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2)] diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index 8ee0e67d97..dbc623a856 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -11,7 +11,10 @@ use crate::{ANNError, ANNResult}; use diskann_linalg::Transpose; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_vector::{Norm, norm::FastL2NormSquared}; +use diskann_vector::{ + Norm, + norm::{FastL2Norm, FastL2NormSquared}, +}; use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; @@ -55,13 +58,11 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { ) -> ANNResult<()>; } -/// Compute L2 squared norms with sequential accumulation. -/// -/// This order fixes rounding of the leader term before GEMM adds the dot term. +/// Compute squared norms with DiskANN's vector implementation. fn l2_squared_norms(vectors: MatrixView<'_, f32>) -> Vec { vectors .row_iter() - .map(|vector| vector.iter().map(|value| value * value).sum()) + .map(|vector| FastL2NormSquared.evaluate(vector)) .collect() } @@ -69,7 +70,7 @@ fn l2_squared_norms(vectors: MatrixView<'_, f32>) -> Vec { fn cosine_norms(vectors: MatrixView<'_, f32>) -> Vec { vectors .row_iter() - .map(|vector| FastL2NormSquared.evaluate(vector).sqrt()) + .map(|vector| FastL2Norm.evaluate(vector)) .collect() } @@ -230,6 +231,39 @@ mod tests { use diskann_vector::distance::Metric; use rstest::rstest; + #[cfg(not(miri))] + #[test] + fn l2_ranking_retains_small_coordinate_contributions() { + // Given: the origin ranks leaders solely by their squared norms. + // 4096^2 + 127 is larger than 4096^2 + 16, regardless of coordinate order. + let points = [0.0; 129]; + let mut values = vec![1.0; 2 * 129]; + values[0] = 4096.0; + values[128] = 0.0; + values[129..].fill(0.0); + values[129] = 4096.0; + values[130] = 4.0; + let leaders = L2::create_leaders(MatrixView::try_from(values.as_slice(), 2, 129).unwrap()); + let mut output = [f32::NAN; 2]; + + L2::compute_distances( + MatrixView::try_from(&points[..], 1, 129).unwrap(), + &leaders, + MutMatrixView::try_from(&mut output[..], 1, 2).unwrap(), + ) + .unwrap(); + + assert!( + output[1] < output[0], + "nearer leader was ranked behind farther leader: {output:?}" + ); + assert!( + (output[0] - 16_777_344.0).abs() <= 16.0, + "lost small squared coordinates: {output:?}" + ); + assert_eq!(output[1], 16_777_232.0); + } + #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2)] From 4c97081ece529b4c2f6eaef99a9d010ade9a5f91 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 18 Sep 2026 07:43:55 +0000 Subject: [PATCH 187/229] test(linalg): share lower-product assertions --- diskann-linalg/src/lib.rs | 90 +++++++++++++++++---------------------- 1 file changed, 39 insertions(+), 51 deletions(-) diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index 8f98ef5b99..89bb31d398 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -726,6 +726,43 @@ mod sgemm_aat_lower_tests { type LowerProduct = fn(usize, usize, f32, &[f32], &mut [f32]) -> Result<(), SgemmError>; + #[track_caller] + fn assert_lower_product_matches_scalar( + shape: (usize, usize), + input: &[f32], + previous: &[f32], + output: &[f32], + add: bool, + ) { + let (rows, dimensions) = shape; + for row in 0..rows { + for column in 0..rows { + let index = row * rows + column; + if column > row { + assert_eq!( + output[index].to_bits(), + previous[index].to_bits(), + "shape={shape:?}, upper ({row},{column})" + ); + } else { + let dot: f64 = (0..dimensions) + .map(|dimension| { + f64::from(input[row * dimensions + dimension]) + * f64::from(input[column * dimensions + dimension]) + }) + .sum(); + // Both sweeps use alpha = -2 and exactly representable arithmetic. + let expected = -2.0 * dot + if add { f64::from(previous[index]) } else { 0.0 }; + assert_eq!( + f64::from(output[index]), + expected, + "shape={shape:?}, lower ({row},{column})" + ); + } + } + } + } + #[rstest] #[case::replace(sgemm_aat_lower, [-8.0, 101.0, 102.0, 8.0, -26.0, 103.0, 12.0, -36.0, -50.0])] #[case::add(sgemm_aat_lower_add, [2.0, 101.0, 102.0, 28.0, 4.0, 103.0, 52.0, 14.0, 10.0])] @@ -759,31 +796,7 @@ mod sgemm_aat_lower_tests { operation(rows, dimensions, -2.0, &input, &mut output).unwrap(); - for row in 0..rows { - for column in 0..rows { - let index = row * rows + column; - if column > row { - assert_eq!( - output[index].to_bits(), - previous[index].to_bits(), - "upper ({row},{column})" - ); - } else { - let dot: f64 = (0..dimensions) - .map(|d| { - f64::from(input[row * dimensions + d]) - * f64::from(input[column * dimensions + d]) - }) - .sum(); - let expected = -2.0 * dot + if add { f64::from(previous[index]) } else { 0.0 }; - assert_eq!( - f64::from(output[index]), - expected, - "lower ({row},{column}), dimensions={dimensions}" - ); - } - } - } + assert_lower_product_matches_scalar((rows, dimensions), &input, &previous, &output, add); } #[rstest] @@ -813,32 +826,7 @@ mod sgemm_aat_lower_tests { operation(rows, dimensions, -2.0, &input, &mut output).unwrap(); - // Check every entry, including the untouched upper triangle. - for row in 0..rows { - for column in 0..rows { - let index = row * rows + column; - if column > row { - assert_eq!( - output[index].to_bits(), - previous[index].to_bits(), - "shape={shape:?}, upper ({row},{column})" - ); - } else { - let dot: f64 = (0..dimensions) - .map(|dimension| { - f64::from(input[row * dimensions + dimension]) - * f64::from(input[column * dimensions + dimension]) - }) - .sum(); - let expected = -2.0 * dot + if add { f64::from(previous[index]) } else { 0.0 }; - assert_eq!( - f64::from(output[index]), - expected, - "shape={shape:?}, lower ({row},{column})" - ); - } - } - } + assert_lower_product_matches_scalar(shape, &input, &previous, &output, add); } #[test] From 5b853e06508625dae59284d4ec8bde6c39fda1a0 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 18 Sep 2026 07:49:57 +0000 Subject: [PATCH 188/229] test(linalg): trim redundant GEMM shapes --- diskann-linalg/src/lib.rs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index 89bb31d398..ae452709c6 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -780,13 +780,17 @@ mod sgemm_aat_lower_tests { } #[rstest] - #[case::replace(sgemm_aat_lower, false)] - #[case::add(sgemm_aat_lower_add, true)] + #[case::replace_single_entry(sgemm_aat_lower, false, 1, 1)] + #[case::add_single_entry(sgemm_aat_lower_add, true, 1, 1)] + #[case::replace_zero_inner_dimension(sgemm_aat_lower, false, 3, 0)] + #[case::add_zero_inner_dimension(sgemm_aat_lower_add, true, 3, 0)] + #[case::replace_rectangular(sgemm_aat_lower, false, 3, 7)] + #[case::add_rectangular(sgemm_aat_lower_add, true, 3, 7)] fn rectangular_inputs_match_scalar_dot_products( #[case] operation: LowerProduct, #[case] add: bool, - #[values(1, 3, 17)] rows: usize, - #[values(0, 1, 7, 8, 9, 16, 17, 33)] dimensions: usize, + #[case] rows: usize, + #[case] dimensions: usize, ) { let input: Vec<_> = (0..rows * dimensions) .map(|i| (i % 9) as f32 - 4.0) @@ -805,13 +809,11 @@ mod sgemm_aat_lower_tests { fn large_dense_inputs_match_scalar_dot_products( #[case] operation: LowerProduct, #[case] add: bool, - #[values((33, 384), (65, 768), (129, 1536), (17, 1537), (513, 129), (17, 4097))] shape: ( - usize, - usize, - ), ) { use rand::{rngs::StdRng, SeedableRng}; + // One representative dense input checks the wrapper's row-major wiring. + let shape = (129, 1536); let (rows, dimensions) = shape; let mut rng = StdRng::seed_from_u64(1287); // Bounded multiples of 1/8 keep products and sums exact in f32. From 9660547d621ff62c25ca7b30de66579d0bf4a673 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 18 Sep 2026 08:02:54 +0000 Subject: [PATCH 189/229] refactor(pipnn): defer matrix view hardening PiPNN callers already bound or check their matrix shapes. Keep the shared constructor change and its tests outside this PR. --- diskann-utils/src/views.rs | 49 +------------------------------------- 1 file changed, 1 insertion(+), 48 deletions(-) diff --git a/diskann-utils/src/views.rs b/diskann-utils/src/views.rs index eb7cfe5d44..a9352918c9 100644 --- a/diskann-utils/src/views.rs +++ b/diskann-utils/src/views.rs @@ -211,7 +211,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 nrows.checked_mul(ncols) != Some(len) { + if len != nrows * ncols { Err(TryFromError { data, nrows, ncols }) } else { Ok(Self { data, nrows, ncols }) @@ -731,53 +731,6 @@ where // Tests // /////////// -#[cfg(test)] -mod matrix_shape_tests { - use super::*; - use rstest::rstest; - - #[rstest] - #[case::wraps_to_empty(usize::MAX / 2 + 1, 2, &[])] - #[case::wraps_to_one(usize::MAX, usize::MAX, &[7])] - fn overflowing_shape_is_rejected_even_when_wrapping_matches_the_buffer( - #[case] rows: usize, - #[case] columns: usize, - #[case] values: &[u32], - ) { - let data = values.to_vec().into_boxed_slice(); - - let error = Matrix::try_from(data, rows, columns).unwrap_err(); - - assert_eq!((error.nrows, error.ncols), (rows, columns)); - assert_eq!(error.into_inner().as_ref(), values); - } - - #[test] - fn rejected_mutable_view_returns_the_original_buffer() { - let mut values = [11_u32]; - - let error = MutMatrixView::try_from(&mut values[..], usize::MAX, usize::MAX).unwrap_err(); - - let returned = error.into_inner(); - assert_eq!(returned, [11]); - returned[0] = 13; - assert_eq!(values, [13]); - } - - #[rstest] - #[case::zero_rows(0, usize::MAX)] - #[case::zero_columns(usize::MAX, 0)] - fn zero_sized_shapes_remain_valid_with_a_large_other_dimension( - #[case] rows: usize, - #[case] columns: usize, - ) { - let matrix = MatrixView::try_from(&[] as &[u32], rows, columns).unwrap(); - - assert_eq!((matrix.nrows(), matrix.ncols()), (rows, columns)); - assert!(matrix.as_slice().is_empty()); - } -} - #[cfg(test)] mod tests { use super::*; From 166ac2523f36ba1b3b601e23765c9a66ec14943b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 18 Sep 2026 08:19:52 +0000 Subject: [PATCH 190/229] fix(pipnn): expect unwired kernels in PR2 The merge with main enables clippy::allow_attributes. Use one scoped dead-code expectation until graph construction wires the kernels. --- diskann/src/graph/pipnn/mod.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 0dff853d77..6a0b167ffb 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -8,17 +8,16 @@ //! Metric modules fill portable distance buffers. Kernel modules use architecture //! `A` to traverse those buffers. Callers select metric `M` once per graph build. -#[allow(dead_code)] +#![expect( + dead_code, + reason = "graph construction integrates these kernels in the next PR" +)] + mod leaf_kernel; -#[allow(dead_code)] mod leaf_metric; -#[allow(dead_code)] mod partition_kernel; -#[allow(dead_code)] mod partition_metric; -#[allow(dead_code)] mod simd; -#[allow(dead_code)] mod topk; pub(super) struct L2; From b59ea3229db11109b168e02d8697ddbe5147ba17 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:38:44 +0000 Subject: [PATCH 191/229] test(pipnn): rewrite core contract tests --- diskann/src/graph/pipnn/conversion.rs | 106 ++- diskann/src/graph/pipnn/finalization.rs | 258 +++---- diskann/src/graph/pipnn/leaf_build.rs | 339 ++++++---- diskann/src/graph/pipnn/mod.rs | 749 ++++++++++---------- diskann/src/graph/pipnn/partitioning.rs | 863 ++++++++++-------------- 5 files changed, 1159 insertions(+), 1156 deletions(-) diff --git a/diskann/src/graph/pipnn/conversion.rs b/diskann/src/graph/pipnn/conversion.rs index f8482ba0c6..cca4ef42fb 100644 --- a/diskann/src/graph/pipnn/conversion.rs +++ b/diskann/src/graph/pipnn/conversion.rs @@ -50,37 +50,87 @@ where } #[cfg(test)] -mod gather_as_f32_tests { +mod tests { use super::*; + use rstest::rstest; + + #[rstest] + #[case::f32([1.5_f32, -2.0, 7.0, 4.0, 0.5, 9.0], [0.5, 9.0, 1.5, -2.0, 0.5, 9.0])] + #[case::i8([1_i8, -2, 7, 4, -5, 9], [-5.0, 9.0, 1.0, -2.0, -5.0, 9.0])] + #[case::u8([1_u8, 2, 7, 4, 5, 9], [5.0, 9.0, 1.0, 2.0, 5.0, 9.0])] + #[case::f16([1.5_f32, -2.0, 7.0, 4.0, 0.5, 9.0].map(f16::from_f32), [0.5, 9.0, 1.5, -2.0, 0.5, 9.0])] + fn gathered_rows_follow_requested_ids( + #[case] values: [T; 6], + #[case] expected: [f32; 6], + ) { + // Given: row IDs are deliberately out of order and contain a repeat. + let data = MatrixView::try_from(&values[..], 3, 2).unwrap(); + let mut output = [f32::NAN; 6]; + + gather_as_f32(data, &[2, 0, 2], &mut output).unwrap(); + + assert_eq!(output, expected); + } + + #[rstest] + #[case::before_vector(7)] + #[case::complete_vector(8)] + #[case::after_vector(9)] + #[case::multiple_vectors(32)] + #[case::multiple_vectors_and_tail(33)] + #[case::embedding_and_tail(1537)] + fn fp16_gather_converts_every_coordinate(#[case] dimensions: usize) { + let mut values: Vec<_> = (0..3 * dimensions) + .map(|index| f16::from_f32((index % 31) as f32 - 15.0)) + .collect(); + for (row, values) in values.chunks_exact_mut(dimensions).enumerate() { + values[dimensions - 1] = f16::from_f32(100.0 + row as f32); + } + let ids = [2, 0, 2]; + let expected: Vec<_> = ids + .iter() + .flat_map(|&row| { + values[row as usize * dimensions..(row as usize + 1) * dimensions] + .iter() + .map(|value| value.to_f32()) + }) + .collect(); + let mut output = vec![f32::NAN; expected.len()]; + + gather_as_f32( + MatrixView::try_from(values.as_slice(), 3, dimensions).unwrap(), + &ids, + &mut output, + ) + .unwrap(); + + assert_eq!(output, expected); + } #[test] - fn fp16_gather_preserves_bits_and_requested_row_order_with_a_tail() { - // Given: repeated, noncontiguous rows of nine elements exercise tails - // and the row stride while overwriting stale destination values. - let points = [ - [-4.0_f32, -3.5, -2.0, -1.5, -0.0, 0.5, 1.0, 2.5, 3.0], - [1.0, 2.5, 3.0, 0.5, -1.0, 4.0, -2.5, 0.0, -3.5], - [8.0, 2.0, -7.0, 0.25, 0.75, 3.5, -4.5, 0.0, -2.0], + fn fp16_gather_preserves_special_values_and_zero_signs() { + let values = [ + f16::ZERO, + f16::NEG_ZERO, + f16::INFINITY, + f16::NEG_INFINITY, + f16::from_bits(1), + f16::NAN, ]; - let half_points = points.map(|row| row.map(f16::from_f32_const)); - let data = MatrixView::try_from(half_points.as_flattened(), 3, 9).unwrap(); - let ids = [2, 0, 2]; - // All coordinates are exactly representable in FP16, including -0.0. - let expected = [points[2], points[0], points[2]]; - let mut actual = [123.0_f32; 27]; - - // When - gather_as_f32(data, &ids, &mut actual).unwrap(); - - // Then - assert_eq!( - actual.map(f32::to_bits).as_slice(), - expected - .as_flattened() - .iter() - .copied() - .map(f32::to_bits) - .collect::>() - ); + let mut output = [42.0; 6]; + + gather_as_f32( + MatrixView::try_from(&values[..], 2, 3).unwrap(), + &[1, 0], + &mut output, + ) + .unwrap(); + + assert_eq!(output[0], f32::NEG_INFINITY); + assert_eq!(output[1], 2.0_f32.powi(-24)); + assert!(output[2].is_nan()); + assert_eq!(output[3].to_bits(), 0.0_f32.to_bits()); + assert_eq!(output[4].to_bits(), (-0.0_f32).to_bits()); + assert_eq!(output[5], f32::INFINITY); } } diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 4c5ef60fc1..13d187a8b0 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -124,173 +124,193 @@ pub(crate) fn prune_overfull( } #[cfg(test)] -mod prune_tests { - use crate::graph::config::{self, MaxDegree}; - use rstest::rstest; - +mod tests { use super::*; + use crate::graph::config; + use rstest::rstest; - fn graph_config(degree: usize) -> Config { - config::Builder::new_with( - degree, - MaxDegree::same(), - degree, - Metric::L2.into(), - |builder| { - builder.alpha(1.2); - }, - ) + fn pruning_config( + degree: usize, + metric: Metric, + alpha: f32, + ) -> Result { + config::Builder::new_with(degree, config::MaxDegree::same(), 16, metric.into(), |b| { + b.alpha(alpha); + }) .build() - .unwrap() + } + + fn candidates_in_order(ids: &[u32]) -> AdjacencyList { + let mut candidates = AdjacencyList::new(); + candidates.extend_from_slice(ids); + candidates } #[rstest] #[case::empty(&[])] - #[case::below_degree(&[3])] + #[case::below_degree(&[2])] #[case::at_degree(&[3, 1])] - fn rows_within_degree_are_unchanged(#[case] ids: &[u32]) { - // Given: preserve the caller's order, including a non-sorted row. - let values = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::column_vector(&values[..]); - let mut candidates = AdjacencyList::new(); - candidates.overwrite_trusted(ids); + fn candidates_within_degree_keep_their_order(#[case] ids: &[u32]) { + let values = [0.0_f32, 1.0, 3.0, -2.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let graph = pruning_config(2, Metric::L2, 1.0).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); - let mut workspace = PruneWorkspace::default(); - // When - let actual = workspace - .prune(data, 0, candidates, &graph_config(2), &distance) + let actual = PruneWorkspace::default() + .prune(data, 0, candidates_in_order(ids), &graph, &distance) .unwrap(); - // Then assert_eq!(&*actual, ids); } - #[test] - fn overfull_rows_keep_nearest_unoccluded_neighbors_at_their_source_id() { - // Given: source 1 is at 0. Its squared distances to IDs 2, 0, 3 are - // 1, 4, 9. ID 2 occludes ID 0 because 4 / (2 - 1)^2 = 4 > alpha. - // ID 3 survives because 9 / (-3 - 1)^2 = 9/16 <= 1. - let values = [2.0_f32, 0.0, 1.0, -3.0]; - let data = MatrixView::column_vector(&values[..]); - let candidates = vec![ - AdjacencyList::new(), - AdjacencyList::from_iter_untrusted([0, 2, 3]), - AdjacencyList::new(), - AdjacencyList::new(), - ]; - let expected = [vec![], vec![2, 3], vec![], vec![]]; - - // When - let actual: Vec> = prune_overfull(data, candidates, &graph_config(2), Metric::L2) - .unwrap() - .into_iter() - .map(Vec::from) - .collect(); - - // Then - assert_eq!(actual, expected); - } - - #[test] - fn self_candidates_do_not_shift_selected_neighbor_ids() { - // Given: self occupies the first sorted position but cannot be selected. - // Of the other points, 1 occludes 2, while -3 remains on the opposite ray. - let values = [0.0_f32, 1.0, 2.0, -3.0]; - let data = MatrixView::column_vector(&values[..]); - let candidates = AdjacencyList::from_iter_untrusted([0, 1, 2, 3]); - let distance = f32::distance(Metric::L2, Some(1)); - let expected = [1, 3]; - let mut workspace = PruneWorkspace::default(); + #[rstest] + #[case::l2(Metric::L2, [0.0, 0.0, 1.0, 0.0, 3.0, 0.0, -2.0, 0.0], [1, 3])] + #[case::cosine(Metric::Cosine, [1.0, 0.0, 2.0, 0.0, 1.0, 1.0, -1.0, 0.0], [1, 2])] + #[case::normalized_cosine(Metric::CosineNormalized, [1.0, 0.0, 0.8, 0.6, -1.0, 0.0, 0.0, -1.0], [1, 3])] + #[case::inner_product(Metric::InnerProduct, [1.0, 0.0, 3.0, 0.0, 2.0, 1.0, -1.0, 4.0], [1, 3])] + fn overfull_candidates_follow_the_metric_pruning_rule( + #[case] metric: Metric, + #[case] values: [f32; 8], + #[case] expected: [u32; 2], + ) { + // ID 1 is nearest. L2 occludes ID 2: 9 / 4 > 1. + // Cosine retains ID 2: equal directions give an occlusion ratio of 1. + // Normalized cosine retains ID 3: 1 / 1.6 < 1. + // Inner product occludes ID 2: its dot with ID 1 is 6, versus 2 with the source. + let data = MatrixView::try_from(&values[..], 4, 2).unwrap(); + let graph = pruning_config(2, metric, 1.0).unwrap(); + let distance = f32::distance(metric, Some(2)); - // When - let actual = workspace - .prune(data, 0, candidates, &graph_config(2), &distance) + let actual = PruneWorkspace::default() + .prune(data, 0, candidates_in_order(&[3, 2, 1]), &graph, &distance) .unwrap(); - // Then - assert_eq!(&*actual, &expected); + assert_eq!(&*actual, expected); } #[test] - fn reused_workspace_clears_previous_candidates_and_occlusion_states() { - // Given: the same workspace prunes four candidates, then three, then four. - let values = [0.0_f32, 1.0, 2.0, -3.0, 4.0]; - let data = MatrixView::column_vector(&values[..]); - let graph = graph_config(2); + fn pruning_excludes_the_source_without_shifting_selected_ids() { + let values = [0.0_f32, 1.0, 3.0, -2.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let graph = pruning_config(2, Metric::L2, 1.0).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); - let mut workspace = PruneWorkspace::default(); - workspace + + let actual = PruneWorkspace::default() .prune( data, 0, - AdjacencyList::from_iter_untrusted([1, 2, 3, 4]), - &graph, - &distance, - ) - .unwrap(); - workspace - .prune( - data, - 4, - AdjacencyList::from_iter_untrusted([0, 1, 3]), + candidates_in_order(&[3, 0, 2, 1]), &graph, &distance, ) .unwrap(); - // For source 2 at coordinate 2, ID 1 at 1 is nearest. It occludes ID 0 - // (4/1 > alpha); ID 4 at 4 survives (4/9 <= 1). - let expected = [1, 4]; - let candidates = AdjacencyList::from_iter_untrusted([0, 1, 3, 4]); - - // When - let actual = workspace - .prune(data, 2, candidates, &graph, &distance) + + assert_eq!(&*actual, [1, 3]); + } + + #[rstest] + #[case::strict(1.0, &[1])] + #[case::relaxed(2.0, &[1, 3])] + fn graph_alpha_controls_which_occluded_neighbors_return( + #[case] alpha: f32, + #[case] expected: &[u32], + ) { + // After selecting x=1, x=3 has ratio 9/4 and x=4 has ratio 16/9. + // Alpha 2 admits x=4 while x=3 remains occluded. + let values = [0.0_f32, 1.0, 3.0, 4.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let graph = pruning_config(2, Metric::L2, alpha).unwrap(); + let distance = f32::distance(Metric::L2, Some(1)); + + let actual = PruneWorkspace::default() + .prune(data, 0, candidates_in_order(&[3, 2, 1]), &graph, &distance) .unwrap(); - // Then - assert_eq!(&*actual, &expected); + assert_eq!(&*actual, expected); + } + + #[test] + fn reused_workspace_does_not_carry_neighbors_or_occlusion_between_sources() { + let values = [0.0_f32, 1.0, 3.0, -2.0, 9.0]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + let graph = pruning_config(2, Metric::L2, 1.0).unwrap(); + let distance = f32::distance(Metric::L2, Some(1)); + let mut workspace = PruneWorkspace::default(); + + // One workspace serves a large list, a smaller list, then a large list again. + for (source, ids, expected) in [ + (0, &[4, 3, 2, 1][..], &[1, 3][..]), + (4, &[3, 2, 1][..], &[2][..]), + (1, &[4, 3, 2, 0][..], &[0, 2][..]), + ] { + let actual = workspace + .prune(data, source, candidates_in_order(ids), &graph, &distance) + .unwrap(); + + assert_eq!(&*actual, expected, "source {source}"); + } } + #[cfg(not(miri))] #[test] - fn maximum_u16_candidate_count_can_be_pruned() { - // Given: all candidates lie on the positive ray, so ID 1 is nearest. + fn the_largest_supported_candidate_list_can_be_pruned() { let count = u16::MAX as usize; - let values: Vec = (0..=count).map(|id| id as f32).collect(); - let data = MatrixView::column_vector(&values[..]); + let values: Vec<_> = (0..=count).map(|i| i as f32).collect(); + let data = MatrixView::try_from(values.as_slice(), count + 1, 1).unwrap(); let candidates = AdjacencyList::from_iter_untrusted(1..=count as u32); + let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); - let mut workspace = PruneWorkspace::default(); - // When - let actual = workspace - .prune(data, 0, candidates, &graph_config(1), &distance) + let actual = PruneWorkspace::default() + .prune(data, 0, candidates, &graph, &distance) .unwrap(); - // Then - assert_eq!(&*actual, &[1]); + assert_eq!(&*actual, [1]); } + #[cfg(not(miri))] #[test] - fn candidate_count_above_u16_limit_is_rejected() { - // Given: the list is one entry longer than RobustPrune's position limit. + fn a_candidate_list_beyond_the_position_limit_returns_its_size() { let count = u16::MAX as usize + 1; - let values = vec![0.0_f32; count + 1]; - let data = MatrixView::column_vector(&values[..]); + let values: Vec<_> = (0..=count).map(|i| i as f32).collect(); + let data = MatrixView::try_from(values.as_slice(), count + 1, 1).unwrap(); let candidates = AdjacencyList::from_iter_untrusted(1..=count as u32); + let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); - let mut workspace = PruneWorkspace::default(); - // When - let error = workspace - .prune(data, 0, candidates, &graph_config(1), &distance) + let error = PruneWorkspace::default() + .prune(data, 0, candidates, &graph, &distance) .unwrap_err(); - // Then - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::TooManyCandidates { actual, max }) - if *actual == count && *max == u16::MAX as usize - )); + let error = error.downcast_ref::().unwrap(); + assert!( + matches!(error, FinalizationError::TooManyCandidates { actual, max } + if *actual == count && *max == u16::MAX as usize) + ); + } + + #[rstest] + #[case::one_worker(1)] + #[case::several_workers(3)] + fn parallel_pruning_keeps_each_result_with_its_source(#[case] workers: usize) { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(workers) + .build() + .unwrap(); + let values = [0.0_f32, 1.0, 4.0, 9.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let candidates = [[3, 2, 1], [3, 2, 0], [3, 1, 0], [2, 1, 0]] + .map(|ids| candidates_in_order(&ids)) + .to_vec(); + let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); + + let actual = pool + .install(|| prune_overfull(data, candidates, &graph, Metric::L2)) + .unwrap(); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [vec![1], vec![0], vec![1], vec![2]] + ); } } diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 2f09f38064..8693cfc63c 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -243,171 +243,256 @@ fn grow(values: &mut Vec, len: usize, value: T) { } } -#[cfg(test)] -mod build_leaf_candidates_tests { - use diskann_wide::arch::Scalar; - use half::f16; +#[cfg(all(test, not(miri)))] +mod tests { + use super::*; + use crate::graph::pipnn::{L2, test_support}; + use diskann_wide::ARCH; use rstest::rstest; - use super::*; - use crate::graph::pipnn::{InnerProduct, L2}; - - fn build_candidates( - data: MatrixView<'_, T>, - leaves: Vec>, - k: usize, - threads: usize, - ) -> Vec> { - rayon::ThreadPoolBuilder::new() - .num_threads(threads) + #[test] + fn selected_neighbors_use_global_ids_and_contribute_both_edge_directions() { + // The leaf lists IDs [5, 1, 3], at coordinates [9, 0, 2]. + // The directed choices are 1 -> 3, 3 -> 1 and 5 -> 3. + let values = [100.0_f32, 0.0, -100.0, 2.0, 200.0, 9.0]; + let data = MatrixView::try_from(&values[..], 6, 1).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) .build() - .unwrap() - .install(|| build_leaf_candidates::<_, M, _>(Scalar, data, leaves, k)) - .unwrap() - .into_iter() - .map(Vec::from) - .collect() + .unwrap(); + + let actual = pool + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, vec![vec![5, 1, 3]], 1)) + .unwrap(); + let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + + assert_eq!( + test_support::sorted_members_per_row(&actual), + [vec![], vec![3], vec![], vec![1, 5], vec![], vec![3]] + ); } #[rstest] - #[case::f32([5.5_f32, 100.0, 1.25, 50.0, 0.25])] - #[case::f16([5.5, 100.0, 1.25, 50.0, 0.25].map(f16::from_f32))] - #[case::u8([5_u8, 100, 1, 50, 0])] - #[case::i8([-6_i8, 100, -10, 50, -11])] - fn selected_local_positions_become_global_ids(#[case] values: [T; 5]) { - // Given: only IDs 0, 2, 4 belong to the leaf. Their coordinates decrease, - // so 0 selects 2, while 2 and 4 select each other. Reverse edges add 2 -> 0. - let data = MatrixView::column_vector(&values[..]); - let leaves = vec![vec![0, 2, 4]]; - let expected = [vec![2], vec![], vec![0, 4], vec![], vec![2]]; - - // When - let actual = build_candidates::<_, L2>(data, leaves, 1, 1); - - // Then - assert_eq!(actual, expected); - } + #[case::one_worker(1)] + #[case::several_workers(3)] + fn overlapping_leaves_merge_each_neighbor_once(#[case] workers: usize) { + let values = [0.0_f32, 1.0, 4.0, 9.0, 16.0]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + let leaves = vec![vec![0, 1, 2], vec![1, 2, 3], vec![2, 3, 4], vec![0, 1, 2]]; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(workers) + .build() + .unwrap(); - #[test] - fn unassigned_kernel_results_add_no_edges() { - // Given: the NaN point has no rankable distances. The two finite points - // select each other, leaving their second output slot unassigned. - let values = [0.0_f32, 1.0, f32::NAN]; - let data = MatrixView::column_vector(&values[..]); - let expected = [vec![1], vec![0], vec![]]; - - // When - let actual = build_candidates::<_, InnerProduct>(data, vec![vec![0, 1, 2]], 2, 1); - - // Then - assert_eq!(actual, expected); + let actual = pool + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves, 1)) + .unwrap(); + let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + + assert_eq!( + test_support::sorted_members_per_row(&actual), + [vec![1], vec![0, 2], vec![1, 3], vec![2, 4], vec![3]] + ); } #[rstest] - #[case::one_worker(1)] - #[case::four_workers(4)] - fn overlapping_leaves_merge_into_sorted_unique_neighbors(#[case] threads: usize) { - // Given: k = 2 connects every pair in each three-point leaf. The union - // contains all pairs except 1 <-> 3, regardless of repeated leaf jobs. - let values = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::column_vector(&values[..]); - let leaves = [vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]] - .into_iter() - .cycle() - .take(48) - .collect(); - let expected = [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]]; - - // When - let actual = build_candidates::<_, L2>(data, leaves, 2, threads); - - // Then - assert_eq!(actual, expected); + #[case::no_leaves(vec![], 2)] + #[case::empty_leaf(vec![vec![]], 2)] + #[case::singleton(vec![vec![2]], 2)] + #[case::zero_neighbors(vec![vec![0, 1, 2]], 0)] + fn leaves_without_selected_pairs_produce_empty_adjacency( + #[case] leaves: Vec>, + #[case] requested_k: usize, + ) { + let values = [0.0_f32, 1.0, 4.0]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + + let actual = pool + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves, requested_k)) + .unwrap(); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [Vec::::new(), vec![], vec![]] + ); } #[test] - fn reverse_edges_can_give_a_point_more_than_twice_k_neighbors() { - // Given: each outer point is distance 1 from the origin and at least - // sqrt(2) from another outer point. Their three reverse edges exceed 2k. - let points = [[0.0_f32, 0.0], [1.0, 0.0], [-1.0, 0.0], [0.0, 1.0]]; - let data = MatrixView::try_from(points.as_flattened(), 4, 2).unwrap(); - let expected = [vec![1, 2, 3], vec![0], vec![0], vec![0]]; - - // When - let actual = build_candidates::<_, L2>(data, vec![vec![0, 1, 2, 3]], 1, 1); - - // Then - assert_eq!(actual, expected); + fn requesting_more_neighbors_than_a_leaf_has_selects_all_other_points() { + let values = [100.0_f32, 0.0, -100.0, 2.0, 200.0, 9.0]; + let data = MatrixView::try_from(&values[..], 6, 1).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + + let actual = pool + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, vec![vec![1, 3, 5]], 99)) + .unwrap(); + let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + + assert_eq!( + test_support::sorted_members_per_row(&actual), + [vec![], vec![3, 5], vec![], vec![1, 5], vec![], vec![1, 3]] + ); } #[test] - fn singleton_leaves_add_no_edges() { - // Given - let values = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::column_vector(&values[..]); - let leaves = vec![vec![0], vec![1], vec![2]]; - let expected: [Vec; 3] = [vec![], vec![], vec![]]; - - // When - let actual = build_candidates::<_, L2>(data, leaves, 1, 1); - - // Then - assert_eq!(actual, expected); + fn unrankable_pairs_do_not_add_unassigned_ids_to_the_graph() { + let values = [0.0_f32, 3.0, f32::NAN]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + + let actual = pool + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, vec![vec![0, 1, 2]], 2)) + .unwrap(); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [vec![1], vec![0], vec![]] + ); } #[test] - fn reused_leaf_buffers_do_not_add_edges_from_previous_leaves() { - // Given: use the same buffers for four points, two points, then three. - // The final leaf excludes ID 1 and connects all pairs among 0, 2, 3. - let values = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::column_vector(&values[..]); + fn reused_leaf_buffers_do_not_carry_edges_between_different_leaf_shapes() { + let values = [0.0_f32, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0]; + let data = MatrixView::try_from(&values[..], 7, 1).unwrap(); let mut buffers = LeafBuffers::default(); - let mut build_leaf = |point_ids: &[u32]| { - let candidates: Vec<_> = (0..data.nrows()) - .map(|_| Mutex::new(AdjacencyList::new())) - .collect(); + + // The same worker processes a leaf, a smaller leaf, then more neighbors + // per point. Every call must use only its current IDs and active rows. + for (ids, requested_k, expected) in [ + ( + vec![0, 1, 2, 3], + 1, + vec![ + vec![1], + vec![0, 2], + vec![1, 3], + vec![2], + vec![], + vec![], + vec![], + ], + ), + ( + vec![4, 6], + 99, + vec![vec![], vec![], vec![], vec![], vec![6], vec![], vec![4]], + ), + ( + vec![1, 3, 5], + 2, + vec![ + vec![], + vec![3, 5], + vec![], + vec![1, 5], + vec![], + vec![1, 3], + vec![], + ], + ), + ] { + let candidates: Vec<_> = (0..7).map(|_| Mutex::new(AdjacencyList::new())).collect(); + add_direct_leaf_candidates::<_, L2, _>( - Scalar, + ARCH, data, 0, - point_ids, - 2, + &ids, + requested_k, &mut buffers, &candidates, ) .unwrap(); - candidates + + let actual: Vec<_> = candidates .into_iter() - .map(|list| { - let mut ids = Vec::from(list.into_inner()); + .map(|row| { + let mut ids = Vec::from(row.into_inner()); ids.sort_unstable(); ids }) - .collect::>() + .collect(); + assert_eq!(actual, expected, "leaf {ids:?}, k={requested_k}"); + } + } + + #[rstest] + #[case::point_values(2, 1)] + #[case::neighbor_slots(1, 2)] + fn an_overflowing_leaf_shape_is_rejected_before_buffers_change( + #[case] dimensions: usize, + #[case] requested_k: usize, + ) { + let old_neighbor = Candidate::new(1, 7.0); + let mut buffers = LeafBuffers { + point_values: vec![3.0, 4.0], + neighbors: vec![old_neighbor], + ..LeafBuffers::default() }; - build_leaf(&[0, 1, 2, 3]); - build_leaf(&[1, 3]); - let expected = [vec![2, 3], vec![], vec![0, 3], vec![0, 2]]; - // When - let actual = build_leaf(&[0, 2, 3]); + let error = buffers + .prepare(7, usize::MAX, dimensions, requested_k) + .unwrap_err(); - // Then - assert_eq!(actual, expected); + assert!(matches!( + error, + LeafBuildError::ShapeOverflow { + leaf: 7, + rows: usize::MAX, + columns: 2, + } + )); + assert_eq!(buffers.point_values, [3.0, 4.0]); + assert_eq!(buffers.neighbors, [old_neighbor]); } #[test] - fn overflowing_leaf_shape_is_rejected() { - // Given - let mut buffers = LeafBuffers::default(); + fn a_kernel_error_retains_the_failed_leaf_and_original_cause() { + #[derive(Debug, thiserror::Error)] + #[error("distance calculation unavailable")] + struct DistanceFailure; + + // This stub supplies an otherwise hard-to-trigger dependency failure. + // Gathering, leaf indexing and error wrapping remain real. + struct UnavailableMetric; + impl LeafMetric for UnavailableMetric { + fn compute_distances(_: MatrixView<'_, f32>, _: &mut [f32]) -> crate::ANNResult<()> { + Err(crate::ANNError::new(DistanceFailure)) + } + } - // When - let result = buffers.prepare(7, usize::MAX, 2, 1); + let values = [0.0_f32, 1.0, 4.0, 9.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); - // Then - assert!(matches!( - result, - Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) - )); + let error = pool + .install(|| { + build_leaf_candidates::<_, UnavailableMetric, _>( + ARCH, + data, + vec![vec![], vec![1, 3]], + 1, + ) + }) + .unwrap_err(); + + let LeafBuildError::Kernel { leaf, source } = error else { + panic!("expected a leaf kernel error, got {error:?}"); + }; + assert_eq!(leaf, 1); + assert!(source.downcast_ref::().is_some()); } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 2b0997e5c6..f4a4203d72 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -357,403 +357,32 @@ fn config_error(message: impl std::fmt::Display) -> ANNError { ANNError::message(format!("PiPNN configuration: {message}")) } -#[cfg(test)] -mod effective_metric_tests { - use super::*; - use half::f16; - use rstest::rstest; - - #[rstest] - #[case::l2(Metric::L2)] - #[case::cosine(Metric::Cosine)] - #[case::normalized_cosine(Metric::CosineNormalized)] - #[case::inner_product(Metric::InnerProduct)] - fn floating_point_formats_keep_the_requested_metric(#[case] metric: Metric) { - // Given / When: floating-point callers define whether their data is normalized. - let actual = ( - effective_metric::(metric), - effective_metric::(metric), - ); - - // Then: both supported floating-point formats preserve the request. - assert_eq!(actual, (metric, metric)); - } - - #[rstest] - #[case::l2(Metric::L2)] - #[case::cosine(Metric::Cosine)] - #[case::inner_product(Metric::InnerProduct)] - fn integer_formats_keep_metrics_that_do_not_assume_normalization(#[case] metric: Metric) { - // Given / When - let actual = ( - effective_metric::(metric), - effective_metric::(metric), - ); - - // Then: only normalized cosine needs the integer-data fallback. - assert_eq!(actual, (metric, metric)); - } -} -#[cfg(test)] -#[allow( - clippy::expect_used, - clippy::unwrap_used, - reason = "deterministic test fixture construction must abort on invalid setup" -)] -mod build_graph_tests { - use super::{PiPNNBuildContext, PiPNNConfig, build_graph}; - use crate::graph::config::{self, MaxDegree}; - use diskann_utils::views::MatrixView; - use diskann_vector::distance::Metric; - use rstest::rstest; - - fn pipnn_config() -> PiPNNConfig { - PiPNNConfig { - c_max: 4, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - leaf_k: 1, - replicas: 1, - } - } - - fn graph_config(metric: Metric, degree: usize) -> crate::graph::Config { - config::Builder::new_with(degree, MaxDegree::same(), 8, metric.into(), |builder| { - builder.alpha(1.2); - }) - .build() - .unwrap() - } - - fn pool(threads: usize) -> rayon::ThreadPool { - rayon::ThreadPoolBuilder::new() - .num_threads(threads) - .build() - .unwrap() - } - - fn rows(graph: Vec>) -> Vec> { - graph.into_iter().map(Vec::from).collect() - } - - #[test] - fn single_leaf_merges_neighbors_in_both_directions() { - // Given: unequal gaps give nearest choices 0->1, 1->0, 2->1, 3->2. - // Symmetric insertion adds the reverse edge for each choice. - let values = [0.0_f32, 1.0, 3.0, 7.0]; - let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let expected = [vec![1], vec![0, 2], vec![1, 3], vec![2]]; - let graph = graph_config(Metric::L2, 2); - let pool = pool(2); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - // When - let actual = build_graph(data, &context).unwrap(); - - // Then - assert_eq!(rows(actual), expected); - } - - #[test] - fn omits_non_rankable_candidates_without_invalid_ids() { - // Given: three points force partitioning at c_max=2. The NaN point - // has no rankable leader or neighbor, so only the finite points connect. - let values = [0.0_f32, 1.0, f32::NAN]; - let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); - let graph = graph_config(Metric::InnerProduct, 2); - let pool = pool(1); - let config = PiPNNConfig { - c_max: 2, - c_min: 1, - p_samp: 1.0, - fanout: vec![2], - leaf_k: 1, - replicas: 1, - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::InnerProduct, &pool).unwrap(); - let expected = [vec![1], vec![0], vec![]]; - - // When - let actual = build_graph(data, &context).unwrap(); - - // Then - assert_eq!(rows(actual), expected); - } - - #[test] - fn finalization_keeps_the_nearest_neighbor_when_degree_is_one() { - // Given: selecting both peers makes each row overfull. With degree one, - // RobustPrune keeps the nearest: 0->1, 1->0, 2->1. - let values = [0.0_f32, 1.0, 3.0]; - let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); - let expected = [vec![1], vec![0], vec![1]]; - let graph = graph_config(Metric::L2, 1); - let pool = pool(2); - let config = PiPNNConfig { - c_max: 3, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - leaf_k: 2, - replicas: 1, - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - // When - let actual = build_graph(data, &context).unwrap(); - - // Then - assert_eq!(rows(actual), expected); - } - - #[rstest] - #[case::no_points(0, 4, "PiPNN requires at least one data point")] - #[case::no_dimensions(4, 0, "PiPNN requires at least one data dimension")] - fn empty_input_shape_returns_an_error( - #[case] points: usize, - #[case] dimensions: usize, - #[case] message: &str, - ) { - // Given - let data = MatrixView::try_from(&[] as &[f32], points, dimensions).unwrap(); - let graph = graph_config(Metric::L2, 2); - let pool = pool(1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - // When - let error = build_graph(data, &context).unwrap_err(); - - // Then - assert!(error.to_string().contains(message), "{error}"); - } - - #[rstest] - #[case::l2(Metric::L2, [vec![1, 2], vec![0], vec![0]])] - #[case::inner_product(Metric::InnerProduct, [vec![1], vec![0, 2], vec![1]])] - fn selected_metric_changes_the_neighbor_graph( - #[case] metric: Metric, - #[case] expected: [Vec; 3], - ) { - // Given: squared L2 scores for pairs (0,1),(0,2),(1,2) are 4,0.8,6.4; - // negative-dot scores are -3,-0.6,-1.8. The nearest choices differ. - let values = [1.0_f32, 0.0, 3.0, 0.0, 0.6, 0.8]; - let data = MatrixView::try_from(&values[..], 3, 2).unwrap(); - let graph = graph_config(metric, 2); - let pool = pool(1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); - - // When - let actual = build_graph(data, &context).unwrap(); - - // Then - assert_eq!(rows(actual), expected); - } - - #[rstest] - #[case::cosine(Metric::Cosine)] - #[case::normalized_cosine(Metric::CosineNormalized)] - fn cosine_metrics_rank_unit_vectors_by_direction(#[case] metric: Metric) { - // Given: pairwise dots are 0.6,-1,-0.6, so 0 chooses 1, 1 chooses 0, - // and 2 chooses 1. Symmetrization gives the expected three-node chain. - let values = [1.0_f32, 0.0, 0.6, 0.8, -1.0, 0.0]; - let data = MatrixView::try_from(&values[..], 3, 2).unwrap(); - let expected = [vec![1], vec![0, 2], vec![1]]; - let graph = graph_config(metric, 2); - let pool = pool(1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); - - // When - let actual = build_graph(data, &context).unwrap(); - - // Then - assert_eq!(rows(actual), expected); - } - - #[test] - fn singleton_has_no_self_neighbor() { - // Given: a singleton has zero non-self candidates even though leaf_k=1. - let values = [3.0_f32, 4.0]; - let data = MatrixView::try_from(&values[..], 1, 2).unwrap(); - let graph = graph_config(Metric::L2, 2); - let pool = pool(1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - // When - let actual = build_graph(data, &context).unwrap(); - - // Then - assert_eq!(rows(actual), [Vec::::new()]); - } - - #[rstest] - #[case::unsigned([1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2])] - #[case::signed([1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2])] - fn integer_normalized_cosine_matches_cosine( - #[case] values: [T; 16], - ) { - // Given: unequal norms distinguish cosine from raw dot ranking. - // Keep one leaf so partition randomness cannot hide metric selection. - let data = MatrixView::try_from(&values[..], 8, 2).unwrap(); - let pool = pool(2); - let build = |metric| { - let graph = graph_config(metric, 2); - let config = PiPNNConfig { - c_max: 8, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - leaf_k: 1, - replicas: 1, - }; - let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); - rows(build_graph(data, &context).unwrap()) - }; - let expected = build(Metric::Cosine); - - // When - let actual = build(Metric::CosineNormalized); - - // Then - assert_eq!(actual, expected); - } - - #[rstest] - #[case::one_worker(1)] - #[case::four_workers(4)] - fn recursive_build_is_deterministic_for_a_fixed_pool(#[case] threads: usize) { - // Given: distinct points require multiple partition levels and two replicas. - let values: Vec = (0..96).flat_map(|i| [i as f32, (i * i) as f32]).collect(); - let data = MatrixView::try_from(values.as_slice(), 96, 2).unwrap(); - let graph = graph_config(Metric::L2, 8); - let pool = pool(threads); - let config = PiPNNConfig { - c_max: 16, - c_min: 4, - p_samp: 0.25, - fanout: vec![3, 2], - leaf_k: 3, - replicas: 2, - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - let expected = build_graph(data, &context).unwrap(); - - // When - let actual = build_graph(data, &context).unwrap(); - - // Then - assert_eq!(actual, expected); - } -} -#[cfg(test)] -#[allow( - clippy::expect_used, - clippy::unwrap_used, - reason = "deterministic test fixture construction must abort on invalid setup" -)] -mod config_tests { - use super::{PiPNNBuildContext, PiPNNConfig}; - use crate::graph::config::{self, MaxDegree}; +#[cfg(all(test, not(miri)))] +mod test_support { use diskann_vector::distance::Metric; - use rstest::rstest; - fn pipnn_config() -> PiPNNConfig { - PiPNNConfig { - c_max: 512, - c_min: 64, - p_samp: 0.01, - fanout: vec![10, 3], - leaf_k: 2, - replicas: 1, - } - } - - fn graph_config(metric: Metric, alpha: f32) -> crate::graph::Config { - config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { - builder.alpha(alpha); - }) - .build() - .unwrap() - } - - fn pool() -> rayon::ThreadPool { - rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap() - } - - #[rstest] - #[case::zero_maximum(PiPNNConfig { c_max: 0, ..pipnn_config() }, "c_max must be greater than zero")] - #[case::zero_minimum(PiPNNConfig { c_min: 0, ..pipnn_config() }, "c_min must be greater than zero")] - #[case::minimum_exceeds_maximum(PiPNNConfig { c_min: 513, ..pipnn_config() }, "c_min (513) must not exceed c_max (512)")] - #[case::zero_sampling(PiPNNConfig { p_samp: 0.0, ..pipnn_config() }, "p_samp (0) must be in (0, 1]")] - #[case::negative_sampling(PiPNNConfig { p_samp: -0.5, ..pipnn_config() }, "p_samp (-0.5) must be in (0, 1]")] - #[case::sampling_above_one(PiPNNConfig { p_samp: 1.5, ..pipnn_config() }, "p_samp (1.5) must be in (0, 1]")] - #[case::nan_sampling(PiPNNConfig { p_samp: f64::NAN, ..pipnn_config() }, "p_samp (NaN) must be in (0, 1]")] - #[case::infinite_sampling(PiPNNConfig { p_samp: f64::INFINITY, ..pipnn_config() }, "p_samp (inf) must be in (0, 1]")] - #[case::empty_fanout(PiPNNConfig { fanout: vec![], ..pipnn_config() }, "fanout must not be empty")] - #[case::zero_fanout(PiPNNConfig { fanout: vec![1, 0], ..pipnn_config() }, "fanout values must be greater than zero")] - #[case::zero_neighbors(PiPNNConfig { leaf_k: 0, ..pipnn_config() }, "leaf_k must be greater than zero")] - #[case::zero_replicas(PiPNNConfig { replicas: 0, ..pipnn_config() }, "replicas must be greater than zero")] - fn invalid_policy_reports_the_rejected_parameter( - #[case] config: PiPNNConfig, - #[case] message: &str, - ) { - // Given: one invalid field in an otherwise valid policy. - let expected = format!("PiPNN configuration: {message}"); - - // When - let error = config.validate().unwrap_err(); - - // Then - assert!(error.to_string().contains(&expected), "{error}"); + // Ignore member order while preserving group positions and duplicate counts. + // Positions identify leaders or graph sources; sorting the outer list would + // hide assignments to the wrong group. Copies keep the observations intact. + pub(super) fn sorted_members_per_row(rows: &[Vec]) -> Vec> { + rows.iter() + .map(|row| { + let mut members = row.clone(); + members.sort_unstable(); + members + }) + .collect() } #[test] - fn context_validates_the_partition_policy() { - // Given: verify the public constructor invokes the policy validation. - let config = PiPNNConfig { - replicas: 0, - ..pipnn_config() - }; - let graph = graph_config(Metric::L2, 1.2); - let pool = pool(); - - // When - let error = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap_err(); - - // Then - assert!( - error - .to_string() - .contains("PiPNN configuration: replicas must be greater than zero"), - "{error}", - ); - } + fn membership_comparison_preserves_group_positions_and_duplicate_counts() { + let rows = [vec![9, 3, 9], vec![], vec![2, 1], vec![2, 1]]; - #[test] - fn rejects_graph_policy_for_a_different_metric() { - // Given: inner-product occlusion cannot prune an L2 build. - let graph = graph_config(Metric::InnerProduct, 1.2); - let pool = pool(); - let expected = format!( - "PiPNN configuration: graph prune kind {:?} is incompatible with metric L2", - graph.prune_kind(), + assert_eq!( + sorted_members_per_row(&rows), + [vec![3, 9, 9], vec![], vec![1, 2], vec![1, 2]] ); - - // When - let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); - - // Then - assert!(error.to_string().contains(&expected), "{error}"); } -} - -#[cfg(all(test, not(miri)))] -mod test_support { - use diskann_vector::distance::Metric; pub(super) fn dense_points(rows: usize, dimensions: usize, seed: u64) -> Vec { use rand::{Rng, SeedableRng, rngs::StdRng}; @@ -865,3 +494,345 @@ mod test_support { ); } } + +#[cfg(all(test, not(miri)))] +mod construction_tests { + use super::*; + use crate::graph::config; + use half::f16; + use rstest::rstest; + + fn partition_policy() -> PiPNNConfig { + PiPNNConfig { + c_max: 4, + c_min: 1, + p_samp: 1.0, + fanout: vec![2], + leaf_k: 1, + replicas: 1, + } + } + + fn graph_policy(degree: usize, metric: Metric) -> Result { + config::Builder::new_with(degree, config::MaxDegree::same(), 16, metric.into(), |b| { + b.alpha(1.0); + }) + .build() + } + + #[rstest] + #[case::zero_maximum(|c: &mut PiPNNConfig| c.c_max = 0, "c_max must be greater than zero")] + #[case::zero_minimum(|c: &mut PiPNNConfig| c.c_min = 0, "c_min must be greater than zero")] + #[case::inverted_leaf_bounds(|c: &mut PiPNNConfig| c.c_min = 5, "c_min (5) must not exceed c_max (4)")] + #[case::zero_probability(|c: &mut PiPNNConfig| c.p_samp = 0.0, "p_samp (0) must be in (0, 1]")] + #[case::negative_probability(|c: &mut PiPNNConfig| c.p_samp = -0.5, "p_samp (-0.5) must be in (0, 1]")] + #[case::probability_above_one(|c: &mut PiPNNConfig| c.p_samp = 1.0 + f64::EPSILON, "must be in (0, 1]")] + #[case::nan_probability(|c: &mut PiPNNConfig| c.p_samp = f64::NAN, "p_samp (NaN) must be in (0, 1]")] + #[case::positive_infinite_probability(|c: &mut PiPNNConfig| c.p_samp = f64::INFINITY, "p_samp (inf) must be in (0, 1]")] + #[case::negative_infinite_probability(|c: &mut PiPNNConfig| c.p_samp = f64::NEG_INFINITY, "p_samp (-inf) must be in (0, 1]")] + #[case::missing_fanout(|c: &mut PiPNNConfig| c.fanout.clear(), "fanout must not be empty")] + #[case::zero_later_fanout(|c: &mut PiPNNConfig| c.fanout = vec![2, 0, 1], "fanout values must be greater than zero")] + #[case::zero_leaf_neighbors(|c: &mut PiPNNConfig| c.leaf_k = 0, "leaf_k must be greater than zero")] + #[case::zero_replicas(|c: &mut PiPNNConfig| c.replicas = 0, "replicas must be greater than zero")] + fn invalid_configuration_identifies_the_rejected_condition( + #[case] invalidate: fn(&mut PiPNNConfig), + #[case] expected_message: &str, + ) { + let mut config = partition_policy(); + invalidate(&mut config); + + let error = config.validate().unwrap_err(); + + let message = error.to_string(); + assert!(message.contains(expected_message), "{message}"); + } + + #[rstest] + #[case::smallest_positive_probability(f64::from_bits(1))] + #[case::all_points_sampled(1.0)] + fn equal_leaf_bounds_and_valid_probability_endpoints_are_accepted(#[case] p_samp: f64) { + let config = PiPNNConfig { + c_min: 4, + p_samp, + ..partition_policy() + }; + + config.validate().unwrap(); + } + + #[test] + fn a_build_context_rejects_an_invalid_partition_policy() { + let config = PiPNNConfig { + leaf_k: 0, + ..partition_policy() + }; + let graph = graph_policy(2, Metric::L2).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + + let error = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap_err(); + + assert!( + error + .to_string() + .contains("leaf_k must be greater than zero"), + "{error}" + ); + } + + #[rstest] + #[case::cosine_uses_triangle_pruning(Metric::L2, Metric::Cosine)] + #[case::normalized_cosine_uses_triangle_pruning(Metric::L2, Metric::CosineNormalized)] + #[case::inner_product_uses_occluding_pruning(Metric::InnerProduct, Metric::InnerProduct)] + fn a_build_context_accepts_metrics_with_the_same_pruning_kind( + #[case] graph_metric: Metric, + #[case] requested_metric: Metric, + ) { + let graph = graph_policy(2, graph_metric).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + + PiPNNBuildContext::new(partition_policy(), &graph, requested_metric, &pool).unwrap(); + } + + #[rstest] + #[case::inner_product_with_triangle_pruning(Metric::L2, Metric::InnerProduct)] + #[case::l2_with_occluding_pruning(Metric::InnerProduct, Metric::L2)] + fn a_build_context_rejects_incompatible_pruning_kinds( + #[case] graph_metric: Metric, + #[case] requested_metric: Metric, + ) { + let graph = graph_policy(2, graph_metric).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + + let error = PiPNNBuildContext::new(partition_policy(), &graph, requested_metric, &pool) + .unwrap_err(); + + assert!( + error.to_string().contains("incompatible with metric"), + "{error}" + ); + } + + #[rstest] + #[case::no_points(0, 2, "at least one data point")] + #[case::no_dimensions(2, 0, "at least one data dimension")] + fn an_empty_dataset_axis_is_rejected_before_building( + #[case] rows: usize, + #[case] columns: usize, + #[case] expected_message: &str, + ) { + let data = MatrixView::try_from(&[] as &[f32], rows, columns).unwrap(); + let graph = graph_policy(2, Metric::L2).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + let context = + PiPNNBuildContext::new(partition_policy(), &graph, Metric::L2, &pool).unwrap(); + + let error = build_graph(data, &context).unwrap_err(); + + assert!(error.to_string().contains(expected_message), "{error}"); + } + + #[rstest] + #[case::float32([0.0_f32, 1.0, 4.0])] + #[case::float16([0.0_f32, 1.0, 4.0].map(f16::from_f32))] + #[case::signed_integer([0_i8, 1, 4])] + #[case::unsigned_integer([0_u8, 1, 4])] + fn native_vector_types_build_the_expected_neighbor_graph( + #[case] values: [T; 3], + ) { + // Nearest choices 0 -> 1, 1 -> 0 and 2 -> 1 become a symmetric chain. + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let graph = graph_policy(2, Metric::L2).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let context = + PiPNNBuildContext::new(partition_policy(), &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + + assert_eq!( + test_support::sorted_members_per_row(&actual), + [vec![1], vec![0, 2], vec![1]] + ); + } + + #[rstest] + #[case::l2(Metric::L2, vec![vec![2], vec![2], vec![0, 1, 3], vec![2]])] + #[case::cosine(Metric::Cosine, vec![vec![1], vec![0], vec![3], vec![2]])] + #[case::normalized_cosine(Metric::CosineNormalized, vec![vec![1], vec![0], vec![3], vec![2]])] + #[case::inner_product(Metric::InnerProduct, vec![vec![1], vec![0, 3], vec![3], vec![1, 2]])] + fn the_requested_metric_determines_leaf_neighbors( + #[case] metric: Metric, + #[case] expected: Vec>, + ) { + // L2's nearest choices are [2, 2, 0, 2]; cosine pairs similar directions + // 0 <-> 1 and 2 <-> 3; dot products give [1, 3, 3, 2]. + // The second coordinate sits in the last dimension of an embedding. + let dimensions = 1537; + let values = test_support::packed_points( + &[[1.0, 0.0], [5.0, 2.0], [1.0, 3.0], [0.0, 9.0]], + dimensions, + metric == Metric::CosineNormalized, + ); + let data = MatrixView::try_from(values.as_slice(), 4, dimensions).unwrap(); + let graph = graph_policy(3, metric).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let context = PiPNNBuildContext::new(partition_policy(), &graph, metric, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + + assert_eq!( + test_support::sorted_members_per_row(&actual), + test_support::sorted_members_per_row(&expected) + ); + } + + #[rstest] + #[case::signed_leaf_ranking([1_i8, 0, 4, 1, 1, 2, 0, 9], 1, 3)] + #[case::unsigned_leaf_ranking([1_u8, 0, 4, 1, 1, 2, 0, 9], 1, 3)] + #[case::signed_final_pruning([1_i8, 0, 4, 1, 1, 2, 0, 9], 3, 1)] + #[case::unsigned_final_pruning([1_u8, 0, 4, 1, 1, 2, 0, 9], 3, 1)] + fn normalized_cosine_requests_on_raw_integers_use_vector_norms( + #[case] values: [T; 8], + #[case] leaf_k: usize, + #[case] degree: usize, + ) { + // Cosine pairs 0 <-> 1 and 2 <-> 3. Unnormalized dot products instead + // connect 1 to 3, so treating raw integers as unit vectors changes edges. + let data = MatrixView::try_from(&values[..], 4, 2).unwrap(); + let config = PiPNNConfig { + leaf_k, + ..partition_policy() + }; + let graph = graph_policy(degree, Metric::CosineNormalized).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let context = + PiPNNBuildContext::new(config, &graph, Metric::CosineNormalized, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [vec![1], vec![0], vec![3], vec![2]] + ); + } + + #[test] + fn final_pruning_enforces_graph_degree_after_leaf_candidates_are_merged() { + let values = [0.0_f32, 1.0, 4.0]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let config = PiPNNConfig { + leaf_k: 2, + ..partition_policy() + }; + let graph = graph_policy(1, Metric::L2).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [vec![1], vec![0], vec![1]] + ); + } + + #[rstest] + #[case::serial_single_replica(1, 1)] + #[case::parallel_replicas(3, 2)] + fn overlapping_partitions_build_the_same_deduplicated_graph( + #[case] workers: usize, + #[case] replicas: usize, + ) { + // Sampling every point with fanout 2 makes two overlapping copies + // of each close pair. Neither replicas nor worker count may duplicate edges. + let values = [0.0_f32, 1.0, 10.0, 11.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let config = PiPNNConfig { + c_max: 2, + replicas, + ..partition_policy() + }; + let graph = graph_policy(2, Metric::L2).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(workers) + .build() + .unwrap(); + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [vec![1], vec![0], vec![3], vec![2]] + ); + } + + #[test] + fn a_single_point_has_one_empty_adjacency_list() { + let values = [3.0_f32]; + let data = MatrixView::try_from(&values[..], 1, 1).unwrap(); + let graph = graph_policy(2, Metric::L2).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + let context = + PiPNNBuildContext::new(partition_policy(), &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [Vec::::new()] + ); + } + + #[test] + fn stalled_partitioning_remains_a_typed_build_error() { + let values = [1.0_f32; 5]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + let config = PiPNNConfig { + c_max: 2, + fanout: vec![1], + ..partition_policy() + }; + let graph = graph_policy(2, Metric::L2).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let error = build_graph(data, &context).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(partitioning::PartitionError::IterationLimit { size: 5, .. }) + )); + } +} diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 9ef05983bc..840f749e82 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -528,561 +528,438 @@ fn assignment_stripe_point_count(leader_count: usize) -> usize { point_count.clamp(MIN_ASSIGNMENT_STRIPE_POINTS, MAX_ASSIGNMENT_STRIPE_POINTS) } -#[cfg(test)] +#[cfg(all(test, not(miri)))] mod tests { - use diskann_utils::views::{Matrix, MatrixView}; - use diskann_vector::{Half, distance::Metric}; - use diskann_wide::arch::{self, Target1}; - use rstest::rstest; - - use super::super::{Cosine, CosineNormalized, InnerProduct, L2}; use super::*; + use crate::graph::pipnn::{L2, test_support}; + use diskann_wide::ARCH; + use rstest::rstest; - struct PartitionCall<'a, T> { - data: MatrixView<'a, T>, - config: &'a PiPNNConfig, + // Neither leaf order nor point order within a leaf is part of the result. + // Preserve repeated leaves and IDs so comparison still detects duplicates. + fn sorted_leaf_memberships(leaves: &[Vec]) -> Vec> { + let mut leaves = test_support::sorted_members_per_row(leaves); + leaves.sort_unstable(); + leaves } - struct DispatchPartition(Metric); - - impl Target1>>, PartitionCall<'_, T>> for DispatchPartition - where - A: PiPNNSIMDSchema, - T: VectorRepr + Send + Sync, - { - fn run(self, arch: A, call: PartitionCall<'_, T>) -> ANNResult>> { - match self.0 { - Metric::L2 => partition::(arch, call.data, call.config), - Metric::Cosine => partition::(arch, call.data, call.config), - Metric::CosineNormalized => { - partition::(arch, call.data, call.config) - } - Metric::InnerProduct => { - partition::(arch, call.data, call.config) - } - } + fn splitting_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 2, + c_min: 1, + p_samp: 1.0, + fanout: vec![2], + leaf_k: 1, + replicas: 1, } } - fn partition_with_runtime_metric( - data: MatrixView<'_, T>, - config: &PiPNNConfig, - metric: Metric, - ) -> ANNResult>> - where - T: VectorRepr + Send + Sync, - { - arch::dispatch1_no_features(DispatchPartition(metric), PartitionCall { data, config }) + #[rstest] + #[case::minimum_sample(7, 0.01, 2)] + #[case::rounded_fraction(7, 0.31, 3)] + #[case::all_points(7, 1.0, 7)] + #[case::below_cap(999, 1.0, 999)] + #[case::at_cap(1000, 1.0, 1000)] + #[case::above_cap(1001, 1.0, 1000)] + fn sampled_leaders_are_a_repeatable_subset_without_replacement( + #[case] point_count: u32, + #[case] fraction: f64, + #[case] expected_count: usize, + ) { + let ids: Vec<_> = (0..point_count).map(|id| 10 + 3 * id).collect(); + + let mut leaders = sample_leaders(&ids, fraction, 1290); + + assert_eq!(leaders.len(), expected_count); + assert_eq!( + leaders.iter().copied().collect::>().len(), + expected_count + ); + assert!(leaders.iter().all(|id| ids.contains(id))); + let mut repeated = sample_leaders(&ids, fraction, 1290); + leaders.sort_unstable(); + repeated.sort_unstable(); + assert_eq!(leaders, repeated); } - fn partition_config( - c_min: usize, - c_max: usize, - fanout: Vec, - replicas: usize, - ) -> PiPNNConfig { - PiPNNConfig { - c_max, - c_min, - p_samp: 0.25, - fanout, - leaf_k: 1, - replicas, - } - } + #[rstest] + #[case::nearest_only(1, vec![vec![5, 1], vec![4]])] + #[case::all_leaders(2, vec![vec![5, 4, 1], vec![5, 4, 1]])] + #[case::fanout_above_leader_count(9, vec![vec![5, 4, 1], vec![5, 4, 1]])] + fn points_join_their_nearest_leaders(#[case] fanout: usize, #[case] expected: Vec>) { + // Leader columns are global IDs [2, 0], at x=10 and x=0. + // The requested points [5, 4, 1] are at x=8, x=2 and x=100. + let values = [0.0_f32, 100.0, 10.0, -100.0, 2.0, 8.0]; + let data = MatrixView::try_from(&values[..], 6, 1).unwrap(); + let buffers = StripeBufferPool::new((), 0, None); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); - fn unit_circle_points(points: usize) -> Matrix { - let coordinates: Vec = (0..points) - .flat_map(|point| { - let angle = std::f32::consts::TAU * point as f32 / points as f32; - [angle.cos(), angle.sin()] + let actual = pool + .install(|| { + assign_to_leaders::<_, L2, _>(ARCH, data, &[5, 4, 1], &[2, 0], fanout, &buffers) }) - .collect(); - Matrix::try_from(coordinates.into_boxed_slice(), points, 2).unwrap() - } + .unwrap(); - fn assert_partition_invariants( - leaves: &[Vec], - points: usize, - c_max: usize, - replicas: usize, - ) { - let mut occurrences = vec![0usize; points]; - for leaf in leaves { - assert!(!leaf.is_empty() && leaf.len() <= c_max); - let unique_ids: HashSet<_> = leaf.iter().copied().collect(); - assert_eq!(unique_ids.len(), leaf.len(), "duplicate ID inside a leaf"); - for &id in leaf { - assert!((id as usize) < points); - occurrences[id as usize] += 1; - } - } - assert!(occurrences.iter().all(|&count| count >= replicas)); + assert_eq!( + test_support::sorted_members_per_row(&actual), + test_support::sorted_members_per_row(&expected) + ); } - mod partition_tests { - use super::*; - - #[rstest] - #[case::single_point(1)] - #[case::below_capacity(7)] - #[case::at_capacity(8)] - fn one_leaf_contains_all_points_when_no_split_is_needed(#[case] point_count: usize) { - // Given - let data = Matrix::new(0.0f32, point_count, 1); - let config = partition_config(2, 8, vec![2], 1); - let expected_leaves = vec![(0..point_count as u32).collect::>()]; - - // When - let leaves = - partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); - - // Then - assert_eq!(leaves, expected_leaves); - } - - #[test] - fn each_replica_emits_its_leaf_when_no_split_is_needed() { - // Given - let data = MatrixView::column_vector(&[0.0f32, 1.0]); - let config = partition_config(1, 2, vec![1], 2); - let expected_leaves = vec![vec![0, 1], vec![0, 1]]; + #[rstest] + #[case::one_worker(1)] + #[case::several_workers(3)] + fn multiple_stripes_and_their_partial_tail_preserve_all_assignments(#[case] workers: usize) { + // More than two maximum-sized stripes, with three points left over. + // Repeated dense rows at 0, 4 and 10 have unambiguous nearest centers. + let point_count = 2 * MAX_ASSIGNMENT_STRIPE_POINTS + 3; + let dimensions = 384; + let values: Vec<_> = (0..point_count) + .flat_map(|point| std::iter::repeat_n([0.0_f32, 4.0, 10.0][point % 3], dimensions)) + .collect(); + let data = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let ids: Vec<_> = (0..point_count as u32).rev().collect(); + let expected = vec![ + ids.iter() + .copied() + .filter(|id| id % 3 == 2) + .collect::>(), + ids.iter() + .copied() + .filter(|id| id % 3 != 2) + .collect::>(), + ]; + let buffers = StripeBufferPool::new((), 0, None); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(workers) + .build() + .unwrap(); - // When - let leaves = partition_with_runtime_metric(data, &config, Metric::L2).unwrap(); + let actual = pool + .install(|| assign_to_leaders::<_, L2, _>(ARCH, data, &ids, &[2, 0], 1, &buffers)) + .unwrap(); - // Then - assert_eq!(leaves, expected_leaves); - } + assert_eq!( + test_support::sorted_members_per_row(&actual), + test_support::sorted_members_per_row(&expected) + ); + } - #[test] - fn leaf_order_and_point_order_are_independent_of_worker_count() { - // Given: both runs split the same replicas through several levels. - let data = unit_circle_points(64); - let config = partition_config(2, 8, vec![3, 2], 2); - let single_worker = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - let three_workers = rayon::ThreadPoolBuilder::new() - .num_threads(3) - .build() - .unwrap(); - - // When - let serial = single_worker.install(|| { - partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap() - }); - let parallel = three_workers.install(|| { - partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap() - }); + #[test] + fn reused_stripe_buffers_replace_previous_points_and_assignment_width() { + let values = [100.0_f32, 0.0, 8.0, 10.0, 2.0]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + let leader_values = [0.0_f32, 10.0, 100.0]; + let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); + let mut buffers = StripeBuffers::default(); + + for (ids, fanout, expected) in [ + (&[4, 2, 0, 3, 1][..], 2, &[0, 1, 1, 0, 2, 1, 1, 0, 0, 1][..]), + (&[2, 1][..], 1, &[1, 0][..]), + (&[1, 0, 4, 2][..], 2, &[0, 1, 2, 1, 0, 1, 1, 0][..]), + ] { + let mut assignments = vec![0; ids.len() * fanout]; + + assign_point_stripe::<_, L2, _>( + ARCH, + data, + ids, + &leaders, + fanout, + &mut buffers, + &mut assignments, + ) + .unwrap(); - // Then: compare the full order; sorting would hide changes to recursive sampling. - assert_eq!(parallel, serial); + let actual: Vec<_> = assignments + .chunks_exact(fanout) + .map(<[u32]>::to_vec) + .collect(); + let expected: Vec<_> = expected.chunks_exact(fanout).map(<[u32]>::to_vec).collect(); + assert_eq!( + test_support::sorted_members_per_row(&actual), + test_support::sorted_members_per_row(&expected), + "points {ids:?}, fanout={fanout}" + ); } + } - #[rstest] - #[case::l2(Metric::L2)] - #[case::cosine(Metric::Cosine)] - #[case::normalized_cosine(Metric::CosineNormalized)] - #[case::inner_product(Metric::InnerProduct)] - fn bounded_leaves_preserve_point_coverage(#[case] metric: Metric) { - // Given: distinct unit vectors allow each metric to make progress. - let data = unit_circle_points(64); - let config = partition_config(2, 16, vec![3, 2], 3); - - // When - let leaves = partition_with_runtime_metric(data.as_view(), &config, metric).unwrap(); - - // Then - assert_partition_invariants(&leaves, 64, 16, 3); - } + #[rstest] + #[case::small_input(7)] + #[case::below_parallel_threshold(PARALLEL_SCATTER_MIN_POINTS - 1)] + #[case::at_parallel_threshold(PARALLEL_SCATTER_MIN_POINTS)] + #[case::uneven_parallel_chunks(PARALLEL_SCATTER_MIN_POINTS + 1)] + fn scatter_groups_assigned_points_by_leader_and_omits_unassigned_slots( + #[case] point_count: usize, + ) { + let ids: Vec<_> = (0..point_count).map(|i| 1_000_000 - i as u32).collect(); + let pattern = [ + [2, 0], + [UNASSIGNED_LEADER, 1], + [0, UNASSIGNED_LEADER], + [UNASSIGNED_LEADER, UNASSIGNED_LEADER], + ]; + let assignments: Vec<_> = (0..point_count).flat_map(|i| pattern[i % 4]).collect(); + let expected = vec![ + ids.iter().copied().step_by(2).collect::>(), + ids.iter().copied().skip(1).step_by(4).collect(), + ids.iter().copied().step_by(4).collect(), + vec![], + ]; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(3) + .build() + .unwrap(); - #[test] - fn identical_points_return_an_iteration_limit_error() { - // Given: every point ties for the same first leader, so splitting cannot shrink it. - let data = Matrix::new(1.0f32, 5, 1); - let config = partition_config(1, 2, vec![1], 1); - - // When - let error = partition_with_runtime_metric(data.as_view(), &config, Metric::L2) - .unwrap_err() - .downcast::() - .unwrap(); - - // Then - assert!(matches!( - error, - PartitionError::IterationLimit { size: 5, .. } - )); - } + let actual = pool.install(|| scatter_assignments(&ids, &assignments, 2, 4)); + + assert_eq!( + test_support::sorted_members_per_row(&actual), + test_support::sorted_members_per_row(&expected) + ); + } - #[test] - fn no_leaves_are_emitted_when_every_distance_is_non_rankable() { - // Given: a split is required, but no point can be assigned to any leader. - let data = Matrix::new(f32::NAN, 3, 1); - let config = partition_config(1, 2, vec![1], 1); + #[rstest] + #[case::configured_overlap(0, vec![vec![0, 1], vec![0, 1], vec![2, 3], vec![2, 3]])] + #[case::after_schedule(1, vec![vec![0], vec![1], vec![2], vec![3]])] + fn a_split_uses_the_configured_fanout_then_falls_back_to_one_leader( + #[case] level: usize, + #[case] expected: Vec>, + ) { + let values = [0.0_f32, 1.0, 10.0, 11.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let config = splitting_config(); + let parent = PendingPartition { + point_ids: vec![0, 1, 2, 3], + level, + seed: 1290, + }; + let buffers = StripeBufferPool::new((), 0, None); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); - // When - let leaves = - partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + let actual = pool + .install(|| split_partition::<_, L2, _>(ARCH, data, &config, parent, &buffers)) + .unwrap(); - // Then - assert!(leaves.is_empty()); - } + // Every point is a leader, so membership is fixed regardless of output order. + assert_eq!( + sorted_leaf_memberships(&actual.leaves), + sorted_leaf_memberships(&expected) + ); + assert!(actual.pending.is_empty()); } #[test] - fn split_assigns_one_leader_after_the_fanout_schedule_ends() { - // Given: all points are leaders; at level one the only configured fanout is exhausted. - let data = MatrixView::column_vector(&[0.0f32, 10.0, 20.0, 30.0]); + fn children_above_cmax_are_queued_at_the_next_level() { + let values = [0.0_f32, 1.0, 10.0, 11.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); let config = PiPNNConfig { - p_samp: 1.0, - ..partition_config(1, 2, vec![2], 1) + c_max: 1, + ..splitting_config() }; let parent = PendingPartition { point_ids: vec![0, 1, 2, 3], - level: 1, - seed: PARTITION_SEED, + level: 0, + seed: 1290, }; - let expected_leaves = vec![vec![0], vec![1], vec![2], vec![3]]; - - // When: each point's unique nearest leader is itself. - let mut split = split_partition::<_, L2, _>( - diskann_wide::ARCH, - data, - &config, - parent, - &StripeBufferPool::new((), 0, None), - ) - .unwrap(); - split.leaves.sort(); - - // Then - assert!(split.pending.is_empty()); - assert_eq!(split.leaves, expected_leaves); - } - - mod assign_to_leaders_tests { - use super::*; - - #[test] - fn clusters_contain_global_point_ids_in_input_order() { - // Given: the leader-column order differs from global ID order. - let data = MatrixView::column_vector(&[0.0f32, 100.0, 4.0, 96.0]); - let points = [3, 2, 1]; - let leaders = [1, 0]; - // Points 3 and 1 are nearest coordinate 100; point 2 is nearest coordinate 0. - let expected_clusters = vec![vec![3, 1], vec![2]]; - - // When - let clusters = assign_to_leaders::<_, L2, _>( - diskann_wide::ARCH, - data, - &points, - &leaders, - 1, - &StripeBufferPool::new((), 0, None), - ) + let buffers = StripeBufferPool::new((), 0, None); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() .unwrap(); - // Then - assert_eq!(clusters, expected_clusters); - } + let actual = pool + .install(|| split_partition::<_, L2, _>(ARCH, data, &config, parent, &buffers)) + .unwrap(); - #[test] - fn fanout_larger_than_the_leader_set_assigns_every_leader_once() { - // Given - let data = MatrixView::column_vector(&[0.0f32, 10.0, 4.0]); - let expected_clusters = vec![vec![2], vec![2]]; + assert!(actual.leaves.is_empty()); + assert!(actual.pending.iter().all(|child| child.level == 1)); + let ids: Vec<_> = actual + .pending + .into_iter() + .map(|child| child.point_ids) + .collect(); + assert_eq!( + sorted_leaf_memberships(&ids), + [vec![0, 1], vec![0, 1], vec![2, 3], vec![2, 3]] + ); + } - // When - let clusters = assign_to_leaders::<_, L2, _>( - diskann_wide::ARCH, - data, - &[2], - &[0, 1], - 3, - &StripeBufferPool::new((), 0, None), - ) + #[rstest] + #[case::below_limit(4)] + #[case::at_limit(3)] + fn data_within_cmax_forms_one_complete_leaf_per_replica(#[case] c_max: usize) { + let values = [0.0_f32, 1.0, 10.0]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let config = PiPNNConfig { + c_max, + replicas: 2, + ..splitting_config() + }; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() .unwrap(); - // Then - assert_eq!(clusters, expected_clusters); - } + let actual = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); - #[test] - fn points_without_rankable_distances_are_omitted() { - // Given: the NaN point has no leader assignments; the finite point has two. - let data = MatrixView::column_vector(&[0.0f32, 10.0, f32::NAN, 1.0]); - let expected_clusters = vec![vec![3], vec![3]]; + assert_eq!( + sorted_leaf_memberships(&actual), + [vec![0, 1, 2], vec![0, 1, 2]] + ); + } - // When - let clusters = assign_to_leaders::<_, L2, _>( - diskann_wide::ARCH, - data, - &[2, 3], - &[0, 1], - 2, - &StripeBufferPool::new((), 0, None), - ) + #[test] + fn recursive_partitioning_preserves_leaf_membership_across_worker_counts() { + let point_count = 129; + let values = test_support::dense_points(point_count, 17, 1290); + let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.2, + fanout: vec![2, 2], + replicas: 2, + ..splitting_config() + }; + let serial = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + let parallel = rayon::ThreadPoolBuilder::new() + .num_threads(3) + .build() .unwrap(); - // Then - assert_eq!(clusters, expected_clusters); - } - - #[rstest] - fn partial_final_stripe_preserves_point_assignments(#[values(1, 2)] worker_count: usize) { - // Given: three stripes, the last containing one point. One worker reuses its lease. - let point_count = 2 * MAX_ASSIGNMENT_STRIPE_POINTS + 1; - let data: Vec = (0..point_count).map(|point| point as f32).collect(); - let points: Vec = (0..point_count as u32).collect(); - let last_point = point_count as u32 - 1; - let leaders = [0, last_point]; - // On the line [0, last_point], the midpoint ties and selects leader column 0. - let midpoint = last_point / 2; - let expected_clusters = vec![ - (0..=midpoint).collect::>(), - (midpoint + 1..=last_point).collect(), - ]; - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(worker_count) - .build() - .unwrap(); - - // When - let clusters = pool.install(|| { - assign_to_leaders::<_, L2, _>( - diskann_wide::ARCH, - MatrixView::column_vector(&data), - &points, - &leaders, - 1, - &StripeBufferPool::new((), 0, None), - ) - .unwrap() - }); + let first = serial + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + let second = parallel + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); - // Then - assert_eq!(clusters, expected_clusters); + assert_eq!( + sorted_leaf_memberships(&first), + sorted_leaf_memberships(&second) + ); + let mut ids: Vec<_> = first.iter().flatten().copied().collect(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids, (0..point_count as u32).collect::>()); + for leaf in first { + assert!( + !leaf.is_empty() && leaf.len() <= config.c_max, + "leaf {leaf:?}" + ); + assert_eq!( + leaf.iter().copied().collect::>().len(), + leaf.len(), + "duplicate IDs in leaf {leaf:?}" + ); } } #[rstest] - #[case::f32([-2.0f32, 1.0, 3.0, -4.0, 5.0, 6.0], [5.0, 6.0, -2.0, 1.0, 5.0, 6.0])] - #[case::f16([-2.0f32, 1.0, 3.0, -4.0, 5.0, 6.0].map(Half::from_f32), [5.0, 6.0, -2.0, 1.0, 5.0, 6.0])] - #[case::i8([-2i8, 1, 3, -4, 5, 6], [5.0, 6.0, -2.0, 1.0, 5.0, 6.0])] - #[case::u8([2u8, 1, 3, 4, 5, 6], [5.0, 6.0, 2.0, 1.0, 5.0, 6.0])] - fn gather_converts_rows_in_requested_id_order( - #[case] values: [T; 6], - #[case] expected: [f32; 6], + #[case::one_unrankable([0.0, 3.0, f32::NAN], vec![vec![0], vec![1]])] + #[case::all_unrankable([f32::NAN; 3], vec![])] + fn unrankable_points_do_not_form_child_leaves( + #[case] values: [f32; 3], + #[case] expected: Vec>, ) { - // Given: request the final row, then the first row, then repeat the final row. - let data = MatrixView::try_from(values.as_slice(), 3, 2).unwrap(); - let mut gathered = [f32::NAN; 6]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let config = PiPNNConfig { + c_max: 1, + fanout: vec![1], + ..splitting_config() + }; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); - // When - gather_vectors(data, &[2, 0, 2], &mut gathered).unwrap(); + let actual = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); - // Then - assert_eq!(gathered, expected); + assert_eq!( + sorted_leaf_memberships(&actual), + sorted_leaf_memberships(&expected) + ); } - mod merge_undersized_leaves_tests { - use super::*; - - #[test] - fn total_below_minimum_is_kept_as_one_leaf() { - // Given: all available points together still fall below the minimum. - let leaves = vec![vec![0], vec![1]]; - let expected_leaves = vec![vec![0, 1]]; - - // When - let merged = merge_undersized_leaves(leaves, 3, 4); - - // Then - assert_eq!(merged, expected_leaves); - } - - #[test] - fn leaves_already_at_minimum_keep_their_order() { - // Given - let leaves = vec![vec![3, 1, 2], vec![7, 5, 6]]; - let expected_leaves = leaves.clone(); - - // When - let merged = merge_undersized_leaves(leaves, 3, 4); - - // Then - assert_eq!(merged, expected_leaves); - } - - #[test] - fn merged_leaves_contain_sorted_unique_ids() { - // Given: the shared ID 1 counts only once toward the minimum of three. - let leaves = vec![vec![3, 1], vec![1, 2]]; - let expected_leaves = vec![vec![1, 2, 3]]; - - // When - let merged = merge_undersized_leaves(leaves, 3, 4); - - // Then - assert_eq!(merged, expected_leaves); - } - - #[test] - fn full_batch_is_emitted_before_another_leaf_would_exceed_capacity() { - // Given: three pairs cannot fit in a five-point leaf; the first two can. - let leaves = vec![vec![0, 1], vec![2, 3], vec![4, 5]]; - let expected_leaves = vec![vec![0, 1, 2, 3], vec![4, 5]]; - - // When - let merged = merge_undersized_leaves(leaves, 5, 5); - - // Then - assert_eq!(merged, expected_leaves); - } - - #[test] - fn leaves_that_exactly_fill_capacity_are_combined() { - // Given - let leaves = vec![vec![0, 1], vec![2, 3]]; - let expected_leaves = vec![vec![0, 1, 2, 3]]; - - // When - let merged = merge_undersized_leaves(leaves, 4, 4); - - // Then - assert_eq!(merged, expected_leaves); - } - - #[test] - fn overlapping_remainder_fits_after_duplicate_ids_are_removed() { - // Given: raw lengths total five, but the union has four IDs and fits. - let leaves = vec![vec![0, 1, 2], vec![2, 3]]; - let expected_leaves = vec![vec![0, 1, 2, 3]]; - - // When - let merged = merge_undersized_leaves(leaves, 3, 4); - - // Then - assert_eq!(merged, expected_leaves); - } - - #[test] - fn new_remainder_ids_form_a_separate_leaf_when_the_last_leaf_is_full() { - // Given: ID 3 already exists, but the new ID 4 cannot fit in the full leaf. - let leaves = vec![vec![0, 1, 2, 3], vec![3, 4]]; - let expected_leaves = vec![vec![0, 1, 2, 3], vec![4]]; - - // When - let merged = merge_undersized_leaves(leaves, 3, 4); - - // Then - assert_eq!(merged, expected_leaves); - } - - #[test] - fn entirely_overlapping_remainder_adds_no_leaf() { - // Given: every remainder ID is already present in the last leaf. - let leaves = vec![vec![0, 1, 2], vec![1, 2]]; - let expected_leaves = vec![vec![0, 1, 2]]; + #[test] + fn indistinguishable_points_report_the_oversized_cluster_when_splitting_stalls() { + let values = [1.0_f32; 5]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + let config = PiPNNConfig { + fanout: vec![1], + ..splitting_config() + }; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); - // When - let merged = merge_undersized_leaves(leaves, 3, 4); + let error = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap_err(); - // Then - assert_eq!(merged, expected_leaves); - } + assert_eq!( + error.downcast_ref::().unwrap(), + &PartitionError::IterationLimit { + size: 5, + level: MAX_PARTITION_ITERATIONS, + limit: MAX_PARTITION_ITERATIONS, + } + ); } #[rstest] - #[case::minimum_two_leaders(10, 0.01, 2)] - #[case::fraction_rounds_up(10, 0.25, 3)] - #[case::leader_cap(50_000, 1.0, LEADER_CAP)] - fn sampled_leader_count_respects_sampling_bounds( - #[case] point_count: usize, - #[case] sampling_fraction: f64, - #[case] expected_leader_count: usize, + #[case::empty(vec![], 2, 4, vec![])] + #[case::large_leaves_unchanged(vec![vec![1, 4], vec![0, 3, 6, 8]], 2, 4, + vec![vec![1, 4], vec![0, 3, 6, 8]])] + #[case::overlapping_small_leaves(vec![vec![4], vec![1], vec![4], vec![2]], 3, 4, + vec![vec![1, 2, 4]])] + #[case::flush_before_exceeding_maximum(vec![vec![4, 5], vec![0, 1]], 3, 3, + vec![vec![4, 5], vec![0, 1]])] + #[case::remainder_fits_after_deduplication(vec![vec![1, 3, 5], vec![3], vec![6]], 3, 4, + vec![vec![1, 3, 5, 6]])] + #[case::remainder_does_not_fit(vec![vec![1, 3, 5], vec![6], vec![8]], 3, 4, + vec![vec![1, 3, 5], vec![6, 8]])] + #[case::only_an_undersized_remainder(vec![vec![7], vec![2]], 3, 4, vec![vec![2, 7]])] + fn small_leaf_merging_deduplicates_ids_without_exceeding_cmax( + #[case] leaves: Vec>, + #[case] c_min: usize, + #[case] c_max: usize, + #[case] expected: Vec>, ) { - // When - let count = sampled_leader_count(point_count, sampling_fraction); + let actual = merge_undersized_leaves(leaves, c_min, c_max); - // Then - assert_eq!(count, expected_leader_count); - } - - mod scatter_assignments_tests { - use super::*; - - #[test] - fn unassigned_slots_are_skipped_without_reordering_points() { - // Given: each assignment pair belongs to one point; leader 2 has no points. - let points = [30, 10, 20]; - let assignments = [1, UNASSIGNED_LEADER, 0, 1, UNASSIGNED_LEADER, 0]; - let expected_clusters = vec![vec![10, 20], vec![30, 10], vec![]]; - - // When - let clusters = scatter_assignments(&points, &assignments, 2, 3); - - // Then - assert_eq!(clusters, expected_clusters); - } - - #[rstest] - fn parallel_stripes_preserve_order_and_skip_unassigned_slots( - #[values(1, 3)] worker_count: usize, - ) { - // Given: three repeating assignments, with a partial final worker chunk. - let points: Vec = (0..(PARALLEL_SCATTER_MIN_POINTS as u32 + 1)) - .rev() - .collect(); - let pattern = [0, UNASSIGNED_LEADER, 1, 0, UNASSIGNED_LEADER, 1]; - let assignments: Vec = - pattern.into_iter().cycle().take(2 * points.len()).collect(); - // Leader 0 receives rows 0 and 1 of each triple; leader 1 receives rows 1 and 2. - let expected_clusters = vec![ - points - .iter() - .enumerate() - .filter(|(row, _)| row % 3 != 2) - .map(|(_, &id)| id) - .collect::>(), - points - .iter() - .enumerate() - .filter(|(row, _)| row % 3 != 0) - .map(|(_, &id)| id) - .collect(), - vec![], - ]; - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(worker_count) - .build() - .unwrap(); - - // When - let clusters = pool.install(|| scatter_assignments(&points, &assignments, 2, 3)); - - // Then - assert_eq!(clusters, expected_clusters); - } + assert_eq!( + sorted_leaf_memberships(&actual), + sorted_leaf_memberships(&expected) + ); } #[test] - fn overflowing_buffer_shape_returns_an_error() { - // Given: this shape cannot be represented by an allocation length. - let rows = usize::MAX; - let columns = 2; - - // When - let error = checked_area("partition assignments", rows, columns).unwrap_err(); - - // Then - assert!(matches!( - error.downcast::().unwrap(), - PartitionError::ShapeOverflow { .. } - )); + fn an_overflowing_buffer_area_identifies_the_failed_shape() { + let error = checked_area("point stripe", usize::MAX, 2).unwrap_err(); + + assert_eq!( + error.downcast_ref::().unwrap(), + &PartitionError::ShapeOverflow { + buffer: "point stripe", + rows: usize::MAX, + cols: 2, + } + ); } } From 83b618d1ad3f6f841238a590d1f3de3f3b5a368e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:39:10 +0000 Subject: [PATCH 192/229] test(pipnn): cover partition composition gaps --- diskann/src/graph/pipnn/mod.rs | 43 ++++++++++++++++++ diskann/src/graph/pipnn/partitioning.rs | 60 ++++++++++++++++++++++++- 2 files changed, 101 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index f4a4203d72..106298638d 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -705,6 +705,49 @@ mod construction_tests { ); } + #[rstest] + #[case::l2(Metric::L2, vec![vec![2], vec![], vec![0], vec![]])] + #[case::cosine(Metric::Cosine, vec![vec![1], vec![0], vec![3], vec![2]])] + #[case::normalized_cosine(Metric::CosineNormalized, vec![vec![1], vec![0], vec![3], vec![2]])] + #[case::inner_product(Metric::InnerProduct, vec![vec![1], vec![0], vec![3], vec![2]])] + fn splitting_uses_the_requested_metric_to_group_points( + #[case] metric: Metric, + #[case] expected: Vec>, + ) { + // Four points force splitting at c_max=2. Every point is sampled, so + // leader order cannot change memberships. With L2, leader 0 gets {0,2} + // and leader 2 gets all four points, which then split into singletons. + // Cosine pairs directions {0,1} and {2,3}. Inner product picks leaders + // {1,2} for points 0/1 and {2,3} for points 2/3; the oversized leader-2 + // cluster then splits into those same pairs using fanout one. + let values = test_support::packed_points( + &[[1.0, 0.0], [6.0, 1.0], [2.0, 4.0], [0.0, 9.0]], + 2, + metric == Metric::CosineNormalized, + ); + let data = MatrixView::try_from(values.as_slice(), 4, 2).unwrap(); + let config = PiPNNConfig { + c_max: 2, + ..partition_policy() + }; + // Every two-point leaf contributes its only pair; degree three retains + // every edge. The graph therefore exposes partition membership directly. + let graph = graph_policy(3, metric).unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + + assert_eq!( + test_support::sorted_members_per_row(&actual), + test_support::sorted_members_per_row(&expected) + ); + } + #[rstest] #[case::signed_leaf_ranking([1_i8, 0, 4, 1, 1, 2, 0, 9], 1, 3)] #[case::unsigned_leaf_ranking([1_u8, 0, 4, 1, 1, 2, 0, 9], 1, 3)] diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 840f749e82..f0ce67fb42 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -614,7 +614,8 @@ mod tests { #[case::several_workers(3)] fn multiple_stripes_and_their_partial_tail_preserve_all_assignments(#[case] workers: usize) { // More than two maximum-sized stripes, with three points left over. - // Repeated dense rows at 0, 4 and 10 have unambiguous nearest centers. + // Each dense row selects two of the three centers at 10, 0 and 4. + // Every row selects 4; its other center is 10 for x=10 and 0 otherwise. let point_count = 2 * MAX_ASSIGNMENT_STRIPE_POINTS + 3; let dimensions = 384; let values: Vec<_> = (0..point_count) @@ -631,6 +632,7 @@ mod tests { .copied() .filter(|id| id % 3 != 2) .collect::>(), + ids.clone(), ]; let buffers = StripeBufferPool::new((), 0, None); let pool = rayon::ThreadPoolBuilder::new() @@ -639,7 +641,7 @@ mod tests { .unwrap(); let actual = pool - .install(|| assign_to_leaders::<_, L2, _>(ARCH, data, &ids, &[2, 0], 1, &buffers)) + .install(|| assign_to_leaders::<_, L2, _>(ARCH, data, &ids, &[2, 0, 1], 2, &buffers)) .unwrap(); assert_eq!( @@ -867,6 +869,60 @@ mod tests { } } + #[test] + fn a_second_replica_adds_new_leaf_memberships() { + let point_count = 33; + let values = test_support::dense_points(point_count, 17, 1290); + let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); + let config = PiPNNConfig { + c_max: 8, + c_min: 1, + p_samp: 0.25, + fanout: vec![1], + ..splitting_config() + }; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let first = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + + let combined = pool + .install(|| { + partition::<_, L2, _>( + ARCH, + data, + &PiPNNConfig { + replicas: 2, + ..config + }, + ) + }) + .unwrap(); + + let first = sorted_leaf_memberships(&first); + let combined = sorted_leaf_memberships(&combined); + assert!(first.iter().all(|leaf| combined.contains(leaf))); + // This fixed cloud and partial sample produce different memberships in + // the second pass. Repeating the first seed would only duplicate leaves. + assert!(combined.iter().any(|leaf| !first.contains(leaf))); + // Fanout one gives every point exactly one membership per replica. + let mut memberships = vec![0; point_count]; + for leaf in &combined { + assert!( + !leaf.is_empty() && leaf.len() <= config.c_max, + "leaf {leaf:?}" + ); + assert_eq!(leaf.iter().collect::>().len(), leaf.len()); + for &id in leaf { + memberships[id as usize] += 1; + } + } + assert_eq!(memberships, vec![2; point_count]); + } + #[rstest] #[case::one_unrankable([0.0, 3.0, f32::NAN], vec![vec![0], vec![1]])] #[case::all_unrankable([f32::NAN; 3], vec![])] From 354918d16b7a0ee18d8ffe28d5fb94463b45e7e6 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:26:11 +0000 Subject: [PATCH 193/229] docs(pipnn): track NEON validation in #1423 --- diskann/src/graph/pipnn/simd.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 687aa19ad7..0cc6b9678f 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -10,9 +10,6 @@ use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDVector}; /// Default SIMD representation used by both PiPNN ranking kernels. /// /// This alias is the single build-time width selection. -/// -/// TODO: This 16-lane choice lacks Neon-specific validation. -/// Add Neon coverage in a follow-up. type DefaultVector = ::f32x16; /// PiPNN SIMD representation for one architecture. From 08f5fd09246e2d29226b111cee7a0f31eac54541 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 02:55:52 +0000 Subject: [PATCH 194/229] docs(linalg): clarify lower-product storage --- diskann-linalg/src/lib.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index ae452709c6..6c63a0bccb 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -191,8 +191,9 @@ pub fn sgemm( /// Replace the lower triangle of `C` with $\alpha A A^\mathsf{T}$. /// -/// `A` is a dense row-major $m \times k$ matrix. The function does not change -/// the upper triangle of `C`. Use this operation when `C` has no prior term. +/// `A` is a dense row-major $m \times k$ matrix. `C` is a dense row-major +/// $m \times m$ matrix backed by exactly `m * m` elements. The function does not +/// change the upper triangle of `C`. Use this operation when `C` has no prior term. /// /// # Errors /// @@ -214,8 +215,9 @@ pub fn sgemm_aat_lower( /// Add the lower triangle of $\alpha A A^\mathsf{T}$ to `C`. /// -/// `A` is a dense row-major $m \times k$ matrix. The function changes only the -/// lower triangle of `C`. Use this operation when `C` contains an initial term. +/// `A` is a dense row-major $m \times k$ matrix. `C` is a dense row-major +/// $m \times m$ matrix backed by exactly `m * m` elements. The function changes +/// only the lower triangle of `C`. Use this operation when `C` contains an initial term. /// /// # Errors /// From 8b85bf7783237bde0e0553d452cf4f9ca962f9a2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 03:27:22 +0000 Subject: [PATCH 195/229] test(linalg): simplify lower-product assertions --- diskann-linalg/src/lib.rs | 36 +++++++++++++++++++++--------------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index 6c63a0bccb..85e3c07cdb 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -728,9 +728,9 @@ mod sgemm_aat_lower_tests { type LowerProduct = fn(usize, usize, f32, &[f32], &mut [f32]) -> Result<(), SgemmError>; - #[track_caller] fn assert_lower_product_matches_scalar( shape: (usize, usize), + alpha: f32, input: &[f32], previous: &[f32], output: &[f32], @@ -742,23 +742,20 @@ mod sgemm_aat_lower_tests { let index = row * rows + column; if column > row { assert_eq!( - output[index].to_bits(), - previous[index].to_bits(), + output[index], previous[index], "shape={shape:?}, upper ({row},{column})" ); } else { - let dot: f64 = (0..dimensions) + let dot: f32 = (0..dimensions) .map(|dimension| { - f64::from(input[row * dimensions + dimension]) - * f64::from(input[column * dimensions + dimension]) + input[row * dimensions + dimension] + * input[column * dimensions + dimension] }) .sum(); - // Both sweeps use alpha = -2 and exactly representable arithmetic. - let expected = -2.0 * dot + if add { f64::from(previous[index]) } else { 0.0 }; + let expected = alpha * dot + if add { previous[index] } else { 0.0 }; assert_eq!( - f64::from(output[index]), - expected, - "shape={shape:?}, lower ({row},{column})" + output[index], expected, + "shape={shape:?}, alpha={alpha}, add={add}, lower ({row},{column})" ); } } @@ -794,15 +791,23 @@ mod sgemm_aat_lower_tests { #[case] rows: usize, #[case] dimensions: usize, ) { + let alpha = -2.0; let input: Vec<_> = (0..rows * dimensions) .map(|i| (i % 9) as f32 - 4.0) .collect(); let previous: Vec<_> = (0..rows * rows).map(|i| 100.0 + i as f32).collect(); let mut output = previous.clone(); - operation(rows, dimensions, -2.0, &input, &mut output).unwrap(); + operation(rows, dimensions, alpha, &input, &mut output).unwrap(); - assert_lower_product_matches_scalar((rows, dimensions), &input, &previous, &output, add); + assert_lower_product_matches_scalar( + (rows, dimensions), + alpha, + &input, + &previous, + &output, + add, + ); } #[rstest] @@ -817,6 +822,7 @@ mod sgemm_aat_lower_tests { // One representative dense input checks the wrapper's row-major wiring. let shape = (129, 1536); let (rows, dimensions) = shape; + let alpha = -2.0; let mut rng = StdRng::seed_from_u64(1287); // Bounded multiples of 1/8 keep products and sums exact in f32. let mut input: Vec<_> = (0..rows * dimensions) @@ -828,9 +834,9 @@ mod sgemm_aat_lower_tests { let previous: Vec<_> = (0..rows * rows).map(|i| (i % 97) as f32 - 48.0).collect(); let mut output = previous.clone(); - operation(rows, dimensions, -2.0, &input, &mut output).unwrap(); + operation(rows, dimensions, alpha, &input, &mut output).unwrap(); - assert_lower_product_matches_scalar(shape, &input, &previous, &output, add); + assert_lower_product_matches_scalar(shape, alpha, &input, &previous, &output, add); } #[test] From 8e2d9d6f8b92fc0521b34fa4472d7e65d522892b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 04:40:36 +0000 Subject: [PATCH 196/229] test(pipnn): consolidate numeric sweeps --- Cargo.lock | 1 - diskann-linalg/Cargo.toml | 1 - diskann-linalg/src/lib.rs | 299 ++++++++++--------- diskann/src/graph/pipnn/leaf_kernel.rs | 313 ++++++++++---------- diskann/src/graph/pipnn/leaf_metric.rs | 147 +++++---- diskann/src/graph/pipnn/partition_kernel.rs | 283 +++++++++--------- diskann/src/graph/pipnn/partition_metric.rs | 245 ++++++++------- diskann/src/graph/pipnn/topk.rs | 190 ++++++------ 8 files changed, 793 insertions(+), 686 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5cc9e8288b..85a46b7e09 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -677,7 +677,6 @@ dependencies = [ "faer", "rand", "rand_distr", - "rstest", "serde", "serde_json", "thiserror 2.0.17", diff --git a/diskann-linalg/Cargo.toml b/diskann-linalg/Cargo.toml index c798185b03..d60cbf3031 100644 --- a/diskann-linalg/Cargo.toml +++ b/diskann-linalg/Cargo.toml @@ -19,7 +19,6 @@ clippy.undocumented_unsafe_blocks = "warn" clippy.uninlined_format_args = "allow" [dev-dependencies] -rstest.workspace = true serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } approx = { workspace = true } diff --git a/diskann-linalg/src/lib.rs b/diskann-linalg/src/lib.rs index 85e3c07cdb..659913e37d 100644 --- a/diskann-linalg/src/lib.rs +++ b/diskann-linalg/src/lib.rs @@ -724,7 +724,6 @@ mod tests { #[cfg(test)] mod sgemm_aat_lower_tests { use super::*; - use rstest::rstest; type LowerProduct = fn(usize, usize, f32, &[f32], &mut [f32]) -> Result<(), SgemmError>; @@ -743,7 +742,7 @@ mod sgemm_aat_lower_tests { if column > row { assert_eq!( output[index], previous[index], - "shape={shape:?}, upper ({row},{column})" + "shape={shape:?}, alpha={alpha}, add={add}, upper ({row},{column})" ); } else { let dot: f32 = (0..dimensions) @@ -762,81 +761,89 @@ mod sgemm_aat_lower_tests { } } - #[rstest] - #[case::replace(sgemm_aat_lower, [-8.0, 101.0, 102.0, 8.0, -26.0, 103.0, 12.0, -36.0, -50.0])] - #[case::add(sgemm_aat_lower_add, [2.0, 101.0, 102.0, 28.0, 4.0, 103.0, 52.0, 14.0, 10.0])] - fn scaled_products_change_only_the_diagonal_and_lower_triangle( - #[case] operation: LowerProduct, - #[case] expected: [f32; 9], - ) { - // Rows [2,0], [-2,3], [-3,4] have lower Gram entries 4,-4,13,-6,18,25. - let input = [2.0, 0.0, -2.0, 3.0, -3.0, 4.0]; - let mut output = [10.0, 101.0, 102.0, 20.0, 30.0, 103.0, 40.0, 50.0, 60.0]; + #[test] + fn scaled_products_change_only_the_diagonal_and_lower_triangle() { + for (name, operation, expected) in [ + ( + "replace", + sgemm_aat_lower as LowerProduct, + [-8.0, 101.0, 102.0, 8.0, -26.0, 103.0, 12.0, -36.0, -50.0], + ), + ( + "add", + sgemm_aat_lower_add as LowerProduct, + [2.0, 101.0, 102.0, 28.0, 4.0, 103.0, 52.0, 14.0, 10.0], + ), + ] { + // Rows [2,0], [-2,3], [-3,4] have lower Gram entries 4,-4,13,-6,18,25. + let input = [2.0, 0.0, -2.0, 3.0, -3.0, 4.0]; + let mut output = [10.0, 101.0, 102.0, 20.0, 30.0, 103.0, 40.0, 50.0, 60.0]; - operation(3, 2, -2.0, &input, &mut output).unwrap(); + operation(3, 2, -2.0, &input, &mut output) + .unwrap_or_else(|error| panic!("{name}: {error}")); - assert_eq!(output, expected); + assert_eq!(output, expected, "{name}"); + } } - #[rstest] - #[case::replace_single_entry(sgemm_aat_lower, false, 1, 1)] - #[case::add_single_entry(sgemm_aat_lower_add, true, 1, 1)] - #[case::replace_zero_inner_dimension(sgemm_aat_lower, false, 3, 0)] - #[case::add_zero_inner_dimension(sgemm_aat_lower_add, true, 3, 0)] - #[case::replace_rectangular(sgemm_aat_lower, false, 3, 7)] - #[case::add_rectangular(sgemm_aat_lower_add, true, 3, 7)] - fn rectangular_inputs_match_scalar_dot_products( - #[case] operation: LowerProduct, - #[case] add: bool, - #[case] rows: usize, - #[case] dimensions: usize, - ) { - let alpha = -2.0; - let input: Vec<_> = (0..rows * dimensions) - .map(|i| (i % 9) as f32 - 4.0) - .collect(); - let previous: Vec<_> = (0..rows * rows).map(|i| 100.0 + i as f32).collect(); - let mut output = previous.clone(); - - operation(rows, dimensions, alpha, &input, &mut output).unwrap(); - - assert_lower_product_matches_scalar( - (rows, dimensions), - alpha, - &input, - &previous, - &output, - add, - ); + #[test] + fn rectangular_inputs_match_scalar_dot_products() { + for (rows, dimensions) in [(1, 1), (3, 0), (3, 7)] { + for (mode, operation, add) in [ + ("replace", sgemm_aat_lower as LowerProduct, false), + ("add", sgemm_aat_lower_add as LowerProduct, true), + ] { + let alpha = -2.0; + let input: Vec<_> = (0..rows * dimensions) + .map(|i| (i % 9) as f32 - 4.0) + .collect(); + let previous: Vec<_> = (0..rows * rows).map(|i| 100.0 + i as f32).collect(); + let mut output = previous.clone(); + + operation(rows, dimensions, alpha, &input, &mut output).unwrap_or_else(|error| { + panic!("{mode}, rows={rows}, dimensions={dimensions}: {error}") + }); + + assert_lower_product_matches_scalar( + (rows, dimensions), + alpha, + &input, + &previous, + &output, + add, + ); + } + } } - #[rstest] - #[case::replace(sgemm_aat_lower, false)] - #[case::add(sgemm_aat_lower_add, true)] - fn large_dense_inputs_match_scalar_dot_products( - #[case] operation: LowerProduct, - #[case] add: bool, - ) { + #[test] + fn large_dense_inputs_match_scalar_dot_products() { use rand::{rngs::StdRng, SeedableRng}; - // One representative dense input checks the wrapper's row-major wiring. - let shape = (129, 1536); - let (rows, dimensions) = shape; - let alpha = -2.0; - let mut rng = StdRng::seed_from_u64(1287); - // Bounded multiples of 1/8 keep products and sums exact in f32. - let mut input: Vec<_> = (0..rows * dimensions) - .map(|_| rng.random_range(-16..=16) as f32 / 8.0) - .collect(); - for (row, values) in input.chunks_exact_mut(dimensions).enumerate() { - values[dimensions - 1] = 8.0 + (row % 7) as f32; - } - let previous: Vec<_> = (0..rows * rows).map(|i| (i % 97) as f32 - 48.0).collect(); - let mut output = previous.clone(); + for (name, operation, add) in [ + ("replace", sgemm_aat_lower as LowerProduct, false), + ("add", sgemm_aat_lower_add as LowerProduct, true), + ] { + // One representative dense input checks the wrapper's row-major wiring. + let shape = (129, 1536); + let (rows, dimensions) = shape; + let alpha = -2.0; + let mut rng = StdRng::seed_from_u64(1287); + // Bounded multiples of 1/8 keep products and sums exact in f32. + let mut input: Vec<_> = (0..rows * dimensions) + .map(|_| rng.random_range(-16..=16) as f32 / 8.0) + .collect(); + for (row, values) in input.chunks_exact_mut(dimensions).enumerate() { + values[dimensions - 1] = 8.0 + (row % 7) as f32; + } + let previous: Vec<_> = (0..rows * rows).map(|i| (i % 97) as f32 - 48.0).collect(); + let mut output = previous.clone(); - operation(rows, dimensions, alpha, &input, &mut output).unwrap(); + operation(rows, dimensions, alpha, &input, &mut output) + .unwrap_or_else(|error| panic!("{name}: {error}")); - assert_lower_product_matches_scalar(shape, alpha, &input, &previous, &output, add); + assert_lower_product_matches_scalar(shape, alpha, &input, &previous, &output, add); + } } #[test] @@ -849,81 +856,107 @@ mod sgemm_aat_lower_tests { assert_eq!(output, [5.0, 17.0, 5.0, 25.0]); } - #[rstest] - #[case::replace(sgemm_aat_lower, [0.0, 17.0, 0.0, 0.0])] - #[case::add(sgemm_aat_lower_add, [3.0, 17.0, 5.0, 7.0])] - fn zero_scale_respects_the_accumulation_mode( - #[case] operation: LowerProduct, - #[case] expected: [f32; 4], - ) { - let mut output = [3.0, 17.0, 5.0, 7.0]; + #[test] + fn zero_scale_respects_the_accumulation_mode() { + for (name, operation, expected) in [ + ( + "replace", + sgemm_aat_lower as LowerProduct, + [0.0, 17.0, 0.0, 0.0], + ), + ( + "add", + sgemm_aat_lower_add as LowerProduct, + [3.0, 17.0, 5.0, 7.0], + ), + ] { + let mut output = [3.0, 17.0, 5.0, 7.0]; - operation(2, 1, 0.0, &[2.0, 4.0], &mut output).unwrap(); + operation(2, 1, 0.0, &[2.0, 4.0], &mut output) + .unwrap_or_else(|error| panic!("{name}: {error}")); - assert_eq!(output, expected); + assert_eq!(output, expected, "{name}"); + } } - #[rstest] - #[case::replace(sgemm_aat_lower)] - #[case::add(sgemm_aat_lower_add)] - fn empty_matrices_require_no_storage(#[case] operation: LowerProduct) { - assert_eq!(operation(0, 3, 2.0, &[], &mut []), Ok(())); + #[test] + fn empty_matrices_require_no_storage() { + for (name, operation) in [ + ("replace", sgemm_aat_lower as LowerProduct), + ("add", sgemm_aat_lower_add as LowerProduct), + ] { + assert_eq!(operation(0, 3, 2.0, &[], &mut []), Ok(()), "{name}"); + } } - #[rstest] - #[case::short_input(3, 4, MatrixName::A, 2, 2, 3)] - #[case::long_input(5, 4, MatrixName::A, 2, 2, 5)] - #[case::short_output(4, 3, MatrixName::C, 2, 2, 3)] - #[case::long_output(4, 5, MatrixName::C, 2, 2, 5)] - fn length_mismatch_reports_the_matrix_and_preserves_output( - #[case] input_len: usize, - #[case] output_len: usize, - #[case] matrix_name: MatrixName, - #[case] expected_rows: usize, - #[case] expected_cols: usize, - #[case] actual_len: usize, - #[values(sgemm_aat_lower, sgemm_aat_lower_add)] operation: LowerProduct, - ) { - let input = vec![1.0; input_len]; - let mut output = vec![19.0; output_len]; - - let error = operation(2, 2, 1.0, &input, &mut output).unwrap_err(); - - assert_eq!( - error, - SgemmError::InvalidMatrixDimensions { - matrix_name, - expected_rows, - expected_cols, - actual_len + #[test] + fn length_mismatch_reports_the_matrix_and_preserves_output() { + for (name, input_len, output_len, matrix_name, expected_rows, expected_cols, actual_len) in [ + ("short_input", 3, 4, MatrixName::A, 2, 2, 3), + ("long_input", 5, 4, MatrixName::A, 2, 2, 5), + ("short_output", 4, 3, MatrixName::C, 2, 2, 3), + ("long_output", 4, 5, MatrixName::C, 2, 2, 5), + ] { + for (mode, operation) in [ + ("replace", sgemm_aat_lower as LowerProduct), + ("add", sgemm_aat_lower_add as LowerProduct), + ] { + let input = vec![1.0; input_len]; + let mut output = vec![19.0; output_len]; + + let Err(error) = operation(2, 2, 1.0, &input, &mut output) else { + panic!("{name}, {mode}: expected a length error"); + }; + + assert_eq!( + error, + SgemmError::InvalidMatrixDimensions { + matrix_name, + expected_rows, + expected_cols, + actual_len + }, + "{name}, {mode}" + ); + assert_eq!(output, vec![19.0; output_len], "{name}, {mode}"); } - ); - assert_eq!(output, vec![19.0; output_len]); + } } - #[rstest] - #[case::input(usize::MAX, 2, MatrixName::A, usize::MAX, 2)] - #[case::output(usize::MAX, 0, MatrixName::C, usize::MAX, usize::MAX)] - fn shape_overflow_reports_the_matrix_and_preserves_output( - #[case] rows: usize, - #[case] dimensions: usize, - #[case] matrix_name: MatrixName, - #[case] error_rows: usize, - #[case] error_cols: usize, - #[values(sgemm_aat_lower, sgemm_aat_lower_add)] operation: LowerProduct, - ) { - let mut output = [23.0, 29.0]; - - let error = operation(rows, dimensions, 1.0, &[], &mut output).unwrap_err(); - - assert_eq!( - error, - SgemmError::DimensionOverflow { - matrix_name, - rows: error_rows, - cols: error_cols + #[test] + fn shape_overflow_reports_the_matrix_and_preserves_output() { + for (name, rows, dimensions, matrix_name, error_rows, error_cols) in [ + ("input", usize::MAX, 2, MatrixName::A, usize::MAX, 2), + ( + "output", + usize::MAX, + 0, + MatrixName::C, + usize::MAX, + usize::MAX, + ), + ] { + for (mode, operation) in [ + ("replace", sgemm_aat_lower as LowerProduct), + ("add", sgemm_aat_lower_add as LowerProduct), + ] { + let mut output = [23.0, 29.0]; + + let Err(error) = operation(rows, dimensions, 1.0, &[], &mut output) else { + panic!("{name}, {mode}: expected an overflow error"); + }; + + assert_eq!( + error, + SgemmError::DimensionOverflow { + matrix_name, + rows: error_rows, + cols: error_cols + }, + "{name}, {mode}" + ); + assert_eq!(output, [23.0, 29.0], "{name}, {mode}"); } - ); - assert_eq!(output, [23.0, 29.0]); + } } } diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index c9b37230d1..ecd58dbc5d 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -187,98 +187,106 @@ mod tests { #[case] _metric: M, #[case] scalar_metric: Metric, #[case] expected_ids: [[u32; 4]; 5], - #[values(0, 1, 2, 3, 4)] neighbors: usize, - #[values(2, 7, 8, 9, 15, 16, 17, 128, 129)] dimensions: usize, ) { - // These five points have distinct scores in each row for every metric. - let coordinates = [ - [2.0, 2.0], - [-1.0, 3.0], - [0.0, -2.0], - [5.0, 0.0], - [-3.0, -4.0], - ]; - let values = test_support::packed_points( - &coordinates, - dimensions, - scalar_metric == Metric::CosineNormalized, - ); - let points = MatrixView::try_from(values.as_slice(), 5, dimensions).unwrap(); - let mut output = vec![Candidate::new(0, -100.0); 5 * neighbors]; - let mut workspace = LeafKernelWorkspace::default(); - - select_leaf_neighbors::<_, M>( - ARCH, - points, - MutMatrixView::try_from(output.as_mut_slice(), 5, neighbors).unwrap(), - &mut workspace, - ) - .unwrap(); - - for point in 0..5 { - let actual = &output[point * neighbors..(point + 1) * neighbors]; - assert_eq!( - actual.iter().map(|c| c.local_idx).collect::>(), - expected_ids[point][..neighbors], - "point={point}" - ); - for candidate in actual { - let expected = test_support::distance( - scalar_metric, - points.row(point), - points.row(candidate.local_idx as usize), - ); - // Only two coordinates are nonzero, so eight f32 ulps at the score scale - // cover the dot/norm rounding without admitting another neighbor. - let tolerance = 8.0 * f64::from(f32::EPSILON) * expected.abs().max(1.0); - assert!( - (f64::from(candidate.distance) - expected).abs() <= tolerance, - "point={point}, candidate={candidate:?}, expected={expected}" + for neighbors in [0, 1, 2, 3, 4] { + for dimensions in [2, 7, 8, 9, 15, 16, 17, 128, 129] { + // These five points have distinct scores in each row for every metric. + let coordinates = [ + [2.0, 2.0], + [-1.0, 3.0], + [0.0, -2.0], + [5.0, 0.0], + [-3.0, -4.0], + ]; + let values = test_support::packed_points( + &coordinates, + dimensions, + scalar_metric == Metric::CosineNormalized, ); + let points = MatrixView::try_from(values.as_slice(), 5, dimensions).unwrap(); + let mut output = vec![Candidate::new(0, -100.0); 5 * neighbors]; + let mut workspace = LeafKernelWorkspace::default(); + + select_leaf_neighbors::<_, M>( + ARCH, + points, + MutMatrixView::try_from(output.as_mut_slice(), 5, neighbors).unwrap(), + &mut workspace, + ) + .unwrap_or_else(|error| { + panic!("neighbors={neighbors}, dimensions={dimensions}: {error}") + }); + + for point in 0..5 { + let actual = &output[point * neighbors..(point + 1) * neighbors]; + assert_eq!( + actual.iter().map(|c| c.local_idx).collect::>(), + expected_ids[point][..neighbors], + "neighbors={neighbors}, dimensions={dimensions}, point={point}" + ); + for candidate in actual { + let expected = test_support::distance( + scalar_metric, + points.row(point), + points.row(candidate.local_idx as usize), + ); + // Only two coordinates are nonzero, so eight f32 ulps at the score scale + // cover the dot/norm rounding without admitting another neighbor. + let tolerance = 8.0 * f64::from(f32::EPSILON) * expected.abs().max(1.0); + assert!( + (f64::from(candidate.distance) - expected).abs() <= tolerance, + "neighbors={neighbors}, dimensions={dimensions}, point={point}, candidate={candidate:?}, expected={expected}" + ); + } + } } } } #[cfg(not(miri))] - #[rstest] - fn neighbors_match_scalar_ranking_across_leaf_sizes_and_counts( - #[values(1, 2, 16, 17, 18, 33, 34)] point_count: usize, - #[values(1, 2, 3, 10, 11, 17)] requested: usize, - ) { - // Slightly increasing gaps avoid ties between the two sides of each point. - // Center the coordinates so intermediate norm sums stay exactly representable. - let mut values: Vec<_> = (0..point_count).map(|i| (64 * i + i * i) as f32).collect(); - let center = values[point_count - 1] / 2.0; - for value in &mut values { - *value -= center; - } - let neighbors = requested.min(point_count - 1); - let mut output = vec![EMPTY; point_count * neighbors]; - - select_leaf_neighbors::<_, L2>( - ARCH, - MatrixView::try_from(values.as_slice(), point_count, 1).unwrap(), - MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors).unwrap(), - &mut LeafKernelWorkspace::default(), - ) - .unwrap(); - - for point in 0..point_count { - let mut expected: Vec<_> = values - .iter() - .enumerate() - .filter(|&(other, _)| other != point) - .map(|(other, &value)| { - Candidate::new(other as u32, (values[point] - value).powi(2)) - }) - .collect(); - expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); - expected.truncate(neighbors); - assert_eq!( - &output[point * neighbors..(point + 1) * neighbors], - expected, - "point={point}, point_count={point_count}, neighbors={neighbors}" - ); + #[test] + fn neighbors_match_scalar_ranking_across_leaf_sizes_and_counts() { + for point_count in [1, 2, 16, 17, 18, 33, 34] { + for requested in [1, 2, 3, 10, 11, 17] { + // Slightly increasing gaps avoid ties between the two sides of each point. + // Center the coordinates so intermediate norm sums stay exactly representable. + let mut values: Vec<_> = + (0..point_count).map(|i| (64 * i + i * i) as f32).collect(); + let center = values[point_count - 1] / 2.0; + for value in &mut values { + *value -= center; + } + let neighbors = requested.min(point_count - 1); + let mut output = vec![EMPTY; point_count * neighbors]; + + select_leaf_neighbors::<_, L2>( + ARCH, + MatrixView::try_from(values.as_slice(), point_count, 1).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap_or_else(|error| { + panic!("point_count={point_count}, requested={requested}: {error}") + }); + + for point in 0..point_count { + let mut expected: Vec<_> = values + .iter() + .enumerate() + .filter(|&(other, _)| other != point) + .map(|(other, &value)| { + Candidate::new(other as u32, (values[point] - value).powi(2)) + }) + .collect(); + expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); + expected.truncate(neighbors); + assert_eq!( + &output[point * neighbors..(point + 1) * neighbors], + expected, + "point_count={point_count}, requested={requested}, point={point}, neighbors={neighbors}" + ); + } + } } } @@ -291,71 +299,78 @@ mod tests { fn large_dense_leaves_select_nearest_non_self_neighbors( #[case] _metric: M, #[case] scalar_metric: Metric, - #[values((33, 384), (65, 768), (129, 1536), (35, 1537))] shape: (usize, usize), - #[values(3, 11)] neighbors: usize, ) { - let (point_count, dimensions) = shape; - let mut values = test_support::dense_points(point_count, dimensions, 1287); - if scalar_metric == Metric::CosineNormalized { - test_support::normalize(&mut values, dimensions); - } - let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); - let mut output = vec![EMPTY; point_count * neighbors]; - - select_leaf_neighbors::<_, M>( - ARCH, - points, - MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors).unwrap(), - &mut LeafKernelWorkspace::default(), - ) - .unwrap(); - - // Unnormalized dyadic sums are exact; cosine rounds sqrt/division, while - // normalized dot products also accumulate coordinate rounding. - let tolerance = match scalar_metric { - Metric::L2 | Metric::InnerProduct => 0.0, - Metric::Cosine => 16.0 * f64::from(f32::EPSILON), - Metric::CosineNormalized => { - let roundoff = dimensions as f64 * f64::from(f32::EPSILON); - roundoff / (1.0 - roundoff) - } - }; - for point in 0..point_count { - let mut expected: Vec<_> = (0..point_count) - .filter(|&other| other != point) - .map(|other| { - ( - other as u32, - test_support::distance(scalar_metric, points.row(point), points.row(other)), - ) - }) - .collect(); - expected.sort_by(|left, right| left.1.total_cmp(&right.1)); - let actual = &output[point * neighbors..(point + 1) * neighbors]; - for (rank, candidate) in actual.iter().enumerate() { - assert!( - !actual[..rank] - .iter() - .any(|previous| previous.local_idx == candidate.local_idx), - "duplicate neighbor for point={point}: {candidate:?}" - ); - let own_score = expected - .iter() - .find(|&&(id, _)| id == candidate.local_idx) - .unwrap_or_else(|| { - panic!("invalid or self neighbor for point={point}: {candidate:?}") - }) - .1; - assert!( - (f64::from(candidate.distance) - own_score).abs() <= tolerance, - "shape={shape:?}, point={point}, candidate={candidate:?}, score={own_score}" - ); - // Equal scores may use either ID, but every rank must be nearest-first. - assert!( - (own_score - expected[rank].1).abs() <= tolerance, - "shape={shape:?}, point={point}, rank={rank}, score={own_score}, expected={:?}", - expected[rank] - ); + for shape in [(33, 384), (65, 768), (129, 1536), (35, 1537)] { + for neighbors in [3, 11] { + let (point_count, dimensions) = shape; + let mut values = test_support::dense_points(point_count, dimensions, 1287); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut values, dimensions); + } + let points = + MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let mut output = vec![EMPTY; point_count * neighbors]; + + select_leaf_neighbors::<_, M>( + ARCH, + points, + MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap_or_else(|error| panic!("shape={shape:?}, neighbors={neighbors}: {error}")); + + // Unnormalized dyadic sums are exact; cosine rounds sqrt/division, while + // normalized dot products also accumulate coordinate rounding. + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + }; + for point in 0..point_count { + let mut expected: Vec<_> = (0..point_count) + .filter(|&other| other != point) + .map(|other| { + ( + other as u32, + test_support::distance( + scalar_metric, + points.row(point), + points.row(other), + ), + ) + }) + .collect(); + expected.sort_by(|left, right| left.1.total_cmp(&right.1)); + let actual = &output[point * neighbors..(point + 1) * neighbors]; + for (rank, candidate) in actual.iter().enumerate() { + assert!( + !actual[..rank] + .iter() + .any(|previous| previous.local_idx == candidate.local_idx), + "shape={shape:?}, neighbors={neighbors}, duplicate neighbor for point={point}: {candidate:?}" + ); + let own_score = expected + .iter() + .find(|&&(id, _)| id == candidate.local_idx) + .unwrap_or_else(|| { + panic!("shape={shape:?}, neighbors={neighbors}, invalid or self neighbor for point={point}: {candidate:?}") + }) + .1; + assert!( + (f64::from(candidate.distance) - own_score).abs() <= tolerance, + "shape={shape:?}, neighbors={neighbors}, point={point}, candidate={candidate:?}, score={own_score}" + ); + // Equal scores may use either ID, but every rank must be nearest-first. + assert!( + (own_score - expected[rank].1).abs() <= tolerance, + "shape={shape:?}, neighbors={neighbors}, point={point}, rank={rank}, score={own_score}, expected={:?}", + expected[rank] + ); + } + } } } } diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index bcbbc374d2..f29ad151b2 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -154,42 +154,50 @@ mod tests { fn lower_triangle_matches_scalar_distances( #[case] _metric: M, #[case] scalar_metric: Metric, - #[values(1, 4, 17)] point_count: usize, - #[values(1, 2, 7, 8, 9, 15, 16, 17, 127, 128, 129)] dimensions: usize, ) { - // Small integers keep L2 and dot products exact; the first coordinate gives - // every vector a nonzero norm. Every dimension contributes to the result. - let mut values: Vec<_> = (0..point_count * dimensions) - .map(|index| (index % 11) as f32 - 5.0) - .collect(); - for (point, row) in values.chunks_exact_mut(dimensions).enumerate() { - row[0] = point as f32 + 1.0; - } - if scalar_metric == Metric::CosineNormalized { - test_support::normalize(&mut values, dimensions); - } - let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); - let mut output = vec![f32::NAN; point_count * point_count]; + for point_count in [1, 4, 17] { + for dimensions in [1, 2, 7, 8, 9, 15, 16, 17, 127, 128, 129] { + // Small integers keep L2 and dot products exact; the first coordinate gives + // every vector a nonzero norm. Every dimension contributes to the result. + let mut values: Vec<_> = (0..point_count * dimensions) + .map(|index| (index % 11) as f32 - 5.0) + .collect(); + for (point, row) in values.chunks_exact_mut(dimensions).enumerate() { + row[0] = point as f32 + 1.0; + } + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut values, dimensions); + } + let points = + MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let mut output = vec![f32::NAN; point_count * point_count]; - M::compute_distances(points, &mut output).unwrap(); + M::compute_distances(points, &mut output).unwrap_or_else(|error| { + panic!("point_count={point_count}, dimensions={dimensions}: {error}") + }); - for source in 0..point_count { - for target in 0..=source { - let expected = - test_support::distance(scalar_metric, points.row(source), points.row(target)); - let actual = f64::from(output[source * point_count + target]); - // Cosine reductions and normalization round in f32; allow eight ulps - // per dimension at unit scale. Integer L2 and dot products are exact. - let tolerance = match scalar_metric { - Metric::L2 | Metric::InnerProduct => 0.0, - Metric::Cosine | Metric::CosineNormalized => { - 8.0 * f64::from(f32::EPSILON) * dimensions as f64 + for source in 0..point_count { + for target in 0..=source { + let expected = test_support::distance( + scalar_metric, + points.row(source), + points.row(target), + ); + let actual = f64::from(output[source * point_count + target]); + // Cosine reductions and normalization round in f32; allow eight ulps + // per dimension at unit scale. Integer L2 and dot products are exact. + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + Metric::Cosine | Metric::CosineNormalized => { + 8.0 * f64::from(f32::EPSILON) * dimensions as f64 + } + }; + assert!( + (actual - expected).abs() <= tolerance, + "point_count={point_count}, dimensions={dimensions}, pair=({source},{target}): {actual} != {expected}" + ); } - }; - assert!( - (actual - expected).abs() <= tolerance, - "pair=({source},{target}), dimensions={dimensions}: {actual} != {expected}" - ); + } } } } @@ -203,42 +211,51 @@ mod tests { fn large_dense_inputs_match_scalar_distances( #[case] _metric: M, #[case] scalar_metric: Metric, - #[values((33, 384), (65, 768), (129, 1536), (17, 1537), (513, 129), (17, 4097))] shape: ( - usize, - usize, - ), ) { - let (point_count, dimensions) = shape; - let mut values = test_support::dense_points(point_count, dimensions, 1287); - if scalar_metric == Metric::CosineNormalized { - test_support::normalize(&mut values, dimensions); - } - let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); - let mut output = vec![f32::NAN; point_count * point_count]; + for shape in [ + (33, 384), + (65, 768), + (129, 1536), + (17, 1537), + (513, 129), + (17, 4097), + ] { + let (point_count, dimensions) = shape; + let mut values = test_support::dense_points(point_count, dimensions, 1287); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut values, dimensions); + } + let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let mut output = vec![f32::NAN; point_count * point_count]; - M::compute_distances(points, &mut output).unwrap(); + M::compute_distances(points, &mut output) + .unwrap_or_else(|error| panic!("shape={shape:?}: {error}")); - let tolerance = match scalar_metric { - Metric::L2 | Metric::InnerProduct => 0.0, - // Dyadic dot/norm sums are exact; only square roots and division round. - Metric::Cosine => 16.0 * f64::from(f32::EPSILON), - Metric::CosineNormalized => { - // Normalized f32 coordinates need a dot-product rounding bound. - // Sum |x*y| is at most approximately one; EPSILON allows both - // product and reduction rounding in gamma = n*u/(1-n*u). - let roundoff = dimensions as f64 * f64::from(f32::EPSILON); - roundoff / (1.0 - roundoff) - } - }; - for source in 0..point_count { - for target in 0..=source { - let expected = - test_support::distance(scalar_metric, points.row(source), points.row(target)); - let actual = f64::from(output[source * point_count + target]); - assert!( - (actual - expected).abs() <= tolerance, - "shape={shape:?}, pair=({source},{target}): {actual} != {expected}, tolerance={tolerance}" - ); + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + // Dyadic dot/norm sums are exact; only square roots and division round. + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + // Normalized f32 coordinates need a dot-product rounding bound. + // Sum |x*y| is at most approximately one; EPSILON allows both + // product and reduction rounding in gamma = n*u/(1-n*u). + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + }; + for source in 0..point_count { + for target in 0..=source { + let expected = test_support::distance( + scalar_metric, + points.row(source), + points.row(target), + ); + let actual = f64::from(output[source * point_count + target]); + assert!( + (actual - expected).abs() <= tolerance, + "shape={shape:?}, pair=({source},{target}): {actual} != {expected}, tolerance={tolerance}" + ); + } } } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index e16a56d976..d3938ef28f 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -138,86 +138,93 @@ mod tests { #[case] _metric: M, #[case] unit_norm: bool, #[case] expected_ids: [[u32; 4]; 3], - #[values(1, 2, 3, 4, 6)] assignments: usize, - #[values(2, 7, 8, 9, 15, 16, 17, 128, 129)] dimensions: usize, ) { - let point_values = test_support::packed_points( - &[[3.0, 2.0], [-1.0, 3.0], [2.0, -4.0]], - dimensions, - unit_norm, - ); - let leader_values = test_support::packed_points( - &[[4.0, 0.0], [0.0, 3.0], [-2.0, 0.0], [0.0, -1.0]], - dimensions, - unit_norm, - ); - let leaders = M::create_leaders( - MatrixView::try_from(leader_values.as_slice(), 4, dimensions).unwrap(), - ); - let mut output = vec![0; 3 * assignments]; + for assignments in [1, 2, 3, 4, 6] { + for dimensions in [2, 7, 8, 9, 15, 16, 17, 128, 129] { + let point_values = test_support::packed_points( + &[[3.0, 2.0], [-1.0, 3.0], [2.0, -4.0]], + dimensions, + unit_norm, + ); + let leader_values = test_support::packed_points( + &[[4.0, 0.0], [0.0, 3.0], [-2.0, 0.0], [0.0, -1.0]], + dimensions, + unit_norm, + ); + let leaders = M::create_leaders( + MatrixView::try_from(leader_values.as_slice(), 4, dimensions).unwrap(), + ); + let mut output = vec![0; 3 * assignments]; - assign_leaders::<_, M>( - ARCH, - MatrixView::try_from(point_values.as_slice(), 3, dimensions).unwrap(), - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), - &mut PartitionKernelWorkspace::default(), - ) - .unwrap(); + assign_leaders::<_, M>( + ARCH, + MatrixView::try_from(point_values.as_slice(), 3, dimensions).unwrap(), + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap_or_else(|error| { + panic!("assignments={assignments}, dimensions={dimensions}: {error}") + }); - for point in 0..3 { - let mut expected = expected_ids[point].to_vec(); - expected.truncate(assignments); - expected.resize(assignments, UNASSIGNED_LEADER); - assert_eq!( - &output[point * assignments..(point + 1) * assignments], - expected, - "point={point}" - ); + for point in 0..3 { + let mut expected = expected_ids[point].to_vec(); + expected.truncate(assignments); + expected.resize(assignments, UNASSIGNED_LEADER); + assert_eq!( + &output[point * assignments..(point + 1) * assignments], + expected, + "assignments={assignments}, dimensions={dimensions}, point={point}" + ); + } + } } } #[cfg(not(miri))] - #[rstest] - fn assignments_match_nearest_leaders_across_counts_and_widths( - #[values(1, 5, 16, 17, 33)] leader_count: usize, - #[values(1, 3, 10, 11, 17)] assignments: usize, - ) { - // Queries avoid midpoints, so each leader has a distinct distance. - let leader_values: Vec<_> = (0..leader_count).map(|i| 2.0 * i as f32).collect(); - let point_values = [-1.0, 1.5, 2.0 * (leader_count - 1) as f32 + 0.25]; - let leaders = L2::create_leaders( - MatrixView::try_from(leader_values.as_slice(), leader_count, 1).unwrap(), - ); - let mut output = vec![0; 3 * assignments]; + #[test] + fn assignments_match_nearest_leaders_across_counts_and_widths() { + for leader_count in [1, 5, 16, 17, 33] { + for assignments in [1, 3, 10, 11, 17] { + // Queries avoid midpoints, so each leader has a distinct distance. + let leader_values: Vec<_> = (0..leader_count).map(|i| 2.0 * i as f32).collect(); + let point_values = [-1.0, 1.5, 2.0 * (leader_count - 1) as f32 + 0.25]; + let leaders = L2::create_leaders( + MatrixView::try_from(leader_values.as_slice(), leader_count, 1).unwrap(), + ); + let mut output = vec![0; 3 * assignments]; - assign_leaders::<_, L2>( - ARCH, - MatrixView::try_from(&point_values[..], 3, 1).unwrap(), - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), - &mut PartitionKernelWorkspace::default(), - ) - .unwrap(); + assign_leaders::<_, L2>( + ARCH, + MatrixView::try_from(&point_values[..], 3, 1).unwrap(), + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap_or_else(|error| { + panic!("leader_count={leader_count}, assignments={assignments}: {error}") + }); - for (point, &coordinate) in point_values.iter().enumerate() { - let mut ordered: Vec<_> = leader_values - .iter() - .enumerate() - .map(|(id, &value)| (id as u32, (coordinate - value).powi(2))) - .collect(); - ordered.sort_by(|left, right| left.1.total_cmp(&right.1)); - let mut expected: Vec<_> = ordered - .iter() - .take(assignments) - .map(|&(id, _)| id) - .collect(); - expected.resize(assignments, UNASSIGNED_LEADER); - assert_eq!( - &output[point * assignments..(point + 1) * assignments], - expected, - "point={point}, leader_count={leader_count}, assignments={assignments}" - ); + for (point, &coordinate) in point_values.iter().enumerate() { + let mut ordered: Vec<_> = leader_values + .iter() + .enumerate() + .map(|(id, &value)| (id as u32, (coordinate - value).powi(2))) + .collect(); + ordered.sort_by(|left, right| left.1.total_cmp(&right.1)); + let mut expected: Vec<_> = ordered + .iter() + .take(assignments) + .map(|&(id, _)| id) + .collect(); + expected.resize(assignments, UNASSIGNED_LEADER); + assert_eq!( + &output[point * assignments..(point + 1) * assignments], + expected, + "point={point}, leader_count={leader_count}, assignments={assignments}" + ); + } + } } } @@ -230,75 +237,77 @@ mod tests { fn large_dense_stripes_select_nearest_leaders( #[case] _metric: M, #[case] scalar_metric: Metric, - #[values((9, 33, 384), (17, 65, 768), (33, 129, 1536), (5, 35, 1537))] shape: ( - usize, - usize, - usize, - ), - #[values(3, 11)] assignments: usize, ) { - let (point_count, leader_count, dimensions) = shape; - let mut point_values = test_support::dense_points(point_count, dimensions, 1287); - let mut leader_values = test_support::dense_points(leader_count, dimensions, 2026); - if scalar_metric == Metric::CosineNormalized { - test_support::normalize(&mut point_values, dimensions); - test_support::normalize(&mut leader_values, dimensions); - } - let points = - MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); - let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); - let leaders = M::create_leaders(leader_matrix); - let mut output = vec![UNASSIGNED_LEADER; point_count * assignments]; + for shape in [(9, 33, 384), (17, 65, 768), (33, 129, 1536), (5, 35, 1537)] { + for assignments in [3, 11] { + let (point_count, leader_count, dimensions) = shape; + let mut point_values = test_support::dense_points(point_count, dimensions, 1287); + let mut leader_values = test_support::dense_points(leader_count, dimensions, 2026); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut point_values, dimensions); + test_support::normalize(&mut leader_values, dimensions); + } + let points = + MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions) + .unwrap(); + let leaders = M::create_leaders(leader_matrix); + let mut output = vec![UNASSIGNED_LEADER; point_count * assignments]; - assign_leaders::<_, M>( - ARCH, - points, - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), point_count, assignments).unwrap(), - &mut PartitionKernelWorkspace::default(), - ) - .unwrap(); + assign_leaders::<_, M>( + ARCH, + points, + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), point_count, assignments) + .unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap_or_else(|error| { + panic!("shape={shape:?}, assignments={assignments}: {error}") + }); - let tolerance = match scalar_metric { - Metric::L2 | Metric::InnerProduct => 0.0, - // Dyadic inputs make the dot/norm sums exact before cosine conversion. - Metric::Cosine => 16.0 * f64::from(f32::EPSILON), - Metric::CosineNormalized => { - let roundoff = dimensions as f64 * f64::from(f32::EPSILON); - roundoff / (1.0 - roundoff) - } - }; - for point in 0..point_count { - let mut expected: Vec<_> = (0..leader_count) - .map(|leader| { - ( - leader as u32, - test_support::distance( - scalar_metric, - points.row(point), - leader_matrix.row(leader), - ), - ) - }) - .collect(); - expected.sort_by(|left, right| left.1.total_cmp(&right.1)); - let actual = &output[point * assignments..(point + 1) * assignments]; - for (rank, &id) in actual.iter().enumerate() { - assert!( - !actual[..rank].contains(&id), - "duplicate leader {id} for point={point}" - ); - let score = expected - .iter() - .find(|&&(leader, _)| leader == id) - .unwrap_or_else(|| panic!("invalid leader {id} for point={point}")) - .1; - assert!( - (score - expected[rank].1).abs() <= tolerance, - "shape={shape:?}, point={point}, rank={rank}, score={score}, expected={:?}", - expected[rank] - ); + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + // Dyadic inputs make the dot/norm sums exact before cosine conversion. + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + }; + for point in 0..point_count { + let mut expected: Vec<_> = (0..leader_count) + .map(|leader| { + ( + leader as u32, + test_support::distance( + scalar_metric, + points.row(point), + leader_matrix.row(leader), + ), + ) + }) + .collect(); + expected.sort_by(|left, right| left.1.total_cmp(&right.1)); + let actual = &output[point * assignments..(point + 1) * assignments]; + for (rank, &id) in actual.iter().enumerate() { + assert!( + !actual[..rank].contains(&id), + "shape={shape:?}, assignments={assignments}, duplicate leader {id} for point={point}" + ); + let score = expected + .iter() + .find(|&&(leader, _)| leader == id) + .unwrap_or_else(|| panic!("shape={shape:?}, assignments={assignments}, invalid leader {id} for point={point}")) + .1; + assert!( + (score - expected[rank].1).abs() <= tolerance, + "shape={shape:?}, assignments={assignments}, point={point}, rank={rank}, score={score}, expected={:?}", + expected[rank] + ); + } + } } } } diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index dbc623a856..f5c20d0581 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -273,70 +273,82 @@ mod tests { fn point_to_leader_scores_match_scalar_distances( #[case] _metric: M, #[case] scalar_metric: Metric, - #[values(1, 3)] point_count: usize, - #[values(1, 4, 17)] leader_count: usize, - #[values(1, 2, 7, 8, 9, 15, 16, 17, 127, 128, 129)] dimensions: usize, ) { - // Small integer coordinates make unnormalized dot products exact. Points - // and leaders use different values, so swapped rows or columns change scores. - let mut point_values: Vec<_> = (0..point_count * dimensions) - .map(|index| (index % 7) as f32 - 3.0) - .collect(); - let mut leader_values: Vec<_> = (0..leader_count * dimensions) - .map(|index| (index % 11) as f32 - 5.0) - .collect(); - for (point, row) in point_values.chunks_exact_mut(dimensions).enumerate() { - row[0] = point as f32 + 1.0; - } - for (leader, row) in leader_values.chunks_exact_mut(dimensions).enumerate() { - row[0] = leader as f32 + 2.0; - } - if scalar_metric == Metric::CosineNormalized { - test_support::normalize(&mut point_values, dimensions); - test_support::normalize(&mut leader_values, dimensions); - } - let points = - MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); - let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); - let leaders = M::create_leaders(leader_matrix); - let mut output = vec![f32::NAN; point_count * leader_count]; - - M::compute_distances( - points, - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), point_count, leader_count).unwrap(), - ) - .unwrap(); - - assert_eq!(M::leader_count(&leaders), leader_count); - for point in 0..point_count { - for leader in 0..leader_count { - let mut expected = test_support::distance( - scalar_metric, - points.row(point), - leader_matrix.row(leader), - ); - if scalar_metric == Metric::L2 { - // Partition L2 omits exactly this constant from every column of the row. - expected -= points - .row(point) - .iter() - .map(|&x| f64::from(x).powi(2)) - .sum::(); - } - let tolerance = match scalar_metric { - Metric::L2 | Metric::InnerProduct => 0.0, - // Bound f32 norm/dot reductions at unit scale, as in the leaf metric. - Metric::Cosine | Metric::CosineNormalized => { - 8.0 * f64::from(f32::EPSILON) * dimensions as f64 + for point_count in [1, 3] { + for leader_count in [1, 4, 17] { + for dimensions in [1, 2, 7, 8, 9, 15, 16, 17, 127, 128, 129] { + // Small integer coordinates make unnormalized dot products exact. Points + // and leaders use different values, so swapped rows or columns change scores. + let mut point_values: Vec<_> = (0..point_count * dimensions) + .map(|index| (index % 7) as f32 - 3.0) + .collect(); + let mut leader_values: Vec<_> = (0..leader_count * dimensions) + .map(|index| (index % 11) as f32 - 5.0) + .collect(); + for (point, row) in point_values.chunks_exact_mut(dimensions).enumerate() { + row[0] = point as f32 + 1.0; + } + for (leader, row) in leader_values.chunks_exact_mut(dimensions).enumerate() { + row[0] = leader as f32 + 2.0; } - }; - let actual = f64::from(output[point * leader_count + leader]); - assert!( - (actual - expected).abs() <= tolerance, - "point={point}, leader={leader}, dimensions={dimensions}: {actual} != {expected}" - ); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut point_values, dimensions); + test_support::normalize(&mut leader_values, dimensions); + } + let points = + MatrixView::try_from(point_values.as_slice(), point_count, dimensions) + .unwrap(); + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions) + .unwrap(); + let leaders = M::create_leaders(leader_matrix); + let mut output = vec![f32::NAN; point_count * leader_count]; + + M::compute_distances( + points, + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), point_count, leader_count) + .unwrap(), + ) + .unwrap_or_else(|error| { + panic!("point_count={point_count}, leader_count={leader_count}, dimensions={dimensions}: {error}") + }); + + assert_eq!( + M::leader_count(&leaders), + leader_count, + "point_count={point_count}, leader_count={leader_count}, dimensions={dimensions}" + ); + for point in 0..point_count { + for leader in 0..leader_count { + let mut expected = test_support::distance( + scalar_metric, + points.row(point), + leader_matrix.row(leader), + ); + if scalar_metric == Metric::L2 { + // Partition L2 omits exactly this constant from every column of the row. + expected -= points + .row(point) + .iter() + .map(|&x| f64::from(x).powi(2)) + .sum::(); + } + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + // Bound f32 norm/dot reductions at unit scale, as in the leaf metric. + Metric::Cosine | Metric::CosineNormalized => { + 8.0 * f64::from(f32::EPSILON) * dimensions as f64 + } + }; + let actual = f64::from(output[point * leader_count + leader]); + assert!( + (actual - expected).abs() <= tolerance, + "point_count={point_count}, leader_count={leader_count}, dimensions={dimensions}, point={point}, leader={leader}: {actual} != {expected}" + ); + } + } + } } } } @@ -350,60 +362,67 @@ mod tests { fn large_dense_inputs_match_scalar_distances( #[case] _metric: M, #[case] scalar_metric: Metric, - #[values((17, 33, 384), (33, 65, 768), (65, 129, 1536), (9, 17, 1537), (257, 513, 129), (5, 17, 4097))] - shape: (usize, usize, usize), ) { - let (point_count, leader_count, dimensions) = shape; - let mut point_values = test_support::dense_points(point_count, dimensions, 1287); - let mut leader_values = test_support::dense_points(leader_count, dimensions, 2026); - if scalar_metric == Metric::CosineNormalized { - test_support::normalize(&mut point_values, dimensions); - test_support::normalize(&mut leader_values, dimensions); - } - let points = - MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); - let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); - let leaders = M::create_leaders(leader_matrix); - let mut output = vec![f32::NAN; point_count * leader_count]; - - M::compute_distances( - points, - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), point_count, leader_count).unwrap(), - ) - .unwrap(); - - let tolerance = match scalar_metric { - Metric::L2 | Metric::InnerProduct => 0.0, - // Dyadic dot/norm sums are exact; square roots and division round. - Metric::Cosine => 16.0 * f64::from(f32::EPSILON), - Metric::CosineNormalized => { - // Bound product and sum rounding for normalized f32 coordinates. - let roundoff = dimensions as f64 * f64::from(f32::EPSILON); - roundoff / (1.0 - roundoff) + for shape in [ + (17, 33, 384), + (33, 65, 768), + (65, 129, 1536), + (9, 17, 1537), + (257, 513, 129), + (5, 17, 4097), + ] { + let (point_count, leader_count, dimensions) = shape; + let mut point_values = test_support::dense_points(point_count, dimensions, 1287); + let mut leader_values = test_support::dense_points(leader_count, dimensions, 2026); + if scalar_metric == Metric::CosineNormalized { + test_support::normalize(&mut point_values, dimensions); + test_support::normalize(&mut leader_values, dimensions); } - }; - for point in 0..point_count { - let point_norm: f64 = points - .row(point) - .iter() - .map(|&x| f64::from(x).powi(2)) - .sum(); - for leader in 0..leader_count { - let mut expected = test_support::distance( - scalar_metric, - points.row(point), - leader_matrix.row(leader), - ); - if scalar_metric == Metric::L2 { - expected -= point_norm; + let points = + MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions).unwrap(); + let leaders = M::create_leaders(leader_matrix); + let mut output = vec![f32::NAN; point_count * leader_count]; + + M::compute_distances( + points, + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), point_count, leader_count).unwrap(), + ) + .unwrap_or_else(|error| panic!("shape={shape:?}: {error}")); + + let tolerance = match scalar_metric { + Metric::L2 | Metric::InnerProduct => 0.0, + // Dyadic dot/norm sums are exact; square roots and division round. + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + // Bound product and sum rounding for normalized f32 coordinates. + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + }; + for point in 0..point_count { + let point_norm: f64 = points + .row(point) + .iter() + .map(|&x| f64::from(x).powi(2)) + .sum(); + for leader in 0..leader_count { + let mut expected = test_support::distance( + scalar_metric, + points.row(point), + leader_matrix.row(leader), + ); + if scalar_metric == Metric::L2 { + expected -= point_norm; + } + let actual = f64::from(output[point * leader_count + leader]); + assert!( + (actual - expected).abs() <= tolerance, + "shape={shape:?}, point={point}, leader={leader}: {actual} != {expected}, tolerance={tolerance}" + ); } - let actual = f64::from(output[point * leader_count + leader]); - assert!( - (actual - expected).abs() <= tolerance, - "shape={shape:?}, point={point}, leader={leader}: {actual} != {expected}, tolerance={tolerance}" - ); } } } diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index abb8d9fc14..d4378087f4 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -481,12 +481,8 @@ mod tests { assert_eq!(output, [Candidate::new(1, 3.0), Candidate::new(2, 9.0)]); } - #[rstest] - fn batch_selection_matches_a_full_sort( - #[values(0, 1, LANES - 1, LANES, LANES + 1, 2 * LANES, 2 * LANES + 3, 3 * LANES + 2)] - count: usize, - #[values(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, LANES + 1)] capacity: usize, - ) { + #[test] + fn batch_selection_matches_a_full_sort() { struct SelectRow<'a> { distances: &'a [f32], output: &'a mut [Candidate], @@ -498,38 +494,51 @@ mod tests { } } - // Each pair offers a nearer score before a farther one. Later pairs improve - // on earlier pairs, so a full result must keep lowering its cutoff mid-vector. - let mut distances: Vec<_> = (0..count).map(|i| -(i as f32) - 1.0).collect(); - for pair in distances.chunks_exact_mut(2) { - pair.swap(0, 1); - } + for count in [ + 0, + 1, + LANES - 1, + LANES, + LANES + 1, + 2 * LANES, + 2 * LANES + 3, + 3 * LANES + 2, + ] { + for capacity in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, LANES + 1] { + // Each pair offers a nearer score before a farther one. Later pairs improve + // on earlier pairs, so a full result must keep lowering its cutoff mid-vector. + let mut distances: Vec<_> = (0..count).map(|i| -(i as f32) - 1.0).collect(); + for pair in distances.chunks_exact_mut(2) { + pair.swap(0, 1); + } - // Sorting the whole input is independent of TopK's insertion and cutoff logic. - let mut expected: Vec<_> = distances - .iter() - .enumerate() - .map(|(index, &distance)| Candidate::new(index as u32, distance)) - .collect(); - expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); - expected.truncate(capacity); - expected.resize(capacity, EMPTY); - let mut output = vec![EMPTY; capacity]; - let mut capacity_reads = 0; - - with_topk( - { - capacity_reads += 1; - capacity - }, - SelectRow { - distances: &distances, - output: &mut output, - }, - ); + // Sorting the whole input is independent of TopK's insertion and cutoff logic. + let mut expected: Vec<_> = distances + .iter() + .enumerate() + .map(|(index, &distance)| Candidate::new(index as u32, distance)) + .collect(); + expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); + expected.truncate(capacity); + expected.resize(capacity, EMPTY); + let mut output = vec![EMPTY; capacity]; + let mut capacity_reads = 0; + + with_topk( + { + capacity_reads += 1; + capacity + }, + SelectRow { + distances: &distances, + output: &mut output, + }, + ); - assert_eq!(output, expected, "count={count}, capacity={capacity}"); - assert_eq!(capacity_reads, 1); + assert_eq!(output, expected, "count={count}, capacity={capacity}"); + assert_eq!(capacity_reads, 1, "count={count}, capacity={capacity}"); + } + } } #[test] @@ -612,12 +621,8 @@ mod tests { assert_eq!(limits, [6.0, 7.0, 4.0, 4.0]); } - #[rstest] - fn pair_scans_match_full_row_sort_across_capacities_and_lengths( - #[values(LANES, LANES + 1, LANES + 2, 2 * LANES, 2 * LANES + 1, 2 * LANES + 2)] - point_count: usize, - #[values(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17)] capacity: usize, - ) { + #[test] + fn pair_scans_match_full_row_sort_across_capacities_and_lengths() { struct ScanPairs<'a> { distances: &'a [Vec], output: MutMatrixView<'a, Candidate>, @@ -639,52 +644,63 @@ mod tests { } } - // Given: XOR gives symmetric distances with distinct, exact scores in each row. - // The final pair row has point_count - 1 distances, straddling vector boundaries. - let distances: Vec> = (0..point_count) - .map(|point| { - (0..point_count) - .map(|other| (point ^ other) as f32) - .collect() - }) - .collect(); - let mut output = vec![EMPTY; point_count * capacity]; - let mut limits = Vec::new(); - let mut rows = - MutMatrixView::try_from(output.as_mut_slice(), point_count, capacity).unwrap(); - - // When: offer every non-self pair once, through the production width dispatch. - with_topk( - capacity, - ScanPairs { - distances: &distances, - output: rows.as_mut_view(), - limits: &mut limits, - }, - ); + for point_count in [ + LANES, + LANES + 1, + LANES + 2, + 2 * LANES, + 2 * LANES + 1, + 2 * LANES + 2, + ] { + for capacity in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17] { + // Given: XOR gives symmetric distances with distinct, exact scores in each row. + // The final pair row has point_count - 1 distances, straddling vector boundaries. + let distances: Vec> = (0..point_count) + .map(|point| { + (0..point_count) + .map(|other| (point ^ other) as f32) + .collect() + }) + .collect(); + let mut output = vec![EMPTY; point_count * capacity]; + let mut limits = Vec::new(); + let mut rows = + MutMatrixView::try_from(output.as_mut_slice(), point_count, capacity).unwrap(); + + // When: offer every non-self pair once, through the production width dispatch. + with_topk( + capacity, + ScanPairs { + distances: &distances, + output: rows.as_mut_view(), + limits: &mut limits, + }, + ); - // Then: independently sort each complete row, excluding only the point itself. - for (point, row) in distances.iter().enumerate() { - let mut expected: Vec<_> = row - .iter() - .enumerate() - .filter(|(other, _)| *other != point) - .map(|(other, &distance)| Candidate::new(other as u32, distance)) - .collect(); - expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); - expected.truncate(capacity); - expected.resize(capacity, EMPTY); - - assert_eq!( - rows.row(point), - expected, - "point={point}, point_count={point_count}, capacity={capacity}" - ); - assert_eq!( - limits[point], - expected[capacity - 1].distance, - "limit for point={point}, point_count={point_count}, capacity={capacity}" - ); + // Then: independently sort each complete row, excluding only the point itself. + for (point, row) in distances.iter().enumerate() { + let mut expected: Vec<_> = row + .iter() + .enumerate() + .filter(|(other, _)| *other != point) + .map(|(other, &distance)| Candidate::new(other as u32, distance)) + .collect(); + expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); + expected.truncate(capacity); + expected.resize(capacity, EMPTY); + + assert_eq!( + rows.row(point), + expected, + "point={point}, point_count={point_count}, capacity={capacity}" + ); + assert_eq!( + limits[point], + expected[capacity - 1].distance, + "limit for point={point}, point_count={point_count}, capacity={capacity}" + ); + } + } } } From a439b6ba8dfe0e5a1157bd52589a6e3476c56f95 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 05:52:35 +0000 Subject: [PATCH 197/229] refactor(pipnn): shorten SIMD trait name --- diskann/src/graph/pipnn/leaf_kernel.rs | 8 ++++---- diskann/src/graph/pipnn/partition_kernel.rs | 8 ++++---- diskann/src/graph/pipnn/simd.rs | 4 ++-- diskann/src/graph/pipnn/topk.rs | 14 +++++++------- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index ecd58dbc5d..16b5a89271 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -13,7 +13,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ leaf_metric::LeafMetric, - simd::PiPNNSIMDSchema, + simd::Simd, topk::{Candidate, TopK, TopKVisitor, Width, with_topk}, }; @@ -66,7 +66,7 @@ pub(super) fn select_leaf_neighbors( workspace: &mut LeafKernelWorkspace, ) -> ANNResult<()> where - A: PiPNNSIMDSchema, + A: Simd, M: LeafMetric, { let point_count = points.nrows(); @@ -86,7 +86,7 @@ where } /// Offer each lower-triangle row to the two endpoint neighbor lists. -fn rank_leaf_distances( +fn rank_leaf_distances( arch: A, distances: MatrixView<'_, f32>, output: MutMatrixView<'_, Candidate>, @@ -110,7 +110,7 @@ struct RankLeaf<'a, A> { thresholds: &'a mut Vec, } -impl TopKVisitor for RankLeaf<'_, A> { +impl TopKVisitor for RankLeaf<'_, A> { #[inline] fn visit(mut self, topk: TopK) { topk.initialize(self.output.as_mut_view(), self.thresholds); diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index d3938ef28f..3b0733c548 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -14,7 +14,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ partition_metric::PartitionMetric, - simd::PiPNNSIMDSchema, + simd::Simd, topk::{Candidate, TopK, TopKVisitor, UNASSIGNED, Width, with_topk}, }; @@ -49,7 +49,7 @@ pub(super) fn assign_leaders( workspace: &mut PartitionKernelWorkspace, ) -> ANNResult<()> where - A: PiPNNSIMDSchema, + A: Simd, M: PartitionMetric, { let point_count = points.nrows(); @@ -80,7 +80,7 @@ where } /// Select leader columns from each point's distance row. -fn rank_leader_distances( +fn rank_leader_distances( arch: A, distances: MatrixView<'_, f32>, output: MutMatrixView<'_, u32>, @@ -105,7 +105,7 @@ struct RankLeaders<'a, A> { candidates: &'a mut Vec, } -impl TopKVisitor for RankLeaders<'_, A> { +impl TopKVisitor for RankLeaders<'_, A> { #[inline] fn visit(mut self, topk: TopK) { for (distances, output) in self.distances.row_iter().zip(self.output.row_iter_mut()) { diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 0cc6b9678f..3b2a80863a 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -13,7 +13,7 @@ use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDVector}; type DefaultVector = ::f32x16; /// PiPNN SIMD representation for one architecture. -pub(super) trait PiPNNSIMDSchema: Architecture { +pub(super) trait Simd: Architecture { /// SIMD vector used by both ranking kernels. type Vector: SIMDVector + SIMDFloat; @@ -21,7 +21,7 @@ pub(super) trait PiPNNSIMDSchema: Architecture { fn active_lanes(mask: ::Mask) -> u64; } -impl PiPNNSIMDSchema for A +impl Simd for A where A: Architecture, DefaultVector: SIMDVector + SIMDFloat, diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index d4378087f4..ad78e9e89b 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -9,7 +9,7 @@ use diskann_utils::views::MutMatrixView; use diskann_wide::{SIMDPartialOrd, SIMDVector}; -use super::simd::PiPNNSIMDSchema; +use super::simd::Simd; /// An output slot with no assigned candidate. pub(super) const UNASSIGNED: u32 = u32::MAX; @@ -125,7 +125,7 @@ impl TopK { /// Select the nearest candidates, replacing the previous contents of output. /// Slice positions become candidate IDs. The threshold is local to this call. #[inline(always)] - pub(super) fn select_topk( + pub(super) fn select_topk( &self, arch: A, distances: &[f32], @@ -154,7 +154,7 @@ impl TopK { /// Call [`Self::initialize`] once before the pair scan; subsequent updates preserve its state. /// Preserve existing results; all slice indexes must be less than point_idx. #[inline] - pub(super) fn update_dual_topk( + pub(super) fn update_dual_topk( &self, arch: A, point_idx: usize, @@ -208,7 +208,7 @@ impl TopK { /// A complete SIMD group or one tail distance. This never leaves the TopK module. /// Both update directions share the loaded vector; scalar reads borrow the input. #[derive(Clone, Copy)] -enum DistanceBlock<'a, A: PiPNNSIMDSchema> { +enum DistanceBlock<'a, A: Simd> { Simd { first_candidate: usize, values: A::Vector, @@ -222,7 +222,7 @@ enum DistanceBlock<'a, A: PiPNNSIMDSchema> { /// Keep all distance loads and tail indexing inside TopK's architecture scope. #[inline(always)] -fn distance_blocks( +fn distance_blocks( arch: A, distances: &[f32], ) -> impl Iterator> { @@ -247,7 +247,7 @@ fn distance_blocks( ) } -impl DistanceBlock<'_, A> { +impl DistanceBlock<'_, A> { /// Update one top-k result from this block's candidates. /// /// Each candidate uses its index in the full distance row as its local ID. @@ -404,7 +404,7 @@ mod tests { use rstest::rstest; const EMPTY: Candidate = Candidate::new(UNASSIGNED, f32::INFINITY); - const LANES: usize = <::Vector as SIMDVector>::LANES; + const LANES: usize = <::Vector as SIMDVector>::LANES; #[rstest] #[case::first_batch(0, 3)] From 12dd73fe9a4f710c3604e9c93846d5b0ad4570b2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 06:03:09 +0000 Subject: [PATCH 198/229] refactor(pipnn): align Width method naming --- diskann/src/graph/pipnn/topk.rs | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index ad78e9e89b..02e27e85c3 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -50,19 +50,19 @@ pub(super) struct Runtime(pub(super) usize); /// Keep fixed capacities available to inlined insertion across architecture scopes. pub(super) trait Width: Copy { - fn capacity(self) -> usize; + fn width(self) -> usize; } impl Width for Fixed { #[inline(always)] - fn capacity(self) -> usize { + fn width(self) -> usize { K } } impl Width for Runtime { #[inline(always)] - fn capacity(self) -> usize { + fn width(self) -> usize { self.0 } } @@ -134,7 +134,7 @@ impl TopK { arch.run2( #[inline(always)] move |distances: &[f32], output: &mut [Candidate]| { - let k = self.width.capacity(); + let k = self.width.width(); if k == 0 { return; } @@ -167,7 +167,7 @@ impl TopK { move |distances: &[f32], mut output: MutMatrixView<'_, Candidate>, thresholds: &mut [f32]| { - let k = self.width.capacity(); + let k = self.width.width(); debug_assert_eq!( output.ncols(), k, @@ -254,7 +254,7 @@ impl DistanceBlock<'_, A> { /// Partition ranking uses leader IDs. Leaf ranking uses IDs of earlier points in the leaf. /// /// Compare candidates against the existing result without clearing its slots. - /// The first `width.capacity()` slots of `nearest` hold this result. + /// The first `width.width()` slots of `nearest` hold this result. /// The distance limit is the last slot's distance, or positive infinity while slots remain empty. /// Each insertion can lower this limit. Return the updated limit for the next block. #[inline(always)] @@ -322,7 +322,7 @@ impl DistanceBlock<'_, A> { /// Each candidate's local ID selects its result row in the full leaf, not its SIMD lane. /// If that row accepts the pair, insert `point_idx` with the pair's distance. /// - /// `candidates` starts at leaf row zero, with `width.capacity()` slots per row. + /// `candidates` starts at leaf row zero, with `width.width()` slots per row. /// `thresholds` holds one distance limit per row. /// The caller initializes both buffers once before the leaf scan. /// These updates retain earlier candidates unless a nearer point displaces them. @@ -338,7 +338,7 @@ impl DistanceBlock<'_, A> { point_idx: u32, width: W, ) { - let k = width.capacity(); + let k = width.width(); match self { Self::Scalar { candidate_idx, @@ -378,11 +378,11 @@ impl DistanceBlock<'_, A> { } /// Insert an eligible candidate in nearest-first order and return the new distance limit. -/// Only the first `width.capacity()` slots belong to this result. +/// Only the first `width.width()` slots belong to this result. /// The caller must check the candidate against the current limit before insertion. #[inline(always)] fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate, width: W) -> f32 { - let k = width.capacity(); + let k = width.width(); // Keep this slice here for performance. // Fixed gives the slice a compile-time constant length. // After inlining, LLVM can unroll the loop and remove bounds checks. From 9d45ac77af3e9f282e31a51a1d615e0d21c6ea83 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 06:29:36 +0000 Subject: [PATCH 199/229] ci(pipnn): focus Miri on TopK boundaries --- .github/workflows/nightly.yml | 13 ++++--------- diskann/src/graph/pipnn/leaf_kernel.rs | 7 ------- diskann/src/graph/pipnn/leaf_metric.rs | 8 -------- diskann/src/graph/pipnn/mod.rs | 2 +- diskann/src/graph/pipnn/partition_kernel.rs | 6 ------ diskann/src/graph/pipnn/partition_metric.rs | 7 ------- 6 files changed, 5 insertions(+), 38 deletions(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 58505d0028..b4de125c71 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -136,15 +136,10 @@ jobs: env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance - - name: PiPNN numerical pointer boundaries + - name: PiPNN TopK SIMD load boundaries env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance run: | - cargo +nightly miri test --locked -p diskann --features pipnn --lib \ - graph::pipnn::leaf_metric - cargo +nightly miri test --locked -p diskann --features pipnn --lib \ - graph::pipnn::partition_metric - cargo +nightly miri test --locked -p diskann --features pipnn --lib \ - graph::pipnn::leaf_kernel::tests::rank_leaf_distances_tests - cargo +nightly miri test --locked -p diskann --features pipnn --lib \ - graph::pipnn::topk + cargo +nightly miri test --locked -p diskann --features pipnn --lib -- \ + graph::pipnn::topk::tests::batch_selection_matches_a_full_sort \ + graph::pipnn::topk::tests::pair_scans_match_full_row_sort_across_capacities_and_lengths diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 16b5a89271..6a2aa94915 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -152,10 +152,8 @@ fn validate_output( #[cfg(test)] mod tests { use super::*; - #[cfg(not(miri))] use crate::graph::pipnn::test_support; use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2}; - #[cfg(not(miri))] use diskann_vector::distance::Metric; use diskann_wide::ARCH; use rstest::rstest; @@ -177,7 +175,6 @@ mod tests { assert_eq!(leaf_neighbor_count(points, requested), expected); } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2, [[1,3,2,4], [0,2,3,4], [4,0,1,3], [0,2,1,4], [2,1,0,3]])] #[case::cosine(Cosine, Metric::Cosine, [[3,1,2,4], [0,3,4,2], [4,3,0,1], [0,2,1,4], [2,1,3,0]])] @@ -243,7 +240,6 @@ mod tests { } } - #[cfg(not(miri))] #[test] fn neighbors_match_scalar_ranking_across_leaf_sizes_and_counts() { for point_count in [1, 2, 16, 17, 18, 33, 34] { @@ -290,7 +286,6 @@ mod tests { } } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2)] #[case::cosine(Cosine, Metric::Cosine)] @@ -375,7 +370,6 @@ mod tests { } } - #[cfg(not(miri))] #[test] fn workspace_reuse_does_not_mix_results_from_different_leaves() { let values = [0.0, 1.0, 4.0, 10.0, 21.0]; @@ -467,7 +461,6 @@ mod tests { assert_eq!(output, [Candidate::new(1, 7.0), Candidate::new(0, 7.0)]); } - // Keep this module path available to the nightly Miri selector. mod rank_leaf_distances_tests { use super::*; diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index f29ad151b2..58080e09f4 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -114,13 +114,10 @@ impl LeafMetric for CosineNormalized { #[cfg(test)] mod tests { use super::*; - #[cfg(not(miri))] use crate::graph::pipnn::test_support; - #[cfg(not(miri))] use diskann_vector::distance::Metric; use rstest::rstest; - #[cfg(not(miri))] #[test] fn l2_distance_retains_small_coordinate_contributions() { // Given: 4096^2 + 127 unit coordinates, orthogonal to a unit vector. @@ -145,7 +142,6 @@ mod tests { ); } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2)] #[case::cosine(Cosine, Metric::Cosine)] @@ -202,7 +198,6 @@ mod tests { } } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2)] #[case::cosine(Cosine, Metric::Cosine)] @@ -260,7 +255,6 @@ mod tests { } } - #[cfg(not(miri))] #[rstest] #[case::squared_l2(L2, &[2.0, 0.0, 0.0, 3.0, -4.0, 0.0], [0.0, 13.0, 0.0, 36.0, 25.0, 0.0])] #[case::negative_dot(InnerProduct, &[2.0, 0.0, 0.0, 3.0, -4.0, 0.0], [-4.0, 0.0, -9.0, 8.0, 0.0, -16.0])] @@ -284,7 +278,6 @@ mod tests { ); } - #[cfg(not(miri))] #[rstest] #[case::zero(0.0)] #[case::squared_norm_underflows(f32::MIN_POSITIVE)] @@ -298,7 +291,6 @@ mod tests { assert_eq!([output[0], output[2], output[3]], [1.0, 1.0, 0.0]); } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, 2.0)] #[case::cosine(Cosine, 1.0)] diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 6a0b167ffb..95bfbd5ccc 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -113,7 +113,7 @@ mod cosine_distance_tests { } } -#[cfg(all(test, not(miri)))] +#[cfg(test)] mod test_support { use diskann_vector::distance::Metric; diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 3b0733c548..0eca811bda 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -120,15 +120,12 @@ impl TopKVisitor for RankLeaders<'_, A> { #[cfg(test)] mod tests { use super::*; - #[cfg(not(miri))] use crate::graph::pipnn::test_support; use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2}; - #[cfg(not(miri))] use diskann_vector::distance::Metric; use diskann_wide::ARCH; use rstest::rstest; - #[cfg(not(miri))] #[rstest] #[case::l2(L2, false, [[0,1,3,2], [1,2,3,0], [3,0,2,1]])] #[case::cosine(Cosine, false, [[0,1,3,2], [1,2,0,3], [3,0,2,1]])] @@ -181,7 +178,6 @@ mod tests { } } - #[cfg(not(miri))] #[test] fn assignments_match_nearest_leaders_across_counts_and_widths() { for leader_count in [1, 5, 16, 17, 33] { @@ -228,7 +224,6 @@ mod tests { } } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2)] #[case::cosine(Cosine, Metric::Cosine)] @@ -312,7 +307,6 @@ mod tests { } } - #[cfg(not(miri))] #[test] fn workspace_reuse_replaces_assignments_for_each_stripe() { let leader_values = [0.0, 5.0, 12.0]; diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index f5c20d0581..ca597ae60d 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -225,13 +225,10 @@ impl PartitionMetric for CosineNormalized { #[cfg(test)] mod tests { use super::*; - #[cfg(not(miri))] use crate::graph::pipnn::test_support; - #[cfg(not(miri))] use diskann_vector::distance::Metric; use rstest::rstest; - #[cfg(not(miri))] #[test] fn l2_ranking_retains_small_coordinate_contributions() { // Given: the origin ranks leaders solely by their squared norms. @@ -264,7 +261,6 @@ mod tests { assert_eq!(output[1], 16_777_232.0); } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2)] #[case::cosine(Cosine, Metric::Cosine)] @@ -353,7 +349,6 @@ mod tests { } } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, Metric::L2)] #[case::cosine(Cosine, Metric::Cosine)] @@ -427,7 +422,6 @@ mod tests { } } - #[cfg(not(miri))] #[rstest] #[case::l2(L2, &[2.0, 0.0, 0.0, 3.0], &[1.0, 0.0, 0.0, -2.0, -3.0, 0.0], [-3.0, 4.0, 21.0, 1.0, 16.0, 9.0])] #[case::cosine(Cosine, &[2.0, 0.0, 0.0, 3.0], &[1.0, 0.0, 0.0, -2.0, -3.0, 0.0], [0.0, 1.0, 2.0, 1.0, 2.0, 1.0])] @@ -454,7 +448,6 @@ mod tests { } } - #[cfg(not(miri))] #[test] fn cosine_assigns_unit_distance_when_either_vector_has_zero_norm() { let point_values = [0.0, 0.0, 0.0, 2.0]; From 6df51c840e891df0e74d4c1b4f46b5cdaf77d0a7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 06:51:13 +0000 Subject: [PATCH 200/229] ci(pipnn): separate ARM runner addition --- .github/workflows/ci.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f8b585fcc4..16a2879ed4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -463,7 +463,6 @@ jobs: os: - windows-latest - ubuntu-latest - - ubuntu-24.04-arm steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 From f673e95c343f112f98d56225e3da18bac1dc997f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 07:31:41 +0000 Subject: [PATCH 201/229] ci(pipnn): remove kernel tests from SDE --- .github/workflows/ci.yml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 16a2879ed4..393c5d66bf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -354,8 +354,6 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ - --package diskann \ - --features pipnn \ -- --skip compile_tests \ --skip pivots::tests::run_test_happy_path @@ -418,8 +416,6 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ - --package diskann \ - --features pipnn \ -- --skip compile_tests test-workspace: From 470915d46a1d203c71dadd6169190a25048ca954 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 07:53:44 +0000 Subject: [PATCH 202/229] perf(pipnn): limit TopK specializations --- diskann/src/graph/pipnn/topk.rs | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 02e27e85c3..9587089e1e 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -75,7 +75,7 @@ pub(super) trait TopKVisitor { /// Select the width specialization and run the visitor once for the batch. /// -/// Specialize capacities 1..=10; zero and larger capacities use the runtime path. +/// Specialize capacities 1, 2, 3, 8, and 10; other capacities use the runtime path. /// Typical partition K is 10 or 3, and leaf K is 3 or 2. #[inline] pub(super) fn with_topk(width: usize, visitor: V) { @@ -83,12 +83,7 @@ pub(super) fn with_topk(width: usize, visitor: V) { 1 => visitor.visit(TopK::new(Fixed::<1>)), 2 => visitor.visit(TopK::new(Fixed::<2>)), 3 => visitor.visit(TopK::new(Fixed::<3>)), - 4 => visitor.visit(TopK::new(Fixed::<4>)), - 5 => visitor.visit(TopK::new(Fixed::<5>)), - 6 => visitor.visit(TopK::new(Fixed::<6>)), - 7 => visitor.visit(TopK::new(Fixed::<7>)), 8 => visitor.visit(TopK::new(Fixed::<8>)), - 9 => visitor.visit(TopK::new(Fixed::<9>)), 10 => visitor.visit(TopK::new(Fixed::<10>)), _ => visitor.visit(TopK::new(Runtime(width))), } From 5ecfd4af2ace588d04961bfe5bfa3e729712def6 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 08:15:07 +0000 Subject: [PATCH 203/229] perf(pipnn): avoid duplicate threshold fills --- diskann/src/graph/pipnn/topk.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index 9587089e1e..de9f082318 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -113,8 +113,9 @@ impl TopK { thresholds: &mut Vec, ) { output.as_mut_slice().fill(Candidate::default()); - thresholds.resize(output.nrows(), f32::INFINITY); + thresholds.truncate(output.nrows()); thresholds.fill(f32::INFINITY); + thresholds.resize(output.nrows(), f32::INFINITY); } /// Select the nearest candidates, replacing the previous contents of output. From e4c3edd567d5c34d6bb947a6ec67e8795af41102 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 22 Sep 2026 11:53:38 +0000 Subject: [PATCH 204/229] refactor(pipnn): bind TopK storage to rankers --- diskann/src/graph/pipnn/leaf_kernel.rs | 30 +- diskann/src/graph/pipnn/partition_kernel.rs | 12 +- diskann/src/graph/pipnn/topk.rs | 682 +++++++++++--------- 3 files changed, 406 insertions(+), 318 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 6a2aa94915..4335ca1c1a 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -14,7 +14,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ leaf_metric::LeafMetric, simd::Simd, - topk::{Candidate, TopK, TopKVisitor, Width, with_topk}, + topk::{Batch, BatchRanker, BatchVisitor, Candidate, with_batch}, }; /// Reusable storage for one leaf numerical pipeline. @@ -92,36 +92,22 @@ fn rank_leaf_distances( output: MutMatrixView<'_, Candidate>, worst: &mut Vec, ) { - with_topk( - output.ncols(), - RankLeaf { - arch, - distances, - output, - thresholds: worst, - }, - ); + with_batch(output, worst, RankLeaf { arch, distances }); } struct RankLeaf<'a, A> { arch: A, distances: MatrixView<'a, f32>, - output: MutMatrixView<'a, Candidate>, - thresholds: &'a mut Vec, } -impl TopKVisitor for RankLeaf<'_, A> { +impl BatchVisitor for RankLeaf<'_, A> { #[inline] - fn visit(mut self, topk: TopK) { - topk.initialize(self.output.as_mut_view(), self.thresholds); + fn visit<'a, Rows>(self, mut batch: Batch<'a, Rows>) + where + Batch<'a, Rows>: BatchRanker, + { for point_idx in 1..self.distances.nrows() { - topk.update_dual_topk( - self.arch, - point_idx, - &self.distances.row(point_idx)[..point_idx], - self.output.as_mut_view(), - self.thresholds.as_mut_slice(), - ); + batch.update_dual_topk(self.arch, &self.distances.row(point_idx)[..point_idx]); } } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index 0eca811bda..b1a39becf9 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -15,7 +15,7 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ partition_metric::PartitionMetric, simd::Simd, - topk::{Candidate, TopK, TopKVisitor, UNASSIGNED, Width, with_topk}, + topk::{Candidate, Ranker, TopKVisitor, UNASSIGNED, with_topk}, }; /// No sampled partition center was rankable for this output slot. @@ -88,12 +88,11 @@ fn rank_leader_distances( ) { candidates.resize(output.ncols(), Candidate::default()); with_topk( - output.ncols(), + candidates.as_mut_slice(), RankLeaders { arch, distances, output, - candidates, }, ); } @@ -102,15 +101,14 @@ struct RankLeaders<'a, A> { arch: A, distances: MatrixView<'a, f32>, output: MutMatrixView<'a, u32>, - candidates: &'a mut Vec, } impl TopKVisitor for RankLeaders<'_, A> { #[inline] - fn visit(mut self, topk: TopK) { + fn visit(mut self, mut ranker: R) { for (distances, output) in self.distances.row_iter().zip(self.output.row_iter_mut()) { - topk.select_topk(self.arch, distances, self.candidates.as_mut_slice()); - for (destination, candidate) in output.iter_mut().zip(self.candidates.iter()) { + ranker.select_topk(self.arch, distances); + for (destination, candidate) in output.iter_mut().zip(ranker.as_ref().iter()) { *destination = candidate.local_idx; } } diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index de9f082318..e22db2bfb4 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -40,163 +40,229 @@ impl Default for Candidate { } } -/// A result capacity known at compile time. -#[derive(Clone, Copy, Debug)] -pub(super) struct Fixed; +/// One borrowed result row whose length is its capacity. +pub(super) trait Ranker: AsRef<[Candidate]> + AsMut<[Candidate]> + Sized { + /// Insert a candidate below this non-empty row's limit and return the new limit. + fn insert(&mut self, candidate: Candidate) -> f32; -/// A result capacity selected at runtime. -#[derive(Clone, Copy, Debug)] -pub(super) struct Runtime(pub(super) usize); - -/// Keep fixed capacities available to inlined insertion across architecture scopes. -pub(super) trait Width: Copy { - fn width(self) -> usize; + /// Replace this row with the nearest input positions. + /// + /// Equal distances need no fixed tie order. Empty rows retain no candidates. + #[inline(always)] + fn select_topk(&mut self, arch: A, distances: &[f32]) { + arch.run2( + #[inline(always)] + move |distances: &[f32], ranker: &mut Self| { + if ranker.as_ref().is_empty() { + return; + } + ranker.as_mut().fill(Candidate::default()); + distance_blocks(arch, distances).fold( + f32::INFINITY, + #[inline(always)] + |limit, block| block.update_topk_from_candidates(ranker, limit), + ); + }, + distances, + self, + ); + } } -impl Width for Fixed { +impl Ranker for &mut [Candidate] { #[inline(always)] - fn width(self) -> usize { - K + fn insert(&mut self, candidate: Candidate) -> f32 { + insert_sorted(self, candidate) } } -impl Width for Runtime { +impl Ranker for &mut [Candidate; K] { #[inline(always)] - fn width(self) -> usize { - self.0 + fn insert(&mut self, candidate: Candidate) -> f32 { + insert_sorted(*self, candidate) } } -/// A batch operation that accepts any TopK width specialization. +/// A batch operation over one bound result row. pub(super) trait TopKVisitor { - /// Consume the operation and run it with the selected width. - fn visit(self, topk: TopK); + fn visit(self, ranker: R); } -/// Select the width specialization and run the visitor once for the batch. +/// Bind one result row and specialize its capacity once for the visitor. /// -/// Specialize capacities 1, 2, 3, 8, and 10; other capacities use the runtime path. +/// Specialize capacities 1, 2, 3, 8, and 10; other capacities use the slice path. /// Typical partition K is 10 or 3, and leaf K is 3 or 2. #[inline] -pub(super) fn with_topk(width: usize, visitor: V) { - match width { - 1 => visitor.visit(TopK::new(Fixed::<1>)), - 2 => visitor.visit(TopK::new(Fixed::<2>)), - 3 => visitor.visit(TopK::new(Fixed::<3>)), - 8 => visitor.visit(TopK::new(Fixed::<8>)), - 10 => visitor.visit(TopK::new(Fixed::<10>)), - _ => visitor.visit(TopK::new(Runtime(width))), +#[expect( + clippy::unwrap_used, + reason = "each match arm proves the candidate slice has the array length" +)] +pub(super) fn with_topk(candidates: &mut [Candidate], visitor: V) { + // Each arm checks this same slice's length, so conversion to its array cannot fail. + match candidates.len() { + 1 => visitor.visit::<&mut [_; 1]>(candidates.try_into().unwrap()), + 2 => visitor.visit::<&mut [_; 2]>(candidates.try_into().unwrap()), + 3 => visitor.visit::<&mut [_; 3]>(candidates.try_into().unwrap()), + 8 => visitor.visit::<&mut [_; 8]>(candidates.try_into().unwrap()), + 10 => visitor.visit::<&mut [_; 10]>(candidates.try_into().unwrap()), + _ => visitor.visit(candidates), } } -/// Top-k selection with a fixed result capacity and caller-owned storage. -/// -/// This object stores no candidates or thresholds. Each operation borrows only -/// the buffers it updates. Candidate IDs are slice positions, not dataset IDs. -/// Callers supply each candidate at most once per result; equal distances need -/// no fixed tie order. [`with_topk`] specializes the width per batch. -/// Selection and dual updates run inside their own architecture scope. -pub(super) struct TopK { - width: W, +/// Access complete result rows with their static or runtime capacity. +pub(super) trait BatchRanker { + type Ranker<'a>: Ranker + where + Self: 'a; + + fn ranker(&mut self, index: usize) -> Option>; } -impl TopK { - /// Configure the result capacity. - pub(super) fn new(width: W) -> Self { - Self { width } - } +/// Candidate rows and their limits stay borrowed together throughout a pair scan. +pub(super) struct Batch<'a, Rows> { + candidates: Rows, + thresholds: &'a mut [f32], +} - /// Clear results and prepare reusable thresholds before a sequence of dual updates. - pub(super) fn initialize( - &self, - mut output: MutMatrixView<'_, Candidate>, - thresholds: &mut Vec, - ) { - output.as_mut_slice().fill(Candidate::default()); - thresholds.truncate(output.nrows()); +impl<'a> Batch<'a, MutMatrixView<'a, Candidate>> { + fn new(mut candidates: MutMatrixView<'a, Candidate>, thresholds: &'a mut Vec) -> Self { + let rows = candidates.nrows(); + candidates.as_mut_slice().fill(Candidate::default()); + thresholds.truncate(rows); thresholds.fill(f32::INFINITY); - thresholds.resize(output.nrows(), f32::INFINITY); + thresholds.resize(rows, f32::INFINITY); + Self { + candidates, + thresholds, + } + } +} + +impl Batch<'_, Rows> { + fn nrows(&self) -> usize { + self.thresholds.len() + } +} + +impl BatchRanker for Batch<'_, &mut [[Candidate; K]]> { + type Ranker<'a> + = &'a mut [Candidate; K] + where + Self: 'a; + + #[inline(always)] + fn ranker(&mut self, index: usize) -> Option> { + self.candidates.get_mut(index) } +} + +impl BatchRanker for Batch<'_, MutMatrixView<'_, Candidate>> { + type Ranker<'a> + = &'a mut [Candidate] + where + Self: 'a; - /// Select the nearest candidates, replacing the previous contents of output. - /// Slice positions become candidate IDs. The threshold is local to this call. #[inline(always)] - pub(super) fn select_topk( - &self, - arch: A, - distances: &[f32], - output: &mut [Candidate], - ) { - arch.run2( - #[inline(always)] - move |distances: &[f32], output: &mut [Candidate]| { - let k = self.width.width(); - if k == 0 { - return; - } - output[..k].fill(Candidate::default()); - distance_blocks(arch, distances).fold( - f32::INFINITY, - #[inline(always)] - |limit, block| block.update_topk_from_candidates(output, limit, self.width), - ); - }, - distances, - output, - ); + fn ranker(&mut self, index: usize) -> Option> { + if index < self.nrows() { + Some(self.candidates.row_mut(index)) + } else { + None + } } +} + +/// A batch operation over initialized pair-scan storage. +pub(super) trait BatchVisitor { + fn visit<'a, Rows>(self, batch: Batch<'a, Rows>) + where + Batch<'a, Rows>: BatchRanker; +} - /// Add each pair (point_idx, j) independently to both endpoints' top-k results. - /// Call [`Self::initialize`] once before the pair scan; subsequent updates preserve its state. - /// Preserve existing results; all slice indexes must be less than point_idx. +/// Initialize a result matrix and specialize its row representation once for the visitor. +#[inline] +pub(super) fn with_batch( + candidates: MutMatrixView<'_, Candidate>, + thresholds: &mut Vec, + visitor: V, +) { + let Batch { + candidates, + thresholds, + } = Batch::new(candidates, thresholds); + // The matrix width selects K, so its storage contains only complete K-element rows. + match candidates.ncols() { + 1 => visitor.visit(Batch { + candidates: candidates.into_inner().as_chunks_mut::<1>().0, + thresholds, + }), + 2 => visitor.visit(Batch { + candidates: candidates.into_inner().as_chunks_mut::<2>().0, + thresholds, + }), + 3 => visitor.visit(Batch { + candidates: candidates.into_inner().as_chunks_mut::<3>().0, + thresholds, + }), + 8 => visitor.visit(Batch { + candidates: candidates.into_inner().as_chunks_mut::<8>().0, + thresholds, + }), + 10 => visitor.visit(Batch { + candidates: candidates.into_inner().as_chunks_mut::<10>().0, + thresholds, + }), + _ => visitor.visit(Batch { + candidates, + thresholds, + }), + } +} + +impl<'a, Rows> Batch<'a, Rows> +where + Self: BatchRanker, +{ + /// Offer one point's distances to all preceding points to both endpoints' results. + /// + /// The source point index is `distances.len()`: each entry describes its pair + /// with the earlier point at that position. Supply each pair once; earlier results are preserved. + /// Equal distances need no fixed tie order. + /// + /// # Panics + /// + /// Panics if `distances.len()` is outside the batch. #[inline] - pub(super) fn update_dual_topk( - &self, - arch: A, - point_idx: usize, - distances: &[f32], - output: MutMatrixView<'_, Candidate>, - thresholds: &mut [f32], - ) { - arch.run3( + #[expect( + clippy::unwrap_used, + reason = "the first ranker lookup rejects an invalid point before scanning" + )] + pub(super) fn update_dual_topk(&mut self, arch: A, distances: &[f32]) { + let point_idx = distances.len(); + // ranker rejects an out-of-bounds point before any results are changed. + if self.ranker(point_idx).unwrap().as_ref().is_empty() { + return; + } + arch.run2( #[inline(always)] - move |distances: &[f32], - mut output: MutMatrixView<'_, Candidate>, - thresholds: &mut [f32]| { - let k = self.width.width(); - debug_assert_eq!( - output.ncols(), - k, - "top-k output width must match its capacity" - ); - if k == 0 { - return; - } - debug_assert!( - distances.len() <= point_idx, - "candidate indexes must exclude the updated point" - ); - // Keep this point's result and local limit independent of reciprocal updates. - let (others, remaining) = output.as_mut_slice().split_at_mut(point_idx * k); + move |distances: &[f32], batch: &mut Self| { let limit = distance_blocks(arch, distances).fold( - thresholds[point_idx], + batch.thresholds[point_idx], #[inline(always)] |limit, block| { - let limit = block.update_topk_from_candidates(remaining, limit, self.width); - block.update_topks_with_point( - others, - thresholds, - point_idx as u32, - self.width, - ); + // End the source row borrow before offering this block to earlier rows. + let limit = { + let mut nearest = batch.ranker(point_idx).unwrap(); + block.update_topk_from_candidates(&mut nearest, limit) + }; + block.update_topks_with_point(batch, point_idx as u32); limit }, ); - thresholds[point_idx] = limit; + batch.thresholds[point_idx] = limit; }, distances, - output, - thresholds, + self, ); } } @@ -249,28 +315,18 @@ impl DistanceBlock<'_, A> { /// Each candidate uses its index in the full distance row as its local ID. /// Partition ranking uses leader IDs. Leaf ranking uses IDs of earlier points in the leaf. /// - /// Compare candidates against the existing result without clearing its slots. - /// The first `width.width()` slots of `nearest` hold this result. + /// Compare candidates against the complete result row without clearing its slots. /// The distance limit is the last slot's distance, or positive infinity while slots remain empty. /// Each insertion can lower this limit. Return the updated limit for the next block. #[inline(always)] - fn update_topk_from_candidates( - self, - nearest: &mut [Candidate], - mut max_distance: f32, - width: W, - ) -> f32 { + fn update_topk_from_candidates(self, nearest: &mut R, mut max_distance: f32) -> f32 { match self { Self::Scalar { candidate_idx, distance, } => { if distance < max_distance { - max_distance = insert_sorted( - nearest, - Candidate::new(candidate_idx as u32, distance), - width, - ); + max_distance = nearest.insert(Candidate::new(candidate_idx as u32, distance)); } } Self::Simd { @@ -283,11 +339,8 @@ impl DistanceBlock<'_, A> { if max_distance == f32::INFINITY { for (lane, &distance) in distances.iter().enumerate() { if distance < max_distance { - max_distance = insert_sorted( - nearest, - Candidate::new((first_candidate + lane) as u32, distance), - width, - ); + max_distance = nearest + .insert(Candidate::new((first_candidate + lane) as u32, distance)); } } } else { @@ -299,11 +352,10 @@ impl DistanceBlock<'_, A> { eligible &= eligible - 1; // Earlier insertions in this block can lower the limit. if distances[lane] < max_distance { - max_distance = insert_sorted( - nearest, - Candidate::new((first_candidate + lane) as u32, distances[lane]), - width, - ); + max_distance = nearest.insert(Candidate::new( + (first_candidate + lane) as u32, + distances[lane], + )); } } } @@ -318,35 +370,33 @@ impl DistanceBlock<'_, A> { /// Each candidate's local ID selects its result row in the full leaf, not its SIMD lane. /// If that row accepts the pair, insert `point_idx` with the pair's distance. /// - /// `candidates` starts at leaf row zero, with `width.width()` slots per row. - /// `thresholds` holds one distance limit per row. - /// The caller initializes both buffers once before the leaf scan. - /// These updates retain earlier candidates unless a nearer point displaces them. + /// Only the earlier rows identified by this block can change. + /// Their candidates are retained unless a nearer point displaces them. /// /// Each row compares the pair distance against its own limit. /// Its decision does not depend on the current point's top-k result. /// Update the row's limit after an insertion. #[inline(always)] - fn update_topks_with_point( - self, - candidates: &mut [Candidate], - thresholds: &mut [f32], - point_idx: u32, - width: W, - ) { - let k = width.width(); + #[expect( + clippy::unwrap_used, + reason = "all block positions precede the checked source point" + )] + fn update_topks_with_point<'a, Rows>(self, batch: &mut Batch<'a, Rows>, point_idx: u32) + where + Batch<'a, Rows>: BatchRanker, + { match self { Self::Scalar { candidate_idx, distance, } => { - let limit = &mut thresholds[candidate_idx]; - if distance < *limit { - *limit = insert_sorted( - &mut candidates[candidate_idx * k..], - Candidate::new(point_idx, distance), - width, - ); + if distance < batch.thresholds[candidate_idx] { + // Block positions are less than the source index checked by update_dual_topk. + let limit = batch + .ranker(candidate_idx) + .unwrap() + .insert(Candidate::new(point_idx, distance)); + batch.thresholds[candidate_idx] = limit; } } Self::Simd { @@ -354,7 +404,7 @@ impl DistanceBlock<'_, A> { values, distances, } => { - let limits = &mut thresholds[first_candidate..][..A::Vector::LANES]; + let limits = &batch.thresholds[first_candidate..][..A::Vector::LANES]; // SAFETY: the slice above has one distance limit per SIMD lane. let limit_values = unsafe { A::Vector::load_simd(values.arch(), limits.as_ptr()) }; let mut eligible = A::active_lanes(values.lt_simd(limit_values)); @@ -362,11 +412,13 @@ impl DistanceBlock<'_, A> { let lane = eligible.trailing_zeros() as usize; eligible &= eligible - 1; // Each lane updates a different result; its limit is current. - limits[lane] = insert_sorted( - &mut candidates[(first_candidate + lane) * k..], - Candidate::new(point_idx, distances[lane]), - width, - ); + let index = first_candidate + lane; + // Every SIMD lane describes an earlier, in-bounds point. + let limit = batch + .ranker(index) + .unwrap() + .insert(Candidate::new(point_idx, distances[lane])); + batch.thresholds[index] = limit; } } } @@ -374,15 +426,10 @@ impl DistanceBlock<'_, A> { } /// Insert an eligible candidate in nearest-first order and return the new distance limit. -/// Only the first `width.width()` slots belong to this result. +/// `nearest` contains exactly one non-empty result row. /// The caller must check the candidate against the current limit before insertion. #[inline(always)] -fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate, width: W) -> f32 { - let k = width.width(); - // Keep this slice here for performance. - // Fixed gives the slice a compile-time constant length. - // After inlining, LLVM can unroll the loop and remove bounds checks. - let nearest = &mut nearest[..k]; +fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate) -> f32 { let last = nearest.len() - 1; let mut slot = last; while slot > 0 && candidate.distance < nearest[slot - 1].distance { @@ -413,12 +460,11 @@ mod tests { #[case] rows: usize, ) { // Given: reusable storage contains assigned candidates and finite limits. - let topk = TopK::new(Fixed::<2>); let mut output = vec![[Candidate::new(0, -12.0), Candidate::new(1, -6.0)]; rows]; let mut limits = vec![-6.0; previous_rows]; // When: prepare storage for a new batch. - topk.initialize( + Batch::new( MutMatrixView::try_from(output.as_flattened_mut(), rows, 2).unwrap(), &mut limits, ); @@ -428,6 +474,55 @@ mod tests { assert_eq!(limits, vec![f32::INFINITY; rows]); } + #[test] + fn batch_rankers_borrow_one_complete_row_without_touching_adjacent_rows() { + struct UpdateMiddleRow; + + impl BatchVisitor for UpdateMiddleRow { + fn visit<'a, Rows>(self, mut batch: Batch<'a, Rows>) + where + Batch<'a, Rows>: BatchRanker, + { + assert!(batch.ranker(3).is_none()); + batch + .ranker(1) + .unwrap() + .as_mut() + .fill(Candidate::new(42, -5.0)); + } + } + + // Exercise array rows and runtime rows through the production dispatcher. + for capacity in [2, 4] { + let mut candidates = vec![Candidate::new(9, -12.0); 3 * capacity]; + let mut thresholds = vec![-12.0; 5]; + + with_batch( + MutMatrixView::try_from(candidates.as_mut_slice(), 3, capacity).unwrap(), + &mut thresholds, + UpdateMiddleRow, + ); + + assert_eq!(&candidates[..capacity], vec![EMPTY; capacity]); + assert_eq!( + &candidates[capacity..2 * capacity], + vec![Candidate::new(42, -5.0); capacity] + ); + assert_eq!(&candidates[2 * capacity..], vec![EMPTY; capacity]); + } + } + + #[test] + #[should_panic] + fn a_pair_scan_rejects_a_point_outside_the_batch() { + let mut candidates = [EMPTY; 3]; + let mut limits = Vec::new(); + let output = MutMatrixView::try_from(&mut candidates[..], 3, 1).unwrap(); + let mut batch = Batch::new(output, &mut limits); + + batch.update_dual_topk(ARCH, &[3.0, 4.0, 5.0]); + } + #[rstest] #[case::empty_input(&[], [EMPTY; 3])] #[case::no_rankable_values(&[f32::INFINITY, f32::NAN], [EMPTY; 3])] @@ -452,27 +547,25 @@ mod tests { #[case] expected: [Candidate; 3], ) { // Given: output still holds results from a different input. - let topk = TopK::new(Fixed::<3>); let mut output = [ Candidate::new(2, -20.0), Candidate::new(4, -10.0), Candidate::new(9, -1.0), ]; - topk.select_topk(ARCH, distances, &mut output); + (&mut output).select_topk(ARCH, distances); assert_eq!(output, expected); } #[test] fn reusing_output_starts_a_selection_with_a_fresh_distance_limit() { - let topk = TopK::new(Fixed::<2>); let mut output = [EMPTY; 2]; - topk.select_topk(ARCH, &[-20.0, -10.0], &mut output); + (&mut output).select_topk(ARCH, &[-20.0, -10.0]); assert_eq!(output, [Candidate::new(0, -20.0), Candidate::new(1, -10.0)]); // Every new distance exceeds the previous selection's limit. - topk.select_topk(ARCH, &[16.0, 3.0, 9.0], &mut output); + (&mut output).select_topk(ARCH, &[16.0, 3.0, 9.0]); assert_eq!(output, [Candidate::new(1, 3.0), Candidate::new(2, 9.0)]); } @@ -481,12 +574,11 @@ mod tests { fn batch_selection_matches_a_full_sort() { struct SelectRow<'a> { distances: &'a [f32], - output: &'a mut [Candidate], } impl TopKVisitor for SelectRow<'_> { - fn visit(self, topk: TopK) { - topk.select_topk(ARCH, self.distances, self.output); + fn visit(self, mut ranker: R) { + ranker.select_topk(ARCH, self.distances); } } @@ -518,21 +610,14 @@ mod tests { expected.truncate(capacity); expected.resize(capacity, EMPTY); let mut output = vec![EMPTY; capacity]; - let mut capacity_reads = 0; - with_topk( - { - capacity_reads += 1; - capacity - }, + &mut output, SelectRow { distances: &distances, - output: &mut output, }, ); assert_eq!(output, expected, "count={count}, capacity={capacity}"); - assert_eq!(capacity_reads, 1, "count={count}, capacity={capacity}"); } } } @@ -544,7 +629,7 @@ mod tests { distances[best_index] = -2.0; let mut output = [EMPTY; 4]; - TopK::new(Fixed::<4>).select_topk(ARCH, &distances, &mut output); + (&mut output).select_topk(ARCH, &distances); assert_eq!(output[0], Candidate::new(best_index as u32, -2.0)); // Any three distinct input positions with distance 8 are valid; tie order is free. @@ -574,68 +659,76 @@ mod tests { distances[index] = distance; let mut output = [EMPTY; 2]; - TopK::new(Runtime(2)).select_topk(ARCH, &distances, &mut output); + output.as_mut_slice().select_topk(ARCH, &distances); assert_eq!(output[0].local_idx, index as u32); assert_eq!(output[0].distance.to_bits(), distance.to_bits()); assert_eq!(output[1], EMPTY); } - #[rstest] - #[case::fixed(Fixed::<2>)] - #[case::runtime(Runtime(2))] - fn successive_pair_rows_accumulate_each_points_nearest_neighbors(#[case] width: impl Width) { - let topk = TopK::new(width); - let mut output = [[EMPTY; 2]; 4]; - let mut limits = Vec::new(); - let mut rows = MutMatrixView::try_from(output.as_flattened_mut(), 4, 2).unwrap(); - topk.initialize(rows.as_mut_view(), &mut limits); - - // Given: pairs 0-1=9, 0-2=2, and 1-2=7; point 3 has no offered pairs. - topk.update_dual_topk(ARCH, 1, &[9.0], rows.as_mut_view(), &mut limits); - topk.update_dual_topk(ARCH, 2, &[2.0, 7.0], rows.as_mut_view(), &mut limits); - let before_last_point = [ - [Candidate::new(2, 2.0), Candidate::new(1, 9.0)], - [Candidate::new(2, 7.0), Candidate::new(0, 9.0)], - [Candidate::new(0, 2.0), Candidate::new(1, 7.0)], - [EMPTY; 2], - ]; - assert_eq!(rows.as_slice(), before_last_point.as_flattened()); - assert_eq!(limits, [9.0, 9.0, 7.0, f32::INFINITY]); - - // When: point 3 is at distances 6, 1, and 4 from points 0, 1, and 2. - topk.update_dual_topk(ARCH, 3, &[6.0, 1.0, 4.0], rows.as_mut_view(), &mut limits); - - // Then: each endpoint keeps its own nearest two, with its farthest retained limit. - let expected = [ - [Candidate::new(2, 2.0), Candidate::new(3, 6.0)], - [Candidate::new(3, 1.0), Candidate::new(2, 7.0)], - [Candidate::new(0, 2.0), Candidate::new(3, 4.0)], - [Candidate::new(1, 1.0), Candidate::new(2, 4.0)], - ]; - assert_eq!(rows.as_slice(), expected.as_flattened()); - assert_eq!(limits, [6.0, 7.0, 4.0, 4.0]); + #[test] + fn successive_pair_rows_accumulate_each_points_nearest_neighbors() { + fn check<'a, Rows>(mut batch: Batch<'a, Rows>) + where + Batch<'a, Rows>: BatchRanker, + { + // Given: pairs 0-1=9, 0-2=2, and 1-2=7; point 3 has no offered pairs. + batch.update_dual_topk(ARCH, &[9.0]); + batch.update_dual_topk(ARCH, &[2.0, 7.0]); + let before_last_point = [ + [Candidate::new(2, 2.0), Candidate::new(1, 9.0)], + [Candidate::new(2, 7.0), Candidate::new(0, 9.0)], + [Candidate::new(0, 2.0), Candidate::new(1, 7.0)], + [EMPTY; 2], + ]; + for (index, expected) in before_last_point.iter().enumerate() { + assert_eq!(batch.ranker(index).unwrap().as_ref(), expected); + } + assert_eq!(batch.thresholds, [9.0, 9.0, 7.0, f32::INFINITY]); + + // When: point 3 is at distances 6, 1, and 4 from points 0, 1, and 2. + batch.update_dual_topk(ARCH, &[6.0, 1.0, 4.0]); + + // Then: each endpoint keeps its own nearest two, with its farthest retained limit. + let expected = [ + [Candidate::new(2, 2.0), Candidate::new(3, 6.0)], + [Candidate::new(3, 1.0), Candidate::new(2, 7.0)], + [Candidate::new(0, 2.0), Candidate::new(3, 4.0)], + [Candidate::new(1, 1.0), Candidate::new(2, 4.0)], + ]; + for (index, expected) in expected.iter().enumerate() { + assert_eq!(batch.ranker(index).unwrap().as_ref(), expected); + } + assert_eq!(batch.thresholds, [6.0, 7.0, 4.0, 4.0]); + } + + let mut fixed = [[EMPTY; 2]; 4]; + let mut fixed_limits = Vec::new(); + fixed_limits.resize(fixed.len(), f32::INFINITY); + check(Batch { + candidates: fixed.as_mut_slice(), + thresholds: &mut fixed_limits, + }); + + let mut dynamic = [EMPTY; 8]; + let mut dynamic_limits = Vec::new(); + let rows = MutMatrixView::try_from(&mut dynamic[..], 4, 2).unwrap(); + check(Batch::new(rows, &mut dynamic_limits)); } #[test] fn pair_scans_match_full_row_sort_across_capacities_and_lengths() { struct ScanPairs<'a> { distances: &'a [Vec], - output: MutMatrixView<'a, Candidate>, - limits: &'a mut Vec, } - impl TopKVisitor for ScanPairs<'_> { - fn visit(mut self, topk: TopK) { - topk.initialize(self.output.as_mut_view(), self.limits); + impl BatchVisitor for ScanPairs<'_> { + fn visit<'a, Rows>(self, mut batch: Batch<'a, Rows>) + where + Batch<'a, Rows>: BatchRanker, + { for (point, row) in self.distances.iter().enumerate().skip(1) { - topk.update_dual_topk( - ARCH, - point, - &row[..point], - self.output.as_mut_view(), - self.limits, - ); + batch.update_dual_topk(ARCH, &row[..point]); } } } @@ -664,12 +757,11 @@ mod tests { MutMatrixView::try_from(output.as_mut_slice(), point_count, capacity).unwrap(); // When: offer every non-self pair once, through the production width dispatch. - with_topk( - capacity, + with_batch( + rows.as_mut_view(), + &mut limits, ScanPairs { distances: &distances, - output: rows.as_mut_view(), - limits: &mut limits, }, ); @@ -705,25 +797,24 @@ mod tests { #[values(1, LANES - 1, LANES, 2 * LANES + 1)] target: usize, ) { let source = 2 * LANES + 2; - let topk = TopK::new(Fixed::<1>); let mut output = vec![EMPTY; source + 1]; let mut limits = Vec::new(); let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); - topk.initialize(rows.as_mut_view(), &mut limits); + let mut batch = Batch::new(rows.as_mut_view(), &mut limits); let mut distances = vec![f32::INFINITY; source]; distances[0] = -6.0; distances[target] = 4.0; // The source fills its only slot with point 0 before reaching this target. - topk.update_dual_topk(ARCH, source, &distances, rows.as_mut_view(), &mut limits); + batch.update_dual_topk(ARCH, &distances); let mut expected = vec![EMPTY; source + 1]; expected[0] = Candidate::new(source as u32, -6.0); expected[target] = Candidate::new(source as u32, 4.0); expected[source] = Candidate::new(0, -6.0); let expected_limits: Vec<_> = expected.iter().map(|c| c.distance).collect(); - assert_eq!(rows.as_slice(), expected, "target={target}"); - assert_eq!(limits, expected_limits); + assert_eq!(batch.candidates.as_mut_slice(), expected, "target={target}"); + assert_eq!(batch.thresholds, expected_limits); } #[rstest] @@ -731,33 +822,26 @@ mod tests { #[values(1, LANES - 1, LANES, 2 * LANES + 1)] target: usize, ) { let source = 2 * LANES + 2; - let topk = TopK::new(Fixed::<1>); let mut output = vec![EMPTY; source + 1]; let mut limits = Vec::new(); let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); - topk.initialize(rows.as_mut_view(), &mut limits); + let mut batch = Batch::new(rows.as_mut_view(), &mut limits); let mut previous_distances = vec![f32::INFINITY; target]; previous_distances[0] = 1.0; - topk.update_dual_topk( - ARCH, - target, - &previous_distances, - rows.as_mut_view(), - &mut limits, - ); + batch.update_dual_topk(ARCH, &previous_distances); let mut distances = vec![f32::INFINITY; source]; distances[target] = 5.0; // The target already has a closer neighbor; the source still needs one. - topk.update_dual_topk(ARCH, source, &distances, rows.as_mut_view(), &mut limits); + batch.update_dual_topk(ARCH, &distances); let mut expected = vec![EMPTY; source + 1]; expected[0] = Candidate::new(target as u32, 1.0); expected[target] = Candidate::new(0, 1.0); expected[source] = Candidate::new(target as u32, 5.0); let expected_limits: Vec<_> = expected.iter().map(|c| c.distance).collect(); - assert_eq!(rows.as_slice(), expected, "target={target}"); - assert_eq!(limits, expected_limits); + assert_eq!(batch.candidates.as_mut_slice(), expected, "target={target}"); + assert_eq!(batch.thresholds, expected_limits); } #[rstest] @@ -769,24 +853,33 @@ mod tests { #[values(LANES - 1, LANES, 2 * LANES + 1)] target: usize, ) { let source = 2 * LANES + 2; - let topk = TopK::new(Fixed::<2>); let mut output = vec![EMPTY; (source + 1) * 2]; let mut limits = Vec::new(); let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 2).unwrap(); - topk.initialize(rows.as_mut_view(), &mut limits); + let mut batch = Batch::new(rows.as_mut_view(), &mut limits); let mut distances = vec![f32::INFINITY; source]; distances[target] = distance; - topk.update_dual_topk(ARCH, source, &distances, rows.as_mut_view(), &mut limits); + batch.update_dual_topk(ARCH, &distances); let mut expected = vec![EMPTY; (source + 1) * 2]; expected[source * 2] = Candidate::new(target as u32, distance); expected[target * 2] = Candidate::new(source as u32, distance); - assert_eq!(rows.as_slice(), expected); - assert_eq!(rows.row(source)[0].distance.to_bits(), distance.to_bits()); - assert_eq!(rows.row(target)[0].distance.to_bits(), distance.to_bits()); + assert_eq!(batch.candidates.as_mut_slice(), expected); + assert_eq!( + batch.candidates.as_mut_slice()[source * 2] + .distance + .to_bits(), + distance.to_bits() + ); + assert_eq!( + batch.candidates.as_mut_slice()[target * 2] + .distance + .to_bits(), + distance.to_bits() + ); // One neighbor leaves each two-slot result open to another candidate. - assert_eq!(limits, vec![f32::INFINITY; source + 1]); + assert_eq!(batch.thresholds, vec![f32::INFINITY; source + 1]); } #[rstest] @@ -795,35 +888,46 @@ mod tests { #[case::infinite_pairs(&[f32::INFINITY; 2 * LANES + 1])] fn an_update_without_rankable_pairs_preserves_existing_results(#[case] distances: &[f32]) { let source = 2 * LANES + 1; - let topk = TopK::new(Fixed::<1>); let mut output = vec![EMPTY; source + 1]; let mut limits = Vec::new(); let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); - topk.initialize(rows.as_mut_view(), &mut limits); - topk.update_dual_topk(ARCH, 1, &[3.0], rows.as_mut_view(), &mut limits); - let previous_output = rows.as_slice().to_vec(); - let previous_limits = limits.clone(); + let mut batch = Batch::new(rows.as_mut_view(), &mut limits); + batch.update_dual_topk(ARCH, &[3.0]); + let previous_output = batch.candidates.as_mut_slice().to_vec(); + let previous_limits = batch.thresholds.to_vec(); - topk.update_dual_topk(ARCH, source, distances, rows.as_mut_view(), &mut limits); + batch.update_dual_topk(ARCH, distances); - assert_eq!(rows.as_slice(), previous_output); - assert_eq!(limits, previous_limits); + assert_eq!(batch.candidates.as_mut_slice(), previous_output); + assert_eq!(batch.thresholds, previous_limits); } - #[rstest] - #[case::fixed(Fixed::<0>)] - #[case::runtime(Runtime(0))] - fn zero_capacity_pair_scans_leave_no_neighbors_or_finite_limits(#[case] width: impl Width) { - let topk = TopK::new(width); - let mut output = []; - let mut limits = Vec::new(); - let mut rows = MutMatrixView::try_from(&mut output[..], 3, 0).unwrap(); - topk.initialize(rows.as_mut_view(), &mut limits); - - topk.update_dual_topk(ARCH, 1, &[4.0], rows.as_mut_view(), &mut limits); - topk.update_dual_topk(ARCH, 2, &[2.0, 7.0], rows.as_mut_view(), &mut limits); + #[test] + fn zero_capacity_pair_scans_leave_no_neighbors_or_finite_limits() { + fn check<'a, Rows>(mut batch: Batch<'a, Rows>) + where + Batch<'a, Rows>: BatchRanker, + { + batch.update_dual_topk(ARCH, &[4.0]); + batch.update_dual_topk(ARCH, &[2.0, 7.0]); + + for index in 0..batch.nrows() { + assert!(batch.ranker(index).unwrap().as_ref().is_empty()); + } + assert_eq!(batch.thresholds, [f32::INFINITY; 3]); + } - assert!(rows.as_slice().is_empty()); - assert_eq!(limits, [f32::INFINITY; 3]); + let mut fixed: [[Candidate; 0]; 3] = [[]; 3]; + let mut fixed_limits = Vec::new(); + fixed_limits.resize(fixed.len(), f32::INFINITY); + check(Batch { + candidates: fixed.as_mut_slice(), + thresholds: &mut fixed_limits, + }); + + let mut dynamic = []; + let mut dynamic_limits = Vec::new(); + let rows = MutMatrixView::try_from(&mut dynamic[..], 3, 0).unwrap(); + check(Batch::new(rows, &mut dynamic_limits)); } } From b0ccf73be46bbcc8b82e851aee1ef21b869e3cd7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 23 Sep 2026 14:56:25 +0000 Subject: [PATCH 205/229] refactor(pipnn): replace TopK rankers with direct selection Each kernel now calls one top-k function: - select_top_k_ids ranks partition rows; select_top_k_symmetric runs the leaf pair scan. Both take k from the output width and choose fixed-size or slice nearest sets internally. - Remove Ranker, TopKVisitor, BatchRanker, Batch, BatchVisitor, with_topk, and with_batch. - Share output-row validation and checked distance scratch between the kernels. Remove leaf_neighbor_count, LeafKernelError, and UNASSIGNED_LEADER. - Load SIMD groups through one unsafe load whose bounds come from &[f32; LANES]. A load from a copied group cost 11% of leaf ranking time on AVX2. - Leaf CosineNormalized delegates to InnerProduct and adds 1. - Top-k and kernel tests run on every supported architecture and cross SIMD groups. Point the nightly Miri step at the renamed tests. --- .github/workflows/nightly.yml | 8 +- diskann/src/graph/pipnn/leaf_kernel.rs | 591 +++----- diskann/src/graph/pipnn/leaf_metric.rs | 167 +-- diskann/src/graph/pipnn/mod.rs | 177 ++- diskann/src/graph/pipnn/partition_kernel.rs | 433 +++--- diskann/src/graph/pipnn/partition_metric.rs | 134 +- diskann/src/graph/pipnn/simd.rs | 46 +- diskann/src/graph/pipnn/topk.rs | 1392 +++++++++---------- 8 files changed, 1366 insertions(+), 1582 deletions(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index b4de125c71..36d3f3f193 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -136,10 +136,12 @@ jobs: env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance - - name: PiPNN TopK SIMD load boundaries + - name: PiPNN top-k group loads env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance + # One selection over two full groups and one pair scan over two full groups + # reach all three unsafe group loads in the top-k module. run: | cargo +nightly miri test --locked -p diskann --features pipnn --lib -- \ - graph::pipnn::topk::tests::batch_selection_matches_a_full_sort \ - graph::pipnn::topk::tests::pair_scans_match_full_row_sort_across_capacities_and_lengths + graph::pipnn::topk::tests::selection_returns_distinct_candidates_when_distances_tie \ + graph::pipnn::topk::tests::a_single_pair_updates_exactly_its_two_endpoints::case_1_finite diff --git a/diskann/src/graph/pipnn/leaf_kernel.rs b/diskann/src/graph/pipnn/leaf_kernel.rs index 4335ca1c1a..ff9bec1409 100644 --- a/diskann/src/graph/pipnn/leaf_kernel.rs +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -3,62 +3,43 @@ * Licensed under the MIT license. */ -//! Leaf-local top-k selection from packed `f32` point vectors. +//! Nearest-neighbor selection inside one PiPNN leaf. //! -//! [`select_leaf_neighbors`] asks the metric to fill a lower-triangle ranking -//! buffer. It scans each point pair once and updates both points' neighbor lists. +//! [`select_leaf_neighbors`] asks the metric for the distances between all point +//! pairs of the leaf. It then runs a pair scan, which reads each pair once and +//! offers it to the nearest sets of both points. -use crate::{ANNError, ANNResult}; +use crate::ANNResult; use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ + check_output_rows, distance_scratch, leaf_metric::LeafMetric, simd::Simd, - topk::{Batch, BatchRanker, BatchVisitor, Candidate, with_batch}, + topk::{Candidate, select_top_k_symmetric}, }; -/// Reusable storage for one leaf numerical pipeline. +/// Reusable storage for [`select_leaf_neighbors`]. Reuse one workspace across +/// leaves to avoid an allocation for each leaf. #[derive(Debug, Default)] pub(super) struct LeafKernelWorkspace { distance_scratch: Vec, - worst: Vec, + kth_distances: Vec, } -/// Invalid output shape for leaf-neighbor selection. -#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)] -pub(super) enum LeafKernelError { - /// The output must contain one row per input point. - #[error("invalid leaf output row count {rows} for {points} points")] - InvalidOutputRows { points: usize, rows: usize }, - /// A source requests more neighbors than the leaf has other points. - #[error("invalid leaf neighbor count {neighbors} for {points} points; maximum is {maximum}")] - InvalidNeighborCount { - points: usize, - neighbors: usize, - maximum: 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)`. -/// -pub(super) fn leaf_neighbor_count(points: usize, requested_k: usize) -> usize { - requested_k.min(points.saturating_sub(1)) -} - -/// Compute local nearest neighbors for one packed leaf matrix. +/// Find the k nearest other points of each point in one leaf. /// -/// `output` contains one row per point, ordered by increasing ranking distance. -/// Candidate IDs are positions in the leaf. Equal distances can select either -/// candidate. NaN and positive infinity are not retained; unfilled slots contain -/// [`Candidate::default`]. +/// `points` holds one leaf vector in each row and must not be empty. k is the +/// column count of `output`, which has one row per point. Each row receives the +/// nearest set of one point: its nearest other points, nearest first, as +/// positions in the leaf. Equal distances can select either point. NaN and +/// positive infinity never enter a row. A row can be wider than the number of +/// other points. Unfilled slots hold [`Candidate::EMPTY`]. /// /// # Errors /// -/// Returns an error for invalid linear-algebra input or output shape. -/// Invalid output shapes leave the output and workspace unchanged. +/// Returns an error if `output` does not have one row per point, or if the +/// distance matrix size overflows `usize`. pub(super) fn select_leaf_neighbors( arch: A, points: MatrixView<'_, f32>, @@ -70,290 +51,202 @@ where M: LeafMetric, { let point_count = points.nrows(); - validate_output(point_count, &output).map_err(ANNError::new)?; - let distance_count = point_count * point_count; - if workspace.distance_scratch.len() < distance_count { - workspace.distance_scratch.resize(distance_count, 0.0); - } - let mut distances = MutMatrixView::try_from( - &mut workspace.distance_scratch[..distance_count], - point_count, - point_count, - )?; - M::compute_distances(points, distances.as_mut_slice())?; - rank_leaf_distances(arch, distances.as_view(), output, &mut workspace.worst); - Ok(()) -} - -/// Offer each lower-triangle row to the two endpoint neighbor lists. -fn rank_leaf_distances( - arch: A, - distances: MatrixView<'_, f32>, - output: MutMatrixView<'_, Candidate>, - worst: &mut Vec, -) { - with_batch(output, worst, RankLeaf { arch, distances }); -} - -struct RankLeaf<'a, A> { - arch: A, - distances: MatrixView<'a, f32>, -} - -impl BatchVisitor for RankLeaf<'_, A> { - #[inline] - fn visit<'a, Rows>(self, mut batch: Batch<'a, Rows>) - where - Batch<'a, Rows>: BatchRanker, - { - for point_idx in 1..self.distances.nrows() { - batch.update_dual_topk(self.arch, &self.distances.row(point_idx)[..point_idx]); - } - } -} - -/// Check for one output row per point and a valid non-self neighbor count. -fn validate_output( - point_count: usize, - output: &MutMatrixView<'_, Candidate>, -) -> Result<(), LeafKernelError> { - if output.nrows() != point_count { - return Err(LeafKernelError::InvalidOutputRows { - points: point_count, - rows: output.nrows(), - }); - } - 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, - }); - } + check_output_rows(point_count, output.nrows())?; + let mut distances = + distance_scratch(&mut workspace.distance_scratch, point_count, point_count)?; + M::compute_distances(points, distances.as_mut_view())?; + select_top_k_symmetric( + arch, + distances.as_view(), + output, + &mut workspace.kth_distances, + ); Ok(()) } #[cfg(test)] mod tests { use super::*; - use crate::graph::pipnn::test_support; + use crate::graph::pipnn::test_support::{self, ArchCheck, for_each_arch}; use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2}; use diskann_vector::distance::Metric; use diskann_wide::ARCH; use rstest::rstest; + use std::marker::PhantomData; - const EMPTY: Candidate = Candidate::new(super::super::topk::UNASSIGNED, f32::INFINITY); - - #[rstest] - #[case::empty(0, 8, 0)] - #[case::singleton(1, 8, 0)] - #[case::zero_requested(5, 0, 0)] - #[case::below_available(5, 2, 2)] - #[case::all_available(5, 4, 4)] - #[case::above_available(5, 9, 4)] - fn neighbor_count_is_limited_to_other_points( - #[case] points: usize, - #[case] requested: usize, - #[case] expected: usize, - ) { - assert_eq!(leaf_neighbor_count(points, requested), expected); - } - + // The metric tests sweep dimensions and the TopK tests sweep widths and lengths. + // This test checks only that the kernel composes them with leaf-local IDs. #[rstest] #[case::l2(L2, Metric::L2, [[1,3,2,4], [0,2,3,4], [4,0,1,3], [0,2,1,4], [2,1,0,3]])] #[case::cosine(Cosine, Metric::Cosine, [[3,1,2,4], [0,3,4,2], [4,3,0,1], [0,2,1,4], [2,1,3,0]])] #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized, [[3,1,2,4], [0,3,4,2], [4,3,0,1], [0,2,1,4], [2,1,3,0]])] #[case::inner_product(InnerProduct, Metric::InnerProduct, [[3,1,2,4], [0,3,2,4], [4,3,0,1], [0,2,1,4], [2,1,0,3]])] - fn neighbors_are_local_non_self_ids_in_metric_order( + fn neighbors_are_other_points_of_the_leaf_in_metric_order( #[case] _metric: M, #[case] scalar_metric: Metric, #[case] expected_ids: [[u32; 4]; 5], ) { - for neighbors in [0, 1, 2, 3, 4] { - for dimensions in [2, 7, 8, 9, 15, 16, 17, 128, 129] { - // These five points have distinct scores in each row for every metric. - let coordinates = [ - [2.0, 2.0], - [-1.0, 3.0], - [0.0, -2.0], - [5.0, 0.0], - [-3.0, -4.0], - ]; - let values = test_support::packed_points( - &coordinates, - dimensions, - scalar_metric == Metric::CosineNormalized, - ); - let points = MatrixView::try_from(values.as_slice(), 5, dimensions).unwrap(); - let mut output = vec![Candidate::new(0, -100.0); 5 * neighbors]; - let mut workspace = LeafKernelWorkspace::default(); - - select_leaf_neighbors::<_, M>( - ARCH, - points, - MutMatrixView::try_from(output.as_mut_slice(), 5, neighbors).unwrap(), - &mut workspace, - ) - .unwrap_or_else(|error| { - panic!("neighbors={neighbors}, dimensions={dimensions}: {error}") - }); - - for point in 0..5 { - let actual = &output[point * neighbors..(point + 1) * neighbors]; - assert_eq!( - actual.iter().map(|c| c.local_idx).collect::>(), - expected_ids[point][..neighbors], - "neighbors={neighbors}, dimensions={dimensions}, point={point}" - ); - for candidate in actual { - let expected = test_support::distance( - scalar_metric, - points.row(point), - points.row(candidate.local_idx as usize), - ); - // Only two coordinates are nonzero, so eight f32 ulps at the score scale - // cover the dot/norm rounding without admitting another neighbor. - let tolerance = 8.0 * f64::from(f32::EPSILON) * expected.abs().max(1.0); - assert!( - (f64::from(candidate.distance) - expected).abs() <= tolerance, - "neighbors={neighbors}, dimensions={dimensions}, point={point}, candidate={candidate:?}, expected={expected}" - ); - } - } - } - } - } - - #[test] - fn neighbors_match_scalar_ranking_across_leaf_sizes_and_counts() { - for point_count in [1, 2, 16, 17, 18, 33, 34] { - for requested in [1, 2, 3, 10, 11, 17] { - // Slightly increasing gaps avoid ties between the two sides of each point. - // Center the coordinates so intermediate norm sums stay exactly representable. - let mut values: Vec<_> = - (0..point_count).map(|i| (64 * i + i * i) as f32).collect(); - let center = values[point_count - 1] / 2.0; - for value in &mut values { - *value -= center; - } - let neighbors = requested.min(point_count - 1); - let mut output = vec![EMPTY; point_count * neighbors]; - - select_leaf_neighbors::<_, L2>( - ARCH, - MatrixView::try_from(values.as_slice(), point_count, 1).unwrap(), - MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors).unwrap(), - &mut LeafKernelWorkspace::default(), - ) - .unwrap_or_else(|error| { - panic!("point_count={point_count}, requested={requested}: {error}") - }); + // These five points have distinct distances in each row for every metric. + let coordinates = [ + [2.0, 2.0], + [-1.0, 3.0], + [0.0, -2.0], + [5.0, 0.0], + [-3.0, -4.0], + ]; + let values = + test_support::packed_points(&coordinates, 2, scalar_metric == Metric::CosineNormalized); + let points = MatrixView::try_from(values.as_slice(), 5, 2).unwrap(); + // Widths above four exceed the other points in the leaf. + for neighbors in [0, 1, 2, 3, 4, 6] { + let mut output = vec![Candidate::new(0, -100.0); 5 * neighbors]; + + select_leaf_neighbors::<_, M>( + ARCH, + points, + MutMatrixView::try_from(output.as_mut_slice(), 5, neighbors).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap_or_else(|error| panic!("neighbors={neighbors}: {error}")); - for point in 0..point_count { - let mut expected: Vec<_> = values + let filled = neighbors.min(4); + for point in 0..5 { + let actual = &output[point * neighbors..(point + 1) * neighbors]; + assert_eq!( + actual[..filled] .iter() - .enumerate() - .filter(|&(other, _)| other != point) - .map(|(other, &value)| { - Candidate::new(other as u32, (values[point] - value).powi(2)) - }) - .collect(); - expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); - expected.truncate(neighbors); - assert_eq!( - &output[point * neighbors..(point + 1) * neighbors], - expected, - "point_count={point_count}, requested={requested}, point={point}, neighbors={neighbors}" + .map(|c| c.local_idx) + .collect::>(), + expected_ids[point][..filled], + "neighbors={neighbors}, point={point}" + ); + assert!( + actual[filled..].iter().all(|&c| c == Candidate::EMPTY), + "neighbors={neighbors}, point={point}: {actual:?}" + ); + for candidate in &actual[..filled] { + let expected = test_support::distance( + scalar_metric, + points.row(point), + points.row(candidate.local_idx as usize), + ); + // Only two coordinates are nonzero, so eight f32 ulps at the distance scale + // cover the dot/norm rounding without admitting another neighbor. + let tolerance = 8.0 * f64::from(f32::EPSILON) * expected.abs().max(1.0); + assert!( + (f64::from(candidate.distance) - expected).abs() <= tolerance, + "neighbors={neighbors}, point={point}, candidate={candidate:?}, expected={expected}" ); } } } } + // The composition test above uses five points, fewer than one SIMD group. This test + // runs leaves whose distance rows cross SIMD groups, with fixed-size nearest sets + // and slices, on each architecture that production can select. #[rstest] #[case::l2(L2, Metric::L2)] #[case::cosine(Cosine, Metric::Cosine)] #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] #[case::inner_product(InnerProduct, Metric::InnerProduct)] - fn large_dense_leaves_select_nearest_non_self_neighbors( + fn large_leaves_select_the_nearest_neighbors_on_every_architecture( #[case] _metric: M, - #[case] scalar_metric: Metric, + #[case] metric: Metric, ) { - for shape in [(33, 384), (65, 768), (129, 1536), (35, 1537)] { - for neighbors in [3, 11] { - let (point_count, dimensions) = shape; - let mut values = test_support::dense_points(point_count, dimensions, 1287); - if scalar_metric == Metric::CosineNormalized { - test_support::normalize(&mut values, dimensions); - } - let points = - MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); - let mut output = vec![EMPTY; point_count * neighbors]; - - select_leaf_neighbors::<_, M>( - ARCH, - points, - MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors).unwrap(), - &mut LeafKernelWorkspace::default(), - ) - .unwrap_or_else(|error| panic!("shape={shape:?}, neighbors={neighbors}: {error}")); + struct LargeLeaves { + metric: Metric, + leaf_metric: PhantomData, + } - // Unnormalized dyadic sums are exact; cosine rounds sqrt/division, while - // normalized dot products also accumulate coordinate rounding. - let tolerance = match scalar_metric { - Metric::L2 | Metric::InnerProduct => 0.0, - Metric::Cosine => 16.0 * f64::from(f32::EPSILON), - Metric::CosineNormalized => { - let roundoff = dimensions as f64 * f64::from(f32::EPSILON); - roundoff / (1.0 - roundoff) + impl ArchCheck for LargeLeaves { + fn check(&self, arch: A) { + let metric = self.metric; + let arch_name = std::any::type_name::(); + // 35 points give distance rows with two SIMD groups and a tail. 129 points + // with 1536 dimensions give many groups and a typical embedding width. + for (point_count, dimensions) in [(35, 129), (129, 1536)] { + let mut values = test_support::dense_points(point_count, dimensions, 1287); + if metric == Metric::CosineNormalized { + test_support::normalize(&mut values, dimensions); } - }; - for point in 0..point_count { - let mut expected: Vec<_> = (0..point_count) - .filter(|&other| other != point) - .map(|other| { - ( - other as u32, - test_support::distance( - scalar_metric, - points.row(point), - points.row(other), - ), - ) + let points = + MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let tolerance = test_support::dense_tolerance(metric, dimensions); + // The oracle sorts every other point by its scalar distance. + let oracle: Vec> = (0..point_count) + .map(|point| { + let mut others: Vec<_> = (0..point_count) + .filter(|&other| other != point) + .map(|other| { + let distance = test_support::distance( + metric, + points.row(point), + points.row(other), + ); + (other as u32, distance) + }) + .collect(); + others.sort_by(|left, right| left.1.total_cmp(&right.1)); + others }) .collect(); - expected.sort_by(|left, right| left.1.total_cmp(&right.1)); - let actual = &output[point * neighbors..(point + 1) * neighbors]; - for (rank, candidate) in actual.iter().enumerate() { - assert!( - !actual[..rank] - .iter() - .any(|previous| previous.local_idx == candidate.local_idx), - "shape={shape:?}, neighbors={neighbors}, duplicate neighbor for point={point}: {candidate:?}" - ); - let own_score = expected - .iter() - .find(|&&(id, _)| id == candidate.local_idx) - .unwrap_or_else(|| { - panic!("shape={shape:?}, neighbors={neighbors}, invalid or self neighbor for point={point}: {candidate:?}") - }) - .1; - assert!( - (f64::from(candidate.distance) - own_score).abs() <= tolerance, - "shape={shape:?}, neighbors={neighbors}, point={point}, candidate={candidate:?}, score={own_score}" - ); - // Equal scores may use either ID, but every rank must be nearest-first. - assert!( - (own_score - expected[rank].1).abs() <= tolerance, - "shape={shape:?}, neighbors={neighbors}, point={point}, rank={rank}, score={own_score}, expected={:?}", - expected[rank] + + // k = 3 uses fixed-size nearest sets. k = 11 uses slices. + for neighbors in [3, 11] { + let context = format!( + "{arch_name}, {metric:?}, shape=({point_count}, {dimensions}), neighbors={neighbors}" ); + let mut output = vec![Candidate::EMPTY; point_count * neighbors]; + + select_leaf_neighbors::( + arch, + points, + MutMatrixView::try_from(output.as_mut_slice(), point_count, neighbors) + .unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap_or_else(|error| panic!("{context}: {error}")); + + for (point, expected) in oracle.iter().enumerate() { + let actual = &output[point * neighbors..(point + 1) * neighbors]; + for (rank, candidate) in actual.iter().enumerate() { + assert!( + !actual[..rank] + .iter() + .any(|previous| previous.local_idx == candidate.local_idx), + "{context}, point={point}: duplicate {candidate:?}" + ); + let own_distance = expected + .iter() + .find(|&&(id, _)| id == candidate.local_idx) + .unwrap_or_else(|| { + panic!("{context}, point={point}: invalid or self {candidate:?}") + }) + .1; + assert!( + (f64::from(candidate.distance) - own_distance).abs() + <= tolerance, + "{context}, point={point}: {candidate:?} != {own_distance}" + ); + // Equal distances can select either ID, but each rank must be + // nearest-first. + assert!( + (own_distance - expected[rank].1).abs() <= tolerance, + "{context}, point={point}, rank={rank}: {own_distance} != {:?}", + expected[rank] + ); + } + } } } } } + + for_each_arch(&LargeLeaves:: { + metric, + leaf_metric: PhantomData, + }); } #[test] @@ -362,8 +255,8 @@ mod tests { let mut workspace = LeafKernelWorkspace::default(); let mut output = Vec::new(); - // Grow, shrink, change K, and finish with a singleton using the same buffers. - for (count, neighbors) in [(4, 1), (2, 1), (5, 3), (1, 0)] { + // Grow, shrink, change K, and finish with a singleton wider than its leaf. + for (count, neighbors) in [(4, 1), (2, 1), (5, 3), (1, 2)] { output.resize(count * neighbors, Candidate::new(4, -100.0)); let points = MatrixView::try_from(&values[..count], count, 1).unwrap(); select_leaf_neighbors::<_, L2>( @@ -382,7 +275,7 @@ mod tests { }) .collect(); expected.sort_by(|a, b| a.distance.total_cmp(&b.distance)); - expected.truncate(neighbors); + expected.resize(neighbors, Candidate::EMPTY); assert_eq!( &output[point * neighbors..(point + 1) * neighbors], expected, @@ -392,119 +285,51 @@ mod tests { } } - #[rstest] - #[case::wrong_rows((2, 1), LeafKernelError::InvalidOutputRows { points: 3, rows: 2 })] - #[case::too_many_neighbors((3, 3), LeafKernelError::InvalidNeighborCount { points: 3, neighbors: 3, maximum: 2 })] - fn invalid_output_shape_preserves_output_and_workspace( - #[case] shape: (usize, usize), - #[case] expected: LeafKernelError, - ) { + #[test] + fn output_without_one_row_per_point_is_rejected() { let values = [1.0, 2.0, 4.0]; - let mut output = vec![Candidate::new(2, 9.0); shape.0 * shape.1]; - let previous_output = output.clone(); - let mut workspace = LeafKernelWorkspace { - distance_scratch: vec![11.0, 13.0], - worst: vec![3.0, 7.0], - }; + let mut output = [Candidate::EMPTY; 2]; let error = select_leaf_neighbors::<_, L2>( ARCH, MatrixView::try_from(&values[..], 3, 1).unwrap(), - MutMatrixView::try_from(output.as_mut_slice(), shape.0, shape.1).unwrap(), - &mut workspace, - ) - .unwrap_err(); - - assert_eq!(error.downcast_ref::(), Some(&expected)); - assert_eq!(output, previous_output); - assert_eq!(workspace.distance_scratch, [11.0, 13.0]); - assert_eq!(workspace.worst, [3.0, 7.0]); - } - - #[test] - fn metric_failure_is_returned_without_publishing_neighbors() { - #[derive(Debug, thiserror::Error)] - #[error("distance computation failed")] - struct MetricFailure; - struct FailingMetric; - impl LeafMetric for FailingMetric { - fn compute_distances(_: MatrixView<'_, f32>, _: &mut [f32]) -> ANNResult<()> { - Err(ANNError::new(MetricFailure)) - } - } - let values = [1.0, 2.0]; - let mut output = [Candidate::new(1, 7.0), Candidate::new(0, 7.0)]; - - let error = select_leaf_neighbors::<_, FailingMetric>( - ARCH, - MatrixView::try_from(&values[..], 2, 1).unwrap(), MutMatrixView::try_from(&mut output[..], 2, 1).unwrap(), &mut LeafKernelWorkspace::default(), ) .unwrap_err(); - assert!(error.is::()); - assert_eq!(output, [Candidate::new(1, 7.0), Candidate::new(0, 7.0)]); + assert!( + error + .to_string() + .contains("invalid kernel output row count 2 for 3 points"), + "{error}" + ); } - mod rank_leaf_distances_tests { - use super::*; - - #[test] - fn ranking_reads_only_pairs_in_the_strict_lower_triangle() { - // Diagonal and upper entries are deliberately better than every real pair. - let distances = [ - -100.0, -100.0, -100.0, -100.0, 7.0, -100.0, -100.0, -100.0, 3.0, 8.0, -100.0, - -100.0, 5.0, 2.0, 6.0, -100.0, - ]; - let mut output = [Candidate::new(0, -200.0); 8]; - let mut limits = vec![-200.0; 4]; - - rank_leaf_distances( - ARCH, - MatrixView::try_from(&distances[..], 4, 4).unwrap(), - MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), - &mut limits, - ); - - let expected = [ - [Candidate::new(2, 3.0), Candidate::new(3, 5.0)], - [Candidate::new(3, 2.0), Candidate::new(0, 7.0)], - [Candidate::new(0, 3.0), Candidate::new(3, 6.0)], - [Candidate::new(1, 2.0), Candidate::new(0, 5.0)], - ]; - assert_eq!(output, expected.as_flattened()); - assert_eq!(limits, [5.0, 7.0, 6.0, 5.0]); - } - - #[rstest] - #[case::nan(f32::NAN)] - #[case::infinity(f32::INFINITY)] - fn unrankable_pairs_leave_unassigned_slots(#[case] unrankable: f32) { - let mut distances = [unrankable; 9]; - distances[6] = 4.0; - let mut output = [Candidate::new(1, -100.0); 6]; - let mut limits = vec![-100.0; 5]; - - rank_leaf_distances( - ARCH, - MatrixView::try_from(&distances[..], 3, 3).unwrap(), - MutMatrixView::try_from(&mut output[..], 3, 2).unwrap(), - &mut limits, - ); - - assert_eq!( - output, - [ - Candidate::new(2, 4.0), - EMPTY, - EMPTY, - EMPTY, - Candidate::new(0, 4.0), - EMPTY - ] - ); - assert_eq!(limits, [f32::INFINITY; 3]); - } + #[test] + fn ranking_reads_only_pairs_in_the_strict_lower_triangle() { + // Diagonal and upper entries are deliberately better than every real pair. + let distances = [ + -100.0, -100.0, -100.0, -100.0, 7.0, -100.0, -100.0, -100.0, 3.0, 8.0, -100.0, -100.0, + 5.0, 2.0, 6.0, -100.0, + ]; + let mut output = [Candidate::new(0, -200.0); 8]; + let mut kth_distances = vec![-200.0; 4]; + + select_top_k_symmetric( + ARCH, + MatrixView::try_from(&distances[..], 4, 4).unwrap(), + MutMatrixView::try_from(&mut output[..], 4, 2).unwrap(), + &mut kth_distances, + ); + + let expected = [ + [Candidate::new(2, 3.0), Candidate::new(3, 5.0)], + [Candidate::new(3, 2.0), Candidate::new(0, 7.0)], + [Candidate::new(0, 3.0), Candidate::new(3, 6.0)], + [Candidate::new(1, 2.0), Candidate::new(0, 5.0)], + ]; + assert_eq!(output, expected.as_flattened()); + assert_eq!(kth_distances, [5.0, 7.0, 6.0, 5.0]); } } diff --git a/diskann/src/graph/pipnn/leaf_metric.rs b/diskann/src/graph/pipnn/leaf_metric.rs index 58080e09f4..6d1fce29fc 100644 --- a/diskann/src/graph/pipnn/leaf_metric.rs +++ b/diskann/src/graph/pipnn/leaf_metric.rs @@ -3,13 +3,14 @@ * Licensed under the MIT license. */ -//! Build ranking distances for one PiPNN leaf. +//! Pairwise distances inside one PiPNN leaf. //! -//! A ranking distance preserves nearest-first order. It does not need to equal -//! the mathematical metric distance. The leaf kernel reads only the lower triangle. +//! Leaf distances equal the DiskANN metric distances, because later graph stages +//! compare and quantize them. The leaf kernel reads only the strict lower +//! triangle, which holds each pair below the diagonal. use crate::{ANNError, ANNResult}; -use diskann_utils::views::MatrixView; +use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::{ Norm, norm::{FastL2Norm, FastL2NormSquared}, @@ -17,62 +18,63 @@ use diskann_vector::{ use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; -/// Fill one flattened lower-triangular ranking buffer. +/// Compute the distances between all point pairs of one leaf. /// -/// An implementation initializes the diagonal and lower triangle. The upper -/// triangle stays unspecified. The input matrix has one point in each row. -/// A zero distance can have either sign. Equal distances can select either candidate. +/// `points` has one point in each row. An implementation writes the diagonal and +/// the lower triangle of `storage`. The upper triangle can hold any value. A zero +/// distance can have either sign. pub(super) trait LeafMetric: Send + Sync + 'static { - /// Compute ranking distances for all unordered point pairs. + /// Compute the metric distance for all unordered point pairs. /// - /// Values preserve nearest-first order, but need not equal metric distances. - /// L2 returns squared distances. Normalized cosine and inner product return - /// `-dot`, omitting the constant in normalized cosine's `1 - dot`. - /// Cosine returns `1 - similarity`. + /// L2 returns squared distances. Cosine and normalized cosine return + /// `1 - similarity`. Inner product returns `-dot`. /// - /// `storage` has `points.nrows() * points.nrows()` elements. - fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()>; + /// `storage` has one row and one column per point. + fn compute_distances( + points: MatrixView<'_, f32>, + storage: MutMatrixView<'_, f32>, + ) -> ANNResult<()>; } impl LeafMetric for L2 { - fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { - let point_count = points.nrows(); + fn compute_distances( + points: MatrixView<'_, f32>, + mut storage: MutMatrixView<'_, f32>, + ) -> ANNResult<()> { // The expanded L2 formula is `||x||² + ||y||² - 2(x·y)`. let squared_norms: Vec = points .row_iter() .map(|point| FastL2NormSquared.evaluate(point)) .collect(); - // Initialize the first two terms before GEMM adds the dot-product term. - for source in 0..point_count { - let row = &mut storage[source * point_count..source * point_count + source + 1]; + // Initialize the norm terms before GEMM adds the dot-product term. + for (source, row) in storage.row_iter_mut().enumerate() { let source_norm = squared_norms[source]; - let mut target = 0; - while target < row.len() { - row[target] = source_norm + squared_norms[target]; - target += 1; + for (distance, &target_norm) in row[..=source].iter_mut().zip(&squared_norms) { + *distance = source_norm + target_norm; } } diskann_linalg::sgemm_aat_lower_add( - point_count, + points.nrows(), points.ncols(), -2.0, points.as_slice(), - storage, + storage.as_mut_slice(), ) - .map_err(ANNError::new)?; - Ok(()) + .map_err(ANNError::new) } } impl LeafMetric for Cosine { - fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { - let point_count = points.nrows(); + fn compute_distances( + points: MatrixView<'_, f32>, + mut storage: MutMatrixView<'_, f32>, + ) -> ANNResult<()> { diskann_linalg::sgemm_aat_lower( - point_count, + points.nrows(), points.ncols(), 1.0, points.as_slice(), - storage, + storage.as_mut_slice(), ) .map_err(ANNError::new)?; let norms: Vec = points @@ -80,10 +82,10 @@ impl LeafMetric for Cosine { .map(|point| FastL2Norm.evaluate(point)) .collect(); // Convert each lower-triangle dot to the bounded cosine distance. - for source in 0..point_count { - for target in 0..=source { - let index = source * point_count + target; - storage[index] = cosine_distance(storage[index], norms[source], norms[target]); + for (source, row) in storage.row_iter_mut().enumerate() { + let source_norm = norms[source]; + for (distance, &target_norm) in row[..=source].iter_mut().zip(&norms) { + *distance = cosine_distance(*distance, source_norm, target_norm); } } Ok(()) @@ -91,23 +93,36 @@ impl LeafMetric for Cosine { } impl LeafMetric for InnerProduct { - fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { + fn compute_distances( + points: MatrixView<'_, f32>, + mut storage: MutMatrixView<'_, f32>, + ) -> ANNResult<()> { diskann_linalg::sgemm_aat_lower( points.nrows(), points.ncols(), -1.0, points.as_slice(), - storage, + storage.as_mut_slice(), ) - .map_err(ANNError::new)?; - Ok(()) + .map_err(ANNError::new) } } impl LeafMetric for CosineNormalized { - fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { - // The constant in `1 - dot` does not change nearest-first order. - InnerProduct::compute_distances(points, storage) + fn compute_distances( + points: MatrixView<'_, f32>, + mut storage: MutMatrixView<'_, f32>, + ) -> ANNResult<()> { + InnerProduct::compute_distances(points, storage.as_mut_view())?; + // Keep the constant of `1 - dot`. Near neighbors then have distances near + // zero, where floating-point spacing is finest. Later stages quantize these + // distances, so the constant keeps more near neighbors distinguishable. + for (source, row) in storage.row_iter_mut().enumerate() { + row[..=source] + .iter_mut() + .for_each(|distance| *distance += 1.0); + } + Ok(()) } } @@ -131,7 +146,11 @@ mod tests { let mut output = [f32::NAN; 4]; let expected = 16_777_344.0; // 4096^2 + 128. - L2::compute_distances(points, &mut output).unwrap(); + L2::compute_distances( + points, + MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), + ) + .unwrap(); // Allow eight f32 ULPs at this scale, far less than the lost 128. let tolerance = 16.0; @@ -168,7 +187,12 @@ mod tests { MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); let mut output = vec![f32::NAN; point_count * point_count]; - M::compute_distances(points, &mut output).unwrap_or_else(|error| { + M::compute_distances( + points, + MutMatrixView::try_from(output.as_mut_slice(), point_count, point_count) + .unwrap(), + ) + .unwrap_or_else(|error| { panic!("point_count={point_count}, dimensions={dimensions}: {error}") }); @@ -223,8 +247,11 @@ mod tests { let points = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); let mut output = vec![f32::NAN; point_count * point_count]; - M::compute_distances(points, &mut output) - .unwrap_or_else(|error| panic!("shape={shape:?}: {error}")); + M::compute_distances( + points, + MutMatrixView::try_from(output.as_mut_slice(), point_count, point_count).unwrap(), + ) + .unwrap_or_else(|error| panic!("shape={shape:?}: {error}")); let tolerance = match scalar_metric { Metric::L2 | Metric::InnerProduct => 0.0, @@ -259,8 +286,8 @@ mod tests { #[case::squared_l2(L2, &[2.0, 0.0, 0.0, 3.0, -4.0, 0.0], [0.0, 13.0, 0.0, 36.0, 25.0, 0.0])] #[case::negative_dot(InnerProduct, &[2.0, 0.0, 0.0, 3.0, -4.0, 0.0], [-4.0, 0.0, -9.0, 8.0, 0.0, -16.0])] #[case::cosine(Cosine, &[2.0, 0.0, 0.0, 3.0, -4.0, 0.0], [0.0, 1.0, 0.0, 2.0, 1.0, 0.0])] - #[case::normalized_cosine(CosineNormalized, &[1.0, 0.0, 0.0, 1.0, -1.0, 0.0], [-1.0, 0.0, -1.0, 1.0, 0.0, -1.0])] - fn ranking_values_follow_the_metric_definition( + #[case::normalized_cosine(CosineNormalized, &[1.0, 0.0, 0.0, 1.0, -1.0, 0.0], [0.0, 1.0, 0.0, 2.0, 1.0, 0.0])] + fn distances_follow_the_metric_definition( #[case] _metric: M, #[case] values: &[f32], #[case] expected: [f32; 6], @@ -268,7 +295,11 @@ mod tests { let points = MatrixView::try_from(values, 3, 2).unwrap(); let mut output = [42.0; 9]; - M::compute_distances(points, &mut output).unwrap(); + M::compute_distances( + points, + MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), + ) + .unwrap(); assert_eq!( [ @@ -286,7 +317,11 @@ mod tests { let points = MatrixView::try_from(&values[..], 2, 2).unwrap(); let mut output = [42.0; 4]; - Cosine::compute_distances(points, &mut output).unwrap(); + Cosine::compute_distances( + points, + MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), + ) + .unwrap(); assert_eq!([output[0], output[2], output[3]], [1.0, 1.0, 0.0]); } @@ -294,7 +329,7 @@ mod tests { #[rstest] #[case::l2(L2, 2.0)] #[case::cosine(Cosine, 1.0)] - #[case::normalized_cosine(CosineNormalized, 0.0)] + #[case::normalized_cosine(CosineNormalized, 1.0)] #[case::inner_product(InnerProduct, 0.0)] fn a_nan_point_does_not_change_other_pair_distances( #[case] _metric: M, @@ -305,37 +340,11 @@ mod tests { M::compute_distances( MatrixView::try_from(&values[..], 3, 2).unwrap(), - &mut output, + MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), ) .unwrap(); assert_eq!(output[3], expected_finite_pair); assert!(output[6..=8].iter().all(|distance| distance.is_nan())); } - - #[rstest] - #[case::l2(L2)] - #[case::cosine(Cosine)] - #[case::normalized_cosine(CosineNormalized)] - #[case::inner_product(InnerProduct)] - fn storage_length_mismatch_reports_the_output_matrix_error(#[case] _metric: M) { - let values = [1.0, 0.0, 0.0, 2.0]; - let mut output = [42.0; 5]; - - let error = M::compute_distances( - MatrixView::try_from(&values[..], 2, 2).unwrap(), - &mut output, - ) - .unwrap_err(); - - assert_eq!( - error.downcast_ref::(), - Some(&diskann_linalg::SgemmError::InvalidMatrixDimensions { - matrix_name: diskann_linalg::MatrixName::C, - expected_rows: 2, - expected_cols: 2, - actual_len: 5, - }) - ); - } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 95bfbd5ccc..f03932c611 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,8 +5,18 @@ //! Numerical kernels for PiPNN graph construction. //! -//! Metric modules fill portable distance buffers. Kernel modules use architecture -//! `A` to traverse those buffers. Callers select metric `M` once per graph build. +//! Each kernel computes a distance matrix with metric `M`, then selects the nearest +//! candidates from it with SIMD architecture `A`: +//! +//! - `partition_kernel` assigns each point to its nearest partition leaders. It +//! uses `partition_metric`. +//! - `leaf_kernel` finds the nearest neighbors of each point in a leaf. It uses +//! `leaf_metric`. +//! - `topk` holds the selection code of both kernels. `simd` sets its vector width. +//! +//! The metric code does not use `A`, because the GEMM and norm routines select +//! their own SIMD code. Callers choose the metric once per graph build, so each +//! kernel compiles for one metric and has no run-time metric check. #![expect( dead_code, @@ -20,16 +30,24 @@ mod partition_metric; mod simd; mod topk; +use crate::{ANNError, ANNResult}; +use diskann_utils::views::MutMatrixView; + +/// Squared Euclidean distance. pub(super) struct L2; +/// Cosine distance, `1 - cos(x, y)`, for vectors of any norm. pub(super) struct Cosine; +/// Cosine distance for unit vectors, computed from the dot product only. pub(super) struct CosineNormalized; +/// Negated inner product: a larger dot product is nearer. pub(super) struct InnerProduct; -/// Convert one dot product and two norms to cosine distance. +/// Convert one dot product and two vector norms to cosine distance. /// -/// Treat a zero or subnormal norm as zero similarity. This rule takes precedence -/// over the dot value. Clamp finite similarity to the cosine range. Otherwise, -/// a NaN input produces a NaN distance. +/// A norm below `√f32::MIN_POSITIVE` counts as zero and gives distance 1 for any +/// dot value, NaN included. Below this cutoff, the product of two norms can +/// underflow. The similarity is clamped to `[-1, 1]` to remove rounding error. +/// When neither norm is below the cutoff, a NaN dot or norm gives a NaN distance. #[inline(always)] fn cosine_distance(dot: f32, source_norm: f32, target_norm: f32) -> f32 { if source_norm < f32::MIN_POSITIVE.sqrt() || target_norm < f32::MIN_POSITIVE.sqrt() { @@ -39,6 +57,40 @@ fn cosine_distance(dot: f32, source_norm: f32, target_norm: f32) -> f32 { } } +/// Return an error if a kernel output does not have one row per input point. +/// +/// The top-k functions panic on this mismatch. The kernel entry points call this +/// check first, so a bad output shape is an error instead of a panic. +fn check_output_rows(points: usize, rows: usize) -> ANNResult<()> { + if rows == points { + Ok(()) + } else { + Err(ANNError::message(format!( + "invalid kernel output row count {rows} for {points} points" + ))) + } +} + +/// Borrow a `rows x columns` prefix of reusable distance storage. +/// +/// The storage grows to the largest shape that it serves and never shrinks. +/// A shape whose element count overflows `usize` is an error, not a wrapped size. +fn distance_scratch( + storage: &mut Vec, + rows: usize, + columns: usize, +) -> ANNResult> { + let len = rows.checked_mul(columns).ok_or_else(|| { + ANNError::message(format!( + "distance matrix size overflows for {rows} x {columns}" + )) + })?; + if storage.len() < len { + storage.resize(len, 0.0); + } + Ok(MutMatrixView::try_from(&mut storage[..len], rows, columns)?) +} + #[cfg(test)] mod cosine_distance_tests { use super::cosine_distance; @@ -50,7 +102,7 @@ mod cosine_distance_tests { #[case::orthogonal(0.0, 2.0, 3.0, 1.0)] #[case::positive_similarity(3.0, 2.0, 3.0, 0.5)] #[case::negative_similarity(-3.0, 2.0, 3.0, 1.5)] - fn distance_uses_both_vector_norms( + fn distance_is_one_minus_the_dot_divided_by_both_norms( #[case] dot: f32, #[case] source_norm: f32, #[case] target_norm: f32, @@ -72,18 +124,17 @@ mod cosine_distance_tests { #[rstest] #[case::zero(0.0)] #[case::below_cutoff(f32::from_bits(f32::MIN_POSITIVE.sqrt().to_bits() - 1))] - fn a_small_norm_takes_precedence_over_the_dot_product( - #[case] small_norm: f32, - #[values(0.75, f32::NAN)] dot: f32, - #[values(false, true)] small_source: bool, - ) { - let (source_norm, target_norm) = if small_source { - (small_norm, 1.0) - } else { - (1.0, small_norm) - }; - - assert_eq!(cosine_distance(dot, source_norm, target_norm), 1.0); + fn a_norm_below_the_cutoff_gives_distance_one_for_any_dot(#[case] small_norm: f32) { + // A NaN dot must not change the result, and either norm can be the small one. + for dot in [0.75, f32::NAN] { + for (source_norm, target_norm) in [(small_norm, 1.0), (1.0, small_norm)] { + assert_eq!( + cosine_distance(dot, source_norm, target_norm), + 1.0, + "dot={dot}, norms=({source_norm:e}, {target_norm:e})" + ); + } + } } #[rstest] @@ -104,7 +155,7 @@ mod cosine_distance_tests { #[case::dot(f32::NAN, 2.0, 3.0)] #[case::source_norm(1.0, f32::NAN, 3.0)] #[case::target_norm(1.0, 2.0, f32::NAN)] - fn nan_propagates_when_neither_norm_is_small( + fn nan_propagates_when_neither_norm_is_below_the_cutoff( #[case] dot: f32, #[case] source_norm: f32, #[case] target_norm: f32, @@ -115,8 +166,55 @@ mod cosine_distance_tests { #[cfg(test)] mod test_support { + use super::simd::Simd; use diskann_vector::distance::Metric; + /// A test body that runs once for each architecture. + pub(super) trait ArchCheck { + fn check(&self, arch: A); + } + + /// Run `test` with `Scalar` and with each SIMD architecture that this CPU supports. + /// + /// The emulated `Scalar` lanes always run. `V3` runs on AVX2 hardware, `V4` runs + /// on AVX-512 hardware or under Miri, and `Neon` runs on aarch64. Production + /// selects one of these architectures at run time, so each one needs its own run. + pub(super) fn for_each_arch(test: &impl ArchCheck) { + test.check(diskann_wide::arch::Scalar); + #[cfg(target_arch = "x86_64")] + { + use diskann_wide::arch::x86_64::{V3, V4}; + if let Some(arch) = V3::new_checked() { + test.check(arch); + } + if let Some(arch) = V4::new_checked_miri() { + test.check(arch); + } + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + test.check(arch); + } + } + + /// Return the largest error between a kernel distance and [`distance`] for + /// [`dense_points`] inputs. + /// + /// L2 and inner product are exact for these inputs. Cosine rounds only in square + /// roots and division. Normalized cosine also rounds each normalized coordinate, + /// so its bound grows with the dimension. + pub(super) fn dense_tolerance(metric: Metric, dimensions: usize) -> f64 { + match metric { + Metric::L2 | Metric::InnerProduct => 0.0, + Metric::Cosine => 16.0 * f64::from(f32::EPSILON), + Metric::CosineNormalized => { + // gamma = n*u / (1 - n*u) bounds product and reduction rounding. + let roundoff = dimensions as f64 * f64::from(f32::EPSILON); + roundoff / (1.0 - roundoff) + } + } + } + pub(super) fn dense_points(rows: usize, dimensions: usize, seed: u64) -> Vec { use rand::{Rng, SeedableRng, rngs::StdRng}; @@ -133,33 +231,27 @@ mod test_support { values } - // These scalar definitions use the actual f32 inputs, with f64 arithmetic. - // Callers supply finite vectors with nonzero norms. L2 includes the point norm. + // These scalar definitions match the DiskANN metric distances. They use the + // actual f32 inputs with f64 arithmetic. Callers supply finite vectors with + // nonzero norms. L2 includes the point norm. pub(super) fn distance(metric: Metric, point: &[f32], target: &[f32]) -> f64 { + let dot = |x: &[f32], y: &[f32]| { + x.iter() + .zip(y) + .map(|(&x, &y)| f64::from(x) * f64::from(y)) + .sum::() + }; match metric { Metric::L2 => point .iter() .zip(target) .map(|(&x, &y)| (f64::from(x) - f64::from(y)).powi(2)) .sum(), - Metric::InnerProduct | Metric::CosineNormalized => -point - .iter() - .zip(target) - .map(|(&x, &y)| f64::from(x) * f64::from(y)) - .sum::(), + Metric::InnerProduct => -dot(point, target), + Metric::CosineNormalized => 1.0 - dot(point, target), Metric::Cosine => { - let dot: f64 = point - .iter() - .zip(target) - .map(|(&x, &y)| f64::from(x) * f64::from(y)) - .sum(); - let norm = |row: &[f32]| { - row.iter() - .map(|&x| f64::from(x).powi(2)) - .sum::() - .sqrt() - }; - 1.0 - (dot / (norm(point) * norm(target))).clamp(-1.0, 1.0) + let norm = |row: &[f32]| dot(row, row).sqrt(); + 1.0 - (dot(point, target) / (norm(point) * norm(target))).clamp(-1.0, 1.0) } } } @@ -195,7 +287,7 @@ mod test_support { } #[test] - fn dense_fixtures_are_replayable_and_exercise_the_last_dimension() { + fn dense_fixtures_are_deterministic_and_nonzero_in_the_last_dimension() { let values = dense_points(3, 17, 1287); assert_eq!(values, dense_points(3, 17, 1287)); @@ -208,6 +300,7 @@ mod test_support { #[rstest::rstest] #[case::squared_l2(Metric::L2, 18.0)] #[case::negative_dot(Metric::InnerProduct, -2.0)] + #[case::one_minus_dot(Metric::CosineNormalized, -1.0)] #[case::cosine(Metric::Cosine, 0.7830695421813438)] fn scalar_reference_matches_hand_calculated_distances( #[case] metric: Metric, @@ -218,12 +311,12 @@ mod test_support { } #[test] - fn normalized_points_keep_their_direction_and_last_coordinate() { + fn normalized_packed_points_put_the_second_coordinate_last() { let actual = packed_points(&[[3.0, 4.0], [0.0, -2.0]], 3, true); assert_eq!(actual, [0.6, 0.0, 0.8, 0.0, 0.0, -1.0]); assert_eq!( distance(Metric::CosineNormalized, &actual[..3], &actual[3..]), - f64::from(0.8_f32) + 1.0 + f64::from(0.8_f32) ); } } diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs index b1a39becf9..0af3df6c28 100644 --- a/diskann/src/graph/pipnn/partition_kernel.rs +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -3,44 +3,47 @@ * Licensed under the MIT license. */ -//! Select partition centers for PiPNN point assignment. +//! Partition assignment for PiPNN. //! -//! [`assign_leaders`] computes point-to-leader ranking values and selects the -//! nearest leader-column IDs for partition scatter. Each sampled leader -//! represents one child partition. +//! [`assign_leaders`] computes the ranking distance from each point of a stripe to +//! each partition leader, then selects the IDs of the nearest leaders. A ranking +//! distance orders the leaders of a point in the same way as the metric distance. +//! Each leader stands for one child partition, so the IDs tell the caller where to +//! send each point. -use crate::{ANNError, ANNResult}; +use crate::ANNResult; use diskann_utils::views::{MatrixView, MutMatrixView}; use super::{ + check_output_rows, distance_scratch, partition_metric::PartitionMetric, simd::Simd, - topk::{Candidate, Ranker, TopKVisitor, UNASSIGNED, with_topk}, + topk::{Candidate, select_top_k_ids}, }; -/// No sampled partition center was rankable for this output slot. -pub(super) const UNASSIGNED_LEADER: u32 = UNASSIGNED; - -/// Reusable storage for one point-stripe numerical pipeline. -#[derive(Default)] +/// Reusable storage for [`assign_leaders`]. Reuse one workspace across stripes to +/// avoid an allocation for each stripe. +#[derive(Debug, Default)] pub(super) struct PartitionKernelWorkspace { distance_scratch: Vec, - ranked_leaders: Vec, + nearest_leaders: Vec, } -/// Assign one packed point stripe to metric-owned partition leaders. +/// Assign each point of one stripe to its k nearest partition leaders. /// -/// Each output row contains leader-column IDs ordered by increasing ranking -/// value. Equal values can select either leader. NaN and positive infinity are -/// not retained. +/// `points` holds one vector in each row. `leaders` must not be empty. k is the +/// column count of `output`, which has one row per point. Each row receives +/// leader IDs, nearest first. A leader ID is the row of the leader in the matrix +/// that created `leaders`. Equal ranking distances can select either leader. /// -/// A point can have fewer assignments than the output width. Each remaining -/// slot contains [`UNASSIGNED_LEADER`]. +/// NaN and positive infinity never enter a row, so a point can have fewer than k +/// leaders. Each remaining slot holds [`UNASSIGNED`](super::topk::UNASSIGNED). /// /// # Errors /// -/// Returns an error for invalid GEMM input or an output row count different from -/// the point count. A row-count mismatch leaves output and workspace unchanged. +/// Returns an error if `output` does not have one row per point, if the +/// distance matrix size overflows `usize`, or if points and leaders have +/// different dimensions. pub(super) fn assign_leaders( arch: A, points: MatrixView<'_, f32>, @@ -53,77 +56,35 @@ where M: PartitionMetric, { let point_count = points.nrows(); - if output.nrows() != point_count { - return Err(ANNError::message(format!( - "invalid partition output row count {} for {point_count} points", - output.nrows() - ))); - } - let leader_count = M::leader_count(leaders); - let distance_count = point_count * leader_count; - if workspace.distance_scratch.len() < distance_count { - workspace.distance_scratch.resize(distance_count, 0.0); - } - let mut distances = MutMatrixView::try_from( - &mut workspace.distance_scratch[..distance_count], + check_output_rows(point_count, output.nrows())?; + let mut distances = distance_scratch( + &mut workspace.distance_scratch, point_count, - leader_count, + M::leader_count(leaders), )?; M::compute_distances(points, leaders, distances.as_mut_view())?; - rank_leader_distances( + select_top_k_ids( arch, distances.as_view(), output, - &mut workspace.ranked_leaders, + &mut workspace.nearest_leaders, ); Ok(()) } -/// Select leader columns from each point's distance row. -fn rank_leader_distances( - arch: A, - distances: MatrixView<'_, f32>, - output: MutMatrixView<'_, u32>, - candidates: &mut Vec, -) { - candidates.resize(output.ncols(), Candidate::default()); - with_topk( - candidates.as_mut_slice(), - RankLeaders { - arch, - distances, - output, - }, - ); -} - -struct RankLeaders<'a, A> { - arch: A, - distances: MatrixView<'a, f32>, - output: MutMatrixView<'a, u32>, -} - -impl TopKVisitor for RankLeaders<'_, A> { - #[inline] - fn visit(mut self, mut ranker: R) { - for (distances, output) in self.distances.row_iter().zip(self.output.row_iter_mut()) { - ranker.select_topk(self.arch, distances); - for (destination, candidate) in output.iter_mut().zip(ranker.as_ref().iter()) { - *destination = candidate.local_idx; - } - } - } -} - #[cfg(test)] mod tests { use super::*; - use crate::graph::pipnn::test_support; + use crate::graph::pipnn::test_support::{self, ArchCheck, for_each_arch}; + use crate::graph::pipnn::topk::UNASSIGNED; use crate::graph::pipnn::{Cosine, CosineNormalized, InnerProduct, L2}; use diskann_vector::distance::Metric; use diskann_wide::ARCH; use rstest::rstest; + use std::marker::PhantomData; + // The metric tests sweep dimensions and the TopK tests sweep widths and lengths. + // This test checks only that the kernel composes them with leader-column IDs. #[rstest] #[case::l2(L2, false, [[0,1,3,2], [1,2,3,0], [3,0,2,1]])] #[case::cosine(Cosine, false, [[0,1,3,2], [1,2,0,3], [3,0,2,1]])] @@ -134,179 +95,156 @@ mod tests { #[case] unit_norm: bool, #[case] expected_ids: [[u32; 4]; 3], ) { + let point_values = + test_support::packed_points(&[[3.0, 2.0], [-1.0, 3.0], [2.0, -4.0]], 2, unit_norm); + let leader_values = test_support::packed_points( + &[[4.0, 0.0], [0.0, 3.0], [-2.0, 0.0], [0.0, -1.0]], + 2, + unit_norm, + ); + let leaders = + M::create_leaders(MatrixView::try_from(leader_values.as_slice(), 4, 2).unwrap()); + // A width of six exceeds the four leaders. for assignments in [1, 2, 3, 4, 6] { - for dimensions in [2, 7, 8, 9, 15, 16, 17, 128, 129] { - let point_values = test_support::packed_points( - &[[3.0, 2.0], [-1.0, 3.0], [2.0, -4.0]], - dimensions, - unit_norm, - ); - let leader_values = test_support::packed_points( - &[[4.0, 0.0], [0.0, 3.0], [-2.0, 0.0], [0.0, -1.0]], - dimensions, - unit_norm, - ); - let leaders = M::create_leaders( - MatrixView::try_from(leader_values.as_slice(), 4, dimensions).unwrap(), - ); - let mut output = vec![0; 3 * assignments]; - - assign_leaders::<_, M>( - ARCH, - MatrixView::try_from(point_values.as_slice(), 3, dimensions).unwrap(), - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), - &mut PartitionKernelWorkspace::default(), - ) - .unwrap_or_else(|error| { - panic!("assignments={assignments}, dimensions={dimensions}: {error}") - }); + let mut output = vec![0; 3 * assignments]; - for point in 0..3 { - let mut expected = expected_ids[point].to_vec(); - expected.truncate(assignments); - expected.resize(assignments, UNASSIGNED_LEADER); - assert_eq!( - &output[point * assignments..(point + 1) * assignments], - expected, - "assignments={assignments}, dimensions={dimensions}, point={point}" - ); - } - } - } - } + assign_leaders::<_, M>( + ARCH, + MatrixView::try_from(point_values.as_slice(), 3, 2).unwrap(), + &leaders, + MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap_or_else(|error| panic!("assignments={assignments}: {error}")); - #[test] - fn assignments_match_nearest_leaders_across_counts_and_widths() { - for leader_count in [1, 5, 16, 17, 33] { - for assignments in [1, 3, 10, 11, 17] { - // Queries avoid midpoints, so each leader has a distinct distance. - let leader_values: Vec<_> = (0..leader_count).map(|i| 2.0 * i as f32).collect(); - let point_values = [-1.0, 1.5, 2.0 * (leader_count - 1) as f32 + 0.25]; - let leaders = L2::create_leaders( - MatrixView::try_from(leader_values.as_slice(), leader_count, 1).unwrap(), + for point in 0..3 { + let mut expected = expected_ids[point].to_vec(); + expected.truncate(assignments); + expected.resize(assignments, UNASSIGNED); + assert_eq!( + &output[point * assignments..(point + 1) * assignments], + expected, + "assignments={assignments}, point={point}" ); - let mut output = vec![0; 3 * assignments]; - - assign_leaders::<_, L2>( - ARCH, - MatrixView::try_from(&point_values[..], 3, 1).unwrap(), - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), 3, assignments).unwrap(), - &mut PartitionKernelWorkspace::default(), - ) - .unwrap_or_else(|error| { - panic!("leader_count={leader_count}, assignments={assignments}: {error}") - }); - - for (point, &coordinate) in point_values.iter().enumerate() { - let mut ordered: Vec<_> = leader_values - .iter() - .enumerate() - .map(|(id, &value)| (id as u32, (coordinate - value).powi(2))) - .collect(); - ordered.sort_by(|left, right| left.1.total_cmp(&right.1)); - let mut expected: Vec<_> = ordered - .iter() - .take(assignments) - .map(|&(id, _)| id) - .collect(); - expected.resize(assignments, UNASSIGNED_LEADER); - assert_eq!( - &output[point * assignments..(point + 1) * assignments], - expected, - "point={point}, leader_count={leader_count}, assignments={assignments}" - ); - } } } } + // The composition test above uses four leaders, fewer than one SIMD group. This test + // runs stripes whose leader rows cross SIMD groups, with a fixed-size nearest set + // and a slice, on each architecture that production can select. #[rstest] #[case::l2(L2, Metric::L2)] #[case::cosine(Cosine, Metric::Cosine)] #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] #[case::inner_product(InnerProduct, Metric::InnerProduct)] - fn large_dense_stripes_select_nearest_leaders( + fn large_stripes_select_the_nearest_leaders_on_every_architecture( #[case] _metric: M, - #[case] scalar_metric: Metric, + #[case] metric: Metric, ) { - for shape in [(9, 33, 384), (17, 65, 768), (33, 129, 1536), (5, 35, 1537)] { - for assignments in [3, 11] { - let (point_count, leader_count, dimensions) = shape; - let mut point_values = test_support::dense_points(point_count, dimensions, 1287); - let mut leader_values = test_support::dense_points(leader_count, dimensions, 2026); - if scalar_metric == Metric::CosineNormalized { - test_support::normalize(&mut point_values, dimensions); - test_support::normalize(&mut leader_values, dimensions); - } - let points = - MatrixView::try_from(point_values.as_slice(), point_count, dimensions).unwrap(); - let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions) - .unwrap(); - let leaders = M::create_leaders(leader_matrix); - let mut output = vec![UNASSIGNED_LEADER; point_count * assignments]; - - assign_leaders::<_, M>( - ARCH, - points, - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), point_count, assignments) - .unwrap(), - &mut PartitionKernelWorkspace::default(), - ) - .unwrap_or_else(|error| { - panic!("shape={shape:?}, assignments={assignments}: {error}") - }); + struct LargeStripes { + metric: Metric, + partition_metric: PhantomData, + } - let tolerance = match scalar_metric { - Metric::L2 | Metric::InnerProduct => 0.0, - // Dyadic inputs make the dot/norm sums exact before cosine conversion. - Metric::Cosine => 16.0 * f64::from(f32::EPSILON), - Metric::CosineNormalized => { - let roundoff = dimensions as f64 * f64::from(f32::EPSILON); - roundoff / (1.0 - roundoff) + impl ArchCheck for LargeStripes { + fn check(&self, arch: A) { + let metric = self.metric; + let arch_name = std::any::type_name::(); + // 35 leaders give distance rows with two SIMD groups and a tail. 129 leaders + // with 1536 dimensions give many groups and a typical embedding width. + for (point_count, leader_count, dimensions) in [(33, 35, 129), (65, 129, 1536)] { + let mut point_values = + test_support::dense_points(point_count, dimensions, 1287); + let mut leader_values = + test_support::dense_points(leader_count, dimensions, 2026); + if metric == Metric::CosineNormalized { + test_support::normalize(&mut point_values, dimensions); + test_support::normalize(&mut leader_values, dimensions); } - }; - for point in 0..point_count { - let mut expected: Vec<_> = (0..leader_count) - .map(|leader| { - ( - leader as u32, - test_support::distance( - scalar_metric, - points.row(point), - leader_matrix.row(leader), - ), - ) + let points = + MatrixView::try_from(point_values.as_slice(), point_count, dimensions) + .unwrap(); + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_count, dimensions) + .unwrap(); + let leaders = M::create_leaders(leader_matrix); + let tolerance = test_support::dense_tolerance(metric, dimensions); + // The oracle sorts every leader by its scalar distance to the point. + let oracle: Vec> = (0..point_count) + .map(|point| { + let mut ranked: Vec<_> = (0..leader_count) + .map(|leader| { + let distance = test_support::distance( + metric, + points.row(point), + leader_matrix.row(leader), + ); + (leader as u32, distance) + }) + .collect(); + ranked.sort_by(|left, right| left.1.total_cmp(&right.1)); + ranked }) .collect(); - expected.sort_by(|left, right| left.1.total_cmp(&right.1)); - let actual = &output[point * assignments..(point + 1) * assignments]; - for (rank, &id) in actual.iter().enumerate() { - assert!( - !actual[..rank].contains(&id), - "shape={shape:?}, assignments={assignments}, duplicate leader {id} for point={point}" - ); - let score = expected - .iter() - .find(|&&(leader, _)| leader == id) - .unwrap_or_else(|| panic!("shape={shape:?}, assignments={assignments}, invalid leader {id} for point={point}")) - .1; - assert!( - (score - expected[rank].1).abs() <= tolerance, - "shape={shape:?}, assignments={assignments}, point={point}, rank={rank}, score={score}, expected={:?}", - expected[rank] + + // k = 3 uses a fixed-size nearest set. k = 11 uses a slice. + for assignments in [3, 11] { + let context = format!( + "{arch_name}, {metric:?}, shape=({point_count}, {leader_count}, {dimensions}), assignments={assignments}" ); + let mut output = vec![UNASSIGNED; point_count * assignments]; + + assign_leaders::( + arch, + points, + &leaders, + MutMatrixView::try_from( + output.as_mut_slice(), + point_count, + assignments, + ) + .unwrap(), + &mut PartitionKernelWorkspace::default(), + ) + .unwrap_or_else(|error| panic!("{context}: {error}")); + + for (point, expected) in oracle.iter().enumerate() { + let actual = &output[point * assignments..(point + 1) * assignments]; + for (rank, &id) in actual.iter().enumerate() { + assert!( + !actual[..rank].contains(&id), + "{context}, point={point}: duplicate leader {id}" + ); + let distance = expected + .iter() + .find(|&&(leader, _)| leader == id) + .unwrap_or_else(|| { + panic!("{context}, point={point}: invalid leader {id}") + }) + .1; + // Equal distances can select either leader, but each rank must + // be nearest-first. + assert!( + (distance - expected[rank].1).abs() <= tolerance, + "{context}, point={point}, rank={rank}: {distance} != {:?}", + expected[rank] + ); + } + } } } } } + + for_each_arch(&LargeStripes:: { + metric, + partition_metric: PhantomData, + }); } #[test] - fn workspace_reuse_replaces_assignments_for_each_stripe() { + fn workspace_reuse_does_not_mix_results_from_different_stripes() { let leader_values = [0.0, 5.0, 12.0]; let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); let point_values = [1.0, 7.0, 11.0]; @@ -335,7 +273,7 @@ mod tests { } #[test] - fn each_distance_row_replaces_previous_assignments_and_unused_slots() { + fn stale_output_and_scratch_do_not_leak_into_new_assignments() { let distances = [ 3.0, 1.0, @@ -353,7 +291,7 @@ mod tests { let mut output = [2; 9]; let mut candidates = vec![Candidate::new(3, -10.0); 4]; - rank_leader_distances( + select_top_k_ids( ARCH, MatrixView::try_from(&distances[..], 3, 4).unwrap(), MutMatrixView::try_from(&mut output[..], 3, 3).unwrap(), @@ -363,75 +301,32 @@ mod tests { assert_eq!( output, [ - 1, - 2, - 0, - UNASSIGNED_LEADER, - UNASSIGNED_LEADER, - UNASSIGNED_LEADER, - 1, - UNASSIGNED_LEADER, - UNASSIGNED_LEADER + 1, 2, 0, UNASSIGNED, UNASSIGNED, UNASSIGNED, 1, UNASSIGNED, UNASSIGNED ] ); } #[test] - fn invalid_output_rows_preserve_assignments_and_workspace() { + fn output_without_one_row_per_point_is_rejected() { let point_values = [1.0, 6.0]; let leader_values = [0.0, 5.0, 12.0]; let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); let mut output = [2]; - let previous_candidates = vec![Candidate::new(1, 7.0)]; - let mut workspace = PartitionKernelWorkspace { - distance_scratch: vec![9.0, 11.0], - ranked_leaders: previous_candidates.clone(), - }; let error = assign_leaders::<_, L2>( ARCH, MatrixView::try_from(&point_values[..], 2, 1).unwrap(), &leaders, MutMatrixView::try_from(&mut output[..], 1, 1).unwrap(), - &mut workspace, + &mut PartitionKernelWorkspace::default(), ) .unwrap_err(); assert!( error .to_string() - .contains("invalid partition output row count 1 for 2 points") - ); - assert_eq!(output, [2]); - assert_eq!(workspace.distance_scratch, [9.0, 11.0]); - assert_eq!(workspace.ranked_leaders, previous_candidates); - } - - #[test] - fn metric_dimension_errors_are_returned_without_publishing_assignments() { - let point_values = [1.0, 2.0, 3.0]; - let leader_values = [0.0, 1.0, 2.0, 3.0]; - let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 2, 2).unwrap()); - let mut output = [1]; - - let error = assign_leaders::<_, L2>( - ARCH, - MatrixView::try_from(&point_values[..], 1, 3).unwrap(), - &leaders, - MutMatrixView::try_from(&mut output[..], 1, 1).unwrap(), - &mut PartitionKernelWorkspace::default(), - ) - .unwrap_err(); - - assert_eq!( - error.downcast_ref::(), - Some(&diskann_linalg::SgemmError::InvalidMatrixDimensions { - matrix_name: diskann_linalg::MatrixName::B, - expected_rows: 3, - expected_cols: 2, - actual_len: 4, - }) + .contains("invalid kernel output row count 1 for 2 points"), + "{error}" ); - assert_eq!(output, [1]); } } diff --git a/diskann/src/graph/pipnn/partition_metric.rs b/diskann/src/graph/pipnn/partition_metric.rs index ca597ae60d..29078bde87 100644 --- a/diskann/src/graph/pipnn/partition_metric.rs +++ b/diskann/src/graph/pipnn/partition_metric.rs @@ -3,10 +3,13 @@ * Licensed under the MIT license. */ -//! Build ranking distances from point stripes to partition leaders. +//! Ranking distances from points to partition leaders. //! -//! A ranking distance preserves nearest-first order. One leader set serves all -//! point stripes in a partition split. +//! A ranking distance orders the leaders of one point in the same way as the +//! metric distance. It can omit terms that are equal for every leader of the +//! point, so its value can differ from the metric distance. A point stripe is a +//! block of consecutive points that one task assigns. One leader set serves all +//! point stripes of a partition split. use crate::{ANNError, ANNResult}; use diskann_linalg::Transpose; @@ -18,39 +21,40 @@ use diskann_vector::{ use super::{Cosine, CosineNormalized, InnerProduct, L2, cosine_distance}; -/// Store leader values with immutable metric data. +/// Leader vectors and the leader norms that a metric needs. /// -/// L2 stores squared norms. Cosine stores norms. Construction computes them once -/// before point stripes share the leader set. Other metrics need no norms. +/// L2 stores squared norms, and Cosine stores norms. The norms are computed once, +/// when the leader set is created, and all point stripes share them. Inner product +/// and normalized cosine store no norms. pub(super) struct PartitionLeaders<'a, Norms> { values: MatrixView<'a, f32>, norms: Norms, } -/// Fill one point-to-leader ranking matrix. +/// Compute ranking distances from points to partition leaders. /// -/// The associated leader type hides metric data from the caller. The caller -/// creates one value and shares it across all point stripes. +/// `Leaders` hides the norms of each metric from the caller. The caller creates +/// one leader set for a partition split and shares it across all point stripes. pub(super) trait PartitionMetric: Send + Sync + 'static { - /// Leader values and immutable metric data for one partition split. + /// Leader vectors and precomputed norms for one partition split. type Leaders<'a>: Sync; - /// Bind one non-empty leader matrix to this metric. + /// Create a leader set from a matrix with one leader in each row. /// - /// Partitioning creates at least one leader before it calls this function. + /// Partitioning always samples at least one leader. The kernels do not support + /// an empty leader set. fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a>; - /// Return the number of leaders in a metric-owned leader set. + /// Return the number of leaders. fn leader_count(leaders: &Self::Leaders<'_>) -> usize; - /// Compute one row-major point-to-leader ranking buffer. + /// Write the ranking distance from each point to each leader. /// - /// Values preserve nearest-first order. L2 omits the point's squared norm, - /// which is constant across its row. Normalized cosine and inner product - /// return `-dot`; cosine returns `1 - similarity`. - /// - /// `storage` has `points.nrows()` rows and `leader_count` columns. - /// A zero distance can have either sign. Equal distances can select either leader. + /// L2 omits the squared norm of the point, which is equal for every leader of + /// the point. Normalized cosine and inner product give `-dot`. Cosine gives + /// `1 - similarity`. `storage` has one row per point and one column per leader. + /// [`assign_leaders`](super::partition_kernel::assign_leaders) creates it with + /// this shape. A zero distance can have either sign. fn compute_distances( points: MatrixView<'_, f32>, leaders: &Self::Leaders<'_>, @@ -58,15 +62,8 @@ pub(super) trait PartitionMetric: Send + Sync + 'static { ) -> ANNResult<()>; } -/// Compute squared norms with DiskANN's vector implementation. -fn l2_squared_norms(vectors: MatrixView<'_, f32>) -> Vec { - vectors - .row_iter() - .map(|vector| FastL2NormSquared.evaluate(vector)) - .collect() -} - -/// Compute cosine norms with the same reduction for points and leaders. +/// Compute the L2 norm of each row. Points and leaders both use this function, so +/// their norms round the same way. fn cosine_norms(vectors: MatrixView<'_, f32>) -> Vec { vectors .row_iter() @@ -80,7 +77,10 @@ impl PartitionMetric for L2 { fn create_leaders<'a>(values: MatrixView<'a, f32>) -> Self::Leaders<'a> { PartitionLeaders { values, - norms: l2_squared_norms(values), + norms: values + .row_iter() + .map(|leader| FastL2NormSquared.evaluate(leader)) + .collect(), } } @@ -93,11 +93,9 @@ impl PartitionMetric for L2 { leaders: &Self::Leaders<'_>, mut storage: MutMatrixView<'_, f32>, ) -> ANNResult<()> { - if storage.nrows() != points.nrows() || storage.ncols() != leaders.values.nrows() { - return Err(ANNError::message("point-to-leader output shape mismatch")); - } - // The point norm is constant across a point row. It cannot change ranking. - // Initialize each point row before GEMM adds the dot-product term. + // The ranking distance is `||l||² - 2(p·l)`: the squared L2 distance without + // `||p||²`, which is equal for every leader of the point. Start each row with + // the leader norms, and let GEMM add the dot-product term. for row in storage.row_iter_mut() { row.copy_from_slice(&leaders.norms); } @@ -137,9 +135,6 @@ impl PartitionMetric for Cosine { leaders: &Self::Leaders<'_>, mut storage: MutMatrixView<'_, f32>, ) -> ANNResult<()> { - if storage.nrows() != points.nrows() || storage.ncols() != leaders.values.nrows() { - return Err(ANNError::message("point-to-leader output shape mismatch")); - } diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -155,7 +150,8 @@ impl PartitionMetric for Cosine { .map_err(ANNError::new)?; let point_norms = cosine_norms(points); let leader_norms = &leaders.norms; - // Convert each dot to cosine distance. Reuse leader norms across stripes. + // Convert each dot to cosine distance. The leader norms come from the leader + // set, so each stripe computes only its point norms. for (row, &point_norm) in storage.row_iter_mut().zip(point_norms.iter()) { for (distance, &leader_norm) in row.iter_mut().zip(leader_norms.iter()) { *distance = cosine_distance(*distance, point_norm, leader_norm); @@ -181,9 +177,6 @@ impl PartitionMetric for InnerProduct { leaders: &Self::Leaders<'_>, mut storage: MutMatrixView<'_, f32>, ) -> ANNResult<()> { - if storage.nrows() != points.nrows() || storage.ncols() != leaders.values.nrows() { - return Err(ANNError::message("point-to-leader output shape mismatch")); - } diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -266,7 +259,7 @@ mod tests { #[case::cosine(Cosine, Metric::Cosine)] #[case::normalized_cosine(CosineNormalized, Metric::CosineNormalized)] #[case::inner_product(InnerProduct, Metric::InnerProduct)] - fn point_to_leader_scores_match_scalar_distances( + fn ranking_distances_match_the_scalar_reference( #[case] _metric: M, #[case] scalar_metric: Metric, ) { @@ -274,7 +267,7 @@ mod tests { for leader_count in [1, 4, 17] { for dimensions in [1, 2, 7, 8, 9, 15, 16, 17, 127, 128, 129] { // Small integer coordinates make unnormalized dot products exact. Points - // and leaders use different values, so swapped rows or columns change scores. + // and leaders use different values, so swapped rows or columns change the distances. let mut point_values: Vec<_> = (0..point_count * dimensions) .map(|index| (index % 7) as f32 - 3.0) .collect(); @@ -322,13 +315,17 @@ mod tests { points.row(point), leader_matrix.row(leader), ); - if scalar_metric == Metric::L2 { - // Partition L2 omits exactly this constant from every column of the row. - expected -= points - .row(point) - .iter() - .map(|&x| f64::from(x).powi(2)) - .sum::(); + // Partition ranking omits one constant from every column of the row. + match scalar_metric { + Metric::L2 => { + expected -= points + .row(point) + .iter() + .map(|&x| f64::from(x).powi(2)) + .sum::(); + } + Metric::CosineNormalized => expected -= 1.0, + Metric::Cosine | Metric::InnerProduct => {} } let tolerance = match scalar_metric { Metric::L2 | Metric::InnerProduct => 0.0, @@ -409,8 +406,10 @@ mod tests { points.row(point), leader_matrix.row(leader), ); - if scalar_metric == Metric::L2 { - expected -= point_norm; + match scalar_metric { + Metric::L2 => expected -= point_norm, + Metric::CosineNormalized => expected -= 1.0, + Metric::Cosine | Metric::InnerProduct => {} } let actual = f64::from(output[point * leader_count + leader]); assert!( @@ -466,35 +465,6 @@ mod tests { assert_eq!(output, [1.0, 1.0, 1.0, 1.0, 1.0, 2.0]); } - #[rstest] - #[case::l2(L2)] - #[case::cosine(Cosine)] - #[case::normalized_cosine(CosineNormalized)] - #[case::inner_product(InnerProduct)] - fn output_shape_mismatch_is_rejected_before_writing( - #[case] _metric: M, - #[values((1, 3), (2, 2), (3, 3), (2, 4))] shape: (usize, usize), - ) { - let point_values = [1.0, 2.0, 3.0, 4.0]; - let leader_values = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; - let leaders = M::create_leaders(MatrixView::try_from(&leader_values[..], 3, 2).unwrap()); - let mut output = vec![17.0; shape.0 * shape.1]; - - let error = M::compute_distances( - MatrixView::try_from(&point_values[..], 2, 2).unwrap(), - &leaders, - MutMatrixView::try_from(output.as_mut_slice(), shape.0, shape.1).unwrap(), - ) - .unwrap_err(); - - assert!( - error - .to_string() - .contains("point-to-leader output shape mismatch") - ); - assert_eq!(output, vec![17.0; shape.0 * shape.1]); - } - #[rstest] #[case::l2(L2)] #[case::cosine(Cosine)] diff --git a/diskann/src/graph/pipnn/simd.rs b/diskann/src/graph/pipnn/simd.rs index 3b2a80863a..4a5784d130 100644 --- a/diskann/src/graph/pipnn/simd.rs +++ b/diskann/src/graph/pipnn/simd.rs @@ -3,28 +3,54 @@ * Licensed under the MIT license. */ -//! SIMD schema for PiPNN numerical kernels. +//! The SIMD vector type of the PiPNN ranking kernels. -use diskann_wide::{Architecture, SIMDFloat, SIMDMask, SIMDVector}; +use diskann_wide::{Architecture, Const, SIMDFloat, SIMDMask, SIMDVector}; -/// Default SIMD representation used by both PiPNN ranking kernels. +/// The number of distances in one SIMD group of both ranking kernels. /// -/// This alias is the single build-time width selection. -type DefaultVector = ::f32x16; +/// Change the SIMD width here only. 16 lanes fill one AVX-512 register. AVX2 uses +/// two registers for a group, and Neon uses four. +pub(super) const LANES: usize = 16; -/// PiPNN SIMD representation for one architecture. +/// Map a lane count `Const` to the `f32` vector of architecture `A` with `N` +/// lanes. +pub(super) trait F32Vector { + type Vector; +} + +impl F32Vector for Const<4> { + type Vector = A::f32x4; +} + +impl F32Vector for Const<8> { + type Vector = A::f32x8; +} + +impl F32Vector for Const<16> { + type Vector = A::f32x16; +} + +/// The `f32` vector of architecture `A` with [`LANES`] lanes. +type DefaultVector = as F32Vector>::Vector; + +/// An architecture with an `f32` vector of [`LANES`] lanes. +/// +/// The blanket implementation below covers `Scalar`, `V3`, `V4`, and `Neon`. pub(super) trait Simd: Architecture { - /// SIMD vector used by both ranking kernels. - type Vector: SIMDVector + SIMDFloat; + /// The vector that holds one group of [`LANES`] distances. + type Vector: SIMDVector> + SIMDFloat; - /// Return one bit for each selected lane. + /// Convert a lane mask to an integer with bit `i` set when lane `i` is set. fn active_lanes(mask: ::Mask) -> u64; } +// The mask must convert to `u64`, so the kernels can visit the set lanes with +// integer bit operations. impl Simd for A where A: Architecture, - DefaultVector: SIMDVector + SIMDFloat, + DefaultVector: SIMDVector> + SIMDFloat, u64: From< << as SIMDVector>::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying, >, diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index e22db2bfb4..f8f9300074 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -3,18 +3,40 @@ * Licensed under the MIT license. */ -//! Top-k updates from caller-selected distance slices. -//! NaN and positive infinity are not retained. - -use diskann_utils::views::MutMatrixView; +//! Top-k selection for the PiPNN ranking kernels. +//! +//! The nearest set of a point holds the k nearest candidates found so far, nearest +//! first. An unfilled slot holds [`Candidate::EMPTY`]. The k-th distance of a +//! nearest set is the distance in its last slot. It stays positive infinity until +//! the set is full. To offer a candidate to a nearest set is to insert the +//! candidate if it is nearer than the k-th distance. NaN and positive infinity are +//! never nearer, so they never enter a nearest set. +//! +//! [`select_top_k_ids`] serves partition assignment. It selects the nearest set of +//! each distance row and writes its IDs. [`select_top_k_symmetric`] serves leaf +//! construction. It runs a pair scan: it reads each unordered point pair once and +//! offers the pair to the nearest sets of both points. The caller sets k through +//! the width of the output. For common values of k, both functions use fixed-size +//! nearest sets, so the compiler can unroll insertion. +//! +//! Both functions compare [`LANES`] distances at a time with the k-th distance and +//! insert only the nearer ones. After a nearest set fills, most groups have no +//! nearer distance, so one SIMD comparison replaces `LANES` scalar comparisons. +//! The SIMD loops run inside `run2` or `run3` of architecture `A`, which compiles +//! them with the target features of `A`. + +use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_wide::{SIMDPartialOrd, SIMDVector}; -use super::simd::Simd; +use super::simd::{LANES, Simd}; -/// An output slot with no assigned candidate. +/// The ID of an output slot that holds no candidate. pub(super) const UNASSIGNED: u32 = u32::MAX; -/// A local candidate index paired with its ranking distance. +/// One slot of a nearest set: a candidate ID and its distance. +/// +/// The ID is a position in the kernel input. Partition assignment uses leader +/// columns. Leaf construction uses point positions in the leaf. #[derive(Clone, Copy, Debug, PartialEq)] pub(super) struct Candidate { pub(super) local_idx: u32, @@ -22,6 +44,9 @@ pub(super) struct Candidate { } impl Candidate { + /// An unfilled slot. Every distance that can enter a nearest set is nearer. + pub(super) const EMPTY: Self = Self::new(UNASSIGNED, f32::INFINITY); + pub(super) const fn new(local_idx: u32, distance: f32) -> Self { Self { local_idx, @@ -29,6 +54,7 @@ impl Candidate { } } + /// Return `true` if this slot holds a candidate. pub(super) const fn is_assigned(self) -> bool { self.local_idx != UNASSIGNED } @@ -36,398 +62,301 @@ impl Candidate { impl Default for Candidate { fn default() -> Self { - Self::new(UNASSIGNED, f32::INFINITY) - } -} - -/// One borrowed result row whose length is its capacity. -pub(super) trait Ranker: AsRef<[Candidate]> + AsMut<[Candidate]> + Sized { - /// Insert a candidate below this non-empty row's limit and return the new limit. - fn insert(&mut self, candidate: Candidate) -> f32; - - /// Replace this row with the nearest input positions. - /// - /// Equal distances need no fixed tie order. Empty rows retain no candidates. - #[inline(always)] - fn select_topk(&mut self, arch: A, distances: &[f32]) { - arch.run2( - #[inline(always)] - move |distances: &[f32], ranker: &mut Self| { - if ranker.as_ref().is_empty() { - return; - } - ranker.as_mut().fill(Candidate::default()); - distance_blocks(arch, distances).fold( - f32::INFINITY, - #[inline(always)] - |limit, block| block.update_topk_from_candidates(ranker, limit), - ); - }, - distances, - self, - ); - } -} - -impl Ranker for &mut [Candidate] { - #[inline(always)] - fn insert(&mut self, candidate: Candidate) -> f32 { - insert_sorted(self, candidate) + Self::EMPTY } } -impl Ranker for &mut [Candidate; K] { - #[inline(always)] - fn insert(&mut self, candidate: Candidate) -> f32 { - insert_sorted(*self, candidate) - } -} - -/// A batch operation over one bound result row. -pub(super) trait TopKVisitor { - fn visit(self, ranker: R); -} - -/// Bind one result row and specialize its capacity once for the visitor. +/// Write the IDs of the k nearest columns of each distance row, nearest first. /// -/// Specialize capacities 1, 2, 3, 8, and 10; other capacities use the slice path. -/// Typical partition K is 10 or 3, and leaf K is 3 or 2. -#[inline] -#[expect( - clippy::unwrap_used, - reason = "each match arm proves the candidate slice has the array length" -)] -pub(super) fn with_topk(candidates: &mut [Candidate], visitor: V) { - // Each arm checks this same slice's length, so conversion to its array cannot fail. - match candidates.len() { - 1 => visitor.visit::<&mut [_; 1]>(candidates.try_into().unwrap()), - 2 => visitor.visit::<&mut [_; 2]>(candidates.try_into().unwrap()), - 3 => visitor.visit::<&mut [_; 3]>(candidates.try_into().unwrap()), - 8 => visitor.visit::<&mut [_; 8]>(candidates.try_into().unwrap()), - 10 => visitor.visit::<&mut [_; 10]>(candidates.try_into().unwrap()), - _ => visitor.visit(candidates), - } -} - -/// Access complete result rows with their static or runtime capacity. -pub(super) trait BatchRanker { - type Ranker<'a>: Ranker - where - Self: 'a; - - fn ranker(&mut self, index: usize) -> Option>; -} - -/// Candidate rows and their limits stay borrowed together throughout a pair scan. -pub(super) struct Batch<'a, Rows> { - candidates: Rows, - thresholds: &'a mut [f32], -} - -impl<'a> Batch<'a, MutMatrixView<'a, Candidate>> { - fn new(mut candidates: MutMatrixView<'a, Candidate>, thresholds: &'a mut Vec) -> Self { - let rows = candidates.nrows(); - candidates.as_mut_slice().fill(Candidate::default()); - thresholds.truncate(rows); - thresholds.fill(f32::INFINITY); - thresholds.resize(rows, f32::INFINITY); - Self { - candidates, - thresholds, +/// k is the column count of `output`, and `output` has one row per distance row. +/// Equal distances can select either column. If a row has fewer than k distances +/// that are not NaN or positive infinity, the remaining slots hold [`UNASSIGNED`]. +/// `scratch` is reusable storage for values of k without a fixed-size nearest set. +/// +/// # Panics +/// +/// Panics if `output` does not have one row per distance row, or if `distances` +/// has no columns and k is not zero. +pub(super) fn select_top_k_ids( + arch: A, + distances: MatrixView<'_, f32>, + output: MutMatrixView<'_, u32>, + scratch: &mut Vec, +) { + assert_eq!( + distances.nrows(), + output.nrows(), + "top-k IDs need one output row per distance row" + ); + // Fixed-size nearest sets serve the common partition fanouts 1, 2, 3, 8, and + // 10. On AVX2 with 1000-leader rows, a slice took about 3x the ranking time of + // a fixed-size set at k = 8 and k = 10. That was 16-25% of `assign_leaders` + // time at 128 dimensions. A branch-free insertion was slower than both. + match output.ncols() { + 0 => {} + 1 => select_ids_with(arch, distances, output, &mut [Candidate::EMPTY; 1]), + 2 => select_ids_with(arch, distances, output, &mut [Candidate::EMPTY; 2]), + 3 => select_ids_with(arch, distances, output, &mut [Candidate::EMPTY; 3]), + 8 => select_ids_with(arch, distances, output, &mut [Candidate::EMPTY; 8]), + 10 => select_ids_with(arch, distances, output, &mut [Candidate::EMPTY; 10]), + k => { + scratch.resize(k, Candidate::EMPTY); + select_ids_with(arch, distances, output, scratch.as_mut_slice()); } } } -impl Batch<'_, Rows> { - fn nrows(&self) -> usize { - self.thresholds.len() - } -} - -impl BatchRanker for Batch<'_, &mut [[Candidate; K]]> { - type Ranker<'a> - = &'a mut [Candidate; K] - where - Self: 'a; - - #[inline(always)] - fn ranker(&mut self, index: usize) -> Option> { - self.candidates.get_mut(index) +/// Fill the nearest set of each point with its k nearest other points. +/// +/// `distances` is a symmetric matrix with one row and one column per point, and +/// `output` has one row per point. k is the column count of `output`. Only the +/// strict lower triangle of `distances` is read: each pair is read once and +/// offered to both points. Candidate IDs are point positions. Equal distances can +/// select either point. Slots that no pair fills hold [`Candidate::EMPTY`]. +/// `kth_distances` is reusable storage. It holds the k-th distance of each point +/// after the scan. +/// +/// # Panics +/// +/// Panics if `distances` does not have one row and one column per output row. +pub(super) fn select_top_k_symmetric( + arch: A, + distances: MatrixView<'_, f32>, + output: MutMatrixView<'_, Candidate>, + kth_distances: &mut Vec, +) { + let points = output.nrows(); + assert!( + distances.nrows() == points && distances.ncols() == points, + "a symmetric scan needs one distance row and column per point" + ); + let k = output.ncols(); + let slots = output.into_inner(); + // Fixed-size nearest sets serve the common `leaf_k` values 1, 2, and 3. On AVX2 + // at k = 3, slices took 1.1x to 2.5x the ranking time of fixed-size sets. + match k { + 0 => {} + 1 => scan_pairs(arch, distances, slots.as_chunks_mut::<1>().0, kth_distances), + 2 => scan_pairs(arch, distances, slots.as_chunks_mut::<2>().0, kth_distances), + 3 => scan_pairs(arch, distances, slots.as_chunks_mut::<3>().0, kth_distances), + k => { + let mut neighborhoods: Vec<&mut [Candidate]> = slots.chunks_exact_mut(k).collect(); + scan_pairs(arch, distances, &mut neighborhoods, kth_distances); + } } } -impl BatchRanker for Batch<'_, MutMatrixView<'_, Candidate>> { - type Ranker<'a> - = &'a mut [Candidate] - where - Self: 'a; - - #[inline(always)] - fn ranker(&mut self, index: usize) -> Option> { - if index < self.nrows() { - Some(self.candidates.row_mut(index)) - } else { - None +/// Select the nearest set of each distance row in `nearest`, then copy its IDs to +/// the matching output row. All rows reuse the storage of `nearest`. +fn select_ids_with( + arch: A, + distances: MatrixView<'_, f32>, + mut output: MutMatrixView<'_, u32>, + nearest: &mut Nearest, +) where + A: Simd, + Nearest: AsMut<[Candidate]> + ?Sized, +{ + for (row, ids) in distances.row_iter().zip(output.row_iter_mut()) { + select_nearest(arch, nearest, row); + for (id, candidate) in ids.iter_mut().zip(nearest.as_mut().iter()) { + *id = candidate.local_idx; } } } -/// A batch operation over initialized pair-scan storage. -pub(super) trait BatchVisitor { - fn visit<'a, Rows>(self, batch: Batch<'a, Rows>) - where - Batch<'a, Rows>: BatchRanker; -} - -/// Initialize a result matrix and specialize its row representation once for the visitor. +/// Replace the contents of a non-empty nearest set with the k nearest entries of +/// `distances`. +/// +/// Each candidate ID is a position in `distances`. #[inline] -pub(super) fn with_batch( - candidates: MutMatrixView<'_, Candidate>, - thresholds: &mut Vec, - visitor: V, -) { - let Batch { - candidates, - thresholds, - } = Batch::new(candidates, thresholds); - // The matrix width selects K, so its storage contains only complete K-element rows. - match candidates.ncols() { - 1 => visitor.visit(Batch { - candidates: candidates.into_inner().as_chunks_mut::<1>().0, - thresholds, - }), - 2 => visitor.visit(Batch { - candidates: candidates.into_inner().as_chunks_mut::<2>().0, - thresholds, - }), - 3 => visitor.visit(Batch { - candidates: candidates.into_inner().as_chunks_mut::<3>().0, - thresholds, - }), - 8 => visitor.visit(Batch { - candidates: candidates.into_inner().as_chunks_mut::<8>().0, - thresholds, - }), - 10 => visitor.visit(Batch { - candidates: candidates.into_inner().as_chunks_mut::<10>().0, - thresholds, - }), - _ => visitor.visit(Batch { - candidates, - thresholds, - }), - } -} - -impl<'a, Rows> Batch<'a, Rows> +fn select_nearest(arch: A, nearest: &mut Nearest, distances: &[f32]) where - Self: BatchRanker, + A: Simd, + Nearest: AsMut<[Candidate]> + ?Sized, { - /// Offer one point's distances to all preceding points to both endpoints' results. - /// - /// The source point index is `distances.len()`: each entry describes its pair - /// with the earlier point at that position. Supply each pair once; earlier results are preserved. - /// Equal distances need no fixed tie order. - /// - /// # Panics - /// - /// Panics if `distances.len()` is outside the batch. - #[inline] - #[expect( - clippy::unwrap_used, - reason = "the first ranker lookup rejects an invalid point before scanning" - )] - pub(super) fn update_dual_topk(&mut self, arch: A, distances: &[f32]) { - let point_idx = distances.len(); - // ranker rejects an out-of-bounds point before any results are changed. - if self.ranker(point_idx).unwrap().as_ref().is_empty() { - return; - } - arch.run2( - #[inline(always)] - move |distances: &[f32], batch: &mut Self| { - let limit = distance_blocks(arch, distances).fold( - batch.thresholds[point_idx], - #[inline(always)] - |limit, block| { - // End the source row borrow before offering this block to earlier rows. - let limit = { - let mut nearest = batch.ranker(point_idx).unwrap(); - block.update_topk_from_candidates(&mut nearest, limit) - }; - block.update_topks_with_point(batch, point_idx as u32); - limit - }, - ); - batch.thresholds[point_idx] = limit; - }, - distances, - self, - ); - } -} - -/// A complete SIMD group or one tail distance. This never leaves the TopK module. -/// Both update directions share the loaded vector; scalar reads borrow the input. -#[derive(Clone, Copy)] -enum DistanceBlock<'a, A: Simd> { - Simd { - first_candidate: usize, - values: A::Vector, - distances: &'a [f32], - }, - Scalar { - candidate_idx: usize, - distance: f32, - }, + arch.run2( + #[inline(always)] + move |distances: &[f32], nearest: &mut Nearest| { + let nearest = nearest.as_mut(); + nearest.fill(Candidate::EMPTY); + let (groups, tail) = distances.as_chunks::(); + let mut kth_distance = f32::INFINITY; + for (group, group_distances) in groups.iter().enumerate() { + let (first, values) = (group * LANES, load_group(arch, group_distances)); + kth_distance = + offer_group(arch, nearest, first, group_distances, values, kth_distance); + } + offer_each(nearest, groups.len() * LANES, tail, kth_distance); + }, + distances, + nearest, + ); } -/// Keep all distance loads and tail indexing inside TopK's architecture scope. -#[inline(always)] -fn distance_blocks( +/// Run a pair scan over every point, and leave the k-th distance of each point in +/// `kth_distances`. +fn scan_pairs( arch: A, - distances: &[f32], -) -> impl Iterator> { - let simd_end = distances.len() - distances.len() % A::Vector::LANES; - distances[..simd_end] - .chunks_exact(A::Vector::LANES) - .enumerate() - .map(move |(group, distances)| DistanceBlock::Simd { - first_candidate: group * A::Vector::LANES, - // SAFETY: chunks_exact supplies one distance for every SIMD lane. - values: unsafe { A::Vector::load_simd(arch, distances.as_ptr()) }, - distances, - }) - .chain( - distances[simd_end..] - .iter() - .enumerate() - .map(move |(tail, &distance)| DistanceBlock::Scalar { - candidate_idx: simd_end + tail, - distance, - }), - ) + distances: MatrixView<'_, f32>, + neighborhoods: &mut [Nearest], + kth_distances: &mut Vec, +) where + A: Simd, + Nearest: AsMut<[Candidate]>, +{ + // The k-th distances live in their own array. One SIMD load then reads the + // k-th distances of `LANES` targets, and one comparison checks `LANES` pairs. + neighborhoods + .iter_mut() + .for_each(|nearest| nearest.as_mut().fill(Candidate::EMPTY)); + kth_distances.clear(); + kth_distances.resize(neighborhoods.len(), f32::INFINITY); + // Point 0 has no earlier point to pair with. + for source in 1..neighborhoods.len() { + let row = distances.row(source); + offer_pairs(arch, source, row, neighborhoods, kth_distances); + } } -impl DistanceBlock<'_, A> { - /// Update one top-k result from this block's candidates. - /// - /// Each candidate uses its index in the full distance row as its local ID. - /// Partition ranking uses leader IDs. Leaf ranking uses IDs of earlier points in the leaf. - /// - /// Compare candidates against the complete result row without clearing its slots. - /// The distance limit is the last slot's distance, or positive infinity while slots remain empty. - /// Each insertion can lower this limit. Return the updated limit for the next block. - #[inline(always)] - fn update_topk_from_candidates(self, nearest: &mut R, mut max_distance: f32) -> f32 { - match self { - Self::Scalar { - candidate_idx, - distance, - } => { - if distance < max_distance { - max_distance = nearest.insert(Candidate::new(candidate_idx as u32, distance)); +/// Offer the pair of point `source` and each earlier point `target` to the nearest +/// sets of both points. +/// +/// `row[target]` is the distance between `source` and `target`. Only the entries +/// before `source` are read. For each point `p`, `kth_distances[p]` must be the +/// k-th distance of `neighborhoods[p]`. This function keeps that true. +#[inline] +fn offer_pairs( + arch: A, + source: usize, + row: &[f32], + neighborhoods: &mut [Nearest], + kth_distances: &mut [f32], +) where + A: Simd, + Nearest: AsMut<[Candidate]>, +{ + arch.run3( + #[inline(always)] + move |distances: &[f32], neighborhoods: &mut [Nearest], kth_distances: &mut [f32]| { + let mut kth_distance = kth_distances[source]; + let (groups, tail) = distances.as_chunks::(); + for (group, group_distances) in groups.iter().enumerate() { + let first = group * LANES; + // The source and its targets have different nearest sets, so both + // directions use the same loaded vector. + let values = load_group(arch, group_distances); + let nearest = neighborhoods[source].as_mut(); + kth_distance = + offer_group(arch, nearest, first, group_distances, values, kth_distance); + // Each lane compares with the k-th distance of its own target. The + // source has no effect on this decision. Each lane also updates a + // different target, so the mask stays correct while lanes insert. + // `offer_group` must check again, because all its lanes share one set. + let target_kth = load_group(arch, &kth_distances.as_chunks::().0[group]); + let mut eligible = A::active_lanes(values.lt_simd(target_kth)); + while eligible != 0 { + let lane = eligible.trailing_zeros() as usize; + eligible &= eligible - 1; + let candidate = Candidate::new(source as u32, group_distances[lane]); + let target = first + lane; + kth_distances[target] = + insert_sorted(neighborhoods[target].as_mut(), candidate); } } - Self::Simd { - first_candidate, - values, - distances, - } => { - // An unfilled result admits all rankable values. Scan them in - // order without constructing and traversing an all-eligible mask. - if max_distance == f32::INFINITY { - for (lane, &distance) in distances.iter().enumerate() { - if distance < max_distance { - max_distance = nearest - .insert(Candidate::new((first_candidate + lane) as u32, distance)); - } - } - } else { - let mut eligible = A::active_lanes( - values.lt_simd(A::Vector::splat(values.arch(), max_distance)), - ); - while eligible != 0 { - let lane = eligible.trailing_zeros() as usize; - eligible &= eligible - 1; - // Earlier insertions in this block can lower the limit. - if distances[lane] < max_distance { - max_distance = nearest.insert(Candidate::new( - (first_candidate + lane) as u32, - distances[lane], - )); - } - } + // The tail has fewer than `LANES` pairs. Offer each pair to both points + // without SIMD. + let first = groups.len() * LANES; + for (offset, &distance) in tail.iter().enumerate() { + let target = first + offset; + if distance < kth_distance { + let candidate = Candidate::new(target as u32, distance); + kth_distance = insert_sorted(neighborhoods[source].as_mut(), candidate); + } + if distance < kth_distances[target] { + let candidate = Candidate::new(source as u32, distance); + kth_distances[target] = + insert_sorted(neighborhoods[target].as_mut(), candidate); } } + kth_distances[source] = kth_distance; + }, + &row[..source], + neighborhoods, + kth_distances, + ); +} + +/// Offer the candidates of one group to a nearest set, and return the new k-th +/// distance. +/// +/// `first_candidate` is the candidate ID of `distances[0]`, and `values` holds +/// the same distances as `distances`. +#[inline(always)] +fn offer_group( + arch: A, + nearest: &mut [Candidate], + first_candidate: usize, + distances: &[f32; LANES], + values: A::Vector, + mut kth_distance: f32, +) -> f32 { + // An unfilled nearest set accepts every distance that is not NaN or positive + // infinity, so a scalar scan skips the mask. On AVX2, removing this path cost + // 3% of leaf ranking time at k = 3 and 20% for 100-leader rows at k = 3. + if kth_distance == f32::INFINITY { + return offer_each(nearest, first_candidate, distances, kth_distance); + } + let mut eligible = A::active_lanes(values.lt_simd(A::Vector::splat(arch, kth_distance))); + while eligible != 0 { + let lane = eligible.trailing_zeros() as usize; + eligible &= eligible - 1; + // An earlier insertion in this group can lower the k-th distance. + if distances[lane] < kth_distance { + let candidate = Candidate::new((first_candidate + lane) as u32, distances[lane]); + kth_distance = insert_sorted(nearest, candidate); } - max_distance } + kth_distance +} - /// Offer one leaf point to the top-k results of earlier leaf points. - /// - /// This block stores distances from `point_idx` to those earlier points. - /// Each candidate's local ID selects its result row in the full leaf, not its SIMD lane. - /// If that row accepts the pair, insert `point_idx` with the pair's distance. - /// - /// Only the earlier rows identified by this block can change. - /// Their candidates are retained unless a nearer point displaces them. - /// - /// Each row compares the pair distance against its own limit. - /// Its decision does not depend on the current point's top-k result. - /// Update the row's limit after an insertion. - #[inline(always)] - #[expect( - clippy::unwrap_used, - reason = "all block positions precede the checked source point" - )] - fn update_topks_with_point<'a, Rows>(self, batch: &mut Batch<'a, Rows>, point_idx: u32) - where - Batch<'a, Rows>: BatchRanker, - { - match self { - Self::Scalar { - candidate_idx, - distance, - } => { - if distance < batch.thresholds[candidate_idx] { - // Block positions are less than the source index checked by update_dual_topk. - let limit = batch - .ranker(candidate_idx) - .unwrap() - .insert(Candidate::new(point_idx, distance)); - batch.thresholds[candidate_idx] = limit; - } - } - Self::Simd { - first_candidate, - values, - distances, - } => { - let limits = &batch.thresholds[first_candidate..][..A::Vector::LANES]; - // SAFETY: the slice above has one distance limit per SIMD lane. - let limit_values = unsafe { A::Vector::load_simd(values.arch(), limits.as_ptr()) }; - let mut eligible = A::active_lanes(values.lt_simd(limit_values)); - while eligible != 0 { - let lane = eligible.trailing_zeros() as usize; - eligible &= eligible - 1; - // Each lane updates a different result; its limit is current. - let index = first_candidate + lane; - // Every SIMD lane describes an earlier, in-bounds point. - let limit = batch - .ranker(index) - .unwrap() - .insert(Candidate::new(point_idx, distances[lane])); - batch.thresholds[index] = limit; - } - } +/// Offer each entry of `distances` to a nearest set, and return the new k-th +/// distance. +/// +/// `first_candidate` is the candidate ID of `distances[0]`. +#[inline(always)] +fn offer_each( + nearest: &mut [Candidate], + first_candidate: usize, + distances: &[f32], + mut kth_distance: f32, +) -> f32 { + for (offset, &distance) in distances.iter().enumerate() { + if distance < kth_distance { + let candidate = Candidate::new((first_candidate + offset) as u32, distance); + kth_distance = insert_sorted(nearest, candidate); } } + kth_distance +} + +/// Load one group of distances without copying it. +/// +/// A load from a copy (`from_array(*group)`) can leave the copy on the stack in a +/// large loop. Each group then pays a store and a dependent reload: on AVX2 that +/// cost 11% of leaf ranking time. +#[inline(always)] +fn load_group(arch: A, group: &[f32; LANES]) -> A::Vector { + // SAFETY: `group` holds exactly `LANES` readable values, and `A::Vector` has + // `LANES` lanes, so the load reads only `group`. + unsafe { A::Vector::load_simd(arch, group.as_ptr()) } } -/// Insert an eligible candidate in nearest-first order and return the new distance limit. -/// `nearest` contains exactly one non-empty result row. -/// The caller must check the candidate against the current limit before insertion. +/// Insert a candidate into a nearest set in nearest-first order, and return the +/// new k-th distance. +/// +/// The set must not be empty, and the candidate must be nearer than the k-th +/// distance. A candidate goes after existing candidates with an equal distance. #[inline(always)] fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate) -> f32 { let last = nearest.len() - 1; @@ -443,92 +372,64 @@ fn insert_sorted(nearest: &mut [Candidate], candidate: Candidate) -> f32 { #[cfg(test)] mod tests { use super::*; - use diskann_wide::{ARCH, arch::Current}; + use crate::graph::pipnn::test_support::{ArchCheck, for_each_arch}; + use diskann_wide::ARCH; + use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom}; use rstest::rstest; - const EMPTY: Candidate = Candidate::new(UNASSIGNED, f32::INFINITY); - const LANES: usize = <::Vector as SIMDVector>::LANES; + /// Split flat storage into nearest sets of `k` slots each. + fn runtime_neighborhoods(storage: &mut [Candidate], k: usize) -> Vec<&mut [Candidate]> { + storage.chunks_exact_mut(k).collect() + } #[rstest] - #[case::first_batch(0, 3)] - #[case::larger_batch(2, 5)] - #[case::smaller_batch(5, 2)] - #[case::same_size_batch(3, 3)] - #[case::empty_batch(3, 0)] - fn starting_a_batch_clears_results_and_sizes_its_limits( - #[case] previous_rows: usize, - #[case] rows: usize, + #[case::first_scan(0, 3)] + #[case::more_points(2, 5)] + #[case::fewer_points(5, 2)] + #[case::same_points(3, 3)] + #[case::no_points(3, 0)] + fn starting_a_pair_scan_clears_neighborhoods_and_sizes_kth_distances( + #[case] previous_points: usize, + #[case] points: usize, ) { - // Given: reusable storage contains assigned candidates and finite limits. - let mut output = vec![[Candidate::new(0, -12.0), Candidate::new(1, -6.0)]; rows]; - let mut limits = vec![-6.0; previous_rows]; - - // When: prepare storage for a new batch. - Batch::new( - MutMatrixView::try_from(output.as_flattened_mut(), rows, 2).unwrap(), - &mut limits, + // Given: reusable storage contains assigned candidates and finite k-th distances. + let mut output = vec![[Candidate::new(0, -12.0), Candidate::new(1, -6.0)]; points]; + let mut kth_distances = vec![-6.0; previous_points]; + // No pair can enter a nearest set, so every slot must come from the new scan. + let distances = vec![f32::INFINITY; points * points]; + + select_top_k_symmetric( + ARCH, + MatrixView::try_from(distances.as_slice(), points, points).unwrap(), + MutMatrixView::try_from(output.as_flattened_mut(), points, 2).unwrap(), + &mut kth_distances, ); - // Then: no old candidate or limit belongs to the new batch. - assert_eq!(output, vec![[EMPTY; 2]; rows]); - assert_eq!(limits, vec![f32::INFINITY; rows]); - } - - #[test] - fn batch_rankers_borrow_one_complete_row_without_touching_adjacent_rows() { - struct UpdateMiddleRow; - - impl BatchVisitor for UpdateMiddleRow { - fn visit<'a, Rows>(self, mut batch: Batch<'a, Rows>) - where - Batch<'a, Rows>: BatchRanker, - { - assert!(batch.ranker(3).is_none()); - batch - .ranker(1) - .unwrap() - .as_mut() - .fill(Candidate::new(42, -5.0)); - } - } - - // Exercise array rows and runtime rows through the production dispatcher. - for capacity in [2, 4] { - let mut candidates = vec![Candidate::new(9, -12.0); 3 * capacity]; - let mut thresholds = vec![-12.0; 5]; - - with_batch( - MutMatrixView::try_from(candidates.as_mut_slice(), 3, capacity).unwrap(), - &mut thresholds, - UpdateMiddleRow, - ); - - assert_eq!(&candidates[..capacity], vec![EMPTY; capacity]); - assert_eq!( - &candidates[capacity..2 * capacity], - vec![Candidate::new(42, -5.0); capacity] - ); - assert_eq!(&candidates[2 * capacity..], vec![EMPTY; capacity]); - } + // Then: no old candidate or k-th distance belongs to the new scan. + assert_eq!(output, vec![[Candidate::EMPTY; 2]; points]); + assert_eq!(kth_distances, vec![f32::INFINITY; points]); } #[test] - #[should_panic] - fn a_pair_scan_rejects_a_point_outside_the_batch() { - let mut candidates = [EMPTY; 3]; - let mut limits = Vec::new(); - let output = MutMatrixView::try_from(&mut candidates[..], 3, 1).unwrap(); - let mut batch = Batch::new(output, &mut limits); - - batch.update_dual_topk(ARCH, &[3.0, 4.0, 5.0]); + #[should_panic(expected = "one distance row and column per point")] + fn a_symmetric_scan_panics_on_a_distance_matrix_of_another_size() { + let distances = [0.0; 4]; + let mut output = [Candidate::EMPTY; 3]; + + select_top_k_symmetric( + ARCH, + MatrixView::try_from(&distances[..], 2, 2).unwrap(), + MutMatrixView::try_from(&mut output[..], 3, 1).unwrap(), + &mut Vec::new(), + ); } #[rstest] - #[case::empty_input(&[], [EMPTY; 3])] - #[case::no_rankable_values(&[f32::INFINITY, f32::NAN], [EMPTY; 3])] + #[case::empty_input(&[], [Candidate::EMPTY; 3])] + #[case::only_nan_and_infinity(&[f32::INFINITY, f32::NAN], [Candidate::EMPTY; 3])] #[case::fewer_candidates_than_slots( &[11.0, -4.0], - [Candidate::new(1, -4.0), Candidate::new(0, 11.0), EMPTY], + [Candidate::new(1, -4.0), Candidate::new(0, 11.0), Candidate::EMPTY], )] #[case::already_nearest_first( &[-7.0, 2.0, 8.0, 13.0], @@ -538,7 +439,7 @@ mod tests { &[8.0, -7.0, 13.0, 2.0], [Candidate::new(1, -7.0), Candidate::new(3, 2.0), Candidate::new(0, 8.0)], )] - #[case::unrankable_values_leave_gaps_in_ids( + #[case::nan_and_infinity_leave_gaps_in_ids( &[f32::NAN, 6.0, f32::INFINITY, -2.0, 1.0], [Candidate::new(3, -2.0), Candidate::new(4, 1.0), Candidate::new(1, 6.0)], )] @@ -553,83 +454,89 @@ mod tests { Candidate::new(9, -1.0), ]; - (&mut output).select_topk(ARCH, distances); + select_nearest(ARCH, &mut output, distances); assert_eq!(output, expected); } #[test] - fn reusing_output_starts_a_selection_with_a_fresh_distance_limit() { - let mut output = [EMPTY; 2]; - (&mut output).select_topk(ARCH, &[-20.0, -10.0]); + fn a_reused_nearest_set_does_not_keep_its_old_kth_distance() { + let mut output = [Candidate::EMPTY; 2]; + select_nearest(ARCH, &mut output, &[-20.0, -10.0]); assert_eq!(output, [Candidate::new(0, -20.0), Candidate::new(1, -10.0)]); - // Every new distance exceeds the previous selection's limit. - (&mut output).select_topk(ARCH, &[16.0, 3.0, 9.0]); + // Every new distance exceeds the previous selection's k-th distance. + select_nearest(ARCH, &mut output, &[16.0, 3.0, 9.0]); assert_eq!(output, [Candidate::new(1, 3.0), Candidate::new(2, 9.0)]); } #[test] - fn batch_selection_matches_a_full_sort() { - struct SelectRow<'a> { - distances: &'a [f32], - } - - impl TopKVisitor for SelectRow<'_> { - fn visit(self, mut ranker: R) { - ranker.select_topk(ARCH, self.distances); - } - } - - for count in [ - 0, - 1, - LANES - 1, - LANES, - LANES + 1, - 2 * LANES, - 2 * LANES + 3, - 3 * LANES + 2, - ] { - for capacity in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, LANES + 1] { - // Each pair offers a nearer score before a farther one. Later pairs improve - // on earlier pairs, so a full result must keep lowering its cutoff mid-vector. - let mut distances: Vec<_> = (0..count).map(|i| -(i as f32) - 1.0).collect(); - for pair in distances.chunks_exact_mut(2) { - pair.swap(0, 1); + fn id_selection_matches_a_full_sort() { + // Production selects the architecture at run time, so the grid runs on each one. + struct SelectionGrid; + + impl ArchCheck for SelectionGrid { + fn check(&self, arch: A) { + let arch_name = std::any::type_name::(); + // An empty distance row has no matrix form here. The `empty_input` case of + // the selection test covers it. + for count in [ + 1, + LANES - 1, + LANES, + LANES + 1, + 2 * LANES, + 2 * LANES + 3, + 3 * LANES + 2, + ] { + // Each pair offers a nearer distance before a farther one. Later pairs + // improve on earlier pairs, so a full nearest set must keep lowering + // its k-th distance inside a group. + let mut distances: Vec<_> = (0..count).map(|i| -(i as f32) - 1.0).collect(); + for pair in distances.chunks_exact_mut(2) { + pair.swap(0, 1); + } + // k = 1, 2, 3, 8, and 10 use fixed-size nearest sets. The others use + // slices. + for k in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, LANES + 1] { + // Sorting the whole input is independent of the insertion and + // k-th distance logic. + let mut order: Vec<_> = (0..count).collect(); + order.sort_by(|&left, &right| distances[left].total_cmp(&distances[right])); + let mut expected: Vec<_> = order + .into_iter() + .take(k) + .map(|index| index as u32) + .collect(); + expected.resize(k, UNASSIGNED); + let mut output = vec![0; k]; + + select_top_k_ids( + arch, + MatrixView::try_from(distances.as_slice(), 1, count).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), 1, k).unwrap(), + &mut Vec::new(), + ); + + assert_eq!(output, expected, "{arch_name}, count={count}, k={k}"); + } } - - // Sorting the whole input is independent of TopK's insertion and cutoff logic. - let mut expected: Vec<_> = distances - .iter() - .enumerate() - .map(|(index, &distance)| Candidate::new(index as u32, distance)) - .collect(); - expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); - expected.truncate(capacity); - expected.resize(capacity, EMPTY); - let mut output = vec![EMPTY; capacity]; - with_topk( - &mut output, - SelectRow { - distances: &distances, - }, - ); - - assert_eq!(output, expected, "count={count}, capacity={capacity}"); } } + + for_each_arch(&SelectionGrid); } + // The nightly Miri step selects this test by name. Rename it there too. #[test] fn selection_returns_distinct_candidates_when_distances_tie() { let best_index = 2 * LANES + 2; let mut distances = vec![8.0; best_index + 1]; distances[best_index] = -2.0; - let mut output = [EMPTY; 4]; + let mut output = [Candidate::EMPTY; 4]; - (&mut output).select_topk(ARCH, &distances); + select_nearest(ARCH, &mut output, &distances); assert_eq!(output[0], Candidate::new(best_index as u32, -2.0)); // Any three distinct input positions with distance 8 are valid; tie order is free. @@ -651,283 +558,340 @@ mod tests { #[case::lowest_finite(f32::MIN)] #[case::highest_finite(f32::MAX)] #[case::negative_infinity(f32::NEG_INFINITY)] - fn a_rankable_value_keeps_its_input_id_and_float_bits( - #[case] distance: f32, - #[values(0, LANES - 1, LANES, 2 * LANES + 1)] index: usize, - ) { - let mut distances = vec![f32::NAN; 2 * LANES + 2]; - distances[index] = distance; - let mut output = [EMPTY; 2]; + fn a_selected_distance_keeps_its_position_and_bit_pattern(#[case] distance: f32) { + // The SIMD comparison differs by architecture, so each one must keep the bits. + struct FloatBits { + distance: f32, + } + + impl ArchCheck for FloatBits { + fn check(&self, arch: A) { + let arch_name = std::any::type_name::(); + // The first and last lanes of a group, the first lane of the next group, + // and the scalar tail. + for index in [0, LANES - 1, LANES, 2 * LANES + 1] { + let mut distances = vec![f32::NAN; 2 * LANES + 2]; + distances[index] = self.distance; + let mut output = [Candidate::EMPTY; 2]; - output.as_mut_slice().select_topk(ARCH, &distances); + select_nearest(arch, output.as_mut_slice(), &distances); - assert_eq!(output[0].local_idx, index as u32); - assert_eq!(output[0].distance.to_bits(), distance.to_bits()); - assert_eq!(output[1], EMPTY); + let context = format!("{arch_name}, index={index}"); + assert_eq!(output[0].local_idx, index as u32, "{context}"); + assert_eq!( + output[0].distance.to_bits(), + self.distance.to_bits(), + "{context}" + ); + assert_eq!(output[1], Candidate::EMPTY, "{context}"); + } + } + } + + for_each_arch(&FloatBits { distance }); } #[test] fn successive_pair_rows_accumulate_each_points_nearest_neighbors() { - fn check<'a, Rows>(mut batch: Batch<'a, Rows>) - where - Batch<'a, Rows>: BatchRanker, - { + fn check>(neighborhoods: &mut [Nearest]) { + let mut kth_distances = vec![f32::INFINITY; 4]; // Given: pairs 0-1=9, 0-2=2, and 1-2=7; point 3 has no offered pairs. - batch.update_dual_topk(ARCH, &[9.0]); - batch.update_dual_topk(ARCH, &[2.0, 7.0]); + offer_pairs(ARCH, 1, &[9.0], neighborhoods, &mut kth_distances); + offer_pairs(ARCH, 2, &[2.0, 7.0], neighborhoods, &mut kth_distances); let before_last_point = [ [Candidate::new(2, 2.0), Candidate::new(1, 9.0)], [Candidate::new(2, 7.0), Candidate::new(0, 9.0)], [Candidate::new(0, 2.0), Candidate::new(1, 7.0)], - [EMPTY; 2], + [Candidate::EMPTY; 2], ]; - for (index, expected) in before_last_point.iter().enumerate() { - assert_eq!(batch.ranker(index).unwrap().as_ref(), expected); + for (nearest, expected) in neighborhoods.iter_mut().zip(&before_last_point) { + assert_eq!(nearest.as_mut(), expected); } - assert_eq!(batch.thresholds, [9.0, 9.0, 7.0, f32::INFINITY]); + assert_eq!(kth_distances, [9.0, 9.0, 7.0, f32::INFINITY]); // When: point 3 is at distances 6, 1, and 4 from points 0, 1, and 2. - batch.update_dual_topk(ARCH, &[6.0, 1.0, 4.0]); + offer_pairs(ARCH, 3, &[6.0, 1.0, 4.0], neighborhoods, &mut kth_distances); - // Then: each endpoint keeps its own nearest two, with its farthest retained limit. + // Then: each endpoint keeps its own nearest two, and its k-th distance is + // the farther one. let expected = [ [Candidate::new(2, 2.0), Candidate::new(3, 6.0)], [Candidate::new(3, 1.0), Candidate::new(2, 7.0)], [Candidate::new(0, 2.0), Candidate::new(3, 4.0)], [Candidate::new(1, 1.0), Candidate::new(2, 4.0)], ]; - for (index, expected) in expected.iter().enumerate() { - assert_eq!(batch.ranker(index).unwrap().as_ref(), expected); + for (nearest, expected) in neighborhoods.iter_mut().zip(&expected) { + assert_eq!(nearest.as_mut(), expected); } - assert_eq!(batch.thresholds, [6.0, 7.0, 4.0, 4.0]); + assert_eq!(kth_distances, [6.0, 7.0, 4.0, 4.0]); } - let mut fixed = [[EMPTY; 2]; 4]; - let mut fixed_limits = Vec::new(); - fixed_limits.resize(fixed.len(), f32::INFINITY); - check(Batch { - candidates: fixed.as_mut_slice(), - thresholds: &mut fixed_limits, - }); - - let mut dynamic = [EMPTY; 8]; - let mut dynamic_limits = Vec::new(); - let rows = MutMatrixView::try_from(&mut dynamic[..], 4, 2).unwrap(); - check(Batch::new(rows, &mut dynamic_limits)); + check(&mut [[Candidate::EMPTY; 2]; 4]); + let mut storage = [Candidate::EMPTY; 8]; + check(&mut runtime_neighborhoods(&mut storage, 2)); } #[test] - fn pair_scans_match_full_row_sort_across_capacities_and_lengths() { - struct ScanPairs<'a> { - distances: &'a [Vec], - } - - impl BatchVisitor for ScanPairs<'_> { - fn visit<'a, Rows>(self, mut batch: Batch<'a, Rows>) - where - Batch<'a, Rows>: BatchRanker, - { - for (point, row) in self.distances.iter().enumerate().skip(1) { - batch.update_dual_topk(ARCH, &row[..point]); - } - } - } - - for point_count in [ - LANES, - LANES + 1, - LANES + 2, - 2 * LANES, - 2 * LANES + 1, - 2 * LANES + 2, - ] { - for capacity in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17] { - // Given: XOR gives symmetric distances with distinct, exact scores in each row. - // The final pair row has point_count - 1 distances, straddling vector boundaries. - let distances: Vec> = (0..point_count) - .map(|point| { - (0..point_count) - .map(|other| (point ^ other) as f32) - .collect() - }) - .collect(); - let mut output = vec![EMPTY; point_count * capacity]; - let mut limits = Vec::new(); - let mut rows = - MutMatrixView::try_from(output.as_mut_slice(), point_count, capacity).unwrap(); - - // When: offer every non-self pair once, through the production width dispatch. - with_batch( - rows.as_mut_view(), - &mut limits, - ScanPairs { - distances: &distances, - }, - ); - - // Then: independently sort each complete row, excluding only the point itself. - for (point, row) in distances.iter().enumerate() { - let mut expected: Vec<_> = row - .iter() - .enumerate() - .filter(|(other, _)| *other != point) - .map(|(other, &distance)| Candidate::new(other as u32, distance)) + fn pair_scans_match_a_full_sort_for_each_k_and_point_count() { + // Production selects the architecture at run time, so the grid runs on each one. + struct PairScanGrid; + + impl ArchCheck for PairScanGrid { + fn check(&self, arch: A) { + let arch_name = std::any::type_name::(); + for point_count in [ + LANES, + LANES + 1, + LANES + 2, + 2 * LANES, + 2 * LANES + 1, + 2 * LANES + 2, + ] { + // Given: each pair has a distinct integer distance in shuffled order, so + // rows see neighbors in an order unrelated to their IDs and no pairs tie. + // The final pair row has point_count - 1 distances, straddling vector + // boundaries. + let mut pair_distances: Vec = (0..point_count * (point_count - 1) / 2) + .map(|value| value as f32) .collect(); - expected.sort_by(|left, right| left.distance.total_cmp(&right.distance)); - expected.truncate(capacity); - expected.resize(capacity, EMPTY); - - assert_eq!( - rows.row(point), - expected, - "point={point}, point_count={point_count}, capacity={capacity}" - ); - assert_eq!( - limits[point], - expected[capacity - 1].distance, - "limit for point={point}, point_count={point_count}, capacity={capacity}" - ); + pair_distances.shuffle(&mut StdRng::seed_from_u64(point_count as u64)); + // Pair (high, low) with low < high has index high * (high - 1) / 2 + low. + let distances: Vec = (0..point_count * point_count) + .map(|index| { + let (point, other) = (index / point_count, index % point_count); + let (high, low) = (point.max(other), point.min(other)); + if high == low { + 0.0 + } else { + pair_distances[high * (high - 1) / 2 + low] + } + }) + .collect(); + let distances = + MatrixView::try_from(distances.as_slice(), point_count, point_count) + .unwrap(); + // k = 1, 2, and 3 use fixed-size nearest sets. The others use slices. + for k in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 17] { + let mut output = vec![Candidate::EMPTY; point_count * k]; + let mut kth_distances = Vec::new(); + let mut neighborhoods = + MutMatrixView::try_from(output.as_mut_slice(), point_count, k).unwrap(); + + // When: offer every non-self pair once, through the production + // width dispatch. + select_top_k_symmetric( + arch, + distances, + neighborhoods.as_mut_view(), + &mut kth_distances, + ); + + // Then: independently sort each complete row, excluding only the + // point itself. + for (point, &kth_distance) in kth_distances.iter().enumerate() { + let mut expected: Vec<_> = distances + .row(point) + .iter() + .enumerate() + .filter(|(other, _)| *other != point) + .map(|(other, &distance)| Candidate::new(other as u32, distance)) + .collect(); + expected + .sort_by(|left, right| left.distance.total_cmp(&right.distance)); + expected.truncate(k); + expected.resize(k, Candidate::EMPTY); + + assert_eq!( + neighborhoods.row(point), + expected, + "{arch_name}, point={point}, point_count={point_count}, k={k}" + ); + assert_eq!( + kth_distance, + expected[k - 1].distance, + "{arch_name}, k-th distance of point={point}, point_count={point_count}, k={k}" + ); + } + } } } } + + for_each_arch(&PairScanGrid); } - #[rstest] - fn a_source_rejection_does_not_prevent_the_target_from_accepting( - #[values(1, LANES - 1, LANES, 2 * LANES + 1)] target: usize, - ) { + #[test] + fn the_target_accepts_a_pair_that_the_source_rejects() { let source = 2 * LANES + 2; - let mut output = vec![EMPTY; source + 1]; - let mut limits = Vec::new(); - let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); - let mut batch = Batch::new(rows.as_mut_view(), &mut limits); - let mut distances = vec![f32::INFINITY; source]; - distances[0] = -6.0; - distances[target] = 4.0; - - // The source fills its only slot with point 0 before reaching this target. - batch.update_dual_topk(ARCH, &distances); - - let mut expected = vec![EMPTY; source + 1]; - expected[0] = Candidate::new(source as u32, -6.0); - expected[target] = Candidate::new(source as u32, 4.0); - expected[source] = Candidate::new(0, -6.0); - let expected_limits: Vec<_> = expected.iter().map(|c| c.distance).collect(); - assert_eq!(batch.candidates.as_mut_slice(), expected, "target={target}"); - assert_eq!(batch.thresholds, expected_limits); + // Targets sit inside a group, at both edges of a group, and in the scalar tail. + for target in [1, LANES - 1, LANES, 2 * LANES + 1] { + let mut storage = vec![Candidate::EMPTY; source + 1]; + let mut neighborhoods = runtime_neighborhoods(&mut storage, 1); + let mut kth_distances = vec![f32::INFINITY; source + 1]; + let mut distances = vec![f32::INFINITY; source]; + distances[0] = -6.0; + distances[target] = 4.0; + + // The source fills its only slot with point 0 before reaching this target. + offer_pairs( + ARCH, + source, + &distances, + &mut neighborhoods, + &mut kth_distances, + ); + + let mut expected = vec![Candidate::EMPTY; source + 1]; + expected[0] = Candidate::new(source as u32, -6.0); + expected[target] = Candidate::new(source as u32, 4.0); + expected[source] = Candidate::new(0, -6.0); + let expected_kth_distances: Vec<_> = expected.iter().map(|c| c.distance).collect(); + drop(neighborhoods); + assert_eq!(storage, expected, "target={target}"); + assert_eq!(kth_distances, expected_kth_distances, "target={target}"); + } } - #[rstest] - fn a_target_rejection_does_not_prevent_the_source_from_accepting( - #[values(1, LANES - 1, LANES, 2 * LANES + 1)] target: usize, - ) { + #[test] + fn the_source_accepts_a_pair_that_the_target_rejects() { let source = 2 * LANES + 2; - let mut output = vec![EMPTY; source + 1]; - let mut limits = Vec::new(); - let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); - let mut batch = Batch::new(rows.as_mut_view(), &mut limits); - let mut previous_distances = vec![f32::INFINITY; target]; - previous_distances[0] = 1.0; - batch.update_dual_topk(ARCH, &previous_distances); - let mut distances = vec![f32::INFINITY; source]; - distances[target] = 5.0; - - // The target already has a closer neighbor; the source still needs one. - batch.update_dual_topk(ARCH, &distances); - - let mut expected = vec![EMPTY; source + 1]; - expected[0] = Candidate::new(target as u32, 1.0); - expected[target] = Candidate::new(0, 1.0); - expected[source] = Candidate::new(target as u32, 5.0); - let expected_limits: Vec<_> = expected.iter().map(|c| c.distance).collect(); - assert_eq!(batch.candidates.as_mut_slice(), expected, "target={target}"); - assert_eq!(batch.thresholds, expected_limits); + // Targets sit inside a group, at both edges of a group, and in the scalar tail. + for target in [1, LANES - 1, LANES, 2 * LANES + 1] { + let mut storage = vec![Candidate::EMPTY; source + 1]; + let mut neighborhoods = runtime_neighborhoods(&mut storage, 1); + let mut kth_distances = vec![f32::INFINITY; source + 1]; + let mut previous_distances = vec![f32::INFINITY; target]; + previous_distances[0] = 1.0; + offer_pairs( + ARCH, + target, + &previous_distances, + &mut neighborhoods, + &mut kth_distances, + ); + let mut distances = vec![f32::INFINITY; source]; + distances[target] = 5.0; + + // The target already has a closer neighbor; the source still needs one. + offer_pairs( + ARCH, + source, + &distances, + &mut neighborhoods, + &mut kth_distances, + ); + + let mut expected = vec![Candidate::EMPTY; source + 1]; + expected[0] = Candidate::new(target as u32, 1.0); + expected[target] = Candidate::new(0, 1.0); + expected[source] = Candidate::new(target as u32, 5.0); + let expected_kth_distances: Vec<_> = expected.iter().map(|c| c.distance).collect(); + drop(neighborhoods); + assert_eq!(storage, expected, "target={target}"); + assert_eq!(kth_distances, expected_kth_distances, "target={target}"); + } } + // The nightly Miri step selects the `finite` case by its generated name, + // `case_1_finite`. Keep this case first, or change the filter there too. #[rstest] #[case::finite(12.0)] #[case::negative_infinity(f32::NEG_INFINITY)] #[case::negative_zero(-0.0)] - fn a_single_pair_updates_exactly_its_two_endpoints( - #[case] distance: f32, - #[values(LANES - 1, LANES, 2 * LANES + 1)] target: usize, - ) { + fn a_single_pair_updates_exactly_its_two_endpoints(#[case] distance: f32) { let source = 2 * LANES + 2; - let mut output = vec![EMPTY; (source + 1) * 2]; - let mut limits = Vec::new(); - let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 2).unwrap(); - let mut batch = Batch::new(rows.as_mut_view(), &mut limits); - let mut distances = vec![f32::INFINITY; source]; - distances[target] = distance; - - batch.update_dual_topk(ARCH, &distances); - - let mut expected = vec![EMPTY; (source + 1) * 2]; - expected[source * 2] = Candidate::new(target as u32, distance); - expected[target * 2] = Candidate::new(source as u32, distance); - assert_eq!(batch.candidates.as_mut_slice(), expected); - assert_eq!( - batch.candidates.as_mut_slice()[source * 2] - .distance - .to_bits(), - distance.to_bits() - ); - assert_eq!( - batch.candidates.as_mut_slice()[target * 2] - .distance - .to_bits(), - distance.to_bits() - ); - // One neighbor leaves each two-slot result open to another candidate. - assert_eq!(batch.thresholds, vec![f32::INFINITY; source + 1]); + // Targets sit at both edges of a group and in the scalar tail. + for target in [LANES - 1, LANES, 2 * LANES + 1] { + let mut storage = vec![Candidate::EMPTY; (source + 1) * 2]; + let mut neighborhoods = runtime_neighborhoods(&mut storage, 2); + let mut kth_distances = vec![f32::INFINITY; source + 1]; + let mut distances = vec![f32::INFINITY; source]; + distances[target] = distance; + + offer_pairs( + ARCH, + source, + &distances, + &mut neighborhoods, + &mut kth_distances, + ); + + drop(neighborhoods); + let mut expected = vec![Candidate::EMPTY; (source + 1) * 2]; + expected[source * 2] = Candidate::new(target as u32, distance); + expected[target * 2] = Candidate::new(source as u32, distance); + assert_eq!(storage, expected, "target={target}"); + assert_eq!( + storage[source * 2].distance.to_bits(), + distance.to_bits(), + "target={target}" + ); + assert_eq!( + storage[target * 2].distance.to_bits(), + distance.to_bits(), + "target={target}" + ); + // One neighbor leaves each two-slot nearest set open to another candidate. + assert_eq!( + kth_distances, + vec![f32::INFINITY; source + 1], + "target={target}" + ); + } } #[rstest] #[case::no_pairs(&[])] #[case::nan_pairs(&[f32::NAN; 2 * LANES + 1])] #[case::infinite_pairs(&[f32::INFINITY; 2 * LANES + 1])] - fn an_update_without_rankable_pairs_preserves_existing_results(#[case] distances: &[f32]) { - let source = 2 * LANES + 1; - let mut output = vec![EMPTY; source + 1]; - let mut limits = Vec::new(); - let mut rows = MutMatrixView::try_from(output.as_mut_slice(), source + 1, 1).unwrap(); - let mut batch = Batch::new(rows.as_mut_view(), &mut limits); - batch.update_dual_topk(ARCH, &[3.0]); - let previous_output = batch.candidates.as_mut_slice().to_vec(); - let previous_limits = batch.thresholds.to_vec(); - - batch.update_dual_topk(ARCH, distances); - - assert_eq!(batch.candidates.as_mut_slice(), previous_output); - assert_eq!(batch.thresholds, previous_limits); + fn nan_infinite_or_missing_pairs_leave_nearest_sets_unchanged(#[case] distances: &[f32]) { + let source = distances.len(); + let mut storage = vec![Candidate::EMPTY; 2 * LANES + 2]; + let mut neighborhoods = runtime_neighborhoods(&mut storage, 1); + let mut kth_distances = vec![f32::INFINITY; 2 * LANES + 2]; + offer_pairs(ARCH, 1, &[3.0], &mut neighborhoods, &mut kth_distances); + let previous_output: Vec<_> = neighborhoods + .iter() + .map(|nearest| nearest.to_vec()) + .collect(); + let previous_kth_distances = kth_distances.clone(); + + // An empty row makes point 0 the source; it has no earlier points. + offer_pairs( + ARCH, + source, + distances, + &mut neighborhoods, + &mut kth_distances, + ); + + let output: Vec<_> = neighborhoods + .iter() + .map(|nearest| nearest.to_vec()) + .collect(); + assert_eq!(output, previous_output); + assert_eq!(kth_distances, previous_kth_distances); } #[test] - fn zero_capacity_pair_scans_leave_no_neighbors_or_finite_limits() { - fn check<'a, Rows>(mut batch: Batch<'a, Rows>) - where - Batch<'a, Rows>: BatchRanker, - { - batch.update_dual_topk(ARCH, &[4.0]); - batch.update_dual_topk(ARCH, &[2.0, 7.0]); - - for index in 0..batch.nrows() { - assert!(batch.ranker(index).unwrap().as_ref().is_empty()); - } - assert_eq!(batch.thresholds, [f32::INFINITY; 3]); - } + fn zero_width_outputs_select_nothing() { + let distances = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]; + let mut kth_distances = vec![-1.0; 2]; + + select_top_k_symmetric( + ARCH, + MatrixView::try_from(&distances[..], 3, 3).unwrap(), + MutMatrixView::try_from(&mut [][..], 3, 0).unwrap(), + &mut kth_distances, + ); + select_top_k_ids( + ARCH, + MatrixView::try_from(&distances[..], 3, 3).unwrap(), + MutMatrixView::try_from(&mut [][..], 3, 0).unwrap(), + &mut Vec::new(), + ); - let mut fixed: [[Candidate; 0]; 3] = [[]; 3]; - let mut fixed_limits = Vec::new(); - fixed_limits.resize(fixed.len(), f32::INFINITY); - check(Batch { - candidates: fixed.as_mut_slice(), - thresholds: &mut fixed_limits, - }); - - let mut dynamic = []; - let mut dynamic_limits = Vec::new(); - let rows = MutMatrixView::try_from(&mut dynamic[..], 3, 0).unwrap(); - check(Batch::new(rows, &mut dynamic_limits)); + // A scan without slots leaves the reusable k-th distances untouched. + assert_eq!(kth_distances, [-1.0; 2]); } } From cc6da38b9c072e3ec917fc986f80649b93891e06 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 23 Sep 2026 15:09:30 +0000 Subject: [PATCH 206/229] refactor(pipnn): check top-k shapes only in debug builds The kernel entry points already return an error for a bad output shape, so the top-k functions now check their shape preconditions with debug assertions only. Select distance rows by index. A distance matrix without columns then leaves every slot UNASSIGNED instead of panicking in row_iter. --- diskann/src/graph/pipnn/mod.rs | 4 +-- diskann/src/graph/pipnn/topk.rs | 50 ++++++++++++++++++++------------- 2 files changed, 32 insertions(+), 22 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index f03932c611..882788dc79 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -59,8 +59,8 @@ fn cosine_distance(dot: f32, source_norm: f32, target_norm: f32) -> f32 { /// Return an error if a kernel output does not have one row per input point. /// -/// The top-k functions panic on this mismatch. The kernel entry points call this -/// check first, so a bad output shape is an error instead of a panic. +/// The top-k functions check this shape only in debug builds. The kernel entry +/// points check it first, so a bad output shape is an error in every build. fn check_output_rows(points: usize, rows: usize) -> ANNResult<()> { if rows == points { Ok(()) diff --git a/diskann/src/graph/pipnn/topk.rs b/diskann/src/graph/pipnn/topk.rs index f8f9300074..b2d1d6d715 100644 --- a/diskann/src/graph/pipnn/topk.rs +++ b/diskann/src/graph/pipnn/topk.rs @@ -68,22 +68,18 @@ impl Default for Candidate { /// Write the IDs of the k nearest columns of each distance row, nearest first. /// -/// k is the column count of `output`, and `output` has one row per distance row. -/// Equal distances can select either column. If a row has fewer than k distances -/// that are not NaN or positive infinity, the remaining slots hold [`UNASSIGNED`]. -/// `scratch` is reusable storage for values of k without a fixed-size nearest set. -/// -/// # Panics -/// -/// Panics if `output` does not have one row per distance row, or if `distances` -/// has no columns and k is not zero. +/// k is the column count of `output`, and `output` must have one row per distance +/// row. Debug builds check this. Equal distances can select either column. If a +/// row has fewer than k distances that are not NaN or positive infinity, the +/// remaining slots hold [`UNASSIGNED`]. `scratch` is reusable storage for values +/// of k without a fixed-size nearest set. pub(super) fn select_top_k_ids( arch: A, distances: MatrixView<'_, f32>, output: MutMatrixView<'_, u32>, scratch: &mut Vec, ) { - assert_eq!( + debug_assert_eq!( distances.nrows(), output.nrows(), "top-k IDs need one output row per distance row" @@ -108,17 +104,14 @@ pub(super) fn select_top_k_ids( /// Fill the nearest set of each point with its k nearest other points. /// -/// `distances` is a symmetric matrix with one row and one column per point, and -/// `output` has one row per point. k is the column count of `output`. Only the +/// `distances` must be a symmetric matrix with one row and one column per point, +/// and `output` has one row per point. Debug builds check the shape. k is the +/// column count of `output`. Only the /// strict lower triangle of `distances` is read: each pair is read once and /// offered to both points. Candidate IDs are point positions. Equal distances can /// select either point. Slots that no pair fills hold [`Candidate::EMPTY`]. /// `kth_distances` is reusable storage. It holds the k-th distance of each point /// after the scan. -/// -/// # Panics -/// -/// Panics if `distances` does not have one row and one column per output row. pub(super) fn select_top_k_symmetric( arch: A, distances: MatrixView<'_, f32>, @@ -126,7 +119,7 @@ pub(super) fn select_top_k_symmetric( kth_distances: &mut Vec, ) { let points = output.nrows(); - assert!( + debug_assert!( distances.nrows() == points && distances.ncols() == points, "a symmetric scan needs one distance row and column per point" ); @@ -157,8 +150,10 @@ fn select_ids_with( A: Simd, Nearest: AsMut<[Candidate]> + ?Sized, { - for (row, ids) in distances.row_iter().zip(output.row_iter_mut()) { - select_nearest(arch, nearest, row); + // `row_iter` panics on a matrix without columns. `row` returns an empty row, + // so every slot of that row becomes `UNASSIGNED`. + for (row, ids) in output.row_iter_mut().enumerate() { + select_nearest(arch, nearest, distances.row(row)); for (id, candidate) in ids.iter_mut().zip(nearest.as_mut().iter()) { *id = candidate.local_idx; } @@ -411,8 +406,9 @@ mod tests { } #[test] + #[cfg(debug_assertions)] #[should_panic(expected = "one distance row and column per point")] - fn a_symmetric_scan_panics_on_a_distance_matrix_of_another_size() { + fn a_symmetric_scan_debug_checks_the_distance_matrix_size() { let distances = [0.0; 4]; let mut output = [Candidate::EMPTY; 3]; @@ -894,4 +890,18 @@ mod tests { // A scan without slots leaves the reusable k-th distances untouched. assert_eq!(kth_distances, [-1.0; 2]); } + + #[test] + fn distance_rows_without_columns_leave_every_slot_unassigned() { + let mut output = [0; 4]; + + select_top_k_ids( + ARCH, + MatrixView::try_from(&[][..], 2, 0).unwrap(), + MutMatrixView::try_from(&mut output[..], 2, 2).unwrap(), + &mut Vec::new(), + ); + + assert_eq!(output, [UNASSIGNED; 4]); + } } From a990ab600981e750ea930943a875ac2ccce63c1e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 28 Sep 2026 04:57:36 +0000 Subject: [PATCH 207/229] fix(pipnn): finish partitioning on tied leaders; drop unreachable errors Partitioning: - Replace the 30-level cap and `PartitionError::IterationLimit` with a progress rule. When a fanout-one split sends every point to one leader, the copies of that leader become equal-size leaves and the rest is re-split. Duplicate vectors no longer fail the build. - Remove `PartitionError::ShapeOverflow` and `checked_area`. Split IDs are distinct dataset rows, so buffer products cannot overflow. - Lease stripe buffers per stripe instead of per worker chunk, and drop the power-of-two stripe rounding and the unreachable 32-point minimum. BigANN-10M fp16 partition, 16 threads, 4 interleaved runs: medians base 51.7 s, per-stripe 52.2 s, no rounding 52.1 s; identical leaves. - Keep the parallel scatter: 10M points at fanout 10 scatter in 0.59 s serially and 0.15 s in parallel; the serial share grows with threads. - Document the reason and measurement for each tuning constant. Leaf build and conversion: - Remove `LeafBuildError`; its shape, conversion and kernel variants are unreachable. Errors propagate as `ANNError`. - Keep the FP16 gather dispatch and document why: `as_f32_into` uses the build target's instructions, and dependents of `diskann` do not inherit the repository's x86-64-v3 flag. Baseline builds gather 4.6-22x slower without it; x86-64-v3 builds are within 10%. Tests: - Replace rstest cases with table loops; PR3 test instances 117 -> 45. - Add seam oracles at production shape: partition assignment against brute force, and leaf candidates against the leaf kernel per leaf. - Add a multi-level build check that 1 and 4 workers give the same graph for every metric. The previous suite passed with unsorted candidates. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/conversion.rs | 121 ++-- diskann/src/graph/pipnn/finalization.rs | 195 +++---- diskann/src/graph/pipnn/leaf_build.rs | 380 +++++------- diskann/src/graph/pipnn/mod.rs | 583 ++++++++++--------- diskann/src/graph/pipnn/partitioning.rs | 739 +++++++++++++----------- 5 files changed, 990 insertions(+), 1028 deletions(-) diff --git a/diskann/src/graph/pipnn/conversion.rs b/diskann/src/graph/pipnn/conversion.rs index cca4ef42fb..581010e3ca 100644 --- a/diskann/src/graph/pipnn/conversion.rs +++ b/diskann/src/graph/pipnn/conversion.rs @@ -3,11 +3,19 @@ * Licensed under the MIT license. */ -//! Convert gathered vectors with the available CPU's FP16 instructions. +//! Gather dataset rows into packed `f32` storage. +//! +//! `f16` rows use the FP16 instructions of the running CPU. The generic +//! `VectorRepr::as_f32_into` path uses the instructions of the build target +//! instead. Crates that depend on `diskann` do not inherit this repository's +//! `x86-64-v3` target flag, so for them that path converts `f16` in software. +//! On an AMD EPYC 7763, software conversion made gathers 4.6 to 22 times slower +//! (128 to 1536 dimensions). With the `x86-64-v3` target, both paths are within +//! 10% of each other. use std::any::TypeId; -use crate::utils::VectorRepr; +use crate::{ANNResult, utils::VectorRepr}; use diskann_utils::views::MatrixView; use diskann_vector::conversion::SliceCast; use diskann_wide::{ @@ -16,17 +24,19 @@ use diskann_wide::{ }; use half::f16; -/// Gather rows into packed f32 storage, dispatching once for the whole FP16 batch. +/// Convert rows `ids` of `data` into consecutive rows of `output`. +/// +/// `f16` input selects the CPU's instructions once for the whole batch. pub(super) fn gather_as_f32( data: MatrixView<'_, T>, ids: &[u32], output: &mut [f32], -) -> Result<(), (u32, T::Error)> { +) -> ANNResult<()> { if TypeId::of::() == TypeId::of::() { arch::dispatch3(GatherFp16, data, ids, output); } else { for (&id, destination) in ids.iter().zip(output.chunks_exact_mut(data.ncols())) { - T::as_f32_into(data.row(id as usize), destination).map_err(|error| (id, error))?; + T::as_f32_into(data.row(id as usize), destination).map_err(Into::into)?; } } Ok(()) @@ -52,59 +62,62 @@ where #[cfg(test)] mod tests { use super::*; - use rstest::rstest; - - #[rstest] - #[case::f32([1.5_f32, -2.0, 7.0, 4.0, 0.5, 9.0], [0.5, 9.0, 1.5, -2.0, 0.5, 9.0])] - #[case::i8([1_i8, -2, 7, 4, -5, 9], [-5.0, 9.0, 1.0, -2.0, -5.0, 9.0])] - #[case::u8([1_u8, 2, 7, 4, 5, 9], [5.0, 9.0, 1.0, 2.0, 5.0, 9.0])] - #[case::f16([1.5_f32, -2.0, 7.0, 4.0, 0.5, 9.0].map(f16::from_f32), [0.5, 9.0, 1.5, -2.0, 0.5, 9.0])] - fn gathered_rows_follow_requested_ids( - #[case] values: [T; 6], - #[case] expected: [f32; 6], - ) { - // Given: row IDs are deliberately out of order and contain a repeat. - let data = MatrixView::try_from(&values[..], 3, 2).unwrap(); - let mut output = [f32::NAN; 6]; - - gather_as_f32(data, &[2, 0, 2], &mut output).unwrap(); - - assert_eq!(output, expected); - } - #[rstest] - #[case::before_vector(7)] - #[case::complete_vector(8)] - #[case::after_vector(9)] - #[case::multiple_vectors(32)] - #[case::multiple_vectors_and_tail(33)] - #[case::embedding_and_tail(1537)] - fn fp16_gather_converts_every_coordinate(#[case] dimensions: usize) { - let mut values: Vec<_> = (0..3 * dimensions) - .map(|index| f16::from_f32((index % 31) as f32 - 15.0)) - .collect(); - for (row, values) in values.chunks_exact_mut(dimensions).enumerate() { - values[dimensions - 1] = f16::from_f32(100.0 + row as f32); + #[test] + fn gathered_rows_follow_requested_ids() { + // Row IDs are out of order and contain a repeat. + fn check(values: [T; 6], expected: [f32; 6]) { + let data = MatrixView::try_from(&values[..], 3, 2).unwrap(); + let mut output = [f32::NAN; 6]; + + gather_as_f32(data, &[2, 0, 2], &mut output).unwrap(); + + assert_eq!(output, expected, "{}", std::any::type_name::()); } - let ids = [2, 0, 2]; - let expected: Vec<_> = ids - .iter() - .flat_map(|&row| { - values[row as usize * dimensions..(row as usize + 1) * dimensions] - .iter() - .map(|value| value.to_f32()) - }) - .collect(); - let mut output = vec![f32::NAN; expected.len()]; - gather_as_f32( - MatrixView::try_from(values.as_slice(), 3, dimensions).unwrap(), - &ids, - &mut output, - ) - .unwrap(); + check( + [1.5_f32, -2.0, 7.0, 4.0, 0.5, 9.0], + [0.5, 9.0, 1.5, -2.0, 0.5, 9.0], + ); + check([1_i8, -2, 7, 4, -5, 9], [-5.0, 9.0, 1.0, -2.0, -5.0, 9.0]); + check([1_u8, 2, 7, 4, 5, 9], [5.0, 9.0, 1.0, 2.0, 5.0, 9.0]); + check( + [1.5_f32, -2.0, 7.0, 4.0, 0.5, 9.0].map(f16::from_f32), + [0.5, 9.0, 1.5, -2.0, 0.5, 9.0], + ); + } - assert_eq!(output, expected); + #[test] + fn fp16_gather_converts_every_coordinate() { + // x86 converts eight values at a time. The dimensions cover a short row, + // full vectors with and without a tail, and an embedding with a tail. + for dimensions in [7, 8, 9, 32, 33, 1537] { + let mut values: Vec<_> = (0..3 * dimensions) + .map(|index| f16::from_f32((index % 31) as f32 - 15.0)) + .collect(); + for (row, values) in values.chunks_exact_mut(dimensions).enumerate() { + values[dimensions - 1] = f16::from_f32(100.0 + row as f32); + } + let ids = [2, 0, 2]; + let expected: Vec<_> = ids + .iter() + .flat_map(|&row| { + values[row as usize * dimensions..(row as usize + 1) * dimensions] + .iter() + .map(|value| value.to_f32()) + }) + .collect(); + let mut output = vec![f32::NAN; expected.len()]; + + gather_as_f32( + MatrixView::try_from(values.as_slice(), 3, dimensions).unwrap(), + &ids, + &mut output, + ) + .unwrap(); + + assert_eq!(output, expected, "{dimensions} dimensions"); + } } #[test] diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 61a0df88ed..03d18c161b 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -125,8 +125,7 @@ pub(crate) fn prune_overfull( #[cfg(test)] mod tests { use super::*; - use crate::graph::config; - use rstest::rstest; + use crate::graph::{config, pipnn::test_support::thread_pool}; fn pruning_config( degree: usize, @@ -145,87 +144,88 @@ mod tests { candidates } - #[rstest] - #[case::empty(&[])] - #[case::below_degree(&[2])] - #[case::at_degree(&[3, 1])] - fn candidates_within_degree_keep_their_order(#[case] ids: &[u32]) { - let values = [0.0_f32, 1.0, 3.0, -2.0]; - let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let graph = pruning_config(2, Metric::L2, 1.0).unwrap(); + /// Prune the candidates of source 0 in one-dimensional L2 data. + fn prune_l2(values: &[f32], ids: &[u32], degree: usize, alpha: f32) -> Vec { + let data = MatrixView::try_from(values, values.len(), 1).unwrap(); + let graph = pruning_config(degree, Metric::L2, alpha).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); - - let actual = PruneWorkspace::default() + let pruned = PruneWorkspace::default() .prune(data, 0, candidates_in_order(ids), &graph, &distance) .unwrap(); + pruned.to_vec() + } - assert_eq!(&*actual, ids); + #[test] + fn candidates_within_degree_keep_their_order() { + let values = [0.0_f32, 1.0, 3.0, -2.0]; + + for ids in [&[][..], &[2][..], &[3, 1][..]] { + assert_eq!(prune_l2(&values, ids, 2, 1.0), ids); + } } - #[rstest] - #[case::l2(Metric::L2, [0.0, 0.0, 1.0, 0.0, 3.0, 0.0, -2.0, 0.0], [1, 3])] - #[case::cosine(Metric::Cosine, [1.0, 0.0, 2.0, 0.0, 1.0, 1.0, -1.0, 0.0], [1, 2])] - #[case::normalized_cosine(Metric::CosineNormalized, [1.0, 0.0, 0.8, 0.6, -1.0, 0.0, 0.0, -1.0], [1, 3])] - #[case::inner_product(Metric::InnerProduct, [1.0, 0.0, 3.0, 0.0, 2.0, 1.0, -1.0, 4.0], [1, 3])] - fn overfull_candidates_follow_the_metric_pruning_rule( - #[case] metric: Metric, - #[case] values: [f32; 8], - #[case] expected: [u32; 2], - ) { - // ID 1 is nearest. L2 occludes ID 2: 9 / 4 > 1. + #[test] + fn overfull_candidates_follow_the_metric_pruning_rule() { + // Source 0 has candidates [3, 2, 1]; ID 1 is nearest for every metric. + // L2 occludes ID 2: 9 / 4 > 1. // Cosine retains ID 2: equal directions give an occlusion ratio of 1. // Normalized cosine retains ID 3: 1 / 1.6 < 1. // Inner product occludes ID 2: its dot with ID 1 is 6, versus 2 with the source. - let data = MatrixView::try_from(&values[..], 4, 2).unwrap(); - let graph = pruning_config(2, metric, 1.0).unwrap(); - let distance = f32::distance(metric, Some(2)); + for (metric, values, expected) in [ + ( + Metric::L2, + [0.0, 0.0, 1.0, 0.0, 3.0, 0.0, -2.0, 0.0], + [1, 3], + ), + ( + Metric::Cosine, + [1.0, 0.0, 2.0, 0.0, 1.0, 1.0, -1.0, 0.0], + [1, 2], + ), + ( + Metric::CosineNormalized, + [1.0, 0.0, 0.8, 0.6, -1.0, 0.0, 0.0, -1.0], + [1, 3], + ), + ( + Metric::InnerProduct, + [1.0, 0.0, 3.0, 0.0, 2.0, 1.0, -1.0, 4.0], + [1, 3], + ), + ] { + let data = MatrixView::try_from(&values[..], 4, 2).unwrap(); + let graph = pruning_config(2, metric, 1.0).unwrap(); + let distance = f32::distance(metric, Some(2)); - let actual = PruneWorkspace::default() - .prune(data, 0, candidates_in_order(&[3, 2, 1]), &graph, &distance) - .unwrap(); + let actual = PruneWorkspace::default() + .prune(data, 0, candidates_in_order(&[3, 2, 1]), &graph, &distance) + .unwrap(); - assert_eq!(&*actual, expected); + assert_eq!(&*actual, expected, "{metric:?}"); + } } #[test] fn pruning_excludes_the_source_without_shifting_selected_ids() { - let values = [0.0_f32, 1.0, 3.0, -2.0]; - let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let graph = pruning_config(2, Metric::L2, 1.0).unwrap(); - let distance = f32::distance(Metric::L2, Some(1)); - - let actual = PruneWorkspace::default() - .prune( - data, - 0, - candidates_in_order(&[3, 0, 2, 1]), - &graph, - &distance, - ) - .unwrap(); - - assert_eq!(&*actual, [1, 3]); + assert_eq!( + prune_l2(&[0.0, 1.0, 3.0, -2.0], &[3, 0, 2, 1], 2, 1.0), + [1, 3] + ); } - #[rstest] - #[case::strict(1.0, &[1])] - #[case::relaxed(2.0, &[1, 3])] - fn graph_alpha_controls_which_occluded_neighbors_return( - #[case] alpha: f32, - #[case] expected: &[u32], - ) { + #[test] + fn graph_alpha_controls_which_occluded_neighbors_return() { // After selecting x=1, x=3 has ratio 9/4 and x=4 has ratio 16/9. // Alpha 2 admits x=4 while x=3 remains occluded. let values = [0.0_f32, 1.0, 3.0, 4.0]; - let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let graph = pruning_config(2, Metric::L2, alpha).unwrap(); - let distance = f32::distance(Metric::L2, Some(1)); - - let actual = PruneWorkspace::default() - .prune(data, 0, candidates_in_order(&[3, 2, 1]), &graph, &distance) - .unwrap(); - assert_eq!(&*actual, expected); + for (alpha, expected) in [(1.0, &[1][..]), (2.0, &[1, 3][..])] { + assert_eq!( + prune_l2(&values, &[3, 2, 1], 2, alpha), + expected, + "alpha {alpha}" + ); + } } #[test] @@ -252,64 +252,49 @@ mod tests { #[cfg(not(miri))] #[test] - fn the_largest_supported_candidate_list_can_be_pruned() { - let count = u16::MAX as usize; - let values: Vec<_> = (0..=count).map(|i| i as f32).collect(); - let data = MatrixView::try_from(values.as_slice(), count + 1, 1).unwrap(); - let candidates = AdjacencyList::from_iter_untrusted(1..=count as u32); + fn the_candidate_position_limit_accepts_u16_max_and_rejects_one_more() { + let values: Vec<_> = (0..=u16::MAX as usize + 1).map(|i| i as f32).collect(); + let data = MatrixView::try_from(values.as_slice(), values.len(), 1).unwrap(); let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); + let largest = AdjacencyList::from_iter_untrusted(1..=u16::MAX as u32); let actual = PruneWorkspace::default() - .prune(data, 0, candidates, &graph, &distance) + .prune(data, 0, largest, &graph, &distance) .unwrap(); - assert_eq!(&*actual, [1]); - } - - #[cfg(not(miri))] - #[test] - fn a_candidate_list_beyond_the_position_limit_returns_its_size() { - let count = u16::MAX as usize + 1; - let values: Vec<_> = (0..=count).map(|i| i as f32).collect(); - let data = MatrixView::try_from(values.as_slice(), count + 1, 1).unwrap(); - let candidates = AdjacencyList::from_iter_untrusted(1..=count as u32); - let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); - let distance = f32::distance(Metric::L2, Some(1)); + let too_many = AdjacencyList::from_iter_untrusted(1..=u16::MAX as u32 + 1); let error = PruneWorkspace::default() - .prune(data, 0, candidates, &graph, &distance) + .prune(data, 0, too_many, &graph, &distance) .unwrap_err(); - - let error = error.downcast_ref::().unwrap(); - assert!( - matches!(error, FinalizationError::TooManyCandidates { actual, max } - if *actual == count && *max == u16::MAX as usize) - ); + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::TooManyCandidates { actual, max }) + if *actual == u16::MAX as usize + 1 && *max == u16::MAX as usize + )); } - #[rstest] - #[case::one_worker(1)] - #[case::several_workers(3)] - fn parallel_pruning_keeps_each_result_with_its_source(#[case] workers: usize) { - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(workers) - .build() - .unwrap(); + #[test] + fn parallel_pruning_keeps_each_result_with_its_source() { let values = [0.0_f32, 1.0, 4.0, 9.0]; let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let candidates = [[3, 2, 1], [3, 2, 0], [3, 1, 0], [2, 1, 0]] - .map(|ids| candidates_in_order(&ids)) - .to_vec(); let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); - let actual = pool - .install(|| prune_overfull(data, candidates, &graph, Metric::L2)) - .unwrap(); + for workers in [1, 3] { + let candidates = [[3, 2, 1], [3, 2, 0], [3, 1, 0], [2, 1, 0]] + .map(|ids| candidates_in_order(&ids)) + .to_vec(); - assert_eq!( - actual.into_iter().map(Vec::from).collect::>(), - [vec![1], vec![0], vec![1], vec![2]] - ); + let actual = thread_pool(workers) + .install(|| prune_overfull(data, candidates, &graph, Metric::L2)) + .unwrap(); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [vec![1], vec![0], vec![1], vec![2]], + "{workers} workers" + ); + } } } diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 2c28c8c93b..5604114db1 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -19,41 +19,18 @@ use parking_lot::Mutex; -use crate::{graph::AdjacencyList, utils::VectorRepr}; +use crate::{ANNResult, graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; use rayon::prelude::*; use super::{ + conversion::gather_as_f32, leaf_kernel::{LeafKernelWorkspace, select_leaf_neighbors}, leaf_metric::LeafMetric, simd::Simd, topk::Candidate, }; -/// Failure while converting leaves into direct graph candidates. -#[derive(Debug, thiserror::Error)] -pub(crate) enum LeafBuildError { - #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] - ShapeOverflow { - leaf: usize, - rows: usize, - columns: usize, - }, - #[error("failed to convert point {point} in leaf {leaf}")] - Conversion { - leaf: usize, - point: u32, - #[source] - source: crate::ANNError, - }, - #[error("nearest-neighbor selection failed for leaf {leaf}")] - Kernel { - leaf: usize, - #[source] - source: crate::ANNError, - }, -} - /// Reusable buffers for one Rayon leaf job. /// /// The numerical vectors keep the largest leaf shape that this job observed. @@ -67,36 +44,18 @@ struct LeafBuffers { } impl LeafBuffers { - fn prepare( - &mut self, - leaf: usize, - point_count: usize, - dimension_count: usize, - requested_k: usize, - ) -> Result<(usize, usize), LeafBuildError> { - let point_value_count = - point_count - .checked_mul(dimension_count) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: point_count, - columns: dimension_count, - })?; + /// Grow the buffers for one leaf and return its effective `k`. + fn prepare(&mut self, point_count: usize, dimension_count: usize, requested_k: usize) -> usize { // A point has at most `point_count - 1` other points in its leaf. Wider rows // would hold only empty slots. let leaf_k = requested_k.min(point_count.saturating_sub(1)); - let neighbor_count = - point_count - .checked_mul(leaf_k) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: point_count, - columns: leaf_k, - })?; - - grow(&mut self.point_values, point_value_count, 0.0); - grow(&mut self.neighbors, neighbor_count, Candidate::default()); - Ok((leaf_k, neighbor_count)) + grow(&mut self.point_values, point_count * dimension_count, 0.0); + grow( + &mut self.neighbors, + point_count * leaf_k, + Candidate::default(), + ); + leaf_k } fn prepare_local_adjacency(&mut self, point_count: usize) { @@ -118,7 +77,7 @@ pub(super) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, requested_k: usize, -) -> Result>, LeafBuildError> +) -> ANNResult>> where A: Simd, M: LeafMetric, @@ -127,20 +86,18 @@ where let candidates: Vec<_> = (0..data.nrows()) .map(|_| Mutex::new(AdjacencyList::new())) .collect(); - leaves.par_iter().enumerate().try_for_each_init( - LeafBuffers::default, - |buffers, (leaf, point_ids)| { + leaves + .par_iter() + .try_for_each_init(LeafBuffers::default, |buffers, point_ids| { add_direct_leaf_candidates::( arch, data, - leaf, point_ids, requested_k, buffers, &candidates, ) - }, - )?; + })?; Ok(candidates .into_iter() .map(Mutex::into_inner) @@ -154,59 +111,44 @@ where /// Add one leaf's symmetric neighbors to the direct candidate lists. /// /// Reusable buffers can be longer than this leaf, so all accesses use the current -/// leaf shape. -#[expect( - clippy::expect_used, - reason = "buffer prefixes have the checked leaf shape" -)] +/// leaf shape. Leaf IDs are distinct dataset rows, so `points x dimensions` is no +/// larger than the dataset. fn add_direct_leaf_candidates( arch: A, data: MatrixView<'_, T>, - leaf: usize, point_ids: &[u32], requested_k: usize, buffers: &mut LeafBuffers, candidates: &[Mutex>], -) -> Result<(), LeafBuildError> +) -> ANNResult<()> where A: Simd, M: LeafMetric, T: VectorRepr, { - let (leaf_k, neighbor_value_count) = - buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; + let point_count = point_ids.len(); + let leaf_k = buffers.prepare(point_count, data.ncols(), requested_k); if leaf_k == 0 { return Ok(()); } - let point_value_count = point_ids.len() * data.ncols(); - let point_values = &mut buffers.point_values[..point_value_count]; - - super::conversion::gather_as_f32(data, point_ids, point_values).map_err( - |(point, source)| LeafBuildError::Conversion { - leaf, - point, - source: source.into(), - }, - )?; - - let points = MatrixView::try_from(&*point_values, point_ids.len(), data.ncols()) - .expect("point buffer prefix has the checked leaf shape"); + let point_values = &mut buffers.point_values[..point_count * data.ncols()]; + gather_as_f32(data, point_ids, point_values)?; + let points = MatrixView::try_from(&*point_values, point_count, data.ncols())?; + let neighbor_count = point_count * leaf_k; let output = MutMatrixView::try_from( - &mut buffers.neighbors[..neighbor_value_count], - point_ids.len(), + &mut buffers.neighbors[..neighbor_count], + point_count, leaf_k, - ) - .expect("neighbor buffer prefix has the checked leaf shape"); - select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + )?; + select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace)?; - buffers.prepare_local_adjacency(point_ids.len()); + buffers.prepare_local_adjacency(point_count); add_symmetric_neighbors( point_ids, leaf_k, - &buffers.neighbors[..neighbor_value_count], - &mut buffers.local_adjacency[..point_ids.len()], + &buffers.neighbors[..neighbor_count], + &mut buffers.local_adjacency[..point_count], ); for (&point_id, additions) in point_ids.iter().zip(&buffers.local_adjacency) { candidates[point_id as usize] @@ -246,97 +188,78 @@ fn grow(values: &mut Vec, len: usize, value: T) { #[cfg(all(test, not(miri)))] mod tests { use super::*; - use crate::graph::pipnn::{L2, test_support}; + use crate::graph::pipnn::{ + L2, + test_support::{dense_points, sorted_members_per_row, thread_pool}, + }; use diskann_wide::ARCH; - use rstest::rstest; + use std::collections::BTreeSet; + + #[expect(clippy::unwrap_used, reason = "an invalid fixture fails the test")] + fn build(values: &[f32], leaves: Vec>, k: usize, workers: usize) -> Vec> { + let data = MatrixView::try_from(values, values.len(), 1).unwrap(); + let candidates = thread_pool(workers) + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves, k)) + .unwrap(); + sorted_members_per_row(&candidates.into_iter().map(Vec::from).collect::>()) + } #[test] fn selected_neighbors_use_global_ids_and_contribute_both_edge_directions() { // The leaf lists IDs [5, 1, 3], at coordinates [9, 0, 2]. // The directed choices are 1 -> 3, 3 -> 1 and 5 -> 3. let values = [100.0_f32, 0.0, -100.0, 2.0, 200.0, 9.0]; - let data = MatrixView::try_from(&values[..], 6, 1).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - let actual = pool - .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, vec![vec![5, 1, 3]], 1)) - .unwrap(); - let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + let actual = build(&values, vec![vec![5, 1, 3]], 1, 1); assert_eq!( - test_support::sorted_members_per_row(&actual), + actual, [vec![], vec![3], vec![], vec![1, 5], vec![], vec![3]] ); } - #[rstest] - #[case::one_worker(1)] - #[case::several_workers(3)] - fn overlapping_leaves_merge_each_neighbor_once(#[case] workers: usize) { + #[test] + fn overlapping_leaves_merge_each_neighbor_once() { let values = [0.0_f32, 1.0, 4.0, 9.0, 16.0]; - let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); let leaves = vec![vec![0, 1, 2], vec![1, 2, 3], vec![2, 3, 4], vec![0, 1, 2]]; - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(workers) - .build() - .unwrap(); - let actual = pool - .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves, 1)) - .unwrap(); - let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + for workers in [1, 3] { + let actual = build(&values, leaves.clone(), 1, workers); - assert_eq!( - test_support::sorted_members_per_row(&actual), - [vec![1], vec![0, 2], vec![1, 3], vec![2, 4], vec![3]] - ); + assert_eq!( + actual, + [vec![1], vec![0, 2], vec![1, 3], vec![2, 4], vec![3]], + "{workers} workers" + ); + } } - #[rstest] - #[case::no_leaves(vec![], 2)] - #[case::empty_leaf(vec![vec![]], 2)] - #[case::singleton(vec![vec![2]], 2)] - #[case::zero_neighbors(vec![vec![0, 1, 2]], 0)] - fn leaves_without_selected_pairs_produce_empty_adjacency( - #[case] leaves: Vec>, - #[case] requested_k: usize, - ) { + #[test] + fn leaves_without_selected_pairs_produce_empty_adjacency() { let values = [0.0_f32, 1.0, 4.0]; - let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - let actual = pool - .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves, requested_k)) - .unwrap(); + for (leaves, k) in [ + (vec![], 2), + (vec![vec![]], 2), + (vec![vec![2]], 2), + (vec![vec![0, 1, 2]], 0), + ] { + let case = format!("leaves {leaves:?}, k={k}"); - assert_eq!( - actual.into_iter().map(Vec::from).collect::>(), - [Vec::::new(), vec![], vec![]] - ); + let actual = build(&values, leaves, k, 1); + + assert_eq!(actual, [Vec::::new(), vec![], vec![]], "{case}"); + } } #[test] fn requesting_more_neighbors_than_a_leaf_has_selects_all_other_points() { let values = [100.0_f32, 0.0, -100.0, 2.0, 200.0, 9.0]; - let data = MatrixView::try_from(&values[..], 6, 1).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - let actual = pool - .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, vec![vec![1, 3, 5]], 99)) - .unwrap(); - let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + let actual = build(&values, vec![vec![1, 3, 5]], 99, 1); assert_eq!( - test_support::sorted_members_per_row(&actual), + actual, [vec![], vec![3, 5], vec![], vec![1, 5], vec![], vec![1, 3]] ); } @@ -344,20 +267,67 @@ mod tests { #[test] fn unrankable_pairs_do_not_add_unassigned_ids_to_the_graph() { let values = [0.0_f32, 3.0, f32::NAN]; - let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - let actual = pool - .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, vec![vec![0, 1, 2]], 2)) - .unwrap(); + let actual = build(&values, vec![vec![0, 1, 2]], 2, 1); - assert_eq!( - actual.into_iter().map(Vec::from).collect::>(), - [vec![1], vec![0], vec![]] - ); + assert_eq!(actual, [vec![1], vec![0], vec![]]); + } + + #[test] + fn candidates_at_production_shape_are_the_symmetric_leaf_kernel_results() { + // Sixty overlapping leaves of 17 to 40 points run on four workers with + // reused buffers. They fill full SIMD groups of the leaf kernel. k = 2 + // uses a fixed top-k width, and k = 5 uses the runtime width. The kernel + // on each leaf alone is the reference, so tie order cannot differ. + let (point_count, dimensions) = (500, 24); + let values = dense_points(point_count, dimensions, 1290); + let data = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let leaves: Vec> = (0..60u32) + .map(|leaf| { + let ids: BTreeSet<_> = (0..17 + leaf % 24) + .map(|i| (leaf * 7 + i * 13) % point_count as u32) + .collect(); + ids.into_iter().collect() + }) + .collect(); + + for k in [2, 5] { + let mut expected = vec![BTreeSet::new(); point_count]; + for leaf in &leaves { + let rows: Vec = leaf + .iter() + .flat_map(|&id| &values[id as usize * dimensions..][..dimensions]) + .copied() + .collect(); + let width = k.min(leaf.len() - 1); + let mut output = vec![Candidate::default(); leaf.len() * width]; + select_leaf_neighbors::<_, L2>( + ARCH, + MatrixView::try_from(rows.as_slice(), leaf.len(), dimensions).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), leaf.len(), width).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap(); + for (source, row) in output.chunks_exact(width).enumerate() { + for neighbor in row.iter().filter(|neighbor| neighbor.is_assigned()) { + let (source, target) = (leaf[source], leaf[neighbor.local_idx as usize]); + expected[source as usize].insert(target); + expected[target as usize].insert(source); + } + } + } + let expected: Vec> = expected + .into_iter() + .map(|ids| ids.into_iter().collect()) + .collect(); + + let actual = thread_pool(4) + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves.clone(), k)) + .unwrap(); + + let actual: Vec> = actual.into_iter().map(Vec::from).collect(); + assert_eq!(actual, expected, "k={k}"); + } } #[test] @@ -406,7 +376,6 @@ mod tests { add_direct_leaf_candidates::<_, L2, _>( ARCH, data, - 0, &ids, requested_k, &mut buffers, @@ -416,86 +385,13 @@ mod tests { let actual: Vec<_> = candidates .into_iter() - .map(|row| { - let mut ids = Vec::from(row.into_inner()); - ids.sort_unstable(); - ids - }) + .map(|row| Vec::from(row.into_inner())) .collect(); - assert_eq!(actual, expected, "leaf {ids:?}, k={requested_k}"); + assert_eq!( + sorted_members_per_row(&actual), + expected, + "leaf {ids:?}, k={requested_k}" + ); } } - - #[rstest] - #[case::point_values(2, 1)] - #[case::neighbor_slots(1, 2)] - fn an_overflowing_leaf_shape_is_rejected_before_buffers_change( - #[case] dimensions: usize, - #[case] requested_k: usize, - ) { - let old_neighbor = Candidate::new(1, 7.0); - let mut buffers = LeafBuffers { - point_values: vec![3.0, 4.0], - neighbors: vec![old_neighbor], - ..LeafBuffers::default() - }; - - let error = buffers - .prepare(7, usize::MAX, dimensions, requested_k) - .unwrap_err(); - - assert!(matches!( - error, - LeafBuildError::ShapeOverflow { - leaf: 7, - rows: usize::MAX, - columns: 2, - } - )); - assert_eq!(buffers.point_values, [3.0, 4.0]); - assert_eq!(buffers.neighbors, [old_neighbor]); - } - - #[test] - fn a_kernel_error_retains_the_failed_leaf_and_original_cause() { - #[derive(Debug, thiserror::Error)] - #[error("distance calculation unavailable")] - struct DistanceFailure; - - // This stub supplies an otherwise hard-to-trigger dependency failure. - // Gathering, leaf indexing and error wrapping remain real. - struct UnavailableMetric; - impl LeafMetric for UnavailableMetric { - fn compute_distances( - _: MatrixView<'_, f32>, - _: MutMatrixView<'_, f32>, - ) -> crate::ANNResult<()> { - Err(crate::ANNError::new(DistanceFailure)) - } - } - - let values = [0.0_f32, 1.0, 4.0, 9.0]; - let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); - - let error = pool - .install(|| { - build_leaf_candidates::<_, UnavailableMetric, _>( - ARCH, - data, - vec![vec![], vec![1, 3]], - 1, - ) - }) - .unwrap_err(); - - let LeafBuildError::Kernel { leaf, source } = error else { - panic!("expected a leaf kernel error, got {error:?}"); - }; - assert_eq!(leaf, 1); - assert!(source.downcast_ref::().is_some()); - } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 1c318ef846..6e1bd9304e 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -370,7 +370,6 @@ where // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) - .map_err(ANNError::new) })?; // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. @@ -413,6 +412,14 @@ mod test_support { .collect() } + /// Build a Rayon pool with `threads` workers. + pub(super) fn thread_pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + } + #[test] fn membership_comparison_preserves_group_positions_and_duplicate_counts() { let rows = [vec![9, 3, 9], vec![], vec![2, 1], vec![2, 1]]; @@ -580,7 +587,7 @@ mod construction_tests { use super::*; use crate::graph::config; use half::f16; - use rstest::rstest; + use test_support::{sorted_members_per_row, thread_pool}; fn partition_policy() -> PiPNNConfig { PiPNNConfig { @@ -600,44 +607,90 @@ mod construction_tests { .build() } - #[rstest] - #[case::zero_maximum(|c: &mut PiPNNConfig| c.c_max = 0, "c_max must be greater than zero")] - #[case::zero_minimum(|c: &mut PiPNNConfig| c.c_min = 0, "c_min must be greater than zero")] - #[case::inverted_leaf_bounds(|c: &mut PiPNNConfig| c.c_min = 5, "c_min (5) must not exceed c_max (4)")] - #[case::zero_probability(|c: &mut PiPNNConfig| c.p_samp = 0.0, "p_samp (0) must be in (0, 1]")] - #[case::negative_probability(|c: &mut PiPNNConfig| c.p_samp = -0.5, "p_samp (-0.5) must be in (0, 1]")] - #[case::probability_above_one(|c: &mut PiPNNConfig| c.p_samp = 1.0 + f64::EPSILON, "must be in (0, 1]")] - #[case::nan_probability(|c: &mut PiPNNConfig| c.p_samp = f64::NAN, "p_samp (NaN) must be in (0, 1]")] - #[case::positive_infinite_probability(|c: &mut PiPNNConfig| c.p_samp = f64::INFINITY, "p_samp (inf) must be in (0, 1]")] - #[case::negative_infinite_probability(|c: &mut PiPNNConfig| c.p_samp = f64::NEG_INFINITY, "p_samp (-inf) must be in (0, 1]")] - #[case::missing_fanout(|c: &mut PiPNNConfig| c.fanout.clear(), "fanout must not be empty")] - #[case::zero_later_fanout(|c: &mut PiPNNConfig| c.fanout = vec![2, 0, 1], "fanout values must be greater than zero")] - #[case::zero_leaf_neighbors(|c: &mut PiPNNConfig| c.leaf_k = 0, "leaf_k must be greater than zero")] - #[case::zero_replicas(|c: &mut PiPNNConfig| c.replicas = 0, "replicas must be greater than zero")] - fn invalid_configuration_identifies_the_rejected_condition( - #[case] invalidate: fn(&mut PiPNNConfig), - #[case] expected_message: &str, - ) { - let mut config = partition_policy(); - invalidate(&mut config); - - let error = config.validate().unwrap_err(); - - let message = error.to_string(); - assert!(message.contains(expected_message), "{message}"); + /// Build the graph of `values` with `threads` workers. + #[expect(clippy::unwrap_used, reason = "an invalid fixture fails the test")] + fn build( + values: &[T], + dimensions: usize, + config: PiPNNConfig, + degree: usize, + metric: Metric, + threads: usize, + ) -> Vec> { + let data = MatrixView::try_from(values, values.len() / dimensions, dimensions).unwrap(); + let graph = graph_policy(degree, metric).unwrap(); + let pool = thread_pool(threads); + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + let graph = build_graph(data, &context).unwrap(); + graph.into_iter().map(Vec::from).collect() } - #[rstest] - #[case::smallest_positive_probability(f64::from_bits(1))] - #[case::all_points_sampled(1.0)] - fn equal_leaf_bounds_and_valid_probability_endpoints_are_accepted(#[case] p_samp: f64) { - let config = PiPNNConfig { - c_min: 4, - p_samp, - ..partition_policy() + #[test] + fn invalid_configuration_identifies_the_rejected_condition() { + let invalid = |change: fn(&mut PiPNNConfig)| { + let mut config = partition_policy(); + change(&mut config); + config }; + for (config, expected_message) in [ + (invalid(|c| c.c_max = 0), "c_max must be greater than zero"), + (invalid(|c| c.c_min = 0), "c_min must be greater than zero"), + ( + invalid(|c| c.c_min = 5), + "c_min (5) must not exceed c_max (4)", + ), + (invalid(|c| c.p_samp = 0.0), "p_samp (0) must be in (0, 1]"), + ( + invalid(|c| c.p_samp = -0.5), + "p_samp (-0.5) must be in (0, 1]", + ), + ( + invalid(|c| c.p_samp = 1.0 + f64::EPSILON), + "must be in (0, 1]", + ), + ( + invalid(|c| c.p_samp = f64::NAN), + "p_samp (NaN) must be in (0, 1]", + ), + ( + invalid(|c| c.p_samp = f64::INFINITY), + "p_samp (inf) must be in (0, 1]", + ), + ( + invalid(|c| c.p_samp = f64::NEG_INFINITY), + "p_samp (-inf) must be in (0, 1]", + ), + (invalid(|c| c.fanout.clear()), "fanout must not be empty"), + ( + invalid(|c| c.fanout = vec![2, 0, 1]), + "fanout values must be greater than zero", + ), + ( + invalid(|c| c.leaf_k = 0), + "leaf_k must be greater than zero", + ), + ( + invalid(|c| c.replicas = 0), + "replicas must be greater than zero", + ), + ] { + let message = config.validate().unwrap_err().to_string(); + + assert!(message.contains(expected_message), "{message}"); + } + } - config.validate().unwrap(); + #[test] + fn equal_leaf_bounds_and_valid_probability_endpoints_are_accepted() { + for p_samp in [f64::from_bits(1), 1.0] { + let config = PiPNNConfig { + c_min: 4, + p_samp, + ..partition_policy() + }; + + config.validate().unwrap(); + } } #[test] @@ -647,10 +700,7 @@ mod construction_tests { ..partition_policy() }; let graph = graph_policy(2, Metric::L2).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); + let pool = thread_pool(1); let error = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap_err(); @@ -662,300 +712,285 @@ mod construction_tests { ); } - #[rstest] - #[case::cosine_uses_triangle_pruning(Metric::L2, Metric::Cosine)] - #[case::normalized_cosine_uses_triangle_pruning(Metric::L2, Metric::CosineNormalized)] - #[case::inner_product_uses_occluding_pruning(Metric::InnerProduct, Metric::InnerProduct)] - fn a_build_context_accepts_metrics_with_the_same_pruning_kind( - #[case] graph_metric: Metric, - #[case] requested_metric: Metric, - ) { - let graph = graph_policy(2, graph_metric).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - - PiPNNBuildContext::new(partition_policy(), &graph, requested_metric, &pool).unwrap(); - } - - #[rstest] - #[case::inner_product_with_triangle_pruning(Metric::L2, Metric::InnerProduct)] - #[case::l2_with_occluding_pruning(Metric::InnerProduct, Metric::L2)] - fn a_build_context_rejects_incompatible_pruning_kinds( - #[case] graph_metric: Metric, - #[case] requested_metric: Metric, - ) { - let graph = graph_policy(2, graph_metric).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - - let error = PiPNNBuildContext::new(partition_policy(), &graph, requested_metric, &pool) - .unwrap_err(); - - assert!( - error.to_string().contains("incompatible with metric"), - "{error}" - ); + #[test] + fn a_build_context_requires_the_pruning_kind_of_its_metric() { + // L2, cosine and normalized cosine share triangle pruning. Inner product + // uses occluding pruning. + let pool = thread_pool(1); + for (graph_metric, metric, compatible) in [ + (Metric::L2, Metric::Cosine, true), + (Metric::L2, Metric::CosineNormalized, true), + (Metric::InnerProduct, Metric::InnerProduct, true), + (Metric::L2, Metric::InnerProduct, false), + (Metric::InnerProduct, Metric::L2, false), + ] { + let graph = graph_policy(2, graph_metric).unwrap(); + + let result = PiPNNBuildContext::new(partition_policy(), &graph, metric, &pool); + + let case = format!("graph {graph_metric:?}, metric {metric:?}"); + match result { + Ok(_) => assert!(compatible, "{case}"), + Err(error) => { + assert!(!compatible, "{case}: {error}"); + assert!( + error.to_string().contains("incompatible with metric"), + "{case}: {error}" + ); + } + } + } } - #[rstest] - #[case::no_points(0, 2, "at least one data point")] - #[case::no_dimensions(2, 0, "at least one data dimension")] - fn an_empty_dataset_axis_is_rejected_before_building( - #[case] rows: usize, - #[case] columns: usize, - #[case] expected_message: &str, - ) { - let data = MatrixView::try_from(&[] as &[f32], rows, columns).unwrap(); + #[test] + fn an_empty_dataset_axis_is_rejected_before_building() { let graph = graph_policy(2, Metric::L2).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); + let pool = thread_pool(1); let context = PiPNNBuildContext::new(partition_policy(), &graph, Metric::L2, &pool).unwrap(); - let error = build_graph(data, &context).unwrap_err(); + for (rows, columns, expected_message) in [ + (0, 2, "at least one data point"), + (2, 0, "at least one data dimension"), + ] { + let data = MatrixView::try_from(&[] as &[f32], rows, columns).unwrap(); + + let error = build_graph(data, &context).unwrap_err(); - assert!(error.to_string().contains(expected_message), "{error}"); + assert!(error.to_string().contains(expected_message), "{error}"); + } } - #[rstest] - #[case::float32([0.0_f32, 1.0, 4.0])] - #[case::float16([0.0_f32, 1.0, 4.0].map(f16::from_f32))] - #[case::signed_integer([0_i8, 1, 4])] - #[case::unsigned_integer([0_u8, 1, 4])] - fn native_vector_types_build_the_expected_neighbor_graph( - #[case] values: [T; 3], - ) { + #[test] + fn native_vector_types_build_the_expected_neighbor_graph() { // Nearest choices 0 -> 1, 1 -> 0 and 2 -> 1 become a symmetric chain. - let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); - let graph = graph_policy(2, Metric::L2).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); - let context = - PiPNNBuildContext::new(partition_policy(), &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); - let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); + fn check(values: [T; 3]) { + let actual = build(&values, 1, partition_policy(), 2, Metric::L2, 2); + + assert_eq!( + sorted_members_per_row(&actual), + [vec![1], vec![0, 2], vec![1]], + "{}", + std::any::type_name::() + ); + } - assert_eq!( - test_support::sorted_members_per_row(&actual), - [vec![1], vec![0, 2], vec![1]] - ); + check([0.0_f32, 1.0, 4.0]); + check([0.0_f32, 1.0, 4.0].map(f16::from_f32)); + check([0_i8, 1, 4]); + check([0_u8, 1, 4]); } - #[rstest] - #[case::l2(Metric::L2, vec![vec![2], vec![2], vec![0, 1, 3], vec![2]])] - #[case::cosine(Metric::Cosine, vec![vec![1], vec![0], vec![3], vec![2]])] - #[case::normalized_cosine(Metric::CosineNormalized, vec![vec![1], vec![0], vec![3], vec![2]])] - #[case::inner_product(Metric::InnerProduct, vec![vec![1], vec![0, 3], vec![3], vec![1, 2]])] - fn the_requested_metric_determines_leaf_neighbors( - #[case] metric: Metric, - #[case] expected: Vec>, - ) { + #[test] + fn the_requested_metric_determines_leaf_neighbors() { // L2's nearest choices are [2, 2, 0, 2]; cosine pairs similar directions // 0 <-> 1 and 2 <-> 3; dot products give [1, 3, 3, 2]. // The second coordinate sits in the last dimension of an embedding. let dimensions = 1537; - let values = test_support::packed_points( - &[[1.0, 0.0], [5.0, 2.0], [1.0, 3.0], [0.0, 9.0]], - dimensions, - metric == Metric::CosineNormalized, - ); - let data = MatrixView::try_from(values.as_slice(), 4, dimensions).unwrap(); - let graph = graph_policy(3, metric).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); - let context = PiPNNBuildContext::new(partition_policy(), &graph, metric, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); - let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); - - assert_eq!( - test_support::sorted_members_per_row(&actual), - test_support::sorted_members_per_row(&expected) - ); + for (metric, expected) in [ + (Metric::L2, vec![vec![2], vec![2], vec![0, 1, 3], vec![2]]), + (Metric::Cosine, vec![vec![1], vec![0], vec![3], vec![2]]), + ( + Metric::CosineNormalized, + vec![vec![1], vec![0], vec![3], vec![2]], + ), + ( + Metric::InnerProduct, + vec![vec![1], vec![0, 3], vec![3], vec![1, 2]], + ), + ] { + let values = test_support::packed_points( + &[[1.0, 0.0], [5.0, 2.0], [1.0, 3.0], [0.0, 9.0]], + dimensions, + metric == Metric::CosineNormalized, + ); + + let actual = build(&values, dimensions, partition_policy(), 3, metric, 2); + + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "{metric:?}" + ); + } } - #[rstest] - #[case::l2(Metric::L2, vec![vec![2], vec![], vec![0], vec![]])] - #[case::cosine(Metric::Cosine, vec![vec![1], vec![0], vec![3], vec![2]])] - #[case::normalized_cosine(Metric::CosineNormalized, vec![vec![1], vec![0], vec![3], vec![2]])] - #[case::inner_product(Metric::InnerProduct, vec![vec![1], vec![0], vec![3], vec![2]])] - fn splitting_uses_the_requested_metric_to_group_points( - #[case] metric: Metric, - #[case] expected: Vec>, - ) { + #[test] + fn splitting_uses_the_requested_metric_to_group_points() { // Four points force splitting at c_max=2. Every point is sampled, so // leader order cannot change memberships. With L2, leader 0 gets {0,2} // and leader 2 gets all four points, which then split into singletons. // Cosine pairs directions {0,1} and {2,3}. Inner product picks leaders // {1,2} for points 0/1 and {2,3} for points 2/3; the oversized leader-2 // cluster then splits into those same pairs using fanout one. - let values = test_support::packed_points( - &[[1.0, 0.0], [6.0, 1.0], [2.0, 4.0], [0.0, 9.0]], - 2, - metric == Metric::CosineNormalized, - ); - let data = MatrixView::try_from(values.as_slice(), 4, 2).unwrap(); + // Every two-point leaf contributes its only pair; degree three retains + // every edge. The graph therefore exposes partition membership directly. let config = PiPNNConfig { c_max: 2, ..partition_policy() }; - // Every two-point leaf contributes its only pair; degree three retains - // every edge. The graph therefore exposes partition membership directly. - let graph = graph_policy(3, metric).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); - let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); - let actual: Vec<_> = actual.into_iter().map(Vec::from).collect(); - - assert_eq!( - test_support::sorted_members_per_row(&actual), - test_support::sorted_members_per_row(&expected) - ); + for (metric, expected) in [ + (Metric::L2, vec![vec![2], vec![], vec![0], vec![]]), + (Metric::Cosine, vec![vec![1], vec![0], vec![3], vec![2]]), + ( + Metric::CosineNormalized, + vec![vec![1], vec![0], vec![3], vec![2]], + ), + ( + Metric::InnerProduct, + vec![vec![1], vec![0], vec![3], vec![2]], + ), + ] { + let values = test_support::packed_points( + &[[1.0, 0.0], [6.0, 1.0], [2.0, 4.0], [0.0, 9.0]], + 2, + metric == Metric::CosineNormalized, + ); + + let actual = build(&values, 2, config.clone(), 3, metric, 2); + + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "{metric:?}" + ); + } } - #[rstest] - #[case::signed_leaf_ranking([1_i8, 0, 4, 1, 1, 2, 0, 9], 1, 3)] - #[case::unsigned_leaf_ranking([1_u8, 0, 4, 1, 1, 2, 0, 9], 1, 3)] - #[case::signed_final_pruning([1_i8, 0, 4, 1, 1, 2, 0, 9], 3, 1)] - #[case::unsigned_final_pruning([1_u8, 0, 4, 1, 1, 2, 0, 9], 3, 1)] - fn normalized_cosine_requests_on_raw_integers_use_vector_norms( - #[case] values: [T; 8], - #[case] leaf_k: usize, - #[case] degree: usize, - ) { + #[test] + fn normalized_cosine_requests_on_raw_integers_use_vector_norms() { // Cosine pairs 0 <-> 1 and 2 <-> 3. Unnormalized dot products instead // connect 1 to 3, so treating raw integers as unit vectors changes edges. - let data = MatrixView::try_from(&values[..], 4, 2).unwrap(); - let config = PiPNNConfig { - leaf_k, - ..partition_policy() - }; - let graph = graph_policy(degree, Metric::CosineNormalized).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); - let context = - PiPNNBuildContext::new(config, &graph, Metric::CosineNormalized, &pool).unwrap(); + // The first build checks leaf ranking. The second build selects three + // leaf neighbors and checks final pruning to one. + fn check(values: [T; 8]) { + for (leaf_k, degree) in [(1, 3), (3, 1)] { + let config = PiPNNConfig { + leaf_k, + ..partition_policy() + }; - let actual = build_graph(data, &context).unwrap(); + let actual = build(&values, 2, config, degree, Metric::CosineNormalized, 2); - assert_eq!( - actual.into_iter().map(Vec::from).collect::>(), - [vec![1], vec![0], vec![3], vec![2]] - ); + assert_eq!( + actual, + [vec![1], vec![0], vec![3], vec![2]], + "{}, leaf_k={leaf_k}, degree={degree}", + std::any::type_name::() + ); + } + } + + check([1_i8, 0, 4, 1, 1, 2, 0, 9]); + check([1_u8, 0, 4, 1, 1, 2, 0, 9]); } #[test] fn final_pruning_enforces_graph_degree_after_leaf_candidates_are_merged() { - let values = [0.0_f32, 1.0, 4.0]; - let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); let config = PiPNNConfig { leaf_k: 2, ..partition_policy() }; - let graph = graph_policy(1, Metric::L2).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - let actual = build_graph(data, &context).unwrap(); + let actual = build(&[0.0_f32, 1.0, 4.0], 1, config, 1, Metric::L2, 2); - assert_eq!( - actual.into_iter().map(Vec::from).collect::>(), - [vec![1], vec![0], vec![1]] - ); + assert_eq!(actual, [vec![1], vec![0], vec![1]]); } - #[rstest] - #[case::serial_single_replica(1, 1)] - #[case::parallel_replicas(3, 2)] - fn overlapping_partitions_build_the_same_deduplicated_graph( - #[case] workers: usize, - #[case] replicas: usize, - ) { + #[test] + fn overlapping_partitions_build_the_same_deduplicated_graph() { // Sampling every point with fanout 2 makes two overlapping copies // of each close pair. Neither replicas nor worker count may duplicate edges. - let values = [0.0_f32, 1.0, 10.0, 11.0]; - let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let config = PiPNNConfig { - c_max: 2, - replicas, - ..partition_policy() - }; - let graph = graph_policy(2, Metric::L2).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(workers) - .build() - .unwrap(); - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); - - assert_eq!( - actual.into_iter().map(Vec::from).collect::>(), - [vec![1], vec![0], vec![3], vec![2]] - ); + for (workers, replicas) in [(1, 1), (3, 2)] { + let config = PiPNNConfig { + c_max: 2, + replicas, + ..partition_policy() + }; + + let actual = build( + &[0.0_f32, 1.0, 10.0, 11.0], + 1, + config, + 2, + Metric::L2, + workers, + ); + + assert_eq!( + actual, + [vec![1], vec![0], vec![3], vec![2]], + "{workers} workers, {replicas} replicas" + ); + } } #[test] fn a_single_point_has_one_empty_adjacency_list() { - let values = [3.0_f32]; - let data = MatrixView::try_from(&values[..], 1, 1).unwrap(); - let graph = graph_policy(2, Metric::L2).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - let context = - PiPNNBuildContext::new(partition_policy(), &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); + let actual = build(&[3.0_f32], 1, partition_policy(), 2, Metric::L2, 1); - assert_eq!( - actual.into_iter().map(Vec::from).collect::>(), - [Vec::::new()] - ); + assert_eq!(actual, [Vec::::new()]); } #[test] - fn stalled_partitioning_remains_a_typed_build_error() { - let values = [1.0_f32; 5]; - let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + fn multi_level_builds_are_deterministic_for_every_metric() { + // With 1,000 points and c_max 32, the build splits over several levels, + // makes overlapping leaves, merges their candidates in parallel and + // prunes lists above the degree. The worker count must not change the + // graph. + let (points, dimensions, degree) = (1000, 16, 8); let config = PiPNNConfig { - c_max: 2, - fanout: vec![1], - ..partition_policy() + c_max: 32, + c_min: 8, + p_samp: 0.05, + fanout: vec![4, 2], + leaf_k: 2, + replicas: 1, }; - let graph = graph_policy(2, Metric::L2).unwrap(); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let mut values = test_support::dense_points(points, dimensions, 1290); + if metric == Metric::CosineNormalized { + test_support::normalize(&mut values, dimensions); + } - let error = build_graph(data, &context).unwrap_err(); + let serial = build(&values, dimensions, config.clone(), degree, metric, 1); + let parallel = build(&values, dimensions, config.clone(), degree, metric, 4); + + assert_eq!(serial, parallel, "{metric:?}"); + for (point, neighbors) in serial.iter().enumerate() { + let mut distinct = neighbors.clone(); + distinct.sort_unstable(); + distinct.dedup(); + assert!( + !neighbors.is_empty() + && neighbors.len() <= degree + && distinct.len() == neighbors.len() + && !neighbors.contains(&(point as u32)), + "{metric:?}, point {point}: {neighbors:?}" + ); + } + } + } - assert!(matches!( - error.downcast_ref::(), - Some(partitioning::PartitionError::IterationLimit { size: 5, .. }) - )); + #[test] + fn copies_of_one_vector_build_leaf_local_edges() { + // Every leader is a copy, so each point ties at every split. The build + // splits the six copies into two leaves of three points. + let actual = build(&[1.0_f32; 6], 1, partition_policy(), 2, Metric::L2, 2); + + for (point, neighbors) in actual.iter().enumerate() { + assert!(!neighbors.is_empty(), "point {point}"); + assert!( + neighbors + .iter() + .all(|&neighbor| neighbor / 3 == point as u32 / 3), + "point {point}: {neighbors:?}" + ); + } } } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 1871c139d8..61528c943a 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -15,10 +15,16 @@ //! A configured level assigns each point to `fanout[level]` leaders. A deeper //! level assigns each point to one leader. Each replica uses a different //! deterministic seed. +//! +//! Buffer sizes use plain multiplication. Point and leader IDs in one split are +//! distinct dataset rows, so a gathered `rows x dimensions` buffer is no larger +//! than the dataset. An assignment buffer holds `points x fanout` IDs. With at +//! most `u32::MAX` points and `LEADER_CAP` leaders, this count fits in a 64-bit +//! `usize`. use std::collections::HashSet; -use crate::{ANNError, ANNResult, utils::VectorRepr}; +use crate::{ANNResult, utils::VectorRepr}; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, @@ -28,41 +34,41 @@ use rayon::prelude::*; use super::{ PiPNNConfig, + conversion::gather_as_f32, partition_kernel::{PartitionKernelWorkspace, assign_leaders}, partition_metric::PartitionMetric, simd::Simd, topk::UNASSIGNED, }; -// These constants control internal batching and deterministic seed generation. +/// Seed of the first replica. Fixed seeds make each build reproducible. const PARTITION_SEED: u64 = 1_000; +/// Seed step between replicas, so each replica samples different leaders. const REPLICA_SEED_STEP: u64 = 7_919; +/// Maximum number of leaders in one split. The PiPNN paper uses the same hard +/// cap ("typically 1000", Section 4.1). A split costs +/// `points x leaders x dimensions` multiply-adds. Without the cap, the first +/// split of 10M points at `p_samp = 0.005` samples 50,000 leaders and does 50 +/// times the work. const LEADER_CAP: usize = 1_000; -const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; -const MIN_ASSIGNMENT_STRIPE_POINTS: usize = 32; +/// Size target of the distance block of one stripe, `points x leaders x 4` +/// bytes. GEMM writes the block and the top-k scan reads it back, so the block +/// must stay in the L2 cache of the core. On an AMD EPYC 7763 (512 KiB L2 per +/// core), this value gave the shortest partition time for 10M points: 128 KiB +/// was about 7% slower and 2 MiB about 5% slower. Current server cores have +/// 512 KiB to 2 MiB of L2 cache. +const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 512 * 1024; +/// Maximum number of points in one stripe. With few leaders, the cache target +/// alone allows very tall stripes. This bound keeps the gathered points of one +/// worker at `1024 x dimensions x 4` bytes or less. const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; +/// Clusters with at least this many points scatter in parallel. Such clusters +/// occur at the first levels, where few splits run at the same time. A serial +/// scatter of 10M points at fanout 10 takes 0.59 s, and a parallel scatter +/// takes 0.15 s on 16 threads (AMD EPYC 7763). The serial time does not +/// decrease with more threads, so its share of the build increases on larger +/// machines. Below this size, a serial scatter takes a few milliseconds. const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; -const MAX_PARTITION_ITERATIONS: usize = 30; - -/// Error from partition shape checks or recursion progress. -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub(crate) enum PartitionError { - #[error("{buffer} shape {rows} x {cols} overflows usize")] - ShapeOverflow { - buffer: &'static str, - rows: usize, - cols: usize, - }, - #[error( - "partition stopped after {limit} iterations with an oversized cluster of size \ - {size} at level {level}" - )] - IterationLimit { - size: usize, - level: usize, - limit: usize, - }, -} struct PendingPartition { point_ids: Vec, @@ -129,6 +135,11 @@ where /// /// The function processes one work queue per recursion level. It merges leaves /// smaller than `c_min` after the queue becomes empty. +/// +/// The queue empties because the fanout schedule is finite. After the schedule, +/// a split puts each point into at most one child. [`split_partition`] makes sure +/// that such a child is smaller than its parent, so the cluster sizes decrease at +/// each later level. fn partition_replica( arch: A, data: MatrixView<'_, T>, @@ -152,12 +163,7 @@ where level: 0, seed, }]; - - for _ in 0..MAX_PARTITION_ITERATIONS { - if pending.is_empty() { - return Ok(merge_undersized_leaves(leaves, config.c_min, config.c_max)); - } - + while !pending.is_empty() { // Indexed parallel collection preserves parent-partition order. let splits: ANNResult> = pending .into_par_iter() @@ -173,18 +179,7 @@ where } pending = next_level; } - - let Some(largest) = pending - .iter() - .max_by_key(|partition| partition.point_ids.len()) - else { - return Ok(merge_undersized_leaves(leaves, config.c_min, config.c_max)); - }; - Err(ANNError::new(PartitionError::IterationLimit { - size: largest.point_ids.len(), - level: largest.level, - limit: MAX_PARTITION_ITERATIONS, - })) + Ok(merge_undersized_leaves(leaves, config.c_min, config.c_max)) } /// Split one oversized cluster into child partitions. @@ -217,7 +212,27 @@ where let mut pending = Vec::new(); let mut leaves = Vec::new(); - for cluster in clusters { + for (&leader, mut cluster) in leaders.iter().zip(clusters) { + if fanout == 1 && cluster.len() == partition.point_ids.len() { + // Every point chose this leader, so the split made no progress. Under + // L2, each leader is nearest to itself. Thus this occurs only when all + // sampled leaders are copies of one vector. Under cosine, vectors with + // the same direction also tie. Under inner product, one large leader + // can be first for every point. + // + // Copies of the leader have the same distance to every point, so any + // grouping of them gives the same leaf neighbors. Put the copies into + // leaves of equal size. The rest does not contain the leader, so it is + // smaller than its parent. Each stalled level therefore removes at + // least one point. + let copy = bytemuck::cast_slice::(data.row(leader as usize)); + let (copies, rest): (Vec<_>, Vec<_>) = cluster + .into_iter() + .partition(|&id| bytemuck::cast_slice::(data.row(id as usize)) == copy); + let leaf_len = copies.len().div_ceil(copies.len().div_ceil(config.c_max)); + leaves.extend(copies.chunks(leaf_len).map(<[u32]>::to_vec)); + cluster = rest; + } if cluster.is_empty() { continue; } @@ -253,8 +268,9 @@ fn replica_seed(replica: usize) -> u64 { PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) } -// This LCG derives child seeds. Wrapping arithmetic gives the same mapping in -// debug and release builds on all supported platforms. +// This LCG step derives a child seed from the parent seed and the cluster size. +// The multiplier is Knuth's MMIX constant. Wrapping arithmetic gives the same +// mapping in debug and release builds on all supported platforms. fn mix_seed(seed: u64, salt: u64) -> u64 { seed.wrapping_mul(6_364_136_223_846_793_005) .wrapping_add(salt) @@ -279,51 +295,33 @@ where T: VectorRepr + Send + Sync, { let dimension_count = data.ncols(); - let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; - let mut leader_values = vec![0.0f32; leader_values_len]; - gather_vectors(data, leader_ids, &mut leader_values)?; + let mut leader_values = vec![0.0f32; leader_ids.len() * dimension_count]; + gather_as_f32(data, leader_ids, &mut leader_values)?; let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count) - .map_err(|error| ANNError::new(error.as_static()))?; + MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count)?; let leaders = M::create_leaders(leader_matrix); let leader_count = M::leader_count(&leaders); let fanout = fanout.min(leader_count); - let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; - let mut assignments = vec![0u32; assignment_len]; + let mut assignments = vec![0u32; point_ids.len() * fanout]; let stripe_points = assignment_stripe_point_count(leader_count); - let stripe_assignment_count = stripe_points * fanout; - let stripe_count = point_ids.len().div_ceil(stripe_points); - let worker_stripe_count = stripe_count.div_ceil(rayon::current_num_threads()); - let worker_point_count = checked_area("assignment worker", worker_stripe_count, stripe_points)?; - let worker_assignment_count = checked_area("assignment worker", worker_point_count, fanout)?; - - // Each worker chunk reuses one buffer lease for all its stripes. // `build_graph` runs this operation in the pool from the build context. + // A buffer lease takes one lock, which costs far less than the GEMM and + // ranking of one stripe. assignments - .par_chunks_mut(worker_assignment_count) - .enumerate() - .try_for_each(|(worker, worker_assignments)| { - let mut buffers = stripe_buffers.get_ref(()); - let worker_first = worker * worker_point_count; - for (stripe, stripe_assignments) in worker_assignments - .chunks_mut(stripe_assignment_count) - .enumerate() - { - let first_point = worker_first + stripe * stripe_points; - let stripe_point_count = stripe_assignments.len() / fanout; - assign_point_stripe::( - arch, - data, - &point_ids[first_point..first_point + stripe_point_count], - &leaders, - fanout, - &mut buffers, - stripe_assignments, - )?; - } - Ok::<(), ANNError>(()) + .par_chunks_mut(stripe_points * fanout) + .zip(point_ids.par_chunks(stripe_points)) + .try_for_each(|(stripe_assignments, stripe_ids)| { + assign_point_stripe::( + arch, + data, + stripe_ids, + &leaders, + fanout, + &mut stripe_buffers.get_ref(()), + stripe_assignments, + ) })?; Ok(scatter_assignments( @@ -356,7 +354,7 @@ where { let point_count = point_ids.len(); let dimensions = data.ncols(); - let point_values_len = checked_area("point stripe", point_count, dimensions)?; + let point_values_len = point_count * dimensions; // Keep each buffer at its largest length. Every operation uses an explicit // active prefix. if buffers.point_values.len() < point_values_len { @@ -370,21 +368,12 @@ where &mut point_values[..point_values_len], point_count, dimensions, - ) - .map_err(|error| ANNError::new(error.as_static()))?; - gather_vectors(data, point_ids, points.as_mut_slice())?; - let output = MutMatrixView::try_from(assignments, point_count, fanout) - .map_err(|error| ANNError::new(error.as_static()))?; + )?; + gather_as_f32(data, point_ids, points.as_mut_slice())?; + let output = MutMatrixView::try_from(assignments, point_count, fanout)?; assign_leaders::(arch, points.as_view(), leaders, output, kernel_workspace) } -fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> -where - T: VectorRepr, -{ - super::conversion::gather_as_f32(data, indices, output).map_err(|(_, error)| error.into()) -} - /// Group assigned point IDs by child partition. /// /// Both the serial and parallel paths preserve point order inside each child. @@ -510,32 +499,28 @@ fn drain_sorted(set: &mut HashSet) -> Vec { values } -fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { - rows.checked_mul(cols) - .ok_or_else(|| ANNError::new(PartitionError::ShapeOverflow { buffer, rows, cols })) -} - +/// Return the number of points in one assignment stripe. +/// +/// `LEADER_CAP` keeps this count at 131 points or more. fn assignment_stripe_point_count(leader_count: usize) -> usize { - let point_count = ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count * size_of::()); - let point_count = if point_count.is_power_of_two() { - point_count - } else { - point_count.next_power_of_two() / 2 - }; - point_count.clamp(MIN_ASSIGNMENT_STRIPE_POINTS, MAX_ASSIGNMENT_STRIPE_POINTS) + (ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count * size_of::())) + .min(MAX_ASSIGNMENT_STRIPE_POINTS) } #[cfg(all(test, not(miri)))] mod tests { use super::*; - use crate::graph::pipnn::{L2, test_support}; + use crate::graph::pipnn::{ + L2, + test_support::{self, dense_points, sorted_members_per_row, thread_pool}, + }; + use diskann_vector::distance::Metric; use diskann_wide::ARCH; - use rstest::rstest; // Neither leaf order nor point order within a leaf is part of the result. // Preserve repeated leaves and IDs so comparison still detects duplicates. fn sorted_leaf_memberships(leaves: &[Vec]) -> Vec> { - let mut leaves = test_support::sorted_members_per_row(leaves); + let mut leaves = sorted_members_per_row(leaves); leaves.sort_unstable(); leaves } @@ -551,65 +536,76 @@ mod tests { } } - #[rstest] - #[case::minimum_sample(7, 0.01, 2)] - #[case::rounded_fraction(7, 0.31, 3)] - #[case::all_points(7, 1.0, 7)] - #[case::below_cap(999, 1.0, 999)] - #[case::at_cap(1000, 1.0, 1000)] - #[case::above_cap(1001, 1.0, 1000)] - fn sampled_leaders_are_a_repeatable_subset_without_replacement( - #[case] point_count: u32, - #[case] fraction: f64, - #[case] expected_count: usize, - ) { - let ids: Vec<_> = (0..point_count).map(|id| 10 + 3 * id).collect(); + #[expect(clippy::unwrap_used, reason = "an invalid fixture fails the test")] + fn split( + data: MatrixView<'_, f32>, + config: &PiPNNConfig, + parent: PendingPartition, + ) -> PartitionSplit { + let buffers = StripeBufferPool::new((), 0, None); + thread_pool(2) + .install(|| split_partition::<_, M, _>(ARCH, data, config, parent, &buffers)) + .unwrap() + } - let mut leaders = sample_leaders(&ids, fraction, 1290); + #[test] + fn sampled_leaders_are_a_repeatable_subset_without_replacement() { + // (points, sampling fraction, expected leaders) + for (point_count, fraction, expected_count) in [ + (7, 0.01, 2), // A split needs two leaders. + (7, 0.31, 3), // 2.17 rounds up. + (7, 1.0, 7), + (999, 1.0, 999), + (1000, 1.0, LEADER_CAP), + (1001, 1.0, LEADER_CAP), + ] { + let case = format!("{point_count} points, fraction {fraction}"); + let ids: Vec<_> = (0..point_count).map(|id| 10 + 3 * id).collect(); + + let mut leaders = sample_leaders(&ids, fraction, 1290); + + assert_eq!(leaders.len(), expected_count, "{case}"); + let distinct: HashSet<_> = leaders.iter().collect(); + assert_eq!(distinct.len(), expected_count, "{case}"); + assert!(leaders.iter().all(|id| ids.contains(id)), "{case}"); + let mut repeated = sample_leaders(&ids, fraction, 1290); + leaders.sort_unstable(); + repeated.sort_unstable(); + assert_eq!(leaders, repeated, "{case}"); + } + } - assert_eq!(leaders.len(), expected_count); - assert_eq!( - leaders.iter().copied().collect::>().len(), - expected_count - ); - assert!(leaders.iter().all(|id| ids.contains(id))); - let mut repeated = sample_leaders(&ids, fraction, 1290); - leaders.sort_unstable(); - repeated.sort_unstable(); - assert_eq!(leaders, repeated); - } - - #[rstest] - #[case::nearest_only(1, vec![vec![5, 1], vec![4]])] - #[case::all_leaders(2, vec![vec![5, 4, 1], vec![5, 4, 1]])] - #[case::fanout_above_leader_count(9, vec![vec![5, 4, 1], vec![5, 4, 1]])] - fn points_join_their_nearest_leaders(#[case] fanout: usize, #[case] expected: Vec>) { + #[test] + fn points_join_their_nearest_leaders() { // Leader columns are global IDs [2, 0], at x=10 and x=0. // The requested points [5, 4, 1] are at x=8, x=2 and x=100. let values = [0.0_f32, 100.0, 10.0, -100.0, 2.0, 8.0]; let data = MatrixView::try_from(&values[..], 6, 1).unwrap(); let buffers = StripeBufferPool::new((), 0, None); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); + let pool = thread_pool(2); - let actual = pool - .install(|| { - assign_to_leaders::<_, L2, _>(ARCH, data, &[5, 4, 1], &[2, 0], fanout, &buffers) - }) - .unwrap(); + for (fanout, expected) in [ + (1, [vec![5, 1], vec![4]]), + (2, [vec![5, 4, 1], vec![5, 4, 1]]), + // Fanout above the leader count assigns every point to every leader. + (9, [vec![5, 4, 1], vec![5, 4, 1]]), + ] { + let actual = pool + .install(|| { + assign_to_leaders::<_, L2, _>(ARCH, data, &[5, 4, 1], &[2, 0], fanout, &buffers) + }) + .unwrap(); - assert_eq!( - test_support::sorted_members_per_row(&actual), - test_support::sorted_members_per_row(&expected) - ); + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "fanout {fanout}" + ); + } } - #[rstest] - #[case::one_worker(1)] - #[case::several_workers(3)] - fn multiple_stripes_and_their_partial_tail_preserve_all_assignments(#[case] workers: usize) { + #[test] + fn multiple_stripes_and_their_partial_tail_preserve_all_assignments() { // More than two maximum-sized stripes, with three points left over. // Each dense row selects two of the three centers at 10, 0 and 4. // Every row selects 4; its other center is 10 for x=10 and 0 otherwise. @@ -620,31 +616,72 @@ mod tests { .collect(); let data = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); let ids: Vec<_> = (0..point_count as u32).rev().collect(); - let expected = vec![ - ids.iter() - .copied() - .filter(|id| id % 3 == 2) - .collect::>(), - ids.iter() - .copied() - .filter(|id| id % 3 != 2) - .collect::>(), + let expected = [ + ids.iter().copied().filter(|id| id % 3 == 2).collect(), + ids.iter().copied().filter(|id| id % 3 != 2).collect(), ids.clone(), ]; let buffers = StripeBufferPool::new((), 0, None); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(workers) - .build() - .unwrap(); - let actual = pool - .install(|| assign_to_leaders::<_, L2, _>(ARCH, data, &ids, &[2, 0, 1], 2, &buffers)) + for workers in [1, 3] { + let actual = thread_pool(workers) + .install(|| { + assign_to_leaders::<_, L2, _>(ARCH, data, &ids, &[2, 0, 1], 2, &buffers) + }) + .unwrap(); + + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "{workers} workers" + ); + } + } + + #[test] + fn assignments_at_production_shape_join_each_point_to_its_nearest_leaders() { + // 300 leaders give stripes of 436 points, so 3,000 points use several + // stripes, a partial tail, several workers and full SIMD leader groups. + // Dense fixtures keep the f32 scores exact, so the check needs no + // tolerance. It accepts any order of tied leaders. + let (point_count, dimensions, fanout) = (3000, 32, 3); + let values = dense_points(point_count, dimensions, 1290); + let data = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let row = |id: u32| &values[id as usize * dimensions..][..dimensions]; + let ids: Vec = (0..point_count as u32).rev().collect(); + let leaders: Vec = (0..point_count as u32).step_by(10).collect(); + let buffers = StripeBufferPool::new((), 0, None); + + let clusters = thread_pool(4) + .install(|| assign_to_leaders::<_, L2, _>(ARCH, data, &ids, &leaders, fanout, &buffers)) .unwrap(); - assert_eq!( - test_support::sorted_members_per_row(&actual), - test_support::sorted_members_per_row(&expected) - ); + let mut chosen = vec![Vec::new(); point_count]; + for (column, cluster) in clusters.iter().enumerate() { + // `ids` is descending, and scatter keeps the input order. + assert!(cluster.is_sorted_by(|a, b| a > b), "leader {column}"); + for &id in cluster { + chosen[id as usize].push(column); + } + } + for (id, columns) in chosen.iter().enumerate() { + let distance = |column: usize| { + test_support::distance(Metric::L2, row(id as u32), row(leaders[column])) + }; + let farthest_chosen = columns + .iter() + .map(|&column| distance(column)) + .fold(f64::MIN, f64::max); + let nearest_other = (0..leaders.len()) + .filter(|column| !columns.contains(column)) + .map(distance) + .fold(f64::INFINITY, f64::min); + assert_eq!(columns.len(), fanout, "point {id}"); + assert!( + farthest_chosen <= nearest_other, + "point {id}: chosen {columns:?} at {farthest_chosen}, other at {nearest_other}" + ); + } } #[test] @@ -679,79 +716,75 @@ mod tests { .collect(); let expected: Vec<_> = expected.chunks_exact(fanout).map(<[u32]>::to_vec).collect(); assert_eq!( - test_support::sorted_members_per_row(&actual), - test_support::sorted_members_per_row(&expected), + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), "points {ids:?}, fanout={fanout}" ); } } - #[rstest] - #[case::small_input(7)] - #[case::below_parallel_threshold(PARALLEL_SCATTER_MIN_POINTS - 1)] - #[case::at_parallel_threshold(PARALLEL_SCATTER_MIN_POINTS)] - #[case::uneven_parallel_chunks(PARALLEL_SCATTER_MIN_POINTS + 1)] - fn scatter_groups_assigned_points_by_leader_and_omits_unassigned_slots( - #[case] point_count: usize, - ) { - let ids: Vec<_> = (0..point_count).map(|i| 1_000_000 - i as u32).collect(); + #[test] + fn scatter_groups_assigned_points_by_leader_and_omits_unassigned_slots() { let pattern = [ [2, 0], [UNASSIGNED, 1], [0, UNASSIGNED], [UNASSIGNED, UNASSIGNED], ]; - let assignments: Vec<_> = (0..point_count).flat_map(|i| pattern[i % 4]).collect(); - let expected = vec![ - ids.iter().copied().step_by(2).collect::>(), - ids.iter().copied().skip(1).step_by(4).collect(), - ids.iter().copied().step_by(4).collect(), - vec![], - ]; - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(3) - .build() - .unwrap(); + let pool = thread_pool(3); + + // Counts around the parallel threshold cover both scatter paths and + // parallel chunks of unequal length. + for point_count in [ + 7, + PARALLEL_SCATTER_MIN_POINTS - 1, + PARALLEL_SCATTER_MIN_POINTS, + PARALLEL_SCATTER_MIN_POINTS + 1, + ] { + let ids: Vec<_> = (0..point_count).map(|i| 1_000_000 - i as u32).collect(); + let assignments: Vec<_> = (0..point_count).flat_map(|i| pattern[i % 4]).collect(); + let expected = [ + ids.iter().copied().step_by(2).collect(), + ids.iter().copied().skip(1).step_by(4).collect(), + ids.iter().copied().step_by(4).collect(), + vec![], + ]; - let actual = pool.install(|| scatter_assignments(&ids, &assignments, 2, 4)); + let actual = pool.install(|| scatter_assignments(&ids, &assignments, 2, 4)); - assert_eq!( - test_support::sorted_members_per_row(&actual), - test_support::sorted_members_per_row(&expected) - ); + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "{point_count} points" + ); + } } - #[rstest] - #[case::configured_overlap(0, vec![vec![0, 1], vec![0, 1], vec![2, 3], vec![2, 3]])] - #[case::after_schedule(1, vec![vec![0], vec![1], vec![2], vec![3]])] - fn a_split_uses_the_configured_fanout_then_falls_back_to_one_leader( - #[case] level: usize, - #[case] expected: Vec>, - ) { + #[test] + fn a_split_uses_the_configured_fanout_then_falls_back_to_one_leader() { + // Every point is a leader, so membership is fixed regardless of output order. let values = [0.0_f32, 1.0, 10.0, 11.0]; let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); - let config = splitting_config(); - let parent = PendingPartition { - point_ids: vec![0, 1, 2, 3], - level, - seed: 1290, - }; - let buffers = StripeBufferPool::new((), 0, None); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); - let actual = pool - .install(|| split_partition::<_, L2, _>(ARCH, data, &config, parent, &buffers)) - .unwrap(); + for (level, expected) in [ + (0, vec![vec![0, 1], vec![0, 1], vec![2, 3], vec![2, 3]]), + (1, vec![vec![0], vec![1], vec![2], vec![3]]), + ] { + let parent = PendingPartition { + point_ids: vec![0, 1, 2, 3], + level, + seed: 1290, + }; - // Every point is a leader, so membership is fixed regardless of output order. - assert_eq!( - sorted_leaf_memberships(&actual.leaves), - sorted_leaf_memberships(&expected) - ); - assert!(actual.pending.is_empty()); + let actual = split::(data, &splitting_config(), parent); + + assert_eq!( + sorted_leaf_memberships(&actual.leaves), + sorted_leaf_memberships(&expected), + "level {level}" + ); + assert!(actual.pending.is_empty(), "level {level}"); + } } #[test] @@ -767,15 +800,8 @@ mod tests { level: 0, seed: 1290, }; - let buffers = StripeBufferPool::new((), 0, None); - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); - let actual = pool - .install(|| split_partition::<_, L2, _>(ARCH, data, &config, parent, &buffers)) - .unwrap(); + let actual = split::(data, &config, parent); assert!(actual.leaves.is_empty()); assert!(actual.pending.iter().all(|child| child.level == 1)); @@ -790,36 +816,71 @@ mod tests { ); } - #[rstest] - #[case::below_limit(4)] - #[case::at_limit(3)] - fn data_within_cmax_forms_one_complete_leaf_per_replica(#[case] c_max: usize) { - let values = [0.0_f32, 1.0, 10.0]; - let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + #[test] + fn a_stalled_split_moves_the_leader_copies_into_leaves() { + // IDs 0 to 5 are copies of one vector, and ID 6 differs. When both + // sampled leaders are copies, all points tie and choose the same leader. + let values = [1.0_f32, 1.0, 1.0, 1.0, 1.0, 1.0, 9.0]; + let data = MatrixView::try_from(&values[..], 7, 1).unwrap(); + let ids: Vec = (0..7).collect(); let config = PiPNNConfig { - c_max, - replicas: 2, + c_max: 4, + p_samp: 0.01, + fanout: vec![1], ..splitting_config() }; - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() + // `split_partition` derives its sampling seed from the parent seed and size. + let seed = (0..) + .find(|&seed| { + let leaders = sample_leaders(&ids, config.p_samp, mix_seed(seed, 7)); + leaders.iter().all(|&id| id < 6) + }) .unwrap(); + let parent = PendingPartition { + point_ids: ids, + level: 0, + seed, + }; - let actual = pool - .install(|| partition::<_, L2, _>(ARCH, data, &config)) - .unwrap(); + let actual = split::(data, &config, parent); + // The copies fill two equal leaves. The distinct point is not grouped with them. assert_eq!( - sorted_leaf_memberships(&actual), - [vec![0, 1, 2], vec![0, 1, 2]] + sorted_leaf_memberships(&actual.leaves), + [vec![0, 1, 2], vec![3, 4, 5], vec![6]] ); + assert!(actual.pending.is_empty()); + } + + #[test] + fn data_within_cmax_forms_one_complete_leaf_per_replica() { + let values = [0.0_f32, 1.0, 10.0]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let pool = thread_pool(1); + + for c_max in [3, 4] { + let config = PiPNNConfig { + c_max, + replicas: 2, + ..splitting_config() + }; + + let actual = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + + assert_eq!( + sorted_leaf_memberships(&actual), + [vec![0, 1, 2], vec![0, 1, 2]], + "c_max {c_max}" + ); + } } #[test] fn recursive_partitioning_preserves_leaf_membership_across_worker_counts() { let point_count = 129; - let values = test_support::dense_points(point_count, 17, 1290); + let values = dense_points(point_count, 17, 1290); let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); let config = PiPNNConfig { c_max: 16, @@ -829,19 +890,11 @@ mod tests { replicas: 2, ..splitting_config() }; - let serial = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - let parallel = rayon::ThreadPoolBuilder::new() - .num_threads(3) - .build() - .unwrap(); - let first = serial + let first = thread_pool(1) .install(|| partition::<_, L2, _>(ARCH, data, &config)) .unwrap(); - let second = parallel + let second = thread_pool(3) .install(|| partition::<_, L2, _>(ARCH, data, &config)) .unwrap(); @@ -869,7 +922,7 @@ mod tests { #[test] fn a_second_replica_adds_new_leaf_memberships() { let point_count = 33; - let values = test_support::dense_points(point_count, 17, 1290); + let values = dense_points(point_count, 17, 1290); let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); let config = PiPNNConfig { c_max: 8, @@ -878,10 +931,7 @@ mod tests { fanout: vec![1], ..splitting_config() }; - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap(); + let pool = thread_pool(2); let first = pool .install(|| partition::<_, L2, _>(ARCH, data, &config)) .unwrap(); @@ -920,99 +970,82 @@ mod tests { assert_eq!(memberships, vec![2; point_count]); } - #[rstest] - #[case::one_unrankable([0.0, 3.0, f32::NAN], vec![vec![0], vec![1]])] - #[case::all_unrankable([f32::NAN; 3], vec![])] - fn unrankable_points_do_not_form_child_leaves( - #[case] values: [f32; 3], - #[case] expected: Vec>, - ) { - let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); - let config = PiPNNConfig { - c_max: 1, - fanout: vec![1], - ..splitting_config() - }; - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); - - let actual = pool - .install(|| partition::<_, L2, _>(ARCH, data, &config)) - .unwrap(); - - assert_eq!( - sorted_leaf_memberships(&actual), - sorted_leaf_memberships(&expected) - ); - } - #[test] - fn indistinguishable_points_report_the_oversized_cluster_when_splitting_stalls() { - let values = [1.0_f32; 5]; - let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + fn unrankable_points_do_not_form_child_leaves() { let config = PiPNNConfig { + c_max: 1, fanout: vec![1], ..splitting_config() }; - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(1) - .build() - .unwrap(); + let pool = thread_pool(1); - let error = pool - .install(|| partition::<_, L2, _>(ARCH, data, &config)) - .unwrap_err(); - - assert_eq!( - error.downcast_ref::().unwrap(), - &PartitionError::IterationLimit { - size: 5, - level: MAX_PARTITION_ITERATIONS, - limit: MAX_PARTITION_ITERATIONS, - } - ); - } + for (values, expected) in [ + ([0.0, 3.0, f32::NAN], vec![vec![0], vec![1]]), + ([f32::NAN; 3], vec![]), + ] { + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); - #[rstest] - #[case::empty(vec![], 2, 4, vec![])] - #[case::large_leaves_unchanged(vec![vec![1, 4], vec![0, 3, 6, 8]], 2, 4, - vec![vec![1, 4], vec![0, 3, 6, 8]])] - #[case::overlapping_small_leaves(vec![vec![4], vec![1], vec![4], vec![2]], 3, 4, - vec![vec![1, 2, 4]])] - #[case::flush_before_exceeding_maximum(vec![vec![4, 5], vec![0, 1]], 3, 3, - vec![vec![4, 5], vec![0, 1]])] - #[case::remainder_fits_after_deduplication(vec![vec![1, 3, 5], vec![3], vec![6]], 3, 4, - vec![vec![1, 3, 5, 6]])] - #[case::remainder_does_not_fit(vec![vec![1, 3, 5], vec![6], vec![8]], 3, 4, - vec![vec![1, 3, 5], vec![6, 8]])] - #[case::only_an_undersized_remainder(vec![vec![7], vec![2]], 3, 4, vec![vec![2, 7]])] - fn small_leaf_merging_deduplicates_ids_without_exceeding_cmax( - #[case] leaves: Vec>, - #[case] c_min: usize, - #[case] c_max: usize, - #[case] expected: Vec>, - ) { - let actual = merge_undersized_leaves(leaves, c_min, c_max); + let actual = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); - assert_eq!( - sorted_leaf_memberships(&actual), - sorted_leaf_memberships(&expected) - ); + assert_eq!( + sorted_leaf_memberships(&actual), + sorted_leaf_memberships(&expected), + "values {values:?}" + ); + } } #[test] - fn an_overflowing_buffer_area_identifies_the_failed_shape() { - let error = checked_area("point stripe", usize::MAX, 2).unwrap_err(); + fn small_leaf_merging_deduplicates_ids_without_exceeding_cmax() { + // (leaves, c_min, c_max, expected leaves) + for (leaves, c_min, c_max, expected) in [ + (vec![], 2, 4, vec![]), + ( + vec![vec![1, 4], vec![0, 3, 6, 8]], + 2, + 4, + vec![vec![1, 4], vec![0, 3, 6, 8]], + ), + // Overlapping small leaves merge into one leaf without repeated IDs. + ( + vec![vec![4], vec![1], vec![4], vec![2]], + 3, + 4, + vec![vec![1, 2, 4]], + ), + // A merge that would exceed c_max flushes the first leaf. + ( + vec![vec![4, 5], vec![0, 1]], + 3, + 3, + vec![vec![4, 5], vec![0, 1]], + ), + // After duplicate removal, the undersized remainder fits the last leaf. + ( + vec![vec![1, 3, 5], vec![3], vec![6]], + 3, + 4, + vec![vec![1, 3, 5, 6]], + ), + ( + vec![vec![1, 3, 5], vec![6], vec![8]], + 3, + 4, + vec![vec![1, 3, 5], vec![6, 8]], + ), + (vec![vec![7], vec![2]], 3, 4, vec![vec![2, 7]]), + ] { + let case = format!("{leaves:?}, c_min={c_min}, c_max={c_max}"); - assert_eq!( - error.downcast_ref::().unwrap(), - &PartitionError::ShapeOverflow { - buffer: "point stripe", - rows: usize::MAX, - cols: 2, - } - ); + let actual = merge_undersized_leaves(leaves, c_min, c_max); + + assert_eq!( + sorted_leaf_memberships(&actual), + sorted_leaf_memberships(&expected), + "{case}" + ); + } } } From 0c1c8d4776afa3622425a3d9f51b8edb32603dd9 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 28 Sep 2026 08:49:03 +0000 Subject: [PATCH 208/229] refactor(pipnn): skip unsplittable levels; drop merge and dead errors Partitioning: - When a level's fanout is not below its leader count, every point would join every leader and each child would repeat the cluster. Requeue the cluster at the next level instead of splitting it. - When a fanout-one split sends every point to one leader, cut the child into equal leaves in point order. Copy detection is gone. - Remove `merge_undersized_leaves` and `c_min`. SIFT-1M and Enron-1M builds (3 runs each) gave the same recall@10 at L=10..80 with and without the merge (largest difference 0.08 pp), and no isolated points. Configuration (review feedback): - `c_max`, `leaf_k`, `replicas` and the fanout entries are `NonZeroUsize`. - `PiPNNConfigError` replaces message errors for configuration and context checks. Finalization: - Keep the `u16::MAX` nearest candidates for RobustPrune, as Vamana limits its pool, instead of failing the build. `FinalizationError` is removed. Tests run under Miri when selected; the `not(miri)` gates are removed. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/finalization.rs | 102 +++--- diskann/src/graph/pipnn/leaf_build.rs | 88 +++-- diskann/src/graph/pipnn/mod.rs | 324 ++++++++---------- diskann/src/graph/pipnn/partitioning.rs | 436 ++++++++++-------------- 4 files changed, 398 insertions(+), 552 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 03d18c161b..c4c620352f 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -3,18 +3,14 @@ * Licensed under the MIT license. */ -//! Graph-degree enforcement with the Vamana RobustPrune kernel. +//! Final pruning: cut each candidate list to the graph degree with RobustPrune. //! -//! Candidate merging can produce more than the graph degree for one point. -//! Lists within the degree limit pass through unchanged. Longer lists are sorted -//! by source distance and passed to RobustPrune, then overwritten with the -//! selected IDs. Reusable per-job scratch stores candidates and pruning state. -//! -//! RobustPrune owns occlusion and alpha-round behavior. This module supplies -//! source vectors and metric distances. +//! A list within the degree passes through unchanged. A longer list is sorted by +//! distance to its source point and pruned with the shared Vamana RobustPrune, +//! which owns occlusion and the alpha rounds. The selected IDs overwrite the list +//! in place. use crate::{ - ANNError, ANNResult, graph::{ AdjacencyList, Config, internal::{SortedNeighbors, prune}, @@ -26,13 +22,7 @@ use diskann_utils::views::MatrixView; use diskann_vector::{DistanceFunction, distance::Metric}; use rayon::prelude::*; -#[derive(Debug, thiserror::Error)] -pub(crate) enum FinalizationError { - #[error("candidate count {actual} exceeds the u16 position limit {max}")] - TooManyCandidates { actual: usize, max: usize }, -} - -/// RobustPrune state for one Rayon job. +/// RobustPrune scratch for one Rayon job. #[derive(Default)] struct PruneWorkspace { candidates: Vec>, @@ -41,6 +31,7 @@ struct PruneWorkspace { } impl PruneWorkspace { + /// Prune the candidates of point `source` to the graph degree. fn prune( &mut self, data: MatrixView<'_, T>, @@ -48,17 +39,10 @@ impl PruneWorkspace { mut source_candidates: AdjacencyList, graph: &Config, distance: &T::Distance, - ) -> ANNResult> { + ) -> AdjacencyList { let degree = graph.pruned_degree().get(); - let candidate_count = source_candidates.len(); - if candidate_count <= degree { - return Ok(source_candidates); - } - if candidate_count > u16::MAX as usize { - return Err(ANNError::new(FinalizationError::TooManyCandidates { - actual: candidate_count, - max: u16::MAX as usize, - })); + if source_candidates.len() <= degree { + return source_candidates; } let source_vector = data.row(source); @@ -70,14 +54,17 @@ impl PruneWorkspace { distance.evaluate_similarity(source_vector, data.row(candidate as usize)), ) })); + // RobustPrune stores candidate positions as `u16`, so keep the `u16::MAX` + // nearest candidates. Vamana limits its candidate pool the same way. + // // Sort before mapping: changing the element type can change the unstable // sort's order for equal distances, and therefore RobustPrune's choices. - let sorted = SortedNeighbors::new(&mut self.candidates, candidate_count); + let sorted = SortedNeighbors::new(&mut self.candidates, u16::MAX as usize); let mapped = sorted.map_in(&mut self.mapped_candidates, |&candidate| { (candidate as usize != source).then_some(candidate) }); self.states.clear(); - self.states.resize(candidate_count, prune::State::default()); + self.states.resize(sorted.len(), prune::State::default()); let selected = prune::robust_prune( mapped, &mut self.states, @@ -94,22 +81,22 @@ impl PruneWorkspace { *destination = *sorted[state.neighbor as usize].id(); } output.finish(selected); - Ok(source_candidates) + source_candidates } } -/// Prune each candidate list that exceeds the graph degree. +/// Prune each candidate list that is longer than the graph degree. /// -/// Candidate builders supply one list per data row and valid dataset IDs. +/// `candidates` holds one list of dataset IDs for each data row. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, graph: &Config, metric: Metric, -) -> ANNResult>> { +) -> Vec> { let distance = T::distance(metric, Some(data.ncols())); - // The build context supplies the Rayon pool for this terminal operation. + // This runs in the pool of the build context (see `build_graph`). candidates .into_par_iter() .enumerate() @@ -149,9 +136,8 @@ mod tests { let data = MatrixView::try_from(values, values.len(), 1).unwrap(); let graph = pruning_config(degree, Metric::L2, alpha).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); - let pruned = PruneWorkspace::default() - .prune(data, 0, candidates_in_order(ids), &graph, &distance) - .unwrap(); + let pruned = + PruneWorkspace::default().prune(data, 0, candidates_in_order(ids), &graph, &distance); pruned.to_vec() } @@ -197,9 +183,13 @@ mod tests { let graph = pruning_config(2, metric, 1.0).unwrap(); let distance = f32::distance(metric, Some(2)); - let actual = PruneWorkspace::default() - .prune(data, 0, candidates_in_order(&[3, 2, 1]), &graph, &distance) - .unwrap(); + let actual = PruneWorkspace::default().prune( + data, + 0, + candidates_in_order(&[3, 2, 1]), + &graph, + &distance, + ); assert_eq!(&*actual, expected, "{metric:?}"); } @@ -229,7 +219,7 @@ mod tests { } #[test] - fn reused_workspace_does_not_carry_neighbors_or_occlusion_between_sources() { + fn a_reused_workspace_does_not_carry_state_between_sources() { let values = [0.0_f32, 1.0, 3.0, -2.0, 9.0]; let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); let graph = pruning_config(2, Metric::L2, 1.0).unwrap(); @@ -242,37 +232,24 @@ mod tests { (4, &[3, 2, 1][..], &[2][..]), (1, &[4, 3, 2, 0][..], &[0, 2][..]), ] { - let actual = workspace - .prune(data, source, candidates_in_order(ids), &graph, &distance) - .unwrap(); + let actual = workspace.prune(data, source, candidates_in_order(ids), &graph, &distance); assert_eq!(&*actual, expected, "source {source}"); } } - #[cfg(not(miri))] #[test] - fn the_candidate_position_limit_accepts_u16_max_and_rejects_one_more() { - let values: Vec<_> = (0..=u16::MAX as usize + 1).map(|i| i as f32).collect(); - let data = MatrixView::try_from(values.as_slice(), values.len(), 1).unwrap(); + fn a_list_above_the_u16_position_limit_keeps_its_nearest_candidates() { + let count = u16::MAX as usize + 1; + let values: Vec<_> = (0..=count).map(|i| i as f32).collect(); + let data = MatrixView::try_from(values.as_slice(), count + 1, 1).unwrap(); let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); + let candidates = AdjacencyList::from_iter_untrusted(1..=count as u32); - let largest = AdjacencyList::from_iter_untrusted(1..=u16::MAX as u32); - let actual = PruneWorkspace::default() - .prune(data, 0, largest, &graph, &distance) - .unwrap(); - assert_eq!(&*actual, [1]); + let actual = PruneWorkspace::default().prune(data, 0, candidates, &graph, &distance); - let too_many = AdjacencyList::from_iter_untrusted(1..=u16::MAX as u32 + 1); - let error = PruneWorkspace::default() - .prune(data, 0, too_many, &graph, &distance) - .unwrap_err(); - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::TooManyCandidates { actual, max }) - if *actual == u16::MAX as usize + 1 && *max == u16::MAX as usize - )); + assert_eq!(&*actual, [1]); } #[test] @@ -287,8 +264,7 @@ mod tests { .to_vec(); let actual = thread_pool(workers) - .install(|| prune_overfull(data, candidates, &graph, Metric::L2)) - .unwrap(); + .install(|| prune_overfull(data, candidates, &graph, Metric::L2)); assert_eq!( actual.into_iter().map(Vec::from).collect::>(), diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 5604114db1..4c777be7ac 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -3,19 +3,14 @@ * Licensed under the MIT license. */ -//! Leaf-local graph construction and candidate accumulation. +//! Leaf construction: select the nearest neighbors inside each leaf and merge +//! them into one candidate list per data point. //! -//! Partitioning supplies sorted, unique global point IDs for each leaf. One leaf -//! job does these steps: -//! -//! 1. Gather each ID and convert its vector to reusable `f32` storage. -//! 2. Call the leaf kernel for ranking-distance construction and local ranking. -//! 3. Convert local positions to global point IDs. -//! 4. Add both edge directions to global candidate lists. -//! -//! Overlapping leaves run concurrently. A worker locks one destination list only -//! while it adds one leaf's IDs. Reusable buffers keep their largest allocation. -//! Each operation uses an explicit active prefix. +//! A leaf job gathers the vectors of its points as `f32`. The leaf kernel then +//! selects the `k` nearest leaf points of each point, and the job adds each +//! selected pair in both directions. Leaves overlap and run in parallel, so each +//! data point has a locked candidate list. A job groups its edges by point first +//! and then locks each list once. use parking_lot::Mutex; @@ -31,10 +26,9 @@ use super::{ topk::Candidate, }; -/// Reusable buffers for one Rayon leaf job. +/// Scratch for one Rayon leaf job. /// -/// The numerical vectors keep the largest leaf shape that this job observed. -/// The job creates local adjacency lists only when the effective `k` is not zero. +/// Each buffer keeps the largest size that the job needed. A leaf uses a prefix. #[derive(Default)] struct LeafBuffers { point_values: Vec, @@ -44,7 +38,8 @@ struct LeafBuffers { } impl LeafBuffers { - /// Grow the buffers for one leaf and return its effective `k`. + /// Grow the buffers for a leaf of `point_count` points and return the + /// effective `k` of the leaf. fn prepare(&mut self, point_count: usize, dimension_count: usize, requested_k: usize) -> usize { // A point has at most `point_count - 1` other points in its leaf. Wider rows // would hold only empty slots. @@ -58,6 +53,7 @@ impl LeafBuffers { leaf_k } + /// Clear the edge lists of the first `point_count` leaf points. fn prepare_local_adjacency(&mut self, point_count: usize) { if self.local_adjacency.len() < point_count { self.local_adjacency.resize_with(point_count, Vec::new); @@ -68,10 +64,12 @@ impl LeafBuffers { } } -/// Build direct graph candidates from all overlapping leaves. +/// Build one candidate list per data point from the nearest neighbors in all +/// leaves. /// -/// Each selected leaf pair contributes both edge directions. Candidate lists use -/// global dataset IDs and contain no duplicate IDs. +/// Each selected pair adds both directions. A list holds global IDs without +/// duplicates. The lists are sorted, so the result does not depend on the order +/// in which the parallel jobs finish. pub(super) fn build_leaf_candidates( arch: A, data: MatrixView<'_, T>, @@ -89,14 +87,7 @@ where leaves .par_iter() .try_for_each_init(LeafBuffers::default, |buffers, point_ids| { - add_direct_leaf_candidates::( - arch, - data, - point_ids, - requested_k, - buffers, - &candidates, - ) + add_leaf_candidates::(arch, data, point_ids, requested_k, buffers, &candidates) })?; Ok(candidates .into_iter() @@ -108,12 +99,11 @@ where .collect()) } -/// Add one leaf's symmetric neighbors to the direct candidate lists. +/// Add the selected pairs of one leaf to the candidate lists. /// -/// Reusable buffers can be longer than this leaf, so all accesses use the current -/// leaf shape. Leaf IDs are distinct dataset rows, so `points x dimensions` is no +/// The IDs of a leaf are distinct dataset rows, so `points x dimensions` is not /// larger than the dataset. -fn add_direct_leaf_candidates( +fn add_leaf_candidates( arch: A, data: MatrixView<'_, T>, point_ids: &[u32], @@ -158,6 +148,8 @@ where Ok(()) } +/// Convert each selected pair from leaf positions to global IDs and add it to +/// the edge lists of both points. fn add_symmetric_neighbors( point_ids: &[u32], leaf_k: usize, @@ -170,11 +162,10 @@ fn add_symmetric_neighbors( continue; } let target = neighbor.local_idx as usize; - let source_id = point_ids[source]; - let target_id = point_ids[target]; - // Leaves contain unique IDs, and the kernel excludes each point itself. - local_adjacency[source].push(target_id); - local_adjacency[target].push(source_id); + // The kernel never selects a point for itself, and the IDs of a leaf + // are unique, so no list gets a self edge. + local_adjacency[source].push(point_ids[target]); + local_adjacency[target].push(point_ids[source]); } } } @@ -185,7 +176,7 @@ fn grow(values: &mut Vec, len: usize, value: T) { } } -#[cfg(all(test, not(miri)))] +#[cfg(test)] mod tests { use super::*; use crate::graph::pipnn::{ @@ -195,7 +186,6 @@ mod tests { use diskann_wide::ARCH; use std::collections::BTreeSet; - #[expect(clippy::unwrap_used, reason = "an invalid fixture fails the test")] fn build(values: &[f32], leaves: Vec>, k: usize, workers: usize) -> Vec> { let data = MatrixView::try_from(values, values.len(), 1).unwrap(); let candidates = thread_pool(workers) @@ -205,7 +195,7 @@ mod tests { } #[test] - fn selected_neighbors_use_global_ids_and_contribute_both_edge_directions() { + fn selected_pairs_add_global_ids_in_both_directions() { // The leaf lists IDs [5, 1, 3], at coordinates [9, 0, 2]. // The directed choices are 1 -> 3, 3 -> 1 and 5 -> 3. let values = [100.0_f32, 0.0, -100.0, 2.0, 200.0, 9.0]; @@ -235,7 +225,7 @@ mod tests { } #[test] - fn leaves_without_selected_pairs_produce_empty_adjacency() { + fn leaves_without_pairs_add_no_candidates() { let values = [0.0_f32, 1.0, 4.0]; for (leaves, k) in [ @@ -253,7 +243,7 @@ mod tests { } #[test] - fn requesting_more_neighbors_than_a_leaf_has_selects_all_other_points() { + fn k_above_the_leaf_size_selects_every_other_point() { let values = [100.0_f32, 0.0, -100.0, 2.0, 200.0, 9.0]; let actual = build(&values, vec![vec![1, 3, 5]], 99, 1); @@ -265,7 +255,7 @@ mod tests { } #[test] - fn unrankable_pairs_do_not_add_unassigned_ids_to_the_graph() { + fn unrankable_pairs_add_no_candidates() { let values = [0.0_f32, 3.0, f32::NAN]; let actual = build(&values, vec![vec![0, 1, 2]], 2, 1); @@ -274,11 +264,11 @@ mod tests { } #[test] - fn candidates_at_production_shape_are_the_symmetric_leaf_kernel_results() { - // Sixty overlapping leaves of 17 to 40 points run on four workers with - // reused buffers. They fill full SIMD groups of the leaf kernel. k = 2 - // uses a fixed top-k width, and k = 5 uses the runtime width. The kernel - // on each leaf alone is the reference, so tie order cannot differ. + fn candidates_match_the_leaf_kernel_on_each_leaf() { + // The reference runs the leaf kernel on each leaf alone, so the tie order + // is the same. Sixty overlapping leaves of 17 to 40 points run on four + // workers with reused buffers and fill whole SIMD groups. k = 2 uses a + // fixed top-k width, and k = 5 uses the runtime width. let (point_count, dimensions) = (500, 24); let values = dense_points(point_count, dimensions, 1290); let data = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); @@ -331,7 +321,7 @@ mod tests { } #[test] - fn reused_leaf_buffers_do_not_carry_edges_between_different_leaf_shapes() { + fn reused_buffers_do_not_carry_edges_between_leaves() { let values = [0.0_f32, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0]; let data = MatrixView::try_from(&values[..], 7, 1).unwrap(); let mut buffers = LeafBuffers::default(); @@ -373,7 +363,7 @@ mod tests { ] { let candidates: Vec<_> = (0..7).map(|_| Mutex::new(AdjacencyList::new())).collect(); - add_direct_leaf_candidates::<_, L2, _>( + add_leaf_candidates::<_, L2, _>( ARCH, data, &ids, diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 6e1bd9304e..ffb353ea1c 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,31 +5,23 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! PiPNN builds graph candidates in three steps. A leader is a sampled dataset -//! point that acts as the center of one child partition. A leaf is a bounded -//! child partition used for local neighbor selection. +//! PiPNN builds a graph without graph search, in three steps. A leaf is a small +//! cluster of at most `c_max` points; the build selects neighbors only inside +//! leaves. //! -//! 1. `partitioning` samples leaders and makes overlapping leaves of at most -//! `c_max` points. Points without rankable leaders contribute no assignments. -//! 2. `leaf_build` computes one ranking-distance buffer for each leaf. It -//! selects local neighbors and merges their global point IDs. -//! 3. `finalization` applies Vamana RobustPrune to each candidate list that is -//! longer than the graph degree. +//! 1. `partitioning` splits the dataset into overlapping leaves. +//! 2. `leaf_build` selects the `leaf_k` nearest neighbors of each point inside +//! each leaf and merges them into one candidate list per point. +//! 3. `finalization` prunes each list that is longer than the graph degree with +//! Vamana RobustPrune. //! -//! `diskann-wide` selects architecture `A` for the ranking kernels. One match -//! selects metric marker `M`. Metric computation stays architecture-neutral. -//! Ranking kernels reapply architecture features inside Rayon jobs. +//! [`build_graph`] selects the SIMD architecture and the metric once, so the +//! steps compile for one architecture and one metric. It borrows one contiguous +//! [`MatrixView`] and returns one adjacency list for each point. Each step takes +//! ownership of the previous output and frees it when it returns. //! -//! [`PiPNNConfig`] contains partition and local-neighbor parameters. -//! [`PiPNNBuildContext`] borrows graph policy and a Rayon pool. [`build_graph`] -//! borrows one contiguous [`MatrixView`]. It returns one adjacency list for each -//! input point. -//! -//! The function does not load providers or select start and frozen points. It -//! also does not quantize, serialize, or search the graph. -//! -//! Partition and leaf work use separate reusable buffers. The build consumes -//! each output before it creates another graph representation. +//! The build does not load providers or select start points. It also does not +//! quantize, serialize, or search the graph. mod conversion; mod finalization; @@ -42,9 +34,11 @@ mod partitioning; mod simd; mod topk; +use std::num::NonZeroUsize; + use crate::{ ANNError, ANNResult, - graph::{AdjacencyList, Config}, + graph::{AdjacencyList, Config, config::PruneKind}, utils::VectorRepr, }; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -185,64 +179,55 @@ mod cosine_distance_tests { } } -/// PiPNN partition and leaf-selection policy. +/// Error for invalid PiPNN parameters. +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum PiPNNConfigError { + #[error("p_samp ({0}) must be in (0, 1]")] + SamplingFraction(f64), + #[error("fanout must not be empty")] + EmptyFanout, + #[error("graph prune kind {prune_kind:?} does not match metric {metric:?}")] + PruneKind { + prune_kind: PruneKind, + metric: Metric, + }, +} + +crate::convert_error!(PiPNNConfigError); + +/// PiPNN partition and leaf parameters. /// -/// DiskANN graph policy separately supplies degree, alpha, and prune metric. +/// The DiskANN graph [`Config`] supplies the degree, alpha, and prune kind. #[derive(Clone, Debug, PartialEq)] pub struct PiPNNConfig { /// Maximum number of points in a leaf. - pub c_max: usize, - /// Minimum leaf size used by global small-leaf merging. - pub c_min: usize, - /// Fraction of a cluster sampled as child-partition centers. + pub c_max: NonZeroUsize, + /// Fraction of the points of a cluster that a split samples as leaders, in + /// `(0, 1]`. pub p_samp: f64, - /// Number of nearest centers assigned at each recursive partition level. - /// Levels after this schedule assign each point to one center. - pub fanout: Vec, - /// Number of nearest neighbors selected within each leaf. - pub leaf_k: usize, - /// Number of independent partition passes over the dataset. - pub replicas: usize, + /// Number of nearest leaders that each point joins, for each split level. + /// Levels after this schedule use one leader. The schedule must not be empty. + pub fanout: Vec, + /// Number of nearest neighbors that each point selects inside a leaf. + pub leaf_k: NonZeroUsize, + /// Number of independent partitionings of the dataset. + pub replicas: NonZeroUsize, } impl PiPNNConfig { - /// Validate the algorithm-specific partition and leaf-build parameters. - pub fn validate(&self) -> ANNResult<()> { - if self.c_max == 0 { - return Err(config_error("c_max must be greater than zero")); - } - if self.c_min == 0 { - return Err(config_error("c_min must be greater than zero")); - } - if self.c_min > self.c_max { - return Err(config_error(format!( - "c_min ({}) must not exceed c_max ({})", - self.c_min, self.c_max - ))); - } + /// Check the parameters that the field types do not constrain. + pub fn validate(&self) -> Result<(), PiPNNConfigError> { if !(0.0 < self.p_samp && self.p_samp <= 1.0) { - return Err(config_error(format!( - "p_samp ({}) must be in (0, 1]", - self.p_samp - ))); + return Err(PiPNNConfigError::SamplingFraction(self.p_samp)); } if self.fanout.is_empty() { - return Err(config_error("fanout must not be empty")); - } - if self.fanout.contains(&0) { - return Err(config_error("fanout values must be greater than zero")); - } - if self.leaf_k == 0 { - return Err(config_error("leaf_k must be greater than zero")); - } - if self.replicas == 0 { - return Err(config_error("replicas must be greater than zero")); + return Err(PiPNNConfigError::EmptyFanout); } Ok(()) } } -/// PiPNN policy and borrowed execution resources for one graph build. +/// PiPNN parameters, graph policy, and the Rayon pool for one graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { config: PiPNNConfig, @@ -252,19 +237,20 @@ pub struct PiPNNBuildContext<'a> { } impl<'a> PiPNNBuildContext<'a> { - /// Check and combine PiPNN configuration with DiskANN graph policy. + /// Check `config` and combine it with the graph policy. + /// + /// Final pruning applies the prune kind of `graph` to distances of `metric`, + /// so the two must match. pub fn new( config: PiPNNConfig, graph: &'a Config, metric: Metric, pool: &'a ThreadPool, - ) -> ANNResult { + ) -> Result { config.validate()?; - if graph.prune_kind() != metric.into() { - return Err(config_error(format!( - "graph prune kind {:?} is incompatible with metric {metric:?}", - graph.prune_kind() - ))); + let prune_kind = graph.prune_kind(); + if prune_kind != metric.into() { + return Err(PiPNNConfigError::PruneKind { prune_kind, metric }); } Ok(Self { @@ -276,13 +262,11 @@ impl<'a> PiPNNBuildContext<'a> { } } -/// Build one PiPNN adjacency list for each point in `data`. -/// -/// This graph contains only real dataset points. Start-point selection and index -/// serialization are separate operations. +/// Build one adjacency list for each point of `data`. /// -/// Raw `u8` and `i8` vectors are not unit-normalized after conversion to `f32`. -/// The build therefore uses norm-aware cosine for these two input types. +/// The graph links dataset points only; the caller selects start points and +/// serializes the index. `CosineNormalized` on `u8` or `i8` data builds with +/// `Cosine`, because converted integer vectors do not have unit norm. pub fn build_graph( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, @@ -312,14 +296,15 @@ where )); } if data.nrows() > u32::MAX as usize { - return Err(config_error(format!( - "dataset point count ({}) exceeds the u32 graph ID limit", + return Err(ANNError::message(format!( + "PiPNN dataset point count ({}) exceeds the u32 graph ID limit", data.nrows() ))); } arch::dispatch2_no_features(BuildGraph, data, context) } +/// Runs the build for the architecture that `diskann-wide` selects at run time. struct BuildGraph; impl Target2>>, MatrixView<'_, T>, &PiPNNBuildContext<'_>> @@ -349,10 +334,7 @@ where } } -/// Run the PiPNN graph pipeline for one selected metric implementation. -/// -/// The function builds overlapping leaves, merges direct candidates, and applies -/// final graph-degree pruning. +/// Run the three build steps for architecture `A` and metric `M`. fn build_graph_for( arch: A, data: MatrixView<'_, T>, @@ -366,17 +348,17 @@ where { let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition::(arch, data, &context.config))?; - // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, - // so its complete allocation drops when leaf construction returns. - let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) - })?; - // Finalization consumes each candidate list. It reuses that list's allocation - // for the final adjacency when the graph policy permits it. - tracing::info_span!("pipnn.finalization") - .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) + // The leaf step takes the leaves by value, so their memory is free before + // final pruning starts. + let leaf_k = context.config.leaf_k.get(); + let candidates = tracing::info_span!("pipnn.leaf_build") + .in_scope(|| leaf_build::build_leaf_candidates::(arch, data, leaves, leaf_k))?; + // Final pruning cuts each candidate list in place. + Ok(tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric))) } +/// Return the metric that the build uses for element type `T`. See [`build_graph`]. fn effective_metric(metric: Metric) -> Metric { use std::any::TypeId; @@ -389,19 +371,14 @@ fn effective_metric(metric: Metric) -> Metric { } } -#[track_caller] -fn config_error(message: impl std::fmt::Display) -> ANNError { - ANNError::message(format!("PiPNN configuration: {message}")) -} - #[cfg(test)] mod test_support { use super::simd::Simd; use diskann_vector::distance::Metric; - // Ignore member order while preserving group positions and duplicate counts. - // Positions identify leaders or graph sources; sorting the outer list would - // hide assignments to the wrong group. Copies keep the observations intact. + // Sort the members of each row but keep the row order. A row position is a + // leader column or a graph source, so sorting the rows would hide a wrong + // assignment. Repeated members stay visible. pub(super) fn sorted_members_per_row(rows: &[Vec]) -> Vec> { rows.iter() .map(|row| { @@ -412,6 +389,10 @@ mod test_support { .collect() } + pub(super) fn nz(value: usize) -> std::num::NonZeroUsize { + std::num::NonZeroUsize::new(value).unwrap() + } + /// Build a Rayon pool with `threads` workers. pub(super) fn thread_pool(threads: usize) -> rayon::ThreadPool { rayon::ThreadPoolBuilder::new() @@ -582,21 +563,20 @@ mod test_support { } } -#[cfg(all(test, not(miri)))] +#[cfg(test)] mod construction_tests { use super::*; use crate::graph::config; use half::f16; - use test_support::{sorted_members_per_row, thread_pool}; + use test_support::{nz, sorted_members_per_row, thread_pool}; fn partition_policy() -> PiPNNConfig { PiPNNConfig { - c_max: 4, - c_min: 1, + c_max: nz(4), p_samp: 1.0, - fanout: vec![2], - leaf_k: 1, - replicas: 1, + fanout: vec![nz(2)], + leaf_k: nz(1), + replicas: nz(1), } } @@ -608,7 +588,6 @@ mod construction_tests { } /// Build the graph of `values` with `threads` workers. - #[expect(clippy::unwrap_used, reason = "an invalid fixture fails the test")] fn build( values: &[T], dimensions: usize, @@ -626,77 +605,50 @@ mod construction_tests { } #[test] - fn invalid_configuration_identifies_the_rejected_condition() { - let invalid = |change: fn(&mut PiPNNConfig)| { - let mut config = partition_policy(); - change(&mut config); - config - }; - for (config, expected_message) in [ - (invalid(|c| c.c_max = 0), "c_max must be greater than zero"), - (invalid(|c| c.c_min = 0), "c_min must be greater than zero"), - ( - invalid(|c| c.c_min = 5), - "c_min (5) must not exceed c_max (4)", - ), - (invalid(|c| c.p_samp = 0.0), "p_samp (0) must be in (0, 1]"), - ( - invalid(|c| c.p_samp = -0.5), - "p_samp (-0.5) must be in (0, 1]", - ), - ( - invalid(|c| c.p_samp = 1.0 + f64::EPSILON), - "must be in (0, 1]", - ), - ( - invalid(|c| c.p_samp = f64::NAN), - "p_samp (NaN) must be in (0, 1]", - ), - ( - invalid(|c| c.p_samp = f64::INFINITY), - "p_samp (inf) must be in (0, 1]", - ), - ( - invalid(|c| c.p_samp = f64::NEG_INFINITY), - "p_samp (-inf) must be in (0, 1]", - ), - (invalid(|c| c.fanout.clear()), "fanout must not be empty"), - ( - invalid(|c| c.fanout = vec![2, 0, 1]), - "fanout values must be greater than zero", - ), - ( - invalid(|c| c.leaf_k = 0), - "leaf_k must be greater than zero", - ), - ( - invalid(|c| c.replicas = 0), - "replicas must be greater than zero", - ), + fn invalid_sampling_fractions_and_an_empty_fanout_are_rejected() { + for p_samp in [ + 0.0, + -0.5, + 1.0 + f64::EPSILON, + f64::NAN, + f64::INFINITY, + f64::NEG_INFINITY, ] { - let message = config.validate().unwrap_err().to_string(); + let config = PiPNNConfig { + p_samp, + ..partition_policy() + }; + + let error = config.validate().unwrap_err(); - assert!(message.contains(expected_message), "{message}"); + assert!( + matches!(error, PiPNNConfigError::SamplingFraction(value) if value.to_bits() == p_samp.to_bits()), + "p_samp {p_samp}: {error}" + ); } + let config = PiPNNConfig { + fanout: vec![], + ..partition_policy() + }; + assert_eq!(config.validate(), Err(PiPNNConfigError::EmptyFanout)); } #[test] - fn equal_leaf_bounds_and_valid_probability_endpoints_are_accepted() { + fn sampling_fraction_endpoints_are_accepted() { for p_samp in [f64::from_bits(1), 1.0] { let config = PiPNNConfig { - c_min: 4, p_samp, ..partition_policy() }; - config.validate().unwrap(); + assert_eq!(config.validate(), Ok(()), "p_samp {p_samp}"); } } #[test] - fn a_build_context_rejects_an_invalid_partition_policy() { + fn a_build_context_checks_the_configuration() { let config = PiPNNConfig { - leaf_k: 0, + fanout: vec![], ..partition_policy() }; let graph = graph_policy(2, Metric::L2).unwrap(); @@ -704,12 +656,7 @@ mod construction_tests { let error = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap_err(); - assert!( - error - .to_string() - .contains("leaf_k must be greater than zero"), - "{error}" - ); + assert_eq!(error, PiPNNConfigError::EmptyFanout); } #[test] @@ -731,13 +678,17 @@ mod construction_tests { let case = format!("graph {graph_metric:?}, metric {metric:?}"); match result { Ok(_) => assert!(compatible, "{case}"), - Err(error) => { - assert!(!compatible, "{case}: {error}"); - assert!( - error.to_string().contains("incompatible with metric"), - "{case}: {error}" - ); - } + Err(error) => assert_eq!( + (compatible, error), + ( + false, + PiPNNConfigError::PruneKind { + prune_kind: graph_metric.into(), + metric + } + ), + "{case}" + ), } } } @@ -826,7 +777,7 @@ mod construction_tests { // Every two-point leaf contributes its only pair; degree three retains // every edge. The graph therefore exposes partition membership directly. let config = PiPNNConfig { - c_max: 2, + c_max: nz(2), ..partition_policy() }; for (metric, expected) in [ @@ -866,7 +817,7 @@ mod construction_tests { fn check(values: [T; 8]) { for (leaf_k, degree) in [(1, 3), (3, 1)] { let config = PiPNNConfig { - leaf_k, + leaf_k: nz(leaf_k), ..partition_policy() }; @@ -886,9 +837,9 @@ mod construction_tests { } #[test] - fn final_pruning_enforces_graph_degree_after_leaf_candidates_are_merged() { + fn final_pruning_cuts_merged_candidates_to_the_degree() { let config = PiPNNConfig { - leaf_k: 2, + leaf_k: nz(2), ..partition_policy() }; @@ -903,8 +854,8 @@ mod construction_tests { // of each close pair. Neither replicas nor worker count may duplicate edges. for (workers, replicas) in [(1, 1), (3, 2)] { let config = PiPNNConfig { - c_max: 2, - replicas, + c_max: nz(2), + replicas: nz(replicas), ..partition_policy() }; @@ -940,12 +891,11 @@ mod construction_tests { // graph. let (points, dimensions, degree) = (1000, 16, 8); let config = PiPNNConfig { - c_max: 32, - c_min: 8, + c_max: nz(32), p_samp: 0.05, - fanout: vec![4, 2], - leaf_k: 2, - replicas: 1, + fanout: vec![nz(4), nz(2)], + leaf_k: nz(2), + replicas: nz(1), }; for metric in [ Metric::L2, diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 61528c943a..cf5d07404b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -3,27 +3,26 @@ * Licensed under the MIT license. */ -//! Deterministic overlapping partition construction for PiPNN. +//! Randomized ball carving: recursive, overlapping partitioning of the dataset +//! into leaves of at most `c_max` points. //! -//! A leader is a sampled point that acts as the center of one child partition. -//! A point can join several leaders, so child partitions can overlap. +//! One split samples leaders from a cluster and assigns each point to its +//! `fanout` nearest leaders. The points of one leader form a child cluster, so +//! children overlap when `fanout > 1`. A child with at most `c_max` points is a +//! leaf; a larger child is split again at the next level. `fanout[level]` gives +//! the fanout of each level, and levels after the schedule use fanout 1. Each +//! replica repeats the process with a different seed. //! -//! This module owns recursive splitting, leader sampling, row gathering, and -//! assignment scatter. The partition kernel owns GEMM, norm preparation, and -//! local ranking. An `ObjectPool` supplies reusable scratch to Rayon workers. +//! The partition kernel computes point-to-leader distances and selects the +//! nearest leaders. This module samples leaders, gathers vectors, and groups the +//! assignments into children. //! -//! A configured level assigns each point to `fanout[level]` leaders. A deeper -//! level assigns each point to one leader. Each replica uses a different -//! deterministic seed. -//! -//! Buffer sizes use plain multiplication. Point and leader IDs in one split are -//! distinct dataset rows, so a gathered `rows x dimensions` buffer is no larger -//! than the dataset. An assignment buffer holds `points x fanout` IDs. With at -//! most `u32::MAX` points and `LEADER_CAP` leaders, this count fits in a 64-bit +//! Buffer sizes use plain multiplication. The IDs of one split are distinct +//! dataset rows, so a gathered `rows x dimensions` buffer is not larger than the +//! dataset. An assignment buffer holds `points x fanout` IDs. With at most +//! `u32::MAX` points and `LEADER_CAP` leaders, this count fits in a 64-bit //! `usize`. -use std::collections::HashSet; - use crate::{ANNResult, utils::VectorRepr}; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, @@ -70,17 +69,23 @@ const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; /// machines. Below this size, a serial scatter takes a few milliseconds. const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; +/// A cluster with more than `c_max` points that waits for its split. struct PendingPartition { point_ids: Vec, + /// Position in the fanout schedule. level: usize, + /// Seed of the parent split. The next split mixes it with the cluster size. seed: u64, } +/// The result of one split: leaves, and clusters for the next level. struct PartitionSplit { pending: Vec, leaves: Vec>, } +/// Scratch for one assignment stripe: the gathered `f32` points and the kernel +/// workspace. #[derive(Default)] struct StripeBuffers { point_values: Vec, @@ -93,23 +98,21 @@ impl AsPooled<()> for StripeBuffers { } fn modify(&mut self, _: ()) { - // Keep the largest allocation across leases. `assign_point_stripe` defines the - // active prefix before each read. + // Keep the largest allocation across leases. `assign_point_stripe` sets + // the active prefix before each read. } } -/// Reusable buffers for point-to-leader assignment. +/// Stripe scratch shared by all splits of a build. /// -/// `ObjectPool` locks only when it gives or receives a lease. Numerical work -/// holds the lease, not the pool lock. +/// `ObjectPool` locks only to give or take back a lease. Numerical work holds the +/// lease, not the pool lock. type StripeBufferPool = ObjectPool; -/// Build overlapping bounded leaves for all configured replicas. +/// Build the leaves of all replicas. /// -/// Each split samples partition centers and assigns every cluster point to its -/// nearest centers. A cluster above `c_max` is split again. A level without a -/// configured fanout assigns each point to one center. Non-rankable distances -/// do not produce assignments. +/// A point whose distances to all leaders of a split are not rankable, for +/// example a NaN vector, joins no child of that split. pub(super) fn partition( arch: A, data: MatrixView<'_, T>, @@ -122,7 +125,7 @@ where { let mut leaves = Vec::new(); let stripe_buffers = StripeBufferPool::new((), 0, None); - for replica in 0..config.replicas { + for replica in 0..config.replicas.get() { let seed = replica_seed(replica); let mut replica_leaves = partition_replica::(arch, data, config, seed, &stripe_buffers)?; @@ -131,15 +134,12 @@ where Ok(leaves) } -/// Partition one replica until each leaf has at most `c_max` points. -/// -/// The function processes one work queue per recursion level. It merges leaves -/// smaller than `c_min` after the queue becomes empty. +/// Split one replica, one level at a time, until every cluster is a leaf. /// -/// The queue empties because the fanout schedule is finite. After the schedule, -/// a split puts each point into at most one child. [`split_partition`] makes sure -/// that such a child is smaller than its parent, so the cluster sizes decrease at -/// each later level. +/// The loop ends. The fanout schedule is finite, and each split sends its +/// clusters one level deeper. After the schedule, fanout is 1, so each point +/// joins one child. A child is then smaller than its parent, or it holds every +/// point and [`split_partition`] cuts it into leaves. fn partition_replica( arch: A, data: MatrixView<'_, T>, @@ -153,7 +153,8 @@ where T: VectorRepr + Send + Sync, { let initial_point_ids = (0..data.nrows() as u32).collect(); - if data.nrows() <= config.c_max { + let c_max = config.c_max.get(); + if data.nrows() <= c_max { return Ok(vec![initial_point_ids]); } @@ -164,7 +165,7 @@ where seed, }]; while !pending.is_empty() { - // Indexed parallel collection preserves parent-partition order. + // Indexed parallel collection keeps the parent order. let splits: ANNResult> = pending .into_par_iter() .map(|partition| { @@ -179,13 +180,10 @@ where } pending = next_level; } - Ok(merge_undersized_leaves(leaves, config.c_min, config.c_max)) + Ok(leaves) } -/// Split one oversized cluster into child partitions. -/// -/// The function samples center points, assigns the cluster points, and returns -/// bounded leaves separately from child clusters that need another split. +/// Split one cluster into leaves and clusters for the next level. fn split_partition( arch: A, data: MatrixView<'_, T>, @@ -199,8 +197,26 @@ where T: VectorRepr + Send + Sync, { let split_seed = mix_seed(partition.seed, partition.point_ids.len() as u64); - let fanout = config.fanout.get(partition.level).copied().unwrap_or(1); - let leaders = sample_leaders(&partition.point_ids, config.p_samp, split_seed); + let fanout = config + .fanout + .get(partition.level) + .map_or(1, |fanout| fanout.get()); + let c_max = config.c_max.get(); + let leader_count = sampled_leader_count(partition.point_ids.len(), config.p_samp); + if fanout > 1 && fanout >= leader_count { + // Every point would join every leader, and each child would repeat this + // cluster. Split the cluster at the next level instead. + return Ok(PartitionSplit { + pending: vec![PendingPartition { + point_ids: partition.point_ids, + level: partition.level + 1, + seed: split_seed, + }], + leaves: Vec::new(), + }); + } + + let leaders = sample_leaders(&partition.point_ids, leader_count, split_seed); let clusters = assign_to_leaders::( arch, data, @@ -212,32 +228,21 @@ where let mut pending = Vec::new(); let mut leaves = Vec::new(); - for (&leader, mut cluster) in leaders.iter().zip(clusters) { - if fanout == 1 && cluster.len() == partition.point_ids.len() { - // Every point chose this leader, so the split made no progress. Under - // L2, each leader is nearest to itself. Thus this occurs only when all - // sampled leaders are copies of one vector. Under cosine, vectors with - // the same direction also tie. Under inner product, one large leader - // can be first for every point. - // - // Copies of the leader have the same distance to every point, so any - // grouping of them gives the same leaf neighbors. Put the copies into - // leaves of equal size. The rest does not contain the leader, so it is - // smaller than its parent. Each stalled level therefore removes at - // least one point. - let copy = bytemuck::cast_slice::(data.row(leader as usize)); - let (copies, rest): (Vec<_>, Vec<_>) = cluster - .into_iter() - .partition(|&id| bytemuck::cast_slice::(data.row(id as usize)) == copy); - let leaf_len = copies.len().div_ceil(copies.len().div_ceil(config.c_max)); - leaves.extend(copies.chunks(leaf_len).map(<[u32]>::to_vec)); - cluster = rest; - } + for cluster in clusters { if cluster.is_empty() { continue; } - if cluster.len() <= config.c_max { + if cluster.len() <= c_max { leaves.push(cluster); + } else if fanout == 1 && cluster.len() == partition.point_ids.len() { + // Every point chose this leader, so the split made no progress. This + // happens when all sampled leaders are copies of one vector: each + // point is equally near to all of them. Another split can stall the + // same way, so cut the cluster into equal leaves in point order. The + // cut groups points arbitrarily, which can lower recall but keeps + // the graph valid. + let leaf_len = cluster.len().div_ceil(cluster.len().div_ceil(c_max)); + leaves.extend(cluster.chunks(leaf_len).map(<[u32]>::to_vec)); } else { pending.push(PendingPartition { point_ids: cluster, @@ -249,17 +254,15 @@ where Ok(PartitionSplit { pending, leaves }) } -/// Sample point IDs that act as centers for one partition split. -fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> Vec { - let count = sampled_leader_count(points.len(), sampling_fraction); +/// Sample `count` distinct leaders from the points of a cluster. +fn sample_leaders(points: &[u32], count: usize, seed: u64) -> Vec { let mut rng = rand::rngs::StdRng::seed_from_u64(seed); points.choose_multiple(&mut rng, count).copied().collect() } -/// Return the number of centers to sample from one cluster. +/// Return the number of leaders for a cluster of `points` points. /// -/// Splitting requires at least two points. The validated sampling fraction is in -/// `(0, 1]`; round up its point count and keep between two and `LEADER_CAP` centers. +/// A split needs two leaders. `LEADER_CAP` bounds the cost of one split. fn sampled_leader_count(points: usize, sampling_fraction: f64) -> usize { ((points as f64 * sampling_fraction).ceil() as usize).clamp(2, LEADER_CAP) } @@ -276,11 +279,11 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { .wrapping_add(salt) } -/// Assign each cluster point to its nearest sampled partition centers. +/// Assign each point to its `fanout` nearest leaders and group the points by +/// leader. /// -/// The function gathers center vectors once and evaluates points in bounded -/// stripes. The assignment matrix keeps point order. Scatter preserves this order -/// inside each child partition, which makes recursive sampling deterministic. +/// Stripes of points run in parallel. Each child keeps the input order of its +/// points, so the next split samples the same leaders for any thread count. fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, @@ -303,12 +306,11 @@ where let leaders = M::create_leaders(leader_matrix); let leader_count = M::leader_count(&leaders); - let fanout = fanout.min(leader_count); let mut assignments = vec![0u32; point_ids.len() * fanout]; let stripe_points = assignment_stripe_point_count(leader_count); - // `build_graph` runs this operation in the pool from the build context. - // A buffer lease takes one lock, which costs far less than the GEMM and - // ranking of one stripe. + // This runs in the pool of the build context (see `build_graph`). A lease + // for each stripe takes one lock, which costs far less than the GEMM and + // ranking of the stripe. assignments .par_chunks_mut(stripe_points * fanout) .zip(point_ids.par_chunks(stripe_points)) @@ -332,11 +334,8 @@ where )) } -/// Assign one point stripe to sampled partition centers. -/// -/// The function gathers point IDs into a packed `f32` matrix. The partition -/// kernel owns ranking-distance construction and ranking. This function writes the -/// returned leader-column IDs for partition scatter. +/// Gather one stripe of points as `f32` and write the leader columns of the +/// `fanout` nearest leaders of each point to `assignments`. #[inline] fn assign_point_stripe( arch: A, @@ -355,8 +354,7 @@ where let point_count = point_ids.len(); let dimensions = data.ncols(); let point_values_len = point_count * dimensions; - // Keep each buffer at its largest length. Every operation uses an explicit - // active prefix. + // The buffer keeps its largest length. This stripe uses a prefix of it. if buffers.point_values.len() < point_values_len { buffers.point_values.resize(point_values_len, 0.0); } @@ -374,10 +372,10 @@ where assign_leaders::(arch, points.as_view(), leaders, output, kernel_workspace) } -/// Group assigned point IDs by child partition. +/// Group the points by assigned leader. Unassigned slots join no child. /// -/// Both the serial and parallel paths preserve point order inside each child. -/// This order is required for deterministic recursive sampling. +/// Both paths keep the input order of the points inside each child, because the +/// next split samples its leaders by position. fn scatter_assignments( points: &[u32], assignments: &[u32], @@ -390,7 +388,8 @@ fn scatter_assignments( let stripe_points = points.len().div_ceil(rayon::current_num_threads()); let stripe_assignment_count = stripe_points * fanout; - // Indexed parallel collection preserves stripe order. + // Scatter each stripe of points on its own. Indexed collection keeps the + // stripe order. let locals: Vec<_> = points .par_chunks(stripe_points) .zip(assignments.par_chunks(stripe_assignment_count)) @@ -404,8 +403,7 @@ fn scatter_assignments( } } - // `build_graph` runs this Rayon operation in the pool from the build context. - // Each worker creates one independent leader cluster. + // Join the stripe parts of each child in stripe order. sizes .into_par_iter() .enumerate() @@ -442,63 +440,6 @@ fn scatter_serial( clusters } -/// Merge leaves smaller than `c_min` without exceeding `c_max`. -/// -/// A `HashSet` removes duplicate point IDs across merged leaves. The function -/// sorts each merged result before it returns. -fn merge_undersized_leaves(leaves: Vec>, c_min: usize, c_max: usize) -> Vec> { - let mut merged = Vec::with_capacity(leaves.len()); - let mut small_leaves = Vec::new(); - for leaf in leaves { - if leaf.len() >= c_min { - merged.push(leaf); - } else { - small_leaves.push(leaf); - } - } - if small_leaves.is_empty() { - return merged; - } - - let mut small = HashSet::with_capacity(c_max); - - for leaf in small_leaves { - let combined = small.len() + leaf.len(); - if combined > c_max { - merged.push(drain_sorted(&mut small)); - } - small.extend(leaf); - if small.len() >= c_min { - merged.push(drain_sorted(&mut small)); - } - } - - if !small.is_empty() { - let mut remainder = drain_sorted(&mut small); - if remainder.len() < c_min - && let Some(last) = merged.last_mut() - { - remainder.retain(|id| !last.contains(id)); - let combined = last.len() + remainder.len(); - if combined <= c_max { - last.append(&mut remainder); - last.sort_unstable(); - } - } - if !remainder.is_empty() { - merged.push(remainder); - } - } - - merged -} - -fn drain_sorted(set: &mut HashSet) -> Vec { - let mut values: Vec<_> = set.drain().collect(); - values.sort_unstable(); - values -} - /// Return the number of points in one assignment stripe. /// /// `LEADER_CAP` keeps this count at 131 points or more. @@ -507,15 +448,16 @@ fn assignment_stripe_point_count(leader_count: usize) -> usize { .min(MAX_ASSIGNMENT_STRIPE_POINTS) } -#[cfg(all(test, not(miri)))] +#[cfg(test)] mod tests { use super::*; use crate::graph::pipnn::{ L2, - test_support::{self, dense_points, sorted_members_per_row, thread_pool}, + test_support::{self, dense_points, nz, sorted_members_per_row, thread_pool}, }; use diskann_vector::distance::Metric; use diskann_wide::ARCH; + use std::collections::HashSet; // Neither leaf order nor point order within a leaf is part of the result. // Preserve repeated leaves and IDs so comparison still detects duplicates. @@ -527,16 +469,14 @@ mod tests { fn splitting_config() -> PiPNNConfig { PiPNNConfig { - c_max: 2, - c_min: 1, + c_max: nz(2), p_samp: 1.0, - fanout: vec![2], - leaf_k: 1, - replicas: 1, + fanout: vec![nz(2)], + leaf_k: nz(1), + replicas: nz(1), } } - #[expect(clippy::unwrap_used, reason = "an invalid fixture fails the test")] fn split( data: MatrixView<'_, f32>, config: &PiPNNConfig, @@ -562,13 +502,14 @@ mod tests { let case = format!("{point_count} points, fraction {fraction}"); let ids: Vec<_> = (0..point_count).map(|id| 10 + 3 * id).collect(); - let mut leaders = sample_leaders(&ids, fraction, 1290); + let count = sampled_leader_count(ids.len(), fraction); + let mut leaders = sample_leaders(&ids, count, 1290); assert_eq!(leaders.len(), expected_count, "{case}"); let distinct: HashSet<_> = leaders.iter().collect(); assert_eq!(distinct.len(), expected_count, "{case}"); assert!(leaders.iter().all(|id| ids.contains(id)), "{case}"); - let mut repeated = sample_leaders(&ids, fraction, 1290); + let mut repeated = sample_leaders(&ids, count, 1290); leaders.sort_unstable(); repeated.sort_unstable(); assert_eq!(leaders, repeated, "{case}"); @@ -587,8 +528,6 @@ mod tests { for (fanout, expected) in [ (1, [vec![5, 1], vec![4]]), (2, [vec![5, 4, 1], vec![5, 4, 1]]), - // Fanout above the leader count assigns every point to every leader. - (9, [vec![5, 4, 1], vec![5, 4, 1]]), ] { let actual = pool .install(|| { @@ -639,7 +578,7 @@ mod tests { } #[test] - fn assignments_at_production_shape_join_each_point_to_its_nearest_leaders() { + fn stripes_assign_each_point_to_its_nearest_leaders() { // 300 leaders give stripes of 436 points, so 3,000 points use several // stripes, a partial tail, several workers and full SIMD leader groups. // Dense fixtures keep the f32 scores exact, so the check needs no @@ -685,7 +624,7 @@ mod tests { } #[test] - fn reused_stripe_buffers_replace_previous_points_and_assignment_width() { + fn reused_stripe_buffers_do_not_carry_points_or_widths() { let values = [100.0_f32, 0.0, 8.0, 10.0, 2.0]; let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); let leader_values = [0.0_f32, 10.0, 100.0]; @@ -724,7 +663,7 @@ mod tests { } #[test] - fn scatter_groups_assigned_points_by_leader_and_omits_unassigned_slots() { + fn scatter_groups_points_by_leader_and_drops_unassigned_slots() { let pattern = [ [2, 0], [UNASSIGNED, 1], @@ -792,7 +731,7 @@ mod tests { let values = [0.0_f32, 1.0, 10.0, 11.0]; let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); let config = PiPNNConfig { - c_max: 1, + c_max: nz(1), ..splitting_config() }; let parent = PendingPartition { @@ -817,41 +756,86 @@ mod tests { } #[test] - fn a_stalled_split_moves_the_leader_copies_into_leaves() { - // IDs 0 to 5 are copies of one vector, and ID 6 differs. When both - // sampled leaders are copies, all points tie and choose the same leader. - let values = [1.0_f32, 1.0, 1.0, 1.0, 1.0, 1.0, 9.0]; - let data = MatrixView::try_from(&values[..], 7, 1).unwrap(); - let ids: Vec = (0..7).collect(); + fn a_fanout_one_split_without_progress_cuts_equal_leaves() { + // Copies of one vector tie for every leader, so all points choose the + // same leader and the child equals its parent. + let values = [1.0_f32; 5]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); let config = PiPNNConfig { - c_max: 4, - p_samp: 0.01, - fanout: vec![1], + fanout: vec![nz(1)], ..splitting_config() }; - // `split_partition` derives its sampling seed from the parent seed and size. - let seed = (0..) - .find(|&seed| { - let leaders = sample_leaders(&ids, config.p_samp, mix_seed(seed, 7)); - leaders.iter().all(|&id| id < 6) - }) - .unwrap(); let parent = PendingPartition { - point_ids: ids, + point_ids: (0..5).collect(), level: 0, - seed, + seed: 1290, }; let actual = split::(data, &config, parent); - // The copies fill two equal leaves. The distinct point is not grouped with them. assert_eq!( sorted_leaf_memberships(&actual.leaves), - [vec![0, 1, 2], vec![3, 4, 5], vec![6]] + [vec![0, 1], vec![2, 3], vec![4]] ); assert!(actual.pending.is_empty()); } + #[test] + fn a_child_that_repeats_its_parent_at_fanout_two_is_split_again() { + // Points at x=0, x=10 and x=4 are all leaders. Each point joins itself and + // x=4, so the child of x=4 holds all points. With fanout above one this + // is overlap, not a stalled split. + let values = [0.0_f32, 10.0, 4.0]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let parent = PendingPartition { + point_ids: vec![0, 1, 2], + level: 0, + seed: 1290, + }; + + let actual = split::(data, &splitting_config(), parent); + + assert_eq!( + sorted_leaf_memberships(&actual.leaves), + [vec![0, 2], vec![1]] + ); + let [child] = &actual.pending[..] else { + panic!("expected one pending cluster"); + }; + assert_eq!( + (child.point_ids.as_slice(), child.level), + (&[0, 1, 2][..], 1) + ); + } + + #[test] + fn a_level_without_more_leaders_than_its_fanout_is_skipped() { + // Four points at p_samp 0.01 give two leaders. Fanout 2 would put every + // point into both children, so the cluster moves to the next level. + let values = [0.0_f32, 1.0, 10.0, 11.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let config = PiPNNConfig { + p_samp: 0.01, + ..splitting_config() + }; + let parent = PendingPartition { + point_ids: vec![0, 1, 2, 3], + level: 0, + seed: 1290, + }; + + let actual = split::(data, &config, parent); + + assert!(actual.leaves.is_empty()); + let [child] = &actual.pending[..] else { + panic!("expected one pending cluster"); + }; + assert_eq!( + (child.point_ids.as_slice(), child.level), + (&[0, 1, 2, 3][..], 1) + ); + } + #[test] fn data_within_cmax_forms_one_complete_leaf_per_replica() { let values = [0.0_f32, 1.0, 10.0]; @@ -860,8 +844,8 @@ mod tests { for c_max in [3, 4] { let config = PiPNNConfig { - c_max, - replicas: 2, + c_max: nz(c_max), + replicas: nz(2), ..splitting_config() }; @@ -883,11 +867,10 @@ mod tests { let values = dense_points(point_count, 17, 1290); let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); let config = PiPNNConfig { - c_max: 16, - c_min: 4, + c_max: nz(16), p_samp: 0.2, - fanout: vec![2, 2], - replicas: 2, + fanout: vec![nz(2), nz(2)], + replicas: nz(2), ..splitting_config() }; @@ -908,7 +891,7 @@ mod tests { assert_eq!(ids, (0..point_count as u32).collect::>()); for leaf in first { assert!( - !leaf.is_empty() && leaf.len() <= config.c_max, + !leaf.is_empty() && leaf.len() <= config.c_max.get(), "leaf {leaf:?}" ); assert_eq!( @@ -925,10 +908,9 @@ mod tests { let values = dense_points(point_count, 17, 1290); let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); let config = PiPNNConfig { - c_max: 8, - c_min: 1, + c_max: nz(8), p_samp: 0.25, - fanout: vec![1], + fanout: vec![nz(1)], ..splitting_config() }; let pool = thread_pool(2); @@ -942,7 +924,7 @@ mod tests { ARCH, data, &PiPNNConfig { - replicas: 2, + replicas: nz(2), ..config }, ) @@ -959,7 +941,7 @@ mod tests { let mut memberships = vec![0; point_count]; for leaf in &combined { assert!( - !leaf.is_empty() && leaf.len() <= config.c_max, + !leaf.is_empty() && leaf.len() <= config.c_max.get(), "leaf {leaf:?}" ); assert_eq!(leaf.iter().collect::>().len(), leaf.len()); @@ -971,10 +953,10 @@ mod tests { } #[test] - fn unrankable_points_do_not_form_child_leaves() { + fn unrankable_points_join_no_leaf() { let config = PiPNNConfig { - c_max: 1, - fanout: vec![1], + c_max: nz(1), + fanout: vec![nz(1)], ..splitting_config() }; let pool = thread_pool(1); @@ -996,56 +978,4 @@ mod tests { ); } } - - #[test] - fn small_leaf_merging_deduplicates_ids_without_exceeding_cmax() { - // (leaves, c_min, c_max, expected leaves) - for (leaves, c_min, c_max, expected) in [ - (vec![], 2, 4, vec![]), - ( - vec![vec![1, 4], vec![0, 3, 6, 8]], - 2, - 4, - vec![vec![1, 4], vec![0, 3, 6, 8]], - ), - // Overlapping small leaves merge into one leaf without repeated IDs. - ( - vec![vec![4], vec![1], vec![4], vec![2]], - 3, - 4, - vec![vec![1, 2, 4]], - ), - // A merge that would exceed c_max flushes the first leaf. - ( - vec![vec![4, 5], vec![0, 1]], - 3, - 3, - vec![vec![4, 5], vec![0, 1]], - ), - // After duplicate removal, the undersized remainder fits the last leaf. - ( - vec![vec![1, 3, 5], vec![3], vec![6]], - 3, - 4, - vec![vec![1, 3, 5, 6]], - ), - ( - vec![vec![1, 3, 5], vec![6], vec![8]], - 3, - 4, - vec![vec![1, 3, 5], vec![6, 8]], - ), - (vec![vec![7], vec![2]], 3, 4, vec![vec![2, 7]]), - ] { - let case = format!("{leaves:?}, c_min={c_min}, c_max={c_max}"); - - let actual = merge_undersized_leaves(leaves, c_min, c_max); - - assert_eq!( - sorted_leaf_memberships(&actual), - sorted_leaf_memberships(&expected), - "{case}" - ); - } - } } From 175f0c8bae87460bba8ecc8fdf2eef615e504637 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 29 Sep 2026 08:16:06 +0000 Subject: [PATCH 209/229] perf(pipnn): stop sorting merged leaf candidates The candidate lists were sorted by point ID on one thread after the leaf jobs. ID order carries no distance information: final pruning ranks the candidates by distance, and search reads every neighbor of a point, so the sort only fixed how distance ties break. On BigANN-10M (c_max 512, fanout [10, 3], k 3, 16 threads) the serial pass took 10.4-12.1 s per build; recall@10 at L 10-100 stayed within 3e-5 without it. The graph now depends on job completion order in distance ties and in neighbor order, so the 1-vs-4-worker equality test becomes a row validity test. The leaves stay reproducible. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/leaf_build.rs | 16 +++++----------- diskann/src/graph/pipnn/mod.rs | 14 ++++++-------- diskann/src/graph/pipnn/partitioning.rs | 2 +- 3 files changed, 12 insertions(+), 20 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 4c777be7ac..e5480de2e9 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -68,8 +68,9 @@ impl LeafBuffers { /// leaves. /// /// Each selected pair adds both directions. A list holds global IDs without -/// duplicates. The lists are sorted, so the result does not depend on the order -/// in which the parallel jobs finish. +/// duplicates, in the order in which the parallel jobs add them. The lists stay +/// unsorted: final pruning ranks candidates by distance, and a graph search reads +/// all neighbors of a point, so the order only breaks distance ties. pub(super) fn build_leaf_candidates( arch: A, data: MatrixView<'_, T>, @@ -89,14 +90,7 @@ where .try_for_each_init(LeafBuffers::default, |buffers, point_ids| { add_leaf_candidates::(arch, data, point_ids, requested_k, buffers, &candidates) })?; - Ok(candidates - .into_iter() - .map(Mutex::into_inner) - .map(|mut neighbors| { - neighbors.sort(); - neighbors - }) - .collect()) + Ok(candidates.into_iter().map(Mutex::into_inner).collect()) } /// Add the selected pairs of one leaf to the candidate lists. @@ -316,7 +310,7 @@ mod tests { .unwrap(); let actual: Vec> = actual.into_iter().map(Vec::from).collect(); - assert_eq!(actual, expected, "k={k}"); + assert_eq!(sorted_members_per_row(&actual), expected, "k={k}"); } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index ffb353ea1c..d125c7ab83 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -884,11 +884,11 @@ mod construction_tests { } #[test] - fn multi_level_builds_are_deterministic_for_every_metric() { + fn multi_level_builds_give_valid_rows_for_every_metric() { // With 1,000 points and c_max 32, the build splits over several levels, - // makes overlapping leaves, merges their candidates in parallel and - // prunes lists above the degree. The worker count must not change the - // graph. + // makes overlapping leaves, merges their candidates on four workers and + // prunes lists above the degree. Each row must be non-empty, within the + // degree, and free of duplicates and self edges. let (points, dimensions, degree) = (1000, 16, 8); let config = PiPNNConfig { c_max: nz(32), @@ -908,11 +908,9 @@ mod construction_tests { test_support::normalize(&mut values, dimensions); } - let serial = build(&values, dimensions, config.clone(), degree, metric, 1); - let parallel = build(&values, dimensions, config.clone(), degree, metric, 4); + let actual = build(&values, dimensions, config.clone(), degree, metric, 4); - assert_eq!(serial, parallel, "{metric:?}"); - for (point, neighbors) in serial.iter().enumerate() { + for (point, neighbors) in actual.iter().enumerate() { let mut distinct = neighbors.clone(); distinct.sort_unstable(); distinct.dedup(); diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index cf5d07404b..e291ee8925 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -40,7 +40,7 @@ use super::{ topk::UNASSIGNED, }; -/// Seed of the first replica. Fixed seeds make each build reproducible. +/// Seed of the first replica. Fixed seeds give the same leaves in every build. const PARTITION_SEED: u64 = 1_000; /// Seed step between replicas, so each replica samples different leaders. const REPLICA_SEED_STEP: u64 = 7_919; From b45927f41e92e593dcbca1610d0d7ab7e22f185a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 12 Oct 2026 08:29:02 +0000 Subject: [PATCH 210/229] perf(pipnn): use a fixed 256-point assignment stripe The stripe height came from a 512 KiB distance-block target justified by the L2 cache. Sweeps of assign_to_leaders on Zen 3, Zen 4, Emerald Rapids and Ice Lake (GCP VMs, 0% steal) showed that L2 is not the limit: one worker runs 16 MiB stripes at full speed, and slowdowns start when the stripes of all workers exceed the last-level cache. The working set also includes the gathered points, which the old rule ignored; with 64 leaders and 1,536 dimensions its 1,024-point cap was up to 2.8x slower than the best height. A fixed 256 points stayed within 11% of the best height for f32, f16 and u8 data, 128-1,536 dimensions, 64-1,000 leaders and 1-16 workers (1.4% on average). A cache-size rule did not do better: VMs report a last-level cache that other tenants share. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 49 ++++++++++++------------- 1 file changed, 23 insertions(+), 26 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index e291ee8925..c88f9a098c 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -50,17 +50,23 @@ const REPLICA_SEED_STEP: u64 = 7_919; /// split of 10M points at `p_samp = 0.005` samples 50,000 leaders and does 50 /// times the work. const LEADER_CAP: usize = 1_000; -/// Size target of the distance block of one stripe, `points x leaders x 4` -/// bytes. GEMM writes the block and the top-k scan reads it back, so the block -/// must stay in the L2 cache of the core. On an AMD EPYC 7763 (512 KiB L2 per -/// core), this value gave the shortest partition time for 10M points: 128 KiB -/// was about 7% slower and 2 MiB about 5% slower. Current server cores have -/// 512 KiB to 2 MiB of L2 cache. -const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 512 * 1024; -/// Maximum number of points in one stripe. With few leaders, the cache target -/// alone allows very tall stripes. This bound keeps the gathered points of one -/// worker at `1024 x dimensions x 4` bytes or less. -const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; +/// Number of points in one assignment stripe. A stripe gathers its points as +/// `f32` and computes their distances to all leaders, so it moves +/// `points x (leaders + dimensions) x 4` bytes through the cache. +/// +/// Each stripe repeats fixed work, such as one GEMM call over the leader matrix, +/// so 16-point stripes took 1.3 to 2 times as long as the best height. Tall +/// stripes slow down when the stripes of all workers exceed the last-level +/// cache. The L2 cache is not the limit: one worker ran 16 MiB stripes at full +/// speed on a CPU with 512 KiB of L2 per core. +/// +/// Sweeps on Zen 3, Zen 4, Emerald Rapids and Ice Lake CPUs covered f32, f16 and +/// u8 data, 128 to 1,536 dimensions, 64 to 1,000 leaders and 1 to 16 workers. +/// 256 points stayed within 11% of the best height, and within 1.4% on average. +/// With 8 leaders and 1,536 dimensions, where the gather dominates, the loss was +/// up to 27%. A rule from the detected cache size did not do better, because a +/// virtual machine reports a last-level cache that other tenants also use. +const ASSIGNMENT_STRIPE_POINTS: usize = 256; /// Clusters with at least this many points scatter in parallel. Such clusters /// occur at the first levels, where few splits run at the same time. A serial /// scatter of 10M points at fanout 10 takes 0.59 s, and a parallel scatter @@ -307,13 +313,12 @@ where let leader_count = M::leader_count(&leaders); let mut assignments = vec![0u32; point_ids.len() * fanout]; - let stripe_points = assignment_stripe_point_count(leader_count); // This runs in the pool of the build context (see `build_graph`). A lease // for each stripe takes one lock, which costs far less than the GEMM and // ranking of the stripe. assignments - .par_chunks_mut(stripe_points * fanout) - .zip(point_ids.par_chunks(stripe_points)) + .par_chunks_mut(ASSIGNMENT_STRIPE_POINTS * fanout) + .zip(point_ids.par_chunks(ASSIGNMENT_STRIPE_POINTS)) .try_for_each(|(stripe_assignments, stripe_ids)| { assign_point_stripe::( arch, @@ -440,14 +445,6 @@ fn scatter_serial( clusters } -/// Return the number of points in one assignment stripe. -/// -/// `LEADER_CAP` keeps this count at 131 points or more. -fn assignment_stripe_point_count(leader_count: usize) -> usize { - (ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count * size_of::())) - .min(MAX_ASSIGNMENT_STRIPE_POINTS) -} - #[cfg(test)] mod tests { use super::*; @@ -545,10 +542,10 @@ mod tests { #[test] fn multiple_stripes_and_their_partial_tail_preserve_all_assignments() { - // More than two maximum-sized stripes, with three points left over. + // Two full stripes, with three points left over. // Each dense row selects two of the three centers at 10, 0 and 4. // Every row selects 4; its other center is 10 for x=10 and 0 otherwise. - let point_count = 2 * MAX_ASSIGNMENT_STRIPE_POINTS + 3; + let point_count = 2 * ASSIGNMENT_STRIPE_POINTS + 3; let dimensions = 384; let values: Vec<_> = (0..point_count) .flat_map(|point| std::iter::repeat_n([0.0_f32, 4.0, 10.0][point % 3], dimensions)) @@ -579,8 +576,8 @@ mod tests { #[test] fn stripes_assign_each_point_to_its_nearest_leaders() { - // 300 leaders give stripes of 436 points, so 3,000 points use several - // stripes, a partial tail, several workers and full SIMD leader groups. + // 3,000 points fill 11 stripes and a partial tail on several workers. + // 300 leaders fill whole SIMD leader groups. // Dense fixtures keep the f32 scores exact, so the check needs no // tolerance. It accepts any order of tied leaders. let (point_count, dimensions, fanout) = (3000, 32, 3); From a2f27eef090fda2129d0c7993aa647f607f304a8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 12 Oct 2026 09:14:40 +0000 Subject: [PATCH 211/229] fix(pipnn): cap the prune pool at max_occlusion_size Final pruning kept the u16::MAX nearest candidates and said that Vamana does the same. Vamana keeps the config.max_occlusion_size() nearest (default 750). Use the graph configuration value so PiPNN and Vamana prune the same pool and long candidate lists cost the same. Lists in the usual configurations (at most l_max with HashPrune, at most 595 for direct BigANN-10M candidates) are below the cap, so their graphs do not change. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/finalization.rs | 45 +++++++++++++++++++------ 1 file changed, 34 insertions(+), 11 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index c4c620352f..a36ec7ac43 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -54,12 +54,13 @@ impl PruneWorkspace { distance.evaluate_similarity(source_vector, data.row(candidate as usize)), ) })); - // RobustPrune stores candidate positions as `u16`, so keep the `u16::MAX` - // nearest candidates. Vamana limits its candidate pool the same way. + // Like Vamana, prune only the `max_occlusion_size` nearest candidates. The + // graph configuration keeps this size within the `u16` positions of + // RobustPrune. // // Sort before mapping: changing the element type can change the unstable // sort's order for equal distances, and therefore RobustPrune's choices. - let sorted = SortedNeighbors::new(&mut self.candidates, u16::MAX as usize); + let sorted = SortedNeighbors::new(&mut self.candidates, graph.max_occlusion_size().get()); let mapped = sorted.map_in(&mut self.mapped_candidates, |&candidate| { (candidate as usize != source).then_some(candidate) }); @@ -239,17 +240,39 @@ mod tests { } #[test] - fn a_list_above_the_u16_position_limit_keeps_its_nearest_candidates() { - let count = u16::MAX as usize + 1; - let values: Vec<_> = (0..=count).map(|i| i as f32).collect(); - let data = MatrixView::try_from(values.as_slice(), count + 1, 1).unwrap(); - let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); + fn pruning_reads_only_the_max_occlusion_size_nearest_candidates() { + // Source 0 is at x=0. Candidates 2 and 3 lie on one side and candidate 1 + // on the other, so 2 occludes 3 but not 1. A pool of two candidates + // drops 1, the farthest. + let values = [0.0_f32, -3.0, 1.0, 2.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); let distance = f32::distance(Metric::L2, Some(1)); - let candidates = AdjacencyList::from_iter_untrusted(1..=count as u32); + for (max_occlusion_size, expected) in [(2, &[2][..]), (3, &[2, 1][..])] { + let graph = config::Builder::new_with( + 2, + config::MaxDegree::same(), + 16, + Metric::L2.into(), + |b| { + b.alpha(1.0).max_occlusion_size(max_occlusion_size); + }, + ) + .build() + .unwrap(); - let actual = PruneWorkspace::default().prune(data, 0, candidates, &graph, &distance); + let actual = PruneWorkspace::default().prune( + data, + 0, + candidates_in_order(&[1, 2, 3]), + &graph, + &distance, + ); - assert_eq!(&*actual, [1]); + assert_eq!( + &*actual, expected, + "max_occlusion_size {max_occlusion_size}" + ); + } } #[test] From 5ceb7901787bd51535eebc00fa3cad36217e7532 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:24:53 +0000 Subject: [PATCH 212/229] docs(pipnn): name the child seed multiplier and explain seed values Reviewers asked why 6364136223846793005 was chosen. Name it CHILD_SEED_MULTIPLIER and state where it comes from (Knuth's MMIX LCG) and why any odd multiplier would do: StdRng::seed_from_u64 mixes the seed again, so the step only has to keep seeds distinct. Mark the partition seed and replica step values as arbitrary. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index c88f9a098c..b2d5c6dbda 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -40,10 +40,19 @@ use super::{ topk::UNASSIGNED, }; -/// Seed of the first replica. Fixed seeds give the same leaves in every build. +/// Seed of the first replica. The value is arbitrary. A fixed seed gives the +/// same leaves in every build. const PARTITION_SEED: u64 = 1_000; -/// Seed step between replicas, so each replica samples different leaders. +/// Seed step between replicas. Any nonzero step gives each replica its own +/// leaders. The value is arbitrary. const REPLICA_SEED_STEP: u64 = 7_919; +/// Multiplier of the step that derives a child seed from its parent seed. +/// +/// This is the multiplier of Knuth's MMIX 64-bit linear congruential generator. +/// It is odd, so multiplication modulo 2^64 maps different parent seeds to +/// different products. `StdRng::seed_from_u64` mixes the child seed again, so +/// the step must only keep seeds distinct. +const CHILD_SEED_MULTIPLIER: u64 = 6_364_136_223_846_793_005; /// Maximum number of leaders in one split. The PiPNN paper uses the same hard /// cap ("typically 1000", Section 4.1). A split costs /// `points x leaders x dimensions` multiply-adds. Without the cap, the first @@ -277,12 +286,11 @@ fn replica_seed(replica: usize) -> u64 { PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) } -// This LCG step derives a child seed from the parent seed and the cluster size. -// The multiplier is Knuth's MMIX constant. Wrapping arithmetic gives the same -// mapping in debug and release builds on all supported platforms. +// Derive a child seed from the parent seed and the cluster size. Wrapping +// arithmetic gives the same mapping in debug and release builds on all +// supported platforms. fn mix_seed(seed: u64, salt: u64) -> u64 { - seed.wrapping_mul(6_364_136_223_846_793_005) - .wrapping_add(salt) + seed.wrapping_mul(CHILD_SEED_MULTIPLIER).wrapping_add(salt) } /// Assign each point to its `fanout` nearest leaders and group the points by From f0e5c3bdb62a1c6979362d1dd7552fb80d4840cd Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:33:56 +0000 Subject: [PATCH 213/229] refactor(pipnn): draw child seeds from the split RNG The child seed came from an LCG step with Knuth's MMIX multiplier, and replica seeds from a prime step of 7,919. Neither value mattered: StdRng::seed_from_u64 already expands the seed with PCG32 and gives independent streams for adjacent seeds. Each split now seeds one StdRng, samples its leaders from it, and draws the seed of each child from the same stream. Replica r uses PARTITION_SEED + r. The leaves stay reproducible, but differ from the previous seeds. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 50 ++++++++----------------- 1 file changed, 15 insertions(+), 35 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index b2d5c6dbda..02281fdf54 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -28,7 +28,7 @@ use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, }; -use rand::{SeedableRng, prelude::IndexedRandom}; +use rand::{Rng, SeedableRng, prelude::IndexedRandom, rngs::StdRng}; use rayon::prelude::*; use super::{ @@ -40,19 +40,10 @@ use super::{ topk::UNASSIGNED, }; -/// Seed of the first replica. The value is arbitrary. A fixed seed gives the -/// same leaves in every build. +/// Seed of the first replica; replica `r` uses `PARTITION_SEED + r`. The value +/// is arbitrary. A fixed seed gives the same leaves in every build, and +/// `StdRng::seed_from_u64` gives independent streams for adjacent seeds. const PARTITION_SEED: u64 = 1_000; -/// Seed step between replicas. Any nonzero step gives each replica its own -/// leaders. The value is arbitrary. -const REPLICA_SEED_STEP: u64 = 7_919; -/// Multiplier of the step that derives a child seed from its parent seed. -/// -/// This is the multiplier of Knuth's MMIX 64-bit linear congruential generator. -/// It is odd, so multiplication modulo 2^64 maps different parent seeds to -/// different products. `StdRng::seed_from_u64` mixes the child seed again, so -/// the step must only keep seeds distinct. -const CHILD_SEED_MULTIPLIER: u64 = 6_364_136_223_846_793_005; /// Maximum number of leaders in one split. The PiPNN paper uses the same hard /// cap ("typically 1000", Section 4.1). A split costs /// `points x leaders x dimensions` multiply-adds. Without the cap, the first @@ -89,7 +80,8 @@ struct PendingPartition { point_ids: Vec, /// Position in the fanout schedule. level: usize, - /// Seed of the parent split. The next split mixes it with the cluster size. + /// Seed of the random stream of this cluster's split. The parent split + /// draws it, so every split samples its leaders from its own stream. seed: u64, } @@ -141,7 +133,7 @@ where let mut leaves = Vec::new(); let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas.get() { - let seed = replica_seed(replica); + let seed = PARTITION_SEED + replica as u64; let mut replica_leaves = partition_replica::(arch, data, config, seed, &stripe_buffers)?; leaves.append(&mut replica_leaves); @@ -211,7 +203,7 @@ where M: PartitionMetric, T: VectorRepr + Send + Sync, { - let split_seed = mix_seed(partition.seed, partition.point_ids.len() as u64); + let mut rng = StdRng::seed_from_u64(partition.seed); let fanout = config .fanout .get(partition.level) @@ -225,13 +217,13 @@ where pending: vec![PendingPartition { point_ids: partition.point_ids, level: partition.level + 1, - seed: split_seed, + seed: rng.random(), }], leaves: Vec::new(), }); } - let leaders = sample_leaders(&partition.point_ids, leader_count, split_seed); + let leaders = sample_leaders(&partition.point_ids, leader_count, &mut rng); let clusters = assign_to_leaders::( arch, data, @@ -262,7 +254,7 @@ where pending.push(PendingPartition { point_ids: cluster, level: partition.level + 1, - seed: split_seed, + seed: rng.random(), }); } } @@ -270,9 +262,8 @@ where } /// Sample `count` distinct leaders from the points of a cluster. -fn sample_leaders(points: &[u32], count: usize, seed: u64) -> Vec { - let mut rng = rand::rngs::StdRng::seed_from_u64(seed); - points.choose_multiple(&mut rng, count).copied().collect() +fn sample_leaders(points: &[u32], count: usize, rng: &mut StdRng) -> Vec { + points.choose_multiple(rng, count).copied().collect() } /// Return the number of leaders for a cluster of `points` points. @@ -282,17 +273,6 @@ fn sampled_leader_count(points: usize, sampling_fraction: f64) -> usize { ((points as f64 * sampling_fraction).ceil() as usize).clamp(2, LEADER_CAP) } -fn replica_seed(replica: usize) -> u64 { - PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) -} - -// Derive a child seed from the parent seed and the cluster size. Wrapping -// arithmetic gives the same mapping in debug and release builds on all -// supported platforms. -fn mix_seed(seed: u64, salt: u64) -> u64 { - seed.wrapping_mul(CHILD_SEED_MULTIPLIER).wrapping_add(salt) -} - /// Assign each point to its `fanout` nearest leaders and group the points by /// leader. /// @@ -508,13 +488,13 @@ mod tests { let ids: Vec<_> = (0..point_count).map(|id| 10 + 3 * id).collect(); let count = sampled_leader_count(ids.len(), fraction); - let mut leaders = sample_leaders(&ids, count, 1290); + let mut leaders = sample_leaders(&ids, count, &mut StdRng::seed_from_u64(1290)); assert_eq!(leaders.len(), expected_count, "{case}"); let distinct: HashSet<_> = leaders.iter().collect(); assert_eq!(distinct.len(), expected_count, "{case}"); assert!(leaders.iter().all(|id| ids.contains(id)), "{case}"); - let mut repeated = sample_leaders(&ids, count, 1290); + let mut repeated = sample_leaders(&ids, count, &mut StdRng::seed_from_u64(1290)); leaders.sort_unstable(); repeated.sort_unstable(); assert_eq!(leaders, repeated, "{case}"); From 2dcf68b16680a15ec1c9bb5eac614ab137f0aeb8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:41:51 +0000 Subject: [PATCH 214/229] feat(pipnn): take the leader sampling seed from the caller PiPNNConfig gains seed: Option, following the random_seed convention of the disk builder: Some(seed) gives the same leaves in every build, and None draws a seed from the operating system. Replica r uses base + r. The fixed PARTITION_SEED is gone. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/mod.rs | 5 +++++ diskann/src/graph/pipnn/partitioning.rs | 10 +++++----- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index d125c7ab83..7d21c04ee7 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -212,6 +212,9 @@ pub struct PiPNNConfig { pub leaf_k: NonZeroUsize, /// Number of independent partitionings of the dataset. pub replicas: NonZeroUsize, + /// Seed of leader sampling. The same seed gives the same leaves. `None` + /// draws a seed from the operating system, so each build samples new leaves. + pub seed: Option, } impl PiPNNConfig { @@ -577,6 +580,7 @@ mod construction_tests { fanout: vec![nz(2)], leaf_k: nz(1), replicas: nz(1), + seed: Some(1290), } } @@ -896,6 +900,7 @@ mod construction_tests { fanout: vec![nz(4), nz(2)], leaf_k: nz(2), replicas: nz(1), + seed: Some(1290), }; for metric in [ Metric::L2, diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 02281fdf54..2734e0b72e 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -40,10 +40,6 @@ use super::{ topk::UNASSIGNED, }; -/// Seed of the first replica; replica `r` uses `PARTITION_SEED + r`. The value -/// is arbitrary. A fixed seed gives the same leaves in every build, and -/// `StdRng::seed_from_u64` gives independent streams for adjacent seeds. -const PARTITION_SEED: u64 = 1_000; /// Maximum number of leaders in one split. The PiPNN paper uses the same hard /// cap ("typically 1000", Section 4.1). A split costs /// `points x leaders x dimensions` multiply-adds. Without the cap, the first @@ -132,8 +128,11 @@ where { let mut leaves = Vec::new(); let stripe_buffers = StripeBufferPool::new((), 0, None); + // Replica `r` uses seed `base + r`. `StdRng::seed_from_u64` gives + // independent streams for adjacent seeds. + let base_seed = config.seed.unwrap_or_else(rand::random); for replica in 0..config.replicas.get() { - let seed = PARTITION_SEED + replica as u64; + let seed = base_seed.wrapping_add(replica as u64); let mut replica_leaves = partition_replica::(arch, data, config, seed, &stripe_buffers)?; leaves.append(&mut replica_leaves); @@ -459,6 +458,7 @@ mod tests { fanout: vec![nz(2)], leaf_k: nz(1), replicas: nz(1), + seed: Some(1290), } } From 18874ae0e4b1c7ceb55b903e4526827cf5c3b6d5 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:59:16 +0000 Subject: [PATCH 215/229] refactor(pipnn): one name per partitioning concept "partition" named the whole process, one waiting cluster and a split result, and "scatter" named the grouping of points by leader. Keep partition/partitioning for the whole process only: - PendingPartition -> Cluster, split_partition -> split_cluster, PartitionSplit -> ClusterSplit with next_level and leaves - scatter_assignments -> group_by_leader, scatter_serial -> group_by_leader_serial, PARALLEL_SCATTER_MIN_POINTS -> PARALLEL_GROUPING_MIN_POINTS - the parallel grouping splits points into chunks, not stripes, so stripe keeps its single meaning (one assignment task) No behavior change. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 149 ++++++++++++------------ 1 file changed, 73 insertions(+), 76 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 2734e0b72e..72652211d7 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -63,16 +63,16 @@ const LEADER_CAP: usize = 1_000; /// up to 27%. A rule from the detected cache size did not do better, because a /// virtual machine reports a last-level cache that other tenants also use. const ASSIGNMENT_STRIPE_POINTS: usize = 256; -/// Clusters with at least this many points scatter in parallel. Such clusters -/// occur at the first levels, where few splits run at the same time. A serial -/// scatter of 10M points at fanout 10 takes 0.59 s, and a parallel scatter -/// takes 0.15 s on 16 threads (AMD EPYC 7763). The serial time does not -/// decrease with more threads, so its share of the build increases on larger -/// machines. Below this size, a serial scatter takes a few milliseconds. -const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; +/// Clusters with at least this many points group their points by leader in +/// parallel. Such clusters occur at the first levels, where few splits run at +/// the same time. Serial grouping of 10M points at fanout 10 takes 0.59 s, and +/// parallel grouping takes 0.15 s on 16 threads (AMD EPYC 7763). The serial time +/// does not decrease with more threads, so its share of the build increases on +/// larger machines. Below this size, serial grouping takes a few milliseconds. +const PARALLEL_GROUPING_MIN_POINTS: usize = 100_000; /// A cluster with more than `c_max` points that waits for its split. -struct PendingPartition { +struct Cluster { point_ids: Vec, /// Position in the fanout schedule. level: usize, @@ -82,8 +82,8 @@ struct PendingPartition { } /// The result of one split: leaves, and clusters for the next level. -struct PartitionSplit { - pending: Vec, +struct ClusterSplit { + next_level: Vec, leaves: Vec>, } @@ -145,7 +145,7 @@ where /// The loop ends. The fanout schedule is finite, and each split sends its /// clusters one level deeper. After the schedule, fanout is 1, so each point /// joins one child. A child is then smaller than its parent, or it holds every -/// point and [`split_partition`] cuts it into leaves. +/// point and [`split_cluster`] cuts it into leaves. fn partition_replica( arch: A, data: MatrixView<'_, T>, @@ -165,7 +165,7 @@ where } let mut leaves = Vec::new(); - let mut pending = vec![PendingPartition { + let mut pending = vec![Cluster { point_ids: initial_point_ids, level: 0, seed, @@ -174,14 +174,12 @@ where // Indexed parallel collection keeps the parent order. let splits: ANNResult> = pending .into_par_iter() - .map(|partition| { - split_partition::(arch, data, config, partition, stripe_buffers) - }) + .map(|cluster| split_cluster::(arch, data, config, cluster, stripe_buffers)) .collect(); let mut next_level = Vec::new(); for mut split in splits? { - next_level.append(&mut split.pending); + next_level.append(&mut split.next_level); leaves.append(&mut split.leaves); } pending = next_level; @@ -190,74 +188,74 @@ where } /// Split one cluster into leaves and clusters for the next level. -fn split_partition( +fn split_cluster( arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, - partition: PendingPartition, + cluster: Cluster, stripe_buffers: &StripeBufferPool, -) -> ANNResult +) -> ANNResult where A: Simd, M: PartitionMetric, T: VectorRepr + Send + Sync, { - let mut rng = StdRng::seed_from_u64(partition.seed); + let mut rng = StdRng::seed_from_u64(cluster.seed); let fanout = config .fanout - .get(partition.level) + .get(cluster.level) .map_or(1, |fanout| fanout.get()); let c_max = config.c_max.get(); - let leader_count = sampled_leader_count(partition.point_ids.len(), config.p_samp); + let leader_count = sampled_leader_count(cluster.point_ids.len(), config.p_samp); if fanout > 1 && fanout >= leader_count { // Every point would join every leader, and each child would repeat this // cluster. Split the cluster at the next level instead. - return Ok(PartitionSplit { - pending: vec![PendingPartition { - point_ids: partition.point_ids, - level: partition.level + 1, + return Ok(ClusterSplit { + next_level: vec![Cluster { + point_ids: cluster.point_ids, + level: cluster.level + 1, seed: rng.random(), }], leaves: Vec::new(), }); } - let leaders = sample_leaders(&partition.point_ids, leader_count, &mut rng); - let clusters = assign_to_leaders::( + let leaders = sample_leaders(&cluster.point_ids, leader_count, &mut rng); + let children = assign_to_leaders::( arch, data, - &partition.point_ids, + &cluster.point_ids, &leaders, fanout, stripe_buffers, )?; - let mut pending = Vec::new(); + let mut next_level = Vec::new(); let mut leaves = Vec::new(); - for cluster in clusters { - if cluster.is_empty() { + for child in children { + if child.is_empty() { continue; } - if cluster.len() <= c_max { - leaves.push(cluster); - } else if fanout == 1 && cluster.len() == partition.point_ids.len() { + if child.len() <= c_max { + leaves.push(child); + } else if fanout == 1 && child.len() == cluster.point_ids.len() { // Every point chose this leader, so the split made no progress. This // happens when all sampled leaders are copies of one vector: each // point is equally near to all of them. Another split can stall the - // same way, so cut the cluster into equal leaves in point order. The + // same way, so cut the child into equal leaves in point order. The // cut groups points arbitrarily, which can lower recall but keeps // the graph valid. - let leaf_len = cluster.len().div_ceil(cluster.len().div_ceil(c_max)); - leaves.extend(cluster.chunks(leaf_len).map(<[u32]>::to_vec)); + let leaf_len = child.len().div_ceil(child.len().div_ceil(c_max)); + leaves.extend(child.chunks(leaf_len).map(<[u32]>::to_vec)); } else { - pending.push(PendingPartition { - point_ids: cluster, - level: partition.level + 1, + next_level.push(Cluster { + point_ids: child, + level: cluster.level + 1, seed: rng.random(), }); } } - Ok(PartitionSplit { pending, leaves }) + Ok(ClusterSplit { next_level, leaves }) } /// Sample `count` distinct leaders from the points of a cluster. @@ -318,7 +316,7 @@ where ) })?; - Ok(scatter_assignments( + Ok(group_by_leader( point_ids, &assignments, fanout, @@ -368,24 +366,23 @@ where /// /// Both paths keep the input order of the points inside each child, because the /// next split samples its leaders by position. -fn scatter_assignments( +fn group_by_leader( points: &[u32], assignments: &[u32], fanout: usize, leaders: usize, ) -> Vec> { - if points.len() < PARALLEL_SCATTER_MIN_POINTS { - return scatter_serial(points, assignments, fanout, leaders); + if points.len() < PARALLEL_GROUPING_MIN_POINTS { + return group_by_leader_serial(points, assignments, fanout, leaders); } - let stripe_points = points.len().div_ceil(rayon::current_num_threads()); - let stripe_assignment_count = stripe_points * fanout; - // Scatter each stripe of points on its own. Indexed collection keeps the - // stripe order. + let chunk_points = points.len().div_ceil(rayon::current_num_threads()); + // Group each chunk of points on its own. Indexed collection keeps the chunk + // order. let locals: Vec<_> = points - .par_chunks(stripe_points) - .zip(assignments.par_chunks(stripe_assignment_count)) - .map(|(points, assignments)| scatter_serial(points, assignments, fanout, leaders)) + .par_chunks(chunk_points) + .zip(assignments.par_chunks(chunk_points * fanout)) + .map(|(points, assignments)| group_by_leader_serial(points, assignments, fanout, leaders)) .collect(); let mut sizes = vec![0usize; leaders]; @@ -395,7 +392,7 @@ fn scatter_assignments( } } - // Join the stripe parts of each child in stripe order. + // Join the chunk parts of each child in chunk order. sizes .into_par_iter() .enumerate() @@ -409,7 +406,7 @@ fn scatter_assignments( .collect() } -fn scatter_serial( +fn group_by_leader_serial( points: &[u32], assignments: &[u32], fanout: usize, @@ -465,11 +462,11 @@ mod tests { fn split( data: MatrixView<'_, f32>, config: &PiPNNConfig, - parent: PendingPartition, - ) -> PartitionSplit { + parent: Cluster, + ) -> ClusterSplit { let buffers = StripeBufferPool::new((), 0, None); thread_pool(2) - .install(|| split_partition::<_, M, _>(ARCH, data, config, parent, &buffers)) + .install(|| split_cluster::<_, M, _>(ARCH, data, config, parent, &buffers)) .unwrap() } @@ -582,7 +579,7 @@ mod tests { let mut chosen = vec![Vec::new(); point_count]; for (column, cluster) in clusters.iter().enumerate() { - // `ids` is descending, and scatter keeps the input order. + // `ids` is descending, and grouping keeps the input order. assert!(cluster.is_sorted_by(|a, b| a > b), "leader {column}"); for &id in cluster { chosen[id as usize].push(column); @@ -648,7 +645,7 @@ mod tests { } #[test] - fn scatter_groups_points_by_leader_and_drops_unassigned_slots() { + fn grouping_by_leader_drops_unassigned_slots() { let pattern = [ [2, 0], [UNASSIGNED, 1], @@ -657,13 +654,13 @@ mod tests { ]; let pool = thread_pool(3); - // Counts around the parallel threshold cover both scatter paths and + // Counts around the parallel threshold cover both grouping paths and // parallel chunks of unequal length. for point_count in [ 7, - PARALLEL_SCATTER_MIN_POINTS - 1, - PARALLEL_SCATTER_MIN_POINTS, - PARALLEL_SCATTER_MIN_POINTS + 1, + PARALLEL_GROUPING_MIN_POINTS - 1, + PARALLEL_GROUPING_MIN_POINTS, + PARALLEL_GROUPING_MIN_POINTS + 1, ] { let ids: Vec<_> = (0..point_count).map(|i| 1_000_000 - i as u32).collect(); let assignments: Vec<_> = (0..point_count).flat_map(|i| pattern[i % 4]).collect(); @@ -674,7 +671,7 @@ mod tests { vec![], ]; - let actual = pool.install(|| scatter_assignments(&ids, &assignments, 2, 4)); + let actual = pool.install(|| group_by_leader(&ids, &assignments, 2, 4)); assert_eq!( sorted_members_per_row(&actual), @@ -694,7 +691,7 @@ mod tests { (0, vec![vec![0, 1], vec![0, 1], vec![2, 3], vec![2, 3]]), (1, vec![vec![0], vec![1], vec![2], vec![3]]), ] { - let parent = PendingPartition { + let parent = Cluster { point_ids: vec![0, 1, 2, 3], level, seed: 1290, @@ -707,7 +704,7 @@ mod tests { sorted_leaf_memberships(&expected), "level {level}" ); - assert!(actual.pending.is_empty(), "level {level}"); + assert!(actual.next_level.is_empty(), "level {level}"); } } @@ -719,7 +716,7 @@ mod tests { c_max: nz(1), ..splitting_config() }; - let parent = PendingPartition { + let parent = Cluster { point_ids: vec![0, 1, 2, 3], level: 0, seed: 1290, @@ -728,9 +725,9 @@ mod tests { let actual = split::(data, &config, parent); assert!(actual.leaves.is_empty()); - assert!(actual.pending.iter().all(|child| child.level == 1)); + assert!(actual.next_level.iter().all(|child| child.level == 1)); let ids: Vec<_> = actual - .pending + .next_level .into_iter() .map(|child| child.point_ids) .collect(); @@ -750,7 +747,7 @@ mod tests { fanout: vec![nz(1)], ..splitting_config() }; - let parent = PendingPartition { + let parent = Cluster { point_ids: (0..5).collect(), level: 0, seed: 1290, @@ -762,7 +759,7 @@ mod tests { sorted_leaf_memberships(&actual.leaves), [vec![0, 1], vec![2, 3], vec![4]] ); - assert!(actual.pending.is_empty()); + assert!(actual.next_level.is_empty()); } #[test] @@ -772,7 +769,7 @@ mod tests { // is overlap, not a stalled split. let values = [0.0_f32, 10.0, 4.0]; let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); - let parent = PendingPartition { + let parent = Cluster { point_ids: vec![0, 1, 2], level: 0, seed: 1290, @@ -784,7 +781,7 @@ mod tests { sorted_leaf_memberships(&actual.leaves), [vec![0, 2], vec![1]] ); - let [child] = &actual.pending[..] else { + let [child] = &actual.next_level[..] else { panic!("expected one pending cluster"); }; assert_eq!( @@ -803,7 +800,7 @@ mod tests { p_samp: 0.01, ..splitting_config() }; - let parent = PendingPartition { + let parent = Cluster { point_ids: vec![0, 1, 2, 3], level: 0, seed: 1290, @@ -812,7 +809,7 @@ mod tests { let actual = split::(data, &config, parent); assert!(actual.leaves.is_empty()); - let [child] = &actual.pending[..] else { + let [child] = &actual.next_level[..] else { panic!("expected one pending cluster"); }; assert_eq!( From cc1d41cd96b3ac24fbacc8b0687e73c1cdbf5e30 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 13 Oct 2026 02:54:35 +0000 Subject: [PATCH 216/229] docs(pipnn): document default PiPNN parameters and their effects Reviewers asked how the parameters are usually set. State the default configuration (c_max 512, p_samp 0.005, fanout [10, 3], leaf_k 2, replicas 1) and what raising each parameter does to candidates, recall and build time. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/mod.rs | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 7d21c04ee7..4e15338347 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -198,19 +198,36 @@ crate::convert_error!(PiPNNConfigError); /// PiPNN partition and leaf parameters. /// /// The DiskANN graph [`Config`] supplies the degree, alpha, and prune kind. +/// +/// The default configuration is `c_max = 512`, `p_samp = 0.005`, +/// `fanout = [10, 3]`, `leaf_k = 2` and `replicas = 1`. Each point then joins +/// about 30 leaves, the product of the fanout values. Larger `c_max`, fanout +/// values, `leaf_k` and `replicas` all give each point more or nearer candidates, +/// which raises recall and makes the build slower. #[derive(Clone, Debug, PartialEq)] pub struct PiPNNConfig { - /// Maximum number of points in a leaf. + /// Maximum number of points in a leaf. A larger leaf holds more of the true + /// neighbors of its points, so the selected neighbors are nearer. The leaf + /// distance work grows with the square of the leaf size. pub c_max: NonZeroUsize, /// Fraction of the points of a cluster that a split samples as leaders, in - /// `(0, 1]`. + /// `(0, 1]`. More leaders give smaller children and fewer split levels, but + /// each split costs `points x leaders x dimensions` multiply-adds. A split + /// samples at most 1,000 leaders. pub p_samp: f64, /// Number of nearest leaders that each point joins, for each split level. /// Levels after this schedule use one leader. The schedule must not be empty. + /// Larger values make leaves overlap more, so each point gets candidates + /// from more leaves. The number of leaves and the leaf build time grow with + /// the product of the values. pub fanout: Vec, - /// Number of nearest neighbors that each point selects inside a leaf. + /// Number of nearest neighbors that each point selects inside a leaf. Each + /// selected pair adds a candidate in both directions. A larger value gives + /// more candidates and makes final pruning slower. pub leaf_k: NonZeroUsize, - /// Number of independent partitionings of the dataset. + /// Number of independent partitionings of the dataset. Each replica repeats + /// the partition and leaf work with new leaders, which adds candidates and + /// build time. pub replicas: NonZeroUsize, /// Seed of leader sampling. The same seed gives the same leaves. `None` /// draws a seed from the operating system, so each build samples new leaves. From db08e2b744a4f7d8b61f24f3f6f471e6c5d7c5a8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 13 Oct 2026 03:32:37 +0000 Subject: [PATCH 217/229] refactor(pipnn): fold buffer growth into LeafBuffers::prepare prepare was the only caller of the grow helper. Grow the two buffers in place and say why they never shrink. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/leaf_build.rs | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index e5480de2e9..d371539ade 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -40,16 +40,21 @@ struct LeafBuffers { impl LeafBuffers { /// Grow the buffers for a leaf of `point_count` points and return the /// effective `k` of the leaf. + /// + /// The buffers never shrink. A smaller leaf uses a prefix, so a later large + /// leaf does not fill the same slots again. fn prepare(&mut self, point_count: usize, dimension_count: usize, requested_k: usize) -> usize { // A point has at most `point_count - 1` other points in its leaf. Wider rows // would hold only empty slots. let leaf_k = requested_k.min(point_count.saturating_sub(1)); - grow(&mut self.point_values, point_count * dimension_count, 0.0); - grow( - &mut self.neighbors, - point_count * leaf_k, - Candidate::default(), - ); + let value_count = point_count * dimension_count; + if self.point_values.len() < value_count { + self.point_values.resize(value_count, 0.0); + } + let neighbor_count = point_count * leaf_k; + if self.neighbors.len() < neighbor_count { + self.neighbors.resize(neighbor_count, Candidate::default()); + } leaf_k } @@ -164,12 +169,6 @@ fn add_symmetric_neighbors( } } -fn grow(values: &mut Vec, len: usize, value: T) { - if values.len() < len { - values.resize(len, value); - } -} - #[cfg(test)] mod tests { use super::*; From 95bf0c7501c4cce1b941d428a0de04ce18755906 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 13 Oct 2026 06:17:33 +0000 Subject: [PATCH 218/229] docs(pipnn): say why leaf rows can have empty slots Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/leaf_build.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index d371539ade..13a4e24cf9 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -157,6 +157,8 @@ fn add_symmetric_neighbors( ) { for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { for neighbor in source_neighbors { + // Skip the empty slots that the kernel leaves for NaN and infinite + // distances. if !neighbor.is_assigned() { continue; } From 8559ba4e4dc3788856dff121f75148d170c27023 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 13 Oct 2026 06:43:17 +0000 Subject: [PATCH 219/229] test(pipnn): cover every input type in leaf build, f16 in multi-level builds The leaf-build oracle test now runs f32, f16, i8 and u8 data with k = 1 to 10 on leaves of 1 to 40 points, so k falls above, equal to and below the other points of a leaf while the buffers are reused. Leaf build has no metric logic; the kernel tests cover metrics, dimensions and widths. The multi-level build test also runs on f16 data. The other f16 builds are too small to split. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/leaf_build.rs | 110 ++++++++++++++++---------- diskann/src/graph/pipnn/mod.rs | 54 +++++++------ 2 files changed, 101 insertions(+), 63 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 13a4e24cf9..3b61957e55 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -179,6 +179,7 @@ mod tests { test_support::{dense_points, sorted_members_per_row, thread_pool}, }; use diskann_wide::ARCH; + use half::f16; use std::collections::BTreeSet; fn build(values: &[f32], leaves: Vec>, k: usize, workers: usize) -> Vec> { @@ -260,59 +261,86 @@ mod tests { #[test] fn candidates_match_the_leaf_kernel_on_each_leaf() { - // The reference runs the leaf kernel on each leaf alone, so the tie order - // is the same. Sixty overlapping leaves of 17 to 40 points run on four - // workers with reused buffers and fill whole SIMD groups. k = 2 uses a - // fixed top-k width, and k = 5 uses the runtime width. - let (point_count, dimensions) = (500, 24); - let values = dense_points(point_count, dimensions, 1290); - let data = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + // The kernel tests cover metrics, dimensions, widths and architectures. + // This test covers what leaf build adds: the gather of each element type, + // k cut to the leaf size, buffers reused across leaf shapes, and the merge + // of overlapping leaves on four workers. + let dimensions = 24; + let values = dense_points(500, dimensions, 1290); + let half: Vec<_> = values.iter().map(|&x| f16::from_f32(x)).collect(); + // Dense points are multiples of 1/8 in [-2, 14], so eight times them fits + // `i8`, and an offset of 16 makes them fit `u8`. + let signed: Vec<_> = values.iter().map(|&x| (8.0 * x) as i8).collect(); + let unsigned: Vec<_> = values.iter().map(|&x| (8.0 * x + 16.0) as u8).collect(); + // Leaves of 1 to 40 points put each k above, equal to and below the other + // points of a leaf. + for k in 1..=10 { + check_against_kernel(&values, dimensions, k); + check_against_kernel(&half, dimensions, k); + check_against_kernel(&signed, dimensions, k); + check_against_kernel(&unsigned, dimensions, k); + } + } + + /// Compare `build_leaf_candidates` with the leaf kernel run on each leaf + /// alone. The kernel sees the same rows in the same order, so ties break the + /// same way and the check is exact. + fn check_against_kernel(values: &[T], dimensions: usize, k: usize) { + let point_count = values.len() / dimensions; + let data = MatrixView::try_from(values, point_count, dimensions).unwrap(); + // Sixty overlapping leaves of 1 to 40 points. let leaves: Vec> = (0..60u32) .map(|leaf| { - let ids: BTreeSet<_> = (0..17 + leaf % 24) + let ids: BTreeSet<_> = (0..1 + leaf % 40) .map(|i| (leaf * 7 + i * 13) % point_count as u32) .collect(); ids.into_iter().collect() }) .collect(); - for k in [2, 5] { - let mut expected = vec![BTreeSet::new(); point_count]; - for leaf in &leaves { - let rows: Vec = leaf - .iter() - .flat_map(|&id| &values[id as usize * dimensions..][..dimensions]) - .copied() - .collect(); - let width = k.min(leaf.len() - 1); - let mut output = vec![Candidate::default(); leaf.len() * width]; - select_leaf_neighbors::<_, L2>( - ARCH, - MatrixView::try_from(rows.as_slice(), leaf.len(), dimensions).unwrap(), - MutMatrixView::try_from(output.as_mut_slice(), leaf.len(), width).unwrap(), - &mut LeafKernelWorkspace::default(), - ) - .unwrap(); - for (source, row) in output.chunks_exact(width).enumerate() { - for neighbor in row.iter().filter(|neighbor| neighbor.is_assigned()) { - let (source, target) = (leaf[source], leaf[neighbor.local_idx as usize]); - expected[source as usize].insert(target); - expected[target as usize].insert(source); - } + let mut expected = vec![BTreeSet::new(); point_count]; + for leaf in &leaves { + let width = k.min(leaf.len() - 1); + if width == 0 { + continue; + } + // Convert with `VectorRepr`, not with the gather under test. + let mut rows = vec![0.0; leaf.len() * dimensions]; + for (&id, row) in leaf.iter().zip(rows.chunks_exact_mut(dimensions)) { + T::as_f32_into(data.row(id as usize), row).unwrap(); + } + let mut output = vec![Candidate::default(); leaf.len() * width]; + select_leaf_neighbors::<_, L2>( + ARCH, + MatrixView::try_from(rows.as_slice(), leaf.len(), dimensions).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), leaf.len(), width).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap(); + for (source, row) in output.chunks_exact(width).enumerate() { + for neighbor in row.iter().filter(|neighbor| neighbor.is_assigned()) { + let (source, target) = (leaf[source], leaf[neighbor.local_idx as usize]); + expected[source as usize].insert(target); + expected[target as usize].insert(source); } } - let expected: Vec> = expected - .into_iter() - .map(|ids| ids.into_iter().collect()) - .collect(); + } + let expected: Vec> = expected + .into_iter() + .map(|ids| ids.into_iter().collect()) + .collect(); - let actual = thread_pool(4) - .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves.clone(), k)) - .unwrap(); + let actual = thread_pool(4) + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves, k)) + .unwrap(); - let actual: Vec> = actual.into_iter().map(Vec::from).collect(); - assert_eq!(sorted_members_per_row(&actual), expected, "k={k}"); - } + let actual: Vec> = actual.into_iter().map(Vec::from).collect(); + assert_eq!( + sorted_members_per_row(&actual), + expected, + "{}, k={k}", + std::any::type_name::() + ); } #[test] diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 4e15338347..0e50e26815 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -909,28 +909,20 @@ mod construction_tests { // With 1,000 points and c_max 32, the build splits over several levels, // makes overlapping leaves, merges their candidates on four workers and // prunes lists above the degree. Each row must be non-empty, within the - // degree, and free of duplicates and self edges. - let (points, dimensions, degree) = (1000, 16, 8); - let config = PiPNNConfig { - c_max: nz(32), - p_samp: 0.05, - fanout: vec![nz(4), nz(2)], - leaf_k: nz(2), - replicas: nz(1), - seed: Some(1290), - }; - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let mut values = test_support::dense_points(points, dimensions, 1290); - if metric == Metric::CosineNormalized { - test_support::normalize(&mut values, dimensions); - } + // degree, and free of duplicates and self edges. The f16 build checks + // that every step reads f16 rows, which the smaller builds never split. + fn check(values: &[T], metric: Metric) { + let (dimensions, degree) = (16, 8); + let config = PiPNNConfig { + c_max: nz(32), + p_samp: 0.05, + fanout: vec![nz(4), nz(2)], + leaf_k: nz(2), + replicas: nz(1), + seed: Some(1290), + }; - let actual = build(&values, dimensions, config.clone(), degree, metric, 4); + let actual = build(values, dimensions, config, degree, metric, 4); for (point, neighbors) in actual.iter().enumerate() { let mut distinct = neighbors.clone(); @@ -941,10 +933,28 @@ mod construction_tests { && neighbors.len() <= degree && distinct.len() == neighbors.len() && !neighbors.contains(&(point as u32)), - "{metric:?}, point {point}: {neighbors:?}" + "{metric:?}, {}, point {point}: {neighbors:?}", + std::any::type_name::() ); } } + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let mut values = test_support::dense_points(1000, 16, 1290); + if metric == Metric::CosineNormalized { + test_support::normalize(&mut values, 16); + } + check(&values, metric); + check( + &values.iter().map(|&x| f16::from_f32(x)).collect::>(), + metric, + ); + } } #[test] From 361d0f979f9eff498c0e58e537ef10624d334737 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 13 Oct 2026 08:04:12 +0000 Subject: [PATCH 220/229] refactor(pipnn): drop the prune kind check, as Vamana does PiPNNBuildContext::new rejected a graph Config whose prune kind did not match the metric. Vamana takes the same Config and metric pair without this check, and every PiPNN caller builds the Config from the metric, so the check never fires. Final pruning keeps using the prune kind of the Config. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/mod.rs | 50 ++-------------------------------- 1 file changed, 3 insertions(+), 47 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 0e50e26815..fd9aab13c1 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -38,7 +38,7 @@ use std::num::NonZeroUsize; use crate::{ ANNError, ANNResult, - graph::{AdjacencyList, Config, config::PruneKind}, + graph::{AdjacencyList, Config}, utils::VectorRepr, }; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -186,11 +186,6 @@ pub enum PiPNNConfigError { SamplingFraction(f64), #[error("fanout must not be empty")] EmptyFanout, - #[error("graph prune kind {prune_kind:?} does not match metric {metric:?}")] - PruneKind { - prune_kind: PruneKind, - metric: Metric, - }, } crate::convert_error!(PiPNNConfigError); @@ -259,8 +254,8 @@ pub struct PiPNNBuildContext<'a> { impl<'a> PiPNNBuildContext<'a> { /// Check `config` and combine it with the graph policy. /// - /// Final pruning applies the prune kind of `graph` to distances of `metric`, - /// so the two must match. + /// Final pruning uses the prune kind of `graph`, as Vamana does. Build + /// `graph` with the prune kind of `metric` (`metric.into()`). pub fn new( config: PiPNNConfig, graph: &'a Config, @@ -268,11 +263,6 @@ impl<'a> PiPNNBuildContext<'a> { pool: &'a ThreadPool, ) -> Result { config.validate()?; - let prune_kind = graph.prune_kind(); - if prune_kind != metric.into() { - return Err(PiPNNConfigError::PruneKind { prune_kind, metric }); - } - Ok(Self { config, graph, @@ -680,40 +670,6 @@ mod construction_tests { assert_eq!(error, PiPNNConfigError::EmptyFanout); } - #[test] - fn a_build_context_requires_the_pruning_kind_of_its_metric() { - // L2, cosine and normalized cosine share triangle pruning. Inner product - // uses occluding pruning. - let pool = thread_pool(1); - for (graph_metric, metric, compatible) in [ - (Metric::L2, Metric::Cosine, true), - (Metric::L2, Metric::CosineNormalized, true), - (Metric::InnerProduct, Metric::InnerProduct, true), - (Metric::L2, Metric::InnerProduct, false), - (Metric::InnerProduct, Metric::L2, false), - ] { - let graph = graph_policy(2, graph_metric).unwrap(); - - let result = PiPNNBuildContext::new(partition_policy(), &graph, metric, &pool); - - let case = format!("graph {graph_metric:?}, metric {metric:?}"); - match result { - Ok(_) => assert!(compatible, "{case}"), - Err(error) => assert_eq!( - (compatible, error), - ( - false, - PiPNNConfigError::PruneKind { - prune_kind: graph_metric.into(), - metric - } - ), - "{case}" - ), - } - } - } - #[test] fn an_empty_dataset_axis_is_rejected_before_building() { let graph = graph_policy(2, Metric::L2).unwrap(); From 375ad3b768ecb7f039d4999aad44caefc08d0ea2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 13 Oct 2026 10:51:42 +0000 Subject: [PATCH 221/229] docs(pipnn): say why only f16 gathers select instructions at run time The baseline x86-64 target has SSE2 instructions for u8 and i8 to f32 conversion, so the generic loop vectorizes without the x86-64-v3 flag. f16 conversion needs F16C, which starts at x86-64-v3. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/conversion.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/diskann/src/graph/pipnn/conversion.rs b/diskann/src/graph/pipnn/conversion.rs index 581010e3ca..b8cbe1beb6 100644 --- a/diskann/src/graph/pipnn/conversion.rs +++ b/diskann/src/graph/pipnn/conversion.rs @@ -12,6 +12,14 @@ //! On an AMD EPYC 7763, software conversion made gathers 4.6 to 22 times slower //! (128 to 1536 dimensions). With the `x86-64-v3` target, both paths are within //! 10% of each other. +//! +//! Other element types use `VectorRepr::as_f32_into` for every target. `f32` +//! rows are copied. The baseline x86-64 target has SSE2 instructions that +//! convert `u8` and `i8` to `f32`, so the compiler vectorizes that loop without +//! the `x86-64-v3` flag. The baseline target has no `f16` conversion +//! instruction; F16C starts at `x86-64-v3`. On the AMD EPYC 7763, the baseline +//! target converted `u8` and `i8` at most 1.8 times slower than `x86-64-v3`, and +//! `f16` 16.5 times slower (rows in L2). use std::any::TypeId; From 665ce5c310beaa8b952232038130ad045d7b5740 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 19 Oct 2026 03:31:29 +0000 Subject: [PATCH 222/229] refactor(pipnn): keep leader values and assignments in matrices assign_to_leaders built flat Vecs, rebuilt views with try_from, and passed fanout down only to recover the row width. Matrix keeps the shape with the data: stripes come from par_window_iter_mut, and assign_point_stripe and group_by_leader read the fanout from the column count. The pipnn feature enables diskann-utils/rayon for the parallel window iterators. Co-Authored-By: Claude Opus 5.5 --- diskann/Cargo.toml | 2 +- diskann/src/graph/pipnn/partitioning.rs | 62 +++++++++++-------------- 2 files changed, 29 insertions(+), 35 deletions(-) diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 30ec5df32f..20745a3f69 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -62,7 +62,7 @@ panic = "warn" default = ["tracing"] # Enable PiPNN batch graph construction. -pipnn = ["dep:diskann-linalg", "dep:parking_lot", "dep:rayon", "tracing"] +pipnn = ["dep:diskann-linalg", "dep:parking_lot", "dep:rayon", "diskann-utils/rayon", "tracing"] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 72652211d7..7698c37a10 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -26,7 +26,7 @@ use crate::{ANNResult, utils::VectorRepr}; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, - views::{MatrixView, MutMatrixView}, + views::{Matrix, MatrixView, MutMatrixView}, }; use rand::{Rng, SeedableRng, prelude::IndexedRandom, rngs::StdRng}; use rayon::prelude::*; @@ -288,21 +288,18 @@ where M: PartitionMetric, T: VectorRepr + Send + Sync, { - let dimension_count = data.ncols(); - let mut leader_values = vec![0.0f32; leader_ids.len() * dimension_count]; - gather_as_f32(data, leader_ids, &mut leader_values)?; + let mut leader_values = Matrix::new(0.0f32, leader_ids.len(), data.ncols()); + gather_as_f32(data, leader_ids, leader_values.as_mut_slice())?; - let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count)?; - let leaders = M::create_leaders(leader_matrix); + let leaders = M::create_leaders(leader_values.as_view()); let leader_count = M::leader_count(&leaders); - let mut assignments = vec![0u32; point_ids.len() * fanout]; + let mut assignments = Matrix::new(0u32, point_ids.len(), fanout); // This runs in the pool of the build context (see `build_graph`). A lease // for each stripe takes one lock, which costs far less than the GEMM and // ranking of the stripe. assignments - .par_chunks_mut(ASSIGNMENT_STRIPE_POINTS * fanout) + .par_window_iter_mut(ASSIGNMENT_STRIPE_POINTS) .zip(point_ids.par_chunks(ASSIGNMENT_STRIPE_POINTS)) .try_for_each(|(stripe_assignments, stripe_ids)| { assign_point_stripe::( @@ -310,7 +307,6 @@ where data, stripe_ids, &leaders, - fanout, &mut stripe_buffers.get_ref(()), stripe_assignments, ) @@ -318,23 +314,21 @@ where Ok(group_by_leader( point_ids, - &assignments, - fanout, + assignments.as_view(), leader_count, )) } /// Gather one stripe of points as `f32` and write the leader columns of the -/// `fanout` nearest leaders of each point to `assignments`. +/// nearest leaders of each point to its row of `assignments`, one per column. #[inline] fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, point_ids: &[u32], leaders: &M::Leaders<'_>, - fanout: usize, buffers: &mut StripeBuffers, - assignments: &mut [u32], + assignments: MutMatrixView<'_, u32>, ) -> ANNResult<()> where A: Simd, @@ -358,8 +352,13 @@ where dimensions, )?; gather_as_f32(data, point_ids, points.as_mut_slice())?; - let output = MutMatrixView::try_from(assignments, point_count, fanout)?; - assign_leaders::(arch, points.as_view(), leaders, output, kernel_workspace) + assign_leaders::( + arch, + points.as_view(), + leaders, + assignments, + kernel_workspace, + ) } /// Group the points by assigned leader. Unassigned slots join no child. @@ -368,12 +367,11 @@ where /// next split samples its leaders by position. fn group_by_leader( points: &[u32], - assignments: &[u32], - fanout: usize, + assignments: MatrixView<'_, u32>, leaders: usize, ) -> Vec> { if points.len() < PARALLEL_GROUPING_MIN_POINTS { - return group_by_leader_serial(points, assignments, fanout, leaders); + return group_by_leader_serial(points, assignments, leaders); } let chunk_points = points.len().div_ceil(rayon::current_num_threads()); @@ -381,8 +379,8 @@ fn group_by_leader( // order. let locals: Vec<_> = points .par_chunks(chunk_points) - .zip(assignments.par_chunks(chunk_points * fanout)) - .map(|(points, assignments)| group_by_leader_serial(points, assignments, fanout, leaders)) + .zip(assignments.par_window_iter(chunk_points)) + .map(|(points, assignments)| group_by_leader_serial(points, assignments, leaders)) .collect(); let mut sizes = vec![0usize; leaders]; @@ -408,18 +406,17 @@ fn group_by_leader( fn group_by_leader_serial( points: &[u32], - assignments: &[u32], - fanout: usize, + assignments: MatrixView<'_, u32>, leaders: usize, ) -> Vec> { let mut sizes = vec![0usize; leaders]; - for &leader in assignments { + for &leader in assignments.as_slice() { if leader != UNASSIGNED { sizes[leader as usize] += 1; } } let mut clusters: Vec> = sizes.into_iter().map(Vec::with_capacity).collect(); - for (&point, point_assignments) in points.iter().zip(assignments.chunks_exact(fanout)) { + for (&point, point_assignments) in points.iter().zip(assignments.row_iter()) { for &leader in point_assignments { if leader != UNASSIGNED { clusters[leader as usize].push(point); @@ -618,23 +615,19 @@ mod tests { (&[2, 1][..], 1, &[1, 0][..]), (&[1, 0, 4, 2][..], 2, &[0, 1, 2, 1, 0, 1, 1, 0][..]), ] { - let mut assignments = vec![0; ids.len() * fanout]; + let mut assignments = Matrix::new(0, ids.len(), fanout); assign_point_stripe::<_, L2, _>( ARCH, data, ids, &leaders, - fanout, &mut buffers, - &mut assignments, + assignments.as_mut_view(), ) .unwrap(); - let actual: Vec<_> = assignments - .chunks_exact(fanout) - .map(<[u32]>::to_vec) - .collect(); + let actual: Vec<_> = assignments.row_iter().map(<[u32]>::to_vec).collect(); let expected: Vec<_> = expected.chunks_exact(fanout).map(<[u32]>::to_vec).collect(); assert_eq!( sorted_members_per_row(&actual), @@ -664,6 +657,7 @@ mod tests { ] { let ids: Vec<_> = (0..point_count).map(|i| 1_000_000 - i as u32).collect(); let assignments: Vec<_> = (0..point_count).flat_map(|i| pattern[i % 4]).collect(); + let assignments = MatrixView::try_from(assignments.as_slice(), point_count, 2).unwrap(); let expected = [ ids.iter().copied().step_by(2).collect(), ids.iter().copied().skip(1).step_by(4).collect(), @@ -671,7 +665,7 @@ mod tests { vec![], ]; - let actual = pool.install(|| group_by_leader(&ids, &assignments, 2, 4)); + let actual = pool.install(|| group_by_leader(&ids, assignments, 4)); assert_eq!( sorted_members_per_row(&actual), From 0f568aa3038110427b8d07e82eea5b74a90d8505 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 19 Oct 2026 05:55:40 +0000 Subject: [PATCH 223/229] refactor(pipnn): pass the leaf kernel output to the merge as a matrix view add_symmetric_neighbors took the neighbor buffer as a slice plus leaf_k to recover the row width. It now takes the kernel output view and walks its rows. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/leaf_build.rs | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 3b61957e55..bb3a01e207 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -121,22 +121,25 @@ where return Ok(()); } + buffers.prepare_local_adjacency(point_count); let point_values = &mut buffers.point_values[..point_count * data.ncols()]; gather_as_f32(data, point_ids, point_values)?; let points = MatrixView::try_from(&*point_values, point_count, data.ncols())?; - let neighbor_count = point_count * leaf_k; - let output = MutMatrixView::try_from( - &mut buffers.neighbors[..neighbor_count], + let mut output = MutMatrixView::try_from( + &mut buffers.neighbors[..point_count * leaf_k], point_count, leaf_k, )?; - select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace)?; + select_leaf_neighbors::( + arch, + points, + output.as_mut_view(), + &mut buffers.kernel_workspace, + )?; - buffers.prepare_local_adjacency(point_count); add_symmetric_neighbors( point_ids, - leaf_k, - &buffers.neighbors[..neighbor_count], + output.as_view(), &mut buffers.local_adjacency[..point_count], ); for (&point_id, additions) in point_ids.iter().zip(&buffers.local_adjacency) { @@ -151,11 +154,10 @@ where /// the edge lists of both points. fn add_symmetric_neighbors( point_ids: &[u32], - leaf_k: usize, - neighbors: &[Candidate], + neighbors: MatrixView<'_, Candidate>, local_adjacency: &mut [Vec], ) { - for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for (source, source_neighbors) in neighbors.row_iter().enumerate() { for neighbor in source_neighbors { // Skip the empty slots that the kernel leaves for NaN and infinite // distances. From 9fe4551767d9998726e82951e068fe1188118831 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 19 Oct 2026 07:02:31 +0000 Subject: [PATCH 224/229] test(pipnn): derive the stripe buffer reuse expectations from distances The test hard-coded leader columns for three stripes with changing fanout. Width changes belong to the kernel workspace, which the kernel tests cover. Keep one stripe that grows the buffer and one that reads a prefix of it, and compute the nearest leaders from the coordinates. Rows are compared in order, so stale rows from the previous stripe fail the test. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 32 +++++++++++++++---------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 7698c37a10..e58207b551 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -603,18 +603,18 @@ mod tests { } #[test] - fn reused_stripe_buffers_do_not_carry_points_or_widths() { + fn reused_stripe_buffers_do_not_carry_points() { + // Points and leaders lie on a line, so the nearest leaders follow from + // the coordinates. The first stripe grows the buffer. The second stripe is + // smaller and reads only a prefix of it. let values = [100.0_f32, 0.0, 8.0, 10.0, 2.0]; let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); let leader_values = [0.0_f32, 10.0, 100.0]; let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); + let fanout = 2; let mut buffers = StripeBuffers::default(); - for (ids, fanout, expected) in [ - (&[4, 2, 0, 3, 1][..], 2, &[0, 1, 1, 0, 2, 1, 1, 0, 0, 1][..]), - (&[2, 1][..], 1, &[1, 0][..]), - (&[1, 0, 4, 2][..], 2, &[0, 1, 2, 1, 0, 1, 1, 0][..]), - ] { + for ids in [&[4, 2, 0, 3, 1][..], &[2, 1][..]] { let mut assignments = Matrix::new(0, ids.len(), fanout); assign_point_stripe::<_, L2, _>( @@ -627,13 +627,21 @@ mod tests { ) .unwrap(); + // The `fanout` leader columns nearest to each point. No two leaders + // are at the same distance from a point. + let distance = + |id: u32, column: u32| (values[id as usize] - leader_values[column as usize]).abs(); + let expected: Vec> = ids + .iter() + .map(|&id| { + let mut columns: Vec = (0..leader_values.len() as u32).collect(); + columns.sort_by(|&a, &b| distance(id, a).total_cmp(&distance(id, b))); + columns.truncate(fanout); + columns + }) + .collect(); let actual: Vec<_> = assignments.row_iter().map(<[u32]>::to_vec).collect(); - let expected: Vec<_> = expected.chunks_exact(fanout).map(<[u32]>::to_vec).collect(); - assert_eq!( - sorted_members_per_row(&actual), - sorted_members_per_row(&expected), - "points {ids:?}, fanout={fanout}" - ); + assert_eq!(actual, expected, "points {ids:?}"); } } From 29d32aec1ee6d9bc7206f30db5fe9d6b8a26df43 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 19 Oct 2026 07:30:07 +0000 Subject: [PATCH 225/229] test(pipnn): say why the stripe test ranks leaders by absolute difference Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index e58207b551..bf47badd53 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -627,8 +627,9 @@ mod tests { ) .unwrap(); - // The `fanout` leader columns nearest to each point. No two leaders - // are at the same distance from a point. + // The `fanout` leader columns nearest to each point. In one dimension, + // |x - leader| orders the leaders the same way as squared L2. No two + // leaders are at the same distance from a point. let distance = |id: u32, column: u32| (values[id as usize] - leader_values[column as usize]).abs(); let expected: Vec> = ids From 467ffe6d7605cb440e271dd8ba7017b28833038e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 19 Oct 2026 07:44:15 +0000 Subject: [PATCH 226/229] test(pipnn): keep only the replica seed check in the replica test The other assertions repeated leaf checks that other partition tests make. The remaining one is the only test that fails when every replica reuses one seed. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 51 ++++++++----------------- 1 file changed, 16 insertions(+), 35 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index bf47badd53..945fc060d1 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -888,7 +888,9 @@ mod tests { } #[test] - fn a_second_replica_adds_new_leaf_memberships() { + fn replicas_use_different_seeds() { + // With one seed for every replica, the second replica would repeat the + // leaves of the first. let point_count = 33; let values = dense_points(point_count, 17, 1290); let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); @@ -899,42 +901,21 @@ mod tests { ..splitting_config() }; let pool = thread_pool(2); - let first = pool - .install(|| partition::<_, L2, _>(ARCH, data, &config)) - .unwrap(); + let leaves = |replicas| { + let config = PiPNNConfig { + replicas: nz(replicas), + ..config.clone() + }; + let leaves = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + sorted_leaf_memberships(&leaves) + }; - let combined = pool - .install(|| { - partition::<_, L2, _>( - ARCH, - data, - &PiPNNConfig { - replicas: nz(2), - ..config - }, - ) - }) - .unwrap(); + let one = leaves(1); + let two = leaves(2); - let first = sorted_leaf_memberships(&first); - let combined = sorted_leaf_memberships(&combined); - assert!(first.iter().all(|leaf| combined.contains(leaf))); - // This fixed cloud and partial sample produce different memberships in - // the second pass. Repeating the first seed would only duplicate leaves. - assert!(combined.iter().any(|leaf| !first.contains(leaf))); - // Fanout one gives every point exactly one membership per replica. - let mut memberships = vec![0; point_count]; - for leaf in &combined { - assert!( - !leaf.is_empty() && leaf.len() <= config.c_max.get(), - "leaf {leaf:?}" - ); - assert_eq!(leaf.iter().collect::>().len(), leaf.len()); - for &id in leaf { - memberships[id as usize] += 1; - } - } - assert_eq!(memberships, vec![2; point_count]); + assert!(two.iter().any(|leaf| !one.contains(leaf))); } #[test] From a2dc064009dfee67969a0c996c7559c5cf65638d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:13:16 +0000 Subject: [PATCH 227/229] docs(pipnn): say when a fanout-1 split can make no progress The comment named only copied leaders. With inner product, a leader with a large norm can also win every point. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 945fc060d1..21dbbc3358 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -239,12 +239,13 @@ where if child.len() <= c_max { leaves.push(child); } else if fanout == 1 && child.len() == cluster.point_ids.len() { - // Every point chose this leader, so the split made no progress. This - // happens when all sampled leaders are copies of one vector: each - // point is equally near to all of them. Another split can stall the - // same way, so cut the child into equal leaves in point order. The - // cut groups points arbitrarily, which can lower recall but keeps - // the graph valid. + // Every point chose this leader, so the split made no progress. With + // L2 or cosine, each leader is nearest to itself, so this needs all + // sampled leaders to be copies of one vector. With inner product, a + // leader with a large norm can also win every point. Another split can + // stall the same way, so cut the child into equal leaves in point + // order. The cut groups points arbitrarily, which can lower recall but + // keeps the graph valid. let leaf_len = child.len().div_ceil(child.len().div_ceil(c_max)); leaves.extend(child.chunks(leaf_len).map(<[u32]>::to_vec)); } else { From 4547a52c9183b5a5ec8e282ef6cac7a5c39cef49 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:21:31 +0000 Subject: [PATCH 228/229] refactor(pipnn): build CosineNormalized integer data as requested PiPNN switched u8 and i8 data from CosineNormalized to Cosine with a TypeId check. CosineNormalized assumes unit-norm vectors, and the caller owns that assumption for every element type. Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/mod.rs | 47 ++-------------------------------- 1 file changed, 2 insertions(+), 45 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index fd9aab13c1..425b7e1ceb 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -275,8 +275,7 @@ impl<'a> PiPNNBuildContext<'a> { /// Build one adjacency list for each point of `data`. /// /// The graph links dataset points only; the caller selects start points and -/// serializes the index. `CosineNormalized` on `u8` or `i8` data builds with -/// `Cosine`, because converted integer vectors do not have unit norm. +/// serializes the index. pub fn build_graph( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, @@ -329,8 +328,7 @@ where data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, ) -> ANNResult>> { - // Converting raw integer coordinates does not normalize their vectors. - let metric = effective_metric::(context.metric); + let metric = context.metric; match metric { Metric::L2 => build_graph_for::(arch, data, context, metric), Metric::Cosine => build_graph_for::(arch, data, context, metric), @@ -368,19 +366,6 @@ where .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric))) } -/// Return the metric that the build uses for element type `T`. See [`build_graph`]. -fn effective_metric(metric: Metric) -> Metric { - use std::any::TypeId; - - if metric == Metric::CosineNormalized - && (TypeId::of::() == TypeId::of::() || TypeId::of::() == TypeId::of::()) - { - Metric::Cosine - } else { - metric - } -} - #[cfg(test)] mod test_support { use super::simd::Simd; @@ -785,34 +770,6 @@ mod construction_tests { } } - #[test] - fn normalized_cosine_requests_on_raw_integers_use_vector_norms() { - // Cosine pairs 0 <-> 1 and 2 <-> 3. Unnormalized dot products instead - // connect 1 to 3, so treating raw integers as unit vectors changes edges. - // The first build checks leaf ranking. The second build selects three - // leaf neighbors and checks final pruning to one. - fn check(values: [T; 8]) { - for (leaf_k, degree) in [(1, 3), (3, 1)] { - let config = PiPNNConfig { - leaf_k: nz(leaf_k), - ..partition_policy() - }; - - let actual = build(&values, 2, config, degree, Metric::CosineNormalized, 2); - - assert_eq!( - actual, - [vec![1], vec![0], vec![3], vec![2]], - "{}, leaf_k={leaf_k}, degree={degree}", - std::any::type_name::() - ); - } - } - - check([1_i8, 0, 4, 1, 1, 2, 0, 9]); - check([1_u8, 0, 4, 1, 1, 2, 0, 9]); - } - #[test] fn final_pruning_cuts_merged_candidates_to_the_degree() { let config = PiPNNConfig { From a724b99e676c59bc0864e635711aa660648d9124 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 19 Oct 2026 08:34:29 +0000 Subject: [PATCH 229/229] docs(pipnn): define the assignment stripe and keep only why it is 256 Co-Authored-By: Claude Opus 5.5 --- diskann/src/graph/pipnn/partitioning.rs | 26 ++++++++++--------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 21dbbc3358..314a1c1d51 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -17,6 +17,10 @@ //! nearest leaders. This module samples leaders, gathers vectors, and groups the //! assignments into children. //! +//! A split assigns its points in parallel stripes. An assignment stripe is a +//! batch of up to `ASSIGNMENT_STRIPE_POINTS` points that one task gathers as +//! `f32` and ranks against all leaders of the split. +//! //! Buffer sizes use plain multiplication. The IDs of one split are distinct //! dataset rows, so a gathered `rows x dimensions` buffer is not larger than the //! dataset. An assignment buffer holds `points x fanout` IDs. With at most @@ -46,22 +50,12 @@ use super::{ /// split of 10M points at `p_samp = 0.005` samples 50,000 leaders and does 50 /// times the work. const LEADER_CAP: usize = 1_000; -/// Number of points in one assignment stripe. A stripe gathers its points as -/// `f32` and computes their distances to all leaders, so it moves -/// `points x (leaders + dimensions) x 4` bytes through the cache. -/// -/// Each stripe repeats fixed work, such as one GEMM call over the leader matrix, -/// so 16-point stripes took 1.3 to 2 times as long as the best height. Tall -/// stripes slow down when the stripes of all workers exceed the last-level -/// cache. The L2 cache is not the limit: one worker ran 16 MiB stripes at full -/// speed on a CPU with 512 KiB of L2 per core. -/// -/// Sweeps on Zen 3, Zen 4, Emerald Rapids and Ice Lake CPUs covered f32, f16 and -/// u8 data, 128 to 1,536 dimensions, 64 to 1,000 leaders and 1 to 16 workers. -/// 256 points stayed within 11% of the best height, and within 1.4% on average. -/// With 8 leaders and 1,536 dimensions, where the gather dominates, the loss was -/// up to 27%. A rule from the detected cache size did not do better, because a -/// virtual machine reports a last-level cache that other tenants also use. +/// Number of points in one assignment stripe. Short stripes repeat the fixed +/// cost of a kernel call, and tall stripes overflow the last-level cache once +/// every worker holds one. 256 points stayed within 1.4% of the best height on +/// average in sweeps over x86 CPUs, data types, dimensions and leader counts. A +/// height from the detected cache size did no better, because a VM reports a +/// last-level cache that other tenants share. const ASSIGNMENT_STRIPE_POINTS: usize = 256; /// Clusters with at least this many points group their points by leader in /// parallel. Such clusters occur at the first levels, where few splits run at