PiPNN 2/6: add numerical kernels

[weiyaoluo (SeliMeli)]

Purpose

This PR adds the numerical kernels that PiPNN uses for partition ranking and leaf-neighbor selection.

It does not add graph construction, providers, serialization, or search.

Module structure

• leaf_metric.rs writes one flattened lower-triangle ranking buffer for a leaf.
• partition_metric.rs writes one flattened point-to-leader ranking buffer for a point stripe.
• leaf_kernel.rs scans each unordered point pair once and retains local top-k neighbors.
• partition_kernel.rs retains the nearest sampled leaders for each point.
• simd.rs selects one dispatched f32 vector type for both ranking kernels.
• diskann-linalg provides lower-triangle replace and add forms of scaled A · Aᵀ.

The metric modules own GEMM configuration, portable norm-prefill loops, and metric formulas. They do not receive an architecture value. The kernels own architecture dispatch, buffer reuse, SIMD/scalar traversal, and top-k insertion.

The leaf entry point validates output width before it changes scratch or evaluates distances. The internal ranking loop returns (). Insertion returns the farthest retained ranking distance.

Metric data flow

Leaf

• L2 initializes the lower triangle with endpoint squared-norm sums. GEMM adds -2 · A · Aᵀ.
• Cosine uses GEMM for dot products. It derives norms from the diagonal and converts the lower triangle in place.
• Normalized cosine and inner product share the GEMM calculation with alpha = -1. They use -dot as the ranking distance because the omitted constant cannot change order. Their metric types remain separate.

Partition

PartitionMetric::create_leaders returns an opaque metric-owned leader set. The caller does not manage its norms.

• L2 computes sequentially reduced leader squared norms during construction.
• Cosine computes leader norms during construction.
• Normalized cosine and inner product share the negative-dot calculation and need no leader norms.
• L2 omits the point norm because it is constant across one leader-ranking row.

Point stripes share immutable leader norms. Construction computes them before parallel stripe work starts; there is no lazy initialization state.

Numerical behavior

• Equal distances can select either candidate. No fixed tie order is required.
• NaN and positive infinity do not enter retained top-k sets.
• A finite negative Gram-expanded L2 value remains a valid internal ranking distance.
• Cosine preserves DiskANN zero-norm handling and clamps finite similarity to [-1, 1].
• GEMM can return either zero sign. Tests check numeric zero, not the bit pattern of a scalar reduction. This PR adds no zero-sign normalization.

Tests

Ranking tests pass distance buffers directly to the ranking kernels. Metric tests call the production metric implementations. The shared cosine helper tests cover clamp, zero-norm, subnormal, and NaN behavior.

The private tests cover:

• all four metrics;
• lower-triangle replace and add behavior;
• fixed top-k widths and the runtime-width path;
• SIMD lane boundaries and scalar tails;
• valid tied candidates, NaN, infinity, numeric zero, and negative ranking distances;
• singleton leaves and invalid neighbor widths;
• workspace reuse and unchanged buffers after width rejection;
• sequential L2 leader-norm reduction.

Local validation for the audit cleanup and leaf-kernel follow-up:

• PR2 cargo test --locked -p diskann --all-features --lib: 463 passed, including 109 PiPNN tests.
• AArch64 Linux musl under QEMU: all 109 PiPNN tests passed. The original scalar-bit-pattern assertions failed in this environment before the fix.
• cargo test --locked -p diskann-linalg --lib sgemm_aat_lower: 6 passed.
• Complete stacked PR6 cargo test --locked -p diskann --all-features --lib: 613 passed in an isolated worktree with the PR2 patch.
• PR2 and PR6 cargo clippy --locked -p diskann --all-features --all-targets -- -D warnings: passed.
• cargo fmt --all --check and git diff --check: passed.

Prefill microbenchmarks

Measured on an AMD EPYC 7763 with portable x86-64 compilation. LLVM vectorized the scalar loops with baseline SSE2.

For Leaf L2, portable scalar prefill was 45–77% faster than the prior explicit SIMD loop across 16–512 points. For Partition L2, results depended on leader count: most tested shapes improved, while 128/256/512/1000 leaders regressed by 1.5–7.9% in the isolated prefill loop.

Earlier E2E A/B

These E2E measurements precede the portable-prefill replacement and the subsequent audit cleanup. The prefill change used the microbenchmarks above. The audit cleanup used unit tests and Clippy; it has no new performance measurements.

Measured locally on 16 logical CPUs of an AMD EPYC 7763. Both binaries used -C target-cpu=native and identical HashPrune configurations.

Dataset	Runs	Build median	Wall median	Peak RSS median	Recall
BigANN 10M baseline	5	124.383 s	141.06 s	11.583 GiB	0.956730
BigANN 10M candidate	5	126.690 s (+1.85%)	140.21 s (-0.60%)	11.581 GiB (-0.01%)	0.956730
Enron 1M baseline	3	22.324 s	28.66 s	2.268 GiB	0.972889
Enron 1M candidate	3	20.993 s (-5.96%)	26.56 s (-7.33%)	2.274 GiB (+0.27%)	0.972911

BigANN paired build-time deltas had a +2.35% median. Every BigANN run had identical recall and mean hops. Enron paired build-time deltas had a -6.11% median.

Review order

1. Review diskann-linalg/src/faer.rs and the lower-triangle API in diskann-linalg/src/lib.rs.
2. Review diskann/src/graph/pipnn/leaf_metric.rs.
3. Review diskann/src/graph/pipnn/leaf_kernel.rs.
4. Review diskann/src/graph/pipnn/partition_metric.rs.
5. Review diskann/src/graph/pipnn/partition_kernel.rs.
6. Review diskann/src/graph/pipnn/simd.rs and the nightly Miri scope.

Stack

Stack 2/6. #1315 is merged into main. #1290 adds provider-independent graph construction.

[Copilot]

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-pipnn crate with partition_kernel and leaf_kernel implementations plus extensive correctness tests and Criterion benchmarks.
• Extend diskann-wide to support Div on 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_simd for lower clamping can erase NaNs on the Scalar/Emulated backend, making NaN distances rankable. Clamp with lt_simd + select to 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 (via f32::max). That contradicts the comment about preserving non-rankable NaNs and can change output ordering. Prefer an lt_simd + select clamp 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.

[Copilot]

Pull request overview

Copilot reviewed 26 out of 27 changed files in this pull request and generated no new comments.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 97.91833% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.53%. Comparing base (600c2b9) to head (acda584).
⚠️ Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann/src/graph/pipnn/leaf_kernel.rs	93.33%	11 Missing ⚠️
diskann/src/graph/pipnn/partition_kernel.rs	94.49%	6 Missing ⚠️
diskann/src/graph/pipnn/partition_metric.rs	96.95%	6 Missing ⚠️
diskann/src/graph/pipnn/topk.rs	99.40%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1287      +/-   ##
==========================================
- Coverage   91.55%   91.53%   -0.02%     
==========================================
  Files         521      529       +8     
  Lines      100302   102313    +2011     
==========================================
+ Hits        91828    93657    +1829     
- Misses       8474     8656     +182     

Flag	Coverage Δ
miri	91.53% <97.91%> (-0.02%)	⬇️
unittests	91.34% <97.91%> (+0.10%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
diskann-linalg/src/faer.rs	100.00% <100.00%> (ø)
diskann-linalg/src/lib.rs	99.75% <100.00%> (+0.06%)	⬆️
diskann-utils/src/views.rs	100.00% <100.00%> (ø)
diskann/src/graph/pipnn/leaf_metric.rs	100.00% <100.00%> (ø)
diskann/src/graph/pipnn/mod.rs	100.00% <100.00%> (ø)
diskann/src/graph/pipnn/simd.rs	100.00% <100.00%> (ø)
diskann/src/graph/pipnn/topk.rs	99.40% <99.40%> (ø)
diskann/src/graph/pipnn/partition_kernel.rs	94.49% <94.49%> (ø)
diskann/src/graph/pipnn/partition_metric.rs	96.95% <96.95%> (ø)
diskann/src/graph/pipnn/leaf_kernel.rs	93.33% <93.33%> (ø)

... and 58 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Copilot reviewed 26 out of 27 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

diskann-pipnn/src/partition_kernel/tests.rs:20

• The PartitionTopK contract for Metric::L2 expects leader_scales to contain squared leader norms (see docs and distance(Metric::L2, ..) test). This helper currently populates unsquared norms, which makes the test data inconsistent with the public API contract and could hide contract-related bugs.

    let leader_scales = match metric {
        Metric::L2 => (0..leaders).map(|leader| (leader + 1) as f32).collect(),
        Metric::Cosine => (0..leaders)
            .map(|leader| {

diskann-pipnn/src/partition_kernel.rs:61

• InvalidFanout’s error message says the maximum is {maximum}, but validation also rejects fanout > leaders. When leaders < maximum this message is misleading (it implies the only limit is {maximum}). Consider spelling out both constraints in the message so callers immediately see why it failed.

    #[error("invalid fanout {fanout} for {leaders} leaders; maximum is {maximum}")]

[Copilot]

Pull request overview

Copilot reviewed 25 out of 26 changed files in this pull request and generated no new comments.

Suppressed comments (1)

diskann-pipnn/src/partition_kernel.rs:294

• For Metric::Cosine, NaN norms currently produce a finite distance (1.0) because denominator.gt_simd(0) is false for NaN, so the lane falls back to cosine = 0. That makes NaN-derived pairs/leaders “rankable”, which contradicts the module’s stated NaN-rejection behavior and differs from diskann-vector cosine semantics (NaN norms propagate to a NaN similarity/distance). Consider explicitly preserving NaN denominators so the resulting distance stays NaN and is ignored by insert_topk.

        let denominator = row_norm * leader_norm;
        let valid = denominator.gt_simd(zero);
        let safe_denominator = valid.select(denominator, one);
        let cosine = valid.select(dot / safe_denominator, zero);
        one - cosine

[juchen-ms (partychen)]

Nice work overall. I found one correctness issue in the cosine handling that should be resolved before merge. The remaining comments are mostly about reducing duplicated or unsafe code and tightening the API contracts.

[Aditya Krishnan (arkrishn94)]

Thanks Weiyao, this is progress from the previous mega-PR. I still have some big-picture comments (we covered most of these offline) -

• Documentation: As I mentioned, we need thorough documentation in the diskann-pipnn crate. The main modules, partition_kernel and leaf_kernel need documentation up top, highlighting the main structures and how they are used - e.g. process_rows_binary/unary and nearest_leaders. Similarly with process_pairs_simd_* and nearest_leaf_neighbors
• Testing: I am concerned about the lack of testing for partition_kernel.rs and leaf_kernel.rs.
  • I notice some e2e integration tests but these kernels should be thoroughly tested, sweeping different input parameters, architectures and edge cases. This is especially needed given the amount of unsafe code.
  • That brings me to miri - there should be miri tests too.
  • I'm curious why are the tests in a separate submodule to the main files (for partition_kernel.rs and leaf_kernel.rs)? Let's try to keep tests along with the code being tested.
• Criterion: Since criterion is not a standard part of our library for benchmarking, let us not introduce it for this crate.
• Kernel dispatch: I left comments about you're disptaching the kernels, please take a look.

[Copilot]

Pull request overview

Copilot reviewed 26 out of 27 changed files in this pull request and generated no new comments.

Suppressed comments (2)

diskann-pipnn/src/partition_kernel/tests.rs:19

• PartitionTopK::leader_scales is documented as "squared leader norms for L2" (and cosine uses unsquared norms), but this test helper feeds unsquared values for the L2 case. That makes the test data inconsistent with the public contract and can mask mistakes in distance computation. Consider squaring the L2 norms here so the tests exercise the intended inputs.

    let leader_scales = match metric {
        Metric::L2 => (0..leaders).map(|leader| (leader + 1) as f32).collect(),
        Metric::Cosine => (0..leaders)

diskann-pipnn/src/partition_kernel.rs:252

• For the L2 path, the SIMD chunk uses mul_add_simd (fused multiply-add) but the scalar tail uses norm - 2.0 * dot (non-fused). This can introduce small rounding differences between SIMD and tail elements, which can change ordering/tie behavior right at SIMD-width boundaries. Use f32::mul_add for the scalar tail so both paths compute the same value shape.

                |dot, norm| F::splat(arch, -2.0).mul_add_simd(dot, norm),
                |dot, norm| norm - 2.0 * dot,

[weiyaoluo (SeliMeli)]

weiyaoluo (@SeliMeli) please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

• (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.

@microsoft-github-policy-service agree

• (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.

@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree company="Microsoft"

[Wei Wu (wuw92)]

Overall, the implementation looks reasonable, I have a few maintainability concerns:

1. Please document non-obvious optimizations and their benchmark evidence near the code. Otherwise, the reasoning is only available in review discussions and commit history.
2. There are many tests with overlapping functional coverage. Could we consolidate cases that verify the same invariant, while keeping separate tests for distinct code paths and boundary conditions?
3. PiPNN does not appear in the repository's Codecov results. Could we confirm that these tests run in the coverage job?