PiPNN 6/6: add HashPrune candidate merging - #1295
Conversation
|
Azure BigANN10M validation (
Raw runs: |
d4c2ba6 to
93cf436
Compare
Direct final-stack QCReviewed directly against current Fixed during QC
Remaining findings
Validation
Local runtime benchmark integration remains blocked by unavailable Git LFS (the fixture is a 123-byte pointer); CI's LFS checkout is the remaining runtime oracle. Full report: |
93cf436 to
efedce4
Compare
QC follow-upThe first post-rebase CI run exposed a real AArch64-only failure in The replacement run is green on:
Remote stack metadata was also rebuilt as one stack, #1301: |
efedce4 to
f8f27bf
Compare
There was a problem hiding this comment.
Pull request overview
Adds an optional HashPrune/LSH-based candidate-merging path to PiPNN builds, wiring it through disk build configuration and benchmarks, and extending SIMD/mask utilities needed by the new kernels.
Changes:
- Introduces
HashPrunereservoirs plus random-hyperplane LSH sketch computation, and integrates them into PiPNN leaf building / extraction (optionally followed by RobustPrune). - Extends disk-build and benchmark pipelines to accept and validate HashPrune parameters for PiPNN.
- Adds supporting utilities (trusted adjacency-list constructor, mask helpers, SIMD eq optimization) and CI Miri coverage for the raw-pointer kernels.
Reviewed changes
Copilot reviewed 25 out of 26 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| diskann/src/graph/adjacencylist.rs | Adds from_vec_trusted for zero-copy construction when uniqueness is guaranteed. |
| diskann-wide/src/doubled.rs | Adds first() support for doubled masks. |
| diskann-wide/src/arch/x86_64/v3/i16x16_.rs | Optimizes SIMD equality mask generation using movemask+pext. |
| diskann-vector/src/lib.rs | Makes x86_64 prefetch helpers available without requiring AVX2. |
| diskann-pipnn/tests/config.rs | Adds validation tests for HashPruneConfig parameters. |
| diskann-pipnn/tests/build_graph.rs | Adds parallel HashPrune build invariant test. |
| diskann-pipnn/src/lsh.rs | New LSH sketch computation (seeded random hyperplanes) and errors. |
| diskann-pipnn/src/lib.rs | Adds HashPrune config/types and integrates candidate-merge selection into build_graph. |
| diskann-pipnn/src/leaf_build/tests.rs | Adds CSR construction tests used by HashPrune leaf streaming. |
| diskann-pipnn/src/leaf_build.rs | Refactors leaf computation; adds CSR edge streaming into HashPrune reservoirs. |
| diskann-pipnn/src/hash_prune.rs | New HashPrune implementation (hot/cold slabs, per-row locking, SIMD hash ops, extraction). |
| diskann-pipnn/src/hash_prune/tests.rs | Adds unit and concurrency tests for HashPrune kernels and reservoir behavior. |
| diskann-pipnn/src/bf16.rs | Adds bf16 packing helpers for compact distance keys. |
| diskann-pipnn/Cargo.toml | Adds new dependencies and a HashPrune benchmark target. |
| diskann-pipnn/benches/hash_prune.rs | Adds criterion benchmark comparing direct vs HashPrune merge paths. |
| diskann-disk/src/lib.rs | Re-exports HashPruneParameters when pipnn feature is enabled. |
| diskann-disk/src/build/mod.rs | Re-exports HashPruneParameters from configuration. |
| diskann-disk/src/build/configuration/mod.rs | Exposes HashPrune parameters in configuration module exports. |
| diskann-disk/src/build/configuration/build_algorithm.rs | Extends PiPNNParameters with hash_prune and serde defaults. |
| diskann-disk/src/build/configuration/disk_index_build_parameter.rs | Switches to returning borrowed PiPNN parameters for build selection. |
| diskann-disk/src/build/builder/build/pipnn.rs | Wires optional HashPrune parameters into PiPNNBuildContext. |
| diskann-disk/src/build/builder/build/pipnn/tests.rs | Updates PiPNN disk builder tests for new parameter passing. |
| diskann-disk/src/build/builder/build.rs | Validates PiPNN + HashPrune config when pipnn is selected. |
| diskann-benchmark/src/index/build.rs | Wires optional HashPrune parameters into benchmark PiPNN builds. |
| Cargo.lock | Adds new transitive dependencies for diskann-pipnn changes. |
| .github/workflows/nightly.yml | Improves feature quoting/formatting and adds Miri strict-provenance coverage for HashPrune kernels. |
Comments suppressed due to low confidence (1)
diskann-disk/src/build/configuration/build_algorithm.rs:157
- If
PiPNNParameters::default()is changed to keephash_pruneopt-in, this serde-defaults test should be updated to expectNoneinstead ofSome(HashPruneParameters::default()).
assert_eq!(config.hash_prune, Some(HashPruneParameters::default()));
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| libc = "0.2" | ||
| parking_lot = "0.12" |
| impl Default for PiPNNParameters { | ||
| fn default() -> Self { | ||
| Self { | ||
| c_max: 256, | ||
| c_min: 16, | ||
| p_samp: 0.005, | ||
| fanout: vec![8, 3], | ||
| k: 2, | ||
| replicas: 1, | ||
| hash_prune: Some(HashPruneParameters::default()), | ||
| } | ||
| } |
f8f27bf to
661365c
Compare
661365c to
25ab02b
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-disk/src/build/configuration/build_algorithm.rs:82
PiPNNParametersis#[serde(default)], so deserializing PiPNN configs that omit the newhash_prunefield will inherit thisDefaultvalue. Settinghash_prune: Some(HashPruneParameters::default())therefore enables HashPrune by default and can change behavior for existing JSON configs that previously used direct candidate merging. If HashPrune is meant to be opt-in (as described), make the defaultNoneand require explicit configuration to enable it.
impl Default for PiPNNParameters {
fn default() -> Self {
Self {
c_max: 256,
c_min: 16,
p_samp: 0.005,
fanout: vec![8, 3],
k: 2,
replicas: 1,
hash_prune: Some(HashPruneParameters::default()),
}
25ab02b to
10bec9e
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 27 out of 28 changed files in this pull request and generated no new comments.
Suppressed comments (3)
diskann-disk/src/build/configuration/build_algorithm.rs:157
- This serde-defaults test currently asserts that omitted
hash_prunedeserializes toSome(default), which would force-enable HashPrune by default. If direct merge is intended to remain the default, this assertion should instead verify thathash_pruneis absent unless explicitly set.
assert_eq!(config.hash_prune, Some(HashPruneParameters::default()));
diskann/src/graph/pipnn/hash_prune.rs:717
collect_sorted_neighbors/collect_neighbor_idsuseVec::reserveandVec::with_capacity, which can panic on allocation failure. Elsewhere in the PiPNN pipeline allocations are handled fallibly (try_reserve*) and surfaced asANNError; this new path introduces an OOM panic in a library build/extraction phase.
let n = hot.len as usize;
scratch.clear();
scratch.reserve(n);
for i in 0..n {
// SAFETY: guaranteed by this function's contract.
scratch.push(unsafe { (*neighbors.add(i), *distances.add(i)) });
}
scratch.sort_unstable_by_key(|&(id, distance)| (distance, id));
let out_len = n.min(cap);
let mut out = Vec::with_capacity(out_len);
for &(id, d) in &scratch[..out_len] {
out.push((id, bf16_to_f32(key_to_bf16(d))));
}
out
}
/// Collect the reservoir's neighbor ids, truncated to `cap`, WITHOUT sorting.
/// Reservoir order is intentionally not preserved. Reading only `neighbors`
/// lets the caller drop the hashes and distances slabs before extraction; any
/// ordering required by a later graph-finalization policy belongs to that caller.
///
/// SAFETY: caller holds the slot lock (or owns the reservoir); `neighbors` is
/// valid for `hot.len` elements.
#[inline]
unsafe fn collect_neighbor_ids(hot: &HotSlot, neighbors: *const u32, cap: usize) -> Vec<u32> {
let out_len = (hot.len as usize).min(cap);
let mut out = Vec::with_capacity(out_len);
for i in 0..out_len {
// SAFETY: guaranteed by this function's contract.
out.push(unsafe { *neighbors.add(i) });
}
out
}
diskann-disk/src/build/configuration/build_algorithm.rs:81
PiPNNParametersnow defaultshash_prunetoSome(...), which enables HashPrune even when callers don’t specify it in JSON. That changes the default candidate-merge behavior and contradicts the PR description that direct merging remains the default unless explicitly opted in.
This issue also appears on line 157 of the same file.
hash_prune: Some(HashPruneParameters::default()),
cbbe4f5 to
91bc385
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 27 out of 28 changed files in this pull request and generated no new comments.
Suppressed comments (1)
diskann-disk/src/build/configuration/build_algorithm.rs:83
PiPNNParameters::default()now enables HashPrune by default (hash_prune: Some(...)). That contradicts the PR description (“Direct merging remains the default”) and changes behavior for JSON configs that omithash_prune, because#[serde(default)]uses the struct Default.
Consider making the default hash_prune: None so HashPrune is only enabled when explicitly requested (e.g., via "hash_prune": { ... } in JSON).
#[cfg(feature = "pipnn")]
impl Default for PiPNNParameters {
fn default() -> Self {
Self {
c_max: 256,
c_min: 16,
p_samp: 0.005,
fanout: vec![8, 3],
k: 2,
replicas: 1,
hash_prune: Some(HashPruneParameters::default()),
}
}
8225da6 to
f611b04
Compare
Route HashPrune integration through diskann::graph::pipnn, remove obsolete Cargo benchmarks, and consolidate tests around exact conversion, reservoir, concurrency, and extraction contracts.
Use the shared diskann IAI target for direct, HashPrune plus final prune, and nearest-only paths. Cover the empty-edge scratch fast path and Rust 2024 Windows FFI declaration.
Preserve legacy direct-candidate behavior when hash_prune is omitted, and reject capacity below graph degree before quantizer training or artifact creation.
Run strict-provenance Miri on dispatched hash and reservoir mutation paths, and keep libc target-specific to Linux.
Adapt HashPrune to current main, fix its concurrent aliasing boundaries, and keep all private/composition tests beside their implementation.
Overlapping leaves can emit many duplicate or directionally redundant candidates for one source. The direct PiPNN path retains every unique candidate until final RobustPrune. This PR adds optional HashPrune merging: each source streams candidates into a bounded reservoir keyed by residual direction, keeping the nearer candidate when hashes collide. Direct merging remains the default.
Concepts
lsh.rsprojects each point onto seeded random hyperplanes. For edgesource → target, signs of target/source projection differences form a relative hash. Similar residual directions collide.l_maxis logical per-point reservoir capacity. Withfinal_prune=true, extraction returns up tol_maxcandidates to shared RobustPrune. Withfalse, extraction returns the nearest graph-degree candidates directly. HashPrune bounds candidate merging; it does not own graph degree, metric, or prune policy.Code map
diskann/src/graph/pipnn/mod.rsHashPruneConfigvalidates structural bounds and effective capacity against graph degree.PiPNNBuildContext::with_hash_pruneopts in and checks effective capacity against graph degree.build_graph_innerpreserves direct mode and selects the two consuming extraction paths.lsh.rscreates deterministic random-hyperplane sketches with per-worker conversion scratch.hash_prune.rsleaf_build.rsconverts symmetric leaf output to deduplicated directed CSR, gathers active sketches, translates local/global IDs, and streams edges into reservoirs.bf16.rssupplies compact ordered distance storage.diskann-wideowns dispatched relative-hash/hash-scan implementations; PiPNN never names an ISA.hash_pruneremains direct mode; an explicit nested object opts in.End-to-end flow
Dataset → seeded point sketches + per-point reservoirs → unchanged overlapping partitioning → unchanged leaf-local nearest pairs → deduplicated directed CSR → gather active leaf sketches → compute relative hash per edge → lock only source reservoir → insert/replace/reject → consume reservoirs:
final_prune=true: full candidate lists → private shared finalization/RobustPrune → degree-bounded graph.final_prune=false: nearestRcandidates → degree-bounded graph.Without
with_hash_prune, direct candidate merging is unchanged.Invariants and boundaries
min(l_max, 2^num_hash_planes)must cover graph degree. Disk builders enforce this before quantizer training or artifact creation.scan_lanesis physical SIMD padding;l_maxis logical capacity. Padded cells are never candidates.(bf16 distance, residual hash, neighbor ID)gives history-independent retention under ties.unsafe extern; Linux mmap and Windows VirtualAlloc preserve zero-backed lazy allocation semantics.Review path
HashPruneConfig,with_hash_prune, and orchestration branches.Send/Syncsafety invariants.with_locked; trace every insertion mutation and farthest-cache update.diskann-widedispatch, including padded tails and signed-zero/NaN buckets.Test architecture
Tests are grouped by source conversion, dispatched hash primitives, slab/configuration, leaf ingestion/scratch, reservoir replacement/order, and concurrency/extraction. Weak or duplicate cases were removed only when replaced by exact stronger oracles:
Names describe behavior without mechanical
test_/should_prefixes. Private and composition tests are co-located in the implementation files that own each seam.Validation
diskannlibrary tests pass withpipnn,testing, including HashPrune configuration, graph composition, and contention cases.Stack relation
Stack 6/6. Depends on #1294 and completes the series. Private RobustPrune (#1315), numerical kernels (#1287), core construction (#1290), disk serialization (#1291), and benchmark lifecycle (#1294) remain the owning layers for those concerns.
Stack 6/6: #1294