PiPNN 6/6: add HashPrune candidate merging

[SeliMeli]

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.rs projects each point onto seeded random hyperplanes. For edge source → target, signs of target/source projection differences form a relative hash. Similar residual directions collide.

l_max is logical per-point reservoir capacity. With final_prune=true, extraction returns up to l_max candidates to shared RobustPrune. With false, extraction returns the nearest graph-degree candidates directly. HashPrune bounds candidate merging; it does not own graph degree, metric, or prune policy.

Code map

1. diskann/src/graph/pipnn/mod.rs
   • HashPruneConfig validates plane/capacity bounds.
   • PiPNNBuildContext::with_hash_prune opts in and checks effective capacity against graph degree.
   • build_graph_inner preserves direct mode and selects the two consuming extraction paths.
2. lsh.rs creates deterministic random-hyperplane sketches with per-worker conversion scratch.
3. hash_prune.rs
   • hot metadata slots plus padded cold hash/distance/neighbor slabs;
   • one raw mutex per source row;
   • empty, same-hash, not-full, early-reject, and eviction insertion branches;
   • full-candidate and nearest-only consuming extraction.
4. leaf_build.rs converts symmetric leaf output to deduplicated directed CSR, gathers active sketches, translates local/global IDs, and streams edges into reservoirs.
5. bf16.rs supplies compact ordered distance storage.
6. diskann-wide owns dispatched relative-hash/hash-scan implementations; PiPNN never names an ISA.
7. Disk and benchmark configuration thread the optional policy through existing adapters.
8. diskann/benches/benchmarks_iai/pipnn_candidate_merge.rs adds direct, HashPrune-plus-final-prune, and nearest-only cases to the single DiskANN IAI target.

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: nearest R candidates → 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.
• scan_lanes is physical SIMD padding; l_max is logical capacity. Padded cells are never candidates.
• Each hot slot owns one source lock. Length, farthest cache, and cold row mutate only under that lock.
• A hash appears at most once per logical row; same-hash replacement requires a closer total key.
• (bf16 distance, residual hash, neighbor ID) gives history-independent retention under ties.
• Leaf-local IDs become global dataset IDs before reservoir insertion.
• Grow-only sketch scratch is read only through the active prefix; empty-edge ingestion preserves existing scratch.
• Consuming extraction releases unused sketch/hash/distance storage before or while IDs materialize.
• Raw-pointer kernels require padded valid ranges and exclusive source mutation; strict-provenance Miri covers them.
• Windows slab FFI uses Rust 2024 unsafe extern; Linux mmap and Windows VirtualAlloc preserve zero-backed lazy allocation semantics.

Review path

1. Start with HashPruneConfig, with_hash_prune, and orchestration branches.
2. Review slab allocation/drop and Send/Sync safety invariants.
3. Establish the single-writer rule in with_locked; trace every insertion mutation and farthest-cache update.
4. Compare scalar references with diskann-wide dispatch, including padded tails and signed-zero/NaN buckets.
5. Trace one leaf through CSR dedup, sketch gather, ID translation, insertion, and both extraction paths.
6. Finish with disk/benchmark configuration and IAI scenario registration.

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:

• seeded serial LSH references replace probabilistic seed inequality;
• exact reservoir victim/list assertions replace length-only eviction tests;
• one distinct-direction fixture proves full extraction and exact nearest truncation;
• parallel insertion compares exact serial neighbor lists;
• repeated full builds compare canonical neighbor sets, not unspecified row order;
• empty-edge scratch reuse and finalization overflow propagation remain explicit.

Names describe behavior without mechanical test_/should_ prefixes.

Validation

• 461 diskann unit tests and 35 PiPNN integration tests pass with pipnn,testing.
• Five production disk-adapter tests and dedicated benchmark lifecycle test pass.
• Strict-provenance Miri passes relative-hash and padded hash-scan kernels.
• Windows GNU and AArch64 cross-target checks pass.
• IAI-Callgrind executes 10 consolidated cases with 5% instruction/cache regression limits; Rayon scenarios use a one-thread current-thread pool so measured work stays inside Callgrind collection.
• No standalone PiPNN crate, Criterion PiPNN target, or PiPNN-specific Cargo bench binary remains.

Stack relation

Stack 6/6. Depends on #1294 and completes the series. Numerical kernels (#1287), private RobustPrune (#1288), core construction (#1290), disk serialization (#1291), and benchmark lifecycle (#1294) remain the owning layers for those concerns.

Stack 6/6: #1294

[SeliMeli]

Azure BigANN10M validation (Standard_D16as_v5, AVX2, 16 threads, 4× interleaved per binary) used R=64, c512/c64, fanout=[10,3], k=2, l_max=72, hp=12, final_prune=true.

• Build: 76.777s clean → 79.601s stack (+3.68%)
• Peak RSS: 12.224 → 11.709 GiB (-4.22%)
• Recall@10 L=50: 0.957825 → 0.957790 (-0.0035pp)
• Mean comparisons/hops L=50: +0.023% / +0.006%

Raw runs: ~/pipnn-pr5-ab/results/final-head-k2-20260729T171245Z on the benchmark VM. The residual build gap is being investigated only on local, unpushed pipnn-stack/07-perf-investigation; current attribution points to repeated per-row FP16 conversion dispatch and CSR bookkeeping, not graph quality or HashPrune semantics.

[SeliMeli]

Direct final-stack QC

Reviewed directly against current origin/main; no subagents were used.

Fixed during QC

• Removed the proposed PiPNN-owned unsafe FP16 gather because its measured ~0.48% gain did not meet the unsafe RFC's safety/CI-evidence bar.
• Added PR6's missing direct bytemuck dependency (caught by a clean coverage build).
• Updated shared RobustPrune for the latest Neighbor accessor API.
• Rebuilt the stack from exact per-layer commit boundaries and rebased it onto current main.
• Preserved current-main nightly feature coverage while retaining strict HashPrune Miri jobs.

Remaining findings

1. Evidence gap: performance-motivated unsafe kernels have Criterion/Azure/Miri evidence, but no CI-integrated performance regression threshold as required by rfcs/00109-unsafe-rust.md.
2. Performance objective not met: allocator-matched BigANN10M 4× interleaved A/B reports build median 75.899s clean vs 77.473s final (+2.07%); wall +1.62%; RSS -3.60%; recall effectively unchanged (0.957805 vs 0.957795). No parity/win claim is made.

Validation

• Formatting and CI-style workspace Clippy: pass.
• PiPNN all targets: pass (81 unit tests plus integration/config/kernel suites and benchmark smoke executables).
• RobustPrune/Vamana tests: 22 pass.
• Focused disk PiPNN tests: 5 pass.
• pipnn and disk-index,pipnn feature compilation/Clippy: pass.
• x86-64 baseline and AArch64 all-target checks: pass.
• Strict-provenance Miri HashPrune kernel checks: pass.
• PiPNN ISA/TLS/RefCell/broadcast-cleanup scans: zero hits.
• Stable production-source coverage: 94.19% lines / 91.88% regions; experimental branch coverage 91.32%.
• Complexity/duplication audit: no maintainability blocker.
• Final Azure BigANN10M smoke: Recall@10 L=50 95.774%, build 77.918s, peak RSS 11.68 GiB.

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: /tmp/pipnn-qc-direct.md in the authoring environment.

[SeliMeli]

QC follow-up

The first post-rebase CI run exposed a real AArch64-only failure in cosine_zero_norm_masks_nan_norm_at_simd_boundaries: NEON and x86 have different max(NaN) behavior, so a NaN pair could become distance zero on ARM. PR1 now preserves NaNs explicitly with a validity mask before selecting the non-negative clamp.

The replacement run is green on:

• Ubuntu AArch64 runtime tests
• x86-64 Nehalem/SDE baseline
• AVX-512/SDE
• Ubuntu and Windows default/all-feature tests
• coverage, CodeQL, formatting, and all Clippy matrices

Remote stack metadata was also rebuilt as one stack, #1301:
#1287 → #1288 → #1290 → #1291 → #1294 → #1295.
The obsolete closed HashPrune PR is no longer a stack member.

[Copilot]

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 HashPrune reservoirs 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 keep hash_prune opt-in, this serde-defaults test should be updated to expect None instead of Some(HashPruneParameters::default()).

        assert_eq!(config.hash_prune, Some(HashPruneParameters::default()));

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

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

diskann-disk/src/build/configuration/build_algorithm.rs:82

• PiPNNParameters is #[serde(default)], so deserializing PiPNN configs that omit the new hash_prune field will inherit this Default value. Setting hash_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 default None and 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()),
        }

[Copilot]

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:82

• PiPNNParameters is documented as allowing hash_prune: None for exact/direct candidate accumulation, but the Default impl sets hash_prune: Some(HashPruneParameters::default()). Because the struct is #[serde(default)], omitting hash_prune in JSON will now silently opt into HashPrune instead of keeping direct merging as the default, which contradicts the PR description and the field’s doc comment. Consider defaulting hash_prune to None (or updating the docs/PR description if HashPrune is intended to become the default).

            fanout: vec![8, 3],
            k: 2,
            replicas: 1,
            hash_prune: Some(HashPruneParameters::default()),
        }

[Copilot]

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_prune deserializes to Some(default), which would force-enable HashPrune by default. If direct merge is intended to remain the default, this assertion should instead verify that hash_prune is 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_ids use Vec::reserve and Vec::with_capacity, which can panic on allocation failure. Elsewhere in the PiPNN pipeline allocations are handled fallibly (try_reserve*) and surfaced as ANNError; 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

• PiPNNParameters now defaults hash_prune to Some(...), 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()),

[Copilot]

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 omit hash_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()),
        }
    }