PiPNN 1/6: add numerical kernels - #1287
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds the first set of PiPNN “kernel” building blocks to the DiskANN Rust workspace: SIMD-accelerated top‑k selection for partition assignment and leaf neighbor selection, along with supporting SIMD division and a new lower-triangular A·Aᵀ helper in diskann-linalg.
Changes:
- Add a new
diskann-pipnncrate withpartition_kernelandleaf_kernelimplementations plus extensive correctness tests and Criterion benchmarks. - Extend
diskann-wideto supportDivon relevant f32 SIMD types (native, doubled, and scalar/emulated) and add a corresponding division test macro. - Add
diskann_linalg::sgemm_aat_lower(lower-triangle-only AAT) and wire new crate/tests/CI/mutants exclusions into the workspace.
Reviewed changes
Copilot reviewed 26 out of 27 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| diskann-wide/src/test_utils/ops.rs | Adds test_div! macro to validate lane-wise SIMD division correctness. |
| diskann-wide/src/emulated.rs | Adds Div for scalar/emulated Emulated<f32, N, A> to support division in scalar dispatch. |
| diskann-wide/src/doubled.rs | Adds Div for Doubled<T> to support composite SIMD widths. |
| diskann-wide/src/arch/x86_64/v4/f32x8_.rs | Adds AVX Div op mapping + division tests. |
| diskann-wide/src/arch/x86_64/v4/f32x4_.rs | Adds SSE Div op mapping + division tests. |
| diskann-wide/src/arch/x86_64/v4/f32x16_.rs | Adds AVX-512 Div op mapping + division tests. |
| diskann-wide/src/arch/x86_64/v3/f32x8_.rs | Adds AVX Div op mapping + division tests for V3. |
| diskann-wide/src/arch/x86_64/v3/f32x4_.rs | Adds SSE Div op mapping + division tests for V3. |
| diskann-wide/src/arch/x86_64/v3/f32x16_.rs | Adds division tests for the f32x16 V3 path (likely via doubled composition). |
| diskann-wide/src/arch/aarch64/f32x4_.rs | Adds Neon Div op mapping + division tests. |
| diskann-wide/src/arch/aarch64/f32x2_.rs | Adds Neon Div op mapping + division tests. |
| diskann-pipnn/tests/partition_kernel.rs | New integration tests for partition top‑k dispatch correctness and edge cases. |
| diskann-pipnn/tests/leaf_kernel.rs | New integration tests for leaf neighbor top‑k dispatch correctness and edge cases. |
| diskann-pipnn/src/partition_kernel/tests.rs | New unit tests comparing scalar reference vs runtime dispatch and metric contracts. |
| diskann-pipnn/src/partition_kernel.rs | New partition-assignment distance + top‑k kernel with validation and SIMD dispatch. |
| diskann-pipnn/src/lib.rs | New crate root exporting PiPNN kernel modules. |
| diskann-pipnn/src/leaf_kernel/tests.rs | New unit tests for scalar reference parity and workspace behavior. |
| diskann-pipnn/src/leaf_kernel.rs | New fused lower-triangle leaf neighbor kernel with SIMD dispatch and workspace support. |
| diskann-pipnn/Cargo.toml | Defines new diskann-pipnn crate, dev-deps, and benches. |
| diskann-pipnn/benches/kernels.rs | Adds benchmarks for partition top‑k, lower AAT, leaf top‑k, and full leaf workflow. |
| diskann-linalg/tests/sgemm_aat_lower.rs | New tests for lower-triangle AAT behavior and validation errors. |
| diskann-linalg/src/lib.rs | Adds public sgemm_aat_lower API with dimension checks. |
| diskann-linalg/src/faer.rs | Implements sgemm_aat_lower_impl using Faer triangular matmul. |
| Cargo.toml | Adds diskann-pipnn to workspace members and workspace dependencies. |
| Cargo.lock | Records the new diskann-pipnn package entry. |
| .github/workflows/ci.yml | Adds diskann-pipnn to CI test package lists. |
| .cargo/mutants.toml | Adds mutation-test exclusions for kernel code paths and equivalent transformations. |
Comments suppressed due to low confidence (2)
diskann-pipnn/src/leaf_kernel.rs:651
- Same issue as the L2 arm: using
max_simdfor lower clamping can erase NaNs on the Scalar/Emulated backend, making NaN distances rankable. Clamp withlt_simd+selectto preserve NaNs consistently.
Metric::CosineNormalized => {
let distance = F::splat(arch, 1.0) - dot;
zero.max_simd(distance)
}
diskann-pipnn/src/leaf_kernel.rs:664
- The cosine path also uses
zero.max_simd(distance)for clamping, which can collapse NaNs to zero on the Scalar/Emulated backend (viaf32::max). That contradicts the comment about preserving non-rankable NaNs and can change output ordering. Prefer anlt_simd+selectclamp here as well.
let distance = one - cosine;
// Comparisons with NaN are false, so this explicit lower clamp
// preserves non-rankable NaNs while matching the existing PiPNN
// distance formulas for finite values.
zero.max_simd(distance)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| Metric::L2 => { | ||
| let distance = row_norm + column_norm - F::splat(arch, 2.0) * dot; | ||
| zero.max_simd(distance) | ||
| } |
| let leader_scales = match metric { | ||
| Metric::L2 => (0..leaders).map(|leader| (leader + 1) as f32).collect(), | ||
| Metric::Cosine => (0..leaders) |
|
@SeliMeli please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
Adds the numerical kernels required by later PiPNN graph-construction layers.
Review key points
diskann-widedispatch; no ISA detection or intrinsics in PiPNN.Stack 1/6 → #1288