PiPNN 1/6: add numerical kernels

[SeliMeli]

Adds the numerical kernels used by later PiPNN layers.

Code map

1. diskann-linalg::sgemm_aat_lower computes A · Aᵀ but writes only the lower triangle. Callers may leave the upper triangle uninitialized; tests assert that it is not touched.
2. partition_kernel.rs converts a row of point/leader dot products into the nearest fanout leader IDs. Metric-specific norm handling happens before the fixed-size top-k insertion.
3. leaf_kernel.rs scans each strict-lower-triangle pair once and updates both endpoint top-k trackers. k <= 3 uses const-sized insertion arms; larger k uses the dynamic fallback.
4. Public kernel functions validate shapes and capacities before dispatch. Architecture selection and vector implementations remain in diskann-wide; PiPNN targets Architecture associated vector types only.

Review path

• Start with the scalar references and shape validation in each kernel.
• Then compare the scalar tail and SIMD chunk paths: both must preserve tie order, NaN rejection, zero-norm cosine behavior, and finite f32::MAX candidates.
• For the leaf kernel, verify that lower-triangle traversal updates both rows exactly once and never reads the untouched upper triangle.
• The AArch64 NaN regression is guarded by cosine_zero_norm_masks_nan_norm_at_simd_boundaries; SIMD max has backend-specific NaN behavior, so the clamp explicitly selects the original NaN.

Validation includes differential boundary tests, x86-64 baseline/AArch64 builds, SDE jobs, Criterion workloads, and full-leaf numerical tests.

Stack 1/6 → #1288

[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-commenter]

Codecov Report

❌ Patch coverage is 97.94643% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.51%. Comparing base (59dd048) to head (3bc430d).
⚠️ Report is 7 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-pipnn/src/partition_kernel.rs	96.48%	14 Missing ⚠️
diskann-pipnn/src/leaf_kernel.rs	98.41%	8 Missing ⚠️
diskann-wide/src/test_utils/ops.rs	94.44%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1287      +/-   ##
==========================================
+ Coverage   90.59%   92.51%   +1.91%     
==========================================
  Files         513      519       +6     
  Lines       99091    99281     +190     
==========================================
+ Hits        89775    91849    +2074     
+ Misses       9316     7432    -1884     

Flag	Coverage Δ
miri	92.51% <97.94%> (+1.91%)	⬆️
unittests	92.48% <97.94%> (+2.20%)	⬆️

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.68% <100.00%> (+1.18%)	⬆️
diskann-pipnn/src/kernel_metric.rs	100.00% <100.00%> (ø)
diskann-wide/src/arch/x86_64/v3/f32x16_.rs	100.00% <ø> (ø)
diskann-wide/src/arch/x86_64/v3/f32x4_.rs	100.00% <ø> (ø)
diskann-wide/src/arch/x86_64/v3/f32x8_.rs	100.00% <ø> (ø)
diskann-wide/src/arch/x86_64/v4/f32x16_.rs	100.00% <ø> (+85.88%)	⬆️
diskann-wide/src/arch/x86_64/v4/f32x4_.rs	100.00% <ø> (+83.09%)	⬆️
diskann-wide/src/arch/x86_64/v4/f32x8_.rs	100.00% <ø> (+83.09%)	⬆️
diskann-wide/src/doubled.rs	86.89% <100.00%> (+0.17%)	⬆️
... and 5 more

... and 287 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

[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.

[partychen]

checked_area, check_length, ShapeOverflow and InvalidBufferLength are duplicated character-for-character with leaf_kernel. Small enough to shrug at now, but with four more PRs coming it's probably worth a src/shape.rs with a shared ShapeError that each kernel error wraps via #[from].

[SeliMeli]

I kept the two tiny checked-area/length adapters local because they construct different public kernel error types and sit immediately before each module's unsafe accesses. MatrixView adoption removed the other duplicated shape state; introducing a shared wrapped error would enlarge the public error interface for two call sites.

[SeliMeli]

Review follow-up at PR1 head 3bc430d2:

Dispatch and API

• Replaced per-call architecture/metric dispatch with prepared LeafKernel and PartitionKernel handles backed by diskann-wide::DispatchedN function pointers.
• Added BYO type erasure through concrete KernelMetric visitors; metric arithmetic is monomorphized before final function-pointer erasure.
• Removed KValue and all duplicate requested/effective-k state. LeafKernel::new(metric) prepares only metric/architecture; each call takes its width from output.ncols(). leaf_neighbor_count(points, requested_k) is the caller-side shape helper.
• Replaced separate fixed/dynamic leaf traversals with one process_pairs traversal over NeighborStorage; widths 1/2/3 use fixed arrays and larger widths use safe dynamic slices. Output insertion no longer uses raw pointers.
• Replaced raw shape fields with MatrixView / MutMatrixView; added typed PartitionScales and retained local checked-area/backing-length checks before unsafe SIMD access.
• Renamed algorithm roles consistently: leaf code uses source/target, partition code uses point/leader, and row terminology is limited to matrix representation.

Documentation

• Reworked crate documentation using the structure of diskann-quantization: it now defines PiPNN and its graph-building goal from first principles, explains Randomized Ball Carving, leaders/fanout/leaves, leaf-local picking, merge/prune, module placement, ownership boundaries, and normal call order.
• Expanded partition_kernel top-level docs with module boundaries, main structures, score formulas/scale units, nearest_leaders/process_points flow, tracker invariants, complexity, allocation behavior, numerical behavior, and a runnable example.
• Expanded leaf_kernel top-level docs with leaf creation context, Gram-matrix input, main structures, nearest_neighbors/process_pairs flow, fixed/dynamic storage, endpoint updates, complexity, allocation behavior, numerical behavior, and a runnable example.
• Expanded kernel_metric docs with shared formulas, concrete marker dispatch, scale representations, exact cosine zero threshold, NaN precedence, and per-formula implementation comments.
• Preserved local algorithm/safety comments; corrected two inaccurate statements about transient source_worst caching and cosine zero-vs-NaN precedence.

Tests and repository fit

• Moved private tests into bottom-of-file #[cfg(test)] modules. tests/*_api.rs now contains only public-interface integration tests with independent formulas and sorting.
• Added coverage for one prepared leaf handle serving multiple widths, SIMD boundaries/tails, all metrics, zero/subnormal/NaN/infinity/signed-zero behavior, shape overflow, stable ties, and scalar-versus-dispatch traversal.
• Removed this crate's Criterion dependency and benchmark target.
• Targeted strict-provenance Miri checks pass for the prepared variable-width leaf path and partition cosine scale path.

Validation:

RUSTDOCFLAGS='-D warnings' cargo test -p diskann-pipnn --doc  pass (2 doctests)
cargo test -p diskann-pipnn --all-targets                         pass
cargo clippy -p diskann-pipnn --all-targets -- -D warnings        pass
cargo check --target aarch64-unknown-linux-gnu \
  -p diskann-pipnn --all-targets                                  pass
MIRIFLAGS=-Zmiri-strict-provenance targeted leaf/partition Miri   pass

Pinned fixed-iteration leaf measurements still show no material kernel regression after converting fixed outputs to array rows once per leaf: k=3 462.0 ms -> 459.9 ms; k=2 remained within run noise. Callgrind instruction count changed by +1.37%.

[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/kernel_metric.rs:152

• L2::partition_distance uses mul_add_simd, which can be a fused multiply-add (e.g. _mm256_fmadd_ps). That means SIMD-chunk distances can differ from the scalar tail formula leader_scale - 2.0 * dot (this file’s own l2_partition_scalar_tail_preserves_non_fused_rounding test demonstrates such a mismatch). Because PartitionKernel mixes SIMD chunks and a scalar tail within the same row, this can change ordering/tie behavior depending on whether a leader lands in the SIMD chunk or tail.

        F::splat(arch, -2.0).mul_add_simd(dot, leader_scale)

[Copilot]

Pull request overview

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

Suppressed comments (1)

.github/workflows/nightly.yml:23

• DISKANN_FEATURES is defined as a folded multi-line string with trailing commas. YAML folding inserts spaces at line breaks, producing feature tokens like "tracing, experimental_diversity_search" (note the space) which can be mis-parsed as invalid feature names. Prefer a whitespace-separated feature list (or a single-line comma-separated list without spaces) to avoid CI flakiness.

  DISKANN_FEATURES: >-
    virtual_storage,spherical-quantization,product-quantization,tracing,
    experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,
    multi-vector,bftree,inmem2,integration-test

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

Suppressed comments (1)

.github/workflows/nightly.yml:23

• DISKANN_FEATURES is defined as a folded scalar with commas at line ends. YAML folding inserts spaces at line breaks, producing a value like tracing, experimental_diversity_search,... which can be parsed as having empty/whitespace-prefixed feature names depending on Cargo’s splitting rules. This is brittle and can break the cargo ... --features "${{ env.DISKANN_FEATURES }}" steps.

  DISKANN_FEATURES: >-
    virtual_storage,spherical-quantization,product-quantization,tracing,
    experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,
    multi-vector,bftree,inmem2,integration-test

[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/leaf_kernel.rs:467

• The comment claims no output or scratch mutation occurs on error, but after validate(...) the call to prepare_workspace(...) can return LeafKernelError::Allocation after partially resizing/filling workspace.norms (before workspace.worst is reserved). This makes the comment/documentation inaccurate and could mislead callers relying on workspace immutability on error.

        // Validation establishes every shape and active-prefix invariant used by
        // unchecked loads below. No output or scratch mutation occurs on error.
        validate(call.input, &call.output)?;