PiPNN 1/6: add numerical kernels

.cargo/mutants.toml

@@ -0,0 +1,11 @@
+# Kernel mutations that cannot be distinguished by native correctness tests.
+# V4 and Neon wrappers are exercised by SDE and AArch64 CI respectively.
+# The remaining transformations preserve outputs and only switch between
+# equivalent scalar-tail or zero-clamp expressions.
+exclude_re = [
+  "diskann-pipnn/src/(leaf|partition)_kernel\\.rs:.*Target<.*V4",
+  "diskann-pipnn/src/(leaf|partition)_kernel\\.rs:.*Target<.*Neon",
+  "diskann-pipnn/src/leaf_kernel\\.rs:.*replace < with <= in pair_distance",
+  "diskann-pipnn/src/partition_kernel\\.rs:.*replace \\* with / in process_(unary|binary)",
+  "diskann-pipnn/src/leaf_kernel\\.rs:.*replace > with >= in .*run_simd",
+]

.github/workflows/ci.yml

@@ -355,6 +355,7 @@ jobs:
               --package diskann-wide \
               --package diskann-vector \
               --package diskann-quantization \
+              --package diskann-pipnn \
               -- --skip compile_tests \
                  --skip pivots::tests::run_test_happy_path
 
@@ -415,6 +416,7 @@ jobs:
             --package diskann-wide \
             --package diskann-vector \
             --package diskann-quantization \
+            --package diskann-pipnn \
             -- --skip compile_tests
 
   test-workspace:

Cargo.toml

@@ -13,6 +13,7 @@ members = [
     "diskann-quantization",
     # Algorithm
     "diskann",
+    "diskann-pipnn",
     # Providers
     "diskann-providers",
     "diskann-disk",
@@ -58,6 +59,7 @@ diskann-utils = { path = "diskann-utils", default-features = false, version = "0
 diskann-quantization = { path = "diskann-quantization", default-features = false, version = "0.55.0" }
 # Algorithm
 diskann = { path = "diskann", version = "0.55.0" }
+diskann-pipnn = { path = "diskann-pipnn", version = "0.55.0" }
 # Providers
 diskann-providers = { path = "diskann-providers", default-features = false, version = "0.55.0" }
 diskann-disk = { path = "diskann-disk", version = "0.55.0" }

diskann-linalg/src/faer.rs

@@ -53,6 +53,29 @@ pub(super) fn sgemm_impl(
     faer::linalg::matmul::matmul(c, beta, a, b, alpha, Par::Seq)
 }
 
+/// Implements the public lower-triangular AAT operation.
+///
+/// The caller has already validated the matrix dimensions.
+pub(super) fn sgemm_aat_lower_impl(m: usize, k: usize, a: &[f32], c: &mut [f32]) {
+    use faer::linalg::matmul::triangular::{matmul, BlockStructure};
+
+    let a = faer::mat::MatRef::from_row_major_slice(a, m, k);
+    let at = a.transpose();
+    let c = faer::mat::MatMut::from_row_major_slice_mut(c, m, m);
+
+    matmul(
+        c,
+        BlockStructure::TriangularLower,
+        faer::Accum::Replace,
+        a,
+        BlockStructure::Rectangular,
+        at,
+        BlockStructure::Rectangular,
+        1.0,
+        Par::Seq,
+    );
+}
+
 /// See the documentation for `svd_into`.
 ///
 /// The implementation may assume the the specified invariants hold for the sizes of the

diskann-linalg/src/lib.rs

@@ -205,6 +205,49 @@ pub fn sgemm(
     Ok(())
 }
 
+/// Computes the lower triangle of $C = A A^\mathsf{T}$ for a dense row-major
+/// $m \times k$ matrix $A$.
+///
+/// The lower triangle, including the diagonal, is overwritten. The upper
+/// triangle of the $m \times m$ destination is left unchanged.
+///
+/// # Errors
+///
+/// Returns an error if a matrix-size calculation overflows or either slice does
+/// not match its declared dimensions.
+pub fn sgemm_aat_lower(a: &[f32], m: usize, k: usize, c: &mut [f32]) -> Result<(), SgemmError> {
+    let expected_a_len = m.checked_mul(k).ok_or(SgemmError::DimensionOverflow {
+        matrix_name: MatrixName::A,
+        rows: m,
+        cols: k,
+    })?;
+    if a.len() != expected_a_len {
+        return Err(SgemmError::InvalidMatrixDimensions {
+            matrix_name: MatrixName::A,
+            expected_rows: m,
+            expected_cols: k,
+            actual_len: a.len(),
+        });
+    }
+
+    let expected_c_len = m.checked_mul(m).ok_or(SgemmError::DimensionOverflow {
+        matrix_name: MatrixName::C,
+        rows: m,
+        cols: m,
+    })?;
+    if c.len() != expected_c_len {
+        return Err(SgemmError::InvalidMatrixDimensions {
+            matrix_name: MatrixName::C,
+            expected_rows: m,
+            expected_cols: m,
+            actual_len: c.len(),
+        });
+    }
+
+    faer::sgemm_aat_lower_impl(m, k, a, c);
+    Ok(())
+}
+
 /// Compute the SVD of the provided matrix implicit row-major matrix `data`.
 ///
 /// * `m`: The number of rows in `a`.

diskann-linalg/tests/sgemm_aat_lower.rs

@@ -0,0 +1,109 @@
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
+
+use diskann_linalg::{sgemm_aat_lower, MatrixName, SgemmError};
+
+#[test]
+fn computes_lower_triangle_and_preserves_upper_triangle() {
+    #[rustfmt::skip]
+    let a = [
+        1.0, 2.0,
+        3.0, 4.0,
+        5.0, 6.0,
+    ];
+    let untouched = -123.0;
+    let mut c = [untouched; 9];
+
+    sgemm_aat_lower(&a, 3, 2, &mut c).unwrap();
+
+    #[rustfmt::skip]
+    assert_eq!(c, [
+         5.0, untouched, untouched,
+        11.0,     25.0, untouched,
+        17.0,     39.0,      61.0,
+    ]);
+}
+
+#[test]
+fn accepts_a_matrix_with_no_rows() {
+    sgemm_aat_lower(&[], 0, 3, &mut []).unwrap();
+}
+
+#[test]
+fn zero_inner_dimension_zeros_only_the_lower_triangle() {
+    let untouched = -123.0;
+    let mut c = [untouched; 9];
+
+    sgemm_aat_lower(&[], 3, 0, &mut c).unwrap();
+
+    #[rustfmt::skip]
+    assert_eq!(c, [
+        0.0, untouched, untouched,
+        0.0,       0.0, untouched,
+        0.0,       0.0,       0.0,
+    ]);
+}
+
+#[test]
+fn rejects_invalid_input_dimensions() {
+    let mut c = [0.0; 4];
+
+    let error = sgemm_aat_lower(&[0.0; 3], 2, 2, &mut c).unwrap_err();
+
+    assert_eq!(
+        error,
+        SgemmError::InvalidMatrixDimensions {
+            matrix_name: MatrixName::A,
+            expected_rows: 2,
+            expected_cols: 2,
+            actual_len: 3,
+        }
+    );
+}
+
+#[test]
+fn rejects_invalid_output_dimensions() {
+    let mut c = [0.0; 3];
+
+    let error = sgemm_aat_lower(&[0.0; 4], 2, 2, &mut c).unwrap_err();
+
+    assert_eq!(
+        error,
+        SgemmError::InvalidMatrixDimensions {
+            matrix_name: MatrixName::C,
+            expected_rows: 2,
+            expected_cols: 2,
+            actual_len: 3,
+        }
+    );
+}
+
+#[test]
+fn rejects_input_size_overflow() {
+    let error = sgemm_aat_lower(&[], usize::MAX, 2, &mut []).unwrap_err();
+
+    assert_eq!(
+        error,
+        SgemmError::DimensionOverflow {
+            matrix_name: MatrixName::A,
+            rows: usize::MAX,
+            cols: 2,
+        }
+    );
+}
+
+#[test]
+fn rejects_output_size_overflow() {
+    let error = sgemm_aat_lower(&[], usize::MAX, 0, &mut []).unwrap_err();
+
+    assert_eq!(
+        error,
+        SgemmError::DimensionOverflow {
+            matrix_name: MatrixName::C,
+            rows: usize::MAX,
+            cols: usize::MAX,
+        }
+    );
+}

diskann-pipnn/Cargo.toml

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT license.
+
+[package]
+name = "diskann-pipnn"
+version.workspace = true
+description = "PiPNN graph construction for DiskANN"
+authors.workspace = true
+repository.workspace = true
+license.workspace = true
+edition.workspace = true
+
+[dependencies]
+diskann-vector.workspace = true
+diskann-wide.workspace = true
+thiserror.workspace = true
+
+[dev-dependencies]
+criterion.workspace = true
+diskann-linalg.workspace = true
+
+[[bench]]
+name = "kernels"
+harness = false
+
+[lints]
+workspace = true