diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c874e1b361..30ee4ae606 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -354,6 +354,8 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ + --package diskann \ + --features pipnn \ -- --skip compile_tests \ --skip pivots::tests::run_test_happy_path @@ -416,6 +418,8 @@ jobs: --package diskann-wide \ --package diskann-vector \ --package diskann-quantization \ + --package diskann \ + --features pipnn \ -- --skip compile_tests test-workspace: @@ -459,6 +463,7 @@ jobs: os: - windows-latest - ubuntu-latest + - ubuntu-24.04-arm steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 90c33e0efd..1df300d528 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/Cargo.lock b/Cargo.lock index 2588eee92a..40240e77c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -438,6 +438,7 @@ dependencies = [ "anyhow", "bytemuck", "dashmap", + "diskann-linalg", "diskann-utils", "diskann-vector", "diskann-wide", @@ -448,6 +449,7 @@ dependencies = [ "pin-project", "rand", "relative-path 2.0.1", + "rstest", "serde", "serde_json", "thiserror 2.0.17", diff --git a/diskann-linalg/src/faer.rs b/diskann-linalg/src/faer.rs index 700396de4e..f4f498a13b 100644 --- a/diskann-linalg/src/faer.rs +++ b/diskann-linalg/src/faer.rs @@ -53,6 +53,33 @@ 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. +/// +/// `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}; + + 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..778fb94cf8 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,57 +180,33 @@ 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); Ok(()) } +/// Compute the lower triangle of $C = 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. +/// +/// # 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> { + 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(()) +} + /// Compute the SVD of the provided matrix implicit row-major matrix `data`. /// /// * `m`: The number of rows in `a`. @@ -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-utils/src/views.rs b/diskann-utils/src/views.rs index a9352918c9..e38b5b6c58 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 len != nrows * ncols { + if nrows.checked_mul(ncols) != Some(len) { Err(TryFromError { data, nrows, ncols }) } else { Ok(Self { data, nrows, ncols }) @@ -1053,6 +1053,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 + 1, 2).is_err()); } #[test] 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; diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 14f6ba0746..c0016416dc 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -30,12 +30,14 @@ 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 } 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"] } @@ -56,6 +58,9 @@ panic = "warn" [features] default = ["tracing"] +# Enable PiPNN numerical kernels. +pipnn = ["dep:diskann-linalg"] + # 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/src/graph/pipnn/kernel_metric.rs b/diskann/src/graph/pipnn/kernel_metric.rs new file mode 100644 index 0000000000..333fb11eff --- /dev/null +++ b/diskann/src/graph/pipnn/kernel_metric.rs @@ -0,0 +1,276 @@ +/* + * 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 new file mode 100644 index 0000000000..d87a16c426 --- /dev/null +++ b/diskann/src/graph/pipnn/kernel_metric/leaf.rs @@ -0,0 +1,487 @@ +/* + * 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 new file mode 100644 index 0000000000..277ece78d5 --- /dev/null +++ b/diskann/src/graph/pipnn/kernel_metric/partition.rs @@ -0,0 +1,485 @@ +/* + * 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()) } +} + +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) { + norms.resize(points.nrows(), 0.0); + for (norm, point) in norms.iter_mut().zip(points.row_iter()) { + *norm = FastL2NormSquared.evaluate(point).sqrt(); + } + } + + 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(); + } + } + + #[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); + } + + 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 new file mode 100644 index 0000000000..71a8b15556 --- /dev/null +++ b/diskann/src/graph/pipnn/leaf_kernel.rs @@ -0,0 +1,1714 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! 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 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. +//! +//! Strict comparisons keep scan order for equal distances. They do not rank NaN. +//! 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. +//! [`LeafKernelWorkspace`] stores reusable numerical scratch. + +use crate::{ANNError, ANNResult}; +use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_wide::{SIMDPartialOrd, SIMDVector}; + +use super::{ + kernel_metric::LeafMetric, + simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, +}; + +/// One leaf-local neighbor and its 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`. + pub(super) distance: f32, +} + +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. + 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 { + dot_scratch: Vec, + norm_scratch: Vec, + worst: Vec, +} + +/// Validation error returned by the dot-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. + #[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. +/// +/// # Errors +/// +/// 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: PiPNNSIMDSchema, + M: LeafMetric, +{ + 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>, + worst: &mut Vec, +) -> Result<(), LeafKernelError> +where + A: PiPNNSIMDSchema, + M: LeafMetric, +{ + validate_neighbor_count(input, &output)?; + let neighbor_count = output.ncols(); + if neighbor_count == 0 { + return Ok(()); + } + + worst.resize(input.nrows(), 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::( + arch, + input, + norms, + output.as_mut_slice(), + neighbor_count, + 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( + input: MatrixView<'_, f32>, + 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 { + return Err(LeafKernelError::InvalidNeighborCount { + points: point_count, + neighbors: neighbor_count, + maximum: maximum_neighbors, + }); + } + Ok(()) +} + +/// Select neighbors with a fixed output width. +fn scan_fixed_width( + arch: A, + input: MatrixView<'_, f32>, + norms: &[f32], + output: &mut [LeafNeighbor], + worst: &mut [f32], +) 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_eligible_neighbor(&mut rows[source], target, distance) + }); +} + +/// Select neighbors with a runtime output width. +fn scan_runtime_width( + arch: A, + input: MatrixView<'_, f32>, + norms: &[f32], + output: &mut [LeafNeighbor], + width: usize, + worst: &mut [f32], +) where + 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) + }); +} + +/// 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. +#[inline(never)] +fn scan_point_pairs( + arch: A, + input: MatrixView<'_, f32>, + norms: &[f32], + 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. + 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)); + // 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 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 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 < source_worst { + source_worst = insert(source, (target + 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 target_source = target + 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 }; + } + } + target += A::Vector::LANES; + } + + 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); + if distance < source_worst { + 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(target, source as u32, distance); + // SAFETY: `target < source < worst.len()`. + unsafe { *worst_ptr.add(target) = new_worst }; + } + target += 1; + } + // SAFETY: `source < worst.len()`. + unsafe { *worst_ptr.add(source) = source_worst }; + } +} + +/// Insert one value that the caller has already found eligible. +/// +/// 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; +} + +impl SortedInsert for [T; 1] { + #[inline(always)] + fn insert_eligible_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_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; + self[1] = first; + first + } else { + self[1] = value; + value + } + } +} + +impl SortedInsert for [T; 3] { + #[inline(always)] + 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; + self[1] = first; + self[2] = second; + second + } else if precedes(value, second) { + self[1] = value; + self[2] = second; + second + } else { + self[2] = value; + value + } + } +} + +impl SortedInsert for [T] { + #[inline(always)] + 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]) { + self[slot] = self[slot - 1]; + slot -= 1; + } + self[slot] = value; + self[last] + } +} + +/// 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: SortedInsert + ?Sized, +{ + neighbors + .insert_eligible_sorted_by( + LeafNeighbor::new(target, distance), + |candidate, retained| candidate.distance < retained.distance, + ) + .distance +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::graph::pipnn::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + use diskann_utils::views::{MatrixView, MutMatrixView}; + use diskann_vector::distance::Metric; + + mod test_support { + use std::cmp::Ordering; + + use super::*; + use diskann_wide::arch::{self, Target1}; + + 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 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(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) + } + + /// 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 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) + } + + 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) + } + } + } + } + + 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]); + } + 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. + 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]; + } + } + 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 + } + } + + mod insert_eligible_neighbor_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 mut actual_neighbors = [retained_neighbor]; + + // When + insert_eligible_neighbor( + &mut actual_neighbors, + nearer_candidate.target, + nearer_candidate.distance, + ); + + // 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); + } + + #[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 mut actual_neighbors = [nearest, farthest]; + + // When + insert_eligible_neighbor( + &mut actual_neighbors, + nearer_candidate.target, + nearer_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[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 mut actual_neighbors = [nearest, farthest]; + + // When + insert_eligible_neighbor( + &mut actual_neighbors, + eligible_candidate.target, + eligible_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[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 mut actual_neighbors = [nearest, middle, farthest]; + + // When + insert_eligible_neighbor( + &mut actual_neighbors, + nearer_candidate.target, + nearer_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[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 mut actual_neighbors = [nearest, middle, farthest]; + + // When + insert_eligible_neighbor( + &mut actual_neighbors, + middle_candidate.target, + middle_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[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 mut actual_neighbors = [nearest, middle, farthest]; + + // When + insert_eligible_neighbor( + &mut actual_neighbors, + eligible_candidate.target, + eligible_candidate.distance, + ); + + // Then + assert_eq!(actual_neighbors, expected_neighbors); + } + + #[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 mut actual_neighbors = [first, second, third, fourth]; + + // When + insert_eligible_neighbor( + actual_neighbors.as_mut_slice(), + candidate.target, + candidate.distance, + ); + + // 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 { + 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 + 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); + } + } + + mod select_leaf_neighbors_tests { + use super::*; + + #[test] + 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(); + 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 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]; + 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 + ); + } + } + + mod rank_leaf_dots_tests { + use super::test_support::*; + 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, + ) { + // 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); + + // When + let actual_neighbors = run_rank_leaf_dots(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 = 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); + } + + #[test] + fn selects_the_highest_similarity_neighbor_for_each_point_with_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)), + ]; + 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(); + + // 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, 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)), + ]; + 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(); + + // 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 = 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); + } + + #[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 = 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); + } + + #[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 = 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()); + 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_dots(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 = 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); + } + + #[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_dots(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 = 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!( + 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 = 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!( + 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; + + // 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(); + + // 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 gram = gram_with_uniform_self_dots(points, f32::MAX); + gram[3 * points] = -f32::MAX; + 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(); + + // Then + assert_eq!( + actual_neighbors[3 * expected_leaf_k + expected_leaf_k - 1], + expected_last_neighbor + ); + } + + #[test] + fn scalar_nan_distance_leaves_the_neighbor_slot_unassigned() { + // Given + let gram = [1.0, f32::NAN, f32::NAN, 1.0]; + 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(); + + // 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 = 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( + #[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, expected_distance); + + // When + let actual_neighbors = run_rank_leaf_dots(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 = run_rank_leaf_dots(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 + 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_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_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(); + + // Then + assert_eq!(actual_neighbors, expected_no_neighbors); + } + + #[test] + fn neighbor_width_equal_to_point_count_is_rejected() { + // 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 expected_error = LeafKernelError::InvalidNeighborCount { + points: point_count, + neighbors: invalid_neighbor_width, + maximum: maximum_non_self_width, + }; + + // 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(); + + // Then + assert_eq!(actual_error, expected_error); + } + } +} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs new file mode 100644 index 0000000000..32588f73c4 --- /dev/null +++ b/diskann/src/graph/pipnn/mod.rs @@ -0,0 +1,33 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! 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; + +#[allow(dead_code)] +mod leaf_kernel; +#[allow(dead_code)] +mod partition_kernel; diff --git a/diskann/src/graph/pipnn/partition_kernel.rs b/diskann/src/graph/pipnn/partition_kernel.rs new file mode 100644 index 0000000000..a41f675cb5 --- /dev/null +++ b/diskann/src/graph/pipnn/partition_kernel.rs @@ -0,0 +1,844 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! 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. +//! +//! 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 +//! unfilled output slot contains [`UNASSIGNED_LEADER`]. + +use std::marker::PhantomData; + +use crate::{ANNError, ANNResult}; +use diskann_linalg::Transpose; +use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_wide::{SIMDMask, SIMDVector}; + +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; + +/// 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)] +struct PartitionInput<'a> { + dots: MatrixView<'a, f32>, + norms: PartitionNorms<'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. +pub(super) fn assign_leaders( + arch: A, + points: MatrixView<'_, f32>, + leaders: &PreparedLeaders<'_, M>, + output: MutMatrixView<'_, u32>, + workspace: &mut PartitionKernelWorkspace, +) -> ANNResult<()> +where + A: PiPNNSIMDSchema, + 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)>, +) 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)); + 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: A, + dots: MatrixView<'_, f32>, + norms: PartitionNorms<'_>, + 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)) + .enumerate() + { + 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); + } + + 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 (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. +/// +/// `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)]) +where + F: PiPNNSIMDVector, +{ + let threshold = F::splat(scores.arch(), ranked_leaders[ranked_leaders.len() - 1].1); + let eligible = scores.lt_simd(threshold); + if eligible.none() { + return; + } + + let values = scores.to_lane_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; + 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. `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(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; + } + + ranked_leaders[threshold] = (leader, score); + 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)] +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; + + mod test_support { + use super::*; + use diskann_wide::arch::{self, Target1}; + + struct KernelCall<'a> { + input: PartitionInput<'a>, + output: MutMatrixView<'a, u32>, + ranked_leaders: &'a mut Vec<(u32, f32)>, + } + + 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, + ), + } + } + } + + 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, + }, + } + } + + 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 { + 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) + } + } + } + } + + 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() + { + 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 + } + + /// 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() + .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_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]; + + // 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() { + // 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(); + + // Then + assert_eq!( + 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 + ); + } + } + + mod rank_leader_dots_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::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] 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 = run_rank_leader_dots(metric, input, nearest_leader_count); + + // Then + assert_eq!(actual_assignments, expected_assignments); + } + + #[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 = 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 + 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 = 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_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]; + + // When + let actual_assignment = run_rank_leader_dots( + Metric::Cosine, + partition_input( + &dots, + point_count, + leader_count, + &point_norms, + &leader_norms, + ), + nearest_leader_count, + ); + + // Then + assert_eq!(actual_assignment, expected_first_leader); + } + + #[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 = 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_all_leaders_in_scan_order); + } + + #[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, &[], &[]); + + // 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_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 + 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/simd.rs b/diskann/src/graph/pipnn/simd.rs new file mode 100644 index 0000000000..621114c761 --- /dev/null +++ b/diskann/src/graph/pipnn/simd.rs @@ -0,0 +1,67 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! SIMD schema for PiPNN numerical kernels. + +use diskann_wide::{ + Architecture, Const, SIMDFloat, SIMDMask, SIMDSelect, SIMDVector, SupportedLaneCount, +}; + +/// 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 +{ + /// 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, + 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()) + } + + #[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. +pub(super) trait PiPNNSIMDSchema: Architecture { + /// SIMD vector used by every numerical stage. + type Vector: PiPNNSIMDVector; +} + +impl PiPNNSIMDSchema for A +where + A: Architecture, + DefaultVector: PiPNNSIMDVector, +{ + type Vector = DefaultVector; +}