PiPNN 2/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.
weiyaoluo (SeliMeli)
changed the title
e204cb9 to
b046174
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1287 +/- ##
==========================================
- Coverage 91.55% 91.50% -0.05%
==========================================
Files 521 528 +7
Lines 100302 101756 +1454
==========================================
+ Hits 91828 93113 +1285
- Misses 8474 8643 +169
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
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
PartitionTopKcontract forMetric::L2expectsleader_scalesto contain squared leader norms (see docs anddistance(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 rejectsfanout > leaders. Whenleaders < maximumthis 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}")]
8fb4e92 to
20ab8a0
Compare
There was a problem hiding this comment.
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) becausedenominator.gt_simd(0)is false for NaN, so the lane falls back tocosine = 0. That makes NaN-derived pairs/leaders “rankable”, which contradicts the module’s stated NaN-rejection behavior and differs fromdiskann-vectorcosine semantics (NaN norms propagate to a NaN similarity/distance). Consider explicitly preserving NaN denominators so the resulting distance stays NaN and is ignored byinsert_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
Aditya Krishnan (arkrishn94)
left a comment
Aditya Krishnan (arkrishn94)
left a comment
There was a problem hiding this comment.
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-pipnncrate. The main modules,partition_kernelandleaf_kernelneed documentation up top, highlighting the main structures and how they are used - e.g.process_rows_binary/unaryandnearest_leaders. Similarly withprocess_pairs_simd_*andnearest_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.
There was a problem hiding this comment.
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_scalesis 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 usesnorm - 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. Usef32::mul_addfor 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,
@microsoft-github-policy-service agree company="Microsoft" |
Use named cases and derived expected values so each failure identifies a ranking behavior that can change graph output.
Derive Gram, dot, and norm inputs from source vectors and name each test by the ranking behavior it protects.
Add direct SIMD metric coverage and make the caller-owned insertion eligibility precondition explicit in names, docs, and tests.
Keep runtime-dispatch adapters only for SIMD checks and pass direct production calls through setup closures for ordinary behavior tests.
Keep NaN scores non-rankable and make SIMD InnerProduct match scalar unary negation for signed zero.
Cover exceptional metric values, exact norm boundaries, all-ineligible SIMD groups, and zero-width partition output.
Metric modules own GEMM setup and leader caches. Kernels consume flat ranking distances without metric-specific norm plumbing.
Portable scalar prefill loops auto-vectorize without exposing SIMD types through metric interfaces.
Test ranking directly from distances and keep metric checks at their production seam. Prepare leader norms before stripe workers and retain fixed-width neighbor insertion without a generic comparator. Equal distances need no fixed candidate order. Accept either GEMM zero sign instead of requiring scalar bit patterns across architectures.
Reject invalid neighbor widths before distance evaluation or scratch mutation. Keep ranking infallible and return the farthest distance directly from insertion.
Wei Wu (wuw92)
left a comment
Wei Wu (wuw92)
left a comment
There was a problem hiding this comment.
Overall, the implementation looks reasonable, I have a few maintainability concerns:
- Please document non-obvious optimizations and their benchmark evidence near the code. Otherwise, the reasoning is only available in review discussions and commit history.
- 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?
- PiPNN does not appear in the repository's Codecov results. Could we confirm that these tests run in the coverage job?
| impl LeafMetric for InnerProduct { | ||
| fn compute_distances(points: MatrixView<'_, f32>, storage: &mut [f32]) -> ANNResult<()> { | ||
| // Both metrics rank with `-dot`. Their graph-pruning policies stay separate. | ||
| CosineNormalized::compute_distances(points, storage) |
There was a problem hiding this comment.
Nit: Would it be clearer for CosineNormalized to delegate to InnerProduct instead? For normalized vectors, cosine similarity is just the inner product.
| /// triangle stays unspecified. The input matrix has one point in each row. | ||
| /// A zero distance can have either sign. Equal distances can select either candidate. | ||
| pub(super) trait LeafMetric: Send + Sync + 'static { | ||
| /// Compute ranking distances for all unordered point pairs. |
There was a problem hiding this comment.
Could this method document that it returns ranking values, not standard metric distances? L2 omits the square root, and normalized cosine returns -dot instead of 1 - dot. Callers should not interpret the buffer as metric distances.
There was a problem hiding this comment.
Fixed in 2c26d8b. The method now documents the ranking formulas: squared L2, -dot for normalized cosine and inner product, and 1 - similarity for cosine. The partition method also documents the omitted point norm.
| let point_count = points.nrows(); | ||
| // The constant in `1 - dot` does not change nearest-first order. | ||
| diskann_linalg::sgemm_aat_lower( | ||
| point_count, |
There was a problem hiding this comment.
nit: use points.nrows() and remove point_count
| } | ||
|
|
||
| #[cfg(test)] | ||
| #[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] |
There was a problem hiding this comment.
This allow appears redundant.
| fn insert_eligible(&mut self, candidate: LeafNeighbor) -> f32; | ||
| } | ||
|
|
||
| impl NeighborInsert for [LeafNeighbor; 1] { |
There was a problem hiding this comment.
Could we document why widths 1–3 have separate insertion implementations and link the benchmark that supports this optimization? Without that context, the extra paths look unnecessary and may be simplified later.
| } | ||
|
|
||
| #[cfg(test)] | ||
| #[allow(clippy::unwrap_used, reason = "test matrices have fixed valid shapes")] |
There was a problem hiding this comment.
This allow appears redundant.
| } | ||
|
|
||
| #[test] | ||
| fn reused_l2_leaders_match_fresh_leaders_for_a_new_point_stripe() { |
There was a problem hiding this comment.
Do these reuse tests still protect distinct behavior? Leader norms are now computed eagerly and stored in immutable Vecs, so reusing leaders has no state transition to verify. A smaller multi-point correctness test may cover the remaining behavior.
| fn compute_distances( | ||
| points: MatrixView<'_, f32>, | ||
| leaders: &Self::Leaders<'_>, | ||
| storage: &mut [f32], |
There was a problem hiding this comment.
Nit: Could compute_distances take a MutMatrixView instead? The caller already knows the shape, and the row operations become simpler:
// L2
for row in storage.row_iter_mut() {
row.copy_from_slice(&leaders.norms);
}
// Cosine
for (row, point_norm) in storage.row_iter_mut().zip(cosine_norms(points)) {
for (distance, &leader_norm) in row.iter_mut().zip(&leaders.norms) {
*distance = cosine_distance(*distance, point_norm, leader_norm);
}
}| A: PiPNNSIMDSchema, | ||
| M: PartitionMetric, | ||
| { | ||
| let point_count = points.nrows(); |
There was a problem hiding this comment.
Could this validate that output.nrows() == point_count? The later zip stops at the shorter matrix, so a short output silently skips points and still returns Ok(()), while extra output rows keep stale values.
There was a problem hiding this comment.
fixed, will throw new error when not match
Keep constant widths in row dispatch so the compiler can optimize one insertion loop. Test shifts directly and share width-dispatch coverage across scalar and SIMD updates.
|
|
||
| /// One group of distances and their column indexes in the supplied slice. | ||
| #[derive(Clone, Copy)] | ||
| pub(super) enum DistanceBlock<'a, F: PiPNNSIMDVector> { |
There was a problem hiding this comment.
Could we reconsider this module boundary? DistanceBlock appears to be an implementation protocol that only TopKRows can interpret: this module performs the SIMD load, while Top-k performs the comparison, mask traversal, lane handling, and scalar insertion. That splits one ranking operation across modules and exposes SIMD details to the leaf and partition kernels.
| return; | ||
| } | ||
| worst.resize(distances.nrows(), f32::INFINITY); | ||
| with_topk_rows!(output, worst, |topks| { |
There was a problem hiding this comment.
Could each kernel call one function that performs the full ranking? The kernels currently manage with_topk_rows!, the architecture scope, distance blocks, and the one/many updates. These details seem to belong inside Top-k.
For example, leaf ranking could call topk::rank_symmetric(...), and partition ranking could call topk::rank_rows(...).
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.rswrites one flattened lower-triangle ranking buffer for a leaf.partition_metric.rswrites one flattened point-to-leader ranking buffer for a point stripe.leaf_kernel.rsscans each unordered point pair once and retains local top-k neighbors.partition_kernel.rsretains the nearest sampled leaders for each point.simd.rsselects one dispatchedf32vector type for both ranking kernels.diskann-linalgprovides lower-triangle replace and add forms of scaledA · 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
-2 · A · Aᵀ.alpha = -1. They use-dotas the ranking distance because the omitted constant cannot change order. Their metric types remain separate.Partition
PartitionMetric::create_leadersreturns an opaque metric-owned leader set. The caller does not manage its norms.Point stripes share immutable leader norms. Construction computes them before parallel stripe work starts; there is no lazy initialization state.
Numerical behavior
[-1, 1].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:
Local validation for the audit cleanup and leaf-kernel follow-up:
cargo test --locked -p diskann --all-features --lib: 463 passed, including 109 PiPNN tests.cargo test --locked -p diskann-linalg --lib sgemm_aat_lower: 6 passed.cargo test --locked -p diskann --all-features --lib: 613 passed in an isolated worktree with the PR2 patch.cargo clippy --locked -p diskann --all-features --all-targets -- -D warnings: passed.cargo fmt --all --checkandgit 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=nativeand identical HashPrune configurations.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
diskann-linalg/src/faer.rsand the lower-triangle API indiskann-linalg/src/lib.rs.diskann/src/graph/pipnn/leaf_metric.rs.diskann/src/graph/pipnn/leaf_kernel.rs.diskann/src/graph/pipnn/partition_metric.rs.diskann/src/graph/pipnn/partition_kernel.rs.diskann/src/graph/pipnn/simd.rsand the nightly Miri scope.Stack
Stack 2/6. #1315 is merged into
main. #1290 adds provider-independent graph construction.