PiPNN 2/6: add numerical kernels

.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

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

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

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<Self, TryFromError<T>> {
         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]

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

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

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

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