PiPNN 2/6: add numerical kernels - #1287
weiyaoluo (SeliMeli) wants to merge 127 commits into
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.
e204cb9 to
b046174
Compare
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1287 +/- ##
==========================================
+ Coverage 91.55% 92.72% +1.17%
==========================================
Files 521 533 +12
Lines 100302 104113 +3811
==========================================
+ Hits 91828 96540 +4712
+ Misses 8474 7573 -901
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
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" |
Keep width dispatch and SIMD distance blocks inside TopK. Reuse the same row storage for replacement and dual updates, and retain the measured unfilled-result path. Consolidate tests around reachable kernel behavior, independent output oracles, and workspace reuse.
Remove unreachable empty-kernel smoke cases and duplicate delegated checks. Keep point replacement fixtures at a fixed leader count, combine cosine formula and multi-row checks, and exercise rectangular AAT inputs.
Keep partition selection thresholds local and initialize reusable leaf state explicitly. Kernels own candidate storage across updates.
Enter the architecture scope per TopK operation with explicit buffer arguments so kernels only traverse rows. Preserve const capacity inside the feature boundary to keep sorted insertion specialized in release.
There was a problem hiding this comment.
🟡 Changes recommended
Both kernels perform unchecked distance-matrix size multiplication before reaching the new overflow validation.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 14/15 changed files
- Comments generated: 2
- Review effort level: Balanced
Aditya Krishnan (arkrishn94)
left a comment
There was a problem hiding this comment.
Thanks Weiyao.
- I'm especially concerned about the coverage of the tests for the kernels - e.g. I don't see the happy paths being sufficiently tested over enough dimensions and widths. I also notice that the tests in partition_kernel and leaf_kernel don't exercise all metrics for the happy paths.
- In terms of documentation, I have a slightly hard time following what exactly
update_dual_top_kandupdate_manyintopk.rsare supposed to be doing. I thinktopk.rsneeds better documentation to describe the places the functions are used and what exactly they are doing.
| /// Specialize small capacities 1..=10; larger capacities use the runtime path. | ||
| /// Typical partition K is 10 or 3, and leaf K is 3 or 2. | ||
| /// The width is evaluated once. The body has ordinary block control flow. | ||
| macro_rules! with_topk { |
There was a problem hiding this comment.
I'm not sure why you need this macro. If the goal here is to support using a topk value instantiated based on const/runtime width, this is not a extensible way to do it. It is hard to read/parse too. Why not just bubble the K const up higher so it can be passed?
There was a problem hiding this comment.
replaced with visitor trait
| /// Default SIMD representation used by both PiPNN ranking kernels. | ||
| /// | ||
| /// This alias is the single build-time width selection. | ||
| type DefaultVector<A> = <A as Architecture>::f32x16; |
There was a problem hiding this comment.
Remind me, are we fixing the width to 16 across architectures? Is this reasonable for Neon?
There was a problem hiding this comment.
NEON support seems not a priority for now. 16 works well on both avx2 and avx512, will keep this fixed for now. added todo comment for neon supporting
Check complete metric matrices and count/width ranking grids. Keep ranking and distance tolerances separate. Clarify test names and reusable-state fixtures. Run GEMM-backed tests natively and keep them out of Miri.
Name the dual-update test and dispatcher after their entry point. Record pending Neon width validation.
| /// Default SIMD representation used by both PiPNN ranking kernels. | ||
| /// | ||
| /// This alias is the single build-time width selection. | ||
| /// |
There was a problem hiding this comment.
Use Github issue to track instead of todo in code.
PiPNN callers already bound or check their matrix shapes. Keep the shared constructor change and its tests outside this PR.
The merge with main enables clippy::allow_attributes. Use one scoped dead-code expectation until graph construction wires the kernels.
Summary
Add GEMM-backed distance computation and SIMD top-k selection for PiPNN behind the opt-in diskann/pipnn feature.
The numerical layer supports two operations:
This is part 2 of the PiPNN stack. Recursive graph construction integrates these kernels in #1290.
Design
Separate metric computation from neighbor selection. Metric implementations handle L2, cosine, normalized cosine, and inner product. Partition leaders retain metric-specific norms for reuse across point stripes. Kernels choose which distance slices
to process; the shared TopK implementation handles SIMD loading, threshold filtering, scalar tails, and sorted insertion.
Reuse each leaf distance for both endpoints. New diskann-linalg helpers compute or accumulate the lower triangle of a Gram matrix using faer. Leaf selection scans the strict lower triangle, excludes self-neighbors, and shares each loaded distance
block between both endpoint updates.
Specialize common top-k capacities. Capacities 1–10 use compile-time specialization to enable insertion-loop unrolling; larger capacities use the runtime path. Typical partition fanouts are 10 and 3, and typical leaf-neighbor counts are 2 or 3.
Architecture-specific selection runs through diskann-wide at the TopK entry points.
Keep storage reusable. TopK borrows caller-owned candidate buffers. Kernels manage distance and ranking scratch without exposing SIMD block types to their callers.
The linear-algebra helpers validate matrix shapes and size-product overflow.