Skip to content

PiPNN 3/6: add core graph construction - #1290

Open
SeliMeli wants to merge 25 commits into
pipnn-stack/02-final-prunefrom
pipnn-stack/03-core
Open

PiPNN 3/6: add core graph construction#1290
SeliMeli wants to merge 25 commits into
pipnn-stack/02-final-prunefrom
pipnn-stack/03-core

Conversation

@SeliMeli

@SeliMeli SeliMeli commented Jul 29, 2026

Copy link
Copy Markdown

PiPNN constructs an ANN graph without beam-searching a partially built graph for every inserted point. It creates overlapping randomized partitions, performs dense leaf-local distance work, merges symmetric candidates, then applies DiskANN graph policy. This PR adds that provider-independent core under diskann::graph::pipnn.

The core consumes a borrowed dense MatrixView<T>, validated PiPNN/graph policy, and a caller-owned Rayon pool. It returns adjacency for real dataset IDs. Providers, start/frozen points, PQ, disk headers, serialization, and search remain outside.

Concepts

  • A sampled leader names a child partition.
  • Each point keeps its nearest fanout leaders, so child partitions overlap.
  • Oversized children recurse until bounded leaves remain.
  • Each leaf selects k local companions per point; these are candidates, not final graph degree R.
  • Independent replicas repeat partitioning with deterministic derived seeds.
  • PiPNNConfig owns partition/leaf policy; DiskANN graph::Config owns metric compatibility, R, alpha, prune kind, and saturation.

Code map

  1. diskann/src/graph/pipnn/mod.rs
    • PiPNNConfig::validate checks leaf bounds, sampling, fanout, k, and replicas.
    • PiPNNBuildContext::new combines algorithm policy with graph policy, metric, and Rayon pool.
    • build_graph installs the full call tree in that pool and orchestrates all stages.
  2. partitioning.rs
    • deterministic replicas, leader sampling, striped GEMM assignment, ordered scatter, recursion, small-leaf merge, and final leaf validation.
  3. leaf_build.rs
    • per-job reusable buffers, lower A · Aᵀ, prepared leaf kernel, local/global ID translation, symmetric candidate insertion, sort/dedup.
  4. finalization.rs
    • moves rows at or below R unchanged and sends only overfull rows through private shared RobustPrune.
  5. Tests are grouped by production seam: kernel, partitioning, leaf assembly/CSR, finalization, public configuration, and complete graph construction.
  6. diskann/benches/benchmarks_iai/pipnn_core.rs adds partition-heavy, single-leaf, and overfull-finalization cases to the shared IAI target.

End-to-end flow

Row-major MatrixView<T> → validate dataset/configuration → create deterministic overlapping leaves for each replica → gather each leaf and compute lower-triangular Gram matrix → select symmetric leaf-local neighbors → translate local positions to global u32 IDs → merge/sort/deduplicate candidates per source → preserve bounded rows and RobustPrune only overfull rows → return one adjacency list per input row.

Invariants and boundaries

  • Input has at least one row and dimension; area must not overflow usize; row count must fit u32 graph IDs.
  • 0 < c_min <= c_max; sampling is finite in (0,1]; fanout is nonempty and bounded; k and replicas are nonzero.
  • Graph prune kind matches metric; alpha remains owned and validated (or not) by the existing graph::Config contract.
  • Every replica covers every point; leaves may overlap but none is empty or exceeds c_max.
  • Seed derivation, scatter order, and final candidate sets are deterministic for fixed input/pool policy.
  • Grow-only scratch may retain a larger high-water allocation; all reads use the active prefix.
  • Partition pool locks cover lease/return only, never GEMM. Leaf jobs own their map_init state. No TLS or cleanup broadcast exists.
  • Core output uses real point IDs only and owns no storage/search lifecycle.

Review path

  1. Start with PiPNNConfig, PiPNNBuildContext, and build_graph to establish policy ownership.
  2. Follow one oversized partition through leader sampling, stripes, ordered scatter, recursion, and small-leaf merge.
  3. Follow one leaf through gather, lower-triangle kernel, symmetric CSR, and global candidate insertion.
  4. Finish in finalization; compare its private RobustPrune policy with graph configuration.
  5. Review ownership/reuse and allocation-error paths after algorithm flow.

Test architecture

The suite keeps one independent differential oracle per numerical kernel. Private tests target branch-heavy state; integration tests exercise only public composition. Removed cases were duplicate formula-sharing differentials, constructor-only success checks, or weaker invariant-only wrappers. Exact retained oracles cover:

  • partition seed/order/coverage, fanout exhaustion, merge flush boundaries, stripes, conversion, and validation;
  • leaf all-pairs reference, symmetric candidate/CSR assembly, duplicate/global IDs, conversion, buffer reuse, and poisoned/allocation errors;
  • finalization exact bounded/overfull output, invalid IDs/list counts, and propagated u16 candidate overflow;
  • public graph shape, IDs, degree, metrics/types, determinism, and fixed-seed randomized invariants.

Rust test names are behavior phrases (rejects_*, preserves_*, *_is_*); no mechanical test_ or should_ prefix is used for new PiPNN tests.

Validation

  • 59 private PiPNN tests plus public configuration/build integration suites at this layer, including a regression that PiPNN adds no alpha validation beyond graph::Config.
  • IAI-Callgrind runs all core scenarios on the instrumented thread via a one-thread current-thread Rayon pool.
  • cargo test -p diskann --features pipnn and Clippy with all targets pass with and without default features.

Stack relation

Stack 3/6. Depends on #1287 numerical kernels and #1288 private RobustPrune. #1291 adds production disk integration; #1294 adds benchmark entry points; #1295 adds optional HashPrune merging.

Stack 3/6: #1288#1291

@SeliMeli
SeliMeli requested review from a team and a lite review from Copilot July 29, 2026 11:52

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR introduces the “core” PiPNN build pipeline in the diskann-pipnn crate, wiring together deterministic partitioning, leaf-local candidate construction, and final pruning into a public build_graph API with a validated build context.

Changes:

  • Adds PiPNNConfig validation and a PiPNNBuildContext that binds PiPNN policy to DiskANN graph pruning policy and a caller-owned Rayon thread pool.
  • Implements the three main stages: partitioning (partitioning.rs), leaf candidate construction (leaf_build.rs), and final pruning via shared Vamana robust prune (finalization.rs).
  • Adds comprehensive unit/integration tests and a Criterion benchmark for core scenarios; updates dependencies, lockfile, and mutation-test exclusions.

Reviewed changes

Copilot reviewed 13 out of 14 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
diskann-pipnn/src/lib.rs Adds public PiPNN API (PiPNNConfig, PiPNNBuildContext, build_graph) and stage orchestration.
diskann-pipnn/src/partitioning.rs Implements deterministic overlapping partition construction and leader assignment/scatter.
diskann-pipnn/src/partitioning/tests.rs Adds unit tests covering partition determinism, invariants, error cases, and helpers.
diskann-pipnn/src/leaf_build.rs Builds leaf-local symmetric k-NN candidates and accumulates global candidates safely in parallel.
diskann-pipnn/src/leaf_build/tests.rs Adds unit tests for candidate correctness, invariants, type support, and error handling.
diskann-pipnn/src/finalization.rs Orders/prunes candidate rows using shared robust_prune and validates candidate IDs/shape.
diskann-pipnn/src/finalization/tests.rs Adds unit tests for pruning behavior and candidate validation failures.
diskann-pipnn/src/tests.rs Tests effective_metric behavior for integer cosine-normalized handling.
diskann-pipnn/tests/config.rs Integration tests for config validation and graph-policy compatibility checks.
diskann-pipnn/tests/build_graph.rs Integration tests for end-to-end graph building, invariants, determinism, and type/metric support.
diskann-pipnn/benches/core.rs Adds a Criterion benchmark for stage-focused core build scenarios.
diskann-pipnn/Cargo.toml Updates crate dependencies/dev-dependencies and registers the new core benchmark target.
Cargo.lock Records dependency graph changes for the updated diskann-pipnn crate dependencies.
.cargo/mutants.toml Adds mutation-test exclusions for key PiPNN public boundary checks and partitioning invariants.
Comments suppressed due to low confidence (1)

diskann-pipnn/src/partitioning.rs:604

  • size_of::<f32>() is used without being in scope (no use std::mem::size_of; and not qualified), so this function won’t compile as written.
fn assignment_stripe_rows(leaders: usize) -> usize {
    (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>()))
        .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS)
}

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

Comment thread diskann-pipnn/src/partitioning.rs Outdated
Comment on lines +285 to +289
*scale = FastL2NormSquared.evaluate(row);
if metric == Metric::Cosine {
*scale = scale.sqrt();
}
}
Copilot AI review requested due to automatic review settings July 29, 2026 13:10
@SeliMeli
SeliMeli force-pushed the pipnn-stack/03-core branch from 5be0c9c to 50047c6 Compare July 29, 2026 13:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 13 out of 14 changed files in this pull request and generated 1 comment.

Comment thread diskann-pipnn/src/partitioning.rs Outdated
Comment on lines +606 to +609
fn assignment_stripe_rows(leaders: usize) -> usize {
(ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>()))
.clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS)
}
Copilot AI review requested due to automatic review settings July 29, 2026 16:38
@SeliMeli SeliMeli changed the title Pipnn stack/03 core PiPNN 3/6: add core graph construction Jul 29, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Comments suppressed due to low confidence (2)

diskann-pipnn/src/partitioning.rs:637

  • size_of is used without being in scope (std::mem::size_of), which will not compile. Qualify the call or import it.
fn assignment_stripe_rows(leaders: usize) -> usize {
    (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>()))
        .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS)
}

diskann-pipnn/src/partitioning.rs:19

  • Norm is imported but never used in this module, which will trip unused_imports warnings (and can become CI failures under -D warnings). Remove it from the import list.
use diskann::{utils::VectorRepr, ANNError, ANNResult};
use diskann_linalg::Transpose;
use diskann_utils::views::MatrixView;
use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm};
use rand::{prelude::IndexedRandom, SeedableRng};
use rayon::prelude::*;

Copilot AI review requested due to automatic review settings July 30, 2026 07:47

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Comments suppressed due to low confidence (1)

diskann-pipnn/src/partitioning.rs:390

  • gather_rows uses TypeId::of::<T>(), which implicitly requires T: 'static. Making that bound explicit here avoids surprising/indirect trait-bound errors later and matches the public build_graph boundary (which already requires 'static).
fn gather_rows<T>(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()>
where
    T: VectorRepr,

Copilot AI review requested due to automatic review settings July 30, 2026 08:26
@SeliMeli
SeliMeli force-pushed the pipnn-stack/03-core branch from 1324668 to 857e200 Compare July 30, 2026 08:26

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 14 out of 15 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

diskann-pipnn/src/partitioning.rs:641

  • size_of::<f32>() is used without being imported or qualified, which will fail to compile. Qualify it with std::mem::size_of (or add an explicit import).
    let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::<f32>());

Comment thread diskann-pipnn/src/partitioning.rs Outdated
use diskann::{utils::VectorRepr, ANNError, ANNResult};
use diskann_linalg::Transpose;
use diskann_utils::views::MatrixView;
use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm};
Copilot AI review requested due to automatic review settings July 30, 2026 08:55
@SeliMeli
SeliMeli force-pushed the pipnn-stack/03-core branch from 857e200 to f642204 Compare July 30, 2026 08:55

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 14 out of 15 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (4)

diskann-pipnn/src/leaf_build.rs:222

  • build_leaf is executed from a Rayon parallel context (via build_leaf_candidates), so it should also explicitly require T: Send + Sync to reflect the actual thread-safety requirement.
where
    T: VectorRepr + 'static,
{

diskann-pipnn/src/leaf_build/tests.rs:154

  • assert_source_type forwards T into the parallel leaf build path, so it should also include Send + Sync bounds to match the production requirements.
fn assert_source_type<T>(data: &[T])
where
    T: diskann::utils::VectorRepr + 'static,
{

diskann-pipnn/src/leaf_build.rs:193

  • build_leaf_candidates uses Rayon parallel iteration over data, so T must be Send + Sync. Making this explicit in the signature avoids confusing trait-bound errors at call sites and documents the thread-safety requirement.

This issue also appears on line 220 of the same file.

where
    T: VectorRepr + 'static,
{

diskann-pipnn/src/leaf_build/tests.rs:35

  • This test helper calls build_leaf_candidates, which (via Rayon) requires T: Send + Sync. Add the bounds here so the test continues to compile once the production signature is tightened.

This issue also appears on line 151 of the same file.

where
    T: diskann::utils::VectorRepr + 'static,
{

fanout: usize,
leaders: usize,
) -> ANNResult<Vec<Vec<u32>>> {
let mut sizes = filled_vec(leaders, 0usize)?;
Copilot AI review requested due to automatic review settings July 30, 2026 11:26
@SeliMeli
SeliMeli force-pushed the pipnn-stack/03-core branch from f642204 to b1181d6 Compare July 30, 2026 11:26

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 15 out of 16 changed files in this pull request and generated 1 comment.

Comment on lines +792 to +800
fn assignment_stripe_point_count(leader_count: usize) -> usize {
let point_count = ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count.max(1) * size_of::<f32>());
let point_count = if point_count.is_power_of_two() {
point_count
} else {
point_count.next_power_of_two() / 2
};
point_count.clamp(MIN_ASSIGNMENT_STRIPE_POINTS, MAX_ASSIGNMENT_STRIPE_POINTS)
}
Copilot AI review requested due to automatic review settings August 5, 2026 11:04

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Suppressed comments (2)

diskann/src/graph/pipnn/leaf_build.rs:344

  • This map_err block is not rustfmt-formatted (indentation inside the struct literal is off). This is likely to fail cargo fmt --check in CI; reformat the block to match the rest of the file.
    let dots = MatrixView::try_from(&buffers.dots[..dot_count], point_ids.len(), point_ids.len())
        .map_err(|error| LeafBuildError::Kernel {
        leaf,
        source: LeafKernelError::InvalidBufferLength {
            buffer: "leaf dot-product matrix",
            expected: dot_count,
            actual: error.into_inner().len(),
        },
    })?;

diskann/src/graph/pipnn/leaf_build.rs:305

  • build_leaf calls buffers.prepare(...) (which can grow the dot-product buffer to point_count^2) before checking whether any neighbors will be produced. For requested_k == 0 (and trivially for singleton leaves), we can return early after ID validation and avoid unnecessary allocation/work.
    let leaf_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?;

SeliMeli added 25 commits August 5, 2026 12:01
Move core construction under diskann::graph::pipnn so finalization can reuse private RobustPrune state. Remove standalone PiPNN Cargo benchmarks, group public tests by PiPNN behavior, and delete duplicate or non-discriminating cases.
Restore the finalization adapter overflow oracle while consolidating stage-focused builds in the shared diskann IAI target.
Keep sorting, workspace allocation, source exclusion, and adjacency rewriting in PiPNN finalization. The shared internal kernel now sees only prepared candidates and state.
Address the renamed internal modules and preserve graph Config alpha behavior without adding PiPNN-specific validation.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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

Suppressed comments (1)

diskann/src/graph/pipnn/partitioning.rs:794

  • size_of::<f32>() isn’t in the Rust prelude; this will fail to compile unless you import it. Qualify it or add use std::mem::size_of;.
fn assignment_stripe_point_count(leader_count: usize) -> usize {
    let point_count = ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count.max(1) * size_of::<f32>());
    let point_count = if point_count.is_power_of_two() {

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants