PiPNN 1/6: add numerical kernels

.github/workflows/ci.yml

@@ -22,7 +22,7 @@ env:
   CARGO_TERM_COLOR: always
   # The features we want to explicitly test. For example, the `flatbuffers-build` feature
   # of `diskann-quantization` requires additional setup and so must not be included by default.
-  DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test"
+  DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test,pipnn,testing"
 
   # Intel SDE version used for baseline and AVX-512 emulation jobs.
   SDE_VERSION: "sde-external-10.8.0-2026-03-15-lin"
@@ -355,6 +355,8 @@ jobs:
               --package diskann-wide \
               --package diskann-vector \
               --package diskann-quantization \
+              --package diskann \
+              --features diskann/pipnn \
               -- --skip compile_tests \
                  --skip pivots::tests::run_test_happy_path
 
@@ -415,6 +417,8 @@ jobs:
             --package diskann-wide \
             --package diskann-vector \
             --package diskann-quantization \
+            --package diskann \
+            --features diskann/pipnn \
             -- --skip compile_tests
 
   test-workspace:
@@ -514,6 +518,7 @@ jobs:
           cargo llvm-cov nextest --locked \
             --config "$RUST_CONFIG" \
             --workspace \
+            --features "${{ env.DISKANN_FEATURES }}" \
             --lcov --output-path lcov.info
 
       - name: Generate miri code coverage

.github/workflows/nightly.yml

@@ -17,7 +17,10 @@ env:
   RUST_CONFIG: 'build.rustflags=["-Dwarnings"]'
   RUST_BACKTRACE: 1
   CARGO_TERM_COLOR: always
-  DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test"
+  DISKANN_FEATURES: >-
+    virtual_storage,spherical-quantization,product-quantization,tracing,
+    experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,
+    multi-vector,bftree,inmem2,integration-test,pipnn,testing
 
 defaults:
   run:
@@ -36,7 +39,9 @@ jobs:
         run: rustup show && rustup component add clippy
       - uses: Swatinem/rust-cache@v2
       - name: "clippy --workspace --all-targets"
-        run: cargo clippy --locked --workspace --all-targets --no-deps --config "$RUST_CONFIG" -- -Dwarnings
+        run: |
+          cargo clippy --locked --workspace --all-targets --no-deps \
+            --config "$RUST_CONFIG" -- -Dwarnings
 
   clippy-features:
     name: clippy-features (macos)
@@ -52,7 +57,7 @@ jobs:
           cargo clippy --locked --workspace \
             --all-targets \
             --no-deps \
-            --features ${{ env.DISKANN_FEATURES }} \
+            --features "${{ env.DISKANN_FEATURES }}" \
             --config "$RUST_CONFIG" \
             -- -Dwarnings
 
@@ -109,7 +114,7 @@ jobs:
           set -euxo pipefail
           cargo nextest run --locked --workspace \
             --config "$RUST_CONFIG" \
-            --features ${{ env.DISKANN_FEATURES }}
+            --features "${{ env.DISKANN_FEATURES }}"
 
           cargo test --locked --doc --workspace --config "$RUST_CONFIG"
 
@@ -133,6 +138,10 @@ jobs:
       - uses: Swatinem/rust-cache@v2
 
       - name: miri
-        run: cargo +nightly miri nextest run --locked --package diskann-quantization
+        run: |
+          cargo +nightly miri nextest run --locked \
+            --package diskann-quantization
+          cargo +nightly miri test --locked \
+            --package diskann --features pipnn --lib
         env:
           MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance

diskann-linalg/src/faer.rs

@@ -53,6 +53,34 @@ pub(super) fn sgemm_impl(
     faer::linalg::matmul::matmul(c, beta, a, b, alpha, Par::Seq)
 }
 
+/// Implements the public lower-triangular AAT operation.
+///
+/// Leaf selection consumes each symmetric pair once and updates both endpoints,
+/// so computing or initializing the upper triangle would be wasted bandwidth.
+/// Faer's triangular block structure is the contract that prevents those stores;
+/// callers may keep unrelated values in the upper triangle. The public wrapper
+/// has already checked `a.len() == m * k`, `c.len() == m * m`, and overflow, so
+/// the unchecked matrix views below cannot escape their backing slices.
+pub(super) fn sgemm_aat_lower_impl(m: usize, k: usize, a: &[f32], c: &mut [f32]) {
+    use faer::linalg::matmul::triangular::{matmul, BlockStructure};
+
+    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

@@ -80,6 +80,30 @@ impl fmt::Display for SgemmError {
 
 impl std::error::Error for SgemmError {}
 
+fn check_matrix(
+    matrix_name: MatrixName,
+    actual_len: usize,
+    rows: usize,
+    cols: usize,
+) -> Result<(), SgemmError> {
+    let expected_len = rows
+        .checked_mul(cols)
+        .ok_or(SgemmError::DimensionOverflow {
+            matrix_name,
+            rows,
+            cols,
+        })?;
+    if actual_len != expected_len {
+        return Err(SgemmError::InvalidMatrixDimensions {
+            matrix_name,
+            expected_rows: rows,
+            expected_cols: cols,
+            actual_len,
+        });
+    }
+    Ok(())
+}
+
 // Make the reference implementation available for internal testing.
 #[cfg(test)]
 mod reference;
@@ -154,57 +178,33 @@ pub fn sgemm(
     beta: Option<f32>,
     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(())
 }
 
+/// 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(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`.

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(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,
+        }
+    );
+}

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());

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());

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());