From 6747c3a97c0afa64649c6e04d640212d60367aec Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:22:48 +0000 Subject: [PATCH 01/58] pipnn: assemble direct-candidate graph builder --- Cargo.lock | 1 + diskann/Cargo.toml | 6 +- diskann/src/graph/pipnn/finalization.rs | 123 ++++ diskann/src/graph/pipnn/finalization/tests.rs | 103 +++ diskann/src/graph/pipnn/leaf_build.rs | 347 ++++++++++ diskann/src/graph/pipnn/leaf_build/tests.rs | 370 +++++++++++ diskann/src/graph/pipnn/mod.rs | 222 ++++++- diskann/src/graph/pipnn/partitioning.rs | 612 ++++++++++++++++++ diskann/src/graph/pipnn/partitioning/tests.rs | 277 ++++++++ diskann/tests/build_graph.rs | 231 +++++++ diskann/tests/config.rs | 133 ++++ 11 files changed, 2402 insertions(+), 23 deletions(-) create mode 100644 diskann/src/graph/pipnn/finalization.rs create mode 100644 diskann/src/graph/pipnn/finalization/tests.rs create mode 100644 diskann/src/graph/pipnn/leaf_build.rs create mode 100644 diskann/src/graph/pipnn/leaf_build/tests.rs create mode 100644 diskann/src/graph/pipnn/partitioning.rs create mode 100644 diskann/src/graph/pipnn/partitioning/tests.rs create mode 100644 diskann/tests/build_graph.rs create mode 100644 diskann/tests/config.rs diff --git a/Cargo.lock b/Cargo.lock index 40240e77c5..be5f9ef094 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -448,6 +448,7 @@ dependencies = [ "num-traits", "pin-project", "rand", + "rayon", "relative-path 2.0.1", "rstest", "serde", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index c0016416dc..2568ac6669 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -14,6 +14,7 @@ targets = ["x86_64-unknown-linux-gnu", "aarch64-pc-windows-msvc", "x86_64-pc-win [dependencies] anyhow.workspace = true bytemuck = { workspace = true, features = ["must_cast"]} +diskann-linalg = { workspace = true, optional = true } diskann-utils = { workspace = true, default-features = false } futures-util = { workspace = true, default-features = false } half = { workspace = true, features = ["bytemuck", "num-traits"] } @@ -22,6 +23,7 @@ half = { workspace = true, features = ["bytemuck", "num-traits"] } hashbrown = { version = "0.16.0", default-features = false, features = ["default-hasher"] } num-traits.workspace = true rand.workspace = true +rayon = { workspace = true, optional = true } thiserror.workspace = true tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } tracing = { workspace = true, optional = true } @@ -58,8 +60,8 @@ panic = "warn" [features] default = ["tracing"] -# Enable PiPNN numerical kernels. -pipnn = ["dep:diskann-linalg"] +# Enable PiPNN batch graph construction. +pipnn = ["dep:diskann-linalg", "dep:rayon", "tracing"] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs new file mode 100644 index 0000000000..dab9419cc8 --- /dev/null +++ b/diskann/src/graph/pipnn/finalization.rs @@ -0,0 +1,123 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Orders complete PiPNN candidate rows and applies shared Vamana RobustPrune. + +use std::convert::Infallible; + +use crate::{ + graph::{prune, AdjacencyList, Config}, + neighbor::Neighbor, + utils::VectorRepr, + ANNError, ANNResult, +}; +use diskann_utils::views::MatrixView; +use diskann_vector::{distance::Metric, DistanceFunction}; +use rayon::prelude::*; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum FinalizationError { + #[error("candidate row count {rows} does not match the dataset point count {points}")] + RowCountMismatch { rows: usize, points: usize }, + #[error("candidate ID {candidate} in row {row} is outside a {points}-point dataset")] + InvalidCandidateId { + row: usize, + candidate: u32, + points: usize, + }, +} + +#[derive(Default)] +struct Workspace { + prune: prune::Scratch, + cache: Vec<(f32, Option)>, +} + +pub(crate) fn prune_overfull( + data: MatrixView<'_, T>, + candidates: Vec>, + graph: &Config, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + validate_candidates(&candidates, data.nrows()).map_err(ANNError::opaque)?; + + let degree = graph.pruned_degree().get(); + let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind(), false); + let distance = T::distance(metric, Some(data.ncols())); + + // build_graph installs the complete call tree in the caller-owned pool. + #[allow(clippy::disallowed_methods)] + candidates + .into_par_iter() + .enumerate() + .map_init(Workspace::default, |workspace, (source, mut row)| { + if row.len() <= degree { + return Ok(row); + } + + let source_id = u32::try_from(source).map_err(ANNError::opaque)?; + let source_vector = data.row(source); + let pool = workspace.prune.candidates_mut(); + pool.clear(); + pool.try_reserve(row.len()).map_err(ANNError::opaque)?; + pool.extend(row.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + let candidate_count = pool.len(); + let mut context = workspace.prune.as_context(candidate_count); + prune::robust_prune( + &mut context, + policy, + &mut workspace.cache, + Some, + |left, right| { + Ok::<_, Infallible>( + distance.evaluate_similarity( + data.row(*left as usize), + data.row(*right as usize), + ), + ) + }, + |id| id == source_id, + ) + .map_err(ANNError::opaque)?; + + row.clear(); + row.extend_from_slice(workspace.prune.neighbors()); + Ok(row) + }) + .collect() +} + +fn validate_candidates( + candidates: &[AdjacencyList], + points: usize, +) -> Result<(), FinalizationError> { + if candidates.len() != points { + return Err(FinalizationError::RowCountMismatch { + rows: candidates.len(), + points, + }); + } + for (row_id, row) in candidates.iter().enumerate() { + if let Some(&candidate) = row.iter().find(|&&id| id as usize >= points) { + return Err(FinalizationError::InvalidCandidateId { + row: row_id, + candidate, + points, + }); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/pipnn/finalization/tests.rs b/diskann/src/graph/pipnn/finalization/tests.rs new file mode 100644 index 0000000000..fde3a6e0ea --- /dev/null +++ b/diskann/src/graph/pipnn/finalization/tests.rs @@ -0,0 +1,103 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use crate::graph::{ + config::{self, MaxDegree}, + AdjacencyList, +}; +use diskann_utils::views::MatrixView; + +use super::*; + +fn graph_config(degree: usize) -> Config { + config::Builder::new_with( + degree, + MaxDegree::same(), + degree, + Metric::L2.into(), + |builder| { + builder.alpha(1.2); + }, + ) + .build() + .unwrap() +} + +fn row(ids: impl IntoIterator) -> AdjacencyList { + AdjacencyList::from_iter_untrusted(ids) +} + +#[test] +fn preserves_rows_within_the_degree_bound() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![row([3, 1]), row([]), row([]), row([])]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert_eq!(&*actual[0], &[1, 3]); +} + +#[test] +fn prunes_an_overfull_row_with_the_vamana_kernel() { + let data = [0.0_f32, 1.0, 2.0, -3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![row([3, 2, 1]), row([]), row([]), row([])]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert!(actual[0].len() <= 2); + assert!(actual[0].contains(1)); +} + +#[test] +fn rejects_invalid_candidate_ids_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![row([1, 3]), row([]), row([])]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::InvalidCandidateId { + row: 0, + candidate: 3, + points: 3, + }) + )); +} + +#[test] +fn rejects_candidate_row_count_mismatch_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![row([]), row([]), row([]), row([])]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::RowCountMismatch { rows: 4, points: 3 }) + )); +} + +#[test] +fn rejects_more_candidates_than_the_shared_position_type_can_represent() { + let count = u16::MAX as usize + 1; + let data = vec![0.0_f32; count + 1]; + let data = MatrixView::try_from(&data[..], count + 1, 1).unwrap(); + let mut candidates = Vec::with_capacity(count + 1); + candidates.push(row(1..=count as u32)); + candidates.resize_with(count + 1, AdjacencyList::new); + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::>(), + Some(prune::RobustPruneError::TooManyCandidates { actual, max }) + if *actual == count && *max == u16::MAX as usize + )); +} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs new file mode 100644 index 0000000000..f5fdb22928 --- /dev/null +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -0,0 +1,347 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Leaf construction and direct candidate accumulation. + +use std::{ + collections::{HashSet, TryReserveError}, + sync::Mutex, +}; + +use crate::{graph::AdjacencyList, utils::VectorRepr}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use rayon::prelude::*; + +use crate::leaf_kernel::{ + nearest_leaf_neighbors, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, +}; + +/// Failure while converting leaves into direct graph candidates. +#[derive(Debug, thiserror::Error)] +pub(crate) enum LeafBuildError { + #[error("leaf build requires at least one dimension")] + EmptyDimensions, + #[error("dataset point count {0} exceeds the u32 ID limit")] + TooManyPoints(usize), + #[error("leaf {leaf} is empty")] + EmptyLeaf { leaf: usize }, + #[error("point ID {point} in leaf {leaf} is outside a {points}-point dataset")] + InvalidPointId { + leaf: usize, + point: u32, + points: usize, + }, + #[error("point ID {point} appears more than once in leaf {leaf}")] + DuplicatePointId { leaf: usize, point: u32 }, + #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] + ShapeOverflow { + leaf: usize, + rows: usize, + columns: usize, + }, + #[error("failed to reserve {additional} values for {buffer}")] + Allocation { + buffer: &'static str, + additional: usize, + #[source] + source: TryReserveError, + }, + #[error("failed to convert point {point} in leaf {leaf}")] + Conversion { + leaf: usize, + point: u32, + #[source] + source: diskann::ANNError, + }, + #[error("lower-AAT failed for leaf {leaf}")] + LowerAat { + leaf: usize, + #[source] + source: diskann_linalg::SgemmError, + }, + #[error("nearest-neighbor selection failed for leaf {leaf}")] + Kernel { + leaf: usize, + #[source] + source: LeafKernelError, + }, + #[error("leaf kernel returned local position {position} for a {points}-point leaf")] + InvalidLocalPosition { position: u32, points: usize }, + #[error("candidate row {point} is poisoned")] + PoisonedCandidateRow { point: u32 }, +} + +#[derive(Default)] +struct LeafBuffers { + points: Vec, + dots: Vec, + nearest: Vec, + local_graph: Vec>, + top_k: LeafTopKWorkspace, + seen_ids: HashSet, +} + +impl LeafBuffers { + fn prepare( + &mut self, + leaf: usize, + points: usize, + dimensions: usize, + k: usize, + ) -> Result { + let point_values = points + .checked_mul(dimensions) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: dimensions, + })?; + let dot_values = points + .checked_mul(points) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: points, + })?; + let actual_k = k.min(points.saturating_sub(1)); + let nearest_values = points + .checked_mul(actual_k) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: points, + columns: actual_k, + })?; + + resize("leaf points", &mut self.points, point_values, 0.0)?; + resize("leaf dot products", &mut self.dots, dot_values, 0.0)?; + resize( + "leaf nearest neighbors", + &mut self.nearest, + nearest_values, + LeafNeighbor::default(), + )?; + if points > self.local_graph.len() { + let additional = points - self.local_graph.len(); + self.local_graph + .try_reserve(additional) + .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; + self.local_graph.resize_with(points, AdjacencyList::new); + } + self.local_graph[..points] + .iter_mut() + .for_each(AdjacencyList::clear); + Ok(actual_k) + } +} + +struct DirectCandidates { + rows: Vec>>, +} + +impl DirectCandidates { + fn new(points: usize) -> Result { + let mut rows = Vec::new(); + rows.try_reserve_exact(points) + .map_err(|source| allocation_error("candidate rows", points, source))?; + rows.resize_with(points, || Mutex::new(AdjacencyList::new())); + Ok(Self { rows }) + } + + fn add_leaf( + &self, + point_ids: &[u32], + local_graph: &[AdjacencyList], + ) -> Result<(), LeafBuildError> { + for (&source, additions) in point_ids.iter().zip(local_graph) { + // Every point ID is validated before leaf-local work begins. + let row = &self.rows[source as usize]; + let mut row = row.lock().map_err(|_| poisoned_row(source))?; + row.extend_from_slice(additions); + } + Ok(()) + } + + fn into_rows(self) -> Result>, LeafBuildError> { + let mut output = Vec::new(); + output + .try_reserve_exact(self.rows.len()) + .map_err(|source| allocation_error("candidate output", self.rows.len(), source))?; + for (point, row) in self.rows.into_iter().enumerate() { + let mut row = row.into_inner().map_err(|_| poisoned_row(point as u32))?; + row.sort(); + output.push(row); + } + Ok(output) + } +} + +/// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. +#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. +pub(crate) fn build_leaf_candidates( + data: MatrixView<'_, T>, + leaves: &[Vec], + k: usize, + metric: Metric, +) -> Result>, LeafBuildError> +where + T: VectorRepr + 'static, +{ + if data.ncols() == 0 { + return Err(LeafBuildError::EmptyDimensions); + } + if data.nrows() > u32::MAX as usize { + return Err(LeafBuildError::TooManyPoints(data.nrows())); + } + + let candidates = DirectCandidates::new(data.nrows())?; + leaves.par_iter().enumerate().try_for_each_init( + LeafBuffers::default, + |buffers, (leaf, point_ids)| { + build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) + }, + )?; + candidates.into_rows() +} + +fn build_leaf( + data: MatrixView<'_, T>, + leaf: usize, + point_ids: &[u32], + k: usize, + metric: Metric, + buffers: &mut LeafBuffers, + candidates: &DirectCandidates, +) -> Result<(), LeafBuildError> +where + T: VectorRepr + 'static, +{ + if point_ids.is_empty() { + return Err(LeafBuildError::EmptyLeaf { leaf }); + } + buffers.seen_ids.clear(); + buffers + .seen_ids + .try_reserve(point_ids.len()) + .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; + for &point in point_ids { + if point as usize >= data.nrows() { + return Err(LeafBuildError::InvalidPointId { + leaf, + point, + points: data.nrows(), + }); + } + if !buffers.seen_ids.insert(point) { + return Err(LeafBuildError::DuplicatePointId { leaf, point }); + } + } + let actual_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), k)?; + if actual_k == 0 { + return Ok(()); + } + + for (&point, output) in point_ids + .iter() + .zip(buffers.points.chunks_exact_mut(data.ncols())) + { + // Point IDs were validated before the zero-k/singleton fast path. + let row = data.row(point as usize); + T::as_f32_into(row, output).map_err(|source| LeafBuildError::Conversion { + leaf, + point, + source: source.into(), + })?; + } + + diskann_linalg::sgemm_aat_lower( + &buffers.points, + point_ids.len(), + data.ncols(), + &mut buffers.dots, + ) + .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; + nearest_leaf_neighbors( + LeafTopK { + dots: &buffers.dots, + points: point_ids.len(), + metric, + }, + k, + &mut buffers.nearest, + &mut buffers.top_k, + ) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + + add_symmetric_edges( + point_ids, + actual_k, + &buffers.nearest, + &mut buffers.local_graph[..point_ids.len()], + )?; + candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) +} + +fn add_symmetric_edges( + point_ids: &[u32], + k: usize, + nearest: &[LeafNeighbor], + local_graph: &mut [AdjacencyList], +) -> Result<(), LeafBuildError> { + for (source, nearest) in nearest.chunks_exact(k).enumerate() { + for neighbor in nearest { + let target = neighbor.position as usize; + let Some(&target_id) = point_ids.get(target) else { + return Err(LeafBuildError::InvalidLocalPosition { + position: neighbor.position, + points: point_ids.len(), + }); + }; + let source_id = point_ids[source]; + if source_id != target_id { + local_graph[source].push(target_id); + local_graph[target].push(source_id); + } + } + } + Ok(()) +} + +fn resize( + buffer: &'static str, + values: &mut Vec, + len: usize, + value: T, +) -> Result<(), LeafBuildError> { + if len > values.len() { + let additional = len - values.len(); + values + .try_reserve(additional) + .map_err(|source| allocation_error(buffer, additional, source))?; + values.resize(len, value); + } else { + values.truncate(len); + } + Ok(()) +} + +fn allocation_error( + buffer: &'static str, + additional: usize, + source: TryReserveError, +) -> LeafBuildError { + LeafBuildError::Allocation { + buffer, + additional, + source, + } +} + +fn poisoned_row(point: u32) -> LeafBuildError { + LeafBuildError::PoisonedCandidateRow { point } +} + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs new file mode 100644 index 0000000000..d3f63263dc --- /dev/null +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -0,0 +1,370 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use half::f16; +use std::collections::BTreeSet; + +use super::{ + add_symmetric_edges, allocation_error, build_leaf_candidates, DirectCandidates, LeafBuffers, + LeafBuildError, +}; + +fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + MatrixView::try_from(data, rows, columns).unwrap() +} + +fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() +} + +fn build( + data: MatrixView<'_, T>, + leaves: &[Vec], + k: usize, + metric: Metric, +) -> Result>, LeafBuildError> +where + T: diskann::utils::VectorRepr + 'static, +{ + pool().install(|| build_leaf_candidates(data, leaves, k, metric)) +} + +fn rows(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() +} + +fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { + let mut graph = vec![BTreeSet::new(); data.len()]; + for (source, left) in data.iter().enumerate() { + let mut nearest: Vec<_> = data + .iter() + .enumerate() + .filter(|(target, _)| *target != source) + .map(|(target, right)| { + let distance = left + .iter() + .zip(right) + .map(|(x, y)| (x - y) * (x - y)) + .sum::(); + (target, distance) + }) + .collect(); + nearest.sort_by(|left, right| { + left.1 + .total_cmp(&right.1) + .then_with(|| left.0.cmp(&right.0)) + }); + for &(target, _) in nearest.iter().take(k) { + graph[source].insert(target as u32); + graph[target].insert(source as u32); + } + } + graph + .into_iter() + .map(|neighbors| neighbors.into_iter().collect()) + .collect() +} + +#[test] +fn leaf_adjacency_matches_an_independent_all_pairs_reference() { + let points = [ + [0.0_f32, 0.0], + [1.0, 0.2], + [3.1, 0.5], + [7.8, 1.4], + [-2.3, 4.1], + [6.7, -3.2], + ]; + let flat: Vec<_> = points.into_iter().flatten().collect(); + + let actual = rows( + build( + view(&flat, points.len(), 2), + &[(0..points.len() as u32).collect()], + 2, + Metric::L2, + ) + .unwrap(), + ); + + assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); +} + +#[test] +fn retains_and_deduplicates_candidates_from_overlapping_leaves() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; + + let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + rows(graph), + [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] + ); +} + +#[test] +fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { + let dimensions = 9; + let mut data = vec![0.0_f32; 10 * dimensions]; + for row in 1..10 { + data[row * dimensions + row - 1] = 1.0; + } + + let graph = build( + view(&data, 10, dimensions), + &[(0..10).collect()], + 1, + Metric::L2, + ) + .unwrap(); + + assert_eq!(&*graph[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); + assert!(graph.iter().enumerate().all(|(source, neighbors)| { + neighbors.iter().all(|&target| target as usize != source) + && neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + })); +} + +#[test] +fn global_id_translation_is_independent_of_leaf_order() { + let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; + let leaves = vec![vec![4, 1, 3]]; + + let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + rows(graph), + [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] + ); +} + +fn assert_source_type(data: &[T]) +where + T: diskann::utils::VectorRepr + 'static, +{ + let leaves = vec![vec![0, 1, 2, 3]]; + let graph = build(view(data, 4, 2), &leaves, 1, Metric::L2).unwrap(); + assert_eq!(rows(graph), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); +} + +#[test] +fn gathers_every_supported_source_type_without_full_dataset_conversion() { + assert_source_type(&[0.0_f32, 0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0]); + assert_source_type(&[0_i8, 0, 1, 0, 2, 0, 3, 0]); + assert_source_type(&[0_u8, 0, 1, 0, 2, 0, 3, 0]); + assert_source_type(&[ + f16::from_f32(0.0), + f16::from_f32(0.0), + f16::from_f32(1.0), + f16::from_f32(0.0), + f16::from_f32(2.0), + f16::from_f32(0.0), + f16::from_f32(3.0), + f16::from_f32(0.0), + ]); +} + +#[test] +fn all_metrics_produce_symmetric_unique_non_self_candidates() { + let data = [1.0_f32, 0.0, 0.8, 0.2, 0.0, 1.0, -1.0, 0.0]; + let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); + for (source, neighbors) in graph.iter().enumerate() { + assert!(neighbors.iter().all(|&target| target as usize != source)); + assert!(neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32))); + assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); + } + } +} + +#[test] +fn parallel_leaf_schedule_does_not_change_candidate_order() { + let data: Vec = (0..64).map(|value| value as f32).collect(); + let leaves: Vec> = (0..32) + .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) + .collect(); + let pool = pool(); + pool.install(|| { + let expected = build_leaf_candidates(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + for _ in 0..8 { + let actual = build_leaf_candidates(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + assert_eq!(actual, expected); + } + }); +} + +#[test] +fn rejects_invalid_shape_inputs_without_panicking() { + let data = [0.0_f32, 1.0]; + let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); + assert!(matches!( + build(no_dimensions, &[], 1, Metric::L2), + Err(LeafBuildError::EmptyDimensions) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), + Err(LeafBuildError::EmptyLeaf { leaf: 0 }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { + leaf: 0, + point: 2, + points: 2 + }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) + )); +} + +#[test] +fn singleton_and_zero_k_leaves_add_no_candidates() { + let data = [0.0_f32, 1.0, 2.0]; + let singleton = build( + view(&data, 3, 1), + &[vec![0], vec![1], vec![2]], + 1, + Metric::L2, + ) + .unwrap(); + let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); + assert!(singleton.iter().chain(&zero_k).all(|row| row.is_empty())); +} + +#[test] +fn reuses_worker_buffers_for_smaller_leaves() { + let mut buffers = LeafBuffers::default(); + buffers.prepare(0, 64, 128, 2).unwrap(); + let points = buffers.points.as_ptr(); + let dots = buffers.dots.as_ptr(); + let nearest = buffers.nearest.as_ptr(); + + buffers.prepare(1, 8, 128, 2).unwrap(); + + assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.dots.as_ptr(), dots); + assert_eq!(buffers.nearest.as_ptr(), nearest); +} + +#[test] +fn reports_shape_overflow_before_allocating() { + let mut buffers = LeafBuffers::default(); + assert!(matches!( + buffers.prepare(7, usize::MAX, 2, 1), + Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) + )); +} + +#[test] +fn rejects_an_invalid_kernel_position() { + let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + let error = add_symmetric_edges( + &[10, 20], + 1, + &[ + crate::leaf_kernel::LeafNeighbor::new(9, 1.0), + crate::leaf_kernel::LeafNeighbor::new(0, 1.0), + ], + &mut graph, + ) + .unwrap_err(); + assert!(matches!( + error, + LeafBuildError::InvalidLocalPosition { + position: 9, + points: 2 + } + )); +} + +#[test] +fn skips_duplicate_global_ids_without_self_edges() { + let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + add_symmetric_edges( + &[7, 7], + 1, + &[ + crate::leaf_kernel::LeafNeighbor::new(1, 0.0), + crate::leaf_kernel::LeafNeighbor::new(0, 0.0), + ], + &mut graph, + ) + .unwrap(); + assert!(graph.iter().all(|row| row.is_empty())); +} + +#[test] +fn poisoned_candidate_rows_return_errors() { + let candidates = DirectCandidates::new(1).unwrap(); + let _ = std::panic::catch_unwind(|| { + let _guard = candidates.rows[0].lock().unwrap(); + panic!("poison candidate row"); + }); + assert!(matches!( + candidates.add_leaf(&[0], &[diskann::graph::AdjacencyList::new()]), + Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + )); + assert!(matches!( + candidates.into_rows(), + Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + )); +} + +#[test] +fn allocation_errors_preserve_buffer_context() { + let mut values = Vec::::new(); + let source = values.try_reserve(usize::MAX).unwrap_err(); + let error = allocation_error("test", 1, source); + assert!(matches!( + error, + LeafBuildError::Allocation { + buffer: "test", + additional: 1, + .. + } + )); +} + +#[test] +fn direct_candidate_accumulator_keeps_unique_sorted_rows() { + let candidates = DirectCandidates::new(2).unwrap(); + candidates + .add_leaf( + &[0, 1], + &[ + diskann::graph::AdjacencyList::from_iter_untrusted([1, 1]), + diskann::graph::AdjacencyList::from_iter_untrusted([0]), + ], + ) + .unwrap(); + assert_eq!(rows(candidates.into_rows().unwrap()), [vec![1], vec![0]]); +} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 092ee7549a..3d9680b6a1 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,31 +3,211 @@ * Licensed under the MIT license. */ -//! Numerical kernels for PiPNN graph construction. +//! Numerical kernels for provider-independent PiPNN graph construction. //! -//! [`partition_kernel`] converts point-to-leader dot products into sorted leader -//! positions. The output width sets the fanout. A scratch vector stores the -//! ranked leaders and reuses its allocation for each point. +//! PiPNN assigns points to overlapping leader partitions. It computes one +//! lower-triangular all-pairs matrix per bounded leaf. It then merges selected +//! leaf neighbors into graph candidates. //! -//! [`leaf_kernel`] reads a lower-triangular Gram matrix. It evaluates each point -//! pair once and updates both points. Each point retains at most three local -//! neighbors. +//! - [`partition_kernel`] ranks leader-column positions. Output width is runtime +//! fanout. A scratch vector retains ranked leaders and reuses its allocation. +//! - [`leaf_kernel`] scans each strict-lower-triangle pair once. It updates both +//! endpoints and retains up to three leaf-local neighbors per point. +//! - `kernel_metric` owns norm preparation and exact metric ranking inputs. //! -//! `kernel_metric` defines metric markers and shared math. Separate leaf and -//! partition traits define norm preparation and ranking formulas. -//! -//! The graph builder selects architecture `A` and metric `M` once. It passes -//! these concrete types to both kernels. -//! -//! Each kernel checks all view and norm relationships before unchecked SIMD -//! access. The kernels borrow their matrices. They write only to caller-owned -//! output and workspace. -#[allow(dead_code)] +//! The graph build selects concrete architecture `A` and metric `M` types once. +//! Both kernels validate views and metric norm layouts before unchecked SIMD +//! access. They mutate only caller-owned output and scratch storage. + +mod finalization; mod kernel_metric; -#[allow(dead_code)] mod simd; - -#[allow(dead_code)] +mod leaf_build; mod leaf_kernel; -#[allow(dead_code)] mod partition_kernel; +mod partitioning; + +use crate::{ + graph::{AdjacencyList, Config}, + utils::VectorRepr, + ANNError, ANNResult, +}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use rayon::ThreadPool; + +/// Configuration of PiPNN's partitioning and local-neighbor algorithm. +/// +/// Graph degree, pruning policy, and alpha belong to DiskANN's graph +/// configuration and are supplied separately through [`PiPNNBuildContext`]. +#[derive(Clone, Debug, PartialEq)] +pub struct PiPNNConfig { + /// Maximum number of points in a leaf. + pub c_max: usize, + /// Minimum leaf size used by global small-leaf merging. + pub c_min: usize, + /// Fraction of a cluster sampled as partition leaders. + pub p_samp: f64, + /// Number of nearest leaders retained at each overlapping partition level. + pub fanout: Vec, + /// Number of nearest neighbors selected within each leaf. + pub k: usize, + /// Number of independent partition passes over the dataset. + pub replicas: usize, +} + +impl PiPNNConfig { + fn validate(&self) -> ANNResult<()> { + if self.c_max == 0 { + return Err(config_error("c_max must be greater than zero")); + } + if self.c_min == 0 { + return Err(config_error("c_min must be greater than zero")); + } + if self.c_min > self.c_max { + return Err(config_error(format!( + "c_min ({}) must not exceed c_max ({})", + self.c_min, self.c_max + ))); + } + if !self.p_samp.is_finite() || !(0.0..=1.0).contains(&self.p_samp) || self.p_samp == 0.0 { + return Err(config_error(format!( + "p_samp ({}) must be finite and in (0, 1]", + self.p_samp + ))); + } + if self.fanout.is_empty() { + return Err(config_error("fanout must not be empty")); + } + if let Some(&fanout) = self + .fanout + .iter() + .find(|&&fanout| !(1..=partition_kernel::MAX_PARTITION_FANOUT).contains(&fanout)) + { + return Err(config_error(format!( + "fanout ({fanout}) must be in [1, {}]", + partition_kernel::MAX_PARTITION_FANOUT + ))); + } + if self.k == 0 { + return Err(config_error("k must be greater than zero")); + } + if self.replicas == 0 { + return Err(config_error("replicas must be greater than zero")); + } + Ok(()) + } +} + +/// Validated, borrowed policy and execution context for one PiPNN graph build. +#[derive(Debug)] +pub struct PiPNNBuildContext<'a> { + pub(crate) config: PiPNNConfig, + pub(crate) graph: &'a Config, + pub(crate) metric: Metric, + pub(crate) pool: &'a ThreadPool, +} + +impl<'a> PiPNNBuildContext<'a> { + /// Validate and combine PiPNN configuration with outer graph policy. + pub fn new( + config: PiPNNConfig, + graph: &'a Config, + metric: Metric, + pool: &'a ThreadPool, + ) -> ANNResult { + config.validate()?; + if !graph.alpha().is_finite() || graph.alpha() < 1.0 { + return Err(config_error(format!( + "graph alpha ({}) must be finite and at least 1", + graph.alpha() + ))); + } + if graph.prune_kind() != metric.into() { + return Err(config_error(format!( + "graph prune kind {:?} is incompatible with metric {metric:?}", + graph.prune_kind() + ))); + } + + Ok(Self { + config, + graph, + metric, + pool, + }) + } +} + +/// Build PiPNN adjacency for real rows in `data`. +/// +/// This is the core algorithm boundary. Search entry-point selection, frozen nodes, +/// providers, serialization, and index writers belong to the outer build pipelines. +pub fn build_graph( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync + 'static, +{ + context.pool.install(|| build_graph_inner(data, context)) +} + +fn build_graph_inner( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync + 'static, +{ + if data.nrows() == 0 { + return Err(ANNError::log_dimension_mismatch_error( + "PiPNN requires at least one data row".into(), + )); + } + if data.ncols() == 0 { + return Err(ANNError::log_dimension_mismatch_error( + "PiPNN requires at least one data dimension".into(), + )); + } + if data.nrows() > u32::MAX as usize { + return Err(config_error(format!( + "dataset row count ({}) exceeds the u32 graph ID limit", + data.nrows() + ))); + } + data.nrows().checked_mul(data.ncols()).ok_or_else(|| { + ANNError::log_dimension_mismatch_error(format!( + "PiPNN dataset shape {} x {} overflows usize", + data.nrows(), + data.ncols() + )) + })?; + let metric = effective_metric::(context.metric); + + let leaves = tracing::info_span!("pipnn.partition") + .in_scope(|| partitioning::partition(data, &context.config, metric))?; + let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::build_leaf_candidates(data, &leaves, context.config.k, metric) + .map_err(ANNError::opaque) + })?; + tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) +} + +fn effective_metric(metric: Metric) -> Metric { + use std::any::TypeId; + + if metric == Metric::CosineNormalized + && (TypeId::of::() == TypeId::of::() || TypeId::of::() == TypeId::of::()) + { + Metric::Cosine + } else { + metric + } +} + +#[track_caller] +fn config_error(message: impl std::fmt::Display) -> ANNError { + ANNError::log_index_config_error("PiPNN".into(), message.to_string()) +} diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs new file mode 100644 index 0000000000..40e321c83d --- /dev/null +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -0,0 +1,612 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Deterministic overlapping partition construction for PiPNN. +//! +//! The stage maps real dataset rows to bounded leaf ID lists. Numerical work +//! reuses the partition kernel and dense GEMM; scratch belongs to the Rayon +//! iterator that uses it, so no thread-local cleanup protocol is required. + +use std::collections::HashSet; + +use crate::{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::*; + +use crate::{ + partition_kernel::{nearest_leaders, PartitionTopK}, + PiPNNConfig, +}; + +// Private algorithm and batching constants live together. None are user policy. +const PARTITION_SEED: u64 = 1_000; +const LEADER_CAP: usize = 1_000; +const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 512 * 1024; +const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; +const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; +const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; +const SCATTER_STRIPE_ROWS: usize = 64 * 1024; +const MAX_PARTITION_ITERATIONS: usize = 30; + +/// A partition failure with enough context to diagnose non-progressing input. +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub(crate) enum PartitionError { + #[error("PiPNN cannot partition an empty dataset")] + EmptyDataset, + #[error("PiPNN cannot partition vectors with zero dimensions")] + EmptyDimensions, + #[error("dataset has {0} rows, which exceeds the u32 ID limit")] + TooManyPoints(usize), + #[error("{buffer} shape {rows} x {cols} overflows usize")] + ShapeOverflow { + buffer: &'static str, + rows: usize, + cols: usize, + }, + #[error( + "partition stopped after {limit} iterations with an oversized cluster of size \ + {size} at level {level}" + )] + IterationLimit { + size: usize, + level: usize, + limit: usize, + }, + #[error("partition produced an invalid leaf of size {size}; expected 1..={limit}")] + InvalidLeaf { size: usize, limit: usize }, + #[error("invalid {buffer} length: expected {expected}, got {actual}")] + InvalidBufferLength { + buffer: &'static str, + expected: usize, + actual: usize, + }, + #[error("partition worker did not publish its result")] + MissingWorkerResult, +} + +struct WorkItem { + indices: Vec, + level: usize, + seed: u64, +} + +#[derive(Default)] +struct StripeBuffers { + points: Vec, + dots: Vec, + row_scales: Vec, +} + +/// Partition every configured replica of `data` into leaves no larger than +/// `config.c_max`. The caller is responsible for installing this operation in +/// its build-owned Rayon pool. +pub(crate) fn partition( + data: MatrixView<'_, T>, + config: &PiPNNConfig, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let points = data.nrows(); + if points == 0 { + return Err(ANNError::opaque(PartitionError::EmptyDataset)); + } + if data.ncols() == 0 { + return Err(ANNError::opaque(PartitionError::EmptyDimensions)); + } + if points > u32::MAX as usize { + return Err(ANNError::opaque(PartitionError::TooManyPoints(points))); + } + + let mut leaves = Vec::new(); + for replica in 0..config.replicas { + let seed = mix_seed(PARTITION_SEED, replica as u64); + let mut replica_leaves = partition_replica(data, config, metric, seed)?; + leaves + .try_reserve(replica_leaves.len()) + .map_err(ANNError::opaque)?; + leaves.append(&mut replica_leaves); + } + validate_leaves(&leaves, config.c_max)?; + Ok(leaves) +} + +fn partition_replica( + data: MatrixView<'_, T>, + config: &PiPNNConfig, + metric: Metric, + seed: u64, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let initial_indices = point_ids(data.nrows())?; + if data.nrows() <= config.c_max { + let mut leaves = Vec::new(); + leaves.try_reserve_exact(1).map_err(ANNError::opaque)?; + leaves.push(initial_indices); + return Ok(leaves); + } + + let mut leaves = Vec::new(); + let mut work = Vec::new(); + work.try_reserve_exact(1).map_err(ANNError::opaque)?; + work.push(WorkItem { + indices: initial_indices, + level: 0, + seed, + }); + + for _ in 0..MAX_PARTITION_ITERATIONS { + if work.is_empty() { + return global_merge_small(leaves, config.c_min, config.c_max); + } + + let mut results = Vec::new(); + results + .try_reserve_exact(work.len()) + .map_err(ANNError::opaque)?; + results.resize_with(work.len(), || None); + // build_graph installs this complete private call tree into the + // caller-owned pool; the indexed fill cannot escape that pool. + #[allow(clippy::disallowed_methods)] + results + .par_iter_mut() + .zip(work.into_par_iter()) + .for_each(|(slot, item)| { + *slot = Some(partition_one_level(data, config, metric, item)); + }); + + let mut next_work = Vec::new(); + for result in results { + let (mut pending, mut finished) = + result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??; + next_work + .try_reserve(pending.len()) + .map_err(ANNError::opaque)?; + leaves + .try_reserve(finished.len()) + .map_err(ANNError::opaque)?; + next_work.append(&mut pending); + leaves.append(&mut finished); + } + work = next_work; + } + + if work.is_empty() { + return global_merge_small(leaves, config.c_min, config.c_max); + } + let mut largest = &work[0]; + for item in &work[1..] { + if item.indices.len() > largest.indices.len() { + largest = item; + } + } + Err(ANNError::opaque(PartitionError::IterationLimit { + size: largest.indices.len(), + level: largest.level, + limit: MAX_PARTITION_ITERATIONS, + })) +} + +fn partition_one_level( + data: MatrixView<'_, T>, + config: &PiPNNConfig, + metric: Metric, + item: WorkItem, +) -> ANNResult<(Vec, Vec>)> +where + T: VectorRepr + Send + Sync, +{ + let points = item.indices.len(); + let fanout = config.fanout.get(item.level).copied().unwrap_or(1); + let leaders = sample_leaders( + &item.indices, + config.p_samp, + mix_seed(item.seed, points as u64), + )?; + let clusters = assign_to_leaders(data, &item.indices, &leaders, fanout, metric)?; + + let mut pending = Vec::new(); + let mut finished = Vec::new(); + pending + .try_reserve(clusters.len()) + .map_err(ANNError::opaque)?; + finished + .try_reserve(clusters.len()) + .map_err(ANNError::opaque)?; + let child_seed = mix_seed(item.seed, points as u64); + for cluster in clusters { + if cluster.is_empty() { + continue; + } + if cluster.len() <= config.c_max { + finished.push(cluster); + } else { + pending.push(WorkItem { + indices: cluster, + level: item.level + 1, + seed: child_seed, + }); + } + } + Ok((pending, finished)) +} + +fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResult> { + let count = sample_num_leaders(points.len(), sampling_fraction); + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let mut leaders = Vec::new(); + leaders.try_reserve_exact(count).map_err(ANNError::opaque)?; + leaders.extend(points.choose_multiple(&mut rng, count).copied()); + Ok(leaders) +} + +fn sample_num_leaders(points: usize, sampling_fraction: f64) -> usize { + ((points as f64 * sampling_fraction).ceil() as usize) + .clamp(2, LEADER_CAP) + .min(points) +} + +// A single LCG mixer derives both replica and recursive seeds. Wrapping makes +// the mapping stable across debug/release builds and supported platforms. +fn mix_seed(seed: u64, salt: u64) -> u64 { + seed.wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(salt) +} + +fn assign_to_leaders( + data: MatrixView<'_, T>, + points: &[u32], + leaders: &[u32], + fanout: usize, + metric: Metric, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync, +{ + let dimensions = data.ncols(); + let leader_values_len = checked_area("leader data", leaders.len(), dimensions)?; + let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; + gather_rows(data, leaders, &mut leader_values)?; + + let mut leader_scales = if matches!(metric, Metric::L2 | Metric::Cosine) { + filled_vec(leaders.len(), 0.0f32)? + } else { + Vec::new() + }; + for (scale, row) in leader_scales + .iter_mut() + .zip(leader_values.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(row); + if metric == Metric::Cosine { + *scale = scale.sqrt(); + } + } + + let fanout = fanout.min(leaders.len()); + let assignment_len = checked_area("partition assignments", points.len(), fanout)?; + let mut assignments = filled_vec(assignment_len, 0u32)?; + let stripe_rows = assignment_stripe_rows(leaders.len()); + let assignment_stripe = checked_area("assignment stripe", stripe_rows, fanout)?; + + // build_graph pins this terminal operation to the caller-owned pool. + #[allow(clippy::disallowed_methods)] + assignments + .par_chunks_mut(assignment_stripe) + .enumerate() + .try_for_each_init(StripeBuffers::default, |buffers, (stripe, output)| { + let first = stripe * stripe_rows; + let rows = output.len() / fanout; + let point_values_len = checked_area("point stripe", rows, dimensions)?; + let dots_len = checked_area("dot-product stripe", rows, leaders.len())?; + resize_fallible(&mut buffers.points, point_values_len, 0.0)?; + resize_fallible(&mut buffers.dots, dots_len, 0.0)?; + gather_rows(data, &points[first..first + rows], &mut buffers.points)?; + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + rows, + leaders.len(), + dimensions, + 1.0, + &buffers.points, + &leader_values, + None, + &mut buffers.dots, + ) + .map_err(ANNError::opaque)?; + + let row_scales = if metric == Metric::Cosine { + resize_fallible(&mut buffers.row_scales, rows, 0.0)?; + for (scale, row) in buffers + .row_scales + .iter_mut() + .zip(buffers.points.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(row); + } + buffers.row_scales.as_slice() + } else { + &[] + }; + nearest_leaders( + PartitionTopK { + dots: &buffers.dots, + rows, + leaders: leaders.len(), + row_scales, + leader_scales: &leader_scales, + metric, + }, + fanout, + output, + ) + .map_err(ANNError::opaque) + })?; + + scatter_assignments(points, &assignments, fanout, leaders.len()) +} + +fn gather_rows(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> +where + T: VectorRepr, +{ + let expected = checked_area("gather output", indices.len(), data.ncols())?; + if output.len() != expected { + return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "gather output", + expected, + actual: output.len(), + })); + } + for (&index, row) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { + T::as_f32_into(data.row(index as usize), row).map_err(Into::::into)?; + } + Ok(()) +} + +fn scatter_assignments( + points: &[u32], + assignments: &[u32], + fanout: usize, + leaders: usize, +) -> ANNResult>> { + if points.len() < PARALLEL_SCATTER_MIN_POINTS { + return scatter_serial(points, assignments, fanout, leaders); + } + + let assignment_stripe = checked_area("scatter assignment stripe", SCATTER_STRIPE_ROWS, fanout)?; + let stripes = points.len().div_ceil(SCATTER_STRIPE_ROWS); + let mut partials = Vec::new(); + partials + .try_reserve_exact(stripes) + .map_err(ANNError::opaque)?; + partials.resize_with(stripes, || None); + // See the pool invariant at the other partition terminal operations. + #[allow(clippy::disallowed_methods)] + partials + .par_iter_mut() + .zip( + points + .par_chunks(SCATTER_STRIPE_ROWS) + .zip(assignments.par_chunks(assignment_stripe)), + ) + .for_each(|(slot, (points, assignments))| { + *slot = Some(scatter_serial(points, assignments, fanout, leaders)); + }); + + let mut locals = Vec::new(); + locals + .try_reserve_exact(stripes) + .map_err(ANNError::opaque)?; + for result in partials { + locals.push(result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??); + } + + let mut sizes = filled_vec(leaders, 0usize)?; + for local in &locals { + for (size, cluster) in sizes.iter_mut().zip(local) { + *size = size.checked_add(cluster.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "cluster size", + rows: *size, + cols: cluster.len(), + }) + })?; + } + } + + let mut clusters = clusters_with_capacities(&sizes)?; + for local in locals { + for (cluster, part) in clusters.iter_mut().zip(local) { + debug_assert!(cluster.capacity().saturating_sub(cluster.len()) >= part.len()); + cluster.extend(part); + } + } + Ok(clusters) +} + +fn scatter_serial( + points: &[u32], + assignments: &[u32], + fanout: usize, + leaders: usize, +) -> ANNResult>> { + let mut sizes = filled_vec(leaders, 0usize)?; + for &leader in assignments { + let Some(size) = sizes.get_mut(leader as usize) else { + return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "leader assignment", + expected: leaders, + actual: leader as usize + 1, + })); + }; + *size = size.checked_add(1).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "cluster size", + rows: *size, + cols: 1, + }) + })?; + } + let mut clusters = clusters_with_capacities(&sizes)?; + for (&point, row) in points.iter().zip(assignments.chunks_exact(fanout)) { + for &leader in row { + clusters[leader as usize].push(point); + } + } + Ok(clusters) +} + +fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { + let mut clusters = Vec::new(); + clusters + .try_reserve_exact(sizes.len()) + .map_err(ANNError::opaque)?; + for &size in sizes { + let mut cluster = Vec::new(); + cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + clusters.push(cluster); + } + Ok(clusters) +} + +fn global_merge_small( + leaves: Vec>, + c_min: usize, + c_max: usize, +) -> ANNResult>> { + let mut merged = Vec::new(); + let mut small_leaves = Vec::new(); + merged.try_reserve(leaves.len()).map_err(ANNError::opaque)?; + small_leaves + .try_reserve(leaves.len()) + .map_err(ANNError::opaque)?; + for leaf in leaves { + if leaf.len() >= c_min { + merged.push(leaf); + } else { + small_leaves.push(leaf); + } + } + if small_leaves.is_empty() { + return Ok(merged); + } + + let mut small = HashSet::new(); + small.try_reserve(c_max).map_err(ANNError::opaque)?; + + for leaf in small_leaves { + let combined = small.len().checked_add(leaf.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "small-leaf merge", + rows: small.len(), + cols: leaf.len(), + }) + })?; + if combined > c_max { + merged.push(drain_sorted(&mut small)?); + } + small.try_reserve(leaf.len()).map_err(ANNError::opaque)?; + small.extend(leaf); + if small.len() >= c_min { + merged.push(drain_sorted(&mut small)?); + } + } + + if !small.is_empty() { + let mut remainder = drain_sorted(&mut small)?; + if remainder.len() < c_min { + if let Some(last) = merged.last_mut() { + remainder.retain(|id| !last.contains(id)); + let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "small-leaf tail merge", + rows: last.len(), + cols: remainder.len(), + }) + })?; + if combined <= c_max { + last.try_reserve(remainder.len()) + .map_err(ANNError::opaque)?; + last.append(&mut remainder); + last.sort_unstable(); + } + } + } + if !remainder.is_empty() { + merged.push(remainder); + } + } + + validate_leaves(&merged, c_max)?; + Ok(merged) +} + +fn drain_sorted(set: &mut HashSet) -> ANNResult> { + let mut values = Vec::new(); + values + .try_reserve_exact(set.len()) + .map_err(ANNError::opaque)?; + values.extend(set.drain()); + values.sort_unstable(); + Ok(values) +} + +fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { + if let Some(leaf) = leaves + .iter() + .find(|leaf| leaf.is_empty() || leaf.len() > c_max) + { + return Err(ANNError::opaque(PartitionError::InvalidLeaf { + size: leaf.len(), + limit: c_max, + })); + } + Ok(()) +} + +fn point_ids(points: usize) -> ANNResult> { + let mut ids = Vec::new(); + ids.try_reserve_exact(points).map_err(ANNError::opaque)?; + ids.extend(0..points as u32); + Ok(ids) +} + +fn filled_vec(len: usize, value: T) -> ANNResult> { + let mut values = Vec::new(); + values.try_reserve_exact(len).map_err(ANNError::opaque)?; + values.resize(len, value); + Ok(values) +} + +fn resize_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { + if len > values.len() { + values + .try_reserve(len - values.len()) + .map_err(ANNError::opaque)?; + } + values.resize(len, value); + Ok(()) +} + +fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { + rows.checked_mul(cols) + .ok_or_else(|| ANNError::opaque(PartitionError::ShapeOverflow { buffer, rows, cols })) +} + +fn assignment_stripe_rows(leaders: usize) -> usize { + (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::())) + .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) +} + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs new file mode 100644 index 0000000000..30eed3a2d1 --- /dev/null +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -0,0 +1,277 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_utils::views::{Matrix, MatrixView}; +use diskann_vector::{distance::Metric, Half}; + +use super::*; + +fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { + PiPNNConfig { + c_max, + c_min, + p_samp: 0.25, + fanout, + k: 2, + replicas, + } +} + +fn clustered_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let row = position / dimensions; + let column = position % dimensions; + position += 1; + (row / 8) as f32 * 10.0 + column as f32 * 0.01 + row as f32 * 0.001 + } + }), + points, + dimensions, + ) +} + +fn directional_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let row = position / dimensions; + let column = position % dimensions; + position += 1; + let angle = std::f32::consts::TAU * row as f32 / points as f32; + match column { + 0 => angle.cos(), + 1 => angle.sin(), + _ => 0.0, + } + } + }), + points, + dimensions, + ) +} + +fn sorted_memberships(leaves: &[Vec]) -> Vec> { + let mut memberships: Vec> = leaves + .iter() + .map(|leaf| { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids + }) + .collect(); + memberships.sort(); + memberships +} + +fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { + assert!(leaves + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max)); + let mut counts = vec![0usize; points]; + for leaf in leaves { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids.len(), leaf.len(), "duplicate ID inside a leaf"); + for &id in leaf { + assert!((id as usize) < points); + counts[id as usize] += 1; + } + } + assert!(counts.iter().all(|&count| count >= replicas)); +} + +#[test] +fn returns_one_leaf_at_and_below_c_max() { + for points in [7, 8] { + let data = clustered_data(points, 3); + let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + assert_eq!(leaves, vec![(0..points as u32).collect::>()]); + } +} + +#[test] +fn partition_is_fixed_seed_deterministic_and_bounded() { + let data = clustered_data(96, 8); + let config = config(4, 16, vec![3, 2], 2); + + let first = partition(data.as_view(), &config, Metric::L2).unwrap(); + let second = partition(data.as_view(), &config, Metric::L2).unwrap(); + + assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); + assert_valid_partition(&first, 96, 16, 2); + assert!(first.iter().map(Vec::len).sum::() > 96 * 2); +} + +#[test] +fn recursion_after_fanout_levels_falls_back_to_one() { + let data = clustered_data(80, 4); + let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + + assert_valid_partition(&leaves, 80, 8, 1); +} + +#[test] +fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { + let data = Matrix::new(1.0f32, 24, 4); + let error = partition(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = error.downcast::().unwrap(); + + assert!(matches!( + error, + PartitionError::IterationLimit { + size: 24, + limit: MAX_PARTITION_ITERATIONS, + .. + } + )); +} + +#[test] +fn global_merge_canonicalizes_small_leaf_membership() { + let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; + + let merged = global_merge_small(leaves, 4, 8).unwrap(); + + assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); +} + +#[test] +fn global_merge_never_overfills_before_reaching_c_min() { + let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; + + let merged = global_merge_small(leaves, 11, 11).unwrap(); + + assert_eq!( + merged, + vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] + ); +} + +#[test] +fn replicas_cover_every_point_once_or_more_per_replica() { + let data = directional_data(72, 5); + let leaves = partition( + data.as_view(), + &config(3, 12, vec![3, 2], 3), + Metric::CosineNormalized, + ) + .unwrap(); + + assert_valid_partition(&leaves, 72, 12, 3); +} + +#[test] +fn supported_source_types_share_partition_contract() { + let f32_data: Vec = (0..64 * 4).map(|value| (value % 23) as f32).collect(); + let half_data: Vec = f32_data.iter().copied().map(Half::from_f32).collect(); + let u8_data: Vec = f32_data.iter().map(|value| *value as u8).collect(); + let i8_data: Vec = u8_data.iter().map(|value| *value as i8 - 11).collect(); + let config = config(2, 16, vec![2, 1], 1); + + let f32_leaves = partition( + MatrixView::try_from(f32_data.as_slice(), 64, 4).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + let half_leaves = partition( + MatrixView::try_from(half_data.as_slice(), 64, 4).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + let u8_leaves = partition( + MatrixView::try_from(u8_data.as_slice(), 64, 4).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + let i8_leaves = partition( + MatrixView::try_from(i8_data.as_slice(), 64, 4).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + + for leaves in [&f32_leaves, &half_leaves, &u8_leaves, &i8_leaves] { + assert_valid_partition(leaves, 64, 16, 1); + } + assert_eq!( + sorted_memberships(&f32_leaves), + sorted_memberships(&half_leaves) + ); + assert_eq!( + sorted_memberships(&f32_leaves), + sorted_memberships(&u8_leaves) + ); + assert_eq!( + sorted_memberships(&u8_leaves), + sorted_memberships(&i8_leaves) + ); +} + +#[test] +fn all_metrics_produce_valid_partitions() { + let data = directional_data(64, 8); + let config = config(2, 20, vec![2], 1); + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let leaves = partition(data.as_view(), &config, metric).unwrap(); + assert_valid_partition(&leaves, 64, 20, 1); + } +} + +#[test] +fn leader_count_is_bounded() { + assert_eq!(sample_num_leaders(1, 1.0), 1); + assert_eq!(sample_num_leaders(10, 0.01), 2); + assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); +} + +#[test] +fn parallel_scatter_matches_serial_order() { + let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); + let assignments: Vec = points + .iter() + .flat_map(|point| [point % 7, (point + 3) % 7]) + .collect(); + + let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); + let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); + + assert_eq!(actual, expected); +} + +#[test] +fn rejects_empty_dataset() { + let data = Matrix::::new(0.0, 0, 4); + let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDataset + ); +} + +#[test] +fn rejects_zero_dimensions() { + let data = Matrix::::new(0.0, 4, 0); + let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDimensions + ); +} diff --git a/diskann/tests/build_graph.rs b/diskann/tests/build_graph.rs new file mode 100644 index 0000000000..0f823778eb --- /dev/null +++ b/diskann/tests/build_graph.rs @@ -0,0 +1,231 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +#![cfg(feature = "pipnn")] +#![allow( + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] + +use diskann::graph::config::{self, MaxDegree}; +use diskann::graph::pipnn::{build_graph, PiPNNBuildContext, PiPNNConfig}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; +use half::f16; +use rand::{rngs::StdRng, Rng, SeedableRng}; + +fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 4, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + } +} + +fn graph_config(metric: Metric, degree: usize) -> diskann::graph::Config { + config::Builder::new_with(degree, MaxDegree::same(), 8, metric.into(), |builder| { + builder.alpha(1.2); + }) + .build() + .unwrap() +} + +fn pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() +} + +fn rows(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() +} + +fn assert_graph_invariants( + graph: &[diskann::graph::AdjacencyList], + points: usize, + degree: usize, +) { + assert_eq!(graph.len(), points); + for (source, row) in graph.iter().enumerate() { + assert!(row.len() <= degree); + let mut sorted = row.to_vec(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!(sorted.len(), row.len()); + assert!(row + .iter() + .all(|&id| (id as usize) < points && id as usize != source)); + } +} + +#[test] +fn builds_a_single_leaf_graph_for_real_dataset_ids() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let graph = graph_config(Metric::L2, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); + + let graph = graph_config(Metric::L2, 1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let pruned = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&pruned, 4, 1); + for (source, neighbors) in pruned.iter().enumerate() { + assert_eq!(source.abs_diff(neighbors[0] as usize), 1); + } +} + +#[test] +fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { + let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; + let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); + let graph = graph_config(Metric::L2, 1); + let pool = pool(2); + let config = PiPNNConfig { + c_max: 5, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 4, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&actual, 5, 1); + assert!(actual.iter().all(|row| row.len() == 1)); +} + +#[test] +fn rejects_empty_dataset_dimensions_at_the_public_boundary() { + let graph = graph_config(Metric::L2, 2); + let pool = pool(1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); + let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); + + assert!(build_graph(no_rows, &context).is_err()); + assert!(build_graph(no_columns, &context).is_err()); +} + +#[test] +fn supports_every_source_type_and_metric() { + fn build(values: &[T], metric: Metric) { + let data = MatrixView::try_from(values, 6, 2).unwrap(); + let graph = graph_config(metric, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); + let actual = build_graph(data, &context).unwrap(); + assert_graph_invariants(&actual, 6, 2); + } + + let values = [ + 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, + ]; + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + build(&values, metric); + } + build(&values.map(f16::from_f32), Metric::L2); + build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); + build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); +} + +#[test] +fn integer_normalized_cosine_matches_cosine() { + fn assert_match(values: &[T]) { + let data = MatrixView::try_from(values, 8, 2).unwrap(); + let pool = pool(2); + let build = |metric| { + let graph = graph_config(metric, 2); + let config = PiPNNConfig { + c_max: 8, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + rows(build_graph(data, &context).unwrap()) + }; + assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); + } + + assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); + assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); +} + +#[test] +fn is_deterministic_for_a_fixed_pool_size() { + let data: Vec = (0..96 * 4) + .map(|value| ((value * 17 + 3) % 101) as f32) + .collect(); + let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); + let graph = graph_config(Metric::L2, 8); + let pool = pool(4); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.25, + fanout: vec![3, 2], + k: 3, + replicas: 2, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let first = build_graph(data, &context).unwrap(); + let second = build_graph(data, &context).unwrap(); + + assert_eq!(first, second); + assert_graph_invariants(&first, 96, 8); +} + +#[test] +fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { + let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); + for case in 0..24 { + let points = rng.random_range(4..=32); + let dimensions = rng.random_range(1..=8); + let c_max = rng.random_range(4..=points.min(12)); + let c_min = rng.random_range(1..=c_max); + let degree = rng.random_range(1..=points.min(8)); + let values: Vec = (0..points * dimensions) + .map(|_| rng.random_range(-10.0..10.0)) + .collect(); + let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); + let graph = graph_config(Metric::L2, degree); + let pool = pool(2); + let config = PiPNNConfig { + c_max, + c_min, + p_samp: 0.5, + fanout: vec![2], + k: rng.random_range(1..=3), + replicas: rng.random_range(1..=2), + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context) + .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); + assert_graph_invariants(&actual, points, degree); + } +} diff --git a/diskann/tests/config.rs b/diskann/tests/config.rs new file mode 100644 index 0000000000..7d46e9a2a5 --- /dev/null +++ b/diskann/tests/config.rs @@ -0,0 +1,133 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +#![cfg(feature = "pipnn")] +#![allow( + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] + +use diskann::graph::config::{self, MaxDegree}; +use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig}; +use diskann_vector::distance::Metric; + +fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 512, + c_min: 64, + p_samp: 0.01, + fanout: vec![10, 3], + k: 2, + replicas: 1, + } +} + +fn graph_config(metric: Metric, alpha: f32) -> diskann::graph::Config { + config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { + builder.alpha(alpha); + }) + .build() + .unwrap() +} + +fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap() +} + +#[test] +fn accepts_the_six_algorithm_parameters_with_outer_graph_policy() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); +} + +#[test] +fn rejects_each_invalid_algorithm_parameter() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + let mut cases = [ + PiPNNConfig { + c_max: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 513, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 0.0, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: -0.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 1.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: f64::NAN, + ..pipnn_config() + }, + PiPNNConfig { + fanout: Vec::new(), + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![1, 0], + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![17], + ..pipnn_config() + }, + PiPNNConfig { + k: 0, + ..pipnn_config() + }, + PiPNNConfig { + replicas: 0, + ..pipnn_config() + }, + ]; + + for config in &mut cases { + let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .expect_err("invalid PiPNN config must be rejected"); + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + } +} + +#[test] +fn rejects_graph_policy_for_a_different_metric() { + let graph = graph_config(Metric::InnerProduct, 1.2); + let pool = pool(); + + let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + assert!(error.to_string().contains("prune kind")); +} + +#[test] +fn rejects_invalid_outer_alpha() { + let pool = pool(); + for alpha in [0.9, f32::NAN, f32::INFINITY] { + let graph = graph_config(Metric::L2, alpha); + let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + assert!(error.to_string().contains("alpha")); + } +} From fb6566c1af7f167715d6cb89d8c9384f50ea1af2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:27:00 +0000 Subject: [PATCH 02/58] pipnn: cover partition validation boundaries --- diskann/src/graph/pipnn/partitioning/tests.rs | 40 +++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 30eed3a2d1..13231db797 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -275,3 +275,43 @@ fn rejects_zero_dimensions() { PartitionError::EmptyDimensions ); } + +#[test] +fn rejects_invalid_gather_output_length() { + let data = Matrix::::new(0.0, 2, 2); + let error = gather_rows(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "gather output", + expected: 4, + actual: 3, + } + ); +} + +#[test] +fn rejects_assignment_to_an_unknown_leader() { + let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "leader assignment", + expected: 2, + actual: 3, + } + ); +} + +#[test] +fn rejects_empty_and_oversized_leaves() { + for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { + let error = validate_leaves(&leaves, 2).unwrap_err(); + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidLeaf { size, limit: 2 } + ); + } +} From 499f1f34a0fe430bbb822746eb365638cf7ba1cf Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 28 Jul 2026 16:16:50 +0000 Subject: [PATCH 03/58] pipnn: harden core graph construction --- diskann/src/graph/pipnn/leaf_build.rs | 26 +++++++----------- diskann/src/graph/pipnn/mod.rs | 3 +++ diskann/src/graph/pipnn/partitioning.rs | 21 ++++++--------- diskann/src/graph/pipnn/partitioning/tests.rs | 26 ++++++++++++++++++ diskann/src/graph/pipnn/tests.rs | 27 +++++++++++++++++++ 5 files changed, 74 insertions(+), 29 deletions(-) create mode 100644 diskann/src/graph/pipnn/tests.rs diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index f5fdb22928..91aaa50f62 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -123,13 +123,11 @@ impl LeafBuffers { nearest_values, LeafNeighbor::default(), )?; - if points > self.local_graph.len() { - let additional = points - self.local_graph.len(); - self.local_graph - .try_reserve(additional) - .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; - self.local_graph.resize_with(points, AdjacencyList::new); - } + let additional = points.saturating_sub(self.local_graph.len()); + self.local_graph + .try_reserve(additional) + .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; + self.local_graph.resize_with(points, AdjacencyList::new); self.local_graph[..points] .iter_mut() .for_each(AdjacencyList::clear); @@ -315,15 +313,11 @@ fn resize( len: usize, value: T, ) -> Result<(), LeafBuildError> { - if len > values.len() { - let additional = len - values.len(); - values - .try_reserve(additional) - .map_err(|source| allocation_error(buffer, additional, source))?; - values.resize(len, value); - } else { - values.truncate(len); - } + let additional = len.saturating_sub(values.len()); + values + .try_reserve(additional) + .map_err(|source| allocation_error(buffer, additional, source))?; + values.resize(len, value); Ok(()) } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 3d9680b6a1..7fddef6ae4 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -211,3 +211,6 @@ fn effective_metric(metric: Metric) -> Metric { fn config_error(message: impl std::fmt::Display) -> ANNError { ANNError::log_index_config_error("PiPNN".into(), message.to_string()) } + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 40e321c83d..d9384edfef 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -26,11 +26,11 @@ use crate::{ // Private algorithm and batching constants live together. None are user policy. const PARTITION_SEED: u64 = 1_000; const LEADER_CAP: usize = 1_000; -const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 512 * 1024; +const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; -const SCATTER_STRIPE_ROWS: usize = 64 * 1024; +const SCATTER_STRIPE_ROWS: usize = 65_536; const MAX_PARTITION_ITERATIONS: usize = 30; /// A partition failure with enough context to diagnose non-progressing input. @@ -182,12 +182,9 @@ where if work.is_empty() { return global_merge_small(leaves, config.c_min, config.c_max); } - let mut largest = &work[0]; - for item in &work[1..] { - if item.indices.len() > largest.indices.len() { - largest = item; - } - } + let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { + return global_merge_small(leaves, config.c_min, config.c_max); + }; Err(ANNError::opaque(PartitionError::IterationLimit { size: largest.indices.len(), level: largest.level, @@ -589,11 +586,9 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { } fn resize_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { - if len > values.len() { - values - .try_reserve(len - values.len()) - .map_err(ANNError::opaque)?; - } + values + .try_reserve(len.saturating_sub(values.len())) + .map_err(ANNError::opaque)?; values.resize(len, value); Ok(()) } diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 13231db797..d5a6fd81b8 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -154,6 +154,13 @@ fn global_merge_never_overfills_before_reaching_c_min() { ); } +#[test] +fn global_merge_fills_exact_capacity_before_flushing() { + let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + + assert_eq!(merged, vec![vec![0, 1, 2, 3]]); +} + #[test] fn replicas_cover_every_point_once_or_more_per_replica() { let data = directional_data(72, 5); @@ -240,6 +247,25 @@ fn leader_count_is_bounded() { assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); } +#[test] +fn replica_seed_derivation_is_stable_and_distinct() { + assert_eq!(mix_seed(PARTITION_SEED, 0), 9_518_416_997_697_480); + assert_eq!(mix_seed(PARTITION_SEED, 1), 9_518_416_997_697_481); +} + +#[test] +fn leader_assignment_handles_multiple_stripes() { + let points = 2_048; + let data: Vec = (0..points).map(|point| point as f32).collect(); + let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); + let point_ids: Vec = (0..points as u32).collect(); + + let clusters = assign_to_leaders(data, &point_ids, &[0, 2_047], 1, Metric::L2).unwrap(); + + assert_eq!(clusters[0], (0..1_024).collect::>()); + assert_eq!(clusters[1], (1_024..2_048).collect::>()); +} + #[test] fn parallel_scatter_matches_serial_order() { let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); diff --git a/diskann/src/graph/pipnn/tests.rs b/diskann/src/graph/pipnn/tests.rs new file mode 100644 index 0000000000..d5689eede6 --- /dev/null +++ b/diskann/src/graph/pipnn/tests.rs @@ -0,0 +1,27 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; +use half::f16; + +#[test] +fn integer_normalized_cosine_uses_unnormalized_cosine() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let expected = if metric == Metric::CosineNormalized { + Metric::Cosine + } else { + metric + }; + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), metric); + assert_eq!(effective_metric::(metric), metric); + } +} From f863ea26711815e7bb90dba6988865ad42b0487f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:36:30 +0000 Subject: [PATCH 04/58] docs(pipnn): describe integer cosine policy --- diskann/src/graph/pipnn/mod.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 7fddef6ae4..2242258515 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -143,6 +143,8 @@ impl<'a> PiPNNBuildContext<'a> { /// /// This is the core algorithm boundary. Search entry-point selection, frozen nodes, /// providers, serialization, and index writers belong to the outer build pipelines. +/// For raw `u8` and `i8` rows, `CosineNormalized` is evaluated as `Cosine` because +/// those representations are converted to f32 scratch but are not unit-normalized. pub fn build_graph( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, From e0b67b3532fc86f43f97f6acc260cb5cbdcecb6d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:39:54 +0000 Subject: [PATCH 05/58] pipnn: expose core config validation --- diskann/src/graph/pipnn/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 2242258515..3d3ac733fa 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -57,7 +57,8 @@ pub struct PiPNNConfig { } impl PiPNNConfig { - fn validate(&self) -> ANNResult<()> { + /// Validate the algorithm-specific partition and leaf-build parameters. + pub fn validate(&self) -> ANNResult<()> { if self.c_max == 0 { return Err(config_error("c_max must be greater than zero")); } From fcbb20f9edd44649327e78de97d01e141b671f0e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 12:23:36 +0000 Subject: [PATCH 06/58] refactor(pipnn): consume leaves during leaf build --- diskann/src/graph/pipnn/leaf_build.rs | 6 +++--- diskann/src/graph/pipnn/leaf_build/tests.rs | 8 +++++--- diskann/src/graph/pipnn/mod.rs | 2 +- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 91aaa50f62..31f1f9e8bd 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -180,7 +180,7 @@ impl DirectCandidates { #[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. pub(crate) fn build_leaf_candidates( data: MatrixView<'_, T>, - leaves: &[Vec], + leaves: Vec>, k: usize, metric: Metric, ) -> Result>, LeafBuildError> @@ -195,10 +195,10 @@ where } let candidates = DirectCandidates::new(data.nrows())?; - leaves.par_iter().enumerate().try_for_each_init( + leaves.into_par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) + build_leaf(data, leaf, &point_ids, k, metric, buffers, &candidates) }, )?; candidates.into_rows() diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index d3f63263dc..efa597dd72 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -33,7 +33,7 @@ fn build( where T: diskann::utils::VectorRepr + 'static, { - pool().install(|| build_leaf_candidates(data, leaves, k, metric)) + pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) } fn rows(graph: Vec>) -> Vec> { @@ -204,9 +204,11 @@ fn parallel_leaf_schedule_does_not_change_candidate_order() { .collect(); let pool = pool(); pool.install(|| { - let expected = build_leaf_candidates(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + let expected = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); for _ in 0..8 { - let actual = build_leaf_candidates(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + let actual = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); assert_eq!(actual, expected); } }); diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 3d3ac733fa..64888f8532 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -191,7 +191,7 @@ where let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition(data, &context.config, metric))?; let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates(data, &leaves, context.config.k, metric) + leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) .map_err(ANNError::opaque) })?; tracing::info_span!("pipnn.finalization") From b0164cbd6ec5c6d28bc080704ada4d4c1caad447 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 12:23:36 +0000 Subject: [PATCH 07/58] fix(pipnn): preserve established replica seeds --- diskann/src/graph/pipnn/partitioning.rs | 11 ++++++++--- diskann/src/graph/pipnn/partitioning/tests.rs | 4 ++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index d9384edfef..6c38e80247 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -25,6 +25,7 @@ use crate::{ // Private algorithm and batching constants live together. None are user policy. const PARTITION_SEED: u64 = 1_000; +const REPLICA_SEED_STEP: u64 = 7_919; const LEADER_CAP: usize = 1_000; const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; @@ -106,7 +107,7 @@ where let mut leaves = Vec::new(); for replica in 0..config.replicas { - let seed = mix_seed(PARTITION_SEED, replica as u64); + let seed = replica_seed(replica); let mut replica_leaves = partition_replica(data, config, metric, seed)?; leaves .try_reserve(replica_leaves.len()) @@ -251,8 +252,12 @@ fn sample_num_leaders(points: usize, sampling_fraction: f64) -> usize { .min(points) } -// A single LCG mixer derives both replica and recursive seeds. Wrapping makes -// the mapping stable across debug/release builds and supported platforms. +fn replica_seed(replica: usize) -> u64 { + PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) +} + +// A single LCG mixer derives recursive seeds. Wrapping makes the mapping stable +// across debug/release builds and supported platforms. fn mix_seed(seed: u64, salt: u64) -> u64 { seed.wrapping_mul(6_364_136_223_846_793_005) .wrapping_add(salt) diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index d5a6fd81b8..56ec5f8f2e 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -249,8 +249,8 @@ fn leader_count_is_bounded() { #[test] fn replica_seed_derivation_is_stable_and_distinct() { - assert_eq!(mix_seed(PARTITION_SEED, 0), 9_518_416_997_697_480); - assert_eq!(mix_seed(PARTITION_SEED, 1), 9_518_416_997_697_481); + assert_eq!(replica_seed(0), 1_000); + assert_eq!(replica_seed(1), 8_919); } #[test] From 4e9c6fe12232c65dcfa3cbe0cd17eaae89317636 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:25:52 +0000 Subject: [PATCH 08/58] docs(pipnn): define graph-construction boundary --- diskann-pipnn/README.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 diskann-pipnn/README.md diff --git a/diskann-pipnn/README.md b/diskann-pipnn/README.md new file mode 100644 index 0000000000..4ac1618f03 --- /dev/null +++ b/diskann-pipnn/README.md @@ -0,0 +1,19 @@ +# PiPNN graph construction + +This crate implements the graph-construction stages from [PiPNN: Pick in Partitions for Fast and Accurate ANN Graph Construction](https://arxiv.org/html/2602.21247v1). + +## Boundary + +PiPNN core consumes a dense `MatrixView`, graph policy, and a caller-owned Rayon pool, then returns adjacency lists for the dataset's real point IDs. It does not own start or frozen points, vector or neighbor providers, PQ, disk headers, serialization, or search. Those concerns remain in the outer in-memory and disk pipelines. + +The dense view is intentional: partition assignment and leaf all-pairs kernels operate over the whole source matrix. Materializing provider state inside the algorithm would couple numerical graph construction to storage lifecycle and would require a second dataset copy. Integrations should finish PiPNN scratch before allocating or populating their searchable provider. + +## Policy ownership + +- `PiPNNConfig` owns partition and leaf-selection parameters: leaf bounds, sampling fraction, fanout levels, leaf `k`, and replicas. +- DiskANN graph configuration owns metric, output degree, build-L, alpha, and prune policy. +- Candidate-merging policies are separate validated options; they must not make graph policy fields redundant or silently cap the requested degree. + +## Execution + +A build runs partitioning, leaf construction, candidate merging, then graph finalization. All parallel work executes in the supplied pool. Per-job scratch is initialized through Rayon and is released through normal ownership when its stage completes; the core has no global thread-local buffers or cleanup broadcasts. From dad9cf5b5aa84386a975df25c6a65ddf5b72f5e6 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:10:42 +0000 Subject: [PATCH 09/58] refactor(pipnn): own partition-stage configuration --- diskann/src/graph/pipnn/mod.rs | 3 +- diskann/src/graph/pipnn/partitioning.rs | 31 ++++++++++++++++--- diskann/src/graph/pipnn/partitioning/tests.rs | 31 +++++++++---------- 3 files changed, 44 insertions(+), 21 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 64888f8532..48b07b22f1 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -188,8 +188,9 @@ where })?; let metric = effective_metric::(context.metric); + let partition = partitioning::PartitionConfig::from(&context.config); let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition(data, &context.config, metric))?; + .in_scope(|| partitioning::partition(data, partition, metric))?; let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) .map_err(ANNError::opaque) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 6c38e80247..bbe4a6a92b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -34,6 +34,29 @@ const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const SCATTER_STRIPE_ROWS: usize = 65_536; const MAX_PARTITION_ITERATIONS: usize = 30; +/// Policy owned by the partition stage. Leaf-neighbor and merge settings do +/// not cross this boundary. +#[derive(Clone, Debug)] +pub(crate) struct PartitionConfig { + c_max: usize, + c_min: usize, + p_samp: f64, + fanout: Vec, + replicas: usize, +} + +impl From<&PiPNNConfig> for PartitionConfig { + fn from(config: &PiPNNConfig) -> Self { + Self { + c_max: config.c_max, + c_min: config.c_min, + p_samp: config.p_samp, + fanout: config.fanout.clone(), + replicas: config.replicas, + } + } +} + /// A partition failure with enough context to diagnose non-progressing input. #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum PartitionError { @@ -88,7 +111,7 @@ struct StripeBuffers { /// its build-owned Rayon pool. pub(crate) fn partition( data: MatrixView<'_, T>, - config: &PiPNNConfig, + config: PartitionConfig, metric: Metric, ) -> ANNResult>> where @@ -108,7 +131,7 @@ where let mut leaves = Vec::new(); for replica in 0..config.replicas { let seed = replica_seed(replica); - let mut replica_leaves = partition_replica(data, config, metric, seed)?; + let mut replica_leaves = partition_replica(data, &config, metric, seed)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::opaque)?; @@ -120,7 +143,7 @@ where fn partition_replica( data: MatrixView<'_, T>, - config: &PiPNNConfig, + config: &PartitionConfig, metric: Metric, seed: u64, ) -> ANNResult>> @@ -195,7 +218,7 @@ where fn partition_one_level( data: MatrixView<'_, T>, - config: &PiPNNConfig, + config: &PartitionConfig, metric: Metric, item: WorkItem, ) -> ANNResult<(Vec, Vec>)> diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 56ec5f8f2e..ef7b4b517e 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -8,13 +8,12 @@ use diskann_vector::{distance::Metric, Half}; use super::*; -fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { - PiPNNConfig { +fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PartitionConfig { + PartitionConfig { c_max, c_min, p_samp: 0.25, fanout, - k: 2, replicas, } } @@ -91,7 +90,7 @@ fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, repl fn returns_one_leaf_at_and_below_c_max() { for points in [7, 8] { let data = clustered_data(points, 3); - let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); assert_eq!(leaves, vec![(0..points as u32).collect::>()]); } } @@ -101,8 +100,8 @@ fn partition_is_fixed_seed_deterministic_and_bounded() { let data = clustered_data(96, 8); let config = config(4, 16, vec![3, 2], 2); - let first = partition(data.as_view(), &config, Metric::L2).unwrap(); - let second = partition(data.as_view(), &config, Metric::L2).unwrap(); + let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); + let second = partition(data.as_view(), config, Metric::L2).unwrap(); assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); assert_valid_partition(&first, 96, 16, 2); @@ -112,7 +111,7 @@ fn partition_is_fixed_seed_deterministic_and_bounded() { #[test] fn recursion_after_fanout_levels_falls_back_to_one() { let data = clustered_data(80, 4); - let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); assert_valid_partition(&leaves, 80, 8, 1); } @@ -120,7 +119,7 @@ fn recursion_after_fanout_levels_falls_back_to_one() { #[test] fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { let data = Matrix::new(1.0f32, 24, 4); - let error = partition(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); let error = error.downcast::().unwrap(); assert!(matches!( @@ -166,7 +165,7 @@ fn replicas_cover_every_point_once_or_more_per_replica() { let data = directional_data(72, 5); let leaves = partition( data.as_view(), - &config(3, 12, vec![3, 2], 3), + config(3, 12, vec![3, 2], 3), Metric::CosineNormalized, ) .unwrap(); @@ -184,25 +183,25 @@ fn supported_source_types_share_partition_contract() { let f32_leaves = partition( MatrixView::try_from(f32_data.as_slice(), 64, 4).unwrap(), - &config, + config.clone(), Metric::L2, ) .unwrap(); let half_leaves = partition( MatrixView::try_from(half_data.as_slice(), 64, 4).unwrap(), - &config, + config.clone(), Metric::L2, ) .unwrap(); let u8_leaves = partition( MatrixView::try_from(u8_data.as_slice(), 64, 4).unwrap(), - &config, + config.clone(), Metric::L2, ) .unwrap(); let i8_leaves = partition( MatrixView::try_from(i8_data.as_slice(), 64, 4).unwrap(), - &config, + config, Metric::L2, ) .unwrap(); @@ -235,7 +234,7 @@ fn all_metrics_produce_valid_partitions() { Metric::CosineNormalized, Metric::InnerProduct, ] { - let leaves = partition(data.as_view(), &config, metric).unwrap(); + let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); assert_valid_partition(&leaves, 64, 20, 1); } } @@ -283,7 +282,7 @@ fn parallel_scatter_matches_serial_order() { #[test] fn rejects_empty_dataset() { let data = Matrix::::new(0.0, 0, 4); - let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); assert_eq!( error.downcast::().unwrap(), @@ -294,7 +293,7 @@ fn rejects_empty_dataset() { #[test] fn rejects_zero_dimensions() { let data = Matrix::::new(0.0, 4, 0); - let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); assert_eq!( error.downcast::().unwrap(), From 5b031c53c5c71776a87dffeb0df8d41e5a804ac2 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:43:33 +0000 Subject: [PATCH 10/58] perf(pipnn): release owned leaves after leaf stage --- diskann/src/graph/pipnn/leaf_build.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 31f1f9e8bd..a0babd3e89 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -195,10 +195,10 @@ where } let candidates = DirectCandidates::new(data.nrows())?; - leaves.into_par_iter().enumerate().try_for_each_init( + leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf(data, leaf, &point_ids, k, metric, buffers, &candidates) + build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) }, )?; candidates.into_rows() From 9755ef99d0b28b0b8325e68e14877144479b55de Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:46:11 +0000 Subject: [PATCH 11/58] docs(pipnn): explain partition stage contract --- diskann/src/graph/pipnn/partitioning.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index bbe4a6a92b..64b8eaeea2 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -106,9 +106,14 @@ struct StripeBuffers { row_scales: Vec, } -/// Partition every configured replica of `data` into leaves no larger than -/// `config.c_max`. The caller is responsible for installing this operation in -/// its build-owned Rayon pool. +/// Partition every configured replica into overlapping bounded leaves. +/// +/// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped +/// to the private leader bound), assigns every point to its nearest `fanout` +/// leaders for the current level, and recurses only on oversized clusters. +/// Levels beyond `fanout.len()` retain one leader assignment. Completed small +/// leaves are merged without exceeding `c_max`; every input point must remain +/// covered once per replica. The caller installs the operation in its pool. pub(crate) fn partition( data: MatrixView<'_, T>, config: PartitionConfig, From 68d26f2053fef461b546e2e5ea81d968b35e54e9 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 30 Jul 2026 07:53:25 +0000 Subject: [PATCH 12/58] perf(pipnn): reduce partition and leaf overhead --- diskann/src/graph/pipnn/leaf_build.rs | 67 ++++++++++++++----- diskann/src/graph/pipnn/leaf_build/tests.rs | 7 ++ diskann/src/graph/pipnn/partitioning.rs | 39 +++++++---- diskann/src/graph/pipnn/partitioning/tests.rs | 7 ++ 4 files changed, 88 insertions(+), 32 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index a0babd3e89..6d8a6ffa3a 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -115,14 +115,18 @@ impl LeafBuffers { columns: actual_k, })?; - resize("leaf points", &mut self.points, point_values, 0.0)?; - resize("leaf dot products", &mut self.dots, dot_values, 0.0)?; - resize( + grow("leaf points", &mut self.points, point_values, 0.0)?; + grow("leaf dot products", &mut self.dots, dot_values, 0.0)?; + grow( "leaf nearest neighbors", &mut self.nearest, nearest_values, LeafNeighbor::default(), )?; + Ok(actual_k) + } + + fn prepare_local_graph(&mut self, points: usize) -> Result<(), LeafBuildError> { let additional = points.saturating_sub(self.local_graph.len()); self.local_graph .try_reserve(additional) @@ -131,7 +135,7 @@ impl LeafBuffers { self.local_graph[..points] .iter_mut() .for_each(AdjacencyList::clear); - Ok(actual_k) + Ok(()) } } @@ -219,11 +223,6 @@ where if point_ids.is_empty() { return Err(LeafBuildError::EmptyLeaf { leaf }); } - buffers.seen_ids.clear(); - buffers - .seen_ids - .try_reserve(point_ids.len()) - .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; for &point in point_ids { if point as usize >= data.nrows() { return Err(LeafBuildError::InvalidPointId { @@ -232,8 +231,24 @@ where points: data.nrows(), }); } - if !buffers.seen_ids.insert(point) { - return Err(LeafBuildError::DuplicatePointId { leaf, point }); + } + if point_ids.is_sorted() { + if let Some(pair) = point_ids.windows(2).find(|pair| pair[0] == pair[1]) { + return Err(LeafBuildError::DuplicatePointId { + leaf, + point: pair[0], + }); + } + } else { + buffers.seen_ids.clear(); + buffers + .seen_ids + .try_reserve(point_ids.len()) + .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; + for &point in point_ids { + if !buffers.seen_ids.insert(point) { + return Err(LeafBuildError::DuplicatePointId { leaf, point }); + } } } let actual_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), k)?; @@ -241,11 +256,14 @@ where return Ok(()); } + let point_values = point_ids.len() * data.ncols(); + let dot_values = point_ids.len() * point_ids.len(); + let nearest_values = point_ids.len() * actual_k; + for (&point, output) in point_ids .iter() - .zip(buffers.points.chunks_exact_mut(data.ncols())) + .zip(buffers.points[..point_values].chunks_exact_mut(data.ncols())) { - // Point IDs were validated before the zero-k/singleton fast path. let row = data.row(point as usize); T::as_f32_into(row, output).map_err(|source| LeafBuildError::Conversion { leaf, @@ -255,28 +273,29 @@ where } diskann_linalg::sgemm_aat_lower( - &buffers.points, + &buffers.points[..point_values], point_ids.len(), data.ncols(), - &mut buffers.dots, + &mut buffers.dots[..dot_values], ) .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; nearest_leaf_neighbors( LeafTopK { - dots: &buffers.dots, + dots: &buffers.dots[..dot_values], points: point_ids.len(), metric, }, k, - &mut buffers.nearest, + &mut buffers.nearest[..nearest_values], &mut buffers.top_k, ) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + buffers.prepare_local_graph(point_ids.len())?; add_symmetric_edges( point_ids, actual_k, - &buffers.nearest, + &buffers.nearest[..nearest_values], &mut buffers.local_graph[..point_ids.len()], )?; candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) @@ -307,6 +326,18 @@ fn add_symmetric_edges( Ok(()) } +fn grow( + buffer: &'static str, + values: &mut Vec, + len: usize, + value: T, +) -> Result<(), LeafBuildError> { + if values.len() < len { + resize(buffer, values, len, value)?; + } + Ok(()) +} + fn resize( buffer: &'static str, values: &mut Vec, diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index efa597dd72..8f91ea559c 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -246,6 +246,10 @@ fn rejects_invalid_shape_inputs_without_panicking() { build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![1, 0, 1]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 1 }) + )); } #[test] @@ -275,6 +279,9 @@ fn reuses_worker_buffers_for_smaller_leaves() { assert_eq!(buffers.points.as_ptr(), points); assert_eq!(buffers.dots.as_ptr(), dots); assert_eq!(buffers.nearest.as_ptr(), nearest); + assert_eq!(buffers.points.len(), 64 * 128); + assert_eq!(buffers.dots.len(), 64 * 64); + assert_eq!(buffers.nearest.len(), 64 * 2); } #[test] diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 64b8eaeea2..122261dc91 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -31,7 +31,6 @@ const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; -const SCATTER_STRIPE_ROWS: usize = 65_536; const MAX_PARTITION_ITERATIONS: usize = 30; /// Policy owned by the partition stage. Leaf-neighbor and merge settings do @@ -413,8 +412,9 @@ fn scatter_assignments( return scatter_serial(points, assignments, fanout, leaders); } - let assignment_stripe = checked_area("scatter assignment stripe", SCATTER_STRIPE_ROWS, fanout)?; - let stripes = points.len().div_ceil(SCATTER_STRIPE_ROWS); + let stripe_rows = points.len().div_ceil(rayon::current_num_threads().max(1)); + let assignment_stripe = checked_area("scatter assignment stripe", stripe_rows, fanout)?; + let stripes = points.len().div_ceil(stripe_rows); let mut partials = Vec::new(); partials .try_reserve_exact(stripes) @@ -426,7 +426,7 @@ fn scatter_assignments( .par_iter_mut() .zip( points - .par_chunks(SCATTER_STRIPE_ROWS) + .par_chunks(stripe_rows) .zip(assignments.par_chunks(assignment_stripe)), ) .for_each(|(slot, (points, assignments))| { @@ -454,14 +454,20 @@ fn scatter_assignments( } } - let mut clusters = clusters_with_capacities(&sizes)?; - for local in locals { - for (cluster, part) in clusters.iter_mut().zip(local) { - debug_assert!(cluster.capacity().saturating_sub(cluster.len()) >= part.len()); - cluster.extend(part); - } - } - Ok(clusters) + // See the pool invariant at the other partition terminal operations. + #[allow(clippy::disallowed_methods)] + sizes + .into_par_iter() + .enumerate() + .map(|(leader, size)| { + let mut cluster = Vec::new(); + cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + for local in &locals { + cluster.extend_from_slice(&local[leader]); + } + Ok(cluster) + }) + .collect() } fn scatter_serial( @@ -632,8 +638,13 @@ fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult usize { - (ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::())) - .clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) + let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::()); + let rows = if rows.is_power_of_two() { + rows + } else { + rows.next_power_of_two() / 2 + }; + rows.clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) } #[cfg(test)] diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index ef7b4b517e..b9714c7f64 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -252,6 +252,13 @@ fn replica_seed_derivation_is_stable_and_distinct() { assert_eq!(replica_seed(1), 8_919); } +#[test] +fn assignment_stripes_use_power_of_two_row_counts() { + assert_eq!(assignment_stripe_rows(1_000), 128); + assert_eq!(assignment_stripe_rows(256), 512); + assert_eq!(assignment_stripe_rows(1), MAX_ASSIGNMENT_STRIPE_ROWS); +} + #[test] fn leader_assignment_handles_multiple_stripes() { let points = 2_048; From 112a03c51666fa56e49338f26c0520378d669e51 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 30 Jul 2026 13:25:27 +0000 Subject: [PATCH 13/58] perf(pipnn): reuse partition scratch across work items --- diskann/src/graph/pipnn/partitioning.rs | 192 +++++++++++++----- diskann/src/graph/pipnn/partitioning/tests.rs | 23 ++- 2 files changed, 160 insertions(+), 55 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 122261dc91..2d9c37d29e 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -6,10 +6,11 @@ //! Deterministic overlapping partition construction for PiPNN. //! //! The stage maps real dataset rows to bounded leaf ID lists. Numerical work -//! reuses the partition kernel and dense GEMM; scratch belongs to the Rayon -//! iterator that uses it, so no thread-local cleanup protocol is required. +//! reuses the partition kernel and dense GEMM. A stage-owned pool leases scratch +//! to Rayon chunks and takes it back after each chunk; computation never holds +//! the pool lock, and no thread-local cleanup protocol is required. -use std::collections::HashSet; +use std::{collections::HashSet, sync::Mutex}; use crate::{utils::VectorRepr, ANNError, ANNResult}; use diskann_linalg::Transpose; @@ -105,6 +106,28 @@ struct StripeBuffers { row_scales: Vec, } +#[derive(Default)] +struct StripeBufferPool { + available: Mutex>, +} + +impl StripeBufferPool { + fn take(&self) -> StripeBuffers { + self.available + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .pop() + .unwrap_or_default() + } + + fn put(&self, buffers: StripeBuffers) { + self.available + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .push(buffers); + } +} + /// Partition every configured replica into overlapping bounded leaves. /// /// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped @@ -133,9 +156,10 @@ where } let mut leaves = Vec::new(); + let stripe_buffers = StripeBufferPool::default(); for replica in 0..config.replicas { let seed = replica_seed(replica); - let mut replica_leaves = partition_replica(data, &config, metric, seed)?; + let mut replica_leaves = partition_replica(data, &config, metric, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::opaque)?; @@ -150,6 +174,7 @@ fn partition_replica( config: &PartitionConfig, metric: Metric, seed: u64, + stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where T: VectorRepr + Send + Sync, @@ -188,7 +213,13 @@ where .par_iter_mut() .zip(work.into_par_iter()) .for_each(|(slot, item)| { - *slot = Some(partition_one_level(data, config, metric, item)); + *slot = Some(partition_one_level( + data, + config, + metric, + item, + stripe_buffers, + )); }); let mut next_work = Vec::new(); @@ -225,6 +256,7 @@ fn partition_one_level( config: &PartitionConfig, metric: Metric, item: WorkItem, + stripe_buffers: &StripeBufferPool, ) -> ANNResult<(Vec, Vec>)> where T: VectorRepr + Send + Sync, @@ -236,7 +268,14 @@ where config.p_samp, mix_seed(item.seed, points as u64), )?; - let clusters = assign_to_leaders(data, &item.indices, &leaders, fanout, metric)?; + let clusters = assign_to_leaders( + data, + &item.indices, + &leaders, + fanout, + metric, + stripe_buffers, + )?; let mut pending = Vec::new(); let mut finished = Vec::new(); @@ -296,6 +335,7 @@ fn assign_to_leaders( leaders: &[u32], fanout: usize, metric: Metric, + stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where T: VectorRepr + Send + Sync, @@ -325,65 +365,109 @@ where let mut assignments = filled_vec(assignment_len, 0u32)?; let stripe_rows = assignment_stripe_rows(leaders.len()); let assignment_stripe = checked_area("assignment stripe", stripe_rows, fanout)?; + let stripes = points.len().div_ceil(stripe_rows); + let worker_stripes = stripes.div_ceil(rayon::current_num_threads().max(1)); + let worker_rows = checked_area("assignment worker", worker_stripes, stripe_rows)?; + let worker_assignment = checked_area("assignment worker", worker_rows, fanout)?; + // Each worker chunk owns one scratch value and reuses it for its stripes. // build_graph pins this terminal operation to the caller-owned pool. #[allow(clippy::disallowed_methods)] assignments - .par_chunks_mut(assignment_stripe) + .par_chunks_mut(worker_assignment) .enumerate() - .try_for_each_init(StripeBuffers::default, |buffers, (stripe, output)| { - let first = stripe * stripe_rows; - let rows = output.len() / fanout; - let point_values_len = checked_area("point stripe", rows, dimensions)?; - let dots_len = checked_area("dot-product stripe", rows, leaders.len())?; - resize_fallible(&mut buffers.points, point_values_len, 0.0)?; - resize_fallible(&mut buffers.dots, dots_len, 0.0)?; - gather_rows(data, &points[first..first + rows], &mut buffers.points)?; - diskann_linalg::sgemm( - Transpose::None, - Transpose::Ordinary, - rows, - leaders.len(), - dimensions, - 1.0, - &buffers.points, - &leader_values, - None, - &mut buffers.dots, - ) - .map_err(ANNError::opaque)?; - - let row_scales = if metric == Metric::Cosine { - resize_fallible(&mut buffers.row_scales, rows, 0.0)?; - for (scale, row) in buffers - .row_scales - .iter_mut() - .zip(buffers.points.chunks_exact(dimensions)) - { - *scale = FastL2NormSquared.evaluate(row); + .try_for_each(|(worker, worker_output)| { + let mut buffers = stripe_buffers.take(); + let result = (|| -> ANNResult<()> { + let worker_first = worker * worker_rows; + for (stripe, output) in worker_output.chunks_mut(assignment_stripe).enumerate() { + let first = worker_first + stripe * stripe_rows; + let rows = output.len() / fanout; + assign_stripe( + data, + &points[first..first + rows], + &leader_values, + &leader_scales, + metric, + fanout, + &mut buffers, + output, + )?; } - buffers.row_scales.as_slice() - } else { - &[] - }; - nearest_leaders( - PartitionTopK { - dots: &buffers.dots, - rows, - leaders: leaders.len(), - row_scales, - leader_scales: &leader_scales, - metric, - }, - fanout, - output, - ) - .map_err(ANNError::opaque) + Ok(()) + })(); + stripe_buffers.put(buffers); + result })?; scatter_assignments(points, &assignments, fanout, leaders.len()) } +#[inline] +#[allow(clippy::too_many_arguments)] +fn assign_stripe( + data: MatrixView<'_, T>, + points: &[u32], + leader_values: &[f32], + leader_scales: &[f32], + metric: Metric, + fanout: usize, + buffers: &mut StripeBuffers, + output: &mut [u32], +) -> ANNResult<()> +where + T: VectorRepr, +{ + let rows = points.len(); + let dimensions = data.ncols(); + let leaders = leader_values.len() / dimensions; + let point_values_len = checked_area("point stripe", rows, dimensions)?; + let dots_len = checked_area("dot-product stripe", rows, leaders)?; + resize_fallible(&mut buffers.points, point_values_len, 0.0)?; + resize_fallible(&mut buffers.dots, dots_len, 0.0)?; + gather_rows(data, points, &mut buffers.points)?; + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + rows, + leaders, + dimensions, + 1.0, + &buffers.points, + leader_values, + None, + &mut buffers.dots, + ) + .map_err(ANNError::opaque)?; + + let row_scales = if metric == Metric::Cosine { + resize_fallible(&mut buffers.row_scales, rows, 0.0)?; + for (scale, row) in buffers + .row_scales + .iter_mut() + .zip(buffers.points.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(row); + } + buffers.row_scales.as_slice() + } else { + &[] + }; + nearest_leaders( + PartitionTopK { + dots: &buffers.dots, + rows, + leaders, + row_scales, + leader_scales, + metric, + }, + fanout, + output, + ) + .map_err(ANNError::opaque) +} + fn gather_rows(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> where T: VectorRepr, diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index b9714c7f64..d3815d533e 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -259,6 +259,19 @@ fn assignment_stripes_use_power_of_two_row_counts() { assert_eq!(assignment_stripe_rows(1), MAX_ASSIGNMENT_STRIPE_ROWS); } +#[test] +fn stripe_buffer_pool_reuses_returned_capacity() { + let pool = StripeBufferPool::default(); + let mut buffers = pool.take(); + buffers.points.resize(16, 0.0); + let points = buffers.points.as_ptr(); + pool.put(buffers); + + let buffers = pool.take(); + assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.points.len(), 16); +} + #[test] fn leader_assignment_handles_multiple_stripes() { let points = 2_048; @@ -266,7 +279,15 @@ fn leader_assignment_handles_multiple_stripes() { let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); let point_ids: Vec = (0..points as u32).collect(); - let clusters = assign_to_leaders(data, &point_ids, &[0, 2_047], 1, Metric::L2).unwrap(); + let clusters = assign_to_leaders( + data, + &point_ids, + &[0, 2_047], + 1, + Metric::L2, + &StripeBufferPool::default(), + ) + .unwrap(); assert_eq!(clusters[0], (0..1_024).collect::>()); assert_eq!(clusters[1], (1_024..2_048).collect::>()); From 30bbba416ccf8134a8805ef0856a048351f2085c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:06:21 +0000 Subject: [PATCH 14/58] docs(pipnn): document core stage invariants --- diskann/src/graph/pipnn/finalization.rs | 24 ++- diskann/src/graph/pipnn/leaf_build.rs | 32 +++- diskann/src/graph/pipnn/leaf_build/tests.rs | 61 ++++--- diskann/src/graph/pipnn/mod.rs | 149 ++++++++++++++++-- diskann/src/graph/pipnn/partitioning.rs | 20 +++ diskann/src/graph/pipnn/partitioning/tests.rs | 112 +++++++------ 6 files changed, 319 insertions(+), 79 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index dab9419cc8..5cae12978a 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -3,7 +3,17 @@ * Licensed under the MIT license. */ -//! Orders complete PiPNN candidate rows and applies shared Vamana RobustPrune. +//! Final graph-degree enforcement through the shared Vamana RobustPrune kernel. +//! +//! Candidate merging may produce more than `R` IDs for a point. This stage first +//! validates every global ID, then processes rows independently in the caller's +//! Rayon pool. Rows already within the degree bound are returned without any +//! distance work. Overfull rows are converted to source-distance candidates, +//! passed through RobustPrune, and rewritten from the selected output. +//! +//! The shared kernel owns occlusion and alpha-round semantics; this adapter owns +//! only contiguous dataset access and distance specialization for the source +//! representation. use std::convert::Infallible; @@ -29,12 +39,19 @@ pub(crate) enum FinalizationError { }, } +/// Per-Rayon-job state retained across rows. +/// +/// `prune` owns candidate/state/output buffers. `cache` stores provider lookup +/// results required by the shared kernel. Reusing both avoids per-node +/// allocations, which would otherwise dominate finalization for millions of +/// short rows. #[derive(Default)] struct Workspace { prune: prune::Scratch, cache: Vec<(f32, Option)>, } +/// Validate candidate IDs and prune only rows whose length exceeds graph degree. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, @@ -56,6 +73,8 @@ where .into_par_iter() .enumerate() .map_init(Workspace::default, |workspace, (source, mut row)| { + // Candidate accumulators already enforce uniqueness. A bounded row + // therefore satisfies the graph policy without distance evaluation. if row.len() <= degree { return Ok(row); } @@ -71,6 +90,9 @@ where distance.evaluate_similarity(source_vector, data.row(candidate as usize)), ) })); + // as_context sorts the active candidate prefix by source distance. + // The callback below is needed only for selected-to-candidate + // occlusion checks; dimension specialization stays in `distance`. let candidate_count = pool.len(); let mut context = workspace.prune.as_context(candidate_count); prune::robust_prune( diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6d8a6ffa3a..f847077ef2 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -3,7 +3,19 @@ * Licensed under the MIT license. */ -//! Leaf construction and direct candidate accumulation. +//! Leaf-local graph construction and candidate accumulation. +//! +//! Partitioning supplies leaves as global point IDs. For each leaf this module: +//! +//! 1. validates IDs and converts only those rows to reusable `f32` scratch; +//! 2. computes the lower triangle of `A · Aᵀ`; +//! 3. runs the dual-endpoint leaf top-k kernel; and +//! 4. translates leaf-local positions back to dataset IDs. +//! +//! The final step merges symmetric adjacency rows under per-point locks because +//! overlapping leaves are processed concurrently. Numeric buffers retain their +//! high-water length; every consumer therefore receives an explicit active +//! prefix rather than treating `Vec::len()` as the current leaf shape. use std::{ collections::{HashSet, TryReserveError}, @@ -74,6 +86,12 @@ pub(crate) enum LeafBuildError { PoisonedCandidateRow { point: u32 }, } +/// Scratch leased to one Rayon job and reused for successive leaves. +/// +/// The three numerical vectors retain their largest observed leaf shape. The +/// adjacency rows are prepared separately because zero-k/singleton leaves never +/// write them, and because later candidate-merging modes do not necessarily use +/// this representation. #[derive(Default)] struct LeafBuffers { points: Vec, @@ -139,6 +157,12 @@ impl LeafBuffers { } } +/// Concurrent accumulator indexed by global dataset ID. +/// +/// A point may appear in several overlapping leaves, so workers lock only the +/// destination row long enough to append one leaf's additions. Sorting and +/// duplicate removal are deferred until all leaves finish; doing either under +/// the lock would lengthen the contended section for no semantic benefit. struct DirectCandidates { rows: Vec>>, } @@ -208,6 +232,12 @@ where candidates.into_rows() } +/// Build and publish one leaf's symmetric nearest-neighbor rows. +/// +/// Validation precedes all dataset indexing. Sorted partition output takes the +/// adjacent-duplicate path, while arbitrary-order callers use `seen_ids`. The +/// active lengths computed after `prepare` must be used for every later slice, +/// because the reusable vectors may still be longer than this leaf. fn build_leaf( data: MatrixView<'_, T>, leaf: usize, diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index 8f91ea559c..0236370dee 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -148,30 +148,55 @@ fn global_id_translation_is_independent_of_leaf_order() { ); } -fn assert_source_type(data: &[T]) +fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> where T: diskann::utils::VectorRepr + 'static, { - let leaves = vec![vec![0, 1, 2, 3]]; - let graph = build(view(data, 4, 2), &leaves, 1, Metric::L2).unwrap(); - assert_eq!(rows(graph), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); + let leaves = vec![(0..points as u32).collect()]; + rows(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) +} + +fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) +where + T: diskann::utils::VectorRepr + 'static, +{ + let points = 8; + // Source dimension controls VectorRepr conversion chunking. Cover tails on + // both sides of 4-, 8-, and 16-element boundaries, then a second 16-lane + // chunk. Input integers remain exact in every tested representation. + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let row = index / dimensions; + let column = index % dimensions; + ((row * 7 + column * 3 + row * column) % 23) as u8 + }) + .collect(); + let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + assert_eq!( + source_graph(&converted, points, dimensions), + source_graph(&f32_data, points, dimensions), + "{label} dimensions={dimensions}" + ); + } +} + +#[test] +fn f16_conversion_matches_f32_across_dimension_boundaries() { + assert_source_conversion_matches_f32("f16", |value| f16::from_f32(value as f32)); +} + +#[test] +fn u8_conversion_matches_f32_across_dimension_boundaries() { + assert_source_conversion_matches_f32("u8", |value| value); } #[test] -fn gathers_every_supported_source_type_without_full_dataset_conversion() { - assert_source_type(&[0.0_f32, 0.0, 1.0, 0.0, 2.0, 0.0, 3.0, 0.0]); - assert_source_type(&[0_i8, 0, 1, 0, 2, 0, 3, 0]); - assert_source_type(&[0_u8, 0, 1, 0, 2, 0, 3, 0]); - assert_source_type(&[ - f16::from_f32(0.0), - f16::from_f32(0.0), - f16::from_f32(1.0), - f16::from_f32(0.0), - f16::from_f32(2.0), - f16::from_f32(0.0), - f16::from_f32(3.0), - f16::from_f32(0.0), - ]); +fn i8_conversion_matches_f32_across_dimension_boundaries() { + // Applying the same translation to every coordinate preserves L2 pair + // ordering while exercising signed conversion. + assert_source_conversion_matches_f32("i8", |value| value as i8 - 11); } #[test] diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 48b07b22f1..05191acb7f 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,25 +3,142 @@ * Licensed under the MIT license. */ -//! Numerical kernels for provider-independent PiPNN graph construction. +//! Provider-independent PiPNN graph construction. //! -//! PiPNN assigns points to overlapping leader partitions. It computes one -//! lower-triangular all-pairs matrix per bounded leaf. It then merges selected -//! leaf neighbors into graph candidates. +//! PiPNN means **Pick-in-Partitions Nearest Neighbors**. It builds a graph for +//! approximate nearest-neighbor search: every input vector becomes one graph +//! vertex, and its adjacency list stores other vectors worth visiting during a +//! later query. This crate constructs that adjacency; it does not execute queries. //! -//! - [`partition_kernel`] ranks leader-column positions. Output width is runtime -//! fanout. A scratch vector retains ranked leaders and reuses its allocation. -//! - [`leaf_kernel`] scans each strict-lower-triangle pair once. It updates both -//! endpoints and retains up to three leaf-local neighbors per point. -//! - `kernel_metric` owns norm preparation and exact metric ranking inputs. +//! Incremental builders such as Vamana find construction candidates by running +//! beam search against a partially built graph: they repeatedly follow graph +//! edges to discover nearby vertices, causing random memory access. PiPNN removes +//! that search from construction and uses three bulk stages instead: //! -//! The graph build selects concrete architecture `A` and metric `M` types once. -//! Both kernels validate views and metric norm layouts before unchecked SIMD -//! access. They mutate only caller-owned output and scratch storage. +//! 1. **Partition.** Randomized Ball Carving samples points called *leaders*. +//! Every point is assigned to its nearest `fanout` leaders. Assigning to more +//! than one leader makes child groups overlap. Oversized groups are processed +//! recursively until bounded groups called *leaves* remain. +//! 2. **Pick within leaves.** Vectors in one leaf are contiguous enough for one +//! dense general matrix multiplication (GEMM) to compute all pair dot +//! products. Each point picks its nearest leaf companions; selected pairs +//! become candidate graph edges. +//! 3. **Merge and prune.** Candidate edges from overlapping leaves are combined +//! into one unique list per source. Vamana RobustPrune then selects a bounded, +//! directionally diverse adjacency list. +//! +//! ```text +//! dataset points +//! │ +//! v +//! sample leaders + point/leader GEMM +//! │ +//! v +//! choose nearest leaders ──> overlapping child groups ──> recurse ──> leaves +//! │ +//! leaf all-pairs GEMM +//! │ +//! v +//! pick local neighbors +//! │ +//! v +//! merge/prune edges +//! │ +//! v +//! search graph +//! ``` +//! +//! This module keeps GEMM separate from score selection: callers compute dense +//! dot-product matrices, then the kernels documented below convert those dots to +//! metric scores and retain top candidates. A *point* is a vector being assigned +//! during partitioning; a *leader* names a child group. In leaf selection, +//! *source* names the point whose output list is being built and *target* names +//! another point in that same leaf. +//! +//! The crate owns overlapping partition generation, leaf-local nearest-neighbor +//! construction, candidate merging, and graph-degree finalization. The caller +//! supplies a contiguous dataset view, DiskANN graph policy, and the Rayon pool. +//! Providers, start/frozen points, quantization, persistence, and search remain +//! outside this algorithm boundary. +//! +//! Numerical kernels include: +//! +//! - [`partition_kernel::PartitionKernel`] converts point-by-leader dot-product +//! tiles into nearest leader positions. +//! - [`leaf_kernel::LeafKernel`] scans each leaf's lower-triangular dot-product +//! matrix once and retains nearest non-self neighbors for both endpoints. +//! +//! Both handles are prepared once per build metric and reused across stripes or +//! leaves. Each output view supplies call-specific fanout or neighbor width. +//! Preparation selects the runtime architecture and returns a direct function +//! pointer; repeated calls do not repeat ISA or metric dispatch. +//! +//! # Main modules and structures +//! +//! ## Public build API +//! +//! - [`PiPNNConfig`] holds Randomized Ball Carving, fanout, leaf size, local `k`, +//! and replication parameters. +//! - [`PiPNNBuildContext`] validates that algorithm parameters, graph pruning +//! policy, metric, and caller-owned Rayon pool agree. +//! - [`build_graph`] runs the full pipeline over a borrowed row-major dataset and +//! returns one dataset-ID adjacency list per input point. +//! +//! ## [`partition_kernel`] +//! +//! Partition callers compute point-by-leader dots with GEMM. +//! [`partition_kernel::PartitionInput`] bundles that tile with typed +//! [`partition_kernel::PartitionScales`]. A prepared +//! [`partition_kernel::PartitionKernel`] writes sorted leader-local positions to +//! caller-owned output. Fanout is output column count and is bounded by +//! [`partition_kernel::MAX_PARTITION_FANOUT`]. Module documentation describes +//! scale units, validation, `process_points`, and tracker insertion. +//! +//! ## [`leaf_kernel`] +//! +//! Leaf callers compute a lower-triangular point-by-point dot matrix with +//! `sgemm_aat_lower`. [`leaf_kernel::LeafKernelWorkspace`] owns reusable +//! per-worker scratch, and +//! [`leaf_kernel::LeafKernel`] writes sorted [`leaf_kernel::LeafNeighbor`] values +//! to caller-owned output. [`leaf_kernel::leaf_neighbor_count`] derives each +//! leaf's width from point count and requested `k`. Module documentation explains +//! fixed-width selection, `process_pairs`, and stable endpoint insertion. +//! +//! ## Private pipeline stages +//! +//! - `partitioning` recursively samples leaders, invokes the partition kernel, +//! scatters points into overlapping children, and returns bounded leaves. +//! - `leaf_build` gathers each leaf, computes its Gram matrix, invokes the leaf +//! kernel, translates local positions to dataset IDs, and merges candidates. +//! - `finalization` applies shared Vamana RobustPrune to overfull candidate lists. +//! - `kernel_metric` owns norm preparation, exact ranking inputs, and numerical +//! edge cases. The graph build selects one concrete metric for both kernels. +//! +//! # Typical use +//! +//! 1. Construct [`PiPNNConfig`] and DiskANN graph [`Config`]. +//! 2. Create [`PiPNNBuildContext`] with metric and caller-owned Rayon pool. +//! 3. Call [`build_graph`] with one row-major [`MatrixView`] of dataset vectors. +//! 4. The outer index builder chooses start/frozen points and serializes returned +//! adjacency; those policies are intentionally not part of this crate. +//! +//! Stage outputs are owned values. Leaves move into candidate construction; +//! candidate lists move into finalization. Ownership releases each stage's large +//! scratch before the next outer allocation. +//! +//! # Ownership and performance boundary +//! +//! Kernels borrow all matrices and mutate only caller-owned output/scratch. They +//! do not own providers, thread pools, GEMM buffers, graph IDs, or persistence. +//! Partition traversal performs one score per point-leader pair; leaf traversal +//! performs one score per unordered point pair. Detailed complexity and scratch +//! costs are documented in each module. PiPNN itself never names instruction +//! sets; `diskann-wide` owns architecture selection. -mod finalization; mod kernel_metric; mod simd; + +mod finalization; mod leaf_build; mod leaf_kernel; mod partition_kernel; @@ -186,15 +303,21 @@ where data.ncols() )) })?; + // Integer source rows are not guaranteed unit-normalized after conversion, + // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); let partition = partitioning::PartitionConfig::from(&context.config); let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition(data, partition, metric))?; + // `leaves` is consumed here. Workers borrow individual ID rows during the + // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) .map_err(ANNError::opaque) })?; + // Finalization consumes candidate rows and reuses their allocations for the + // resulting adjacency where possible. tracing::info_span!("pipnn.finalization") .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 2d9c37d29e..248b8cb1bb 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -106,6 +106,13 @@ struct StripeBuffers { row_scales: Vec, } +/// Stage-owned high-water scratch storage for partition assignment. +/// +/// Rayon initializes `map_init` state per split job, not per physical worker. +/// Large dimensions would therefore allocate and zero the point/dot buffers many +/// times during recursive partitioning. Chunks instead take ownership of one +/// buffer here and return it afterward. The mutex protects only the short +/// pop/push operations; gather, GEMM, and top-k run without holding it. #[derive(Default)] struct StripeBufferPool { available: Mutex>, @@ -329,6 +336,13 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { .wrapping_add(salt) } +/// Assign each point to its nearest `fanout` sampled leaders. +/// +/// Leader rows are gathered once. Point rows are processed in cache-sized +/// stripes, while a worker chunk retains one leased scratch buffer across all of +/// its stripes. The flat assignment matrix preserves point order and is then +/// scattered into per-leader clusters; preserving order is required for fixed +/// seed determinism in later recursion levels. fn assign_to_leaders( data: MatrixView<'_, T>, points: &[u32], @@ -486,6 +500,12 @@ where Ok(()) } +/// Convert the flat point-major assignment matrix into leader-major clusters. +/// +/// Small inputs use one serial exact-capacity pass. Large inputs form at most +/// one partial cluster set per Rayon worker, then merge each leader independently. +/// Concatenating partials in stripe order keeps the same member order as the +/// serial implementation while removing a large serial copy tail. fn scatter_assignments( points: &[u32], assignments: &[u32], diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index d3815d533e..acaffd1c73 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -173,54 +173,60 @@ fn replicas_cover_every_point_once_or_more_per_replica() { assert_valid_partition(&leaves, 72, 12, 3); } +fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) +where + T: diskann::utils::VectorRepr + Send + Sync, +{ + let points = 64; + // Partition gathering converts source rows before GEMM. Exercise conversion + // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let row = index / dimensions; + let column = index % dimensions; + ((row * 5 + column * 7 + row * column) % 23) as u8 + }) + .collect(); + let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + let config = config(2, 16, vec![2, 1], 1); + let expected = partition( + MatrixView::try_from(&f32_data, points, dimensions).unwrap(), + config.clone(), + Metric::L2, + ) + .unwrap(); + let actual = partition( + MatrixView::try_from(&converted, points, dimensions).unwrap(), + config, + Metric::L2, + ) + .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); + + assert_valid_partition(&actual, points, 16, 1); + assert_eq!( + sorted_memberships(&actual), + sorted_memberships(&expected), + "{label} dimensions={dimensions}" + ); + } +} + #[test] -fn supported_source_types_share_partition_contract() { - let f32_data: Vec = (0..64 * 4).map(|value| (value % 23) as f32).collect(); - let half_data: Vec = f32_data.iter().copied().map(Half::from_f32).collect(); - let u8_data: Vec = f32_data.iter().map(|value| *value as u8).collect(); - let i8_data: Vec = u8_data.iter().map(|value| *value as i8 - 11).collect(); - let config = config(2, 16, vec![2, 1], 1); - - let f32_leaves = partition( - MatrixView::try_from(f32_data.as_slice(), 64, 4).unwrap(), - config.clone(), - Metric::L2, - ) - .unwrap(); - let half_leaves = partition( - MatrixView::try_from(half_data.as_slice(), 64, 4).unwrap(), - config.clone(), - Metric::L2, - ) - .unwrap(); - let u8_leaves = partition( - MatrixView::try_from(u8_data.as_slice(), 64, 4).unwrap(), - config.clone(), - Metric::L2, - ) - .unwrap(); - let i8_leaves = partition( - MatrixView::try_from(i8_data.as_slice(), 64, 4).unwrap(), - config, - Metric::L2, - ) - .unwrap(); +fn f16_partition_matches_f32_across_dimension_boundaries() { + assert_partition_conversion_matches_f32("f16", |value| Half::from_f32(value as f32)); +} - for leaves in [&f32_leaves, &half_leaves, &u8_leaves, &i8_leaves] { - assert_valid_partition(leaves, 64, 16, 1); - } - assert_eq!( - sorted_memberships(&f32_leaves), - sorted_memberships(&half_leaves) - ); - assert_eq!( - sorted_memberships(&f32_leaves), - sorted_memberships(&u8_leaves) - ); - assert_eq!( - sorted_memberships(&u8_leaves), - sorted_memberships(&i8_leaves) - ); +#[test] +fn u8_partition_matches_f32_across_dimension_boundaries() { + assert_partition_conversion_matches_f32("u8", |value| value); +} + +#[test] +fn i8_partition_matches_f32_across_dimension_boundaries() { + // The same translation in every coordinate preserves L2 ordering. + assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); } #[test] @@ -272,6 +278,20 @@ fn stripe_buffer_pool_reuses_returned_capacity() { assert_eq!(buffers.points.len(), 16); } +#[test] +fn stripe_buffer_pool_recovers_after_lock_poisoning() { + let pool = StripeBufferPool::default(); + let _ = std::panic::catch_unwind(|| { + let _guard = pool.available.lock().unwrap(); + panic!("poison scratch pool"); + }); + + let mut buffers = pool.take(); + buffers.dots.push(1.0); + pool.put(buffers); + assert_eq!(pool.take().dots, [1.0]); +} + #[test] fn leader_assignment_handles_multiple_stripes() { let points = 2_048; From 0aa68eba8f66b6a69177b43c36ca10aa41c88fe8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:37:22 +0000 Subject: [PATCH 15/58] docs(pipnn): diagram core stage ownership --- diskann/src/graph/pipnn/finalization.rs | 20 ++++++++++++++++++ diskann/src/graph/pipnn/partitioning.rs | 27 +++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 5cae12978a..7388a8ff0b 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -14,6 +14,26 @@ //! The shared kernel owns occlusion and alpha-round semantics; this adapter owns //! only contiguous dataset access and distance specialization for the source //! representation. +//! +//! ```text +//! candidate rows ──> validate row count and every global ID +//! │ +//! ┌────────────┴────────────┐ +//! v v +//! len <= R len > R +//! return row source-distance candidates +//! │ +//! v +//! shared RobustPrune +//! │ +//! v +//! rewrite same row owner +//! ``` +//! +//! | Path | Distance evaluations | Allocation behavior | +//! | --- | --- | --- | +//! | bounded row | none | move row directly to output | +//! | overfull row | source and occlusion distances | reuse Rayon-job workspace | use std::convert::Infallible; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 248b8cb1bb..84c6836ab4 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -9,6 +9,33 @@ //! reuses the partition kernel and dense GEMM. A stage-owned pool leases scratch //! to Rayon chunks and takes it back after each chunk; computation never holds //! the pool lock, and no thread-local cleanup protocol is required. +//! +//! ```text +//! replica root IDs ──> work queue +//! │ +//! v +//! sample leaders +//! │ +//! gather stripes ─> GEMM distances ─> nearest leaders +//! │ +//! v +//! stable scatter by leader +//! │ │ +//! size <= c_max oversized cluster +//! │ │ +//! completed leaf next recursion level +//! └──────────┬────────┘ +//! v +//! global small-leaf merge +//! v +//! coverage/bound validation +//! ``` +//! +//! | Recursion level | Assignment multiplicity | +//! | --- | --- | +//! | `level < fanout.len()` | `fanout[level]` nearest leaders | +//! | later levels | one nearest leader until bounded | +//! | replica boundary | independent deterministic seed | use std::{collections::HashSet, sync::Mutex}; From 41b8037f61b58d44d22a9e868b47e9af81f5e6f4 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 02:11:33 +0000 Subject: [PATCH 16/58] fix(pipnn): preserve partition quality and scratch reuse --- diskann/src/graph/pipnn/finalization.rs | 7 +- diskann/src/graph/pipnn/mod.rs | 3 +- diskann/src/graph/pipnn/partitioning.rs | 167 +++++++++--------- diskann/src/graph/pipnn/partitioning/tests.rs | 74 +++++--- 4 files changed, 135 insertions(+), 116 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 7388a8ff0b..8acb5fda46 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -132,8 +132,11 @@ where ) .map_err(ANNError::opaque)?; - row.clear(); - row.extend_from_slice(workspace.prune.neighbors()); + // RobustPrune selects distinct candidate positions, so its output IDs + // are unique by construction. `extend_from_slice` would re-derive that + // with an O(degree^2) membership scan per row; the trusted overwrite is + // a copy and still verifies uniqueness under debug assertions. + row.overwrite_trusted(workspace.prune.neighbors()); Ok(row) }) .collect() diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 05191acb7f..2b0c555be9 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -307,9 +307,8 @@ where // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); - let partition = partitioning::PartitionConfig::from(&context.config); let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition(data, partition, metric))?; + .in_scope(|| partitioning::partition(data, context.config.clone(), metric))?; // `leaves` is consumed here. Workers borrow individual ID rows during the // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 84c6836ab4..84659095ba 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -37,11 +37,14 @@ //! | later levels | one nearest leader until bounded | //! | replica boundary | independent deterministic seed | -use std::{collections::HashSet, sync::Mutex}; +use std::collections::HashSet; use crate::{utils::VectorRepr, ANNError, ANNResult}; use diskann_linalg::Transpose; -use diskann_utils::views::MatrixView; +use diskann_utils::{ + object_pool::{AsPooled, ObjectPool}, + views::MatrixView, +}; use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; use rand::{prelude::IndexedRandom, SeedableRng}; use rayon::prelude::*; @@ -61,29 +64,6 @@ const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const MAX_PARTITION_ITERATIONS: usize = 30; -/// Policy owned by the partition stage. Leaf-neighbor and merge settings do -/// not cross this boundary. -#[derive(Clone, Debug)] -pub(crate) struct PartitionConfig { - c_max: usize, - c_min: usize, - p_samp: f64, - fanout: Vec, - replicas: usize, -} - -impl From<&PiPNNConfig> for PartitionConfig { - fn from(config: &PiPNNConfig) -> Self { - Self { - c_max: config.c_max, - c_min: config.c_min, - p_samp: config.p_samp, - fanout: config.fanout.clone(), - replicas: config.replicas, - } - } -} - /// A partition failure with enough context to diagnose non-progressing input. #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum PartitionError { @@ -133,35 +113,24 @@ struct StripeBuffers { row_scales: Vec, } -/// Stage-owned high-water scratch storage for partition assignment. -/// -/// Rayon initializes `map_init` state per split job, not per physical worker. -/// Large dimensions would therefore allocate and zero the point/dot buffers many -/// times during recursive partitioning. Chunks instead take ownership of one -/// buffer here and return it afterward. The mutex protects only the short -/// pop/push operations; gather, GEMM, and top-k run without holding it. -#[derive(Default)] -struct StripeBufferPool { - available: Mutex>, -} - -impl StripeBufferPool { - fn take(&self) -> StripeBuffers { - self.available - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .pop() - .unwrap_or_default() +impl AsPooled<()> for StripeBuffers { + fn create(_: ()) -> Self { + Self::default() } - fn put(&self, buffers: StripeBuffers) { - self.available - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) - .push(buffers); + fn modify(&mut self, _: ()) { + // Scratch retains its high-water allocation across leases; active + // prefixes are established by `assign_stripe` before every read. } } +/// Stage-owned high-water scratch storage for partition assignment. +/// +/// `ObjectPool` owns the short pop/push lock and returns leases through RAII, +/// including error and panic paths. Gather, GEMM, and top-k run while only the +/// leased `StripeBuffers` is held. +type StripeBufferPool = ObjectPool; + /// Partition every configured replica into overlapping bounded leaves. /// /// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped @@ -172,7 +141,7 @@ impl StripeBufferPool { /// covered once per replica. The caller installs the operation in its pool. pub(crate) fn partition( data: MatrixView<'_, T>, - config: PartitionConfig, + config: PiPNNConfig, metric: Metric, ) -> ANNResult>> where @@ -190,7 +159,7 @@ where } let mut leaves = Vec::new(); - let stripe_buffers = StripeBufferPool::default(); + let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = partition_replica(data, &config, metric, seed, &stripe_buffers)?; @@ -205,7 +174,7 @@ where fn partition_replica( data: MatrixView<'_, T>, - config: &PartitionConfig, + config: &PiPNNConfig, metric: Metric, seed: u64, stripe_buffers: &StripeBufferPool, @@ -287,7 +256,7 @@ where fn partition_one_level( data: MatrixView<'_, T>, - config: &PartitionConfig, + config: &PiPNNConfig, metric: Metric, item: WorkItem, stripe_buffers: &StripeBufferPool, @@ -395,7 +364,11 @@ where .iter_mut() .zip(leader_values.chunks_exact(dimensions)) { - *scale = FastL2NormSquared.evaluate(row); + // Leader norms participate in the top-k ordering. Preserve the original + // scalar reduction order: reassociating this short setup pass through a + // SIMD norm changes low bits and can send near-tied points down different + // recursive partition paths. + *scale = row.iter().map(|value| value * value).sum(); if metric == Metric::Cosine { *scale = scale.sqrt(); } @@ -418,27 +391,23 @@ where .par_chunks_mut(worker_assignment) .enumerate() .try_for_each(|(worker, worker_output)| { - let mut buffers = stripe_buffers.take(); - let result = (|| -> ANNResult<()> { - let worker_first = worker * worker_rows; - for (stripe, output) in worker_output.chunks_mut(assignment_stripe).enumerate() { - let first = worker_first + stripe * stripe_rows; - let rows = output.len() / fanout; - assign_stripe( - data, - &points[first..first + rows], - &leader_values, - &leader_scales, - metric, - fanout, - &mut buffers, - output, - )?; - } - Ok(()) - })(); - stripe_buffers.put(buffers); - result + let mut buffers = stripe_buffers.get_ref(()); + let worker_first = worker * worker_rows; + for (stripe, output) in worker_output.chunks_mut(assignment_stripe).enumerate() { + let first = worker_first + stripe * stripe_rows; + let rows = output.len() / fanout; + assign_stripe( + data, + &points[first..first + rows], + &leader_values, + &leader_scales, + metric, + fanout, + &mut buffers, + output, + )?; + } + Ok::<(), ANNError>(()) })?; scatter_assignments(points, &assignments, fanout, leaders.len()) @@ -464,9 +433,23 @@ where let leaders = leader_values.len() / dimensions; let point_values_len = checked_area("point stripe", rows, dimensions)?; let dots_len = checked_area("dot-product stripe", rows, leaders)?; - resize_fallible(&mut buffers.points, point_values_len, 0.0)?; - resize_fallible(&mut buffers.dots, dots_len, 0.0)?; - gather_rows(data, points, &mut buffers.points)?; + // Scratch keeps its high-water length and every consumer receives an + // explicit active prefix. Resizing to the exact stripe shape would be + // correct but re-zeroes the buffer whenever a pooled value moves between + // work items with different leader counts: `stripe_rows` is derived from + // `leaders`, so the point buffer swings between roughly 768 KiB and 6 MiB + // and `Vec::resize` only truncates on the way down, then memsets the whole + // delta on the way back up. + grow_fallible(&mut buffers.points, point_values_len, 0.0)?; + grow_fallible(&mut buffers.dots, dots_len, 0.0)?; + let StripeBuffers { + points: point_buffer, + dots: dot_buffer, + row_scales: row_scale_buffer, + } = buffers; + let point_values = &mut point_buffer[..point_values_len]; + let dots = &mut dot_buffer[..dots_len]; + gather_rows(data, points, point_values)?; diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -474,29 +457,29 @@ where leaders, dimensions, 1.0, - &buffers.points, + point_values, leader_values, None, - &mut buffers.dots, + dots, ) .map_err(ANNError::opaque)?; let row_scales = if metric == Metric::Cosine { - resize_fallible(&mut buffers.row_scales, rows, 0.0)?; - for (scale, row) in buffers - .row_scales + grow_fallible(row_scale_buffer, rows, 0.0)?; + let row_scales = &mut row_scale_buffer[..rows]; + for (scale, row) in row_scales .iter_mut() - .zip(buffers.points.chunks_exact(dimensions)) + .zip(point_values.chunks_exact(dimensions)) { *scale = FastL2NormSquared.evaluate(row); } - buffers.row_scales.as_slice() + &*row_scales } else { &[] }; nearest_leaders( PartitionTopK { - dots: &buffers.dots, + dots, rows, leaders, row_scales, @@ -755,9 +738,17 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { Ok(values) } -fn resize_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { +/// Grow `values` to at least `len` elements, never shrinking it. +/// +/// Callers slice the active prefix themselves. Shrinking would force the next +/// larger stripe to re-zero the reclaimed tail, which is the dominant cost when +/// one pooled buffer serves work items with different stripe shapes. +fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { + if values.len() >= len { + return Ok(()); + } values - .try_reserve(len.saturating_sub(values.len())) + .try_reserve(len - values.len()) .map_err(ANNError::opaque)?; values.resize(len, value); Ok(()) diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index acaffd1c73..31f3e4bcad 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -8,12 +8,13 @@ use diskann_vector::{distance::Metric, Half}; use super::*; -fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PartitionConfig { - PartitionConfig { +fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { + PiPNNConfig { c_max, c_min, p_samp: 0.25, fanout, + k: 1, replicas, } } @@ -229,6 +230,44 @@ fn i8_partition_matches_f32_across_dimension_boundaries() { assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); } +#[test] +fn l2_leader_norms_preserve_scalar_reduction_order() { + fn next(state: &mut u64) -> f32 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 + } + + // This fixed case sits on opposite sides of the top-1 boundary depending + // on whether leader norms use the original scalar reduction or a SIMD + // reassociation. Point/leader dot products still go through the production + // GEMM; only the setup norm calculation is under test. + let dimensions = 129; + let mut state = 0x3a85_f952_c718_6e49; + let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let leader_zero: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let leader_one: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let data: Vec = leader_zero + .into_iter() + .chain(leader_one) + .chain(point) + .collect(); + let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); + + let clusters = assign_to_leaders( + data, + &[2], + &[0, 1], + 1, + Metric::L2, + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); + + assert_eq!(clusters, [vec![], vec![2]]); +} + #[test] fn all_metrics_produce_valid_partitions() { let data = directional_data(64, 8); @@ -267,31 +306,18 @@ fn assignment_stripes_use_power_of_two_row_counts() { #[test] fn stripe_buffer_pool_reuses_returned_capacity() { - let pool = StripeBufferPool::default(); - let mut buffers = pool.take(); - buffers.points.resize(16, 0.0); - let points = buffers.points.as_ptr(); - pool.put(buffers); - - let buffers = pool.take(); + let pool = StripeBufferPool::new((), 0, None); + let points = { + let mut buffers = pool.get_ref(()); + buffers.points.resize(16, 0.0); + buffers.points.as_ptr() + }; + + let buffers = pool.get_ref(()); assert_eq!(buffers.points.as_ptr(), points); assert_eq!(buffers.points.len(), 16); } -#[test] -fn stripe_buffer_pool_recovers_after_lock_poisoning() { - let pool = StripeBufferPool::default(); - let _ = std::panic::catch_unwind(|| { - let _guard = pool.available.lock().unwrap(); - panic!("poison scratch pool"); - }); - - let mut buffers = pool.take(); - buffers.dots.push(1.0); - pool.put(buffers); - assert_eq!(pool.take().dots, [1.0]); -} - #[test] fn leader_assignment_handles_multiple_stripes() { let points = 2_048; @@ -305,7 +331,7 @@ fn leader_assignment_handles_multiple_stripes() { &[0, 2_047], 1, Metric::L2, - &StripeBufferPool::default(), + &StripeBufferPool::new((), 0, None), ) .unwrap(); From b42a8568530501c91ad56c7f1b3b0b4c774a5f0f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 10:14:35 +0000 Subject: [PATCH 17/58] refactor(pipnn): reuse prepared kernels --- diskann/src/graph/pipnn/leaf_build.rs | 274 ++++++++++-------- diskann/src/graph/pipnn/leaf_build/tests.rs | 68 +++-- diskann/src/graph/pipnn/partitioning.rs | 192 ++++++------ diskann/src/graph/pipnn/partitioning/tests.rs | 37 ++- 4 files changed, 325 insertions(+), 246 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index f847077ef2..37cf5e3d14 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -7,12 +7,12 @@ //! //! Partitioning supplies leaves as global point IDs. For each leaf this module: //! -//! 1. validates IDs and converts only those rows to reusable `f32` scratch; +//! 1. validates IDs and converts only those point vectors to reusable `f32` scratch; //! 2. computes the lower triangle of `A · Aᵀ`; //! 3. runs the dual-endpoint leaf top-k kernel; and //! 4. translates leaf-local positions back to dataset IDs. //! -//! The final step merges symmetric adjacency rows under per-point locks because +//! The final step merges symmetric adjacency lists under per-point locks because //! overlapping leaves are processed concurrently. Numeric buffers retain their //! high-water length; every consumer therefore receives an explicit active //! prefix rather than treating `Vec::len()` as the current leaf shape. @@ -23,12 +23,13 @@ use std::{ }; use crate::{graph::AdjacencyList, utils::VectorRepr}; -use diskann_utils::views::MatrixView; +use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use rayon::prelude::*; use crate::leaf_kernel::{ - nearest_leaf_neighbors, LeafKernelError, LeafNeighbor, LeafTopK, LeafTopKWorkspace, + leaf_neighbor_count, leaf_output_len, LeafInput, LeafKernel, LeafKernelError, + LeafKernelWorkspace, LeafNeighbor, }; /// Failure while converting leaves into direct graph candidates. @@ -80,25 +81,25 @@ pub(crate) enum LeafBuildError { #[source] source: LeafKernelError, }, - #[error("leaf kernel returned local position {position} for a {points}-point leaf")] - InvalidLocalPosition { position: u32, points: usize }, - #[error("candidate row {point} is poisoned")] - PoisonedCandidateRow { point: u32 }, + #[error("leaf kernel returned local target {target} for a {points}-point leaf")] + InvalidLocalTarget { target: u32, points: usize }, + #[error("candidate list for point {point} is poisoned")] + PoisonedCandidateList { point: u32 }, } /// Scratch leased to one Rayon job and reused for successive leaves. /// /// The three numerical vectors retain their largest observed leaf shape. The -/// adjacency rows are prepared separately because zero-k/singleton leaves never +/// adjacency lists are prepared separately because zero-k/singleton leaves never /// write them, and because later candidate-merging modes do not necessarily use /// this representation. #[derive(Default)] struct LeafBuffers { - points: Vec, + point_values: Vec, dots: Vec, - nearest: Vec, - local_graph: Vec>, - top_k: LeafTopKWorkspace, + neighbors: Vec, + local_adjacency: Vec>, + kernel_workspace: LeafKernelWorkspace, seen_ids: HashSet, } @@ -106,51 +107,55 @@ impl LeafBuffers { fn prepare( &mut self, leaf: usize, - points: usize, - dimensions: usize, - k: usize, + point_count: usize, + dimension_count: usize, + requested_k: usize, ) -> Result { - let point_values = points - .checked_mul(dimensions) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: points, - columns: dimensions, - })?; - let dot_values = points - .checked_mul(points) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: points, - columns: points, - })?; - let actual_k = k.min(points.saturating_sub(1)); - let nearest_values = points - .checked_mul(actual_k) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: points, - columns: actual_k, - })?; + let point_value_count = + point_count + .checked_mul(dimension_count) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: point_count, + columns: dimension_count, + })?; + let dot_count = + point_count + .checked_mul(point_count) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: point_count, + columns: point_count, + })?; + let leaf_k = leaf_neighbor_count(point_count, requested_k) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + let neighbor_count = leaf_output_len(point_count, requested_k) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - grow("leaf points", &mut self.points, point_values, 0.0)?; - grow("leaf dot products", &mut self.dots, dot_values, 0.0)?; grow( - "leaf nearest neighbors", - &mut self.nearest, - nearest_values, + "leaf point values", + &mut self.point_values, + point_value_count, + 0.0, + )?; + grow("leaf dot products", &mut self.dots, dot_count, 0.0)?; + grow( + "leaf neighbors", + &mut self.neighbors, + neighbor_count, LeafNeighbor::default(), )?; - Ok(actual_k) + Ok(leaf_k) } - fn prepare_local_graph(&mut self, points: usize) -> Result<(), LeafBuildError> { - let additional = points.saturating_sub(self.local_graph.len()); - self.local_graph + fn prepare_local_adjacency(&mut self, point_count: usize) -> Result<(), LeafBuildError> { + let additional = point_count.saturating_sub(self.local_adjacency.len()); + self.local_adjacency .try_reserve(additional) - .map_err(|source| allocation_error("leaf adjacency rows", additional, source))?; - self.local_graph.resize_with(points, AdjacencyList::new); - self.local_graph[..points] + .map_err(|source| allocation_error("leaf adjacency lists", additional, source))?; + self.local_adjacency + .resize_with(point_count, AdjacencyList::new); + self.local_adjacency[..point_count] .iter_mut() .for_each(AdjacencyList::clear); Ok(()) @@ -160,45 +165,50 @@ impl LeafBuffers { /// Concurrent accumulator indexed by global dataset ID. /// /// A point may appear in several overlapping leaves, so workers lock only the -/// destination row long enough to append one leaf's additions. Sorting and +/// destination list long enough to append one leaf's additions. Sorting and /// duplicate removal are deferred until all leaves finish; doing either under /// the lock would lengthen the contended section for no semantic benefit. struct DirectCandidates { - rows: Vec>>, + lists: Vec>>, } impl DirectCandidates { - fn new(points: usize) -> Result { - let mut rows = Vec::new(); - rows.try_reserve_exact(points) - .map_err(|source| allocation_error("candidate rows", points, source))?; - rows.resize_with(points, || Mutex::new(AdjacencyList::new())); - Ok(Self { rows }) + fn new(point_count: usize) -> Result { + let mut lists = Vec::new(); + lists + .try_reserve_exact(point_count) + .map_err(|source| allocation_error("candidate lists", point_count, source))?; + lists.resize_with(point_count, || Mutex::new(AdjacencyList::new())); + Ok(Self { lists }) } fn add_leaf( &self, point_ids: &[u32], - local_graph: &[AdjacencyList], + local_adjacency: &[AdjacencyList], ) -> Result<(), LeafBuildError> { - for (&source, additions) in point_ids.iter().zip(local_graph) { + for (&source, additions) in point_ids.iter().zip(local_adjacency) { // Every point ID is validated before leaf-local work begins. - let row = &self.rows[source as usize]; - let mut row = row.lock().map_err(|_| poisoned_row(source))?; - row.extend_from_slice(additions); + let candidates = &self.lists[source as usize]; + let mut candidates = candidates + .lock() + .map_err(|_| poisoned_candidate_list(source))?; + candidates.extend_from_slice(additions); } Ok(()) } - fn into_rows(self) -> Result>, LeafBuildError> { + fn into_lists(self) -> Result>, LeafBuildError> { let mut output = Vec::new(); output - .try_reserve_exact(self.rows.len()) - .map_err(|source| allocation_error("candidate output", self.rows.len(), source))?; - for (point, row) in self.rows.into_iter().enumerate() { - let mut row = row.into_inner().map_err(|_| poisoned_row(point as u32))?; - row.sort(); - output.push(row); + .try_reserve_exact(self.lists.len()) + .map_err(|source| allocation_error("candidate output", self.lists.len(), source))?; + for (point, candidates) in self.lists.into_iter().enumerate() { + let mut candidates = candidates + .into_inner() + .map_err(|_| poisoned_candidate_list(point as u32))?; + candidates.sort(); + output.push(candidates); } Ok(output) } @@ -209,7 +219,7 @@ impl DirectCandidates { pub(crate) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, - k: usize, + requested_k: usize, metric: Metric, ) -> Result>, LeafBuildError> where @@ -223,16 +233,27 @@ where } let candidates = DirectCandidates::new(data.nrows())?; + // Metric and ISA are selected before Rayon workers start. Workers share + // this Copy handle; each output view supplies its leaf-specific width. + let kernel = LeafKernel::new(metric); leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf(data, leaf, point_ids, k, metric, buffers, &candidates) + build_leaf( + data, + leaf, + point_ids, + requested_k, + &kernel, + buffers, + &candidates, + ) }, )?; - candidates.into_rows() + candidates.into_lists() } -/// Build and publish one leaf's symmetric nearest-neighbor rows. +/// Build and publish one leaf's symmetric neighbor lists. /// /// Validation precedes all dataset indexing. Sorted partition output takes the /// adjacent-duplicate path, while arbitrary-order callers use `seen_ids`. The @@ -242,8 +263,8 @@ fn build_leaf( data: MatrixView<'_, T>, leaf: usize, point_ids: &[u32], - k: usize, - metric: Metric, + requested_k: usize, + kernel: &LeafKernel, buffers: &mut LeafBuffers, candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> @@ -281,75 +302,92 @@ where } } } - let actual_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), k)?; - if actual_k == 0 { + let leaf_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; + if leaf_k == 0 { return Ok(()); } - let point_values = point_ids.len() * data.ncols(); - let dot_values = point_ids.len() * point_ids.len(); - let nearest_values = point_ids.len() * actual_k; + let point_value_count = point_ids.len() * data.ncols(); + let dot_count = point_ids.len() * point_ids.len(); + let neighbor_value_count = leaf_output_len(point_ids.len(), requested_k) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - for (&point, output) in point_ids + for (&point, point_output) in point_ids .iter() - .zip(buffers.points[..point_values].chunks_exact_mut(data.ncols())) + .zip(buffers.point_values[..point_value_count].chunks_exact_mut(data.ncols())) { - let row = data.row(point as usize); - T::as_f32_into(row, output).map_err(|source| LeafBuildError::Conversion { - leaf, - point, - source: source.into(), + let source_values = data.row(point as usize); + T::as_f32_into(source_values, point_output).map_err(|source| { + LeafBuildError::Conversion { + leaf, + point, + source: source.into(), + } })?; } diskann_linalg::sgemm_aat_lower( - &buffers.points[..point_values], point_ids.len(), data.ncols(), - &mut buffers.dots[..dot_values], + &buffers.point_values[..point_value_count], + &mut buffers.dots[..dot_count], ) .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; - nearest_leaf_neighbors( - LeafTopK { - dots: &buffers.dots[..dot_values], - points: point_ids.len(), - metric, + 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(), }, - k, - &mut buffers.nearest[..nearest_values], - &mut buffers.top_k, + })?; + let output = MutMatrixView::try_from( + &mut buffers.neighbors[..neighbor_value_count], + point_ids.len(), + leaf_k, ) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + .map_err(|error| LeafBuildError::Kernel { + leaf, + source: LeafKernelError::InvalidBufferLength { + buffer: "output", + expected: neighbor_value_count, + actual: error.into_inner().len(), + }, + })?; + kernel + .nearest_neighbors(LeafInput { dots }, output, &mut buffers.kernel_workspace) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - buffers.prepare_local_graph(point_ids.len())?; - add_symmetric_edges( + buffers.prepare_local_adjacency(point_ids.len())?; + add_symmetric_neighbors( point_ids, - actual_k, - &buffers.nearest[..nearest_values], - &mut buffers.local_graph[..point_ids.len()], + leaf_k, + &buffers.neighbors[..neighbor_value_count], + &mut buffers.local_adjacency[..point_ids.len()], )?; - candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) + candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]) } -fn add_symmetric_edges( +fn add_symmetric_neighbors( point_ids: &[u32], - k: usize, - nearest: &[LeafNeighbor], - local_graph: &mut [AdjacencyList], + leaf_k: usize, + neighbors: &[LeafNeighbor], + local_adjacency: &mut [AdjacencyList], ) -> Result<(), LeafBuildError> { - for (source, nearest) in nearest.chunks_exact(k).enumerate() { - for neighbor in nearest { - let target = neighbor.position as usize; + for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for neighbor in source_neighbors { + let target = neighbor.target as usize; let Some(&target_id) = point_ids.get(target) else { - return Err(LeafBuildError::InvalidLocalPosition { - position: neighbor.position, + return Err(LeafBuildError::InvalidLocalTarget { + target: neighbor.target, points: point_ids.len(), }); }; let source_id = point_ids[source]; if source_id != target_id { - local_graph[source].push(target_id); - local_graph[target].push(source_id); + local_adjacency[source].push(target_id); + local_adjacency[target].push(source_id); } } } @@ -394,8 +432,8 @@ fn allocation_error( } } -fn poisoned_row(point: u32) -> LeafBuildError { - LeafBuildError::PoisonedCandidateRow { point } +fn poisoned_candidate_list(point: u32) -> LeafBuildError { + LeafBuildError::PoisonedCandidateList { point } } #[cfg(test)] diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index 0236370dee..2d4e66c2e6 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -9,8 +9,8 @@ use half::f16; use std::collections::BTreeSet; use super::{ - add_symmetric_edges, allocation_error, build_leaf_candidates, DirectCandidates, LeafBuffers, - LeafBuildError, + add_symmetric_neighbors, allocation_error, build_leaf_candidates, DirectCandidates, + LeafBuffers, LeafBuildError, }; fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { @@ -36,7 +36,7 @@ where pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) } -fn rows(graph: Vec>) -> Vec> { +fn adjacency_lists(graph: Vec>) -> Vec> { graph.into_iter().map(Vec::from).collect() } @@ -84,7 +84,7 @@ fn leaf_adjacency_matches_an_independent_all_pairs_reference() { ]; let flat: Vec<_> = points.into_iter().flatten().collect(); - let actual = rows( + let actual = adjacency_lists( build( view(&flat, points.len(), 2), &[(0..points.len() as u32).collect()], @@ -105,7 +105,7 @@ fn retains_and_deduplicates_candidates_from_overlapping_leaves() { let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); assert_eq!( - rows(graph), + adjacency_lists(graph), [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] ); } @@ -114,8 +114,8 @@ fn retains_and_deduplicates_candidates_from_overlapping_leaves() { fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { let dimensions = 9; let mut data = vec![0.0_f32; 10 * dimensions]; - for row in 1..10 { - data[row * dimensions + row - 1] = 1.0; + for source in 1..10 { + data[source * dimensions + source - 1] = 1.0; } let graph = build( @@ -143,7 +143,7 @@ fn global_id_translation_is_independent_of_leaf_order() { let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); assert_eq!( - rows(graph), + adjacency_lists(graph), [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] ); } @@ -153,7 +153,7 @@ where T: diskann::utils::VectorRepr + 'static, { let leaves = vec![(0..points as u32).collect()]; - rows(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) + adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) } fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) @@ -167,9 +167,9 @@ where for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let raw: Vec = (0..points * dimensions) .map(|index| { - let row = index / dimensions; - let column = index % dimensions; - ((row * 7 + column * 3 + row * column) % 23) as u8 + let source = index / dimensions; + let dimension = index % dimensions; + ((source * 7 + dimension * 3 + source * dimension) % 23) as u8 }) .collect(); let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); @@ -288,25 +288,28 @@ fn singleton_and_zero_k_leaves_add_no_candidates() { ) .unwrap(); let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); - assert!(singleton.iter().chain(&zero_k).all(|row| row.is_empty())); + assert!(singleton + .iter() + .chain(&zero_k) + .all(|candidates| candidates.is_empty())); } #[test] fn reuses_worker_buffers_for_smaller_leaves() { let mut buffers = LeafBuffers::default(); buffers.prepare(0, 64, 128, 2).unwrap(); - let points = buffers.points.as_ptr(); + let point_values = buffers.point_values.as_ptr(); let dots = buffers.dots.as_ptr(); - let nearest = buffers.nearest.as_ptr(); + let neighbors = buffers.neighbors.as_ptr(); buffers.prepare(1, 8, 128, 2).unwrap(); - assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.point_values.as_ptr(), point_values); assert_eq!(buffers.dots.as_ptr(), dots); - assert_eq!(buffers.nearest.as_ptr(), nearest); - assert_eq!(buffers.points.len(), 64 * 128); + assert_eq!(buffers.neighbors.as_ptr(), neighbors); + assert_eq!(buffers.point_values.len(), 64 * 128); assert_eq!(buffers.dots.len(), 64 * 64); - assert_eq!(buffers.nearest.len(), 64 * 2); + assert_eq!(buffers.neighbors.len(), 64 * 2); } #[test] @@ -321,7 +324,7 @@ fn reports_shape_overflow_before_allocating() { #[test] fn rejects_an_invalid_kernel_position() { let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; - let error = add_symmetric_edges( + let error = add_symmetric_neighbors( &[10, 20], 1, &[ @@ -333,8 +336,8 @@ fn rejects_an_invalid_kernel_position() { .unwrap_err(); assert!(matches!( error, - LeafBuildError::InvalidLocalPosition { - position: 9, + LeafBuildError::InvalidLocalTarget { + target: 9, points: 2 } )); @@ -343,7 +346,7 @@ fn rejects_an_invalid_kernel_position() { #[test] fn skips_duplicate_global_ids_without_self_edges() { let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; - add_symmetric_edges( + add_symmetric_neighbors( &[7, 7], 1, &[ @@ -353,23 +356,23 @@ fn skips_duplicate_global_ids_without_self_edges() { &mut graph, ) .unwrap(); - assert!(graph.iter().all(|row| row.is_empty())); + assert!(graph.iter().all(|neighbors| neighbors.is_empty())); } #[test] -fn poisoned_candidate_rows_return_errors() { +fn poisoned_candidate_lists_return_errors() { let candidates = DirectCandidates::new(1).unwrap(); let _ = std::panic::catch_unwind(|| { - let _guard = candidates.rows[0].lock().unwrap(); - panic!("poison candidate row"); + let _guard = candidates.lists[0].lock().unwrap(); + panic!("poison candidate list"); }); assert!(matches!( candidates.add_leaf(&[0], &[diskann::graph::AdjacencyList::new()]), - Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + Err(LeafBuildError::PoisonedCandidateList { point: 0 }) )); assert!(matches!( - candidates.into_rows(), - Err(LeafBuildError::PoisonedCandidateRow { point: 0 }) + candidates.into_lists(), + Err(LeafBuildError::PoisonedCandidateList { point: 0 }) )); } @@ -400,5 +403,8 @@ fn direct_candidate_accumulator_keeps_unique_sorted_rows() { ], ) .unwrap(); - assert_eq!(rows(candidates.into_rows().unwrap()), [vec![1], vec![0]]); + assert_eq!( + adjacency_lists(candidates.into_lists().unwrap()), + [vec![1], vec![0]] + ); } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 84659095ba..88867ce108 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -5,7 +5,7 @@ //! Deterministic overlapping partition construction for PiPNN. //! -//! The stage maps real dataset rows to bounded leaf ID lists. Numerical work +//! The stage maps real dataset points to bounded leaf ID lists. Numerical work //! reuses the partition kernel and dense GEMM. A stage-owned pool leases scratch //! to Rayon chunks and takes it back after each chunk; computation never holds //! the pool lock, and no thread-local cleanup protocol is required. @@ -43,14 +43,14 @@ use crate::{utils::VectorRepr, ANNError, ANNResult}; use diskann_linalg::Transpose; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, - views::MatrixView, + views::{MatrixView, MutMatrixView}, }; use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; use rand::{prelude::IndexedRandom, SeedableRng}; use rayon::prelude::*; use crate::{ - partition_kernel::{nearest_leaders, PartitionTopK}, + partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, PiPNNConfig, }; @@ -59,8 +59,8 @@ const PARTITION_SEED: u64 = 1_000; const REPLICA_SEED_STEP: u64 = 7_919; const LEADER_CAP: usize = 1_000; const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; -const MIN_ASSIGNMENT_STRIPE_ROWS: usize = 32; -const MAX_ASSIGNMENT_STRIPE_ROWS: usize = 1_024; +const MIN_ASSIGNMENT_STRIPE_POINTS: usize = 32; +const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const MAX_PARTITION_ITERATIONS: usize = 30; @@ -71,7 +71,7 @@ pub(crate) enum PartitionError { EmptyDataset, #[error("PiPNN cannot partition vectors with zero dimensions")] EmptyDimensions, - #[error("dataset has {0} rows, which exceeds the u32 ID limit")] + #[error("dataset has {0} points, which exceeds the u32 ID limit")] TooManyPoints(usize), #[error("{buffer} shape {rows} x {cols} overflows usize")] ShapeOverflow { @@ -110,7 +110,7 @@ struct WorkItem { struct StripeBuffers { points: Vec, dots: Vec, - row_scales: Vec, + point_scales: Vec, } impl AsPooled<()> for StripeBuffers { @@ -159,10 +159,15 @@ where } let mut leaves = Vec::new(); + // Prepare metric and ISA dispatch before replicas spawn Rayon work. The + // Copy handle is shared read-only; every stripe calls its direct function + // pointer instead of redispatching in the recursive hot path. + let kernel = PartitionKernel::new(metric); let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas { let seed = replica_seed(replica); - let mut replica_leaves = partition_replica(data, &config, metric, seed, &stripe_buffers)?; + let mut replica_leaves = + partition_replica(data, &config, metric, &kernel, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::opaque)?; @@ -176,6 +181,7 @@ fn partition_replica( data: MatrixView<'_, T>, config: &PiPNNConfig, metric: Metric, + kernel: &PartitionKernel, seed: u64, stripe_buffers: &StripeBufferPool, ) -> ANNResult>> @@ -220,6 +226,7 @@ where data, config, metric, + kernel, item, stripe_buffers, )); @@ -258,6 +265,7 @@ fn partition_one_level( data: MatrixView<'_, T>, config: &PiPNNConfig, metric: Metric, + kernel: &PartitionKernel, item: WorkItem, stripe_buffers: &StripeBufferPool, ) -> ANNResult<(Vec, Vec>)> @@ -277,6 +285,7 @@ where &leaders, fanout, metric, + kernel, stripe_buffers, )?; @@ -334,110 +343,117 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { /// Assign each point to its nearest `fanout` sampled leaders. /// -/// Leader rows are gathered once. Point rows are processed in cache-sized +/// Leader vectors are gathered once. Points are processed in cache-sized /// stripes, while a worker chunk retains one leased scratch buffer across all of /// its stripes. The flat assignment matrix preserves point order and is then /// scattered into per-leader clusters; preserving order is required for fixed /// seed determinism in later recursion levels. fn assign_to_leaders( data: MatrixView<'_, T>, - points: &[u32], - leaders: &[u32], + point_ids: &[u32], + leader_ids: &[u32], fanout: usize, metric: Metric, + kernel: &PartitionKernel, stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where T: VectorRepr + Send + Sync, { - let dimensions = data.ncols(); - let leader_values_len = checked_area("leader data", leaders.len(), dimensions)?; + let dimension_count = data.ncols(); + let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; - gather_rows(data, leaders, &mut leader_values)?; + gather_vectors(data, leader_ids, &mut leader_values)?; let mut leader_scales = if matches!(metric, Metric::L2 | Metric::Cosine) { - filled_vec(leaders.len(), 0.0f32)? + filled_vec(leader_ids.len(), 0.0f32)? } else { Vec::new() }; - for (scale, row) in leader_scales + for (scale, leader_vector) in leader_scales .iter_mut() - .zip(leader_values.chunks_exact(dimensions)) + .zip(leader_values.chunks_exact(dimension_count)) { // Leader norms participate in the top-k ordering. Preserve the original // scalar reduction order: reassociating this short setup pass through a // SIMD norm changes low bits and can send near-tied points down different // recursive partition paths. - *scale = row.iter().map(|value| value * value).sum(); + *scale = leader_vector.iter().map(|value| value * value).sum(); if metric == Metric::Cosine { *scale = scale.sqrt(); } } - let fanout = fanout.min(leaders.len()); - let assignment_len = checked_area("partition assignments", points.len(), fanout)?; + let fanout = fanout.min(leader_ids.len()); + let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; let mut assignments = filled_vec(assignment_len, 0u32)?; - let stripe_rows = assignment_stripe_rows(leaders.len()); - let assignment_stripe = checked_area("assignment stripe", stripe_rows, fanout)?; - let stripes = points.len().div_ceil(stripe_rows); - let worker_stripes = stripes.div_ceil(rayon::current_num_threads().max(1)); - let worker_rows = checked_area("assignment worker", worker_stripes, stripe_rows)?; - let worker_assignment = checked_area("assignment worker", worker_rows, fanout)?; + let stripe_points = assignment_stripe_point_count(leader_ids.len()); + let stripe_assignment_count = checked_area("assignment stripe", stripe_points, fanout)?; + let stripe_count = point_ids.len().div_ceil(stripe_points); + let worker_stripe_count = stripe_count.div_ceil(rayon::current_num_threads().max(1)); + let worker_point_count = checked_area("assignment worker", worker_stripe_count, stripe_points)?; + let worker_assignment_count = checked_area("assignment worker", worker_point_count, fanout)?; // Each worker chunk owns one scratch value and reuses it for its stripes. // build_graph pins this terminal operation to the caller-owned pool. #[allow(clippy::disallowed_methods)] assignments - .par_chunks_mut(worker_assignment) + .par_chunks_mut(worker_assignment_count) .enumerate() - .try_for_each(|(worker, worker_output)| { + .try_for_each(|(worker, worker_assignments)| { let mut buffers = stripe_buffers.get_ref(()); - let worker_first = worker * worker_rows; - for (stripe, output) in worker_output.chunks_mut(assignment_stripe).enumerate() { - let first = worker_first + stripe * stripe_rows; - let rows = output.len() / fanout; + let worker_first = worker * worker_point_count; + for (stripe, stripe_assignments) in worker_assignments + .chunks_mut(stripe_assignment_count) + .enumerate() + { + let first_point = worker_first + stripe * stripe_points; + let stripe_point_count = stripe_assignments.len() / fanout; assign_stripe( data, - &points[first..first + rows], + &point_ids[first_point..first_point + stripe_point_count], &leader_values, &leader_scales, metric, + kernel, fanout, &mut buffers, - output, + stripe_assignments, )?; } Ok::<(), ANNError>(()) })?; - scatter_assignments(points, &assignments, fanout, leaders.len()) + scatter_assignments(point_ids, &assignments, fanout, leader_ids.len()) } #[inline] #[allow(clippy::too_many_arguments)] fn assign_stripe( data: MatrixView<'_, T>, - points: &[u32], + point_ids: &[u32], leader_values: &[f32], leader_scales: &[f32], metric: Metric, + kernel: &PartitionKernel, fanout: usize, buffers: &mut StripeBuffers, - output: &mut [u32], + assignments: &mut [u32], ) -> ANNResult<()> where T: VectorRepr, { - let rows = points.len(); + let point_count = point_ids.len(); let dimensions = data.ncols(); - let leaders = leader_values.len() / dimensions; - let point_values_len = checked_area("point stripe", rows, dimensions)?; - let dots_len = checked_area("dot-product stripe", rows, leaders)?; + let leader_count = leader_values.len() / dimensions; + let point_values_len = checked_area("point stripe", point_count, dimensions)?; + let dots_len = checked_area("dot-product stripe", point_count, leader_count)?; + let output_len = checked_area("partition assignments", point_count, fanout)?; // Scratch keeps its high-water length and every consumer receives an // explicit active prefix. Resizing to the exact stripe shape would be // correct but re-zeroes the buffer whenever a pooled value moves between - // work items with different leader counts: `stripe_rows` is derived from - // `leaders`, so the point buffer swings between roughly 768 KiB and 6 MiB + // work items with different leader counts: `stripe_points` is derived from + // `leader_count`, so the point buffer swings between roughly 768 KiB and 6 MiB // and `Vec::resize` only truncates on the way down, then memsets the whole // delta on the way back up. grow_fallible(&mut buffers.points, point_values_len, 0.0)?; @@ -445,16 +461,16 @@ where let StripeBuffers { points: point_buffer, dots: dot_buffer, - row_scales: row_scale_buffer, + point_scales: point_scale_buffer, } = buffers; let point_values = &mut point_buffer[..point_values_len]; let dots = &mut dot_buffer[..dots_len]; - gather_rows(data, points, point_values)?; + gather_vectors(data, point_ids, point_values)?; diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, - rows, - leaders, + point_count, + leader_count, dimensions, 1.0, point_values, @@ -464,35 +480,49 @@ where ) .map_err(ANNError::opaque)?; - let row_scales = if metric == Metric::Cosine { - grow_fallible(row_scale_buffer, rows, 0.0)?; - let row_scales = &mut row_scale_buffer[..rows]; - for (scale, row) in row_scales + let point_scales = if metric == Metric::Cosine { + grow_fallible(point_scale_buffer, point_count, 0.0)?; + let point_scales = &mut point_scale_buffer[..point_count]; + for (scale, point_values) in point_scales .iter_mut() .zip(point_values.chunks_exact(dimensions)) { - *scale = FastL2NormSquared.evaluate(row); + *scale = FastL2NormSquared.evaluate(point_values); } - &*row_scales + &*point_scales } else { &[] }; - nearest_leaders( - PartitionTopK { - dots, - rows, - leaders, - row_scales, - leader_scales, - metric, + let scales = match metric { + Metric::L2 => PartitionScales::L2 { + leader_squared_norms: leader_scales, }, - fanout, - output, - ) - .map_err(ANNError::opaque) + Metric::Cosine => PartitionScales::Cosine { + point_squared_norms: point_scales, + leader_norms: leader_scales, + }, + Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, + }; + let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { + ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "dot-product stripe", + expected: dots_len, + actual: dots.len(), + }) + })?; + let output = MutMatrixView::try_from(assignments, point_count, fanout).map_err(|error| { + ANNError::opaque(PartitionError::InvalidBufferLength { + buffer: "partition assignments", + expected: output_len, + actual: error.into_inner().len(), + }) + })?; + kernel + .nearest_leaders(PartitionInput { dots, scales }, output) + .map_err(ANNError::opaque) } -fn gather_rows(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> +fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> where T: VectorRepr, { @@ -504,8 +534,8 @@ where actual: output.len(), })); } - for (&index, row) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { - T::as_f32_into(data.row(index as usize), row).map_err(Into::::into)?; + for (&index, vector_output) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { + T::as_f32_into(data.row(index as usize), vector_output).map_err(Into::::into)?; } Ok(()) } @@ -526,9 +556,9 @@ fn scatter_assignments( return scatter_serial(points, assignments, fanout, leaders); } - let stripe_rows = points.len().div_ceil(rayon::current_num_threads().max(1)); - let assignment_stripe = checked_area("scatter assignment stripe", stripe_rows, fanout)?; - let stripes = points.len().div_ceil(stripe_rows); + let stripe_points = points.len().div_ceil(rayon::current_num_threads().max(1)); + let stripe_assignment_count = checked_area("scatter assignment stripe", stripe_points, fanout)?; + let stripes = points.len().div_ceil(stripe_points); let mut partials = Vec::new(); partials .try_reserve_exact(stripes) @@ -540,8 +570,8 @@ fn scatter_assignments( .par_iter_mut() .zip( points - .par_chunks(stripe_rows) - .zip(assignments.par_chunks(assignment_stripe)), + .par_chunks(stripe_points) + .zip(assignments.par_chunks(stripe_assignment_count)), ) .for_each(|(slot, (points, assignments))| { *slot = Some(scatter_serial(points, assignments, fanout, leaders)); @@ -608,8 +638,8 @@ fn scatter_serial( })?; } let mut clusters = clusters_with_capacities(&sizes)?; - for (&point, row) in points.iter().zip(assignments.chunks_exact(fanout)) { - for &leader in row { + for (&point, point_assignments) in points.iter().zip(assignments.chunks_exact(fanout)) { + for &leader in point_assignments { clusters[leader as usize].push(point); } } @@ -759,14 +789,14 @@ fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult usize { - let rows = ASSIGNMENT_CACHE_TARGET_BYTES / (leaders.max(1) * size_of::()); - let rows = if rows.is_power_of_two() { - rows +fn assignment_stripe_point_count(leader_count: usize) -> usize { + let point_count = ASSIGNMENT_CACHE_TARGET_BYTES / (leader_count.max(1) * size_of::()); + let point_count = if point_count.is_power_of_two() { + point_count } else { - rows.next_power_of_two() / 2 + point_count.next_power_of_two() / 2 }; - rows.clamp(MIN_ASSIGNMENT_STRIPE_ROWS, MAX_ASSIGNMENT_STRIPE_ROWS) + point_count.clamp(MIN_ASSIGNMENT_STRIPE_POINTS, MAX_ASSIGNMENT_STRIPE_POINTS) } #[cfg(test)] diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 31f3e4bcad..7254ed25d9 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -24,10 +24,10 @@ fn clustered_data(points: usize, dimensions: usize) -> Matrix { diskann_utils::views::Init({ let mut position = 0usize; move || { - let row = position / dimensions; - let column = position % dimensions; + let point = position / dimensions; + let dimension = position % dimensions; position += 1; - (row / 8) as f32 * 10.0 + column as f32 * 0.01 + row as f32 * 0.001 + (point / 8) as f32 * 10.0 + dimension as f32 * 0.01 + point as f32 * 0.001 } }), points, @@ -40,11 +40,11 @@ fn directional_data(points: usize, dimensions: usize) -> Matrix { diskann_utils::views::Init({ let mut position = 0usize; move || { - let row = position / dimensions; - let column = position % dimensions; + let point = position / dimensions; + let dimension = position % dimensions; position += 1; - let angle = std::f32::consts::TAU * row as f32 / points as f32; - match column { + let angle = std::f32::consts::TAU * point as f32 / points as f32; + match dimension { 0 => angle.cos(), 1 => angle.sin(), _ => 0.0, @@ -179,14 +179,14 @@ where T: diskann::utils::VectorRepr + Send + Sync, { let points = 64; - // Partition gathering converts source rows before GEMM. Exercise conversion + // Partition gathering converts source vectors before GEMM. Exercise conversion // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let raw: Vec = (0..points * dimensions) .map(|index| { - let row = index / dimensions; - let column = index % dimensions; - ((row * 5 + column * 7 + row * column) % 23) as u8 + let point = index / dimensions; + let dimension = index % dimensions; + ((point * 5 + dimension * 7 + point * dimension) % 23) as u8 }) .collect(); let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); @@ -261,6 +261,7 @@ fn l2_leader_norms_preserve_scalar_reduction_order() { &[0, 1], 1, Metric::L2, + &PartitionKernel::new(Metric::L2), &StripeBufferPool::new((), 0, None), ) .unwrap(); @@ -298,10 +299,13 @@ fn replica_seed_derivation_is_stable_and_distinct() { } #[test] -fn assignment_stripes_use_power_of_two_row_counts() { - assert_eq!(assignment_stripe_rows(1_000), 128); - assert_eq!(assignment_stripe_rows(256), 512); - assert_eq!(assignment_stripe_rows(1), MAX_ASSIGNMENT_STRIPE_ROWS); +fn assignment_stripes_use_power_of_two_point_counts() { + assert_eq!(assignment_stripe_point_count(1_000), 128); + assert_eq!(assignment_stripe_point_count(256), 512); + assert_eq!( + assignment_stripe_point_count(1), + MAX_ASSIGNMENT_STRIPE_POINTS + ); } #[test] @@ -331,6 +335,7 @@ fn leader_assignment_handles_multiple_stripes() { &[0, 2_047], 1, Metric::L2, + &PartitionKernel::new(Metric::L2), &StripeBufferPool::new((), 0, None), ) .unwrap(); @@ -378,7 +383,7 @@ fn rejects_zero_dimensions() { #[test] fn rejects_invalid_gather_output_length() { let data = Matrix::::new(0.0, 2, 2); - let error = gather_rows(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); + let error = gather_vectors(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); assert_eq!( error.downcast::().unwrap(), From c8b45c135b9497c3383cdd9f3322c97f86491272 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:51:03 +0000 Subject: [PATCH 18/58] refactor(pipnn): name candidate lists --- diskann/src/graph/pipnn/finalization.rs | 134 +++++++++--------- diskann/src/graph/pipnn/finalization/tests.rs | 44 ++++-- diskann/src/graph/pipnn/leaf_build/tests.rs | 4 +- diskann/src/graph/pipnn/mod.rs | 14 +- 4 files changed, 111 insertions(+), 85 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 8acb5fda46..e714d9336d 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -6,9 +6,9 @@ //! Final graph-degree enforcement through the shared Vamana RobustPrune kernel. //! //! Candidate merging may produce more than `R` IDs for a point. This stage first -//! validates every global ID, then processes rows independently in the caller's -//! Rayon pool. Rows already within the degree bound are returned without any -//! distance work. Overfull rows are converted to source-distance candidates, +//! validates every global ID, then processes each point's candidate list in the +//! caller's Rayon pool. Lists already within the degree bound are returned without +//! distance work. Overfull lists are converted to source-distance candidates, //! passed through RobustPrune, and rewritten from the selected output. //! //! The shared kernel owns occlusion and alpha-round semantics; this adapter owns @@ -16,24 +16,24 @@ //! representation. //! //! ```text -//! candidate rows ──> validate row count and every global ID +//! candidate lists ──> validate point count and every global ID //! │ //! ┌────────────┴────────────┐ //! v v //! len <= R len > R -//! return row source-distance candidates +//! return list source-distance candidates //! │ //! v //! shared RobustPrune //! │ //! v -//! rewrite same row owner +//! rewrite same list owner //! ``` //! //! | Path | Distance evaluations | Allocation behavior | //! | --- | --- | --- | -//! | bounded row | none | move row directly to output | -//! | overfull row | source and occlusion distances | reuse Rayon-job workspace | +//! | bounded list | none | move list directly to output | +//! | overfull list | source and occlusion distances | reuse Rayon-job workspace | use std::convert::Infallible; @@ -49,29 +49,31 @@ use rayon::prelude::*; #[derive(Debug, thiserror::Error)] pub(crate) enum FinalizationError { - #[error("candidate row count {rows} does not match the dataset point count {points}")] - RowCountMismatch { rows: usize, points: usize }, - #[error("candidate ID {candidate} in row {row} is outside a {points}-point dataset")] + #[error("candidate list count {lists} does not match the dataset point count {points}")] + CandidateListCountMismatch { lists: usize, points: usize }, + #[error( + "candidate ID {candidate} for source {source_index} is outside a {points}-point dataset" + )] InvalidCandidateId { - row: usize, + source_index: usize, candidate: u32, points: usize, }, } -/// Per-Rayon-job state retained across rows. +/// Per-Rayon-job state retained across source points. /// /// `prune` owns candidate/state/output buffers. `cache` stores provider lookup /// results required by the shared kernel. Reusing both avoids per-node /// allocations, which would otherwise dominate finalization for millions of -/// short rows. +/// short candidate lists. #[derive(Default)] struct Workspace { prune: prune::Scratch, cache: Vec<(f32, Option)>, } -/// Validate candidate IDs and prune only rows whose length exceeds graph degree. +/// Validate candidate IDs and prune only lists whose length exceeds graph degree. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, @@ -81,7 +83,7 @@ pub(crate) fn prune_overfull( where T: VectorRepr + Send + Sync, { - validate_candidates(&candidates, data.nrows()).map_err(ANNError::opaque)?; + validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::opaque)?; let degree = graph.pruned_degree().get(); let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind(), false); @@ -92,70 +94,72 @@ where candidates .into_par_iter() .enumerate() - .map_init(Workspace::default, |workspace, (source, mut row)| { - // Candidate accumulators already enforce uniqueness. A bounded row - // therefore satisfies the graph policy without distance evaluation. - if row.len() <= degree { - return Ok(row); - } + .map_init( + Workspace::default, + |workspace, (source, mut source_candidates)| { + // Candidate accumulators already enforce uniqueness. A bounded list + // therefore satisfies the graph policy without distance evaluation. + if source_candidates.len() <= degree { + return Ok(source_candidates); + } - let source_id = u32::try_from(source).map_err(ANNError::opaque)?; - let source_vector = data.row(source); - let pool = workspace.prune.candidates_mut(); - pool.clear(); - pool.try_reserve(row.len()).map_err(ANNError::opaque)?; - pool.extend(row.iter().copied().map(|candidate| { - Neighbor::new( - candidate, - distance.evaluate_similarity(source_vector, data.row(candidate as usize)), - ) - })); - // as_context sorts the active candidate prefix by source distance. - // The callback below is needed only for selected-to-candidate - // occlusion checks; dimension specialization stays in `distance`. - let candidate_count = pool.len(); - let mut context = workspace.prune.as_context(candidate_count); - prune::robust_prune( - &mut context, - policy, - &mut workspace.cache, - Some, - |left, right| { - Ok::<_, Infallible>( - distance.evaluate_similarity( + let source_id = u32::try_from(source).map_err(ANNError::opaque)?; + let source_vector = data.row(source); + let pool = workspace.prune.candidates_mut(); + pool.clear(); + pool.try_reserve(source_candidates.len()) + .map_err(ANNError::opaque)?; + pool.extend(source_candidates.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + // as_context sorts the active candidate prefix by source distance. + // The callback below is needed only for selected-to-candidate + // occlusion checks; dimension specialization stays in `distance`. + let candidate_count = pool.len(); + let mut context = workspace.prune.as_context(candidate_count); + prune::robust_prune( + &mut context, + policy, + &mut workspace.cache, + Some, + |left, right| { + Ok::<_, Infallible>(distance.evaluate_similarity( data.row(*left as usize), data.row(*right as usize), - ), - ) - }, - |id| id == source_id, - ) - .map_err(ANNError::opaque)?; + )) + }, + |id| id == source_id, + ) + .map_err(ANNError::opaque)?; - // RobustPrune selects distinct candidate positions, so its output IDs - // are unique by construction. `extend_from_slice` would re-derive that - // with an O(degree^2) membership scan per row; the trusted overwrite is - // a copy and still verifies uniqueness under debug assertions. - row.overwrite_trusted(workspace.prune.neighbors()); - Ok(row) - }) + // RobustPrune selects distinct candidate positions, so its output IDs + // are unique by construction. `extend_from_slice` would re-derive that + // with an O(degree^2) membership scan per list; the trusted overwrite + // is a copy and still verifies uniqueness under debug assertions. + source_candidates.overwrite_trusted(workspace.prune.neighbors()); + Ok(source_candidates) + }, + ) .collect() } -fn validate_candidates( +fn validate_candidate_lists( candidates: &[AdjacencyList], points: usize, ) -> Result<(), FinalizationError> { if candidates.len() != points { - return Err(FinalizationError::RowCountMismatch { - rows: candidates.len(), + return Err(FinalizationError::CandidateListCountMismatch { + lists: candidates.len(), points, }); } - for (row_id, row) in candidates.iter().enumerate() { - if let Some(&candidate) = row.iter().find(|&&id| id as usize >= points) { + for (source, source_candidates) in candidates.iter().enumerate() { + if let Some(&candidate) = source_candidates.iter().find(|&&id| id as usize >= points) { return Err(FinalizationError::InvalidCandidateId { - row: row_id, + source_index: source, candidate, points, }); diff --git a/diskann/src/graph/pipnn/finalization/tests.rs b/diskann/src/graph/pipnn/finalization/tests.rs index fde3a6e0ea..42a534bb45 100644 --- a/diskann/src/graph/pipnn/finalization/tests.rs +++ b/diskann/src/graph/pipnn/finalization/tests.rs @@ -25,15 +25,20 @@ fn graph_config(degree: usize) -> Config { .unwrap() } -fn row(ids: impl IntoIterator) -> AdjacencyList { +fn candidate_list(ids: impl IntoIterator) -> AdjacencyList { AdjacencyList::from_iter_untrusted(ids) } #[test] -fn preserves_rows_within_the_degree_bound() { +fn preserves_lists_within_the_degree_bound() { let data = [0.0_f32, 1.0, 2.0, 3.0]; let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let candidates = vec![row([3, 1]), row([]), row([]), row([])]; + let candidates = vec![ + candidate_list([3, 1]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); @@ -41,10 +46,15 @@ fn preserves_rows_within_the_degree_bound() { } #[test] -fn prunes_an_overfull_row_with_the_vamana_kernel() { +fn prunes_an_overfull_list_with_the_vamana_kernel() { let data = [0.0_f32, 1.0, 2.0, -3.0]; let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let candidates = vec![row([3, 2, 1]), row([]), row([]), row([])]; + let candidates = vec![ + candidate_list([3, 2, 1]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); @@ -56,14 +66,18 @@ fn prunes_an_overfull_row_with_the_vamana_kernel() { fn rejects_invalid_candidate_ids_without_panicking() { let data = [0.0_f32, 1.0, 2.0]; let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![row([1, 3]), row([]), row([])]; + let candidates = vec![ + candidate_list([1, 3]), + candidate_list([]), + candidate_list([]), + ]; let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); assert!(matches!( error.downcast_ref::(), Some(FinalizationError::InvalidCandidateId { - row: 0, + source_index: 0, candidate: 3, points: 3, }) @@ -71,16 +85,24 @@ fn rejects_invalid_candidate_ids_without_panicking() { } #[test] -fn rejects_candidate_row_count_mismatch_without_panicking() { +fn rejects_candidate_list_count_mismatch_without_panicking() { let data = [0.0_f32, 1.0, 2.0]; let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![row([]), row([]), row([]), row([])]; + let candidates = vec![ + candidate_list([]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); assert!(matches!( error.downcast_ref::(), - Some(FinalizationError::RowCountMismatch { rows: 4, points: 3 }) + Some(FinalizationError::CandidateListCountMismatch { + lists: 4, + points: 3 + }) )); } @@ -90,7 +112,7 @@ fn rejects_more_candidates_than_the_shared_position_type_can_represent() { let data = vec![0.0_f32; count + 1]; let data = MatrixView::try_from(&data[..], count + 1, 1).unwrap(); let mut candidates = Vec::with_capacity(count + 1); - candidates.push(row(1..=count as u32)); + candidates.push(candidate_list(1..=count as u32)); candidates.resize_with(count + 1, AdjacencyList::new); let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index 2d4e66c2e6..5d6cd92e94 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -322,7 +322,7 @@ fn reports_shape_overflow_before_allocating() { } #[test] -fn rejects_an_invalid_kernel_position() { +fn rejects_an_invalid_kernel_target() { let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; let error = add_symmetric_neighbors( &[10, 20], @@ -392,7 +392,7 @@ fn allocation_errors_preserve_buffer_context() { } #[test] -fn direct_candidate_accumulator_keeps_unique_sorted_rows() { +fn direct_candidate_accumulator_keeps_unique_sorted_lists() { let candidates = DirectCandidates::new(2).unwrap(); candidates .add_leaf( diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 2b0c555be9..49b5e4ec25 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -257,11 +257,11 @@ impl<'a> PiPNNBuildContext<'a> { } } -/// Build PiPNN adjacency for real rows in `data`. +/// Build PiPNN adjacency for real points in `data`. /// /// This is the core algorithm boundary. Search entry-point selection, frozen nodes, /// providers, serialization, and index writers belong to the outer build pipelines. -/// For raw `u8` and `i8` rows, `CosineNormalized` is evaluated as `Cosine` because +/// For raw `u8` and `i8` vectors, `CosineNormalized` is evaluated as `Cosine` because /// those representations are converted to f32 scratch but are not unit-normalized. pub fn build_graph( data: MatrixView<'_, T>, @@ -282,7 +282,7 @@ where { if data.nrows() == 0 { return Err(ANNError::log_dimension_mismatch_error( - "PiPNN requires at least one data row".into(), + "PiPNN requires at least one data point".into(), )); } if data.ncols() == 0 { @@ -292,7 +292,7 @@ where } if data.nrows() > u32::MAX as usize { return Err(config_error(format!( - "dataset row count ({}) exceeds the u32 graph ID limit", + "dataset point count ({}) exceeds the u32 graph ID limit", data.nrows() ))); } @@ -303,19 +303,19 @@ where data.ncols() )) })?; - // Integer source rows are not guaranteed unit-normalized after conversion, + // Integer source vectors are not guaranteed unit-normalized after conversion, // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition(data, context.config.clone(), metric))?; - // `leaves` is consumed here. Workers borrow individual ID rows during the + // `leaves` is consumed here. Workers borrow individual ID lists during the // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) .map_err(ANNError::opaque) })?; - // Finalization consumes candidate rows and reuses their allocations for the + // Finalization consumes candidate lists and reuses their allocations for the // resulting adjacency where possible. tracing::info_span!("pipnn.finalization") .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) From 7d2906628df780b5d874dac3b85da9100ba88ee1 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 07:27:37 +0000 Subject: [PATCH 19/58] refactor(pipnn): complete graph module migration 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. --- diskann-pipnn/README.md | 19 ------- diskann/src/graph/pipnn/finalization.rs | 6 +-- diskann/src/graph/pipnn/finalization/tests.rs | 23 +------- diskann/src/graph/pipnn/leaf_build.rs | 8 +-- diskann/src/graph/pipnn/leaf_build/tests.rs | 52 ++++++++++--------- diskann/src/graph/pipnn/mod.rs | 6 +-- diskann/src/graph/pipnn/partitioning.rs | 42 +++++++-------- diskann/src/graph/pipnn/partitioning/tests.rs | 14 ++--- .../{build_graph.rs => pipnn_build_graph.rs} | 12 +++-- diskann/tests/{config.rs => pipnn_config.rs} | 9 +--- 10 files changed, 77 insertions(+), 114 deletions(-) delete mode 100644 diskann-pipnn/README.md rename diskann/tests/{build_graph.rs => pipnn_build_graph.rs} (96%) rename diskann/tests/{config.rs => pipnn_config.rs} (93%) diff --git a/diskann-pipnn/README.md b/diskann-pipnn/README.md deleted file mode 100644 index 4ac1618f03..0000000000 --- a/diskann-pipnn/README.md +++ /dev/null @@ -1,19 +0,0 @@ -# PiPNN graph construction - -This crate implements the graph-construction stages from [PiPNN: Pick in Partitions for Fast and Accurate ANN Graph Construction](https://arxiv.org/html/2602.21247v1). - -## Boundary - -PiPNN core consumes a dense `MatrixView`, graph policy, and a caller-owned Rayon pool, then returns adjacency lists for the dataset's real point IDs. It does not own start or frozen points, vector or neighbor providers, PQ, disk headers, serialization, or search. Those concerns remain in the outer in-memory and disk pipelines. - -The dense view is intentional: partition assignment and leaf all-pairs kernels operate over the whole source matrix. Materializing provider state inside the algorithm would couple numerical graph construction to storage lifecycle and would require a second dataset copy. Integrations should finish PiPNN scratch before allocating or populating their searchable provider. - -## Policy ownership - -- `PiPNNConfig` owns partition and leaf-selection parameters: leaf bounds, sampling fraction, fanout levels, leaf `k`, and replicas. -- DiskANN graph configuration owns metric, output degree, build-L, alpha, and prune policy. -- Candidate-merging policies are separate validated options; they must not make graph policy fields redundant or silently cap the requested degree. - -## Execution - -A build runs partitioning, leaf construction, candidate merging, then graph finalization. All parallel work executes in the supplied pool. Per-job scratch is initialized through Rayon and is released through normal ownership when its stage completes; the core has no global thread-local buffers or cleanup broadcasts. diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index e714d9336d..2d0451dc73 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -38,13 +38,13 @@ use std::convert::Infallible; use crate::{ - graph::{prune, AdjacencyList, Config}, + ANNError, ANNResult, + graph::{AdjacencyList, Config, prune}, neighbor::Neighbor, utils::VectorRepr, - ANNError, ANNResult, }; use diskann_utils::views::MatrixView; -use diskann_vector::{distance::Metric, DistanceFunction}; +use diskann_vector::{DistanceFunction, distance::Metric}; use rayon::prelude::*; #[derive(Debug, thiserror::Error)] diff --git a/diskann/src/graph/pipnn/finalization/tests.rs b/diskann/src/graph/pipnn/finalization/tests.rs index 42a534bb45..66562049c2 100644 --- a/diskann/src/graph/pipnn/finalization/tests.rs +++ b/diskann/src/graph/pipnn/finalization/tests.rs @@ -4,8 +4,8 @@ */ use crate::graph::{ - config::{self, MaxDegree}, AdjacencyList, + config::{self, MaxDegree}, }; use diskann_utils::views::MatrixView; @@ -58,8 +58,7 @@ fn prunes_an_overfull_list_with_the_vamana_kernel() { let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - assert!(actual[0].len() <= 2); - assert!(actual[0].contains(1)); + assert_eq!(&*actual[0], &[1, 3]); } #[test] @@ -105,21 +104,3 @@ fn rejects_candidate_list_count_mismatch_without_panicking() { }) )); } - -#[test] -fn rejects_more_candidates_than_the_shared_position_type_can_represent() { - let count = u16::MAX as usize + 1; - let data = vec![0.0_f32; count + 1]; - let data = MatrixView::try_from(&data[..], count + 1, 1).unwrap(); - let mut candidates = Vec::with_capacity(count + 1); - candidates.push(candidate_list(1..=count as u32)); - candidates.resize_with(count + 1, AdjacencyList::new); - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::>(), - Some(prune::RobustPruneError::TooManyCandidates { actual, max }) - if *actual == count && *max == u16::MAX as usize - )); -} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 37cf5e3d14..ec8763b190 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -27,9 +27,9 @@ use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; use rayon::prelude::*; -use crate::leaf_kernel::{ - leaf_neighbor_count, leaf_output_len, LeafInput, LeafKernel, LeafKernelError, - LeafKernelWorkspace, LeafNeighbor, +use super::leaf_kernel::{ + LeafInput, LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, + leaf_output_len, }; /// Failure while converting leaves into direct graph candidates. @@ -67,7 +67,7 @@ pub(crate) enum LeafBuildError { leaf: usize, point: u32, #[source] - source: diskann::ANNError, + source: crate::ANNError, }, #[error("lower-AAT failed for leaf {leaf}")] LowerAat { diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs index 5d6cd92e94..6c55e6a577 100644 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ b/diskann/src/graph/pipnn/leaf_build/tests.rs @@ -9,8 +9,8 @@ use half::f16; use std::collections::BTreeSet; use super::{ - add_symmetric_neighbors, allocation_error, build_leaf_candidates, DirectCandidates, - LeafBuffers, LeafBuildError, + DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, + build_leaf_candidates, }; fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { @@ -29,14 +29,14 @@ fn build( leaves: &[Vec], k: usize, metric: Metric, -) -> Result>, LeafBuildError> +) -> Result>, LeafBuildError> where - T: diskann::utils::VectorRepr + 'static, + T: crate::utils::VectorRepr + 'static, { pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) } -fn adjacency_lists(graph: Vec>) -> Vec> { +fn adjacency_lists(graph: Vec>) -> Vec> { graph.into_iter().map(Vec::from).collect() } @@ -150,7 +150,7 @@ fn global_id_translation_is_independent_of_leaf_order() { fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> where - T: diskann::utils::VectorRepr + 'static, + T: crate::utils::VectorRepr + 'static, { let leaves = vec![(0..points as u32).collect()]; adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) @@ -158,7 +158,7 @@ where fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) where - T: diskann::utils::VectorRepr + 'static, + T: crate::utils::VectorRepr + 'static, { let points = 8; // Source dimension controls VectorRepr conversion chunking. Cover tails on @@ -213,9 +213,11 @@ fn all_metrics_produce_symmetric_unique_non_self_candidates() { let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); for (source, neighbors) in graph.iter().enumerate() { assert!(neighbors.iter().all(|&target| target as usize != source)); - assert!(neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32))); + assert!( + neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + ); assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); } } @@ -240,7 +242,7 @@ fn parallel_leaf_schedule_does_not_change_candidate_order() { } #[test] -fn rejects_invalid_shape_inputs_without_panicking() { +fn rejects_invalid_dimensions_and_leaf_membership() { let data = [0.0_f32, 1.0]; let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); assert!(matches!( @@ -288,10 +290,12 @@ fn singleton_and_zero_k_leaves_add_no_candidates() { ) .unwrap(); let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); - assert!(singleton - .iter() - .chain(&zero_k) - .all(|candidates| candidates.is_empty())); + assert!( + singleton + .iter() + .chain(&zero_k) + .all(|candidates| candidates.is_empty()) + ); } #[test] @@ -323,13 +327,13 @@ fn reports_shape_overflow_before_allocating() { #[test] fn rejects_an_invalid_kernel_target() { - let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; let error = add_symmetric_neighbors( &[10, 20], 1, &[ - crate::leaf_kernel::LeafNeighbor::new(9, 1.0), - crate::leaf_kernel::LeafNeighbor::new(0, 1.0), + super::super::leaf_kernel::LeafNeighbor::new(9, 1.0), + super::super::leaf_kernel::LeafNeighbor::new(0, 1.0), ], &mut graph, ) @@ -345,13 +349,13 @@ fn rejects_an_invalid_kernel_target() { #[test] fn skips_duplicate_global_ids_without_self_edges() { - let mut graph = vec![diskann::graph::AdjacencyList::new(); 2]; + let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; add_symmetric_neighbors( &[7, 7], 1, &[ - crate::leaf_kernel::LeafNeighbor::new(1, 0.0), - crate::leaf_kernel::LeafNeighbor::new(0, 0.0), + super::super::leaf_kernel::LeafNeighbor::new(1, 0.0), + super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), ], &mut graph, ) @@ -367,7 +371,7 @@ fn poisoned_candidate_lists_return_errors() { panic!("poison candidate list"); }); assert!(matches!( - candidates.add_leaf(&[0], &[diskann::graph::AdjacencyList::new()]), + candidates.add_leaf(&[0], &[crate::graph::AdjacencyList::new()]), Err(LeafBuildError::PoisonedCandidateList { point: 0 }) )); assert!(matches!( @@ -398,8 +402,8 @@ fn direct_candidate_accumulator_keeps_unique_sorted_lists() { .add_leaf( &[0, 1], &[ - diskann::graph::AdjacencyList::from_iter_untrusted([1, 1]), - diskann::graph::AdjacencyList::from_iter_untrusted([0]), + crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), + crate::graph::AdjacencyList::from_iter_untrusted([0]), ], ) .unwrap(); diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 49b5e4ec25..7a5b53839b 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,12 +3,12 @@ * Licensed under the MIT license. */ -//! Provider-independent PiPNN graph construction. +//! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! //! PiPNN means **Pick-in-Partitions Nearest Neighbors**. It builds a graph for //! approximate nearest-neighbor search: every input vector becomes one graph //! vertex, and its adjacency list stores other vectors worth visiting during a -//! later query. This crate constructs that adjacency; it does not execute queries. +//! later query. This module constructs that adjacency; it does not execute queries. //! //! Incremental builders such as Vamana find construction candidates by running //! beam search against a partially built graph: they repeatedly follow graph @@ -145,9 +145,9 @@ mod partition_kernel; mod partitioning; use crate::{ + ANNError, ANNResult, graph::{AdjacencyList, Config}, utils::VectorRepr, - ANNError, ANNResult, }; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 88867ce108..96ed91878b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -39,19 +39,19 @@ use std::collections::HashSet; -use crate::{utils::VectorRepr, ANNError, ANNResult}; +use crate::{ANNError, ANNResult, utils::VectorRepr}; use diskann_linalg::Transpose; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, }; -use diskann_vector::{distance::Metric, norm::FastL2NormSquared, Norm}; -use rand::{prelude::IndexedRandom, SeedableRng}; +use diskann_vector::{Norm, distance::Metric, norm::FastL2NormSquared}; +use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; -use crate::{ - partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, +use super::{ PiPNNConfig, + partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, }; // Private algorithm and batching constants live together. None are user policy. @@ -704,22 +704,22 @@ fn global_merge_small( if !small.is_empty() { let mut remainder = drain_sorted(&mut small)?; - if remainder.len() < c_min { - if let Some(last) = merged.last_mut() { - remainder.retain(|id| !last.contains(id)); - let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { - buffer: "small-leaf tail merge", - rows: last.len(), - cols: remainder.len(), - }) - })?; - if combined <= c_max { - last.try_reserve(remainder.len()) - .map_err(ANNError::opaque)?; - last.append(&mut remainder); - last.sort_unstable(); - } + if remainder.len() < c_min + && let Some(last) = merged.last_mut() + { + remainder.retain(|id| !last.contains(id)); + let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { + ANNError::opaque(PartitionError::ShapeOverflow { + buffer: "small-leaf tail merge", + rows: last.len(), + cols: remainder.len(), + }) + })?; + if combined <= c_max { + last.try_reserve(remainder.len()) + .map_err(ANNError::opaque)?; + last.append(&mut remainder); + last.sort_unstable(); } } if !remainder.is_empty() { diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs index 7254ed25d9..8780078ca5 100644 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ b/diskann/src/graph/pipnn/partitioning/tests.rs @@ -4,7 +4,7 @@ */ use diskann_utils::views::{Matrix, MatrixView}; -use diskann_vector::{distance::Metric, Half}; +use diskann_vector::{Half, distance::Metric}; use super::*; @@ -70,9 +70,11 @@ fn sorted_memberships(leaves: &[Vec]) -> Vec> { } fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { - assert!(leaves - .iter() - .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max)); + assert!( + leaves + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) + ); let mut counts = vec![0usize; points]; for leaf in leaves { let mut ids = leaf.clone(); @@ -110,7 +112,7 @@ fn partition_is_fixed_seed_deterministic_and_bounded() { } #[test] -fn recursion_after_fanout_levels_falls_back_to_one() { +fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { let data = clustered_data(80, 4); let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); @@ -176,7 +178,7 @@ fn replicas_cover_every_point_once_or_more_per_replica() { fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) where - T: diskann::utils::VectorRepr + Send + Sync, + T: crate::utils::VectorRepr + Send + Sync, { let points = 64; // Partition gathering converts source vectors before GEMM. Exercise conversion diff --git a/diskann/tests/build_graph.rs b/diskann/tests/pipnn_build_graph.rs similarity index 96% rename from diskann/tests/build_graph.rs rename to diskann/tests/pipnn_build_graph.rs index 0f823778eb..12b8c0ec9a 100644 --- a/diskann/tests/build_graph.rs +++ b/diskann/tests/pipnn_build_graph.rs @@ -5,16 +5,17 @@ #![cfg(feature = "pipnn")] #![allow( + clippy::expect_used, clippy::unwrap_used, reason = "deterministic test fixture construction must abort on invalid setup" )] use diskann::graph::config::{self, MaxDegree}; -use diskann::graph::pipnn::{build_graph, PiPNNBuildContext, PiPNNConfig}; +use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig, build_graph}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; use half::f16; -use rand::{rngs::StdRng, Rng, SeedableRng}; +use rand::{Rng, SeedableRng, rngs::StdRng}; fn pipnn_config() -> PiPNNConfig { PiPNNConfig { @@ -58,9 +59,10 @@ fn assert_graph_invariants( sorted.sort_unstable(); sorted.dedup(); assert_eq!(sorted.len(), row.len()); - assert!(row - .iter() - .all(|&id| (id as usize) < points && id as usize != source)); + assert!( + row.iter() + .all(|&id| (id as usize) < points && id as usize != source) + ); } } diff --git a/diskann/tests/config.rs b/diskann/tests/pipnn_config.rs similarity index 93% rename from diskann/tests/config.rs rename to diskann/tests/pipnn_config.rs index 7d46e9a2a5..23e569c770 100644 --- a/diskann/tests/config.rs +++ b/diskann/tests/pipnn_config.rs @@ -5,6 +5,7 @@ #![cfg(feature = "pipnn")] #![allow( + clippy::expect_used, clippy::unwrap_used, reason = "deterministic test fixture construction must abort on invalid setup" )] @@ -39,14 +40,6 @@ fn pool() -> rayon::ThreadPool { .unwrap() } -#[test] -fn accepts_the_six_algorithm_parameters_with_outer_graph_policy() { - let graph = graph_config(Metric::L2, 1.2); - let pool = pool(); - - PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); -} - #[test] fn rejects_each_invalid_algorithm_parameter() { let graph = graph_config(Metric::L2, 1.2); From 88c9c972c955084f650bc44aa4f578da8f7f2c56 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 10:46:04 +0000 Subject: [PATCH 20/58] refactor(pipnn): prepare RobustPrune inputs locally Keep sorting, workspace allocation, source exclusion, and adjacency rewriting in PiPNN finalization. The shared internal kernel now sees only prepared candidates and state. --- diskann/src/graph/pipnn/finalization.rs | 97 +++++++++++++++++-------- 1 file changed, 65 insertions(+), 32 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 2d0451dc73..9adf8c92c8 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -39,7 +39,10 @@ use std::convert::Infallible; use crate::{ ANNError, ANNResult, - graph::{AdjacencyList, Config, prune}, + graph::{ + AdjacencyList, Config, + internal::{SortedNeighbors, prune}, + }, neighbor::Neighbor, utils::VectorRepr, }; @@ -61,16 +64,15 @@ pub(crate) enum FinalizationError { }, } -/// Per-Rayon-job state retained across source points. +/// Per-Rayon-job preparation and kernel state retained across source points. /// -/// `prune` owns candidate/state/output buffers. `cache` stores provider lookup -/// results required by the shared kernel. Reusing both avoids per-node -/// allocations, which would otherwise dominate finalization for millions of -/// short candidate lists. +/// PiPNN owns sorting, allocation, and ID translation. The shared internal +/// kernel receives only the prepared candidates and an exactly sized state slice. #[derive(Default)] struct Workspace { - prune: prune::Scratch, - cache: Vec<(f32, Option)>, + pool: Vec>, + prepared: Vec>, + states: Vec, } /// Validate candidate IDs and prune only lists whose length exceeds graph degree. @@ -86,7 +88,7 @@ where validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::opaque)?; let degree = graph.pruned_degree().get(); - let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind(), false); + let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind()); let distance = T::distance(metric, Some(data.ncols())); // build_graph installs the complete call tree in the caller-owned pool. @@ -105,41 +107,72 @@ where let source_id = u32::try_from(source).map_err(ANNError::opaque)?; let source_vector = data.row(source); - let pool = workspace.prune.candidates_mut(); - pool.clear(); - pool.try_reserve(source_candidates.len()) + workspace.pool.clear(); + workspace + .pool + .try_reserve(source_candidates.len()) + .map_err(ANNError::opaque)?; + workspace + .pool + .extend(source_candidates.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance + .evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + + let candidate_count = workspace.pool.len(); + prune::validate_candidate_count::(candidate_count) + .map_err(ANNError::opaque)?; + workspace.prepared.clear(); + workspace + .prepared + .try_reserve(candidate_count) .map_err(ANNError::opaque)?; - pool.extend(source_candidates.iter().copied().map(|candidate| { - Neighbor::new( - candidate, - distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + { + // Sorting/capping precedes source exclusion so filtering cannot + // backfill with farther candidates. + let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); + workspace + .prepared + .extend(sorted.iter().filter_map(|neighbor| { + let id = *neighbor.id(); + (id != source_id) + .then(|| prune::Candidate::new(id, neighbor.distance(), id)) + })); + } + workspace + .states + .try_reserve( + workspace + .prepared + .len() + .saturating_sub(workspace.states.len()), ) - })); - // as_context sorts the active candidate prefix by source distance. - // The callback below is needed only for selected-to-candidate - // occlusion checks; dimension specialization stays in `distance`. - let candidate_count = pool.len(); - let mut context = workspace.prune.as_context(candidate_count); - prune::robust_prune( - &mut context, + .map_err(ANNError::opaque)?; + workspace + .states + .resize(workspace.prepared.len(), prune::State::default()); + + let selected = prune::robust_prune( + &workspace.prepared, + workspace.states.as_mut_slice(), policy, - &mut workspace.cache, - Some, |left, right| { Ok::<_, Infallible>(distance.evaluate_similarity( data.row(*left as usize), data.row(*right as usize), )) }, - |id| id == source_id, ) .map_err(ANNError::opaque)?; - // RobustPrune selects distinct candidate positions, so its output IDs - // are unique by construction. `extend_from_slice` would re-derive that - // with an O(degree^2) membership scan per list; the trusted overwrite - // is a copy and still verifies uniqueness under debug assertions. - source_candidates.overwrite_trusted(workspace.prune.neighbors()); + let mut guard = source_candidates.resize(selected); + for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { + *destination = *workspace.prepared[state.selected_position()].id(); + } + guard.finish(selected); Ok(source_candidates) }, ) From 63e0368e9ed62742e691bea28566a14c6fdfd29e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 5 Aug 2026 12:04:57 +0000 Subject: [PATCH 21/58] refactor(pipnn): use shared robust prune core Address the renamed internal modules and preserve graph Config alpha behavior without adding PiPNN-specific validation. --- diskann/src/graph/pipnn/finalization.rs | 7 ++++--- diskann/src/graph/pipnn/mod.rs | 6 ------ diskann/tests/pipnn_config.rs | 7 ++----- 3 files changed, 6 insertions(+), 14 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 9adf8c92c8..d46b8993a5 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -41,7 +41,7 @@ use crate::{ ANNError, ANNResult, graph::{ AdjacencyList, Config, - internal::{SortedNeighbors, prune}, + internal::{SortedNeighbors, robust_prune as prune}, }, neighbor::Neighbor, utils::VectorRepr, @@ -88,7 +88,6 @@ where validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::opaque)?; let degree = graph.pruned_degree().get(); - let policy = prune::Policy::new(degree, graph.alpha(), graph.prune_kind()); let distance = T::distance(metric, Some(data.ncols())); // build_graph installs the complete call tree in the caller-owned pool. @@ -158,7 +157,9 @@ where let selected = prune::robust_prune( &workspace.prepared, workspace.states.as_mut_slice(), - policy, + degree, + graph.alpha(), + graph.prune_kind(), |left, right| { Ok::<_, Infallible>(distance.evaluate_similarity( data.row(*left as usize), diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 7a5b53839b..7bda42c664 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -235,12 +235,6 @@ impl<'a> PiPNNBuildContext<'a> { pool: &'a ThreadPool, ) -> ANNResult { config.validate()?; - if !graph.alpha().is_finite() || graph.alpha() < 1.0 { - return Err(config_error(format!( - "graph alpha ({}) must be finite and at least 1", - graph.alpha() - ))); - } if graph.prune_kind() != metric.into() { return Err(config_error(format!( "graph prune kind {:?} is incompatible with metric {metric:?}", diff --git a/diskann/tests/pipnn_config.rs b/diskann/tests/pipnn_config.rs index 23e569c770..19a0d1ce8a 100644 --- a/diskann/tests/pipnn_config.rs +++ b/diskann/tests/pipnn_config.rs @@ -114,13 +114,10 @@ fn rejects_graph_policy_for_a_different_metric() { } #[test] -fn rejects_invalid_outer_alpha() { +fn does_not_add_alpha_validation_beyond_graph_config() { let pool = pool(); for alpha in [0.9, f32::NAN, f32::INFINITY] { let graph = graph_config(Metric::L2, alpha); - let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); - - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); - assert!(error.to_string().contains("alpha")); + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); } } From e5dfe1cdea39408d8d55091fdbc184260474a57a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:52:21 +0000 Subject: [PATCH 22/58] test(pipnn): adapt and colocate core tests Use current main error APIs and keep private plus graph/config composition tests beside their owning implementation modules. --- diskann/src/graph/pipnn/finalization.rs | 120 +++- diskann/src/graph/pipnn/finalization/tests.rs | 106 ---- diskann/src/graph/pipnn/leaf_build.rs | 413 +++++++++++++- diskann/src/graph/pipnn/leaf_build/tests.rs | 414 -------------- diskann/src/graph/pipnn/mod.rs | 390 ++++++++++++- diskann/src/graph/pipnn/partitioning.rs | 514 ++++++++++++++++-- diskann/src/graph/pipnn/partitioning/tests.rs | 423 -------------- diskann/src/graph/pipnn/tests.rs | 27 - diskann/tests/pipnn_build_graph.rs | 233 -------- diskann/tests/pipnn_config.rs | 123 ----- 10 files changed, 1365 insertions(+), 1398 deletions(-) delete mode 100644 diskann/src/graph/pipnn/finalization/tests.rs delete mode 100644 diskann/src/graph/pipnn/leaf_build/tests.rs delete mode 100644 diskann/src/graph/pipnn/partitioning/tests.rs delete mode 100644 diskann/src/graph/pipnn/tests.rs delete mode 100644 diskann/tests/pipnn_build_graph.rs delete mode 100644 diskann/tests/pipnn_config.rs diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index d46b8993a5..a10d4df027 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -85,7 +85,7 @@ pub(crate) fn prune_overfull( where T: VectorRepr + Send + Sync, { - validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::opaque)?; + validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::new)?; let degree = graph.pruned_degree().get(); let distance = T::distance(metric, Some(data.ncols())); @@ -104,13 +104,13 @@ where return Ok(source_candidates); } - let source_id = u32::try_from(source).map_err(ANNError::opaque)?; + let source_id = u32::try_from(source).map_err(ANNError::new)?; let source_vector = data.row(source); workspace.pool.clear(); workspace .pool .try_reserve(source_candidates.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; workspace .pool .extend(source_candidates.iter().copied().map(|candidate| { @@ -123,12 +123,12 @@ where let candidate_count = workspace.pool.len(); prune::validate_candidate_count::(candidate_count) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; workspace.prepared.clear(); workspace .prepared .try_reserve(candidate_count) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; { // Sorting/capping precedes source exclusion so filtering cannot // backfill with farther candidates. @@ -138,7 +138,7 @@ where .extend(sorted.iter().filter_map(|neighbor| { let id = *neighbor.id(); (id != source_id) - .then(|| prune::Candidate::new(id, neighbor.distance(), id)) + .then(|| prune::Candidate::new(id, *neighbor.distance(), id)) })); } workspace @@ -149,7 +149,7 @@ where .len() .saturating_sub(workspace.states.len()), ) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; workspace .states .resize(workspace.prepared.len(), prune::State::default()); @@ -167,7 +167,7 @@ where )) }, ) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; let mut guard = source_candidates.resize(selected); for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { @@ -203,4 +203,106 @@ fn validate_candidate_lists( } #[cfg(test)] -mod tests; +mod tests { + use crate::graph::{ + AdjacencyList, + config::{self, MaxDegree}, + }; + use diskann_utils::views::MatrixView; + + use super::*; + + fn graph_config(degree: usize) -> Config { + config::Builder::new_with( + degree, + MaxDegree::same(), + degree, + Metric::L2.into(), + |builder| { + builder.alpha(1.2); + }, + ) + .build() + .unwrap() + } + + fn candidate_list(ids: impl IntoIterator) -> AdjacencyList { + AdjacencyList::from_iter_untrusted(ids) + } + + #[test] + fn preserves_lists_within_the_degree_bound() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![ + candidate_list([3, 1]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert_eq!(&*actual[0], &[1, 3]); + } + + #[test] + fn prunes_an_overfull_list_with_the_vamana_kernel() { + let data = [0.0_f32, 1.0, 2.0, -3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let candidates = vec![ + candidate_list([3, 2, 1]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; + + let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + + assert_eq!(&*actual[0], &[1, 3]); + } + + #[test] + fn rejects_invalid_candidate_ids_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![ + candidate_list([1, 3]), + candidate_list([]), + candidate_list([]), + ]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::InvalidCandidateId { + source_index: 0, + candidate: 3, + points: 3, + }) + )); + } + + #[test] + fn rejects_candidate_list_count_mismatch_without_panicking() { + let data = [0.0_f32, 1.0, 2.0]; + let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); + let candidates = vec![ + candidate_list([]), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ]; + + let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); + + assert!(matches!( + error.downcast_ref::(), + Some(FinalizationError::CandidateListCountMismatch { + lists: 4, + points: 3 + }) + )); + } +} diff --git a/diskann/src/graph/pipnn/finalization/tests.rs b/diskann/src/graph/pipnn/finalization/tests.rs deleted file mode 100644 index 66562049c2..0000000000 --- a/diskann/src/graph/pipnn/finalization/tests.rs +++ /dev/null @@ -1,106 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use crate::graph::{ - AdjacencyList, - config::{self, MaxDegree}, -}; -use diskann_utils::views::MatrixView; - -use super::*; - -fn graph_config(degree: usize) -> Config { - config::Builder::new_with( - degree, - MaxDegree::same(), - degree, - Metric::L2.into(), - |builder| { - builder.alpha(1.2); - }, - ) - .build() - .unwrap() -} - -fn candidate_list(ids: impl IntoIterator) -> AdjacencyList { - AdjacencyList::from_iter_untrusted(ids) -} - -#[test] -fn preserves_lists_within_the_degree_bound() { - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let candidates = vec![ - candidate_list([3, 1]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - - let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - - assert_eq!(&*actual[0], &[1, 3]); -} - -#[test] -fn prunes_an_overfull_list_with_the_vamana_kernel() { - let data = [0.0_f32, 1.0, 2.0, -3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let candidates = vec![ - candidate_list([3, 2, 1]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - - let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - - assert_eq!(&*actual[0], &[1, 3]); -} - -#[test] -fn rejects_invalid_candidate_ids_without_panicking() { - let data = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![ - candidate_list([1, 3]), - candidate_list([]), - candidate_list([]), - ]; - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::InvalidCandidateId { - source_index: 0, - candidate: 3, - points: 3, - }) - )); -} - -#[test] -fn rejects_candidate_list_count_mismatch_without_panicking() { - let data = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![ - candidate_list([]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::CandidateListCountMismatch { - lists: 4, - points: 3 - }) - )); -} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index ec8763b190..1c60a40c4e 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -437,4 +437,415 @@ fn poisoned_candidate_list(point: u32) -> LeafBuildError { } #[cfg(test)] -mod tests; +mod tests { + use diskann_utils::views::MatrixView; + use diskann_vector::distance::Metric; + use half::f16; + use std::collections::BTreeSet; + + use super::{ + DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, + build_leaf_candidates, + }; + + fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + MatrixView::try_from(data, rows, columns).unwrap() + } + + fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() + } + + fn build( + data: MatrixView<'_, T>, + leaves: &[Vec], + k: usize, + metric: Metric, + ) -> Result>, LeafBuildError> + where + T: crate::utils::VectorRepr + 'static, + { + pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) + } + + fn adjacency_lists(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() + } + + fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { + let mut graph = vec![BTreeSet::new(); data.len()]; + for (source, left) in data.iter().enumerate() { + let mut nearest: Vec<_> = data + .iter() + .enumerate() + .filter(|(target, _)| *target != source) + .map(|(target, right)| { + let distance = left + .iter() + .zip(right) + .map(|(x, y)| (x - y) * (x - y)) + .sum::(); + (target, distance) + }) + .collect(); + nearest.sort_by(|left, right| { + left.1 + .total_cmp(&right.1) + .then_with(|| left.0.cmp(&right.0)) + }); + for &(target, _) in nearest.iter().take(k) { + graph[source].insert(target as u32); + graph[target].insert(source as u32); + } + } + graph + .into_iter() + .map(|neighbors| neighbors.into_iter().collect()) + .collect() + } + + #[test] + fn leaf_adjacency_matches_an_independent_all_pairs_reference() { + let points = [ + [0.0_f32, 0.0], + [1.0, 0.2], + [3.1, 0.5], + [7.8, 1.4], + [-2.3, 4.1], + [6.7, -3.2], + ]; + let flat: Vec<_> = points.into_iter().flatten().collect(); + + let actual = adjacency_lists( + build( + view(&flat, points.len(), 2), + &[(0..points.len() as u32).collect()], + 2, + Metric::L2, + ) + .unwrap(), + ); + + assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); + } + + #[test] + fn retains_and_deduplicates_candidates_from_overlapping_leaves() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; + + let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + adjacency_lists(graph), + [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] + ); + } + + #[test] + fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { + let dimensions = 9; + let mut data = vec![0.0_f32; 10 * dimensions]; + for source in 1..10 { + data[source * dimensions + source - 1] = 1.0; + } + + let graph = build( + view(&data, 10, dimensions), + &[(0..10).collect()], + 1, + Metric::L2, + ) + .unwrap(); + + assert_eq!(&*graph[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); + assert!(graph.iter().enumerate().all(|(source, neighbors)| { + neighbors.iter().all(|&target| target as usize != source) + && neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + })); + } + + #[test] + fn global_id_translation_is_independent_of_leaf_order() { + let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; + let leaves = vec![vec![4, 1, 3]]; + + let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); + + assert_eq!( + adjacency_lists(graph), + [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] + ); + } + + fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> + where + T: crate::utils::VectorRepr + 'static, + { + let leaves = vec![(0..points as u32).collect()]; + adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) + } + + fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) + where + T: crate::utils::VectorRepr + 'static, + { + let points = 8; + // Source dimension controls VectorRepr conversion chunking. Cover tails on + // both sides of 4-, 8-, and 16-element boundaries, then a second 16-lane + // chunk. Input integers remain exact in every tested representation. + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let source = index / dimensions; + let dimension = index % dimensions; + ((source * 7 + dimension * 3 + source * dimension) % 23) as u8 + }) + .collect(); + let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + assert_eq!( + source_graph(&converted, points, dimensions), + source_graph(&f32_data, points, dimensions), + "{label} dimensions={dimensions}" + ); + } + } + + #[test] + fn f16_conversion_matches_f32_across_dimension_boundaries() { + assert_source_conversion_matches_f32("f16", |value| f16::from_f32(value as f32)); + } + + #[test] + fn u8_conversion_matches_f32_across_dimension_boundaries() { + assert_source_conversion_matches_f32("u8", |value| value); + } + + #[test] + fn i8_conversion_matches_f32_across_dimension_boundaries() { + // Applying the same translation to every coordinate preserves L2 pair + // ordering while exercising signed conversion. + assert_source_conversion_matches_f32("i8", |value| value as i8 - 11); + } + + #[test] + fn all_metrics_produce_symmetric_unique_non_self_candidates() { + let data = [1.0_f32, 0.0, 0.8, 0.2, 0.0, 1.0, -1.0, 0.0]; + let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); + for (source, neighbors) in graph.iter().enumerate() { + assert!(neighbors.iter().all(|&target| target as usize != source)); + assert!( + neighbors + .iter() + .all(|&target| graph[target as usize].contains(source as u32)) + ); + assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); + } + } + } + + #[test] + fn parallel_leaf_schedule_does_not_change_candidate_order() { + let data: Vec = (0..64).map(|value| value as f32).collect(); + let leaves: Vec> = (0..32) + .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) + .collect(); + let pool = pool(); + pool.install(|| { + let expected = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); + for _ in 0..8 { + let actual = + build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2) + .unwrap(); + assert_eq!(actual, expected); + } + }); + } + + #[test] + fn rejects_invalid_dimensions_and_leaf_membership() { + let data = [0.0_f32, 1.0]; + let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); + assert!(matches!( + build(no_dimensions, &[], 1, Metric::L2), + Err(LeafBuildError::EmptyDimensions) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), + Err(LeafBuildError::EmptyLeaf { leaf: 0 }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { + leaf: 0, + point: 2, + points: 2 + }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), + Err(LeafBuildError::InvalidPointId { point: 2, .. }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) + )); + assert!(matches!( + build(view(&data, 2, 1), &[vec![1, 0, 1]], 1, Metric::L2), + Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 1 }) + )); + } + + #[test] + fn singleton_and_zero_k_leaves_add_no_candidates() { + let data = [0.0_f32, 1.0, 2.0]; + let singleton = build( + view(&data, 3, 1), + &[vec![0], vec![1], vec![2]], + 1, + Metric::L2, + ) + .unwrap(); + let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); + assert!( + singleton + .iter() + .chain(&zero_k) + .all(|candidates| candidates.is_empty()) + ); + } + + #[test] + fn reuses_worker_buffers_for_smaller_leaves() { + let mut buffers = LeafBuffers::default(); + buffers.prepare(0, 64, 128, 2).unwrap(); + let point_values = buffers.point_values.as_ptr(); + let dots = buffers.dots.as_ptr(); + let neighbors = buffers.neighbors.as_ptr(); + + buffers.prepare(1, 8, 128, 2).unwrap(); + + assert_eq!(buffers.point_values.as_ptr(), point_values); + assert_eq!(buffers.dots.as_ptr(), dots); + assert_eq!(buffers.neighbors.as_ptr(), neighbors); + assert_eq!(buffers.point_values.len(), 64 * 128); + assert_eq!(buffers.dots.len(), 64 * 64); + assert_eq!(buffers.neighbors.len(), 64 * 2); + } + + #[test] + fn reports_shape_overflow_before_allocating() { + let mut buffers = LeafBuffers::default(); + assert!(matches!( + buffers.prepare(7, usize::MAX, 2, 1), + Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) + )); + } + + #[test] + fn rejects_an_invalid_kernel_target() { + let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; + let error = add_symmetric_neighbors( + &[10, 20], + 1, + &[ + super::super::leaf_kernel::LeafNeighbor::new(9, 1.0), + super::super::leaf_kernel::LeafNeighbor::new(0, 1.0), + ], + &mut graph, + ) + .unwrap_err(); + assert!(matches!( + error, + LeafBuildError::InvalidLocalTarget { + target: 9, + points: 2 + } + )); + } + + #[test] + fn skips_duplicate_global_ids_without_self_edges() { + let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; + add_symmetric_neighbors( + &[7, 7], + 1, + &[ + super::super::leaf_kernel::LeafNeighbor::new(1, 0.0), + super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), + ], + &mut graph, + ) + .unwrap(); + assert!(graph.iter().all(|neighbors| neighbors.is_empty())); + } + + #[test] + fn poisoned_candidate_lists_return_errors() { + let candidates = DirectCandidates::new(1).unwrap(); + let _ = std::panic::catch_unwind(|| { + let _guard = candidates.lists[0].lock().unwrap(); + panic!("poison candidate list"); + }); + assert!(matches!( + candidates.add_leaf(&[0], &[crate::graph::AdjacencyList::new()]), + Err(LeafBuildError::PoisonedCandidateList { point: 0 }) + )); + assert!(matches!( + candidates.into_lists(), + Err(LeafBuildError::PoisonedCandidateList { point: 0 }) + )); + } + + #[test] + fn allocation_errors_preserve_buffer_context() { + let mut values = Vec::::new(); + let source = values.try_reserve(usize::MAX).unwrap_err(); + let error = allocation_error("test", 1, source); + assert!(matches!( + error, + LeafBuildError::Allocation { + buffer: "test", + additional: 1, + .. + } + )); + } + + #[test] + fn direct_candidate_accumulator_keeps_unique_sorted_lists() { + let candidates = DirectCandidates::new(2).unwrap(); + candidates + .add_leaf( + &[0, 1], + &[ + crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), + crate::graph::AdjacencyList::from_iter_untrusted([0]), + ], + ) + .unwrap(); + assert_eq!( + adjacency_lists(candidates.into_lists().unwrap()), + [vec![1], vec![0]] + ); + } +} diff --git a/diskann/src/graph/pipnn/leaf_build/tests.rs b/diskann/src/graph/pipnn/leaf_build/tests.rs deleted file mode 100644 index 6c55e6a577..0000000000 --- a/diskann/src/graph/pipnn/leaf_build/tests.rs +++ /dev/null @@ -1,414 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use diskann_utils::views::MatrixView; -use diskann_vector::distance::Metric; -use half::f16; -use std::collections::BTreeSet; - -use super::{ - DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, - build_leaf_candidates, -}; - -fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { - MatrixView::try_from(data, rows, columns).unwrap() -} - -fn pool() -> rayon::ThreadPool { - rayon::ThreadPoolBuilder::new() - .num_threads(4) - .build() - .unwrap() -} - -fn build( - data: MatrixView<'_, T>, - leaves: &[Vec], - k: usize, - metric: Metric, -) -> Result>, LeafBuildError> -where - T: crate::utils::VectorRepr + 'static, -{ - pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) -} - -fn adjacency_lists(graph: Vec>) -> Vec> { - graph.into_iter().map(Vec::from).collect() -} - -fn brute_force_symmetric_l2(data: &[[f32; 2]], k: usize) -> Vec> { - let mut graph = vec![BTreeSet::new(); data.len()]; - for (source, left) in data.iter().enumerate() { - let mut nearest: Vec<_> = data - .iter() - .enumerate() - .filter(|(target, _)| *target != source) - .map(|(target, right)| { - let distance = left - .iter() - .zip(right) - .map(|(x, y)| (x - y) * (x - y)) - .sum::(); - (target, distance) - }) - .collect(); - nearest.sort_by(|left, right| { - left.1 - .total_cmp(&right.1) - .then_with(|| left.0.cmp(&right.0)) - }); - for &(target, _) in nearest.iter().take(k) { - graph[source].insert(target as u32); - graph[target].insert(source as u32); - } - } - graph - .into_iter() - .map(|neighbors| neighbors.into_iter().collect()) - .collect() -} - -#[test] -fn leaf_adjacency_matches_an_independent_all_pairs_reference() { - let points = [ - [0.0_f32, 0.0], - [1.0, 0.2], - [3.1, 0.5], - [7.8, 1.4], - [-2.3, 4.1], - [6.7, -3.2], - ]; - let flat: Vec<_> = points.into_iter().flatten().collect(); - - let actual = adjacency_lists( - build( - view(&flat, points.len(), 2), - &[(0..points.len() as u32).collect()], - 2, - Metric::L2, - ) - .unwrap(), - ); - - assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); -} - -#[test] -fn retains_and_deduplicates_candidates_from_overlapping_leaves() { - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; - - let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); - - assert_eq!( - adjacency_lists(graph), - [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] - ); -} - -#[test] -fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { - let dimensions = 9; - let mut data = vec![0.0_f32; 10 * dimensions]; - for source in 1..10 { - data[source * dimensions + source - 1] = 1.0; - } - - let graph = build( - view(&data, 10, dimensions), - &[(0..10).collect()], - 1, - Metric::L2, - ) - .unwrap(); - - assert_eq!(&*graph[0], &[1, 2, 3, 4, 5, 6, 7, 8, 9]); - assert!(graph.iter().enumerate().all(|(source, neighbors)| { - neighbors.iter().all(|&target| target as usize != source) - && neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32)) - })); -} - -#[test] -fn global_id_translation_is_independent_of_leaf_order() { - let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; - let leaves = vec![vec![4, 1, 3]]; - - let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); - - assert_eq!( - adjacency_lists(graph), - [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] - ); -} - -fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> -where - T: crate::utils::VectorRepr + 'static, -{ - let leaves = vec![(0..points as u32).collect()]; - adjacency_lists(build(view(data, points, dimensions), &leaves, 2, Metric::L2).unwrap()) -} - -fn assert_source_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) -where - T: crate::utils::VectorRepr + 'static, -{ - let points = 8; - // Source dimension controls VectorRepr conversion chunking. Cover tails on - // both sides of 4-, 8-, and 16-element boundaries, then a second 16-lane - // chunk. Input integers remain exact in every tested representation. - for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let raw: Vec = (0..points * dimensions) - .map(|index| { - let source = index / dimensions; - let dimension = index % dimensions; - ((source * 7 + dimension * 3 + source * dimension) % 23) as u8 - }) - .collect(); - let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); - let converted: Vec = raw.iter().copied().map(&convert).collect(); - assert_eq!( - source_graph(&converted, points, dimensions), - source_graph(&f32_data, points, dimensions), - "{label} dimensions={dimensions}" - ); - } -} - -#[test] -fn f16_conversion_matches_f32_across_dimension_boundaries() { - assert_source_conversion_matches_f32("f16", |value| f16::from_f32(value as f32)); -} - -#[test] -fn u8_conversion_matches_f32_across_dimension_boundaries() { - assert_source_conversion_matches_f32("u8", |value| value); -} - -#[test] -fn i8_conversion_matches_f32_across_dimension_boundaries() { - // Applying the same translation to every coordinate preserves L2 pair - // ordering while exercising signed conversion. - assert_source_conversion_matches_f32("i8", |value| value as i8 - 11); -} - -#[test] -fn all_metrics_produce_symmetric_unique_non_self_candidates() { - let data = [1.0_f32, 0.0, 0.8, 0.2, 0.0, 1.0, -1.0, 0.0]; - let leaves = vec![vec![0, 1, 2, 3], vec![0, 1, 2, 3]]; - - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let graph = build(view(&data, 4, 2), &leaves, 2, metric).unwrap(); - for (source, neighbors) in graph.iter().enumerate() { - assert!(neighbors.iter().all(|&target| target as usize != source)); - assert!( - neighbors - .iter() - .all(|&target| graph[target as usize].contains(source as u32)) - ); - assert!(neighbors.windows(2).all(|pair| pair[0] < pair[1])); - } - } -} - -#[test] -fn parallel_leaf_schedule_does_not_change_candidate_order() { - let data: Vec = (0..64).map(|value| value as f32).collect(); - let leaves: Vec> = (0..32) - .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) - .collect(); - let pool = pool(); - pool.install(|| { - let expected = - build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); - for _ in 0..8 { - let actual = - build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); - assert_eq!(actual, expected); - } - }); -} - -#[test] -fn rejects_invalid_dimensions_and_leaf_membership() { - let data = [0.0_f32, 1.0]; - let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); - assert!(matches!( - build(no_dimensions, &[], 1, Metric::L2), - Err(LeafBuildError::EmptyDimensions) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), - Err(LeafBuildError::EmptyLeaf { leaf: 0 }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), - Err(LeafBuildError::InvalidPointId { - leaf: 0, - point: 2, - points: 2 - }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), - Err(LeafBuildError::InvalidPointId { point: 2, .. }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), - Err(LeafBuildError::InvalidPointId { point: 2, .. }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), - Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![1, 0, 1]], 1, Metric::L2), - Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 1 }) - )); -} - -#[test] -fn singleton_and_zero_k_leaves_add_no_candidates() { - let data = [0.0_f32, 1.0, 2.0]; - let singleton = build( - view(&data, 3, 1), - &[vec![0], vec![1], vec![2]], - 1, - Metric::L2, - ) - .unwrap(); - let zero_k = build(view(&data, 3, 1), &[vec![0, 1, 2]], 0, Metric::L2).unwrap(); - assert!( - singleton - .iter() - .chain(&zero_k) - .all(|candidates| candidates.is_empty()) - ); -} - -#[test] -fn reuses_worker_buffers_for_smaller_leaves() { - let mut buffers = LeafBuffers::default(); - buffers.prepare(0, 64, 128, 2).unwrap(); - let point_values = buffers.point_values.as_ptr(); - let dots = buffers.dots.as_ptr(); - let neighbors = buffers.neighbors.as_ptr(); - - buffers.prepare(1, 8, 128, 2).unwrap(); - - assert_eq!(buffers.point_values.as_ptr(), point_values); - assert_eq!(buffers.dots.as_ptr(), dots); - assert_eq!(buffers.neighbors.as_ptr(), neighbors); - assert_eq!(buffers.point_values.len(), 64 * 128); - assert_eq!(buffers.dots.len(), 64 * 64); - assert_eq!(buffers.neighbors.len(), 64 * 2); -} - -#[test] -fn reports_shape_overflow_before_allocating() { - let mut buffers = LeafBuffers::default(); - assert!(matches!( - buffers.prepare(7, usize::MAX, 2, 1), - Err(LeafBuildError::ShapeOverflow { leaf: 7, .. }) - )); -} - -#[test] -fn rejects_an_invalid_kernel_target() { - let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; - let error = add_symmetric_neighbors( - &[10, 20], - 1, - &[ - super::super::leaf_kernel::LeafNeighbor::new(9, 1.0), - super::super::leaf_kernel::LeafNeighbor::new(0, 1.0), - ], - &mut graph, - ) - .unwrap_err(); - assert!(matches!( - error, - LeafBuildError::InvalidLocalTarget { - target: 9, - points: 2 - } - )); -} - -#[test] -fn skips_duplicate_global_ids_without_self_edges() { - let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; - add_symmetric_neighbors( - &[7, 7], - 1, - &[ - super::super::leaf_kernel::LeafNeighbor::new(1, 0.0), - super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), - ], - &mut graph, - ) - .unwrap(); - assert!(graph.iter().all(|neighbors| neighbors.is_empty())); -} - -#[test] -fn poisoned_candidate_lists_return_errors() { - let candidates = DirectCandidates::new(1).unwrap(); - let _ = std::panic::catch_unwind(|| { - let _guard = candidates.lists[0].lock().unwrap(); - panic!("poison candidate list"); - }); - assert!(matches!( - candidates.add_leaf(&[0], &[crate::graph::AdjacencyList::new()]), - Err(LeafBuildError::PoisonedCandidateList { point: 0 }) - )); - assert!(matches!( - candidates.into_lists(), - Err(LeafBuildError::PoisonedCandidateList { point: 0 }) - )); -} - -#[test] -fn allocation_errors_preserve_buffer_context() { - let mut values = Vec::::new(); - let source = values.try_reserve(usize::MAX).unwrap_err(); - let error = allocation_error("test", 1, source); - assert!(matches!( - error, - LeafBuildError::Allocation { - buffer: "test", - additional: 1, - .. - } - )); -} - -#[test] -fn direct_candidate_accumulator_keeps_unique_sorted_lists() { - let candidates = DirectCandidates::new(2).unwrap(); - candidates - .add_leaf( - &[0, 1], - &[ - crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), - crate::graph::AdjacencyList::from_iter_untrusted([0]), - ], - ) - .unwrap(); - assert_eq!( - adjacency_lists(candidates.into_lists().unwrap()), - [vec![1], vec![0]] - ); -} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 7bda42c664..e715498690 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -275,13 +275,11 @@ where T: VectorRepr + Send + Sync + 'static, { if data.nrows() == 0 { - return Err(ANNError::log_dimension_mismatch_error( - "PiPNN requires at least one data point".into(), - )); + return Err(ANNError::message("PiPNN requires at least one data point")); } if data.ncols() == 0 { - return Err(ANNError::log_dimension_mismatch_error( - "PiPNN requires at least one data dimension".into(), + return Err(ANNError::message( + "PiPNN requires at least one data dimension", )); } if data.nrows() > u32::MAX as usize { @@ -291,7 +289,7 @@ where ))); } data.nrows().checked_mul(data.ncols()).ok_or_else(|| { - ANNError::log_dimension_mismatch_error(format!( + ANNError::message(format!( "PiPNN dataset shape {} x {} overflows usize", data.nrows(), data.ncols() @@ -307,7 +305,7 @@ where // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) - .map_err(ANNError::opaque) + .map_err(ANNError::new) })?; // Finalization consumes candidate lists and reuses their allocations for the // resulting adjacency where possible. @@ -329,8 +327,382 @@ fn effective_metric(metric: Metric) -> Metric { #[track_caller] fn config_error(message: impl std::fmt::Display) -> ANNError { - ANNError::log_index_config_error("PiPNN".into(), message.to_string()) + ANNError::message(format!("PiPNN configuration: {message}")) +} + +#[cfg(test)] +mod tests { + use super::*; + use half::f16; + + #[test] + fn integer_normalized_cosine_uses_unnormalized_cosine() { + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let expected = if metric == Metric::CosineNormalized { + Metric::Cosine + } else { + metric + }; + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), expected); + assert_eq!(effective_metric::(metric), metric); + assert_eq!(effective_metric::(metric), metric); + } + } } +#[cfg(test)] +#[allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] +mod build_graph_tests { + use super::{PiPNNBuildContext, PiPNNConfig, build_graph}; + use crate::graph::config::{self, MaxDegree}; + use diskann_utils::views::MatrixView; + use diskann_vector::distance::Metric; + use half::f16; + use rand::{Rng, SeedableRng, rngs::StdRng}; + + fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 4, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + } + } + + fn graph_config(metric: Metric, degree: usize) -> crate::graph::Config { + config::Builder::new_with(degree, MaxDegree::same(), 8, metric.into(), |builder| { + builder.alpha(1.2); + }) + .build() + .unwrap() + } + + fn pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + } + + fn rows(graph: Vec>) -> Vec> { + graph.into_iter().map(Vec::from).collect() + } + + fn assert_graph_invariants( + graph: &[crate::graph::AdjacencyList], + points: usize, + degree: usize, + ) { + assert_eq!(graph.len(), points); + for (source, row) in graph.iter().enumerate() { + assert!(row.len() <= degree); + let mut sorted = row.to_vec(); + sorted.sort_unstable(); + sorted.dedup(); + assert_eq!(sorted.len(), row.len()); + assert!( + row.iter() + .all(|&id| (id as usize) < points && id as usize != source) + ); + } + } + + #[test] + fn builds_a_single_leaf_graph_for_real_dataset_ids() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + let graph = graph_config(Metric::L2, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); + + let graph = graph_config(Metric::L2, 1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let pruned = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&pruned, 4, 1); + for (source, neighbors) in pruned.iter().enumerate() { + assert_eq!(source.abs_diff(neighbors[0] as usize), 1); + } + } + + #[test] + fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { + let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; + let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); + let graph = graph_config(Metric::L2, 1); + let pool = pool(2); + let config = PiPNNConfig { + c_max: 5, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 4, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&actual, 5, 1); + assert!(actual.iter().all(|row| row.len() == 1)); + } + + #[test] + fn rejects_empty_dataset_dimensions_at_the_public_boundary() { + let graph = graph_config(Metric::L2, 2); + let pool = pool(1); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); + let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); + + assert!(build_graph(no_rows, &context).is_err()); + assert!(build_graph(no_columns, &context).is_err()); + } + #[test] + fn supports_every_source_type_and_metric() { + fn build( + values: &[T], + metric: Metric, + ) { + let data = MatrixView::try_from(values, 6, 2).unwrap(); + let graph = graph_config(metric, 2); + let pool = pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); + let actual = build_graph(data, &context).unwrap(); + assert_graph_invariants(&actual, 6, 2); + } + + let values = [ + 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, + ]; + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + build(&values, metric); + } + build(&values.map(f16::from_f32), Metric::L2); + build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); + build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); + } + + #[test] + fn integer_normalized_cosine_matches_cosine() { + fn assert_match(values: &[T]) { + let data = MatrixView::try_from(values, 8, 2).unwrap(); + let pool = pool(2); + let build = |metric| { + let graph = graph_config(metric, 2); + let config = PiPNNConfig { + c_max: 8, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: 1, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + rows(build_graph(data, &context).unwrap()) + }; + assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); + } + + assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); + assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); + } + + #[test] + fn is_deterministic_for_a_fixed_pool_size() { + let data: Vec = (0..96 * 4) + .map(|value| ((value * 17 + 3) % 101) as f32) + .collect(); + let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); + let graph = graph_config(Metric::L2, 8); + let pool = pool(4); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.25, + fanout: vec![3, 2], + k: 3, + replicas: 2, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let first = build_graph(data, &context).unwrap(); + let second = build_graph(data, &context).unwrap(); + + assert_eq!(first, second); + assert_graph_invariants(&first, 96, 8); + } + + #[test] + fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { + let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); + for case in 0..24 { + let points = rng.random_range(4..=32); + let dimensions = rng.random_range(1..=8); + let c_max = rng.random_range(4..=points.min(12)); + let c_min = rng.random_range(1..=c_max); + let degree = rng.random_range(1..=points.min(8)); + let values: Vec = (0..points * dimensions) + .map(|_| rng.random_range(-10.0..10.0)) + .collect(); + let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); + let graph = graph_config(Metric::L2, degree); + let pool = pool(2); + let config = PiPNNConfig { + c_max, + c_min, + p_samp: 0.5, + fanout: vec![2], + k: rng.random_range(1..=3), + replicas: rng.random_range(1..=2), + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); + + let actual = build_graph(data, &context) + .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); + assert_graph_invariants(&actual, points, degree); + } + } +} #[cfg(test)] -mod tests; +#[allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] +mod config_tests { + use super::{PiPNNBuildContext, PiPNNConfig}; + use crate::graph::config::{self, MaxDegree}; + use diskann_vector::distance::Metric; + + fn pipnn_config() -> PiPNNConfig { + PiPNNConfig { + c_max: 512, + c_min: 64, + p_samp: 0.01, + fanout: vec![10, 3], + k: 2, + replicas: 1, + } + } + + fn graph_config(metric: Metric, alpha: f32) -> crate::graph::Config { + config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { + builder.alpha(alpha); + }) + .build() + .unwrap() + } + + fn pool() -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap() + } + + #[test] + fn rejects_each_invalid_algorithm_parameter() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + let mut cases = [ + PiPNNConfig { + c_max: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 0, + ..pipnn_config() + }, + PiPNNConfig { + c_min: 513, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 0.0, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: -0.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: 1.01, + ..pipnn_config() + }, + PiPNNConfig { + p_samp: f64::NAN, + ..pipnn_config() + }, + PiPNNConfig { + fanout: Vec::new(), + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![1, 0], + ..pipnn_config() + }, + PiPNNConfig { + fanout: vec![17], + ..pipnn_config() + }, + PiPNNConfig { + k: 0, + ..pipnn_config() + }, + PiPNNConfig { + replicas: 0, + ..pipnn_config() + }, + ]; + + for config in &mut cases { + let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .expect_err("invalid PiPNN config must be rejected"); + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + } + } + + #[test] + fn rejects_graph_policy_for_a_different_metric() { + let graph = graph_config(Metric::InnerProduct, 1.2); + let pool = pool(); + + let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); + assert!(error.to_string().contains("prune kind")); + } + + #[test] + fn does_not_add_alpha_validation_beyond_graph_config() { + let pool = pool(); + for alpha in [0.9, f32::NAN, f32::INFINITY] { + let graph = graph_config(Metric::L2, alpha); + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + } + } +} diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 96ed91878b..5de9a68cf5 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -149,13 +149,13 @@ where { let points = data.nrows(); if points == 0 { - return Err(ANNError::opaque(PartitionError::EmptyDataset)); + return Err(ANNError::new(PartitionError::EmptyDataset)); } if data.ncols() == 0 { - return Err(ANNError::opaque(PartitionError::EmptyDimensions)); + return Err(ANNError::new(PartitionError::EmptyDimensions)); } if points > u32::MAX as usize { - return Err(ANNError::opaque(PartitionError::TooManyPoints(points))); + return Err(ANNError::new(PartitionError::TooManyPoints(points))); } let mut leaves = Vec::new(); @@ -170,7 +170,7 @@ where partition_replica(data, &config, metric, &kernel, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; leaves.append(&mut replica_leaves); } validate_leaves(&leaves, config.c_max)?; @@ -191,14 +191,14 @@ where let initial_indices = point_ids(data.nrows())?; if data.nrows() <= config.c_max { let mut leaves = Vec::new(); - leaves.try_reserve_exact(1).map_err(ANNError::opaque)?; + leaves.try_reserve_exact(1).map_err(ANNError::new)?; leaves.push(initial_indices); return Ok(leaves); } let mut leaves = Vec::new(); let mut work = Vec::new(); - work.try_reserve_exact(1).map_err(ANNError::opaque)?; + work.try_reserve_exact(1).map_err(ANNError::new)?; work.push(WorkItem { indices: initial_indices, level: 0, @@ -213,7 +213,7 @@ where let mut results = Vec::new(); results .try_reserve_exact(work.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; results.resize_with(work.len(), || None); // build_graph installs this complete private call tree into the // caller-owned pool; the indexed fill cannot escape that pool. @@ -235,13 +235,11 @@ where let mut next_work = Vec::new(); for result in results { let (mut pending, mut finished) = - result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??; + result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))??; next_work .try_reserve(pending.len()) - .map_err(ANNError::opaque)?; - leaves - .try_reserve(finished.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; + leaves.try_reserve(finished.len()).map_err(ANNError::new)?; next_work.append(&mut pending); leaves.append(&mut finished); } @@ -254,7 +252,7 @@ where let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { return global_merge_small(leaves, config.c_min, config.c_max); }; - Err(ANNError::opaque(PartitionError::IterationLimit { + Err(ANNError::new(PartitionError::IterationLimit { size: largest.indices.len(), level: largest.level, limit: MAX_PARTITION_ITERATIONS, @@ -291,12 +289,10 @@ where let mut pending = Vec::new(); let mut finished = Vec::new(); - pending - .try_reserve(clusters.len()) - .map_err(ANNError::opaque)?; + pending.try_reserve(clusters.len()).map_err(ANNError::new)?; finished .try_reserve(clusters.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; let child_seed = mix_seed(item.seed, points as u64); for cluster in clusters { if cluster.is_empty() { @@ -319,7 +315,7 @@ fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResul let count = sample_num_leaders(points.len(), sampling_fraction); let mut rng = rand::rngs::StdRng::seed_from_u64(seed); let mut leaders = Vec::new(); - leaders.try_reserve_exact(count).map_err(ANNError::opaque)?; + leaders.try_reserve_exact(count).map_err(ANNError::new)?; leaders.extend(points.choose_multiple(&mut rng, count).copied()); Ok(leaders) } @@ -478,7 +474,7 @@ where None, dots, ) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; let point_scales = if metric == Metric::Cosine { grow_fallible(point_scale_buffer, point_count, 0.0)?; @@ -504,14 +500,14 @@ where Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, }; let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { - ANNError::opaque(PartitionError::InvalidBufferLength { + ANNError::new(PartitionError::InvalidBufferLength { buffer: "dot-product stripe", expected: dots_len, actual: dots.len(), }) })?; let output = MutMatrixView::try_from(assignments, point_count, fanout).map_err(|error| { - ANNError::opaque(PartitionError::InvalidBufferLength { + ANNError::new(PartitionError::InvalidBufferLength { buffer: "partition assignments", expected: output_len, actual: error.into_inner().len(), @@ -519,7 +515,7 @@ where })?; kernel .nearest_leaders(PartitionInput { dots, scales }, output) - .map_err(ANNError::opaque) + .map_err(ANNError::new) } fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> @@ -528,7 +524,7 @@ where { let expected = checked_area("gather output", indices.len(), data.ncols())?; if output.len() != expected { - return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + return Err(ANNError::new(PartitionError::InvalidBufferLength { buffer: "gather output", expected, actual: output.len(), @@ -560,9 +556,7 @@ fn scatter_assignments( let stripe_assignment_count = checked_area("scatter assignment stripe", stripe_points, fanout)?; let stripes = points.len().div_ceil(stripe_points); let mut partials = Vec::new(); - partials - .try_reserve_exact(stripes) - .map_err(ANNError::opaque)?; + partials.try_reserve_exact(stripes).map_err(ANNError::new)?; partials.resize_with(stripes, || None); // See the pool invariant at the other partition terminal operations. #[allow(clippy::disallowed_methods)] @@ -578,18 +572,16 @@ fn scatter_assignments( }); let mut locals = Vec::new(); - locals - .try_reserve_exact(stripes) - .map_err(ANNError::opaque)?; + locals.try_reserve_exact(stripes).map_err(ANNError::new)?; for result in partials { - locals.push(result.ok_or_else(|| ANNError::opaque(PartitionError::MissingWorkerResult))??); + locals.push(result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))??); } let mut sizes = filled_vec(leaders, 0usize)?; for local in &locals { for (size, cluster) in sizes.iter_mut().zip(local) { *size = size.checked_add(cluster.len()).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { + ANNError::new(PartitionError::ShapeOverflow { buffer: "cluster size", rows: *size, cols: cluster.len(), @@ -605,7 +597,7 @@ fn scatter_assignments( .enumerate() .map(|(leader, size)| { let mut cluster = Vec::new(); - cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + cluster.try_reserve_exact(size).map_err(ANNError::new)?; for local in &locals { cluster.extend_from_slice(&local[leader]); } @@ -623,14 +615,14 @@ fn scatter_serial( let mut sizes = filled_vec(leaders, 0usize)?; for &leader in assignments { let Some(size) = sizes.get_mut(leader as usize) else { - return Err(ANNError::opaque(PartitionError::InvalidBufferLength { + return Err(ANNError::new(PartitionError::InvalidBufferLength { buffer: "leader assignment", expected: leaders, actual: leader as usize + 1, })); }; *size = size.checked_add(1).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { + ANNError::new(PartitionError::ShapeOverflow { buffer: "cluster size", rows: *size, cols: 1, @@ -650,10 +642,10 @@ fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { let mut clusters = Vec::new(); clusters .try_reserve_exact(sizes.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; for &size in sizes { let mut cluster = Vec::new(); - cluster.try_reserve_exact(size).map_err(ANNError::opaque)?; + cluster.try_reserve_exact(size).map_err(ANNError::new)?; clusters.push(cluster); } Ok(clusters) @@ -666,10 +658,10 @@ fn global_merge_small( ) -> ANNResult>> { let mut merged = Vec::new(); let mut small_leaves = Vec::new(); - merged.try_reserve(leaves.len()).map_err(ANNError::opaque)?; + merged.try_reserve(leaves.len()).map_err(ANNError::new)?; small_leaves .try_reserve(leaves.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; for leaf in leaves { if leaf.len() >= c_min { merged.push(leaf); @@ -682,11 +674,11 @@ fn global_merge_small( } let mut small = HashSet::new(); - small.try_reserve(c_max).map_err(ANNError::opaque)?; + small.try_reserve(c_max).map_err(ANNError::new)?; for leaf in small_leaves { let combined = small.len().checked_add(leaf.len()).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { + ANNError::new(PartitionError::ShapeOverflow { buffer: "small-leaf merge", rows: small.len(), cols: leaf.len(), @@ -695,7 +687,7 @@ fn global_merge_small( if combined > c_max { merged.push(drain_sorted(&mut small)?); } - small.try_reserve(leaf.len()).map_err(ANNError::opaque)?; + small.try_reserve(leaf.len()).map_err(ANNError::new)?; small.extend(leaf); if small.len() >= c_min { merged.push(drain_sorted(&mut small)?); @@ -709,15 +701,14 @@ fn global_merge_small( { remainder.retain(|id| !last.contains(id)); let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { - ANNError::opaque(PartitionError::ShapeOverflow { + ANNError::new(PartitionError::ShapeOverflow { buffer: "small-leaf tail merge", rows: last.len(), cols: remainder.len(), }) })?; if combined <= c_max { - last.try_reserve(remainder.len()) - .map_err(ANNError::opaque)?; + last.try_reserve(remainder.len()).map_err(ANNError::new)?; last.append(&mut remainder); last.sort_unstable(); } @@ -733,9 +724,7 @@ fn global_merge_small( fn drain_sorted(set: &mut HashSet) -> ANNResult> { let mut values = Vec::new(); - values - .try_reserve_exact(set.len()) - .map_err(ANNError::opaque)?; + values.try_reserve_exact(set.len()).map_err(ANNError::new)?; values.extend(set.drain()); values.sort_unstable(); Ok(values) @@ -746,7 +735,7 @@ fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { .iter() .find(|leaf| leaf.is_empty() || leaf.len() > c_max) { - return Err(ANNError::opaque(PartitionError::InvalidLeaf { + return Err(ANNError::new(PartitionError::InvalidLeaf { size: leaf.len(), limit: c_max, })); @@ -756,14 +745,14 @@ fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { fn point_ids(points: usize) -> ANNResult> { let mut ids = Vec::new(); - ids.try_reserve_exact(points).map_err(ANNError::opaque)?; + ids.try_reserve_exact(points).map_err(ANNError::new)?; ids.extend(0..points as u32); Ok(ids) } fn filled_vec(len: usize, value: T) -> ANNResult> { let mut values = Vec::new(); - values.try_reserve_exact(len).map_err(ANNError::opaque)?; + values.try_reserve_exact(len).map_err(ANNError::new)?; values.resize(len, value); Ok(values) } @@ -779,14 +768,14 @@ fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResu } values .try_reserve(len - values.len()) - .map_err(ANNError::opaque)?; + .map_err(ANNError::new)?; values.resize(len, value); Ok(()) } fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { rows.checked_mul(cols) - .ok_or_else(|| ANNError::opaque(PartitionError::ShapeOverflow { buffer, rows, cols })) + .ok_or_else(|| ANNError::new(PartitionError::ShapeOverflow { buffer, rows, cols })) } fn assignment_stripe_point_count(leader_count: usize) -> usize { @@ -800,4 +789,423 @@ fn assignment_stripe_point_count(leader_count: usize) -> usize { } #[cfg(test)] -mod tests; +mod tests { + use diskann_utils::views::{Matrix, MatrixView}; + use diskann_vector::{Half, distance::Metric}; + + use super::*; + + fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { + PiPNNConfig { + c_max, + c_min, + p_samp: 0.25, + fanout, + k: 1, + replicas, + } + } + + fn clustered_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let point = position / dimensions; + let dimension = position % dimensions; + position += 1; + (point / 8) as f32 * 10.0 + dimension as f32 * 0.01 + point as f32 * 0.001 + } + }), + points, + dimensions, + ) + } + + fn directional_data(points: usize, dimensions: usize) -> Matrix { + Matrix::new( + diskann_utils::views::Init({ + let mut position = 0usize; + move || { + let point = position / dimensions; + let dimension = position % dimensions; + position += 1; + let angle = std::f32::consts::TAU * point as f32 / points as f32; + match dimension { + 0 => angle.cos(), + 1 => angle.sin(), + _ => 0.0, + } + } + }), + points, + dimensions, + ) + } + + fn sorted_memberships(leaves: &[Vec]) -> Vec> { + let mut memberships: Vec> = leaves + .iter() + .map(|leaf| { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids + }) + .collect(); + memberships.sort(); + memberships + } + + fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { + assert!( + leaves + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) + ); + let mut counts = vec![0usize; points]; + for leaf in leaves { + let mut ids = leaf.clone(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids.len(), leaf.len(), "duplicate ID inside a leaf"); + for &id in leaf { + assert!((id as usize) < points); + counts[id as usize] += 1; + } + } + assert!(counts.iter().all(|&count| count >= replicas)); + } + + #[test] + fn returns_one_leaf_at_and_below_c_max() { + for points in [7, 8] { + let data = clustered_data(points, 3); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + assert_eq!(leaves, vec![(0..points as u32).collect::>()]); + } + } + + #[test] + fn partition_is_fixed_seed_deterministic_and_bounded() { + let data = clustered_data(96, 8); + let config = config(4, 16, vec![3, 2], 2); + + let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); + let second = partition(data.as_view(), config, Metric::L2).unwrap(); + + assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); + assert_valid_partition(&first, 96, 16, 2); + assert!(first.iter().map(Vec::len).sum::() > 96 * 2); + } + + #[test] + fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { + let data = clustered_data(80, 4); + let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + + assert_valid_partition(&leaves, 80, 8, 1); + } + + #[test] + fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { + let data = Matrix::new(1.0f32, 24, 4); + let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = error.downcast::().unwrap(); + + assert!(matches!( + error, + PartitionError::IterationLimit { + size: 24, + limit: MAX_PARTITION_ITERATIONS, + .. + } + )); + } + + #[test] + fn global_merge_canonicalizes_small_leaf_membership() { + let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; + + let merged = global_merge_small(leaves, 4, 8).unwrap(); + + assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); + } + + #[test] + fn global_merge_never_overfills_before_reaching_c_min() { + let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; + + let merged = global_merge_small(leaves, 11, 11).unwrap(); + + assert_eq!( + merged, + vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] + ); + } + + #[test] + fn global_merge_fills_exact_capacity_before_flushing() { + let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + + assert_eq!(merged, vec![vec![0, 1, 2, 3]]); + } + + #[test] + fn replicas_cover_every_point_once_or_more_per_replica() { + let data = directional_data(72, 5); + let leaves = partition( + data.as_view(), + config(3, 12, vec![3, 2], 3), + Metric::CosineNormalized, + ) + .unwrap(); + + assert_valid_partition(&leaves, 72, 12, 3); + } + + fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) + where + T: crate::utils::VectorRepr + Send + Sync, + { + let points = 64; + // Partition gathering converts source vectors before GEMM. Exercise conversion + // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let point = index / dimensions; + let dimension = index % dimensions; + ((point * 5 + dimension * 7 + point * dimension) % 23) as u8 + }) + .collect(); + let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + let config = config(2, 16, vec![2, 1], 1); + let expected = partition( + MatrixView::try_from(&f32_data, points, dimensions).unwrap(), + config.clone(), + Metric::L2, + ) + .unwrap(); + let actual = partition( + MatrixView::try_from(&converted, points, dimensions).unwrap(), + config, + Metric::L2, + ) + .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); + + assert_valid_partition(&actual, points, 16, 1); + assert_eq!( + sorted_memberships(&actual), + sorted_memberships(&expected), + "{label} dimensions={dimensions}" + ); + } + } + + #[test] + fn f16_partition_matches_f32_across_dimension_boundaries() { + assert_partition_conversion_matches_f32("f16", |value| Half::from_f32(value as f32)); + } + + #[test] + fn u8_partition_matches_f32_across_dimension_boundaries() { + assert_partition_conversion_matches_f32("u8", |value| value); + } + + #[test] + fn i8_partition_matches_f32_across_dimension_boundaries() { + // The same translation in every coordinate preserves L2 ordering. + assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); + } + + #[test] + fn l2_leader_norms_preserve_scalar_reduction_order() { + fn next(state: &mut u64) -> f32 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 + } + + // This fixed case sits on opposite sides of the top-1 boundary depending + // on whether leader norms use the original scalar reduction or a SIMD + // reassociation. Point/leader dot products still go through the production + // GEMM; only the setup norm calculation is under test. + let dimensions = 129; + let mut state = 0x3a85_f952_c718_6e49; + let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let leader_zero: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let leader_one: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); + let data: Vec = leader_zero + .into_iter() + .chain(leader_one) + .chain(point) + .collect(); + let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); + + let clusters = assign_to_leaders( + data, + &[2], + &[0, 1], + 1, + Metric::L2, + &PartitionKernel::new(Metric::L2), + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); + + assert_eq!(clusters, [vec![], vec![2]]); + } + + #[test] + fn all_metrics_produce_valid_partitions() { + let data = directional_data(64, 8); + let config = config(2, 20, vec![2], 1); + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); + assert_valid_partition(&leaves, 64, 20, 1); + } + } + + #[test] + fn leader_count_is_bounded() { + assert_eq!(sample_num_leaders(1, 1.0), 1); + assert_eq!(sample_num_leaders(10, 0.01), 2); + assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); + } + + #[test] + fn replica_seed_derivation_is_stable_and_distinct() { + assert_eq!(replica_seed(0), 1_000); + assert_eq!(replica_seed(1), 8_919); + } + + #[test] + fn assignment_stripes_use_power_of_two_point_counts() { + assert_eq!(assignment_stripe_point_count(1_000), 128); + assert_eq!(assignment_stripe_point_count(256), 512); + assert_eq!( + assignment_stripe_point_count(1), + MAX_ASSIGNMENT_STRIPE_POINTS + ); + } + + #[test] + fn stripe_buffer_pool_reuses_returned_capacity() { + let pool = StripeBufferPool::new((), 0, None); + let points = { + let mut buffers = pool.get_ref(()); + buffers.points.resize(16, 0.0); + buffers.points.as_ptr() + }; + + let buffers = pool.get_ref(()); + assert_eq!(buffers.points.as_ptr(), points); + assert_eq!(buffers.points.len(), 16); + } + + #[test] + fn leader_assignment_handles_multiple_stripes() { + let points = 2_048; + let data: Vec = (0..points).map(|point| point as f32).collect(); + let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); + let point_ids: Vec = (0..points as u32).collect(); + + let clusters = assign_to_leaders( + data, + &point_ids, + &[0, 2_047], + 1, + Metric::L2, + &PartitionKernel::new(Metric::L2), + &StripeBufferPool::new((), 0, None), + ) + .unwrap(); + + assert_eq!(clusters[0], (0..1_024).collect::>()); + assert_eq!(clusters[1], (1_024..2_048).collect::>()); + } + + #[test] + fn parallel_scatter_matches_serial_order() { + let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); + let assignments: Vec = points + .iter() + .flat_map(|point| [point % 7, (point + 3) % 7]) + .collect(); + + let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); + let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); + + assert_eq!(actual, expected); + } + + #[test] + fn rejects_empty_dataset() { + let data = Matrix::::new(0.0, 0, 4); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDataset + ); + } + + #[test] + fn rejects_zero_dimensions() { + let data = Matrix::::new(0.0, 4, 0); + let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::EmptyDimensions + ); + } + + #[test] + fn rejects_invalid_gather_output_length() { + let data = Matrix::::new(0.0, 2, 2); + let error = gather_vectors(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "gather output", + expected: 4, + actual: 3, + } + ); + } + + #[test] + fn rejects_assignment_to_an_unknown_leader() { + let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "leader assignment", + expected: 2, + actual: 3, + } + ); + } + + #[test] + fn rejects_empty_and_oversized_leaves() { + for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { + let error = validate_leaves(&leaves, 2).unwrap_err(); + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidLeaf { size, limit: 2 } + ); + } + } +} diff --git a/diskann/src/graph/pipnn/partitioning/tests.rs b/diskann/src/graph/pipnn/partitioning/tests.rs deleted file mode 100644 index 8780078ca5..0000000000 --- a/diskann/src/graph/pipnn/partitioning/tests.rs +++ /dev/null @@ -1,423 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use diskann_utils::views::{Matrix, MatrixView}; -use diskann_vector::{Half, distance::Metric}; - -use super::*; - -fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { - PiPNNConfig { - c_max, - c_min, - p_samp: 0.25, - fanout, - k: 1, - replicas, - } -} - -fn clustered_data(points: usize, dimensions: usize) -> Matrix { - Matrix::new( - diskann_utils::views::Init({ - let mut position = 0usize; - move || { - let point = position / dimensions; - let dimension = position % dimensions; - position += 1; - (point / 8) as f32 * 10.0 + dimension as f32 * 0.01 + point as f32 * 0.001 - } - }), - points, - dimensions, - ) -} - -fn directional_data(points: usize, dimensions: usize) -> Matrix { - Matrix::new( - diskann_utils::views::Init({ - let mut position = 0usize; - move || { - let point = position / dimensions; - let dimension = position % dimensions; - position += 1; - let angle = std::f32::consts::TAU * point as f32 / points as f32; - match dimension { - 0 => angle.cos(), - 1 => angle.sin(), - _ => 0.0, - } - } - }), - points, - dimensions, - ) -} - -fn sorted_memberships(leaves: &[Vec]) -> Vec> { - let mut memberships: Vec> = leaves - .iter() - .map(|leaf| { - let mut ids = leaf.clone(); - ids.sort_unstable(); - ids - }) - .collect(); - memberships.sort(); - memberships -} - -fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { - assert!( - leaves - .iter() - .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) - ); - let mut counts = vec![0usize; points]; - for leaf in leaves { - let mut ids = leaf.clone(); - ids.sort_unstable(); - ids.dedup(); - assert_eq!(ids.len(), leaf.len(), "duplicate ID inside a leaf"); - for &id in leaf { - assert!((id as usize) < points); - counts[id as usize] += 1; - } - } - assert!(counts.iter().all(|&count| count >= replicas)); -} - -#[test] -fn returns_one_leaf_at_and_below_c_max() { - for points in [7, 8] { - let data = clustered_data(points, 3); - let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); - assert_eq!(leaves, vec![(0..points as u32).collect::>()]); - } -} - -#[test] -fn partition_is_fixed_seed_deterministic_and_bounded() { - let data = clustered_data(96, 8); - let config = config(4, 16, vec![3, 2], 2); - - let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); - let second = partition(data.as_view(), config, Metric::L2).unwrap(); - - assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); - assert_valid_partition(&first, 96, 16, 2); - assert!(first.iter().map(Vec::len).sum::() > 96 * 2); -} - -#[test] -fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { - let data = clustered_data(80, 4); - let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); - - assert_valid_partition(&leaves, 80, 8, 1); -} - -#[test] -fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { - let data = Matrix::new(1.0f32, 24, 4); - let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); - let error = error.downcast::().unwrap(); - - assert!(matches!( - error, - PartitionError::IterationLimit { - size: 24, - limit: MAX_PARTITION_ITERATIONS, - .. - } - )); -} - -#[test] -fn global_merge_canonicalizes_small_leaf_membership() { - let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; - - let merged = global_merge_small(leaves, 4, 8).unwrap(); - - assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); -} - -#[test] -fn global_merge_never_overfills_before_reaching_c_min() { - let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; - - let merged = global_merge_small(leaves, 11, 11).unwrap(); - - assert_eq!( - merged, - vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] - ); -} - -#[test] -fn global_merge_fills_exact_capacity_before_flushing() { - let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); - - assert_eq!(merged, vec![vec![0, 1, 2, 3]]); -} - -#[test] -fn replicas_cover_every_point_once_or_more_per_replica() { - let data = directional_data(72, 5); - let leaves = partition( - data.as_view(), - config(3, 12, vec![3, 2], 3), - Metric::CosineNormalized, - ) - .unwrap(); - - assert_valid_partition(&leaves, 72, 12, 3); -} - -fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) -where - T: crate::utils::VectorRepr + Send + Sync, -{ - let points = 64; - // Partition gathering converts source vectors before GEMM. Exercise conversion - // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. - for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { - let raw: Vec = (0..points * dimensions) - .map(|index| { - let point = index / dimensions; - let dimension = index % dimensions; - ((point * 5 + dimension * 7 + point * dimension) % 23) as u8 - }) - .collect(); - let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); - let converted: Vec = raw.iter().copied().map(&convert).collect(); - let config = config(2, 16, vec![2, 1], 1); - let expected = partition( - MatrixView::try_from(&f32_data, points, dimensions).unwrap(), - config.clone(), - Metric::L2, - ) - .unwrap(); - let actual = partition( - MatrixView::try_from(&converted, points, dimensions).unwrap(), - config, - Metric::L2, - ) - .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); - - assert_valid_partition(&actual, points, 16, 1); - assert_eq!( - sorted_memberships(&actual), - sorted_memberships(&expected), - "{label} dimensions={dimensions}" - ); - } -} - -#[test] -fn f16_partition_matches_f32_across_dimension_boundaries() { - assert_partition_conversion_matches_f32("f16", |value| Half::from_f32(value as f32)); -} - -#[test] -fn u8_partition_matches_f32_across_dimension_boundaries() { - assert_partition_conversion_matches_f32("u8", |value| value); -} - -#[test] -fn i8_partition_matches_f32_across_dimension_boundaries() { - // The same translation in every coordinate preserves L2 ordering. - assert_partition_conversion_matches_f32("i8", |value| value as i8 - 11); -} - -#[test] -fn l2_leader_norms_preserve_scalar_reduction_order() { - fn next(state: &mut u64) -> f32 { - *state ^= *state << 13; - *state ^= *state >> 7; - *state ^= *state << 17; - (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 - } - - // This fixed case sits on opposite sides of the top-1 boundary depending - // on whether leader norms use the original scalar reduction or a SIMD - // reassociation. Point/leader dot products still go through the production - // GEMM; only the setup norm calculation is under test. - let dimensions = 129; - let mut state = 0x3a85_f952_c718_6e49; - let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); - let leader_zero: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); - let leader_one: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); - let data: Vec = leader_zero - .into_iter() - .chain(leader_one) - .chain(point) - .collect(); - let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); - - let clusters = assign_to_leaders( - data, - &[2], - &[0, 1], - 1, - Metric::L2, - &PartitionKernel::new(Metric::L2), - &StripeBufferPool::new((), 0, None), - ) - .unwrap(); - - assert_eq!(clusters, [vec![], vec![2]]); -} - -#[test] -fn all_metrics_produce_valid_partitions() { - let data = directional_data(64, 8); - let config = config(2, 20, vec![2], 1); - - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); - assert_valid_partition(&leaves, 64, 20, 1); - } -} - -#[test] -fn leader_count_is_bounded() { - assert_eq!(sample_num_leaders(1, 1.0), 1); - assert_eq!(sample_num_leaders(10, 0.01), 2); - assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); -} - -#[test] -fn replica_seed_derivation_is_stable_and_distinct() { - assert_eq!(replica_seed(0), 1_000); - assert_eq!(replica_seed(1), 8_919); -} - -#[test] -fn assignment_stripes_use_power_of_two_point_counts() { - assert_eq!(assignment_stripe_point_count(1_000), 128); - assert_eq!(assignment_stripe_point_count(256), 512); - assert_eq!( - assignment_stripe_point_count(1), - MAX_ASSIGNMENT_STRIPE_POINTS - ); -} - -#[test] -fn stripe_buffer_pool_reuses_returned_capacity() { - let pool = StripeBufferPool::new((), 0, None); - let points = { - let mut buffers = pool.get_ref(()); - buffers.points.resize(16, 0.0); - buffers.points.as_ptr() - }; - - let buffers = pool.get_ref(()); - assert_eq!(buffers.points.as_ptr(), points); - assert_eq!(buffers.points.len(), 16); -} - -#[test] -fn leader_assignment_handles_multiple_stripes() { - let points = 2_048; - let data: Vec = (0..points).map(|point| point as f32).collect(); - let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); - let point_ids: Vec = (0..points as u32).collect(); - - let clusters = assign_to_leaders( - data, - &point_ids, - &[0, 2_047], - 1, - Metric::L2, - &PartitionKernel::new(Metric::L2), - &StripeBufferPool::new((), 0, None), - ) - .unwrap(); - - assert_eq!(clusters[0], (0..1_024).collect::>()); - assert_eq!(clusters[1], (1_024..2_048).collect::>()); -} - -#[test] -fn parallel_scatter_matches_serial_order() { - let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); - let assignments: Vec = points - .iter() - .flat_map(|point| [point % 7, (point + 3) % 7]) - .collect(); - - let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); - let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); - - assert_eq!(actual, expected); -} - -#[test] -fn rejects_empty_dataset() { - let data = Matrix::::new(0.0, 0, 4); - let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::EmptyDataset - ); -} - -#[test] -fn rejects_zero_dimensions() { - let data = Matrix::::new(0.0, 4, 0); - let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::EmptyDimensions - ); -} - -#[test] -fn rejects_invalid_gather_output_length() { - let data = Matrix::::new(0.0, 2, 2); - let error = gather_vectors(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "gather output", - expected: 4, - actual: 3, - } - ); -} - -#[test] -fn rejects_assignment_to_an_unknown_leader() { - let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "leader assignment", - expected: 2, - actual: 3, - } - ); -} - -#[test] -fn rejects_empty_and_oversized_leaves() { - for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { - let error = validate_leaves(&leaves, 2).unwrap_err(); - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidLeaf { size, limit: 2 } - ); - } -} diff --git a/diskann/src/graph/pipnn/tests.rs b/diskann/src/graph/pipnn/tests.rs deleted file mode 100644 index d5689eede6..0000000000 --- a/diskann/src/graph/pipnn/tests.rs +++ /dev/null @@ -1,27 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use super::*; -use half::f16; - -#[test] -fn integer_normalized_cosine_uses_unnormalized_cosine() { - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let expected = if metric == Metric::CosineNormalized { - Metric::Cosine - } else { - metric - }; - assert_eq!(effective_metric::(metric), expected); - assert_eq!(effective_metric::(metric), expected); - assert_eq!(effective_metric::(metric), metric); - assert_eq!(effective_metric::(metric), metric); - } -} diff --git a/diskann/tests/pipnn_build_graph.rs b/diskann/tests/pipnn_build_graph.rs deleted file mode 100644 index 12b8c0ec9a..0000000000 --- a/diskann/tests/pipnn_build_graph.rs +++ /dev/null @@ -1,233 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -#![cfg(feature = "pipnn")] -#![allow( - clippy::expect_used, - clippy::unwrap_used, - reason = "deterministic test fixture construction must abort on invalid setup" -)] - -use diskann::graph::config::{self, MaxDegree}; -use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig, build_graph}; -use diskann_utils::views::MatrixView; -use diskann_vector::distance::Metric; -use half::f16; -use rand::{Rng, SeedableRng, rngs::StdRng}; - -fn pipnn_config() -> PiPNNConfig { - PiPNNConfig { - c_max: 4, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - k: 1, - replicas: 1, - } -} - -fn graph_config(metric: Metric, degree: usize) -> diskann::graph::Config { - config::Builder::new_with(degree, MaxDegree::same(), 8, metric.into(), |builder| { - builder.alpha(1.2); - }) - .build() - .unwrap() -} - -fn pool(threads: usize) -> rayon::ThreadPool { - rayon::ThreadPoolBuilder::new() - .num_threads(threads) - .build() - .unwrap() -} - -fn rows(graph: Vec>) -> Vec> { - graph.into_iter().map(Vec::from).collect() -} - -fn assert_graph_invariants( - graph: &[diskann::graph::AdjacencyList], - points: usize, - degree: usize, -) { - assert_eq!(graph.len(), points); - for (source, row) in graph.iter().enumerate() { - assert!(row.len() <= degree); - let mut sorted = row.to_vec(); - sorted.sort_unstable(); - sorted.dedup(); - assert_eq!(sorted.len(), row.len()); - assert!( - row.iter() - .all(|&id| (id as usize) < points && id as usize != source) - ); - } -} - -#[test] -fn builds_a_single_leaf_graph_for_real_dataset_ids() { - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); - let graph = graph_config(Metric::L2, 2); - let pool = pool(2); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); - - assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); - - let graph = graph_config(Metric::L2, 1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - let pruned = build_graph(data, &context).unwrap(); - - assert_graph_invariants(&pruned, 4, 1); - for (source, neighbors) in pruned.iter().enumerate() { - assert_eq!(source.abs_diff(neighbors[0] as usize), 1); - } -} - -#[test] -fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { - let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; - let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); - let graph = graph_config(Metric::L2, 1); - let pool = pool(2); - let config = PiPNNConfig { - c_max: 5, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - k: 4, - replicas: 1, - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context).unwrap(); - - assert_graph_invariants(&actual, 5, 1); - assert!(actual.iter().all(|row| row.len() == 1)); -} - -#[test] -fn rejects_empty_dataset_dimensions_at_the_public_boundary() { - let graph = graph_config(Metric::L2, 2); - let pool = pool(1); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - - let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); - let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); - - assert!(build_graph(no_rows, &context).is_err()); - assert!(build_graph(no_columns, &context).is_err()); -} - -#[test] -fn supports_every_source_type_and_metric() { - fn build(values: &[T], metric: Metric) { - let data = MatrixView::try_from(values, 6, 2).unwrap(); - let graph = graph_config(metric, 2); - let pool = pool(2); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); - let actual = build_graph(data, &context).unwrap(); - assert_graph_invariants(&actual, 6, 2); - } - - let values = [ - 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, - ]; - for metric in [ - Metric::L2, - Metric::Cosine, - Metric::CosineNormalized, - Metric::InnerProduct, - ] { - build(&values, metric); - } - build(&values.map(f16::from_f32), Metric::L2); - build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); - build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); -} - -#[test] -fn integer_normalized_cosine_matches_cosine() { - fn assert_match(values: &[T]) { - let data = MatrixView::try_from(values, 8, 2).unwrap(); - let pool = pool(2); - let build = |metric| { - let graph = graph_config(metric, 2); - let config = PiPNNConfig { - c_max: 8, - c_min: 1, - p_samp: 0.5, - fanout: vec![2], - k: 1, - replicas: 1, - }; - let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); - rows(build_graph(data, &context).unwrap()) - }; - assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); - } - - assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); - assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); -} - -#[test] -fn is_deterministic_for_a_fixed_pool_size() { - let data: Vec = (0..96 * 4) - .map(|value| ((value * 17 + 3) % 101) as f32) - .collect(); - let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); - let graph = graph_config(Metric::L2, 8); - let pool = pool(4); - let config = PiPNNConfig { - c_max: 16, - c_min: 4, - p_samp: 0.25, - fanout: vec![3, 2], - k: 3, - replicas: 2, - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - let first = build_graph(data, &context).unwrap(); - let second = build_graph(data, &context).unwrap(); - - assert_eq!(first, second); - assert_graph_invariants(&first, 96, 8); -} - -#[test] -fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { - let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); - for case in 0..24 { - let points = rng.random_range(4..=32); - let dimensions = rng.random_range(1..=8); - let c_max = rng.random_range(4..=points.min(12)); - let c_min = rng.random_range(1..=c_max); - let degree = rng.random_range(1..=points.min(8)); - let values: Vec = (0..points * dimensions) - .map(|_| rng.random_range(-10.0..10.0)) - .collect(); - let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); - let graph = graph_config(Metric::L2, degree); - let pool = pool(2); - let config = PiPNNConfig { - c_max, - c_min, - p_samp: 0.5, - fanout: vec![2], - k: rng.random_range(1..=3), - replicas: rng.random_range(1..=2), - }; - let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - - let actual = build_graph(data, &context) - .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); - assert_graph_invariants(&actual, points, degree); - } -} diff --git a/diskann/tests/pipnn_config.rs b/diskann/tests/pipnn_config.rs deleted file mode 100644 index 19a0d1ce8a..0000000000 --- a/diskann/tests/pipnn_config.rs +++ /dev/null @@ -1,123 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -#![cfg(feature = "pipnn")] -#![allow( - clippy::expect_used, - clippy::unwrap_used, - reason = "deterministic test fixture construction must abort on invalid setup" -)] - -use diskann::graph::config::{self, MaxDegree}; -use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig}; -use diskann_vector::distance::Metric; - -fn pipnn_config() -> PiPNNConfig { - PiPNNConfig { - c_max: 512, - c_min: 64, - p_samp: 0.01, - fanout: vec![10, 3], - k: 2, - replicas: 1, - } -} - -fn graph_config(metric: Metric, alpha: f32) -> diskann::graph::Config { - config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { - builder.alpha(alpha); - }) - .build() - .unwrap() -} - -fn pool() -> rayon::ThreadPool { - rayon::ThreadPoolBuilder::new() - .num_threads(2) - .build() - .unwrap() -} - -#[test] -fn rejects_each_invalid_algorithm_parameter() { - let graph = graph_config(Metric::L2, 1.2); - let pool = pool(); - let mut cases = [ - PiPNNConfig { - c_max: 0, - ..pipnn_config() - }, - PiPNNConfig { - c_min: 0, - ..pipnn_config() - }, - PiPNNConfig { - c_min: 513, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: 0.0, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: -0.01, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: 1.01, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: f64::NAN, - ..pipnn_config() - }, - PiPNNConfig { - fanout: Vec::new(), - ..pipnn_config() - }, - PiPNNConfig { - fanout: vec![1, 0], - ..pipnn_config() - }, - PiPNNConfig { - fanout: vec![17], - ..pipnn_config() - }, - PiPNNConfig { - k: 0, - ..pipnn_config() - }, - PiPNNConfig { - replicas: 0, - ..pipnn_config() - }, - ]; - - for config in &mut cases { - let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) - .expect_err("invalid PiPNN config must be rejected"); - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); - } -} - -#[test] -fn rejects_graph_policy_for_a_different_metric() { - let graph = graph_config(Metric::InnerProduct, 1.2); - let pool = pool(); - - let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); - - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); - assert!(error.to_string().contains("prune kind")); -} - -#[test] -fn does_not_add_alpha_validation_beyond_graph_config() { - let pool = pool(); - for alpha in [0.9, f32::NAN, f32::INFINITY] { - let graph = graph_config(Metric::L2, alpha); - PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - } -} From 2b963c97f79d7a8b5b2f2b28490e27ff27665195 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:53:28 +0000 Subject: [PATCH 23/58] test(pipnn): adapt assertions to main errors --- diskann/src/graph/pipnn/mod.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index e715498690..d4a34f7fcc 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -680,9 +680,8 @@ mod config_tests { ]; for config in &mut cases { - let error = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) .expect_err("invalid PiPNN config must be rejected"); - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); } } @@ -693,7 +692,6 @@ mod config_tests { let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); - assert_eq!(error.kind(), diskann::ANNErrorKind::IndexConfigError); assert!(error.to_string().contains("prune kind")); } From 4760af80b199d02e5d7654b2718704f14c7766e7 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:14:50 +0000 Subject: [PATCH 24/58] refactor(pipnn): consume positional robust prune Adapt PiPNN-owned preparation and ID translation to the behavior-preserving positional kernel introduced by #1315. --- diskann/src/graph/pipnn/finalization.rs | 36 ++++++++++++++----------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index a10d4df027..e060ad621b 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -35,14 +35,9 @@ //! | bounded list | none | move list directly to output | //! | overfull list | source and occlusion distances | reuse Rayon-job workspace | -use std::convert::Infallible; - use crate::{ ANNError, ANNResult, - graph::{ - AdjacencyList, Config, - internal::{SortedNeighbors, robust_prune as prune}, - }, + graph::{AdjacencyList, Config, internal::SortedNeighbors, robust_prune as prune}, neighbor::Neighbor, utils::VectorRepr, }; @@ -62,6 +57,8 @@ pub(crate) enum FinalizationError { candidate: u32, points: usize, }, + #[error("candidate count {actual} exceeds the u16 position limit {max}")] + TooManyCandidates { actual: usize, max: usize }, } /// Per-Rayon-job preparation and kernel state retained across source points. @@ -71,7 +68,7 @@ pub(crate) enum FinalizationError { #[derive(Default)] struct Workspace { pool: Vec>, - prepared: Vec>, + prepared: Vec<(f32, Option)>, states: Vec, } @@ -122,8 +119,12 @@ where })); let candidate_count = workspace.pool.len(); - prune::validate_candidate_count::(candidate_count) - .map_err(ANNError::new)?; + if candidate_count > u16::MAX as usize { + return Err(ANNError::new(FinalizationError::TooManyCandidates { + actual: candidate_count, + max: u16::MAX as usize, + })); + } workspace.prepared.clear(); workspace .prepared @@ -137,8 +138,7 @@ where .prepared .extend(sorted.iter().filter_map(|neighbor| { let id = *neighbor.id(); - (id != source_id) - .then(|| prune::Candidate::new(id, *neighbor.distance(), id)) + (id != source_id).then(|| (*neighbor.distance(), Some(id))) })); } workspace @@ -161,17 +161,21 @@ where graph.alpha(), graph.prune_kind(), |left, right| { - Ok::<_, Infallible>(distance.evaluate_similarity( + distance.evaluate_similarity( data.row(*left as usize), data.row(*right as usize), - )) + ) }, - ) - .map_err(ANNError::new)?; + ); let mut guard = source_candidates.resize(selected); for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { - *destination = *workspace.prepared[state.selected_position()].id(); + *destination = + workspace.prepared[state.neighbor as usize] + .1 + .ok_or_else(|| { + ANNError::message("RobustPrune selected an unavailable candidate") + })?; } guard.finish(selected); Ok(source_candidates) From 1a18a8923c9a3d26497a793b8094eb8d359f184f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 02:57:30 +0000 Subject: [PATCH 25/58] refactor(pipnn): use sorted prune input Pass the existing SortedNeighbors witness into internal prune so source-distance ordering is enforced by type rather than caller documentation. --- diskann/src/graph/pipnn/finalization.rs | 33 +++++++++++-------------- 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index e060ad621b..7d0ccb722d 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -37,7 +37,10 @@ use crate::{ ANNError, ANNResult, - graph::{AdjacencyList, Config, internal::SortedNeighbors, robust_prune as prune}, + graph::{ + AdjacencyList, Config, + internal::{SortedNeighbors, prune}, + }, neighbor::Neighbor, utils::VectorRepr, }; @@ -130,17 +133,15 @@ where .prepared .try_reserve(candidate_count) .map_err(ANNError::new)?; - { - // Sorting/capping precedes source exclusion so filtering cannot - // backfill with farther candidates. - let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); - workspace - .prepared - .extend(sorted.iter().filter_map(|neighbor| { - let id = *neighbor.id(); - (id != source_id).then(|| (*neighbor.distance(), Some(id))) - })); - } + + // Sorting/capping precedes source exclusion so filtering cannot + // backfill with farther candidates. Passing this witness into + // RobustPrune makes source-distance order part of its input type. + let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); + workspace.prepared.extend(sorted.iter().map(|neighbor| { + let id = *neighbor.id(); + (*neighbor.distance(), (id != source_id).then_some(id)) + })); workspace .states .try_reserve( @@ -155,6 +156,7 @@ where .resize(workspace.prepared.len(), prune::State::default()); let selected = prune::robust_prune( + &sorted, &workspace.prepared, workspace.states.as_mut_slice(), degree, @@ -170,12 +172,7 @@ where let mut guard = source_candidates.resize(selected); for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { - *destination = - workspace.prepared[state.neighbor as usize] - .1 - .ok_or_else(|| { - ANNError::message("RobustPrune selected an unavailable candidate") - })?; + *destination = *sorted[state.neighbor as usize].id(); } guard.finish(selected); Ok(source_candidates) From d4524b7beaff20949de73aa6b0539d634ef6dbf4 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:47:35 +0000 Subject: [PATCH 26/58] refactor(pipnn): use direct leaf matrix input --- diskann/src/graph/pipnn/leaf_build.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 1c60a40c4e..3667c2f1e6 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -28,7 +28,7 @@ use diskann_vector::distance::Metric; use rayon::prelude::*; use super::leaf_kernel::{ - LeafInput, LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, + LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, leaf_output_len, }; @@ -356,7 +356,7 @@ where }, })?; kernel - .nearest_neighbors(LeafInput { dots }, output, &mut buffers.kernel_workspace) + .nearest_neighbors(dots, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; buffers.prepare_local_adjacency(point_ids.len())?; From c73e20dca25c4e1ce5d3672fe3efdaaf541592c0 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 05:56:05 +0000 Subject: [PATCH 27/58] fix(pipnn): validate leaf k capacity Reject k outside 1..=3 at context construction so production never reaches an unsupported leaf-kernel width. --- diskann/src/graph/pipnn/mod.rs | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index d4a34f7fcc..5972bc15de 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -167,7 +167,7 @@ pub struct PiPNNConfig { pub p_samp: f64, /// Number of nearest leaders retained at each overlapping partition level. pub fanout: Vec, - /// Number of nearest neighbors selected within each leaf. + /// Number of nearest neighbors selected within each leaf (`1..=3`). pub k: usize, /// Number of independent partition passes over the dataset. pub replicas: usize, @@ -207,8 +207,12 @@ impl PiPNNConfig { partition_kernel::MAX_PARTITION_FANOUT ))); } - if self.k == 0 { - return Err(config_error("k must be greater than zero")); + if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.k) { + return Err(config_error(format!( + "k ({}) must be in [1, {}]", + self.k, + leaf_kernel::MAX_LEAF_NEIGHBORS + ))); } if self.replicas == 0 { return Err(config_error("replicas must be greater than zero")); @@ -362,7 +366,7 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod build_graph_tests { - use super::{PiPNNBuildContext, PiPNNConfig, build_graph}; + use super::{PiPNNBuildContext, PiPNNConfig, build_graph, leaf_kernel}; use crate::graph::config::{self, MaxDegree}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; @@ -442,7 +446,7 @@ mod build_graph_tests { } #[test] - fn prunes_complete_single_leaf_candidates_to_the_graph_degree() { + fn prunes_overfull_single_leaf_candidates_to_the_graph_degree() { let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); let graph = graph_config(Metric::L2, 1); @@ -452,7 +456,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - k: 4, + k: leaf_kernel::MAX_LEAF_NEIGHBORS, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -594,7 +598,7 @@ mod build_graph_tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod config_tests { - use super::{PiPNNBuildContext, PiPNNConfig}; + use super::{PiPNNBuildContext, PiPNNConfig, leaf_kernel}; use crate::graph::config::{self, MaxDegree}; use diskann_vector::distance::Metric; @@ -673,6 +677,10 @@ mod config_tests { k: 0, ..pipnn_config() }, + PiPNNConfig { + k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, + ..pipnn_config() + }, PiPNNConfig { replicas: 0, ..pipnn_config() From e2374d4f7310c533cab8055886e9d7df7220877d Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:41:17 +0000 Subject: [PATCH 28/58] refactor(pipnn): require sorted leaf IDs Partitioning is the only production source and emits sorted unique IDs. Reject unsorted input linearly and remove the HashSet fallback. --- diskann/src/graph/pipnn/leaf_build.rs | 62 +++++++-------------------- 1 file changed, 16 insertions(+), 46 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 3667c2f1e6..6d54913748 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -5,7 +5,8 @@ //! Leaf-local graph construction and candidate accumulation. //! -//! Partitioning supplies leaves as global point IDs. For each leaf this module: +//! Partitioning supplies strictly increasing, unique global point IDs per leaf. +//! For each leaf this module: //! //! 1. validates IDs and converts only those point vectors to reusable `f32` scratch; //! 2. computes the lower triangle of `A · Aᵀ`; @@ -17,10 +18,7 @@ //! high-water length; every consumer therefore receives an explicit active //! prefix rather than treating `Vec::len()` as the current leaf shape. -use std::{ - collections::{HashSet, TryReserveError}, - sync::Mutex, -}; +use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -49,12 +47,16 @@ pub(crate) enum LeafBuildError { }, #[error("point ID {point} appears more than once in leaf {leaf}")] DuplicatePointId { leaf: usize, point: u32 }, + #[error("point IDs in leaf {leaf} are not strictly increasing")] + UnsortedPointIds { leaf: usize }, #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] ShapeOverflow { leaf: usize, rows: usize, columns: usize, }, + #[error("failed to form {buffer} view for leaf {leaf}")] + InvalidView { leaf: usize, buffer: &'static str }, #[error("failed to reserve {additional} values for {buffer}")] Allocation { buffer: &'static str, @@ -100,7 +102,6 @@ struct LeafBuffers { neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, - seen_ids: HashSet, } impl LeafBuffers { @@ -283,24 +284,14 @@ where }); } } - if point_ids.is_sorted() { - if let Some(pair) = point_ids.windows(2).find(|pair| pair[0] == pair[1]) { + if let Some(pair) = point_ids.windows(2).find(|pair| pair[0] >= pair[1]) { + if pair[0] == pair[1] { return Err(LeafBuildError::DuplicatePointId { leaf, point: pair[0], }); } - } else { - buffers.seen_ids.clear(); - buffers - .seen_ids - .try_reserve(point_ids.len()) - .map_err(|source| allocation_error("leaf ID set", point_ids.len(), source))?; - for &point in point_ids { - if !buffers.seen_ids.insert(point) { - return Err(LeafBuildError::DuplicatePointId { leaf, point }); - } - } + return Err(LeafBuildError::UnsortedPointIds { leaf }); } let leaf_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; if leaf_k == 0 { @@ -334,26 +325,18 @@ where ) .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; let dots = MatrixView::try_from(&buffers.dots[..dot_count], point_ids.len(), point_ids.len()) - .map_err(|error| LeafBuildError::Kernel { + .map_err(|_| LeafBuildError::InvalidView { leaf, - source: LeafKernelError::InvalidBufferLength { - buffer: "leaf dot-product matrix", - expected: dot_count, - actual: error.into_inner().len(), - }, + buffer: "leaf dot-product matrix", })?; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), leaf_k, ) - .map_err(|error| LeafBuildError::Kernel { + .map_err(|_| LeafBuildError::InvalidView { leaf, - source: LeafKernelError::InvalidBufferLength { - buffer: "output", - expected: neighbor_value_count, - actual: error.into_inner().len(), - }, + buffer: "leaf output", })?; kernel .nearest_neighbors(dots, output, &mut buffers.kernel_workspace) @@ -570,19 +553,6 @@ mod tests { })); } - #[test] - fn global_id_translation_is_independent_of_leaf_order() { - let data = [0.0_f32, 10.0, 20.0, 30.0, 40.0]; - let leaves = vec![vec![4, 1, 3]]; - - let graph = build(view(&data, 5, 1), &leaves, 2, Metric::L2).unwrap(); - - assert_eq!( - adjacency_lists(graph), - [vec![], vec![3, 4], vec![], vec![1, 4], vec![1, 3]] - ); - } - fn source_graph(data: &[T], points: usize, dimensions: usize) -> Vec> where T: crate::utils::VectorRepr + 'static, @@ -710,8 +680,8 @@ mod tests { Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) )); assert!(matches!( - build(view(&data, 2, 1), &[vec![1, 0, 1]], 1, Metric::L2), - Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 1 }) + build(view(&data, 2, 1), &[vec![1, 0]], 1, Metric::L2), + Err(LeafBuildError::UnsortedPointIds { leaf: 0 }) )); } From 9ce0b534376ce919542fc15ffbb58a1230ad777b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 06:41:17 +0000 Subject: [PATCH 29/58] refactor(pipnn): borrow partition configuration Rely on validated MatrixView shape and avoid cloning PiPNNConfig and its fanout vector. --- diskann/src/graph/pipnn/mod.rs | 9 +-------- diskann/src/graph/pipnn/partitioning.rs | 26 ++++++++++++------------- 2 files changed, 14 insertions(+), 21 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 5972bc15de..dff75df36a 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -292,19 +292,12 @@ where data.nrows() ))); } - data.nrows().checked_mul(data.ncols()).ok_or_else(|| { - ANNError::message(format!( - "PiPNN dataset shape {} x {} overflows usize", - data.nrows(), - data.ncols() - )) - })?; // Integer source vectors are not guaranteed unit-normalized after conversion, // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition(data, context.config.clone(), metric))?; + .in_scope(|| partitioning::partition(data, &context.config, metric))?; // `leaves` is consumed here. Workers borrow individual ID lists during the // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 5de9a68cf5..5fd341d9cb 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -141,7 +141,7 @@ type StripeBufferPool = ObjectPool; /// covered once per replica. The caller installs the operation in its pool. pub(crate) fn partition( data: MatrixView<'_, T>, - config: PiPNNConfig, + config: &PiPNNConfig, metric: Metric, ) -> ANNResult>> where @@ -167,7 +167,7 @@ where for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = - partition_replica(data, &config, metric, &kernel, seed, &stripe_buffers)?; + partition_replica(data, config, metric, &kernel, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::new)?; @@ -880,7 +880,7 @@ mod tests { fn returns_one_leaf_at_and_below_c_max() { for points in [7, 8] { let data = clustered_data(points, 3); - let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); assert_eq!(leaves, vec![(0..points as u32).collect::>()]); } } @@ -890,8 +890,8 @@ mod tests { let data = clustered_data(96, 8); let config = config(4, 16, vec![3, 2], 2); - let first = partition(data.as_view(), config.clone(), Metric::L2).unwrap(); - let second = partition(data.as_view(), config, Metric::L2).unwrap(); + let first = partition(data.as_view(), &config, Metric::L2).unwrap(); + let second = partition(data.as_view(), &config, Metric::L2).unwrap(); assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); assert_valid_partition(&first, 96, 16, 2); @@ -901,7 +901,7 @@ mod tests { #[test] fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { let data = clustered_data(80, 4); - let leaves = partition(data.as_view(), config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); assert_valid_partition(&leaves, 80, 8, 1); } @@ -909,7 +909,7 @@ mod tests { #[test] fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { let data = Matrix::new(1.0f32, 24, 4); - let error = partition(data.as_view(), config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); let error = error.downcast::().unwrap(); assert!(matches!( @@ -955,7 +955,7 @@ mod tests { let data = directional_data(72, 5); let leaves = partition( data.as_view(), - config(3, 12, vec![3, 2], 3), + &config(3, 12, vec![3, 2], 3), Metric::CosineNormalized, ) .unwrap(); @@ -983,13 +983,13 @@ mod tests { let config = config(2, 16, vec![2, 1], 1); let expected = partition( MatrixView::try_from(&f32_data, points, dimensions).unwrap(), - config.clone(), + &config, Metric::L2, ) .unwrap(); let actual = partition( MatrixView::try_from(&converted, points, dimensions).unwrap(), - config, + &config, Metric::L2, ) .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); @@ -1069,7 +1069,7 @@ mod tests { Metric::CosineNormalized, Metric::InnerProduct, ] { - let leaves = partition(data.as_view(), config.clone(), metric).unwrap(); + let leaves = partition(data.as_view(), &config, metric).unwrap(); assert_valid_partition(&leaves, 64, 20, 1); } } @@ -1150,7 +1150,7 @@ mod tests { #[test] fn rejects_empty_dataset() { let data = Matrix::::new(0.0, 0, 4); - let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); assert_eq!( error.downcast::().unwrap(), @@ -1161,7 +1161,7 @@ mod tests { #[test] fn rejects_zero_dimensions() { let data = Matrix::::new(0.0, 4, 0); - let error = partition(data.as_view(), config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); assert_eq!( error.downcast::().unwrap(), From 83109b4cad19625a91d621aa3556ad9d1e6bd3c3 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:49:02 +0000 Subject: [PATCH 30/58] refactor(pipnn): dispatch partition stages once Run partition orchestration under one architecture/metric specialization and reuse a runtime-sized tracker instead of imposing a fanout cap. --- diskann/src/graph/pipnn/mod.rs | 20 +-- diskann/src/graph/pipnn/partitioning.rs | 170 +++++++++++++++++------- 2 files changed, 128 insertions(+), 62 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index dff75df36a..5788767659 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -90,9 +90,8 @@ //! [`partition_kernel::PartitionInput`] bundles that tile with typed //! [`partition_kernel::PartitionScales`]. A prepared //! [`partition_kernel::PartitionKernel`] writes sorted leader-local positions to -//! caller-owned output. Fanout is output column count and is bounded by -//! [`partition_kernel::MAX_PARTITION_FANOUT`]. Module documentation describes -//! scale units, validation, `process_points`, and tracker insertion. +//! caller-owned output. Fanout is the output column count and cannot exceed the +//! leaders available for that partition call. //! //! ## [`leaf_kernel`] //! @@ -197,15 +196,8 @@ impl PiPNNConfig { if self.fanout.is_empty() { return Err(config_error("fanout must not be empty")); } - if let Some(&fanout) = self - .fanout - .iter() - .find(|&&fanout| !(1..=partition_kernel::MAX_PARTITION_FANOUT).contains(&fanout)) - { - return Err(config_error(format!( - "fanout ({fanout}) must be in [1, {}]", - partition_kernel::MAX_PARTITION_FANOUT - ))); + if self.fanout.contains(&0) { + return Err(config_error("fanout values must be greater than zero")); } if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.k) { return Err(config_error(format!( @@ -662,10 +654,6 @@ mod config_tests { fanout: vec![1, 0], ..pipnn_config() }, - PiPNNConfig { - fanout: vec![17], - ..pipnn_config() - }, PiPNNConfig { k: 0, ..pipnn_config() diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 5fd341d9cb..4bc7cd164e 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -46,12 +46,19 @@ use diskann_utils::{ views::{MatrixView, MutMatrixView}, }; use diskann_vector::{Norm, distance::Metric, norm::FastL2NormSquared}; +use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, +}; use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; use super::{ PiPNNConfig, - partition_kernel::{PartitionInput, PartitionKernel, PartitionScales}, + kernel_metric::{KernelMetric, MetricVisitor, visit_metric}, + partition_kernel::{ + PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders_for, + }, }; // Private algorithm and batching constants live together. None are user policy. @@ -111,6 +118,7 @@ struct StripeBuffers { points: Vec, dots: Vec, point_scales: Vec, + kernel: PartitionKernelWorkspace, } impl AsPooled<()> for StripeBuffers { @@ -146,6 +154,70 @@ pub(crate) fn partition( ) -> ANNResult>> where T: VectorRepr + Send + Sync, +{ + arch::dispatch1_no_features( + RunPartitionStage, + PartitionStageCall { + data, + config, + metric, + }, + ) +} + +struct PartitionStageCall<'a, T> { + data: MatrixView<'a, T>, + config: &'a PiPNNConfig, + metric: Metric, +} + +struct RunPartitionStage; + +impl Target1>>, PartitionStageCall<'_, T>> for RunPartitionStage +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + Send + Sync, +{ + fn run(self, arch: A, call: PartitionStageCall<'_, T>) -> ANNResult>> { + visit_metric(call.metric, ExecutePartitionStage { arch, call }) + } +} + +struct ExecutePartitionStage<'a, A, T> { + arch: A, + call: PartitionStageCall<'a, T>, +} + +impl MetricVisitor for ExecutePartitionStage<'_, A, T> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + Send + Sync, +{ + type Output = ANNResult>>; + + fn visit(self) -> Self::Output { + partition_for::(self.arch, self.call.data, self.call.config) + } +} + +fn partition_for( + arch: A, + data: MatrixView<'_, T>, + config: &PiPNNConfig, +) -> ANNResult>> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, + T: VectorRepr + Send + Sync, { let points = data.nrows(); if points == 0 { @@ -159,15 +231,11 @@ where } let mut leaves = Vec::new(); - // Prepare metric and ISA dispatch before replicas spawn Rayon work. The - // Copy handle is shared read-only; every stripe calls its direct function - // pointer instead of redispatching in the recursive hot path. - let kernel = PartitionKernel::new(metric); let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = - partition_replica(data, config, metric, &kernel, seed, &stripe_buffers)?; + partition_replica::(arch, data, config, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::new)?; @@ -177,15 +245,19 @@ where Ok(leaves) } -fn partition_replica( +fn partition_replica( + arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, - metric: Metric, - kernel: &PartitionKernel, seed: u64, stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr + Send + Sync, { let initial_indices = point_ids(data.nrows())?; @@ -222,11 +294,10 @@ where .par_iter_mut() .zip(work.into_par_iter()) .for_each(|(slot, item)| { - *slot = Some(partition_one_level( + *slot = Some(partition_one_level::( + arch, data, config, - metric, - kernel, item, stripe_buffers, )); @@ -259,15 +330,19 @@ where })) } -fn partition_one_level( +fn partition_one_level( + arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, - metric: Metric, - kernel: &PartitionKernel, item: WorkItem, stripe_buffers: &StripeBufferPool, ) -> ANNResult<(Vec, Vec>)> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr + Send + Sync, { let points = item.indices.len(); @@ -277,15 +352,8 @@ where config.p_samp, mix_seed(item.seed, points as u64), )?; - let clusters = assign_to_leaders( - data, - &item.indices, - &leaders, - fanout, - metric, - kernel, - stripe_buffers, - )?; + let clusters = + assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; let mut pending = Vec::new(); let mut finished = Vec::new(); @@ -344,16 +412,20 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { /// its stripes. The flat assignment matrix preserves point order and is then /// scattered into per-leader clusters; preserving order is required for fixed /// seed determinism in later recursion levels. -fn assign_to_leaders( +fn assign_to_leaders( + arch: A, data: MatrixView<'_, T>, point_ids: &[u32], leader_ids: &[u32], fanout: usize, - metric: Metric, - kernel: &PartitionKernel, stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr + Send + Sync, { let dimension_count = data.ncols(); @@ -361,7 +433,7 @@ where let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; gather_vectors(data, leader_ids, &mut leader_values)?; - let mut leader_scales = if matches!(metric, Metric::L2 | Metric::Cosine) { + let mut leader_scales = if matches!(M::METRIC, Metric::L2 | Metric::Cosine) { filled_vec(leader_ids.len(), 0.0f32)? } else { Vec::new() @@ -375,7 +447,7 @@ where // SIMD norm changes low bits and can send near-tied points down different // recursive partition paths. *scale = leader_vector.iter().map(|value| value * value).sum(); - if metric == Metric::Cosine { + if M::METRIC == Metric::Cosine { *scale = scale.sqrt(); } } @@ -405,13 +477,12 @@ where { let first_point = worker_first + stripe * stripe_points; let stripe_point_count = stripe_assignments.len() / fanout; - assign_stripe( + assign_stripe::( + arch, data, &point_ids[first_point..first_point + stripe_point_count], &leader_values, &leader_scales, - metric, - kernel, fanout, &mut buffers, stripe_assignments, @@ -425,18 +496,22 @@ where #[inline] #[allow(clippy::too_many_arguments)] -fn assign_stripe( +fn assign_stripe( + arch: A, data: MatrixView<'_, T>, point_ids: &[u32], leader_values: &[f32], leader_scales: &[f32], - metric: Metric, - kernel: &PartitionKernel, fanout: usize, buffers: &mut StripeBuffers, assignments: &mut [u32], ) -> ANNResult<()> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr, { let point_count = point_ids.len(); @@ -458,6 +533,7 @@ where points: point_buffer, dots: dot_buffer, point_scales: point_scale_buffer, + kernel: kernel_workspace, } = buffers; let point_values = &mut point_buffer[..point_values_len]; let dots = &mut dot_buffer[..dots_len]; @@ -476,7 +552,7 @@ where ) .map_err(ANNError::new)?; - let point_scales = if metric == Metric::Cosine { + let point_scales = if M::METRIC == Metric::Cosine { grow_fallible(point_scale_buffer, point_count, 0.0)?; let point_scales = &mut point_scale_buffer[..point_count]; for (scale, point_values) in point_scales @@ -489,7 +565,7 @@ where } else { &[] }; - let scales = match metric { + let scales = match M::METRIC { Metric::L2 => PartitionScales::L2 { leader_squared_norms: leader_scales, }, @@ -513,9 +589,13 @@ where actual: error.into_inner().len(), }) })?; - kernel - .nearest_leaders(PartitionInput { dots, scales }, output) - .map_err(ANNError::new) + nearest_leaders_for::( + arch, + PartitionInput { dots, scales }, + output, + kernel_workspace, + ) + .map_err(ANNError::new) } fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> @@ -1044,13 +1124,12 @@ mod tests { .collect(); let data = MatrixView::try_from(data.as_slice(), 3, dimensions).unwrap(); - let clusters = assign_to_leaders( + let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( + diskann_wide::ARCH, data, &[2], &[0, 1], 1, - Metric::L2, - &PartitionKernel::new(Metric::L2), &StripeBufferPool::new((), 0, None), ) .unwrap(); @@ -1118,13 +1197,12 @@ mod tests { let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); let point_ids: Vec = (0..points as u32).collect(); - let clusters = assign_to_leaders( + let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( + diskann_wide::ARCH, data, &point_ids, &[0, 2_047], 1, - Metric::L2, - &PartitionKernel::new(Metric::L2), &StripeBufferPool::new((), 0, None), ) .unwrap(); From 8d69d6d52fee6c832c232cbd67a712389a7db9ba Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:49:02 +0000 Subject: [PATCH 31/58] refactor(pipnn): dispatch leaf stages once Keep architecture and metric concrete across the Rayon leaf pass and call the generic kernel directly. --- diskann/src/graph/pipnn/leaf_build.rs | 118 ++++++++++++++++++++++---- 1 file changed, 101 insertions(+), 17 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6d54913748..c5cc541a1b 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -23,11 +23,18 @@ use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; use diskann_vector::distance::Metric; +use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, +}; use rayon::prelude::*; -use super::leaf_kernel::{ - LeafKernel, LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, - leaf_output_len, +use super::{ + kernel_metric::{KernelMetric, MetricVisitor, visit_metric}, + leaf_kernel::{ + LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, leaf_output_len, + nearest_neighbors_for, + }, }; /// Failure while converting leaves into direct graph candidates. @@ -216,7 +223,6 @@ impl DirectCandidates { } /// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. -#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. pub(crate) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, @@ -225,6 +231,84 @@ pub(crate) fn build_leaf_candidates( ) -> Result>, LeafBuildError> where T: VectorRepr + 'static, +{ + arch::dispatch1_no_features( + RunLeafStage, + LeafStageCall { + data, + leaves, + requested_k, + metric, + }, + ) +} + +struct LeafStageCall<'a, T> { + data: MatrixView<'a, T>, + leaves: Vec>, + requested_k: usize, + metric: Metric, +} + +struct RunLeafStage; + +impl Target1>, LeafBuildError>, LeafStageCall<'_, T>> + for RunLeafStage +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + 'static, +{ + fn run( + self, + arch: A, + call: LeafStageCall<'_, T>, + ) -> Result>, LeafBuildError> { + visit_metric(call.metric, ExecuteLeafStage { arch, call }) + } +} + +struct ExecuteLeafStage<'a, A, T> { + arch: A, + call: LeafStageCall<'a, T>, +} + +impl MetricVisitor for ExecuteLeafStage<'_, A, T> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + 'static, +{ + type Output = Result>, LeafBuildError>; + + fn visit(self) -> Self::Output { + build_leaf_candidates_for::( + self.arch, + self.call.data, + self.call.leaves, + self.call.requested_k, + ) + } +} + +#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. +fn build_leaf_candidates_for( + arch: A, + data: MatrixView<'_, T>, + leaves: Vec>, + requested_k: usize, +) -> Result>, LeafBuildError> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, + T: VectorRepr + 'static, { if data.ncols() == 0 { return Err(LeafBuildError::EmptyDimensions); @@ -234,18 +318,15 @@ where } let candidates = DirectCandidates::new(data.nrows())?; - // Metric and ISA are selected before Rayon workers start. Workers share - // this Copy handle; each output view supplies its leaf-specific width. - let kernel = LeafKernel::new(metric); leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf( + build_leaf::( + arch, data, leaf, point_ids, requested_k, - &kernel, buffers, &candidates, ) @@ -256,20 +337,24 @@ where /// Build and publish one leaf's symmetric neighbor lists. /// -/// Validation precedes all dataset indexing. Sorted partition output takes the -/// adjacent-duplicate path, while arbitrary-order callers use `seen_ids`. The -/// active lengths computed after `prepare` must be used for every later slice, -/// because the reusable vectors may still be longer than this leaf. -fn build_leaf( +/// Validation precedes all dataset indexing and rejects IDs that are not strictly +/// increasing. Active lengths computed after `prepare` must be used for every +/// later slice because reusable vectors may remain longer than this leaf. +fn build_leaf( + arch: A, data: MatrixView<'_, T>, leaf: usize, point_ids: &[u32], requested_k: usize, - kernel: &LeafKernel, buffers: &mut LeafBuffers, candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, T: VectorRepr + 'static, { if point_ids.is_empty() { @@ -338,8 +423,7 @@ where leaf, buffer: "leaf output", })?; - kernel - .nearest_neighbors(dots, output, &mut buffers.kernel_workspace) + nearest_neighbors_for::(arch, dots, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; buffers.prepare_local_adjacency(point_ids.len())?; From 0b27e69e2a9d0cbe3e7070c05a74bf220b7a2fda Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:08:19 +0000 Subject: [PATCH 32/58] refactor(pipnn): dispatch once per graph build --- diskann/src/graph/pipnn/leaf_build.rs | 170 ++++++++---------- diskann/src/graph/pipnn/mod.rs | 222 +++++++++++------------- diskann/src/graph/pipnn/partitioning.rs | 173 +++++++++--------- 3 files changed, 263 insertions(+), 302 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index c5cc541a1b..3ff70284d9 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -22,18 +22,14 @@ use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_vector::distance::Metric; -use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, -}; +use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::{KernelMetric, MetricVisitor, visit_metric}, + kernel_metric::KernelMetric, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, leaf_output_len, - nearest_neighbors_for, + nearest_neighbors, }, }; @@ -223,80 +219,10 @@ impl DirectCandidates { } /// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. -pub(crate) fn build_leaf_candidates( - data: MatrixView<'_, T>, - leaves: Vec>, - requested_k: usize, - metric: Metric, -) -> Result>, LeafBuildError> -where - T: VectorRepr + 'static, -{ - arch::dispatch1_no_features( - RunLeafStage, - LeafStageCall { - data, - leaves, - requested_k, - metric, - }, - ) -} - -struct LeafStageCall<'a, T> { - data: MatrixView<'a, T>, - leaves: Vec>, - requested_k: usize, - metric: Metric, -} - -struct RunLeafStage; - -impl Target1>, LeafBuildError>, LeafStageCall<'_, T>> - for RunLeafStage -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - T: VectorRepr + 'static, -{ - fn run( - self, - arch: A, - call: LeafStageCall<'_, T>, - ) -> Result>, LeafBuildError> { - visit_metric(call.metric, ExecuteLeafStage { arch, call }) - } -} - -struct ExecuteLeafStage<'a, A, T> { - arch: A, - call: LeafStageCall<'a, T>, -} - -impl MetricVisitor for ExecuteLeafStage<'_, A, T> -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - T: VectorRepr + 'static, -{ - type Output = Result>, LeafBuildError>; - - fn visit(self) -> Self::Output { - build_leaf_candidates_for::( - self.arch, - self.call.data, - self.call.leaves, - self.call.requested_k, - ) - } -} - +/// +/// The caller has already selected `A` and `M` for the complete graph build. #[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. -fn build_leaf_candidates_for( +pub(super) fn build_leaf_candidates( arch: A, data: MatrixView<'_, T>, leaves: Vec>, @@ -423,7 +349,7 @@ where leaf, buffer: "leaf output", })?; - nearest_neighbors_for::(arch, dots, output, &mut buffers.kernel_workspace) + nearest_neighbors::(arch, dots, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; buffers.prepare_local_adjacency(point_ids.len())?; @@ -507,6 +433,10 @@ fn poisoned_candidate_list(point: u32) -> LeafBuildError { mod tests { use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; + use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, + }; use half::f16; use std::collections::BTreeSet; @@ -526,6 +456,57 @@ mod tests { .unwrap() } + struct LeafBuildCall<'a, T> { + data: MatrixView<'a, T>, + leaves: Vec>, + k: usize, + } + + struct DispatchLeafBuild(Metric); + + impl + Target1< + A, + Result>, LeafBuildError>, + LeafBuildCall<'_, T>, + > for DispatchLeafBuild + where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: crate::utils::VectorRepr + 'static, + { + fn run( + self, + arch: A, + call: LeafBuildCall<'_, T>, + ) -> Result>, LeafBuildError> { + use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + + match self.0 { + Metric::L2 => { + build_leaf_candidates::(arch, call.data, call.leaves, call.k) + } + Metric::Cosine => { + build_leaf_candidates::(arch, call.data, call.leaves, call.k) + } + Metric::CosineNormalized => build_leaf_candidates::( + arch, + call.data, + call.leaves, + call.k, + ), + Metric::InnerProduct => build_leaf_candidates::( + arch, + call.data, + call.leaves, + call.k, + ), + } + } + } + fn build( data: MatrixView<'_, T>, leaves: &[Vec], @@ -535,7 +516,16 @@ mod tests { where T: crate::utils::VectorRepr + 'static, { - pool().install(|| build_leaf_candidates(data, leaves.to_vec(), k, metric)) + pool().install(|| { + arch::dispatch1_no_features( + DispatchLeafBuild(metric), + LeafBuildCall { + data, + leaves: leaves.to_vec(), + k, + }, + ) + }) } fn adjacency_lists(graph: Vec>) -> Vec> { @@ -718,17 +708,11 @@ mod tests { let leaves: Vec> = (0..32) .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) .collect(); - let pool = pool(); - pool.install(|| { - let expected = - build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2).unwrap(); - for _ in 0..8 { - let actual = - build_leaf_candidates(view(&data, 64, 1), leaves.clone(), 2, Metric::L2) - .unwrap(); - assert_eq!(actual, expected); - } - }); + let expected = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + for _ in 0..8 { + let actual = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + assert_eq!(actual, expected); + } } #[test] diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 5788767659..e1e9929692 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,134 +5,45 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! PiPNN means **Pick-in-Partitions Nearest Neighbors**. It builds a graph for -//! approximate nearest-neighbor search: every input vector becomes one graph -//! vertex, and its adjacency list stores other vectors worth visiting during a -//! later query. This module constructs that adjacency; it does not execute queries. +//! PiPNN builds graph candidates in bulk instead of searching a partially built +//! graph for every insertion: //! -//! Incremental builders such as Vamana find construction candidates by running -//! beam search against a partially built graph: they repeatedly follow graph -//! edges to discover nearby vertices, causing random memory access. PiPNN removes -//! that search from construction and uses three bulk stages instead: -//! -//! 1. **Partition.** Randomized Ball Carving samples points called *leaders*. -//! Every point is assigned to its nearest `fanout` leaders. Assigning to more -//! than one leader makes child groups overlap. Oversized groups are processed -//! recursively until bounded groups called *leaves* remain. -//! 2. **Pick within leaves.** Vectors in one leaf are contiguous enough for one -//! dense general matrix multiplication (GEMM) to compute all pair dot -//! products. Each point picks its nearest leaf companions; selected pairs -//! become candidate graph edges. -//! 3. **Merge and prune.** Candidate edges from overlapping leaves are combined -//! into one unique list per source. Vamana RobustPrune then selects a bounded, -//! directionally diverse adjacency list. +//! 1. `partitioning` recursively samples leaders and assigns each point to its +//! nearest configured fanout, producing overlapping leaves bounded by `c_max`. +//! 2. `leaf_build` gathers each leaf, computes the lower triangle of `A · Aᵀ`, +//! selects up to three local neighbors per point, and merges symmetric global +//! candidate IDs across overlapping leaves. +//! 3. `finalization` applies the shared Vamana RobustPrune policy only to +//! candidate lists that exceed the configured graph degree. //! //! ```text -//! dataset points -//! │ -//! v -//! sample leaders + point/leader GEMM -//! │ -//! v -//! choose nearest leaders ──> overlapping child groups ──> recurse ──> leaves -//! │ -//! leaf all-pairs GEMM -//! │ -//! v -//! pick local neighbors -//! │ -//! v -//! merge/prune edges -//! │ -//! v -//! search graph +//! build_graph +//! └─ caller Rayon pool +//! └─ architecture dispatch + metric match once per build +//! └─ build_graph_for +//! ├─ partitioning::partition +//! │ └─ partition_kernel::nearest_leaders every stripe +//! ├─ leaf_build::build_leaf_candidates +//! │ └─ leaf_kernel::nearest_neighbors every leaf +//! └─ finalization::prune_overfull //! ``` //! -//! This module keeps GEMM separate from score selection: callers compute dense -//! dot-product matrices, then the kernels documented below convert those dots to -//! metric scores and retain top candidates. A *point* is a vector being assigned -//! during partitioning; a *leader* names a child group. In leaf selection, -//! *source* names the point whose output list is being built and *target* names -//! another point in that same leaf. -//! -//! The crate owns overlapping partition generation, leaf-local nearest-neighbor -//! construction, candidate merging, and graph-degree finalization. The caller -//! supplies a contiguous dataset view, DiskANN graph policy, and the Rayon pool. -//! Providers, start/frozen points, quantization, persistence, and search remain -//! outside this algorithm boundary. -//! -//! Numerical kernels include: -//! -//! - [`partition_kernel::PartitionKernel`] converts point-by-leader dot-product -//! tiles into nearest leader positions. -//! - [`leaf_kernel::LeafKernel`] scans each leaf's lower-triangular dot-product -//! matrix once and retains nearest non-self neighbors for both endpoints. -//! -//! Both handles are prepared once per build metric and reused across stripes or -//! leaves. Each output view supplies call-specific fanout or neighbor width. -//! Preparation selects the runtime architecture and returns a direct function -//! pointer; repeated calls do not repeat ISA or metric dispatch. -//! -//! # Main modules and structures -//! -//! ## Public build API -//! -//! - [`PiPNNConfig`] holds Randomized Ball Carving, fanout, leaf size, local `k`, -//! and replication parameters. -//! - [`PiPNNBuildContext`] validates that algorithm parameters, graph pruning -//! policy, metric, and caller-owned Rayon pool agree. -//! - [`build_graph`] runs the full pipeline over a borrowed row-major dataset and -//! returns one dataset-ID adjacency list per input point. -//! -//! ## [`partition_kernel`] -//! -//! Partition callers compute point-by-leader dots with GEMM. -//! [`partition_kernel::PartitionInput`] bundles that tile with typed -//! [`partition_kernel::PartitionScales`]. A prepared -//! [`partition_kernel::PartitionKernel`] writes sorted leader-local positions to -//! caller-owned output. Fanout is the output column count and cannot exceed the -//! leaders available for that partition call. +//! `diskann-wide` selects concrete architecture `A`; one four-way match selects +//! concrete metric marker `M`. Both types are then carried through every replica, +//! recursive partition, Rayon job, stripe, and leaf. Numerical kernels therefore +//! contain no runtime metric match, visitor, trait object, stored function pointer, +//! or repeated ISA dispatch. //! -//! ## [`leaf_kernel`] +//! [`PiPNNConfig`] owns only partition and local-neighbor parameters. +//! [`PiPNNBuildContext`] borrows DiskANN graph policy and the caller-owned Rayon +//! pool. [`build_graph`] borrows a contiguous [`MatrixView`] and returns one +//! dataset-ID adjacency list per real point. Providers, start/frozen points, +//! quantization, persistence, and search remain outside this module. //! -//! Leaf callers compute a lower-triangular point-by-point dot matrix with -//! `sgemm_aat_lower`. [`leaf_kernel::LeafKernelWorkspace`] owns reusable -//! per-worker scratch, and -//! [`leaf_kernel::LeafKernel`] writes sorted [`leaf_kernel::LeafNeighbor`] values -//! to caller-owned output. [`leaf_kernel::leaf_neighbor_count`] derives each -//! leaf's width from point count and requested `k`. Module documentation explains -//! fixed-width selection, `process_pairs`, and stable endpoint insertion. -//! -//! ## Private pipeline stages -//! -//! - `partitioning` recursively samples leaders, invokes the partition kernel, -//! scatters points into overlapping children, and returns bounded leaves. -//! - `leaf_build` gathers each leaf, computes its Gram matrix, invokes the leaf -//! kernel, translates local positions to dataset IDs, and merges candidates. -//! - `finalization` applies shared Vamana RobustPrune to overfull candidate lists. -//! - `kernel_metric` owns norm preparation, exact ranking inputs, and numerical -//! edge cases. The graph build selects one concrete metric for both kernels. -//! -//! # Typical use -//! -//! 1. Construct [`PiPNNConfig`] and DiskANN graph [`Config`]. -//! 2. Create [`PiPNNBuildContext`] with metric and caller-owned Rayon pool. -//! 3. Call [`build_graph`] with one row-major [`MatrixView`] of dataset vectors. -//! 4. The outer index builder chooses start/frozen points and serializes returned -//! adjacency; those policies are intentionally not part of this crate. -//! -//! Stage outputs are owned values. Leaves move into candidate construction; -//! candidate lists move into finalization. Ownership releases each stage's large -//! scratch before the next outer allocation. -//! -//! # Ownership and performance boundary -//! -//! Kernels borrow all matrices and mutate only caller-owned output/scratch. They -//! do not own providers, thread pools, GEMM buffers, graph IDs, or persistence. -//! Partition traversal performs one score per point-leader pair; leaf traversal -//! performs one score per unordered point pair. Detailed complexity and scratch -//! costs are documented in each module. PiPNN itself never names instruction -//! sets; `diskann-wide` owns architecture selection. +//! The partition and leaf stages own disjoint reusable scratch. Stage outputs move +//! forward (`leaves → candidates → adjacency`) so large temporary allocations can +//! drop at their consumption boundary. Kernel modules document validation, +//! numerical edge cases, tie order, scalar tails, and unchecked SIMD preconditions. mod kernel_metric; mod simd; @@ -150,8 +61,14 @@ use crate::{ }; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; +use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, +}; use rayon::ThreadPool; +use self::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; + /// Configuration of PiPNN's partitioning and local-neighbor algorithm. /// /// Graph degree, pruning policy, and alpha belong to DiskANN's graph @@ -287,19 +204,76 @@ where // Integer source vectors are not guaranteed unit-normalized after conversion, // so their normalized-cosine request must use the norm-aware formula. let metric = effective_metric::(context.metric); + arch::dispatch1_no_features( + RunBuildGraph, + BuildGraphCall { + data, + context, + metric, + }, + ) +} + +struct BuildGraphCall<'data, 'context, 'policy, T> { + data: MatrixView<'data, T>, + context: &'context PiPNNBuildContext<'policy>, + metric: Metric, +} + +struct RunBuildGraph; + +impl Target1>>, BuildGraphCall<'_, '_, '_, T>> + for RunBuildGraph +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + Send + Sync + 'static, +{ + fn run( + self, + arch: A, + call: BuildGraphCall<'_, '_, '_, T>, + ) -> ANNResult>> { + match call.metric { + Metric::L2 => build_graph_for::(arch, call.data, call.context), + Metric::Cosine => build_graph_for::(arch, call.data, call.context), + Metric::CosineNormalized => { + build_graph_for::(arch, call.data, call.context) + } + Metric::InnerProduct => { + build_graph_for::(arch, call.data, call.context) + } + } + } +} +fn build_graph_for( + arch: A, + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, + T: VectorRepr + Send + Sync + 'static, +{ let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition(data, &context.config, metric))?; + .in_scope(|| partitioning::partition::(arch, data, &context.config))?; // `leaves` is consumed here. Workers borrow individual ID lists during the // parallel pass, and the complete partition allocation drops on return. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) + leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.k) .map_err(ANNError::new) })?; // Finalization consumes candidate lists and reuses their allocations for the // resulting adjacency where possible. tracing::info_span!("pipnn.finalization") - .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, M::METRIC)) } fn effective_metric(metric: Metric) -> Metric { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 4bc7cd164e..4c9fa56842 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -46,18 +46,15 @@ use diskann_utils::{ views::{MatrixView, MutMatrixView}, }; use diskann_vector::{Norm, distance::Metric, norm::FastL2NormSquared}; -use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, -}; +use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::{KernelMetric, MetricVisitor, visit_metric}, + kernel_metric::KernelMetric, partition_kernel::{ - PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders_for, + PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders, }, }; @@ -146,67 +143,9 @@ type StripeBufferPool = ObjectPool; /// leaders for the current level, and recurses only on oversized clusters. /// Levels beyond `fanout.len()` retain one leader assignment. Completed small /// leaves are merged without exceeding `c_max`; every input point must remain -/// covered once per replica. The caller installs the operation in its pool. -pub(crate) fn partition( - data: MatrixView<'_, T>, - config: &PiPNNConfig, - metric: Metric, -) -> ANNResult>> -where - T: VectorRepr + Send + Sync, -{ - arch::dispatch1_no_features( - RunPartitionStage, - PartitionStageCall { - data, - config, - metric, - }, - ) -} - -struct PartitionStageCall<'a, T> { - data: MatrixView<'a, T>, - config: &'a PiPNNConfig, - metric: Metric, -} - -struct RunPartitionStage; - -impl Target1>>, PartitionStageCall<'_, T>> for RunPartitionStage -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - T: VectorRepr + Send + Sync, -{ - fn run(self, arch: A, call: PartitionStageCall<'_, T>) -> ANNResult>> { - visit_metric(call.metric, ExecutePartitionStage { arch, call }) - } -} - -struct ExecutePartitionStage<'a, A, T> { - arch: A, - call: PartitionStageCall<'a, T>, -} - -impl MetricVisitor for ExecutePartitionStage<'_, A, T> -where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - T: VectorRepr + Send + Sync, -{ - type Output = ANNResult>>; - - fn visit(self) -> Self::Output { - partition_for::(self.arch, self.call.data, self.call.config) - } -} - -fn partition_for( +/// covered once per replica. The caller has already selected `A` and `M` and +/// installed the operation in its pool. +pub(super) fn partition( arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, @@ -589,7 +528,7 @@ where actual: error.into_inner().len(), }) })?; - nearest_leaders_for::( + nearest_leaders::( arch, PartitionInput { dots, scales }, output, @@ -872,9 +811,55 @@ fn assignment_stripe_point_count(leader_count: usize) -> usize { mod tests { use diskann_utils::views::{Matrix, MatrixView}; use diskann_vector::{Half, distance::Metric}; + use diskann_wide::{ + Architecture, SIMDMask, SIMDSelect, SIMDVector, + arch::{self, Target1}, + }; use super::*; + struct PartitionCall<'a, T> { + data: MatrixView<'a, T>, + config: &'a PiPNNConfig, + } + + struct DispatchPartition(Metric); + + impl Target1>>, PartitionCall<'_, T>> for DispatchPartition + where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + T: VectorRepr + Send + Sync, + { + fn run(self, arch: A, call: PartitionCall<'_, T>) -> ANNResult>> { + use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; + + match self.0 { + Metric::L2 => partition::(arch, call.data, call.config), + Metric::Cosine => partition::(arch, call.data, call.config), + Metric::CosineNormalized => { + partition::(arch, call.data, call.config) + } + Metric::InnerProduct => { + partition::(arch, call.data, call.config) + } + } + } + } + + fn partition_with_runtime_metric( + data: MatrixView<'_, T>, + config: &PiPNNConfig, + metric: Metric, + ) -> ANNResult>> + where + T: VectorRepr + Send + Sync, + { + arch::dispatch1_no_features(DispatchPartition(metric), PartitionCall { data, config }) + } + fn config(c_min: usize, c_max: usize, fanout: Vec, replicas: usize) -> PiPNNConfig { PiPNNConfig { c_max, @@ -936,7 +921,12 @@ mod tests { memberships } - fn assert_valid_partition(leaves: &[Vec], points: usize, c_max: usize, replicas: usize) { + fn assert_valid_partition_with_runtime_metric( + leaves: &[Vec], + points: usize, + c_max: usize, + replicas: usize, + ) { assert!( leaves .iter() @@ -960,7 +950,12 @@ mod tests { fn returns_one_leaf_at_and_below_c_max() { for points in [7, 8] { let data = clustered_data(points, 3); - let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = partition_with_runtime_metric( + data.as_view(), + &config(2, 8, vec![2], 1), + Metric::L2, + ) + .unwrap(); assert_eq!(leaves, vec![(0..points as u32).collect::>()]); } } @@ -970,26 +965,30 @@ mod tests { let data = clustered_data(96, 8); let config = config(4, 16, vec![3, 2], 2); - let first = partition(data.as_view(), &config, Metric::L2).unwrap(); - let second = partition(data.as_view(), &config, Metric::L2).unwrap(); + let first = partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); + let second = partition_with_runtime_metric(data.as_view(), &config, Metric::L2).unwrap(); assert_eq!(sorted_memberships(&first), sorted_memberships(&second)); - assert_valid_partition(&first, 96, 16, 2); + assert_valid_partition_with_runtime_metric(&first, 96, 16, 2); assert!(first.iter().map(Vec::len).sum::() > 96 * 2); } #[test] fn partition_remains_bounded_after_the_fanout_schedule_is_exhausted() { let data = clustered_data(80, 4); - let leaves = partition(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2).unwrap(); + let leaves = + partition_with_runtime_metric(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2) + .unwrap(); - assert_valid_partition(&leaves, 80, 8, 1); + assert_valid_partition_with_runtime_metric(&leaves, 80, 8, 1); } #[test] fn duplicate_points_return_iteration_limit_instead_of_oversized_leaf() { let data = Matrix::new(1.0f32, 24, 4); - let error = partition(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = + partition_with_runtime_metric(data.as_view(), &config(2, 4, vec![1], 1), Metric::L2) + .unwrap_err(); let error = error.downcast::().unwrap(); assert!(matches!( @@ -1033,14 +1032,14 @@ mod tests { #[test] fn replicas_cover_every_point_once_or_more_per_replica() { let data = directional_data(72, 5); - let leaves = partition( + let leaves = partition_with_runtime_metric( data.as_view(), &config(3, 12, vec![3, 2], 3), Metric::CosineNormalized, ) .unwrap(); - assert_valid_partition(&leaves, 72, 12, 3); + assert_valid_partition_with_runtime_metric(&leaves, 72, 12, 3); } fn assert_partition_conversion_matches_f32(label: &str, convert: impl Fn(u8) -> T) @@ -1061,20 +1060,20 @@ mod tests { let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); let converted: Vec = raw.iter().copied().map(&convert).collect(); let config = config(2, 16, vec![2, 1], 1); - let expected = partition( + let expected = partition_with_runtime_metric( MatrixView::try_from(&f32_data, points, dimensions).unwrap(), &config, Metric::L2, ) .unwrap(); - let actual = partition( + let actual = partition_with_runtime_metric( MatrixView::try_from(&converted, points, dimensions).unwrap(), &config, Metric::L2, ) .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); - assert_valid_partition(&actual, points, 16, 1); + assert_valid_partition_with_runtime_metric(&actual, points, 16, 1); assert_eq!( sorted_memberships(&actual), sorted_memberships(&expected), @@ -1148,8 +1147,8 @@ mod tests { Metric::CosineNormalized, Metric::InnerProduct, ] { - let leaves = partition(data.as_view(), &config, metric).unwrap(); - assert_valid_partition(&leaves, 64, 20, 1); + let leaves = partition_with_runtime_metric(data.as_view(), &config, metric).unwrap(); + assert_valid_partition_with_runtime_metric(&leaves, 64, 20, 1); } } @@ -1228,7 +1227,9 @@ mod tests { #[test] fn rejects_empty_dataset() { let data = Matrix::::new(0.0, 0, 4); - let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = + partition_with_runtime_metric(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2) + .unwrap_err(); assert_eq!( error.downcast::().unwrap(), @@ -1239,7 +1240,9 @@ mod tests { #[test] fn rejects_zero_dimensions() { let data = Matrix::::new(0.0, 4, 0); - let error = partition(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2).unwrap_err(); + let error = + partition_with_runtime_metric(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2) + .unwrap_err(); assert_eq!( error.downcast::().unwrap(), From b92226996f7ff5d6a3cb32d7b005d899ff76b39c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:58:45 +0000 Subject: [PATCH 33/58] docs(pipnn): describe the active core flow --- diskann/src/graph/pipnn/finalization.rs | 40 +++--- diskann/src/graph/pipnn/leaf_build.rs | 51 ++++---- diskann/src/graph/pipnn/mod.rs | 82 ++++++------ diskann/src/graph/pipnn/partitioning.rs | 158 ++++++++++++------------ 4 files changed, 162 insertions(+), 169 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 7d0ccb722d..cf1f31d234 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -3,17 +3,18 @@ * Licensed under the MIT license. */ -//! Final graph-degree enforcement through the shared Vamana RobustPrune kernel. +//! Graph-degree enforcement with the Vamana RobustPrune kernel. //! -//! Candidate merging may produce more than `R` IDs for a point. This stage first -//! validates every global ID, then processes each point's candidate list in the -//! caller's Rayon pool. Lists already within the degree bound are returned without -//! distance work. Overfull lists are converted to source-distance candidates, -//! passed through RobustPrune, and rewritten from the selected output. +//! Candidate merging can produce more than `R` IDs for one point. This module +//! checks every global ID before parallel work starts. It returns a list with at +//! most `R` IDs without distance work. //! -//! The shared kernel owns occlusion and alpha-round semantics; this adapter owns -//! only contiguous dataset access and distance specialization for the source -//! representation. +//! For a longer list, the module computes each source distance. It sorts the +//! candidates and calls RobustPrune. The module then writes the selected IDs into +//! the original list allocation. +//! +//! RobustPrune defines occlusion and alpha-round behavior. This module supplies +//! contiguous vector access and a distance function for input type `T`. //! //! ```text //! candidate lists ──> validate point count and every global ID @@ -64,10 +65,11 @@ pub(crate) enum FinalizationError { TooManyCandidates { actual: usize, max: usize }, } -/// Per-Rayon-job preparation and kernel state retained across source points. +/// Reusable buffers for one Rayon job. /// -/// PiPNN owns sorting, allocation, and ID translation. The shared internal -/// kernel receives only the prepared candidates and an exactly sized state slice. +/// `pool` stores candidates with source distances. `prepared` stores each +/// distance and optional non-self ID. `states` stores one RobustPrune state for +/// each sorted candidate. #[derive(Default)] struct Workspace { pool: Vec>, @@ -75,7 +77,7 @@ struct Workspace { states: Vec, } -/// Validate candidate IDs and prune only lists whose length exceeds graph degree. +/// Check candidate IDs and prune each list that exceeds the graph degree. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, @@ -90,7 +92,7 @@ where let degree = graph.pruned_degree().get(); let distance = T::distance(metric, Some(data.ncols())); - // build_graph installs the complete call tree in the caller-owned pool. + // `build_graph` runs this Rayon operation in the pool from the build context. #[allow(clippy::disallowed_methods)] candidates .into_par_iter() @@ -98,8 +100,8 @@ where .map_init( Workspace::default, |workspace, (source, mut source_candidates)| { - // Candidate accumulators already enforce uniqueness. A bounded list - // therefore satisfies the graph policy without distance evaluation. + // Candidate merging already removes duplicate IDs. A list within + // the degree limit needs no distance calculation. if source_candidates.len() <= degree { return Ok(source_candidates); } @@ -134,9 +136,9 @@ where .try_reserve(candidate_count) .map_err(ANNError::new)?; - // Sorting/capping precedes source exclusion so filtering cannot - // backfill with farther candidates. Passing this witness into - // RobustPrune makes source-distance order part of its input type. + // Sort all candidates before the code marks a self-edge as absent. + // Thus, self-edge removal cannot add a farther candidate. The + // `SortedNeighbors` value carries this order into RobustPrune. let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); workspace.prepared.extend(sorted.iter().map(|neighbor| { let id = *neighbor.id(); diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 3ff70284d9..c2143ddb6e 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -5,18 +5,18 @@ //! Leaf-local graph construction and candidate accumulation. //! -//! Partitioning supplies strictly increasing, unique global point IDs per leaf. -//! For each leaf this module: +//! Partitioning supplies sorted, unique global point IDs for each leaf. One leaf +//! job does these steps: //! -//! 1. validates IDs and converts only those point vectors to reusable `f32` scratch; -//! 2. computes the lower triangle of `A · Aᵀ`; -//! 3. runs the dual-endpoint leaf top-k kernel; and -//! 4. translates leaf-local positions back to dataset IDs. +//! 1. Check each ID and convert its vector to reusable `f32` storage. +//! 2. Compute the lower triangle of `A · Aᵀ`. +//! 3. Select local neighbors for both points of each pair. +//! 4. Convert local positions to global point IDs. +//! 5. Add both edge directions to global candidate lists. //! -//! The final step merges symmetric adjacency lists under per-point locks because -//! overlapping leaves are processed concurrently. Numeric buffers retain their -//! high-water length; every consumer therefore receives an explicit active -//! prefix rather than treating `Vec::len()` as the current leaf shape. +//! Overlapping leaves run concurrently. A worker locks one destination list only +//! while it adds one leaf's IDs. Reusable buffers keep their largest allocation. +//! Each operation uses an explicit active prefix. use std::{collections::TryReserveError, sync::Mutex}; @@ -92,12 +92,10 @@ pub(crate) enum LeafBuildError { PoisonedCandidateList { point: u32 }, } -/// Scratch leased to one Rayon job and reused for successive leaves. +/// Reusable buffers for one Rayon leaf job. /// -/// The three numerical vectors retain their largest observed leaf shape. The -/// adjacency lists are prepared separately because zero-k/singleton leaves never -/// write them, and because later candidate-merging modes do not necessarily use -/// this representation. +/// The numerical vectors keep the largest leaf shape that this job observed. +/// The job creates local adjacency lists only when the effective `k` is not zero. #[derive(Default)] struct LeafBuffers { point_values: Vec, @@ -166,12 +164,11 @@ impl LeafBuffers { } } -/// Concurrent accumulator indexed by global dataset ID. +/// Concurrent candidate lists indexed by global point ID. /// -/// A point may appear in several overlapping leaves, so workers lock only the -/// destination list long enough to append one leaf's additions. Sorting and -/// duplicate removal are deferred until all leaves finish; doing either under -/// the lock would lengthen the contended section for no semantic benefit. +/// A point can occur in several overlapping leaves. A worker locks one point's +/// list and adds all IDs from one leaf. `AdjacencyList` removes duplicates during +/// this append. The function sorts each list after all leaf jobs finish. struct DirectCandidates { lists: Vec>>, } @@ -192,7 +189,7 @@ impl DirectCandidates { local_adjacency: &[AdjacencyList], ) -> Result<(), LeafBuildError> { for (&source, additions) in point_ids.iter().zip(local_adjacency) { - // Every point ID is validated before leaf-local work begins. + // `build_leaf` checks every point ID before this append. let candidates = &self.lists[source as usize]; let mut candidates = candidates .lock() @@ -218,9 +215,9 @@ impl DirectCandidates { } } -/// Build symmetric leaf-local k-NN graphs and retain every unique global candidate. +/// Build symmetric leaf-local k-NN graphs and return unique global candidates. /// -/// The caller has already selected `A` and `M` for the complete graph build. +/// The caller supplies concrete architecture `A` and metric `M`. #[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. pub(super) fn build_leaf_candidates( arch: A, @@ -261,11 +258,11 @@ where candidates.into_lists() } -/// Build and publish one leaf's symmetric neighbor lists. +/// Build and publish the symmetric neighbor lists for one leaf. /// -/// Validation precedes all dataset indexing and rejects IDs that are not strictly -/// increasing. Active lengths computed after `prepare` must be used for every -/// later slice because reusable vectors may remain longer than this leaf. +/// The function checks all IDs before it indexes the dataset. IDs must be +/// strictly increasing. Reusable vectors can be longer than this leaf. Every +/// read and write uses the active length from `prepare`. fn build_leaf( arch: A, data: MatrixView<'_, T>, diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index e1e9929692..865c29e4ce 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,45 +5,41 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! PiPNN builds graph candidates in bulk instead of searching a partially built -//! graph for every insertion: +//! PiPNN builds graph candidates in three steps: //! -//! 1. `partitioning` recursively samples leaders and assigns each point to its -//! nearest configured fanout, producing overlapping leaves bounded by `c_max`. -//! 2. `leaf_build` gathers each leaf, computes the lower triangle of `A · Aᵀ`, -//! selects up to three local neighbors per point, and merges symmetric global -//! candidate IDs across overlapping leaves. -//! 3. `finalization` applies the shared Vamana RobustPrune policy only to -//! candidate lists that exceed the configured graph degree. +//! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has +//! at most `c_max` points. +//! 2. `leaf_build` computes a lower-triangular Gram matrix for each leaf. It +//! selects local neighbors and merges their global point IDs. +//! 3. `finalization` applies Vamana RobustPrune to each candidate list that is +//! longer than the graph degree. //! //! ```text //! build_graph //! └─ caller Rayon pool -//! └─ architecture dispatch + metric match once per build +//! └─ select architecture A and metric M once //! └─ build_graph_for //! ├─ partitioning::partition -//! │ └─ partition_kernel::nearest_leaders every stripe +//! │ └─ partition_kernel::nearest_leaders per stripe //! ├─ leaf_build::build_leaf_candidates -//! │ └─ leaf_kernel::nearest_neighbors every leaf +//! │ └─ leaf_kernel::nearest_neighbors per leaf //! └─ finalization::prune_overfull //! ``` //! -//! `diskann-wide` selects concrete architecture `A`; one four-way match selects -//! concrete metric marker `M`. Both types are then carried through every replica, -//! recursive partition, Rayon job, stripe, and leaf. Numerical kernels therefore -//! contain no runtime metric match, visitor, trait object, stored function pointer, -//! or repeated ISA dispatch. +//! `diskann-wide` selects architecture `A`. One match selects metric marker `M`. +//! The build passes both concrete types through all replicas, recursive +//! partitions, stripes, and leaves. The numerical loops do not dispatch again. //! -//! [`PiPNNConfig`] owns only partition and local-neighbor parameters. -//! [`PiPNNBuildContext`] borrows DiskANN graph policy and the caller-owned Rayon -//! pool. [`build_graph`] borrows a contiguous [`MatrixView`] and returns one -//! dataset-ID adjacency list per real point. Providers, start/frozen points, -//! quantization, persistence, and search remain outside this module. +//! [`PiPNNConfig`] contains partition and local-neighbor parameters. +//! [`PiPNNBuildContext`] borrows graph policy and a Rayon pool. [`build_graph`] +//! borrows one contiguous [`MatrixView`]. It returns one adjacency list for each +//! input point. //! -//! The partition and leaf stages own disjoint reusable scratch. Stage outputs move -//! forward (`leaves → candidates → adjacency`) so large temporary allocations can -//! drop at their consumption boundary. Kernel modules document validation, -//! numerical edge cases, tie order, scalar tails, and unchecked SIMD preconditions. +//! The function does not load providers or select start and frozen points. It +//! also does not quantize, serialize, or search the graph. +//! +//! Partition and leaf work use separate reusable buffers. The build consumes +//! each stage output before it creates the next graph representation. mod kernel_metric; mod simd; @@ -69,10 +65,10 @@ use rayon::ThreadPool; use self::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; -/// Configuration of PiPNN's partitioning and local-neighbor algorithm. +/// Configuration for PiPNN partitioning and local-neighbor selection. /// -/// Graph degree, pruning policy, and alpha belong to DiskANN's graph -/// configuration and are supplied separately through [`PiPNNBuildContext`]. +/// [`PiPNNBuildContext`] supplies graph degree, prune policy, alpha, metric, and +/// the Rayon pool. #[derive(Clone, Debug, PartialEq)] pub struct PiPNNConfig { /// Maximum number of points in a leaf. @@ -130,7 +126,7 @@ impl PiPNNConfig { } } -/// Validated, borrowed policy and execution context for one PiPNN graph build. +/// Checked policy and execution inputs for one PiPNN graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { pub(crate) config: PiPNNConfig, @@ -140,7 +136,7 @@ pub struct PiPNNBuildContext<'a> { } impl<'a> PiPNNBuildContext<'a> { - /// Validate and combine PiPNN configuration with outer graph policy. + /// Check and combine PiPNN configuration with DiskANN graph policy. pub fn new( config: PiPNNConfig, graph: &'a Config, @@ -164,12 +160,14 @@ impl<'a> PiPNNBuildContext<'a> { } } -/// Build PiPNN adjacency for real points in `data`. +/// Build PiPNN adjacency for all points in `data`. +/// +/// The function does not select start or frozen points. It does not load a +/// provider or write an index. /// -/// This is the core algorithm boundary. Search entry-point selection, frozen nodes, -/// providers, serialization, and index writers belong to the outer build pipelines. -/// For raw `u8` and `i8` vectors, `CosineNormalized` is evaluated as `Cosine` because -/// those representations are converted to f32 scratch but are not unit-normalized. +/// Raw `u8` and `i8` vectors are not unit-normalized after conversion to `f32`. +/// Therefore, the function evaluates `CosineNormalized` as `Cosine` for these +/// two input types. pub fn build_graph( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, @@ -201,8 +199,8 @@ where data.nrows() ))); } - // Integer source vectors are not guaranteed unit-normalized after conversion, - // so their normalized-cosine request must use the norm-aware formula. + // Conversion does not make integer vectors unit length. Use the norm-aware + // cosine formula for these vectors. let metric = effective_metric::(context.metric); arch::dispatch1_no_features( RunBuildGraph, @@ -264,14 +262,14 @@ where { let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition::(arch, data, &context.config))?; - // `leaves` is consumed here. Workers borrow individual ID lists during the - // parallel pass, and the complete partition allocation drops on return. + // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, + // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.k) .map_err(ANNError::new) })?; - // Finalization consumes candidate lists and reuses their allocations for the - // resulting adjacency where possible. + // Finalization consumes each candidate list. It reuses that list's allocation + // for the final adjacency when the graph policy permits it. tracing::info_span!("pipnn.finalization") .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, M::METRIC)) } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 4c9fa56842..c208a9685b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -5,37 +5,38 @@ //! Deterministic overlapping partition construction for PiPNN. //! -//! The stage maps real dataset points to bounded leaf ID lists. Numerical work -//! reuses the partition kernel and dense GEMM. A stage-owned pool leases scratch -//! to Rayon chunks and takes it back after each chunk; computation never holds -//! the pool lock, and no thread-local cleanup protocol is required. +//! This module converts dataset point IDs into bounded leaf ID lists. It uses +//! dense GEMM and the partition kernel for leader assignment. An `ObjectPool` +//! supplies one reusable buffer set to each Rayon worker chunk. The pool lock is +//! not held during gather, GEMM, or top-k selection. //! //! ```text -//! replica root IDs ──> work queue -//! │ -//! v -//! sample leaders -//! │ -//! gather stripes ─> GEMM distances ─> nearest leaders -//! │ -//! v -//! stable scatter by leader -//! │ │ -//! size <= c_max oversized cluster -//! │ │ -//! completed leaf next recursion level -//! └──────────┬────────┘ -//! v -//! global small-leaf merge -//! v -//! coverage/bound validation +//! replica point IDs ──> work queue +//! │ +//! v +//! sample leaders +//! │ +//! gather stripes ─> GEMM ─> nearest leaders +//! │ +//! v +//! stable scatter by leader +//! │ │ +//! size <= c_max oversized cluster +//! │ │ +//! completed leaf work queue +//! └────────┬────────┘ +//! v +//! merge undersized leaves +//! │ +//! v +//! check leaf bounds //! ``` //! -//! | Recursion level | Assignment multiplicity | +//! | Condition | Assignments per point | //! | --- | --- | //! | `level < fanout.len()` | `fanout[level]` nearest leaders | -//! | later levels | one nearest leader until bounded | -//! | replica boundary | independent deterministic seed | +//! | `level >= fanout.len()` | one nearest leader | +//! | new replica | independent deterministic seed | use std::collections::HashSet; @@ -58,7 +59,7 @@ use super::{ }, }; -// Private algorithm and batching constants live together. None are user policy. +// These constants control internal batching and deterministic seed generation. const PARTITION_SEED: u64 = 1_000; const REPLICA_SEED_STEP: u64 = 7_919; const LEADER_CAP: usize = 1_000; @@ -68,7 +69,7 @@ const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const MAX_PARTITION_ITERATIONS: usize = 30; -/// A partition failure with enough context to diagnose non-progressing input. +/// Error from partition input checks, allocation, or recursion progress. #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum PartitionError { #[error("PiPNN cannot partition an empty dataset")] @@ -124,27 +125,27 @@ impl AsPooled<()> for StripeBuffers { } fn modify(&mut self, _: ()) { - // Scratch retains its high-water allocation across leases; active - // prefixes are established by `assign_stripe` before every read. + // Keep the largest allocation across leases. `assign_stripe` defines the + // active prefix before each read. } } -/// Stage-owned high-water scratch storage for partition assignment. +/// Reusable partition buffers owned by this partition call. /// -/// `ObjectPool` owns the short pop/push lock and returns leases through RAII, -/// including error and panic paths. Gather, GEMM, and top-k run while only the -/// leased `StripeBuffers` is held. +/// `ObjectPool` locks only when it gives or receives a lease. RAII returns each +/// lease after success, error, or panic. Numerical work holds only the lease. type StripeBufferPool = ObjectPool; -/// Partition every configured replica into overlapping bounded leaves. +/// Partition all configured replicas into overlapping bounded leaves. /// -/// Each oversized work item samples `ceil(p_samp * points)` leaders (clamped -/// to the private leader bound), assigns every point to its nearest `fanout` -/// leaders for the current level, and recurses only on oversized clusters. -/// Levels beyond `fanout.len()` retain one leader assignment. Completed small -/// leaves are merged without exceeding `c_max`; every input point must remain -/// covered once per replica. The caller has already selected `A` and `M` and -/// installed the operation in its pool. +/// An oversized work item samples `ceil(p_samp * points)` leaders, up to +/// `LEADER_CAP`. It assigns each point to the configured number of nearest +/// leaders. It adds each cluster above `c_max` to the work queue. +/// +/// A level without a configured fanout assigns each point to one leader. The +/// final merge does not make a leaf larger than `c_max`. Each replica covers +/// every input point. The caller supplies concrete architecture `A` and metric +/// `M`. pub(super) fn partition( arch: A, data: MatrixView<'_, T>, @@ -226,8 +227,8 @@ where .try_reserve_exact(work.len()) .map_err(ANNError::new)?; results.resize_with(work.len(), || None); - // build_graph installs this complete private call tree into the - // caller-owned pool; the indexed fill cannot escape that pool. + // `build_graph` runs this Rayon operation in the pool from the build + // context. Each worker writes only to its indexed result slot. #[allow(clippy::disallowed_methods)] results .par_iter_mut() @@ -337,20 +338,20 @@ fn replica_seed(replica: usize) -> u64 { PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) } -// A single LCG mixer derives recursive seeds. Wrapping makes the mapping stable -// across debug/release builds and supported platforms. +// This LCG derives child seeds. Wrapping arithmetic gives the same mapping in +// debug and release builds on all supported platforms. fn mix_seed(seed: u64, salt: u64) -> u64 { seed.wrapping_mul(6_364_136_223_846_793_005) .wrapping_add(salt) } -/// Assign each point to its nearest `fanout` sampled leaders. +/// Assign each point to its nearest sampled leaders. +/// +/// The function gathers leader vectors once. It divides points into cache-sized +/// stripes. Each worker chunk reuses one leased buffer set for all its stripes. /// -/// Leader vectors are gathered once. Points are processed in cache-sized -/// stripes, while a worker chunk retains one leased scratch buffer across all of -/// its stripes. The flat assignment matrix preserves point order and is then -/// scattered into per-leader clusters; preserving order is required for fixed -/// seed determinism in later recursion levels. +/// The flat assignment matrix keeps point order. The scatter step keeps the same +/// order in each leader cluster. Recursive sampling depends on this order. fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, @@ -381,10 +382,8 @@ where .iter_mut() .zip(leader_values.chunks_exact(dimension_count)) { - // Leader norms participate in the top-k ordering. Preserve the original - // scalar reduction order: reassociating this short setup pass through a - // SIMD norm changes low bits and can send near-tied points down different - // recursive partition paths. + // Leader norms affect top-k order. Use this scalar reduction order. + // SIMD reassociation changes low bits and can change a near-tie branch. *scale = leader_vector.iter().map(|value| value * value).sum(); if M::METRIC == Metric::Cosine { *scale = scale.sqrt(); @@ -401,8 +400,8 @@ where let worker_point_count = checked_area("assignment worker", worker_stripe_count, stripe_points)?; let worker_assignment_count = checked_area("assignment worker", worker_point_count, fanout)?; - // Each worker chunk owns one scratch value and reuses it for its stripes. - // build_graph pins this terminal operation to the caller-owned pool. + // Each worker chunk reuses one buffer lease for all its stripes. + // `build_graph` runs this operation in the pool from the build context. #[allow(clippy::disallowed_methods)] assignments .par_chunks_mut(worker_assignment_count) @@ -459,13 +458,10 @@ where let point_values_len = checked_area("point stripe", point_count, dimensions)?; let dots_len = checked_area("dot-product stripe", point_count, leader_count)?; let output_len = checked_area("partition assignments", point_count, fanout)?; - // Scratch keeps its high-water length and every consumer receives an - // explicit active prefix. Resizing to the exact stripe shape would be - // correct but re-zeroes the buffer whenever a pooled value moves between - // work items with different leader counts: `stripe_points` is derived from - // `leader_count`, so the point buffer swings between roughly 768 KiB and 6 MiB - // and `Vec::resize` only truncates on the way down, then memsets the whole - // delta on the way back up. + // Keep each buffer at its largest length and use an explicit active prefix. + // Different leader counts change the point buffer from about 768 KiB to + // 6 MiB. Exact resizing truncates the buffer and then zeros the full growth + // when another work item needs the larger shape. grow_fallible(&mut buffers.points, point_values_len, 0.0)?; grow_fallible(&mut buffers.dots, dots_len, 0.0)?; let StripeBuffers { @@ -555,12 +551,11 @@ where Ok(()) } -/// Convert the flat point-major assignment matrix into leader-major clusters. +/// Convert point-major assignments into leader-major clusters. /// -/// Small inputs use one serial exact-capacity pass. Large inputs form at most -/// one partial cluster set per Rayon worker, then merge each leader independently. -/// Concatenating partials in stripe order keeps the same member order as the -/// serial implementation while removing a large serial copy tail. +/// A small input uses one serial pass with exact capacities. A large input makes +/// at most one partial cluster set for each Rayon worker. It then merges each +/// leader independently. Stripe-order concatenation matches serial member order. fn scatter_assignments( points: &[u32], assignments: &[u32], @@ -577,7 +572,8 @@ fn scatter_assignments( let mut partials = Vec::new(); partials.try_reserve_exact(stripes).map_err(ANNError::new)?; partials.resize_with(stripes, || None); - // See the pool invariant at the other partition terminal operations. + // `build_graph` runs this Rayon operation in the pool from the build context. + // Each worker writes only to its indexed partial result. #[allow(clippy::disallowed_methods)] partials .par_iter_mut() @@ -609,7 +605,8 @@ fn scatter_assignments( } } - // See the pool invariant at the other partition terminal operations. + // `build_graph` runs this Rayon operation in the pool from the build context. + // Each worker creates one independent leader cluster. #[allow(clippy::disallowed_methods)] sizes .into_par_iter() @@ -776,11 +773,11 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { Ok(values) } -/// Grow `values` to at least `len` elements, never shrinking it. +/// Grow `values` to at least `len` elements and do not shrink it. /// -/// Callers slice the active prefix themselves. Shrinking would force the next -/// larger stripe to re-zero the reclaimed tail, which is the dominant cost when -/// one pooled buffer serves work items with different stripe shapes. +/// Callers use an explicit active prefix. Shrinking and regrowing the buffer +/// zeros the reclaimed tail. This zeroing dominates buffer reuse across +/// different stripe shapes. fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { if values.len() >= len { return Ok(()); @@ -1047,8 +1044,8 @@ mod tests { T: crate::utils::VectorRepr + Send + Sync, { let points = 64; - // Partition gathering converts source vectors before GEMM. Exercise conversion - // tails around 4-, 8-, and 16-element boundaries and a second 16-lane chunk. + // Partition gather converts source vectors before GEMM. Test conversion + // tails around 4, 8, 16, and 32 elements. for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { let raw: Vec = (0..points * dimensions) .map(|index| { @@ -1107,10 +1104,9 @@ mod tests { (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 } - // This fixed case sits on opposite sides of the top-1 boundary depending - // on whether leader norms use the original scalar reduction or a SIMD - // reassociation. Point/leader dot products still go through the production - // GEMM; only the setup norm calculation is under test. + // Scalar and SIMD-reassociated leader norms select different top-1 + // leaders for this case. Dot products still use the production GEMM. The + // test changes only the leader-norm reduction. let dimensions = 129; let mut state = 0x3a85_f952_c718_6e49; let point: Vec = (0..dimensions).map(|_| next(&mut state)).collect(); From 2f3443018fc7f0e4c33aac51b1983ce2a7904814 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 12:03:46 +0000 Subject: [PATCH 34/58] docs(pipnn): remove core diagrams and tuning notes --- diskann/src/graph/pipnn/finalization.rs | 20 ------------ diskann/src/graph/pipnn/mod.rs | 15 ++------- diskann/src/graph/pipnn/partitioning.rs | 42 +++++-------------------- 3 files changed, 9 insertions(+), 68 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index cf1f31d234..c44c35bebe 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -15,26 +15,6 @@ //! //! RobustPrune defines occlusion and alpha-round behavior. This module supplies //! contiguous vector access and a distance function for input type `T`. -//! -//! ```text -//! candidate lists ──> validate point count and every global ID -//! │ -//! ┌────────────┴────────────┐ -//! v v -//! len <= R len > R -//! return list source-distance candidates -//! │ -//! v -//! shared RobustPrune -//! │ -//! v -//! rewrite same list owner -//! ``` -//! -//! | Path | Distance evaluations | Allocation behavior | -//! | --- | --- | --- | -//! | bounded list | none | move list directly to output | -//! | overfull list | source and occlusion distances | reuse Rayon-job workspace | use crate::{ ANNError, ANNResult, diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 865c29e4ce..6aa63f5e58 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -13,18 +13,7 @@ //! selects local neighbors and merges their global point IDs. //! 3. `finalization` applies Vamana RobustPrune to each candidate list that is //! longer than the graph degree. -//! -//! ```text -//! build_graph -//! └─ caller Rayon pool -//! └─ select architecture A and metric M once -//! └─ build_graph_for -//! ├─ partitioning::partition -//! │ └─ partition_kernel::nearest_leaders per stripe -//! ├─ leaf_build::build_leaf_candidates -//! │ └─ leaf_kernel::nearest_neighbors per leaf -//! └─ finalization::prune_overfull -//! ``` + //! //! `diskann-wide` selects architecture `A`. One match selects metric marker `M`. //! The build passes both concrete types through all replicas, recursive @@ -39,7 +28,7 @@ //! also does not quantize, serialize, or search the graph. //! //! Partition and leaf work use separate reusable buffers. The build consumes -//! each stage output before it creates the next graph representation. +//! each output before it creates another graph representation. mod kernel_metric; mod simd; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index c208a9685b..8fb2c5bdb5 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -10,33 +10,9 @@ //! supplies one reusable buffer set to each Rayon worker chunk. The pool lock is //! not held during gather, GEMM, or top-k selection. //! -//! ```text -//! replica point IDs ──> work queue -//! │ -//! v -//! sample leaders -//! │ -//! gather stripes ─> GEMM ─> nearest leaders -//! │ -//! v -//! stable scatter by leader -//! │ │ -//! size <= c_max oversized cluster -//! │ │ -//! completed leaf work queue -//! └────────┬────────┘ -//! v -//! merge undersized leaves -//! │ -//! v -//! check leaf bounds -//! ``` -//! -//! | Condition | Assignments per point | -//! | --- | --- | -//! | `level < fanout.len()` | `fanout[level]` nearest leaders | -//! | `level >= fanout.len()` | one nearest leader | -//! | new replica | independent deterministic seed | +//! A configured level assigns each point to `fanout[level]` leaders. A deeper +//! level assigns each point to one leader. Each replica uses a different +//! deterministic seed. use std::collections::HashSet; @@ -347,7 +323,7 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { /// Assign each point to its nearest sampled leaders. /// -/// The function gathers leader vectors once. It divides points into cache-sized +/// The function gathers leader vectors once. It divides points into bounded /// stripes. Each worker chunk reuses one leased buffer set for all its stripes. /// /// The flat assignment matrix keeps point order. The scatter step keeps the same @@ -458,10 +434,8 @@ where let point_values_len = checked_area("point stripe", point_count, dimensions)?; let dots_len = checked_area("dot-product stripe", point_count, leader_count)?; let output_len = checked_area("partition assignments", point_count, fanout)?; - // Keep each buffer at its largest length and use an explicit active prefix. - // Different leader counts change the point buffer from about 768 KiB to - // 6 MiB. Exact resizing truncates the buffer and then zeros the full growth - // when another work item needs the larger shape. + // Keep each buffer at its largest length. Every operation uses an explicit + // active prefix. grow_fallible(&mut buffers.points, point_values_len, 0.0)?; grow_fallible(&mut buffers.dots, dots_len, 0.0)?; let StripeBuffers { @@ -775,9 +749,7 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { /// Grow `values` to at least `len` elements and do not shrink it. /// -/// Callers use an explicit active prefix. Shrinking and regrowing the buffer -/// zeros the reclaimed tail. This zeroing dominates buffer reuse across -/// different stripe shapes. +/// Callers use an explicit active prefix. fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { if values.len() >= len { return Ok(()); From 41b2981a245729ff445b13e71fd015a2d8bcedbd Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:44:03 +0000 Subject: [PATCH 35/58] refactor(pipnn): remove duplicate partition checks --- diskann/src/graph/pipnn/leaf_build.rs | 25 ++++++++++++++++--------- diskann/src/graph/pipnn/partitioning.rs | 1 - 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index c2143ddb6e..062896d161 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -28,8 +28,7 @@ use rayon::prelude::*; use super::{ kernel_metric::KernelMetric, leaf_kernel::{ - LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, leaf_output_len, - nearest_neighbors, + LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, }; @@ -112,7 +111,7 @@ impl LeafBuffers { point_count: usize, dimension_count: usize, requested_k: usize, - ) -> Result { + ) -> Result<(usize, usize), LeafBuildError> { let point_value_count = point_count .checked_mul(dimension_count) @@ -131,8 +130,17 @@ impl LeafBuffers { })?; let leaf_k = leaf_neighbor_count(point_count, requested_k) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - let neighbor_count = leaf_output_len(point_count, requested_k) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + let neighbor_count = + point_count + .checked_mul(leaf_k) + .ok_or_else(|| LeafBuildError::Kernel { + leaf, + source: LeafKernelError::ShapeOverflow { + buffer: "output", + rows: point_count, + cols: leaf_k, + }, + })?; grow( "leaf point values", @@ -147,7 +155,7 @@ impl LeafBuffers { neighbor_count, LeafNeighbor::default(), )?; - Ok(leaf_k) + Ok((leaf_k, neighbor_count)) } fn prepare_local_adjacency(&mut self, point_count: usize) -> Result<(), LeafBuildError> { @@ -301,15 +309,14 @@ where } return Err(LeafBuildError::UnsortedPointIds { leaf }); } - let leaf_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; + let (leaf_k, neighbor_value_count) = + buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; if leaf_k == 0 { return Ok(()); } let point_value_count = point_ids.len() * data.ncols(); let dot_count = point_ids.len() * point_ids.len(); - let neighbor_value_count = leaf_output_len(point_ids.len(), requested_k) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; for (&point, point_output) in point_ids .iter() diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 8fb2c5bdb5..83568138dc 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -157,7 +157,6 @@ where .map_err(ANNError::new)?; leaves.append(&mut replica_leaves); } - validate_leaves(&leaves, config.c_max)?; Ok(leaves) } From 1389f7798178fe3e8538a8d7dded57172a964d5a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:48:34 +0000 Subject: [PATCH 36/58] refactor(pipnn): keep leaf shape validation local --- diskann/src/graph/pipnn/leaf_build.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 062896d161..6e787c02bf 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -133,13 +133,10 @@ impl LeafBuffers { let neighbor_count = point_count .checked_mul(leaf_k) - .ok_or_else(|| LeafBuildError::Kernel { + .ok_or(LeafBuildError::ShapeOverflow { leaf, - source: LeafKernelError::ShapeOverflow { - buffer: "output", - rows: point_count, - cols: leaf_k, - }, + rows: point_count, + columns: leaf_k, })?; grow( From 7efb0775575954cff9f38d3c0ff5657d418b8081 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:54:25 +0000 Subject: [PATCH 37/58] refactor(pipnn): propagate partition worker errors --- diskann/src/graph/pipnn/partitioning.rs | 50 ++++++++----------------- 1 file changed, 15 insertions(+), 35 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 83568138dc..d5376c4033 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -69,8 +69,6 @@ pub(crate) enum PartitionError { level: usize, limit: usize, }, - #[error("partition produced an invalid leaf of size {size}; expected 1..={limit}")] - InvalidLeaf { size: usize, limit: usize }, #[error("invalid {buffer} length: expected {expected}, got {actual}")] InvalidBufferLength { buffer: &'static str, @@ -208,20 +206,21 @@ where results .par_iter_mut() .zip(work.into_par_iter()) - .for_each(|(slot, item)| { + .try_for_each(|(slot, item)| { *slot = Some(partition_one_level::( arch, data, config, item, stripe_buffers, - )); - }); + )?); + Ok::<(), ANNError>(()) + })?; let mut next_work = Vec::new(); for result in results { let (mut pending, mut finished) = - result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))??; + result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))?; next_work .try_reserve(pending.len()) .map_err(ANNError::new)?; @@ -555,14 +554,15 @@ fn scatter_assignments( .par_chunks(stripe_points) .zip(assignments.par_chunks(stripe_assignment_count)), ) - .for_each(|(slot, (points, assignments))| { - *slot = Some(scatter_serial(points, assignments, fanout, leaders)); - }); + .try_for_each(|(slot, (points, assignments))| { + *slot = Some(scatter_serial(points, assignments, fanout, leaders)?); + Ok::<(), ANNError>(()) + })?; let mut locals = Vec::new(); locals.try_reserve_exact(stripes).map_err(ANNError::new)?; for result in partials { - locals.push(result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))??); + locals.push(result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))?); } let mut sizes = filled_vec(leaders, 0usize)?; @@ -707,7 +707,11 @@ fn global_merge_small( } } - validate_leaves(&merged, c_max)?; + debug_assert!( + merged + .iter() + .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) + ); Ok(merged) } @@ -719,19 +723,6 @@ fn drain_sorted(set: &mut HashSet) -> ANNResult> { Ok(values) } -fn validate_leaves(leaves: &[Vec], c_max: usize) -> ANNResult<()> { - if let Some(leaf) = leaves - .iter() - .find(|leaf| leaf.is_empty() || leaf.len() > c_max) - { - return Err(ANNError::new(PartitionError::InvalidLeaf { - size: leaf.len(), - limit: c_max, - })); - } - Ok(()) -} - fn point_ids(points: usize) -> ANNResult> { let mut ids = Vec::new(); ids.try_reserve_exact(points).map_err(ANNError::new)?; @@ -1245,15 +1236,4 @@ mod tests { } ); } - - #[test] - fn rejects_empty_and_oversized_leaves() { - for (leaves, size) in [(vec![vec![]], 0), (vec![vec![0, 1, 2]], 3)] { - let error = validate_leaves(&leaves, 2).unwrap_err(); - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidLeaf { size, limit: 2 } - ); - } - } } From 73c22845f3f6a5723529e86a0296af82f1944db4 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:59:51 +0000 Subject: [PATCH 38/58] refactor(pipnn): remove partition assertions --- diskann/src/graph/pipnn/partitioning.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index d5376c4033..2817c38ca3 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -707,11 +707,6 @@ fn global_merge_small( } } - debug_assert!( - merged - .iter() - .all(|leaf| !leaf.is_empty() && leaf.len() <= c_max) - ); Ok(merged) } From 4b9c34efcdfbd4df4dc2f740ed8c38853ba16cd5 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 14:45:44 +0000 Subject: [PATCH 39/58] docs(pipnn): state core function contracts --- diskann/src/graph/pipnn/mod.rs | 4 ++++ diskann/src/graph/pipnn/partitioning.rs | 16 ++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 6aa63f5e58..0ad85d1657 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -236,6 +236,10 @@ where } } +/// Build the graph with concrete architecture `A` and metric `M`. +/// +/// The function partitions the dataset, builds direct candidates, and applies +/// final graph-degree pruning. fn build_graph_for( arch: A, data: MatrixView<'_, T>, diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 2817c38ca3..483098bd4b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -158,6 +158,10 @@ where Ok(leaves) } +/// Partition one replica until each leaf has at most `c_max` points. +/// +/// The function processes one work queue per recursion level. It merges leaves +/// smaller than `c_min` after the queue becomes empty. fn partition_replica( arch: A, data: MatrixView<'_, T>, @@ -244,6 +248,10 @@ where })) } +/// Process one partition work item. +/// +/// The function samples leaders, assigns all points, and separates complete +/// leaves from clusters that require another recursion level. fn partition_one_level( arch: A, data: MatrixView<'_, T>, @@ -406,6 +414,10 @@ where scatter_assignments(point_ids, &assignments, fanout, leader_ids.len()) } +/// Assign one point stripe to the selected leaders. +/// +/// The function gathers point vectors, computes point-to-leader dot products, +/// and writes nearest leader-column IDs to `assignments`. #[inline] #[allow(clippy::too_many_arguments)] fn assign_stripe( @@ -640,6 +652,10 @@ fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { Ok(clusters) } +/// Merge leaves smaller than `c_min` without exceeding `c_max`. +/// +/// A `HashSet` removes duplicate point IDs across merged leaves. The function +/// sorts each merged result before it returns. fn global_merge_small( leaves: Vec>, c_min: usize, From e334e6818a9ca2027d1e2d5293b42d6fbe4fa392 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:05:49 +0000 Subject: [PATCH 40/58] refactor(pipnn): use domain names in core flow --- diskann/src/graph/pipnn/finalization.rs | 64 ++++++++--------- diskann/src/graph/pipnn/leaf_build.rs | 19 ++--- diskann/src/graph/pipnn/mod.rs | 31 +++++---- diskann/src/graph/pipnn/partitioning.rs | 92 +++++++++++++------------ 4 files changed, 106 insertions(+), 100 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index c44c35bebe..55d86103e6 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -6,15 +6,15 @@ //! Graph-degree enforcement with the Vamana RobustPrune kernel. //! //! Candidate merging can produce more than `R` IDs for one point. This module -//! checks every global ID before parallel work starts. It returns a list with at -//! most `R` IDs without distance work. +//! checks every global ID before parallel work starts. A list at or below `R` +//! returns without distance calculations. //! //! For a longer list, the module computes each source distance. It sorts the //! candidates and calls RobustPrune. The module then writes the selected IDs into //! the original list allocation. //! //! RobustPrune defines occlusion and alpha-round behavior. This module supplies -//! contiguous vector access and a distance function for input type `T`. +//! source vectors and metric distances. use crate::{ ANNError, ANNResult, @@ -45,16 +45,15 @@ pub(crate) enum FinalizationError { TooManyCandidates { actual: usize, max: usize }, } -/// Reusable buffers for one Rayon job. +/// RobustPrune state for one Rayon job. /// -/// `pool` stores candidates with source distances. `prepared` stores each -/// distance and optional non-self ID. `states` stores one RobustPrune state for -/// each sorted candidate. +/// `candidate_slots` and `prune_states` stay positionally aligned with +/// `sorted_candidates`. #[derive(Default)] -struct Workspace { - pool: Vec>, - prepared: Vec<(f32, Option)>, - states: Vec, +struct PruneWorkspace { + sorted_candidates: Vec>, + candidate_slots: Vec<(f32, Option)>, + prune_states: Vec, } /// Check candidate IDs and prune each list that exceeds the graph degree. @@ -78,7 +77,7 @@ where .into_par_iter() .enumerate() .map_init( - Workspace::default, + PruneWorkspace::default, |workspace, (source, mut source_candidates)| { // Candidate merging already removes duplicate IDs. A list within // the degree limit needs no distance calculation. @@ -88,13 +87,13 @@ where let source_id = u32::try_from(source).map_err(ANNError::new)?; let source_vector = data.row(source); - workspace.pool.clear(); + workspace.sorted_candidates.clear(); workspace - .pool + .sorted_candidates .try_reserve(source_candidates.len()) .map_err(ANNError::new)?; workspace - .pool + .sorted_candidates .extend(source_candidates.iter().copied().map(|candidate| { Neighbor::new( candidate, @@ -103,44 +102,47 @@ where ) })); - let candidate_count = workspace.pool.len(); + let candidate_count = workspace.sorted_candidates.len(); if candidate_count > u16::MAX as usize { return Err(ANNError::new(FinalizationError::TooManyCandidates { actual: candidate_count, max: u16::MAX as usize, })); } - workspace.prepared.clear(); + workspace.candidate_slots.clear(); workspace - .prepared + .candidate_slots .try_reserve(candidate_count) .map_err(ANNError::new)?; // Sort all candidates before the code marks a self-edge as absent. // Thus, self-edge removal cannot add a farther candidate. The // `SortedNeighbors` value carries this order into RobustPrune. - let sorted = SortedNeighbors::new(&mut workspace.pool, candidate_count); - workspace.prepared.extend(sorted.iter().map(|neighbor| { - let id = *neighbor.id(); - (*neighbor.distance(), (id != source_id).then_some(id)) - })); + let sorted = + SortedNeighbors::new(&mut workspace.sorted_candidates, candidate_count); workspace - .states + .candidate_slots + .extend(sorted.iter().map(|neighbor| { + let id = *neighbor.id(); + (*neighbor.distance(), (id != source_id).then_some(id)) + })); + workspace + .prune_states .try_reserve( workspace - .prepared + .candidate_slots .len() - .saturating_sub(workspace.states.len()), + .saturating_sub(workspace.prune_states.len()), ) .map_err(ANNError::new)?; workspace - .states - .resize(workspace.prepared.len(), prune::State::default()); + .prune_states + .resize(workspace.candidate_slots.len(), prune::State::default()); let selected = prune::robust_prune( &sorted, - &workspace.prepared, - workspace.states.as_mut_slice(), + &workspace.candidate_slots, + workspace.prune_states.as_mut_slice(), degree, graph.alpha(), graph.prune_kind(), @@ -153,7 +155,7 @@ where ); let mut guard = source_candidates.resize(selected); - for (destination, state) in guard.iter_mut().zip(workspace.states.iter()) { + for (destination, state) in guard.iter_mut().zip(workspace.prune_states.iter()) { *destination = *sorted[state.neighbor as usize].id(); } guard.finish(selected); diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6e787c02bf..25d7df8178 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -194,7 +194,7 @@ impl DirectCandidates { local_adjacency: &[AdjacencyList], ) -> Result<(), LeafBuildError> { for (&source, additions) in point_ids.iter().zip(local_adjacency) { - // `build_leaf` checks every point ID before this append. + // `add_direct_leaf_candidates` checks every point ID before this append. let candidates = &self.lists[source as usize]; let mut candidates = candidates .lock() @@ -220,9 +220,10 @@ impl DirectCandidates { } } -/// Build symmetric leaf-local k-NN graphs and return unique global candidates. +/// Build direct graph candidates from all overlapping leaves. /// -/// The caller supplies concrete architecture `A` and metric `M`. +/// Each selected leaf pair contributes both edge directions. Candidate lists use +/// global dataset IDs and contain no duplicate IDs. #[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. pub(super) fn build_leaf_candidates( arch: A, @@ -249,7 +250,7 @@ where leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { - build_leaf::( + add_direct_leaf_candidates::( arch, data, leaf, @@ -263,12 +264,12 @@ where candidates.into_lists() } -/// Build and publish the symmetric neighbor lists for one leaf. +/// Add one leaf's symmetric neighbors to the direct candidate lists. /// -/// The function checks all IDs before it indexes the dataset. IDs must be -/// strictly increasing. Reusable vectors can be longer than this leaf. Every -/// read and write uses the active length from `prepare`. -fn build_leaf( +/// The function rejects empty, duplicate, unsorted, or out-of-range point IDs. +/// Reusable buffers can be longer than this leaf, so all accesses use the current +/// leaf shape. +fn add_direct_leaf_candidates( arch: A, data: MatrixView<'_, T>, leaf: usize, diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 0ad85d1657..19c9f73897 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -5,7 +5,8 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! PiPNN builds graph candidates in three steps: +//! PiPNN builds graph candidates in three steps. A leader is a sampled dataset +//! point that acts as the center of one child partition. //! //! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has //! at most `c_max` points. @@ -13,7 +14,6 @@ //! selects local neighbors and merges their global point IDs. //! 3. `finalization` applies Vamana RobustPrune to each candidate list that is //! longer than the graph degree. - //! //! `diskann-wide` selects architecture `A`. One match selects metric marker `M`. //! The build passes both concrete types through all replicas, recursive @@ -54,19 +54,18 @@ use rayon::ThreadPool; use self::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; -/// Configuration for PiPNN partitioning and local-neighbor selection. +/// PiPNN partition and leaf-selection policy. /// -/// [`PiPNNBuildContext`] supplies graph degree, prune policy, alpha, metric, and -/// the Rayon pool. +/// DiskANN graph policy separately supplies degree, alpha, and prune metric. #[derive(Clone, Debug, PartialEq)] pub struct PiPNNConfig { /// Maximum number of points in a leaf. pub c_max: usize, /// Minimum leaf size used by global small-leaf merging. pub c_min: usize, - /// Fraction of a cluster sampled as partition leaders. + /// Fraction of a cluster sampled as child-partition centers. pub p_samp: f64, - /// Number of nearest leaders retained at each overlapping partition level. + /// Number of nearest partition centers assigned to each point at each level. pub fanout: Vec, /// Number of nearest neighbors selected within each leaf (`1..=3`). pub k: usize, @@ -115,7 +114,7 @@ impl PiPNNConfig { } } -/// Checked policy and execution inputs for one PiPNN graph build. +/// PiPNN policy and borrowed execution resources for one graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { pub(crate) config: PiPNNConfig, @@ -149,14 +148,13 @@ impl<'a> PiPNNBuildContext<'a> { } } -/// Build PiPNN adjacency for all points in `data`. +/// Build one PiPNN adjacency list for each point in `data`. /// -/// The function does not select start or frozen points. It does not load a -/// provider or write an index. +/// This graph contains only real dataset points. Start-point selection and index +/// serialization are separate operations. /// /// Raw `u8` and `i8` vectors are not unit-normalized after conversion to `f32`. -/// Therefore, the function evaluates `CosineNormalized` as `Cosine` for these -/// two input types. +/// The build therefore uses norm-aware cosine for these two input types. pub fn build_graph( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, @@ -164,10 +162,13 @@ pub fn build_graph( where T: VectorRepr + Send + Sync + 'static, { - context.pool.install(|| build_graph_inner(data, context)) + context + .pool + .install(|| validate_and_dispatch_build(data, context)) } -fn build_graph_inner( +/// Check dataset bounds and select the architecture and metric implementation. +fn validate_and_dispatch_build( data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, ) -> ANNResult>> diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 483098bd4b..3d953ee461 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -5,10 +5,12 @@ //! Deterministic overlapping partition construction for PiPNN. //! -//! This module converts dataset point IDs into bounded leaf ID lists. It uses -//! dense GEMM and the partition kernel for leader assignment. An `ObjectPool` -//! supplies one reusable buffer set to each Rayon worker chunk. The pool lock is -//! not held during gather, GEMM, or top-k selection. +//! A leader is a sampled point that acts as the center of one child partition. +//! A point can join several leaders, so child partitions can overlap. +//! +//! This module recursively splits dataset point IDs into bounded leaves. It uses +//! dense GEMM to compare points with sampled leaders. An `ObjectPool` supplies +//! reusable buffers to Rayon worker chunks. //! //! A configured level assigns each point to `fanout[level]` leaders. A deeper //! level assigns each point to one leader. Each replica uses a different @@ -99,27 +101,23 @@ impl AsPooled<()> for StripeBuffers { } fn modify(&mut self, _: ()) { - // Keep the largest allocation across leases. `assign_stripe` defines the + // Keep the largest allocation across leases. `assign_point_stripe` defines the // active prefix before each read. } } -/// Reusable partition buffers owned by this partition call. +/// Reusable buffers for point-to-leader assignment. /// -/// `ObjectPool` locks only when it gives or receives a lease. RAII returns each -/// lease after success, error, or panic. Numerical work holds only the lease. +/// `ObjectPool` locks only when it gives or receives a lease. Numerical work +/// holds the lease, not the pool lock. type StripeBufferPool = ObjectPool; -/// Partition all configured replicas into overlapping bounded leaves. -/// -/// An oversized work item samples `ceil(p_samp * points)` leaders, up to -/// `LEADER_CAP`. It assigns each point to the configured number of nearest -/// leaders. It adds each cluster above `c_max` to the work queue. +/// Build overlapping bounded leaves for all configured replicas. /// -/// A level without a configured fanout assigns each point to one leader. The -/// final merge does not make a leaf larger than `c_max`. Each replica covers -/// every input point. The caller supplies concrete architecture `A` and metric -/// `M`. +/// Each split samples partition centers and assigns every cluster point to its +/// nearest centers. A cluster above `c_max` is split again. A level without a +/// configured fanout assigns each point to one center. Each replica covers every +/// input point. pub(super) fn partition( arch: A, data: MatrixView<'_, T>, @@ -196,7 +194,7 @@ where for _ in 0..MAX_PARTITION_ITERATIONS { if work.is_empty() { - return global_merge_small(leaves, config.c_min, config.c_max); + return merge_undersized_leaves(leaves, config.c_min, config.c_max); } let mut results = Vec::new(); @@ -211,7 +209,7 @@ where .par_iter_mut() .zip(work.into_par_iter()) .try_for_each(|(slot, item)| { - *slot = Some(partition_one_level::( + *slot = Some(partition_work_item::( arch, data, config, @@ -236,10 +234,10 @@ where } if work.is_empty() { - return global_merge_small(leaves, config.c_min, config.c_max); + return merge_undersized_leaves(leaves, config.c_min, config.c_max); } let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { - return global_merge_small(leaves, config.c_min, config.c_max); + return merge_undersized_leaves(leaves, config.c_min, config.c_max); }; Err(ANNError::new(PartitionError::IterationLimit { size: largest.indices.len(), @@ -248,11 +246,11 @@ where })) } -/// Process one partition work item. +/// Split one oversized cluster into child partitions. /// -/// The function samples leaders, assigns all points, and separates complete -/// leaves from clusters that require another recursion level. -fn partition_one_level( +/// The function samples center points, assigns the cluster points, and returns +/// bounded leaves separately from child clusters that need another split. +fn partition_work_item( arch: A, data: MatrixView<'_, T>, config: &PiPNNConfig, @@ -301,8 +299,9 @@ where Ok((pending, finished)) } +/// Sample point IDs that act as centers for one partition split. fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResult> { - let count = sample_num_leaders(points.len(), sampling_fraction); + let count = sampled_leader_count(points.len(), sampling_fraction); let mut rng = rand::rngs::StdRng::seed_from_u64(seed); let mut leaders = Vec::new(); leaders.try_reserve_exact(count).map_err(ANNError::new)?; @@ -310,7 +309,12 @@ fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResul Ok(leaders) } -fn sample_num_leaders(points: usize, sampling_fraction: f64) -> usize { +/// Return the number of centers to sample from one cluster. +/// +/// The count is `ceil(points * sampling_fraction)`, limited by `LEADER_CAP` and +/// the number of available points. A cluster with at least two points uses at +/// least two centers. +fn sampled_leader_count(points: usize, sampling_fraction: f64) -> usize { ((points as f64 * sampling_fraction).ceil() as usize) .clamp(2, LEADER_CAP) .min(points) @@ -327,13 +331,11 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { .wrapping_add(salt) } -/// Assign each point to its nearest sampled leaders. -/// -/// The function gathers leader vectors once. It divides points into bounded -/// stripes. Each worker chunk reuses one leased buffer set for all its stripes. +/// Assign each cluster point to its nearest sampled partition centers. /// -/// The flat assignment matrix keeps point order. The scatter step keeps the same -/// order in each leader cluster. Recursive sampling depends on this order. +/// The function gathers center vectors once and evaluates points in bounded +/// stripes. The assignment matrix keeps point order. Scatter preserves this order +/// inside each child partition, which makes recursive sampling deterministic. fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, @@ -397,7 +399,7 @@ where { let first_point = worker_first + stripe * stripe_points; let stripe_point_count = stripe_assignments.len() / fanout; - assign_stripe::( + assign_point_stripe::( arch, data, &point_ids[first_point..first_point + stripe_point_count], @@ -414,13 +416,13 @@ where scatter_assignments(point_ids, &assignments, fanout, leader_ids.len()) } -/// Assign one point stripe to the selected leaders. +/// Assign one point stripe to sampled partition centers. /// -/// The function gathers point vectors, computes point-to-leader dot products, -/// and writes nearest leader-column IDs to `assignments`. +/// The function gathers point vectors and computes point-to-center dot products. +/// It writes center-column IDs for partition scatter. #[inline] #[allow(clippy::too_many_arguments)] -fn assign_stripe( +fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, point_ids: &[u32], @@ -656,7 +658,7 @@ fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { /// /// A `HashSet` removes duplicate point IDs across merged leaves. The function /// sorts each merged result before it returns. -fn global_merge_small( +fn merge_undersized_leaves( leaves: Vec>, c_min: usize, c_max: usize, @@ -975,7 +977,7 @@ mod tests { fn global_merge_canonicalizes_small_leaf_membership() { let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; - let merged = global_merge_small(leaves, 4, 8).unwrap(); + let merged = merge_undersized_leaves(leaves, 4, 8).unwrap(); assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); } @@ -984,7 +986,7 @@ mod tests { fn global_merge_never_overfills_before_reaching_c_min() { let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; - let merged = global_merge_small(leaves, 11, 11).unwrap(); + let merged = merge_undersized_leaves(leaves, 11, 11).unwrap(); assert_eq!( merged, @@ -994,7 +996,7 @@ mod tests { #[test] fn global_merge_fills_exact_capacity_before_flushing() { - let merged = global_merge_small(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + let merged = merge_undersized_leaves(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); assert_eq!(merged, vec![vec![0, 1, 2, 3]]); } @@ -1123,9 +1125,9 @@ mod tests { #[test] fn leader_count_is_bounded() { - assert_eq!(sample_num_leaders(1, 1.0), 1); - assert_eq!(sample_num_leaders(10, 0.01), 2); - assert_eq!(sample_num_leaders(50_000, 1.0), LEADER_CAP); + assert_eq!(sampled_leader_count(1, 1.0), 1); + assert_eq!(sampled_leader_count(10, 0.01), 2); + assert_eq!(sampled_leader_count(50_000, 1.0), LEADER_CAP); } #[test] From aec3f2a7f3a472d65838396bbfdfcef5076856d6 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:39:55 +0000 Subject: [PATCH 41/58] docs(pipnn): remove layout-restatement comments --- diskann/src/graph/pipnn/mod.rs | 4 ++-- diskann/src/graph/pipnn/partitioning.rs | 10 +++------- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 19c9f73897..a04157897d 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -237,9 +237,9 @@ where } } -/// Build the graph with concrete architecture `A` and metric `M`. +/// Run the PiPNN graph pipeline for one selected metric implementation. /// -/// The function partitions the dataset, builds direct candidates, and applies +/// The function builds overlapping leaves, merges direct candidates, and applies /// final graph-degree pruning. fn build_graph_for( arch: A, diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 3d953ee461..c988b38c3b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -537,11 +537,10 @@ where Ok(()) } -/// Convert point-major assignments into leader-major clusters. +/// Group assigned point IDs by child partition. /// -/// A small input uses one serial pass with exact capacities. A large input makes -/// at most one partial cluster set for each Rayon worker. It then merges each -/// leader independently. Stripe-order concatenation matches serial member order. +/// Both the serial and parallel paths preserve point order inside each child. +/// This order is required for deterministic recursive sampling. fn scatter_assignments( points: &[u32], assignments: &[u32], @@ -750,9 +749,6 @@ fn filled_vec(len: usize, value: T) -> ANNResult> { Ok(values) } -/// Grow `values` to at least `len` elements and do not shrink it. -/// -/// Callers use an explicit active prefix. fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { if values.len() >= len { return Ok(()); From 6c6b0e83fc7298249fb1a0c4af96954eec39ce86 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:48:39 +0000 Subject: [PATCH 42/58] docs(pipnn): define leaf domain term --- diskann/src/graph/pipnn/mod.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index a04157897d..8b94b67625 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -6,7 +6,8 @@ //! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! //! PiPNN builds graph candidates in three steps. A leader is a sampled dataset -//! point that acts as the center of one child partition. +//! point that acts as the center of one child partition. A leaf is a bounded +//! child partition used for local neighbor selection. //! //! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has //! at most `c_max` points. From 3c16bec5c98665cef1c499ef59c584abc4b49c7a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:05:33 +0000 Subject: [PATCH 43/58] refactor(pipnn): inline poisoned-list errors --- diskann/src/graph/pipnn/leaf_build.rs | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 25d7df8178..6870c05900 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -198,7 +198,7 @@ impl DirectCandidates { let candidates = &self.lists[source as usize]; let mut candidates = candidates .lock() - .map_err(|_| poisoned_candidate_list(source))?; + .map_err(|_| LeafBuildError::PoisonedCandidateList { point: source })?; candidates.extend_from_slice(additions); } Ok(()) @@ -210,9 +210,12 @@ impl DirectCandidates { .try_reserve_exact(self.lists.len()) .map_err(|source| allocation_error("candidate output", self.lists.len(), source))?; for (point, candidates) in self.lists.into_iter().enumerate() { - let mut candidates = candidates - .into_inner() - .map_err(|_| poisoned_candidate_list(point as u32))?; + let mut candidates = + candidates + .into_inner() + .map_err(|_| LeafBuildError::PoisonedCandidateList { + point: point as u32, + })?; candidates.sort(); output.push(candidates); } @@ -427,10 +430,6 @@ fn allocation_error( } } -fn poisoned_candidate_list(point: u32) -> LeafBuildError { - LeafBuildError::PoisonedCandidateList { point } -} - #[cfg(test)] mod tests { use diskann_utils::views::MatrixView; From 7b32795ab63cf943c936081b5544d13e7d408cab Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:55:56 +0000 Subject: [PATCH 44/58] refactor(pipnn): clarify leaf policy names --- diskann/src/graph/pipnn/mod.rs | 33 +++++++++++++------------ diskann/src/graph/pipnn/partitioning.rs | 2 +- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 8b94b67625..6f804ff9f0 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -66,10 +66,11 @@ pub struct PiPNNConfig { pub c_min: usize, /// Fraction of a cluster sampled as child-partition centers. pub p_samp: f64, - /// Number of nearest partition centers assigned to each point at each level. + /// Number of nearest centers assigned at each recursive partition level. + /// Levels after this schedule assign each point to one center. pub fanout: Vec, /// Number of nearest neighbors selected within each leaf (`1..=3`). - pub k: usize, + pub leaf_k: usize, /// Number of independent partition passes over the dataset. pub replicas: usize, } @@ -89,9 +90,9 @@ impl PiPNNConfig { self.c_min, self.c_max ))); } - if !self.p_samp.is_finite() || !(0.0..=1.0).contains(&self.p_samp) || self.p_samp == 0.0 { + if !(0.0 < self.p_samp && self.p_samp <= 1.0) { return Err(config_error(format!( - "p_samp ({}) must be finite and in (0, 1]", + "p_samp ({}) must be in (0, 1]", self.p_samp ))); } @@ -101,10 +102,10 @@ impl PiPNNConfig { if self.fanout.contains(&0) { return Err(config_error("fanout values must be greater than zero")); } - if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.k) { + if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.leaf_k) { return Err(config_error(format!( - "k ({}) must be in [1, {}]", - self.k, + "leaf_k ({}) must be in [1, {}]", + self.leaf_k, leaf_kernel::MAX_LEAF_NEIGHBORS ))); } @@ -260,7 +261,7 @@ where // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.k) + leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) .map_err(ANNError::new) })?; // Finalization consumes each candidate list. It reuses that list's allocation @@ -331,7 +332,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - k: 1, + leaf_k: 1, replicas: 1, } } @@ -408,7 +409,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - k: leaf_kernel::MAX_LEAF_NEIGHBORS, + leaf_k: leaf_kernel::MAX_LEAF_NEIGHBORS, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -474,7 +475,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - k: 1, + leaf_k: 1, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); @@ -500,7 +501,7 @@ mod build_graph_tests { c_min: 4, p_samp: 0.25, fanout: vec![3, 2], - k: 3, + leaf_k: 3, replicas: 2, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -532,7 +533,7 @@ mod build_graph_tests { c_min, p_samp: 0.5, fanout: vec![2], - k: rng.random_range(1..=3), + leaf_k: rng.random_range(1..=3), replicas: rng.random_range(1..=2), }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -560,7 +561,7 @@ mod config_tests { c_min: 64, p_samp: 0.01, fanout: vec![10, 3], - k: 2, + leaf_k: 2, replicas: 1, } } @@ -622,11 +623,11 @@ mod config_tests { ..pipnn_config() }, PiPNNConfig { - k: 0, + leaf_k: 0, ..pipnn_config() }, PiPNNConfig { - k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, + leaf_k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, ..pipnn_config() }, PiPNNConfig { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index c988b38c3b..b00e5bd02b 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -834,7 +834,7 @@ mod tests { c_min, p_samp: 0.25, fanout, - k: 1, + leaf_k: 1, replicas, } } From 716af335d06d52b228cb99328399c58fcb720a71 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 07:07:10 +0000 Subject: [PATCH 45/58] fix(pipnn): reject malformed assignments --- diskann/src/graph/pipnn/partitioning.rs | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index b00e5bd02b..431107fff1 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -614,6 +614,15 @@ fn scatter_serial( fanout: usize, leaders: usize, ) -> ANNResult>> { + let expected = checked_area("scatter assignments", points.len(), fanout)?; + if assignments.len() != expected { + return Err(ANNError::new(PartitionError::InvalidBufferLength { + buffer: "scatter assignments", + expected, + actual: assignments.len(), + })); + } + let mut sizes = filled_vec(leaders, 0usize)?; for &leader in assignments { let Some(size) = sizes.get_mut(leader as usize) else { @@ -1232,6 +1241,20 @@ mod tests { ); } + #[test] + fn rejects_invalid_assignment_length() { + let error = scatter_serial(&[7, 8], &[0, 1, 0], 2, 2).unwrap_err(); + + assert_eq!( + error.downcast::().unwrap(), + PartitionError::InvalidBufferLength { + buffer: "scatter assignments", + expected: 4, + actual: 3, + } + ); + } + #[test] fn rejects_assignment_to_an_unknown_leader() { let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); From a5ed7ce1d36f89db0127cdd19dbdad619dbeaeca Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:50:58 +0000 Subject: [PATCH 46/58] refactor(pipnn): use stage metric contracts --- diskann/src/graph/pipnn/leaf_build.rs | 6 +-- diskann/src/graph/pipnn/mod.rs | 11 +++-- diskann/src/graph/pipnn/partitioning.rs | 64 ++++++++++++------------- 3 files changed, 42 insertions(+), 39 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6870c05900..04129b98c4 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -26,7 +26,7 @@ use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::KernelMetric, + kernel_metric::LeafKernelMetric, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, @@ -239,7 +239,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: LeafKernelMetric, T: VectorRepr + 'static, { if data.ncols() == 0 { @@ -286,7 +286,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: LeafKernelMetric, T: VectorRepr + 'static, { if point_ids.is_empty() { diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 6f804ff9f0..c1a09b7d46 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -53,7 +53,9 @@ use diskann_wide::{ }; use rayon::ThreadPool; -use self::kernel_metric::{Cosine, CosineNormalized, InnerProduct, KernelMetric, L2}; +use self::kernel_metric::{ + Cosine, CosineNormalized, InnerProduct, L2, LeafKernelMetric, MetricTag, PartitionKernelMetric, +}; /// PiPNN partition and leaf-selection policy. /// @@ -253,7 +255,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: LeafKernelMetric + PartitionKernelMetric, T: VectorRepr + Send + Sync + 'static, { let leaves = tracing::info_span!("pipnn.partition") @@ -266,8 +268,9 @@ where })?; // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. - tracing::info_span!("pipnn.finalization") - .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, M::METRIC)) + tracing::info_span!("pipnn.finalization").in_scope(|| { + finalization::prune_overfull(data, candidates, context.graph, ::METRIC) + }) } fn effective_metric(metric: Metric) -> Metric { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 431107fff1..811d3d369d 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -31,10 +31,8 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::KernelMetric, - partition_kernel::{ - PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders, - }, + kernel_metric::{MetricTag, PartitionKernelMetric}, + partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, }; // These constants control internal batching and deterministic seed generation. @@ -91,7 +89,7 @@ struct WorkItem { struct StripeBuffers { points: Vec, dots: Vec, - point_scales: Vec, + point_squared_norms: Vec, kernel: PartitionKernelWorkspace, } @@ -128,7 +126,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { let points = data.nrows(); @@ -172,7 +170,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { let initial_indices = point_ids(data.nrows())?; @@ -262,7 +260,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { let points = item.indices.len(); @@ -349,28 +347,29 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { + let metric = ::METRIC; let dimension_count = data.ncols(); let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; gather_vectors(data, leader_ids, &mut leader_values)?; - let mut leader_scales = if matches!(M::METRIC, Metric::L2 | Metric::Cosine) { + let mut leader_norm_values = if matches!(metric, Metric::L2 | Metric::Cosine) { filled_vec(leader_ids.len(), 0.0f32)? } else { Vec::new() }; - for (scale, leader_vector) in leader_scales + for (norm_value, leader_vector) in leader_norm_values .iter_mut() .zip(leader_values.chunks_exact(dimension_count)) { // Leader norms affect top-k order. Use this scalar reduction order. // SIMD reassociation changes low bits and can change a near-tie branch. - *scale = leader_vector.iter().map(|value| value * value).sum(); - if M::METRIC == Metric::Cosine { - *scale = scale.sqrt(); + *norm_value = leader_vector.iter().map(|value| value * value).sum(); + if metric == Metric::Cosine { + *norm_value = norm_value.sqrt(); } } @@ -404,7 +403,7 @@ where data, &point_ids[first_point..first_point + stripe_point_count], &leader_values, - &leader_scales, + &leader_norm_values, fanout, &mut buffers, stripe_assignments, @@ -427,7 +426,7 @@ fn assign_point_stripe( data: MatrixView<'_, T>, point_ids: &[u32], leader_values: &[f32], - leader_scales: &[f32], + leader_norm_values: &[f32], fanout: usize, buffers: &mut StripeBuffers, assignments: &mut [u32], @@ -437,7 +436,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: KernelMetric, + M: PartitionKernelMetric, T: VectorRepr, { let point_count = point_ids.len(); @@ -453,7 +452,7 @@ where let StripeBuffers { points: point_buffer, dots: dot_buffer, - point_scales: point_scale_buffer, + point_squared_norms: point_squared_norm_buffer, kernel: kernel_workspace, } = buffers; let point_values = &mut point_buffer[..point_values_len]; @@ -473,28 +472,29 @@ where ) .map_err(ANNError::new)?; - let point_scales = if M::METRIC == Metric::Cosine { - grow_fallible(point_scale_buffer, point_count, 0.0)?; - let point_scales = &mut point_scale_buffer[..point_count]; - for (scale, point_values) in point_scales + let metric = ::METRIC; + let point_squared_norms = if metric == Metric::Cosine { + grow_fallible(point_squared_norm_buffer, point_count, 0.0)?; + let point_squared_norms = &mut point_squared_norm_buffer[..point_count]; + for (norm_value, point_values) in point_squared_norms .iter_mut() .zip(point_values.chunks_exact(dimensions)) { - *scale = FastL2NormSquared.evaluate(point_values); + *norm_value = FastL2NormSquared.evaluate(point_values); } - &*point_scales + &*point_squared_norms } else { &[] }; - let scales = match M::METRIC { - Metric::L2 => PartitionScales::L2 { - leader_squared_norms: leader_scales, + let norms = match metric { + Metric::L2 => PartitionNorms::L2 { + leader_squared_norms: leader_norm_values, }, - Metric::Cosine => PartitionScales::Cosine { - point_squared_norms: point_scales, - leader_norms: leader_scales, + Metric::Cosine => PartitionNorms::Cosine { + point_squared_norms, + leader_norms: leader_norm_values, }, - Metric::CosineNormalized | Metric::InnerProduct => PartitionScales::None, + Metric::CosineNormalized | Metric::InnerProduct => PartitionNorms::None, }; let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { ANNError::new(PartitionError::InvalidBufferLength { @@ -512,7 +512,7 @@ where })?; nearest_leaders::( arch, - PartitionInput { dots, scales }, + PartitionInput { dots, norms }, output, kernel_workspace, ) From 92e6616a56a2ff736628ed033c777840e963d70e Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:43:58 +0000 Subject: [PATCH 47/58] refactor(pipnn): use partition metric identity --- diskann/src/graph/pipnn/mod.rs | 9 +++++++-- diskann/src/graph/pipnn/partitioning.rs | 6 +++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index c1a09b7d46..d83cd74744 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -54,7 +54,7 @@ use diskann_wide::{ use rayon::ThreadPool; use self::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, LeafKernelMetric, MetricTag, PartitionKernelMetric, + Cosine, CosineNormalized, InnerProduct, L2, LeafKernelMetric, PartitionKernelMetric, }; /// PiPNN partition and leaf-selection policy. @@ -269,7 +269,12 @@ where // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. tracing::info_span!("pipnn.finalization").in_scope(|| { - finalization::prune_overfull(data, candidates, context.graph, ::METRIC) + finalization::prune_overfull( + data, + candidates, + context.graph, + ::METRIC, + ) }) } diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 811d3d369d..11660adafd 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -31,7 +31,7 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::{MetricTag, PartitionKernelMetric}, + kernel_metric::PartitionKernelMetric, partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, }; @@ -350,7 +350,7 @@ where M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { - let metric = ::METRIC; + let metric = M::METRIC; let dimension_count = data.ncols(); let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; @@ -472,7 +472,7 @@ where ) .map_err(ANNError::new)?; - let metric = ::METRIC; + let metric = M::METRIC; let point_squared_norms = if metric == Metric::Cosine { grow_fallible(point_squared_norm_buffer, point_count, 0.0)?; let point_squared_norms = &mut point_squared_norm_buffer[..point_count]; From bc889d518f1a0bd16c6310b771a26f77c5852799 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:04:04 +0000 Subject: [PATCH 48/58] refactor(pipnn): prepare norms before kernel calls --- diskann/src/graph/pipnn/leaf_build.rs | 47 +++++++-- diskann/src/graph/pipnn/mod.rs | 45 +++++---- diskann/src/graph/pipnn/partitioning.rs | 122 +++++++++++++++--------- 3 files changed, 141 insertions(+), 73 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 04129b98c4..5bca5b6778 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -22,11 +22,12 @@ use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_vector::distance::Metric; use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::LeafKernelMetric, + kernel_metric::{LeafKernelMetric, norm_from_squared}, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, @@ -99,6 +100,7 @@ pub(crate) enum LeafBuildError { struct LeafBuffers { point_values: Vec, dots: Vec, + norms: Vec, neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, @@ -233,6 +235,7 @@ pub(super) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, requested_k: usize, + metric: Metric, ) -> Result>, LeafBuildError> where A: Architecture, @@ -256,6 +259,7 @@ where add_direct_leaf_candidates::( arch, data, + metric, leaf, point_ids, requested_k, @@ -272,9 +276,11 @@ where /// The function rejects empty, duplicate, unsorted, or out-of-range point IDs. /// Reusable buffers can be longer than this leaf, so all accesses use the current /// leaf shape. +#[allow(clippy::too_many_arguments)] fn add_direct_leaf_candidates( arch: A, data: MatrixView<'_, T>, + metric: Metric, leaf: usize, point_ids: &[u32], requested_k: usize, @@ -345,6 +351,21 @@ where leaf, buffer: "leaf dot-product matrix", })?; + let norm_count = if matches!(metric, Metric::L2 | Metric::Cosine) { + point_ids.len() + } else { + 0 + }; + grow("leaf norms", &mut buffers.norms, norm_count, 0.0)?; + for (point, norm) in buffers.norms[..norm_count].iter_mut().enumerate() { + let squared_norm = dots[(point, point)]; + *norm = if metric == Metric::Cosine { + norm_from_squared(squared_norm) + } else { + squared_norm + }; + } + let norms = &buffers.norms[..norm_count]; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), @@ -354,7 +375,7 @@ where leaf, buffer: "leaf output", })?; - nearest_neighbors::(arch, dots, output, &mut buffers.kernel_workspace) + nearest_neighbors::(arch, dots, norms, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; buffers.prepare_local_adjacency(point_ids.len())?; @@ -486,23 +507,33 @@ mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => { - build_leaf_candidates::(arch, call.data, call.leaves, call.k) - } - Metric::Cosine => { - build_leaf_candidates::(arch, call.data, call.leaves, call.k) - } + Metric::L2 => build_leaf_candidates::( + arch, + call.data, + call.leaves, + call.k, + Metric::L2, + ), + Metric::Cosine => build_leaf_candidates::( + arch, + call.data, + call.leaves, + call.k, + Metric::Cosine, + ), Metric::CosineNormalized => build_leaf_candidates::( arch, call.data, call.leaves, call.k, + Metric::CosineNormalized, ), Metric::InnerProduct => build_leaf_candidates::( arch, call.data, call.leaves, call.k, + Metric::InnerProduct, ), } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index d83cd74744..2d7e088c7e 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -229,14 +229,22 @@ where call: BuildGraphCall<'_, '_, '_, T>, ) -> ANNResult>> { match call.metric { - Metric::L2 => build_graph_for::(arch, call.data, call.context), - Metric::Cosine => build_graph_for::(arch, call.data, call.context), - Metric::CosineNormalized => { - build_graph_for::(arch, call.data, call.context) - } - Metric::InnerProduct => { - build_graph_for::(arch, call.data, call.context) + Metric::L2 => build_graph_for::(arch, call.data, call.context, Metric::L2), + Metric::Cosine => { + build_graph_for::(arch, call.data, call.context, Metric::Cosine) } + Metric::CosineNormalized => build_graph_for::( + arch, + call.data, + call.context, + Metric::CosineNormalized, + ), + Metric::InnerProduct => build_graph_for::( + arch, + call.data, + call.context, + Metric::InnerProduct, + ), } } } @@ -249,6 +257,7 @@ fn build_graph_for( arch: A, data: MatrixView<'_, T>, context: &PiPNNBuildContext<'_>, + metric: Metric, ) -> ANNResult>> where A: Architecture, @@ -259,23 +268,23 @@ where T: VectorRepr + Send + Sync + 'static, { let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition::(arch, data, &context.config))?; + .in_scope(|| partitioning::partition::(arch, data, metric, &context.config))?; // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) - .map_err(ANNError::new) + leaf_build::build_leaf_candidates::( + arch, + data, + leaves, + context.config.leaf_k, + metric, + ) + .map_err(ANNError::new) })?; // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. - tracing::info_span!("pipnn.finalization").in_scope(|| { - finalization::prune_overfull( - data, - candidates, - context.graph, - ::METRIC, - ) - }) + tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) } fn effective_metric(metric: Metric) -> Metric { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 11660adafd..4254978205 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -31,7 +31,7 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::PartitionKernelMetric, + kernel_metric::{PartitionKernelMetric, norm_from_squared}, partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, }; @@ -87,10 +87,10 @@ struct WorkItem { #[derive(Default)] struct StripeBuffers { - points: Vec, - dots: Vec, - point_squared_norms: Vec, - kernel: PartitionKernelWorkspace, + point_values: Vec, + dot_values: Vec, + point_norms: Vec, + kernel_workspace: PartitionKernelWorkspace, } impl AsPooled<()> for StripeBuffers { @@ -119,6 +119,7 @@ type StripeBufferPool = ObjectPool; pub(super) fn partition( arch: A, data: MatrixView<'_, T>, + metric: Metric, config: &PiPNNConfig, ) -> ANNResult>> where @@ -145,7 +146,7 @@ where for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = - partition_replica::(arch, data, config, seed, &stripe_buffers)?; + partition_replica::(arch, data, metric, config, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::new)?; @@ -161,6 +162,7 @@ where fn partition_replica( arch: A, data: MatrixView<'_, T>, + metric: Metric, config: &PiPNNConfig, seed: u64, stripe_buffers: &StripeBufferPool, @@ -210,6 +212,7 @@ where *slot = Some(partition_work_item::( arch, data, + metric, config, item, stripe_buffers, @@ -251,6 +254,7 @@ where fn partition_work_item( arch: A, data: MatrixView<'_, T>, + metric: Metric, config: &PiPNNConfig, item: WorkItem, stripe_buffers: &StripeBufferPool, @@ -270,8 +274,15 @@ where config.p_samp, mix_seed(item.seed, points as u64), )?; - let clusters = - assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; + let clusters = assign_to_leaders::( + arch, + data, + metric, + &item.indices, + &leaders, + fanout, + stripe_buffers, + )?; let mut pending = Vec::new(); let mut finished = Vec::new(); @@ -337,6 +348,7 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, + metric: Metric, point_ids: &[u32], leader_ids: &[u32], fanout: usize, @@ -350,7 +362,6 @@ where M: PartitionKernelMetric, T: VectorRepr + Send + Sync, { - let metric = M::METRIC; let dimension_count = data.ncols(); let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; @@ -369,7 +380,7 @@ where // SIMD reassociation changes low bits and can change a near-tie branch. *norm_value = leader_vector.iter().map(|value| value * value).sum(); if metric == Metric::Cosine { - *norm_value = norm_value.sqrt(); + *norm_value = norm_from_squared(*norm_value); } } @@ -401,6 +412,7 @@ where assign_point_stripe::( arch, data, + metric, &point_ids[first_point..first_point + stripe_point_count], &leader_values, &leader_norm_values, @@ -424,6 +436,7 @@ where fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, + metric: Metric, point_ids: &[u32], leader_values: &[f32], leader_norm_values: &[f32], @@ -447,17 +460,28 @@ where let output_len = checked_area("partition assignments", point_count, fanout)?; // Keep each buffer at its largest length. Every operation uses an explicit // active prefix. - grow_fallible(&mut buffers.points, point_values_len, 0.0)?; - grow_fallible(&mut buffers.dots, dots_len, 0.0)?; + grow_fallible(&mut buffers.point_values, point_values_len, 0.0)?; + grow_fallible(&mut buffers.dot_values, dots_len, 0.0)?; let StripeBuffers { - points: point_buffer, - dots: dot_buffer, - point_squared_norms: point_squared_norm_buffer, - kernel: kernel_workspace, + point_values: point_buffer, + dot_values: dot_buffer, + point_norms: point_norm_buffer, + kernel_workspace, } = buffers; - let point_values = &mut point_buffer[..point_values_len]; + let mut point_matrix = MutMatrixView::try_from( + &mut point_buffer[..point_values_len], + point_count, + dimensions, + ) + .map_err(|error| { + ANNError::new(PartitionError::InvalidBufferLength { + buffer: "point stripe", + expected: point_values_len, + actual: error.into_inner().len(), + }) + })?; let dots = &mut dot_buffer[..dots_len]; - gather_vectors(data, point_ids, point_values)?; + gather_vectors(data, point_ids, point_matrix.as_mut_slice())?; diskann_linalg::sgemm( Transpose::None, Transpose::Ordinary, @@ -465,36 +489,29 @@ where leader_count, dimensions, 1.0, - point_values, + point_matrix.as_slice(), leader_values, None, dots, ) .map_err(ANNError::new)?; - let metric = M::METRIC; - let point_squared_norms = if metric == Metric::Cosine { - grow_fallible(point_squared_norm_buffer, point_count, 0.0)?; - let point_squared_norms = &mut point_squared_norm_buffer[..point_count]; - for (norm_value, point_values) in point_squared_norms + let point_norms = if metric == Metric::Cosine { + grow_fallible(point_norm_buffer, point_count, 0.0)?; + let point_norms = &mut point_norm_buffer[..point_count]; + for (norm_value, point_values) in point_norms .iter_mut() - .zip(point_values.chunks_exact(dimensions)) + .zip(point_matrix.as_slice().chunks_exact(dimensions)) { - *norm_value = FastL2NormSquared.evaluate(point_values); + *norm_value = norm_from_squared(FastL2NormSquared.evaluate(point_values)); } - &*point_squared_norms + &*point_norms } else { &[] }; - let norms = match metric { - Metric::L2 => PartitionNorms::L2 { - leader_squared_norms: leader_norm_values, - }, - Metric::Cosine => PartitionNorms::Cosine { - point_squared_norms, - leader_norms: leader_norm_values, - }, - Metric::CosineNormalized | Metric::InnerProduct => PartitionNorms::None, + let norms = PartitionNorms { + point_norms, + leader_norms: leader_norm_values, }; let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { ANNError::new(PartitionError::InvalidBufferLength { @@ -512,6 +529,7 @@ where })?; nearest_leaders::( arch, + metric, PartitionInput { dots, norms }, output, kernel_workspace, @@ -814,14 +832,22 @@ mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => partition::(arch, call.data, call.config), - Metric::Cosine => partition::(arch, call.data, call.config), - Metric::CosineNormalized => { - partition::(arch, call.data, call.config) - } - Metric::InnerProduct => { - partition::(arch, call.data, call.config) + Metric::L2 => partition::(arch, call.data, Metric::L2, call.config), + Metric::Cosine => { + partition::(arch, call.data, Metric::Cosine, call.config) } + Metric::CosineNormalized => partition::( + arch, + call.data, + Metric::CosineNormalized, + call.config, + ), + Metric::InnerProduct => partition::( + arch, + call.data, + Metric::InnerProduct, + call.config, + ), } } } @@ -1102,6 +1128,7 @@ mod tests { let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, + Metric::L2, &[2], &[0, 1], 1, @@ -1156,13 +1183,13 @@ mod tests { let pool = StripeBufferPool::new((), 0, None); let points = { let mut buffers = pool.get_ref(()); - buffers.points.resize(16, 0.0); - buffers.points.as_ptr() + buffers.point_values.resize(16, 0.0); + buffers.point_values.as_ptr() }; let buffers = pool.get_ref(()); - assert_eq!(buffers.points.as_ptr(), points); - assert_eq!(buffers.points.len(), 16); + assert_eq!(buffers.point_values.as_ptr(), points); + assert_eq!(buffers.point_values.len(), 16); } #[test] @@ -1175,6 +1202,7 @@ mod tests { let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, + Metric::L2, &point_ids, &[0, 2_047], 1, From 49827c36be20872d907ff00663eed72487b15a39 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:42:04 +0000 Subject: [PATCH 49/58] refactor(pipnn): dispatch norm preparation through metric types --- diskann/src/graph/pipnn/leaf_build.rs | 50 +++--------- diskann/src/graph/pipnn/mod.rs | 16 ++-- diskann/src/graph/pipnn/partitioning.rs | 101 +++++++----------------- 3 files changed, 47 insertions(+), 120 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 5bca5b6778..273ca54f2c 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -22,12 +22,11 @@ use std::{collections::TryReserveError, sync::Mutex}; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_vector::distance::Metric; use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::{LeafKernelMetric, norm_from_squared}, + kernel_metric::LeafMetric, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, @@ -235,14 +234,13 @@ pub(super) fn build_leaf_candidates( data: MatrixView<'_, T>, leaves: Vec>, requested_k: usize, - metric: Metric, ) -> Result>, LeafBuildError> where A: Architecture, A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: LeafKernelMetric, + M: LeafMetric, T: VectorRepr + 'static, { if data.ncols() == 0 { @@ -259,7 +257,6 @@ where add_direct_leaf_candidates::( arch, data, - metric, leaf, point_ids, requested_k, @@ -280,7 +277,6 @@ where fn add_direct_leaf_candidates( arch: A, data: MatrixView<'_, T>, - metric: Metric, leaf: usize, point_ids: &[u32], requested_k: usize, @@ -292,7 +288,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: LeafKernelMetric, + M: LeafMetric, T: VectorRepr + 'static, { if point_ids.is_empty() { @@ -351,21 +347,9 @@ where leaf, buffer: "leaf dot-product matrix", })?; - let norm_count = if matches!(metric, Metric::L2 | Metric::Cosine) { - point_ids.len() - } else { - 0 - }; - grow("leaf norms", &mut buffers.norms, norm_count, 0.0)?; - for (point, norm) in buffers.norms[..norm_count].iter_mut().enumerate() { - let squared_norm = dots[(point, point)]; - *norm = if metric == Metric::Cosine { - norm_from_squared(squared_norm) - } else { - squared_norm - }; - } - let norms = &buffers.norms[..norm_count]; + M::prepare_norms(dots, &mut buffers.norms) + .map_err(|source| allocation_error("leaf norms", point_ids.len(), source))?; + let norms = &*buffers.norms; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), @@ -507,33 +491,23 @@ mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => build_leaf_candidates::( - arch, - call.data, - call.leaves, - call.k, - Metric::L2, - ), - Metric::Cosine => build_leaf_candidates::( - arch, - call.data, - call.leaves, - call.k, - Metric::Cosine, - ), + Metric::L2 => { + build_leaf_candidates::(arch, call.data, call.leaves, call.k) + } + Metric::Cosine => { + build_leaf_candidates::(arch, call.data, call.leaves, call.k) + } Metric::CosineNormalized => build_leaf_candidates::( arch, call.data, call.leaves, call.k, - Metric::CosineNormalized, ), Metric::InnerProduct => build_leaf_candidates::( arch, call.data, call.leaves, call.k, - Metric::InnerProduct, ), } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 2d7e088c7e..8b3f11243b 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -54,7 +54,7 @@ use diskann_wide::{ use rayon::ThreadPool; use self::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, LeafKernelMetric, PartitionKernelMetric, + Cosine, CosineNormalized, InnerProduct, L2, LeafMetric, PartitionMetric, }; /// PiPNN partition and leaf-selection policy. @@ -264,22 +264,16 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: LeafKernelMetric + PartitionKernelMetric, + M: LeafMetric + PartitionMetric, T: VectorRepr + Send + Sync + 'static, { let leaves = tracing::info_span!("pipnn.partition") - .in_scope(|| partitioning::partition::(arch, data, metric, &context.config))?; + .in_scope(|| partitioning::partition::(arch, data, &context.config))?; // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, // so its complete allocation drops when leaf construction returns. let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates::( - arch, - data, - leaves, - context.config.leaf_k, - metric, - ) - .map_err(ANNError::new) + leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) + .map_err(ANNError::new) })?; // Finalization consumes each candidate list. It reuses that list's allocation // for the final adjacency when the graph policy permits it. diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 4254978205..cdd8849919 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -24,14 +24,13 @@ use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, }; -use diskann_vector::{Norm, distance::Metric, norm::FastL2NormSquared}; use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::{PartitionKernelMetric, norm_from_squared}, + kernel_metric::PartitionMetric, partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, }; @@ -119,7 +118,6 @@ type StripeBufferPool = ObjectPool; pub(super) fn partition( arch: A, data: MatrixView<'_, T>, - metric: Metric, config: &PiPNNConfig, ) -> ANNResult>> where @@ -127,7 +125,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr + Send + Sync, { let points = data.nrows(); @@ -146,7 +144,7 @@ where for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = - partition_replica::(arch, data, metric, config, seed, &stripe_buffers)?; + partition_replica::(arch, data, config, seed, &stripe_buffers)?; leaves .try_reserve(replica_leaves.len()) .map_err(ANNError::new)?; @@ -162,7 +160,6 @@ where fn partition_replica( arch: A, data: MatrixView<'_, T>, - metric: Metric, config: &PiPNNConfig, seed: u64, stripe_buffers: &StripeBufferPool, @@ -172,7 +169,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr + Send + Sync, { let initial_indices = point_ids(data.nrows())?; @@ -212,7 +209,6 @@ where *slot = Some(partition_work_item::( arch, data, - metric, config, item, stripe_buffers, @@ -254,7 +250,6 @@ where fn partition_work_item( arch: A, data: MatrixView<'_, T>, - metric: Metric, config: &PiPNNConfig, item: WorkItem, stripe_buffers: &StripeBufferPool, @@ -264,7 +259,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr + Send + Sync, { let points = item.indices.len(); @@ -274,15 +269,8 @@ where config.p_samp, mix_seed(item.seed, points as u64), )?; - let clusters = assign_to_leaders::( - arch, - data, - metric, - &item.indices, - &leaders, - fanout, - stripe_buffers, - )?; + let clusters = + assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; let mut pending = Vec::new(); let mut finished = Vec::new(); @@ -348,7 +336,6 @@ fn mix_seed(seed: u64, salt: u64) -> u64 { fn assign_to_leaders( arch: A, data: MatrixView<'_, T>, - metric: Metric, point_ids: &[u32], leader_ids: &[u32], fanout: usize, @@ -359,7 +346,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr + Send + Sync, { let dimension_count = data.ncols(); @@ -367,22 +354,18 @@ where let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; gather_vectors(data, leader_ids, &mut leader_values)?; - let mut leader_norm_values = if matches!(metric, Metric::L2 | Metric::Cosine) { - filled_vec(leader_ids.len(), 0.0f32)? - } else { - Vec::new() - }; - for (norm_value, leader_vector) in leader_norm_values - .iter_mut() - .zip(leader_values.chunks_exact(dimension_count)) - { - // Leader norms affect top-k order. Use this scalar reduction order. - // SIMD reassociation changes low bits and can change a near-tie branch. - *norm_value = leader_vector.iter().map(|value| value * value).sum(); - if metric == Metric::Cosine { - *norm_value = norm_from_squared(*norm_value); - } - } + let leader_matrix = + MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count).map_err( + |error| { + ANNError::new(PartitionError::InvalidBufferLength { + buffer: "leader data", + expected: leader_values_len, + actual: error.into_inner().len(), + }) + }, + )?; + let mut leader_norm_values = Vec::new(); + M::prepare_leader_norms(leader_matrix, &mut leader_norm_values).map_err(ANNError::new)?; let fanout = fanout.min(leader_ids.len()); let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; @@ -412,7 +395,6 @@ where assign_point_stripe::( arch, data, - metric, &point_ids[first_point..first_point + stripe_point_count], &leader_values, &leader_norm_values, @@ -436,7 +418,6 @@ where fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, - metric: Metric, point_ids: &[u32], leader_values: &[f32], leader_norm_values: &[f32], @@ -449,7 +430,7 @@ where A::f32x16: std::ops::Div, ::Mask: SIMDSelect, u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, - M: PartitionKernelMetric, + M: PartitionMetric, T: VectorRepr, { let point_count = point_ids.len(); @@ -496,19 +477,8 @@ where ) .map_err(ANNError::new)?; - let point_norms = if metric == Metric::Cosine { - grow_fallible(point_norm_buffer, point_count, 0.0)?; - let point_norms = &mut point_norm_buffer[..point_count]; - for (norm_value, point_values) in point_norms - .iter_mut() - .zip(point_matrix.as_slice().chunks_exact(dimensions)) - { - *norm_value = norm_from_squared(FastL2NormSquared.evaluate(point_values)); - } - &*point_norms - } else { - &[] - }; + M::prepare_point_norms(point_matrix.as_view(), point_norm_buffer).map_err(ANNError::new)?; + let point_norms = &*point_norm_buffer; let norms = PartitionNorms { point_norms, leader_norms: leader_norm_values, @@ -529,7 +499,6 @@ where })?; nearest_leaders::( arch, - metric, PartitionInput { dots, norms }, output, kernel_workspace, @@ -832,22 +801,14 @@ mod tests { use super::super::kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2}; match self.0 { - Metric::L2 => partition::(arch, call.data, Metric::L2, call.config), - Metric::Cosine => { - partition::(arch, call.data, Metric::Cosine, call.config) + Metric::L2 => partition::(arch, call.data, call.config), + Metric::Cosine => partition::(arch, call.data, call.config), + Metric::CosineNormalized => { + partition::(arch, call.data, call.config) + } + Metric::InnerProduct => { + partition::(arch, call.data, call.config) } - Metric::CosineNormalized => partition::( - arch, - call.data, - Metric::CosineNormalized, - call.config, - ), - Metric::InnerProduct => partition::( - arch, - call.data, - Metric::InnerProduct, - call.config, - ), } } } @@ -1128,7 +1089,6 @@ mod tests { let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, - Metric::L2, &[2], &[0, 1], 1, @@ -1202,7 +1162,6 @@ mod tests { let clusters = assign_to_leaders::<_, super::super::kernel_metric::L2, _>( diskann_wide::ARCH, data, - Metric::L2, &point_ids, &[0, 2_047], 1, From 33f29f803bb66c6f90ad94525bbd065510fd346c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:22:00 +0000 Subject: [PATCH 50/58] refactor(pipnn): pass metric-owned ranking context --- diskann/src/graph/pipnn/leaf_build.rs | 9 ++++--- diskann/src/graph/pipnn/partitioning.rs | 33 ++++++++++++++----------- 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 273ca54f2c..603ef78ae7 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -26,7 +26,7 @@ use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::LeafMetric, + kernel_metric::{LeafMetric, NormPreparation}, leaf_kernel::{ LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, }, @@ -347,8 +347,11 @@ where leaf, buffer: "leaf dot-product matrix", })?; - M::prepare_norms(dots, &mut buffers.norms) - .map_err(|source| allocation_error("leaf norms", point_ids.len(), source))?; + M::prepare_norms(NormPreparation { + values: dots, + norms: &mut buffers.norms, + }) + .map_err(|source| allocation_error("leaf norms", point_ids.len(), source))?; let norms = &*buffers.norms; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index cdd8849919..052ec36608 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -30,8 +30,8 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::PartitionMetric, - partition_kernel::{PartitionInput, PartitionKernelWorkspace, PartitionNorms, nearest_leaders}, + kernel_metric::{NormPreparation, PartitionMetric, PartitionNorms}, + partition_kernel::{PartitionInput, nearest_leaders}, }; // These constants control internal batching and deterministic seed generation. @@ -89,7 +89,7 @@ struct StripeBuffers { point_values: Vec, dot_values: Vec, point_norms: Vec, - kernel_workspace: PartitionKernelWorkspace, + ranked_leaders: Vec<(u32, f32)>, } impl AsPooled<()> for StripeBuffers { @@ -365,7 +365,11 @@ where }, )?; let mut leader_norm_values = Vec::new(); - M::prepare_leader_norms(leader_matrix, &mut leader_norm_values).map_err(ANNError::new)?; + M::prepare_leader_norms(NormPreparation { + values: leader_matrix, + norms: &mut leader_norm_values, + }) + .map_err(ANNError::new)?; let fanout = fanout.min(leader_ids.len()); let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; @@ -447,7 +451,7 @@ where point_values: point_buffer, dot_values: dot_buffer, point_norms: point_norm_buffer, - kernel_workspace, + ranked_leaders, } = buffers; let mut point_matrix = MutMatrixView::try_from( &mut point_buffer[..point_values_len], @@ -477,7 +481,11 @@ where ) .map_err(ANNError::new)?; - M::prepare_point_norms(point_matrix.as_view(), point_norm_buffer).map_err(ANNError::new)?; + M::prepare_point_norms(NormPreparation { + values: point_matrix.as_view(), + norms: point_norm_buffer, + }) + .map_err(ANNError::new)?; let point_norms = &*point_norm_buffer; let norms = PartitionNorms { point_norms, @@ -497,13 +505,8 @@ where actual: error.into_inner().len(), }) })?; - nearest_leaders::( - arch, - PartitionInput { dots, norms }, - output, - kernel_workspace, - ) - .map_err(ANNError::new) + nearest_leaders::(arch, PartitionInput { dots, norms }, output, ranked_leaders) + .map_err(ANNError::new) } fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> @@ -1063,7 +1066,7 @@ mod tests { } #[test] - fn l2_leader_norms_preserve_scalar_reduction_order() { + fn l2_leader_norms_preserve_sequential_reduction_order() { fn next(state: &mut u64) -> f32 { *state ^= *state << 13; *state ^= *state >> 7; @@ -1071,7 +1074,7 @@ mod tests { (((*state >> 40) as f32 / 8_388_608.0) - 1.0) * 1_000.0 } - // Scalar and SIMD-reassociated leader norms select different top-1 + // Sequential and SIMD-reassociated leader norms select different top-1 // leaders for this case. Dot products still use the production GEMM. The // test changes only the leader-norm reduction. let dimensions = 129; From 75798ff435d2267db7afc81c0396804f647101f8 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:07:09 +0000 Subject: [PATCH 51/58] fix(pipnn): reset reused prune states --- diskann/src/graph/pipnn/finalization.rs | 38 +++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 55d86103e6..62a3bb9fbc 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -138,6 +138,9 @@ where workspace .prune_states .resize(workspace.candidate_slots.len(), prune::State::default()); + // Each candidate list starts a separate RobustPrune state machine. + // Reset retained entries because resize initializes only new entries. + workspace.prune_states.fill(prune::State::default()); let selected = prune::robust_prune( &sorted, @@ -247,6 +250,41 @@ mod tests { assert_eq!(&*actual[0], &[1, 3]); } + #[test] + fn reused_workspace_matches_fresh_pruning() { + let data = [0.0_f32, 1.0, 2.0, -3.0, 4.0]; + let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); + let first = [3, 2, 1]; + let second = [4, 3, 2]; + let candidates = |first: &[u32], second: &[u32]| { + vec![ + candidate_list(first.iter().copied()), + candidate_list(second.iter().copied()), + candidate_list([]), + candidate_list([]), + candidate_list([]), + ] + }; + let graph = graph_config(2); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + let fresh_first = pool + .install(|| prune_overfull(data, candidates(&first, &[]), &graph, Metric::L2)) + .unwrap(); + let fresh_second = pool + .install(|| prune_overfull(data, candidates(&[], &second), &graph, Metric::L2)) + .unwrap(); + + let reused = pool + .install(|| prune_overfull(data, candidates(&first, &second), &graph, Metric::L2)) + .unwrap(); + + assert_eq!(&*reused[0], &*fresh_first[0]); + assert_eq!(&*reused[1], &*fresh_second[1]); + } + #[test] fn rejects_invalid_candidate_ids_without_panicking() { let data = [0.0_f32, 1.0, 2.0]; From 7e15798f54fb34586aabf547eb77dae2ca622703 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 11 Aug 2026 16:05:12 +0000 Subject: [PATCH 52/58] refactor(pipnn): prepare leaf norms in metric policy --- diskann/src/graph/pipnn/leaf_build.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 603ef78ae7..644e05a1c2 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -347,7 +347,7 @@ where leaf, buffer: "leaf dot-product matrix", })?; - M::prepare_norms(NormPreparation { + M::prepare_leaf_norms(NormPreparation { values: dots, norms: &mut buffers.norms, }) From dc7a40bd1730061a37a1ce8efffcdfd94e4700f6 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:20:46 +0000 Subject: [PATCH 53/58] refactor(pipnn): simplify graph build stages Delegate GEMM, norm preparation, and ranking to numerical kernels. Keep partition and leaf stages focused on IDs, scheduling, and graph mapping. --- Cargo.lock | 1 + diskann/Cargo.toml | 4 +- diskann/src/graph/pipnn/finalization.rs | 125 +------ diskann/src/graph/pipnn/leaf_build.rs | 373 +++---------------- diskann/src/graph/pipnn/mod.rs | 22 +- diskann/src/graph/pipnn/partitioning.rs | 458 +++++------------------- 6 files changed, 162 insertions(+), 821 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index be5f9ef094..72af940896 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -446,6 +446,7 @@ dependencies = [ "half", "hashbrown 0.16.1", "num-traits", + "parking_lot", "pin-project", "rand", "rayon", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 2568ac6669..d8dcf211b8 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -22,6 +22,7 @@ half = { workspace = true, features = ["bytemuck", "num-traits"] } # while other crates use default-features = true. Keeping version 0.16.0 consistent. hashbrown = { version = "0.16.0", default-features = false, features = ["default-hasher"] } num-traits.workspace = true +parking_lot = { version = "0.12.5", optional = true } rand.workspace = true rayon = { workspace = true, optional = true } thiserror.workspace = true @@ -32,7 +33,6 @@ diskann-wide = { workspace = true } # Optional Dependencies dashmap = { workspace = true, optional = true } -diskann-linalg = { workspace = true, optional = true } [dev-dependencies] futures-util = { workspace = true, default-features = false } @@ -61,7 +61,7 @@ panic = "warn" default = ["tracing"] # Enable PiPNN batch graph construction. -pipnn = ["dep:diskann-linalg", "dep:rayon", "tracing"] +pipnn = ["dep:diskann-linalg", "dep:parking_lot", "dep:rayon", "tracing"] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 62a3bb9fbc..6f01c3dd25 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -31,32 +31,24 @@ use rayon::prelude::*; #[derive(Debug, thiserror::Error)] pub(crate) enum FinalizationError { - #[error("candidate list count {lists} does not match the dataset point count {points}")] - CandidateListCountMismatch { lists: usize, points: usize }, - #[error( - "candidate ID {candidate} for source {source_index} is outside a {points}-point dataset" - )] - InvalidCandidateId { - source_index: usize, - candidate: u32, - points: usize, - }, #[error("candidate count {actual} exceeds the u16 position limit {max}")] TooManyCandidates { actual: usize, max: usize }, } /// RobustPrune state for one Rayon job. /// -/// `candidate_slots` and `prune_states` stay positionally aligned with +/// `sorted_cache` and `prune_states` stay positionally aligned with /// `sorted_candidates`. #[derive(Default)] struct PruneWorkspace { sorted_candidates: Vec>, - candidate_slots: Vec<(f32, Option)>, + sorted_cache: Vec<(f32, Option)>, prune_states: Vec, } -/// Check candidate IDs and prune each list that exceeds the graph degree. +/// Prune each candidate list that exceeds the graph degree. +/// +/// Candidate builders supply one list per data row and valid dataset IDs. pub(crate) fn prune_overfull( data: MatrixView<'_, T>, candidates: Vec>, @@ -66,8 +58,6 @@ pub(crate) fn prune_overfull( where T: VectorRepr + Send + Sync, { - validate_candidate_lists(&candidates, data.nrows()).map_err(ANNError::new)?; - let degree = graph.pruned_degree().get(); let distance = T::distance(metric, Some(data.ncols())); @@ -88,10 +78,6 @@ where let source_id = u32::try_from(source).map_err(ANNError::new)?; let source_vector = data.row(source); workspace.sorted_candidates.clear(); - workspace - .sorted_candidates - .try_reserve(source_candidates.len()) - .map_err(ANNError::new)?; workspace .sorted_candidates .extend(source_candidates.iter().copied().map(|candidate| { @@ -109,42 +95,25 @@ where max: u16::MAX as usize, })); } - workspace.candidate_slots.clear(); - workspace - .candidate_slots - .try_reserve(candidate_count) - .map_err(ANNError::new)?; - + workspace.sorted_cache.clear(); // Sort all candidates before the code marks a self-edge as absent. - // Thus, self-edge removal cannot add a farther candidate. The - // `SortedNeighbors` value carries this order into RobustPrune. + // Thus, self-edge removal cannot add a farther candidate. Cache + // construction preserves this order for RobustPrune. let sorted = SortedNeighbors::new(&mut workspace.sorted_candidates, candidate_count); - workspace - .candidate_slots - .extend(sorted.iter().map(|neighbor| { - let id = *neighbor.id(); - (*neighbor.distance(), (id != source_id).then_some(id)) - })); + workspace.sorted_cache.extend(sorted.iter().map(|neighbor| { + let id = *neighbor.id(); + (*neighbor.distance(), (id != source_id).then_some(id)) + })); workspace .prune_states - .try_reserve( - workspace - .candidate_slots - .len() - .saturating_sub(workspace.prune_states.len()), - ) - .map_err(ANNError::new)?; - workspace - .prune_states - .resize(workspace.candidate_slots.len(), prune::State::default()); + .resize(workspace.sorted_cache.len(), prune::State::default()); // Each candidate list starts a separate RobustPrune state machine. // Reset retained entries because resize initializes only new entries. workspace.prune_states.fill(prune::State::default()); let selected = prune::robust_prune( - &sorted, - &workspace.candidate_slots, + &workspace.sorted_cache, workspace.prune_states.as_mut_slice(), degree, graph.alpha(), @@ -168,28 +137,6 @@ where .collect() } -fn validate_candidate_lists( - candidates: &[AdjacencyList], - points: usize, -) -> Result<(), FinalizationError> { - if candidates.len() != points { - return Err(FinalizationError::CandidateListCountMismatch { - lists: candidates.len(), - points, - }); - } - for (source, source_candidates) in candidates.iter().enumerate() { - if let Some(&candidate) = source_candidates.iter().find(|&&id| id as usize >= points) { - return Err(FinalizationError::InvalidCandidateId { - source_index: source, - candidate, - points, - }); - } - } - Ok(()) -} - #[cfg(test)] mod tests { use crate::graph::{ @@ -284,48 +231,4 @@ mod tests { assert_eq!(&*reused[0], &*fresh_first[0]); assert_eq!(&*reused[1], &*fresh_second[1]); } - - #[test] - fn rejects_invalid_candidate_ids_without_panicking() { - let data = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![ - candidate_list([1, 3]), - candidate_list([]), - candidate_list([]), - ]; - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::InvalidCandidateId { - source_index: 0, - candidate: 3, - points: 3, - }) - )); - } - - #[test] - fn rejects_candidate_list_count_mismatch_without_panicking() { - let data = [0.0_f32, 1.0, 2.0]; - let data = MatrixView::try_from(&data[..], 3, 1).unwrap(); - let candidates = vec![ - candidate_list([]), - candidate_list([]), - candidate_list([]), - candidate_list([]), - ]; - - let error = prune_overfull(data, candidates, &graph_config(1), Metric::L2).unwrap_err(); - - assert!(matches!( - error.downcast_ref::(), - Some(FinalizationError::CandidateListCountMismatch { - lists: 4, - points: 3 - }) - )); - } } diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 644e05a1c2..5f02967125 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -8,17 +8,16 @@ //! Partitioning supplies sorted, unique global point IDs for each leaf. One leaf //! job does these steps: //! -//! 1. Check each ID and convert its vector to reusable `f32` storage. -//! 2. Compute the lower triangle of `A · Aᵀ`. -//! 3. Select local neighbors for both points of each pair. -//! 4. Convert local positions to global point IDs. -//! 5. Add both edge directions to global candidate lists. +//! 1. Gather each ID and convert its vector to reusable `f32` storage. +//! 2. Call the leaf kernel for Gram construction, norms, and local ranking. +//! 3. Convert local positions to global point IDs. +//! 4. Add both edge directions to global candidate lists. //! //! Overlapping leaves run concurrently. A worker locks one destination list only //! while it adds one leaf's IDs. Reusable buffers keep their largest allocation. //! Each operation uses an explicit active prefix. -use std::{collections::TryReserveError, sync::Mutex}; +use parking_lot::Mutex; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; @@ -26,31 +25,13 @@ use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ - kernel_metric::{LeafMetric, NormPreparation}, - leaf_kernel::{ - LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_neighbors, - }, + kernel_metric::LeafMetric, + leaf_kernel::{LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, select_leaf_neighbors}, }; /// Failure while converting leaves into direct graph candidates. #[derive(Debug, thiserror::Error)] pub(crate) enum LeafBuildError { - #[error("leaf build requires at least one dimension")] - EmptyDimensions, - #[error("dataset point count {0} exceeds the u32 ID limit")] - TooManyPoints(usize), - #[error("leaf {leaf} is empty")] - EmptyLeaf { leaf: usize }, - #[error("point ID {point} in leaf {leaf} is outside a {points}-point dataset")] - InvalidPointId { - leaf: usize, - point: u32, - points: usize, - }, - #[error("point ID {point} appears more than once in leaf {leaf}")] - DuplicatePointId { leaf: usize, point: u32 }, - #[error("point IDs in leaf {leaf} are not strictly increasing")] - UnsortedPointIds { leaf: usize }, #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] ShapeOverflow { leaf: usize, @@ -59,13 +40,6 @@ pub(crate) enum LeafBuildError { }, #[error("failed to form {buffer} view for leaf {leaf}")] InvalidView { leaf: usize, buffer: &'static str }, - #[error("failed to reserve {additional} values for {buffer}")] - Allocation { - buffer: &'static str, - additional: usize, - #[source] - source: TryReserveError, - }, #[error("failed to convert point {point} in leaf {leaf}")] Conversion { leaf: usize, @@ -73,22 +47,12 @@ pub(crate) enum LeafBuildError { #[source] source: crate::ANNError, }, - #[error("lower-AAT failed for leaf {leaf}")] - LowerAat { - leaf: usize, - #[source] - source: diskann_linalg::SgemmError, - }, #[error("nearest-neighbor selection failed for leaf {leaf}")] Kernel { leaf: usize, #[source] - source: LeafKernelError, + source: crate::ANNError, }, - #[error("leaf kernel returned local target {target} for a {points}-point leaf")] - InvalidLocalTarget { target: u32, points: usize }, - #[error("candidate list for point {point} is poisoned")] - PoisonedCandidateList { point: u32 }, } /// Reusable buffers for one Rayon leaf job. @@ -98,8 +62,6 @@ pub(crate) enum LeafBuildError { #[derive(Default)] struct LeafBuffers { point_values: Vec, - dots: Vec, - norms: Vec, neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, @@ -121,16 +83,7 @@ impl LeafBuffers { rows: point_count, columns: dimension_count, })?; - let dot_count = - point_count - .checked_mul(point_count) - .ok_or(LeafBuildError::ShapeOverflow { - leaf, - rows: point_count, - columns: point_count, - })?; - let leaf_k = leaf_neighbor_count(point_count, requested_k) - .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + let leaf_k = leaf_neighbor_count(point_count, requested_k); let neighbor_count = point_count .checked_mul(leaf_k) @@ -140,33 +93,17 @@ impl LeafBuffers { columns: leaf_k, })?; - grow( - "leaf point values", - &mut self.point_values, - point_value_count, - 0.0, - )?; - grow("leaf dot products", &mut self.dots, dot_count, 0.0)?; - grow( - "leaf neighbors", - &mut self.neighbors, - neighbor_count, - LeafNeighbor::default(), - )?; + grow(&mut self.point_values, point_value_count, 0.0); + grow(&mut self.neighbors, neighbor_count, LeafNeighbor::default()); Ok((leaf_k, neighbor_count)) } - fn prepare_local_adjacency(&mut self, point_count: usize) -> Result<(), LeafBuildError> { - let additional = point_count.saturating_sub(self.local_adjacency.len()); - self.local_adjacency - .try_reserve(additional) - .map_err(|source| allocation_error("leaf adjacency lists", additional, source))?; + fn prepare_local_adjacency(&mut self, point_count: usize) { self.local_adjacency .resize_with(point_count, AdjacencyList::new); self.local_adjacency[..point_count] .iter_mut() .for_each(AdjacencyList::clear); - Ok(()) } } @@ -180,47 +117,31 @@ struct DirectCandidates { } impl DirectCandidates { - fn new(point_count: usize) -> Result { - let mut lists = Vec::new(); - lists - .try_reserve_exact(point_count) - .map_err(|source| allocation_error("candidate lists", point_count, source))?; - lists.resize_with(point_count, || Mutex::new(AdjacencyList::new())); - Ok(Self { lists }) + fn new(point_count: usize) -> Self { + let lists = (0..point_count) + .map(|_| Mutex::new(AdjacencyList::new())) + .collect(); + Self { lists } } - fn add_leaf( - &self, - point_ids: &[u32], - local_adjacency: &[AdjacencyList], - ) -> Result<(), LeafBuildError> { + fn add_leaf(&self, point_ids: &[u32], local_adjacency: &[AdjacencyList]) { for (&source, additions) in point_ids.iter().zip(local_adjacency) { // `add_direct_leaf_candidates` checks every point ID before this append. - let candidates = &self.lists[source as usize]; - let mut candidates = candidates + self.lists[source as usize] .lock() - .map_err(|_| LeafBuildError::PoisonedCandidateList { point: source })?; - candidates.extend_from_slice(additions); + .extend_from_slice(additions); } - Ok(()) } - fn into_lists(self) -> Result>, LeafBuildError> { - let mut output = Vec::new(); - output - .try_reserve_exact(self.lists.len()) - .map_err(|source| allocation_error("candidate output", self.lists.len(), source))?; - for (point, candidates) in self.lists.into_iter().enumerate() { - let mut candidates = + fn into_lists(self) -> Vec> { + self.lists + .into_iter() + .map(Mutex::into_inner) + .map(|mut candidates| { + candidates.sort(); candidates - .into_inner() - .map_err(|_| LeafBuildError::PoisonedCandidateList { - point: point as u32, - })?; - candidates.sort(); - output.push(candidates); - } - Ok(output) + }) + .collect() } } @@ -243,14 +164,7 @@ where M: LeafMetric, T: VectorRepr + 'static, { - if data.ncols() == 0 { - return Err(LeafBuildError::EmptyDimensions); - } - if data.nrows() > u32::MAX as usize { - return Err(LeafBuildError::TooManyPoints(data.nrows())); - } - - let candidates = DirectCandidates::new(data.nrows())?; + let candidates = DirectCandidates::new(data.nrows()); leaves.par_iter().enumerate().try_for_each_init( LeafBuffers::default, |buffers, (leaf, point_ids)| { @@ -265,12 +179,11 @@ where ) }, )?; - candidates.into_lists() + Ok(candidates.into_lists()) } /// Add one leaf's symmetric neighbors to the direct candidate lists. /// -/// The function rejects empty, duplicate, unsorted, or out-of-range point IDs. /// Reusable buffers can be longer than this leaf, so all accesses use the current /// leaf shape. #[allow(clippy::too_many_arguments)] @@ -291,27 +204,6 @@ where M: LeafMetric, T: VectorRepr + 'static, { - if point_ids.is_empty() { - return Err(LeafBuildError::EmptyLeaf { leaf }); - } - for &point in point_ids { - if point as usize >= data.nrows() { - return Err(LeafBuildError::InvalidPointId { - leaf, - point, - points: data.nrows(), - }); - } - } - if let Some(pair) = point_ids.windows(2).find(|pair| pair[0] >= pair[1]) { - if pair[0] == pair[1] { - return Err(LeafBuildError::DuplicatePointId { - leaf, - point: pair[0], - }); - } - return Err(LeafBuildError::UnsortedPointIds { leaf }); - } let (leaf_k, neighbor_value_count) = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; if leaf_k == 0 { @@ -319,11 +211,11 @@ where } let point_value_count = point_ids.len() * data.ncols(); - let dot_count = point_ids.len() * point_ids.len(); + let point_values = &mut buffers.point_values[..point_value_count]; for (&point, point_output) in point_ids .iter() - .zip(buffers.point_values[..point_value_count].chunks_exact_mut(data.ncols())) + .zip(point_values.chunks_exact_mut(data.ncols())) { let source_values = data.row(point as usize); T::as_f32_into(source_values, point_output).map_err(|source| { @@ -335,24 +227,13 @@ where })?; } - diskann_linalg::sgemm_aat_lower( - point_ids.len(), - data.ncols(), - &buffers.point_values[..point_value_count], - &mut buffers.dots[..dot_count], - ) - .map_err(|source| LeafBuildError::LowerAat { leaf, source })?; - let dots = MatrixView::try_from(&buffers.dots[..dot_count], point_ids.len(), point_ids.len()) - .map_err(|_| LeafBuildError::InvalidView { - leaf, - buffer: "leaf dot-product matrix", - })?; - M::prepare_leaf_norms(NormPreparation { - values: dots, - norms: &mut buffers.norms, - }) - .map_err(|source| allocation_error("leaf norms", point_ids.len(), source))?; - let norms = &*buffers.norms; + let points = + MatrixView::try_from(&*point_values, point_ids.len(), data.ncols()).map_err(|_| { + LeafBuildError::InvalidView { + leaf, + buffer: "leaf point matrix", + } + })?; let output = MutMatrixView::try_from( &mut buffers.neighbors[..neighbor_value_count], point_ids.len(), @@ -362,17 +243,18 @@ where leaf, buffer: "leaf output", })?; - nearest_neighbors::(arch, dots, norms, output, &mut buffers.kernel_workspace) + select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - buffers.prepare_local_adjacency(point_ids.len())?; + buffers.prepare_local_adjacency(point_ids.len()); add_symmetric_neighbors( point_ids, leaf_k, &buffers.neighbors[..neighbor_value_count], &mut buffers.local_adjacency[..point_ids.len()], - )?; - candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]) + ); + candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]); + Ok(()) } fn add_symmetric_neighbors( @@ -380,61 +262,23 @@ fn add_symmetric_neighbors( leaf_k: usize, neighbors: &[LeafNeighbor], local_adjacency: &mut [AdjacencyList], -) -> Result<(), LeafBuildError> { +) { for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { for neighbor in source_neighbors { let target = neighbor.target as usize; - let Some(&target_id) = point_ids.get(target) else { - return Err(LeafBuildError::InvalidLocalTarget { - target: neighbor.target, - points: point_ids.len(), - }); - }; let source_id = point_ids[source]; + let target_id = point_ids[target]; if source_id != target_id { local_adjacency[source].push(target_id); local_adjacency[target].push(source_id); } } } - Ok(()) } -fn grow( - buffer: &'static str, - values: &mut Vec, - len: usize, - value: T, -) -> Result<(), LeafBuildError> { +fn grow(values: &mut Vec, len: usize, value: T) { if values.len() < len { - resize(buffer, values, len, value)?; - } - Ok(()) -} - -fn resize( - buffer: &'static str, - values: &mut Vec, - len: usize, - value: T, -) -> Result<(), LeafBuildError> { - let additional = len.saturating_sub(values.len()); - values - .try_reserve(additional) - .map_err(|source| allocation_error(buffer, additional, source))?; - values.resize(len, value); - Ok(()) -} - -fn allocation_error( - buffer: &'static str, - additional: usize, - source: TryReserveError, -) -> LeafBuildError { - LeafBuildError::Allocation { - buffer, - additional, - source, + values.resize(len, value); } } @@ -450,7 +294,7 @@ mod tests { use std::collections::BTreeSet; use super::{ - DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, + DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, build_leaf_candidates, }; @@ -724,44 +568,6 @@ mod tests { } } - #[test] - fn rejects_invalid_dimensions_and_leaf_membership() { - let data = [0.0_f32, 1.0]; - let no_dimensions = MatrixView::try_from(&data[..0], 2, 0).unwrap(); - assert!(matches!( - build(no_dimensions, &[], 1, Metric::L2), - Err(LeafBuildError::EmptyDimensions) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![]], 1, Metric::L2), - Err(LeafBuildError::EmptyLeaf { leaf: 0 }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 2]], 1, Metric::L2), - Err(LeafBuildError::InvalidPointId { - leaf: 0, - point: 2, - points: 2 - }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![2]], 1, Metric::L2), - Err(LeafBuildError::InvalidPointId { point: 2, .. }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 2]], 0, Metric::L2), - Err(LeafBuildError::InvalidPointId { point: 2, .. }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![0, 0]], 1, Metric::L2), - Err(LeafBuildError::DuplicatePointId { leaf: 0, point: 0 }) - )); - assert!(matches!( - build(view(&data, 2, 1), &[vec![1, 0]], 1, Metric::L2), - Err(LeafBuildError::UnsortedPointIds { leaf: 0 }) - )); - } - #[test] fn singleton_and_zero_k_leaves_add_no_candidates() { let data = [0.0_f32, 1.0, 2.0]; @@ -786,16 +592,13 @@ mod tests { let mut buffers = LeafBuffers::default(); buffers.prepare(0, 64, 128, 2).unwrap(); let point_values = buffers.point_values.as_ptr(); - let dots = buffers.dots.as_ptr(); let neighbors = buffers.neighbors.as_ptr(); buffers.prepare(1, 8, 128, 2).unwrap(); assert_eq!(buffers.point_values.as_ptr(), point_values); - assert_eq!(buffers.dots.as_ptr(), dots); assert_eq!(buffers.neighbors.as_ptr(), neighbors); assert_eq!(buffers.point_values.len(), 64 * 128); - assert_eq!(buffers.dots.len(), 64 * 64); assert_eq!(buffers.neighbors.len(), 64 * 2); } @@ -808,28 +611,6 @@ mod tests { )); } - #[test] - fn rejects_an_invalid_kernel_target() { - let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; - let error = add_symmetric_neighbors( - &[10, 20], - 1, - &[ - super::super::leaf_kernel::LeafNeighbor::new(9, 1.0), - super::super::leaf_kernel::LeafNeighbor::new(0, 1.0), - ], - &mut graph, - ) - .unwrap_err(); - assert!(matches!( - error, - LeafBuildError::InvalidLocalTarget { - target: 9, - points: 2 - } - )); - } - #[test] fn skips_duplicate_global_ids_without_self_edges() { let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; @@ -841,58 +622,20 @@ mod tests { super::super::leaf_kernel::LeafNeighbor::new(0, 0.0), ], &mut graph, - ) - .unwrap(); + ); assert!(graph.iter().all(|neighbors| neighbors.is_empty())); } - #[test] - fn poisoned_candidate_lists_return_errors() { - let candidates = DirectCandidates::new(1).unwrap(); - let _ = std::panic::catch_unwind(|| { - let _guard = candidates.lists[0].lock().unwrap(); - panic!("poison candidate list"); - }); - assert!(matches!( - candidates.add_leaf(&[0], &[crate::graph::AdjacencyList::new()]), - Err(LeafBuildError::PoisonedCandidateList { point: 0 }) - )); - assert!(matches!( - candidates.into_lists(), - Err(LeafBuildError::PoisonedCandidateList { point: 0 }) - )); - } - - #[test] - fn allocation_errors_preserve_buffer_context() { - let mut values = Vec::::new(); - let source = values.try_reserve(usize::MAX).unwrap_err(); - let error = allocation_error("test", 1, source); - assert!(matches!( - error, - LeafBuildError::Allocation { - buffer: "test", - additional: 1, - .. - } - )); - } - #[test] fn direct_candidate_accumulator_keeps_unique_sorted_lists() { - let candidates = DirectCandidates::new(2).unwrap(); - candidates - .add_leaf( - &[0, 1], - &[ - crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), - crate::graph::AdjacencyList::from_iter_untrusted([0]), - ], - ) - .unwrap(); - assert_eq!( - adjacency_lists(candidates.into_lists().unwrap()), - [vec![1], vec![0]] + let candidates = DirectCandidates::new(2); + candidates.add_leaf( + &[0, 1], + &[ + crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), + crate::graph::AdjacencyList::from_iter_untrusted([0]), + ], ); + assert_eq!(adjacency_lists(candidates.into_lists()), [vec![1], vec![0]]); } } diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 8b3f11243b..9c0e888c49 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -71,7 +71,7 @@ pub struct PiPNNConfig { /// Number of nearest centers assigned at each recursive partition level. /// Levels after this schedule assign each point to one center. pub fanout: Vec, - /// Number of nearest neighbors selected within each leaf (`1..=3`). + /// Number of nearest neighbors selected within each leaf. pub leaf_k: usize, /// Number of independent partition passes over the dataset. pub replicas: usize, @@ -104,12 +104,8 @@ impl PiPNNConfig { if self.fanout.contains(&0) { return Err(config_error("fanout values must be greater than zero")); } - if !(1..=leaf_kernel::MAX_LEAF_NEIGHBORS).contains(&self.leaf_k) { - return Err(config_error(format!( - "leaf_k ({}) must be in [1, {}]", - self.leaf_k, - leaf_kernel::MAX_LEAF_NEIGHBORS - ))); + if self.leaf_k == 0 { + return Err(config_error("leaf_k must be greater than zero")); } if self.replicas == 0 { return Err(config_error("replicas must be greater than zero")); @@ -330,7 +326,7 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod build_graph_tests { - use super::{PiPNNBuildContext, PiPNNConfig, build_graph, leaf_kernel}; + use super::{PiPNNBuildContext, PiPNNConfig, build_graph}; use crate::graph::config::{self, MaxDegree}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; @@ -420,7 +416,7 @@ mod build_graph_tests { c_min: 1, p_samp: 0.5, fanout: vec![2], - leaf_k: leaf_kernel::MAX_LEAF_NEIGHBORS, + leaf_k: 4, replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -544,7 +540,7 @@ mod build_graph_tests { c_min, p_samp: 0.5, fanout: vec![2], - leaf_k: rng.random_range(1..=3), + leaf_k: rng.random_range(1..=7), replicas: rng.random_range(1..=2), }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); @@ -562,7 +558,7 @@ mod build_graph_tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod config_tests { - use super::{PiPNNBuildContext, PiPNNConfig, leaf_kernel}; + use super::{PiPNNBuildContext, PiPNNConfig}; use crate::graph::config::{self, MaxDegree}; use diskann_vector::distance::Metric; @@ -637,10 +633,6 @@ mod config_tests { leaf_k: 0, ..pipnn_config() }, - PiPNNConfig { - leaf_k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, - ..pipnn_config() - }, PiPNNConfig { replicas: 0, ..pipnn_config() diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 052ec36608..d1a899239c 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -8,9 +8,9 @@ //! A leader is a sampled point that acts as the center of one child partition. //! A point can join several leaders, so child partitions can overlap. //! -//! This module recursively splits dataset point IDs into bounded leaves. It uses -//! dense GEMM to compare points with sampled leaders. An `ObjectPool` supplies -//! reusable buffers to Rayon worker chunks. +//! This module owns recursive splitting, leader sampling, row gathering, and +//! assignment scatter. The partition kernel owns GEMM, norm preparation, and +//! local ranking. An `ObjectPool` supplies reusable scratch to Rayon workers. //! //! A configured level assigns each point to `fanout[level]` leaders. A deeper //! level assigns each point to one leader. Each replica uses a different @@ -19,7 +19,6 @@ use std::collections::HashSet; use crate::{ANNError, ANNResult, utils::VectorRepr}; -use diskann_linalg::Transpose; use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, @@ -30,8 +29,8 @@ use rayon::prelude::*; use super::{ PiPNNConfig, - kernel_metric::{NormPreparation, PartitionMetric, PartitionNorms}, - partition_kernel::{PartitionInput, nearest_leaders}, + kernel_metric::PartitionMetric, + partition_kernel::{PartitionKernelWorkspace, PreparedLeaders, assign_leaders}, }; // These constants control internal batching and deterministic seed generation. @@ -44,15 +43,9 @@ const MAX_ASSIGNMENT_STRIPE_POINTS: usize = 1_024; const PARALLEL_SCATTER_MIN_POINTS: usize = 100_000; const MAX_PARTITION_ITERATIONS: usize = 30; -/// Error from partition input checks, allocation, or recursion progress. +/// Error from partition shape checks or recursion progress. #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub(crate) enum PartitionError { - #[error("PiPNN cannot partition an empty dataset")] - EmptyDataset, - #[error("PiPNN cannot partition vectors with zero dimensions")] - EmptyDimensions, - #[error("dataset has {0} points, which exceeds the u32 ID limit")] - TooManyPoints(usize), #[error("{buffer} shape {rows} x {cols} overflows usize")] ShapeOverflow { buffer: &'static str, @@ -68,14 +61,6 @@ pub(crate) enum PartitionError { level: usize, limit: usize, }, - #[error("invalid {buffer} length: expected {expected}, got {actual}")] - InvalidBufferLength { - buffer: &'static str, - expected: usize, - actual: usize, - }, - #[error("partition worker did not publish its result")] - MissingWorkerResult, } struct WorkItem { @@ -87,9 +72,7 @@ struct WorkItem { #[derive(Default)] struct StripeBuffers { point_values: Vec, - dot_values: Vec, - point_norms: Vec, - ranked_leaders: Vec<(u32, f32)>, + kernel_workspace: PartitionKernelWorkspace, } impl AsPooled<()> for StripeBuffers { @@ -128,26 +111,12 @@ where M: PartitionMetric, T: VectorRepr + Send + Sync, { - let points = data.nrows(); - if points == 0 { - return Err(ANNError::new(PartitionError::EmptyDataset)); - } - if data.ncols() == 0 { - return Err(ANNError::new(PartitionError::EmptyDimensions)); - } - if points > u32::MAX as usize { - return Err(ANNError::new(PartitionError::TooManyPoints(points))); - } - let mut leaves = Vec::new(); let stripe_buffers = StripeBufferPool::new((), 0, None); for replica in 0..config.replicas { let seed = replica_seed(replica); let mut replica_leaves = partition_replica::(arch, data, config, seed, &stripe_buffers)?; - leaves - .try_reserve(replica_leaves.len()) - .map_err(ANNError::new)?; leaves.append(&mut replica_leaves); } Ok(leaves) @@ -172,67 +141,38 @@ where M: PartitionMetric, T: VectorRepr + Send + Sync, { - let initial_indices = point_ids(data.nrows())?; + let initial_indices = point_ids(data.nrows()); if data.nrows() <= config.c_max { - let mut leaves = Vec::new(); - leaves.try_reserve_exact(1).map_err(ANNError::new)?; - leaves.push(initial_indices); - return Ok(leaves); + return Ok(vec![initial_indices]); } let mut leaves = Vec::new(); - let mut work = Vec::new(); - work.try_reserve_exact(1).map_err(ANNError::new)?; - work.push(WorkItem { + let mut work = vec![WorkItem { indices: initial_indices, level: 0, seed, - }); + }]; for _ in 0..MAX_PARTITION_ITERATIONS { if work.is_empty() { return merge_undersized_leaves(leaves, config.c_min, config.c_max); } - let mut results = Vec::new(); - results - .try_reserve_exact(work.len()) - .map_err(ANNError::new)?; - results.resize_with(work.len(), || None); - // `build_graph` runs this Rayon operation in the pool from the build - // context. Each worker writes only to its indexed result slot. + // Indexed parallel collection preserves work-item order. #[allow(clippy::disallowed_methods)] - results - .par_iter_mut() - .zip(work.into_par_iter()) - .try_for_each(|(slot, item)| { - *slot = Some(partition_work_item::( - arch, - data, - config, - item, - stripe_buffers, - )?); - Ok::<(), ANNError>(()) - })?; + let results: ANNResult> = work + .into_par_iter() + .map(|item| partition_work_item::(arch, data, config, item, stripe_buffers)) + .collect(); let mut next_work = Vec::new(); - for result in results { - let (mut pending, mut finished) = - result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))?; - next_work - .try_reserve(pending.len()) - .map_err(ANNError::new)?; - leaves.try_reserve(finished.len()).map_err(ANNError::new)?; + for (mut pending, mut finished) in results? { next_work.append(&mut pending); leaves.append(&mut finished); } work = next_work; } - if work.is_empty() { - return merge_undersized_leaves(leaves, config.c_min, config.c_max); - } let Some(largest) = work.iter().max_by_key(|item| item.indices.len()) else { return merge_undersized_leaves(leaves, config.c_min, config.c_max); }; @@ -268,16 +208,12 @@ where &item.indices, config.p_samp, mix_seed(item.seed, points as u64), - )?; + ); let clusters = assign_to_leaders::(arch, data, &item.indices, &leaders, fanout, stripe_buffers)?; let mut pending = Vec::new(); let mut finished = Vec::new(); - pending.try_reserve(clusters.len()).map_err(ANNError::new)?; - finished - .try_reserve(clusters.len()) - .map_err(ANNError::new)?; let child_seed = mix_seed(item.seed, points as u64); for cluster in clusters { if cluster.is_empty() { @@ -297,13 +233,10 @@ where } /// Sample point IDs that act as centers for one partition split. -fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResult> { +fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> Vec { let count = sampled_leader_count(points.len(), sampling_fraction); let mut rng = rand::rngs::StdRng::seed_from_u64(seed); - let mut leaders = Vec::new(); - leaders.try_reserve_exact(count).map_err(ANNError::new)?; - leaders.extend(points.choose_multiple(&mut rng, count).copied()); - Ok(leaders) + points.choose_multiple(&mut rng, count).copied().collect() } /// Return the number of centers to sample from one cluster. @@ -351,30 +284,18 @@ where { let dimension_count = data.ncols(); let leader_values_len = checked_area("leader data", leader_ids.len(), dimension_count)?; - let mut leader_values = filled_vec(leader_values_len, 0.0f32)?; + let mut leader_values = vec![0.0f32; leader_values_len]; gather_vectors(data, leader_ids, &mut leader_values)?; let leader_matrix = - MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count).map_err( - |error| { - ANNError::new(PartitionError::InvalidBufferLength { - buffer: "leader data", - expected: leader_values_len, - actual: error.into_inner().len(), - }) - }, - )?; - let mut leader_norm_values = Vec::new(); - M::prepare_leader_norms(NormPreparation { - values: leader_matrix, - norms: &mut leader_norm_values, - }) - .map_err(ANNError::new)?; - - let fanout = fanout.min(leader_ids.len()); + MatrixView::try_from(leader_values.as_slice(), leader_ids.len(), dimension_count) + .map_err(|error| ANNError::new(error.as_static()))?; + let leaders = PreparedLeaders::::new(leader_matrix); + + let fanout = fanout.min(leaders.len()); let assignment_len = checked_area("partition assignments", point_ids.len(), fanout)?; - let mut assignments = filled_vec(assignment_len, 0u32)?; - let stripe_points = assignment_stripe_point_count(leader_ids.len()); + let mut assignments = vec![0u32; assignment_len]; + let stripe_points = assignment_stripe_point_count(leaders.len()); let stripe_assignment_count = checked_area("assignment stripe", stripe_points, fanout)?; let stripe_count = point_ids.len().div_ceil(stripe_points); let worker_stripe_count = stripe_count.div_ceil(rayon::current_num_threads().max(1)); @@ -400,8 +321,7 @@ where arch, data, &point_ids[first_point..first_point + stripe_point_count], - &leader_values, - &leader_norm_values, + &leaders, fanout, &mut buffers, stripe_assignments, @@ -415,16 +335,15 @@ where /// Assign one point stripe to sampled partition centers. /// -/// The function gathers point vectors and computes point-to-center dot products. -/// It writes center-column IDs for partition scatter. +/// The function gathers point IDs into a packed `f32` matrix. The partition +/// kernel owns dot products, point norms, and ranking. This function writes the +/// returned leader-column IDs for partition scatter. #[inline] -#[allow(clippy::too_many_arguments)] fn assign_point_stripe( arch: A, data: MatrixView<'_, T>, point_ids: &[u32], - leader_values: &[f32], - leader_norm_values: &[f32], + leaders: &PreparedLeaders<'_, M>, fanout: usize, buffers: &mut StripeBuffers, assignments: &mut [u32], @@ -439,88 +358,30 @@ where { let point_count = point_ids.len(); let dimensions = data.ncols(); - let leader_count = leader_values.len() / dimensions; let point_values_len = checked_area("point stripe", point_count, dimensions)?; - let dots_len = checked_area("dot-product stripe", point_count, leader_count)?; - let output_len = checked_area("partition assignments", point_count, fanout)?; // Keep each buffer at its largest length. Every operation uses an explicit // active prefix. - grow_fallible(&mut buffers.point_values, point_values_len, 0.0)?; - grow_fallible(&mut buffers.dot_values, dots_len, 0.0)?; + grow(&mut buffers.point_values, point_values_len, 0.0); let StripeBuffers { - point_values: point_buffer, - dot_values: dot_buffer, - point_norms: point_norm_buffer, - ranked_leaders, + point_values, + kernel_workspace, } = buffers; - let mut point_matrix = MutMatrixView::try_from( - &mut point_buffer[..point_values_len], - point_count, - dimensions, - ) - .map_err(|error| { - ANNError::new(PartitionError::InvalidBufferLength { - buffer: "point stripe", - expected: point_values_len, - actual: error.into_inner().len(), - }) - })?; - let dots = &mut dot_buffer[..dots_len]; - gather_vectors(data, point_ids, point_matrix.as_mut_slice())?; - diskann_linalg::sgemm( - Transpose::None, - Transpose::Ordinary, + let mut points = MutMatrixView::try_from( + &mut point_values[..point_values_len], point_count, - leader_count, dimensions, - 1.0, - point_matrix.as_slice(), - leader_values, - None, - dots, ) - .map_err(ANNError::new)?; - - M::prepare_point_norms(NormPreparation { - values: point_matrix.as_view(), - norms: point_norm_buffer, - }) - .map_err(ANNError::new)?; - let point_norms = &*point_norm_buffer; - let norms = PartitionNorms { - point_norms, - leader_norms: leader_norm_values, - }; - let dots = MatrixView::try_from(&*dots, point_count, leader_count).map_err(|_| { - ANNError::new(PartitionError::InvalidBufferLength { - buffer: "dot-product stripe", - expected: dots_len, - actual: dots.len(), - }) - })?; - let output = MutMatrixView::try_from(assignments, point_count, fanout).map_err(|error| { - ANNError::new(PartitionError::InvalidBufferLength { - buffer: "partition assignments", - expected: output_len, - actual: error.into_inner().len(), - }) - })?; - nearest_leaders::(arch, PartitionInput { dots, norms }, output, ranked_leaders) - .map_err(ANNError::new) + .map_err(|error| ANNError::new(error.as_static()))?; + gather_vectors(data, point_ids, points.as_mut_slice())?; + let output = MutMatrixView::try_from(assignments, point_count, fanout) + .map_err(|error| ANNError::new(error.as_static()))?; + assign_leaders::(arch, points.as_view(), leaders, output, kernel_workspace) } fn gather_vectors(data: MatrixView<'_, T>, indices: &[u32], output: &mut [f32]) -> ANNResult<()> where T: VectorRepr, { - let expected = checked_area("gather output", indices.len(), data.ncols())?; - if output.len() != expected { - return Err(ANNError::new(PartitionError::InvalidBufferLength { - buffer: "gather output", - expected, - actual: output.len(), - })); - } for (&index, vector_output) in indices.iter().zip(output.chunks_exact_mut(data.ncols())) { T::as_f32_into(data.row(index as usize), vector_output).map_err(Into::::into)?; } @@ -538,64 +399,41 @@ fn scatter_assignments( leaders: usize, ) -> ANNResult>> { if points.len() < PARALLEL_SCATTER_MIN_POINTS { - return scatter_serial(points, assignments, fanout, leaders); + return Ok(scatter_serial(points, assignments, fanout, leaders)); } let stripe_points = points.len().div_ceil(rayon::current_num_threads().max(1)); let stripe_assignment_count = checked_area("scatter assignment stripe", stripe_points, fanout)?; - let stripes = points.len().div_ceil(stripe_points); - let mut partials = Vec::new(); - partials.try_reserve_exact(stripes).map_err(ANNError::new)?; - partials.resize_with(stripes, || None); - // `build_graph` runs this Rayon operation in the pool from the build context. - // Each worker writes only to its indexed partial result. + // Indexed parallel collection preserves stripe order. #[allow(clippy::disallowed_methods)] - partials - .par_iter_mut() - .zip( - points - .par_chunks(stripe_points) - .zip(assignments.par_chunks(stripe_assignment_count)), - ) - .try_for_each(|(slot, (points, assignments))| { - *slot = Some(scatter_serial(points, assignments, fanout, leaders)?); - Ok::<(), ANNError>(()) - })?; + let locals: Vec<_> = points + .par_chunks(stripe_points) + .zip(assignments.par_chunks(stripe_assignment_count)) + .map(|(points, assignments)| scatter_serial(points, assignments, fanout, leaders)) + .collect(); - let mut locals = Vec::new(); - locals.try_reserve_exact(stripes).map_err(ANNError::new)?; - for result in partials { - locals.push(result.ok_or_else(|| ANNError::new(PartitionError::MissingWorkerResult))?); - } - - let mut sizes = filled_vec(leaders, 0usize)?; + let mut sizes = vec![0usize; leaders]; for local in &locals { for (size, cluster) in sizes.iter_mut().zip(local) { - *size = size.checked_add(cluster.len()).ok_or_else(|| { - ANNError::new(PartitionError::ShapeOverflow { - buffer: "cluster size", - rows: *size, - cols: cluster.len(), - }) - })?; + *size += cluster.len(); } } // `build_graph` runs this Rayon operation in the pool from the build context. // Each worker creates one independent leader cluster. #[allow(clippy::disallowed_methods)] - sizes + let clusters = sizes .into_par_iter() .enumerate() .map(|(leader, size)| { - let mut cluster = Vec::new(); - cluster.try_reserve_exact(size).map_err(ANNError::new)?; + let mut cluster = Vec::with_capacity(size); for local in &locals { cluster.extend_from_slice(&local[leader]); } - Ok(cluster) + cluster }) - .collect() + .collect(); + Ok(clusters) } fn scatter_serial( @@ -603,53 +441,22 @@ fn scatter_serial( assignments: &[u32], fanout: usize, leaders: usize, -) -> ANNResult>> { - let expected = checked_area("scatter assignments", points.len(), fanout)?; - if assignments.len() != expected { - return Err(ANNError::new(PartitionError::InvalidBufferLength { - buffer: "scatter assignments", - expected, - actual: assignments.len(), - })); - } - - let mut sizes = filled_vec(leaders, 0usize)?; +) -> Vec> { + let mut sizes = vec![0usize; leaders]; for &leader in assignments { - let Some(size) = sizes.get_mut(leader as usize) else { - return Err(ANNError::new(PartitionError::InvalidBufferLength { - buffer: "leader assignment", - expected: leaders, - actual: leader as usize + 1, - })); - }; - *size = size.checked_add(1).ok_or_else(|| { - ANNError::new(PartitionError::ShapeOverflow { - buffer: "cluster size", - rows: *size, - cols: 1, - }) - })?; + sizes[leader as usize] += 1; } - let mut clusters = clusters_with_capacities(&sizes)?; + let mut clusters = clusters_with_capacities(&sizes); for (&point, point_assignments) in points.iter().zip(assignments.chunks_exact(fanout)) { for &leader in point_assignments { clusters[leader as usize].push(point); } } - Ok(clusters) + clusters } -fn clusters_with_capacities(sizes: &[usize]) -> ANNResult>> { - let mut clusters = Vec::new(); - clusters - .try_reserve_exact(sizes.len()) - .map_err(ANNError::new)?; - for &size in sizes { - let mut cluster = Vec::new(); - cluster.try_reserve_exact(size).map_err(ANNError::new)?; - clusters.push(cluster); - } - Ok(clusters) +fn clusters_with_capacities(sizes: &[usize]) -> Vec> { + sizes.iter().map(|&size| Vec::with_capacity(size)).collect() } /// Merge leaves smaller than `c_min` without exceeding `c_max`. @@ -661,12 +468,8 @@ fn merge_undersized_leaves( c_min: usize, c_max: usize, ) -> ANNResult>> { - let mut merged = Vec::new(); + let mut merged = Vec::with_capacity(leaves.len()); let mut small_leaves = Vec::new(); - merged.try_reserve(leaves.len()).map_err(ANNError::new)?; - small_leaves - .try_reserve(leaves.len()) - .map_err(ANNError::new)?; for leaf in leaves { if leaf.len() >= c_min { merged.push(leaf); @@ -678,42 +481,27 @@ fn merge_undersized_leaves( return Ok(merged); } - let mut small = HashSet::new(); - small.try_reserve(c_max).map_err(ANNError::new)?; + let mut small = HashSet::with_capacity(c_max); for leaf in small_leaves { - let combined = small.len().checked_add(leaf.len()).ok_or_else(|| { - ANNError::new(PartitionError::ShapeOverflow { - buffer: "small-leaf merge", - rows: small.len(), - cols: leaf.len(), - }) - })?; + let combined = small.len() + leaf.len(); if combined > c_max { - merged.push(drain_sorted(&mut small)?); + merged.push(drain_sorted(&mut small)); } - small.try_reserve(leaf.len()).map_err(ANNError::new)?; small.extend(leaf); if small.len() >= c_min { - merged.push(drain_sorted(&mut small)?); + merged.push(drain_sorted(&mut small)); } } if !small.is_empty() { - let mut remainder = drain_sorted(&mut small)?; + let mut remainder = drain_sorted(&mut small); if remainder.len() < c_min && let Some(last) = merged.last_mut() { remainder.retain(|id| !last.contains(id)); - let combined = last.len().checked_add(remainder.len()).ok_or_else(|| { - ANNError::new(PartitionError::ShapeOverflow { - buffer: "small-leaf tail merge", - rows: last.len(), - cols: remainder.len(), - }) - })?; + let combined = last.len() + remainder.len(); if combined <= c_max { - last.try_reserve(remainder.len()).map_err(ANNError::new)?; last.append(&mut remainder); last.sort_unstable(); } @@ -726,37 +514,20 @@ fn merge_undersized_leaves( Ok(merged) } -fn drain_sorted(set: &mut HashSet) -> ANNResult> { - let mut values = Vec::new(); - values.try_reserve_exact(set.len()).map_err(ANNError::new)?; - values.extend(set.drain()); +fn drain_sorted(set: &mut HashSet) -> Vec { + let mut values: Vec<_> = set.drain().collect(); values.sort_unstable(); - Ok(values) -} - -fn point_ids(points: usize) -> ANNResult> { - let mut ids = Vec::new(); - ids.try_reserve_exact(points).map_err(ANNError::new)?; - ids.extend(0..points as u32); - Ok(ids) + values } -fn filled_vec(len: usize, value: T) -> ANNResult> { - let mut values = Vec::new(); - values.try_reserve_exact(len).map_err(ANNError::new)?; - values.resize(len, value); - Ok(values) +fn point_ids(points: usize) -> Vec { + (0..points as u32).collect() } -fn grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { - if values.len() >= len { - return Ok(()); +fn grow(values: &mut Vec, len: usize, value: T) { + if values.len() < len { + values.resize(len, value); } - values - .try_reserve(len - values.len()) - .map_err(ANNError::new)?; - values.resize(len, value); - Ok(()) } fn checked_area(buffer: &'static str, rows: usize, cols: usize) -> ANNResult { @@ -1184,78 +955,9 @@ mod tests { .flat_map(|point| [point % 7, (point + 3) % 7]) .collect(); - let expected = scatter_serial(&points, &assignments, 2, 7).unwrap(); + let expected = scatter_serial(&points, &assignments, 2, 7); let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); assert_eq!(actual, expected); } - - #[test] - fn rejects_empty_dataset() { - let data = Matrix::::new(0.0, 0, 4); - let error = - partition_with_runtime_metric(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2) - .unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::EmptyDataset - ); - } - - #[test] - fn rejects_zero_dimensions() { - let data = Matrix::::new(0.0, 4, 0); - let error = - partition_with_runtime_metric(data.as_view(), &config(1, 4, vec![1], 1), Metric::L2) - .unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::EmptyDimensions - ); - } - - #[test] - fn rejects_invalid_gather_output_length() { - let data = Matrix::::new(0.0, 2, 2); - let error = gather_vectors(data.as_view(), &[0, 1], &mut [0.0; 3]).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "gather output", - expected: 4, - actual: 3, - } - ); - } - - #[test] - fn rejects_invalid_assignment_length() { - let error = scatter_serial(&[7, 8], &[0, 1, 0], 2, 2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "scatter assignments", - expected: 4, - actual: 3, - } - ); - } - - #[test] - fn rejects_assignment_to_an_unknown_leader() { - let error = scatter_serial(&[7], &[2], 1, 2).unwrap_err(); - - assert_eq!( - error.downcast::().unwrap(), - PartitionError::InvalidBufferLength { - buffer: "leader assignment", - expected: 2, - actual: 3, - } - ); - } } From 8bb030f874d9e4bed526c15daf7a8901b6cb4c90 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:43:50 +0000 Subject: [PATCH 54/58] fix(pipnn): omit unrankable graph candidates --- diskann/src/graph/pipnn/leaf_build.rs | 12 ++++++++++++ diskann/src/graph/pipnn/partitioning.rs | 25 ++++++++++++++++++++++--- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 5f02967125..20d276d180 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -265,6 +265,9 @@ fn add_symmetric_neighbors( ) { for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { for neighbor in source_neighbors { + if !neighbor.is_assigned() { + continue; + } let target = neighbor.target as usize; let source_id = point_ids[source]; let target_id = point_ids[target]; @@ -442,6 +445,15 @@ mod tests { assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); } + #[test] + fn non_rankable_neighbors_are_omitted() { + let data = [0.0_f32, 1.0, f32::NAN]; + + let graph = build(view(&data, 3, 1), &[vec![0, 1, 2]], 2, Metric::InnerProduct).unwrap(); + + assert_eq!(adjacency_lists(graph), [vec![1], vec![0], vec![]]); + } + #[test] fn retains_and_deduplicates_candidates_from_overlapping_leaves() { let data = [0.0_f32, 1.0, 2.0, 3.0]; diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index d1a899239c..9d13826d83 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -30,7 +30,9 @@ use rayon::prelude::*; use super::{ PiPNNConfig, kernel_metric::PartitionMetric, - partition_kernel::{PartitionKernelWorkspace, PreparedLeaders, assign_leaders}, + partition_kernel::{ + PartitionKernelWorkspace, PreparedLeaders, UNASSIGNED_LEADER, assign_leaders, + }, }; // These constants control internal batching and deterministic seed generation. @@ -444,12 +446,16 @@ fn scatter_serial( ) -> Vec> { let mut sizes = vec![0usize; leaders]; for &leader in assignments { - sizes[leader as usize] += 1; + if leader != UNASSIGNED_LEADER { + sizes[leader as usize] += 1; + } } let mut clusters = clusters_with_capacities(&sizes); for (&point, point_assignments) in points.iter().zip(assignments.chunks_exact(fanout)) { for &leader in point_assignments { - clusters[leader as usize].push(point); + if leader != UNASSIGNED_LEADER { + clusters[leader as usize].push(point); + } } } clusters @@ -947,6 +953,19 @@ mod tests { assert_eq!(clusters[1], (1_024..2_048).collect::>()); } + #[test] + fn scatter_omits_unassigned_slots() { + assert_eq!( + scatter_serial( + &[10, 11], + &[0, UNASSIGNED_LEADER, UNASSIGNED_LEADER, 1], + 2, + 2 + ), + [vec![10], vec![11]] + ); + } + #[test] fn parallel_scatter_matches_serial_order() { let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); From 89f7a85e319ecbe328be261583e9f464fd9eb57c Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:50:19 +0000 Subject: [PATCH 55/58] test(pipnn): cover non-rankable build candidates --- diskann/src/graph/pipnn/mod.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 9c0e888c49..661a08bcd8 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -405,6 +405,28 @@ mod build_graph_tests { } } + #[test] + fn omits_non_rankable_candidates_without_invalid_ids() { + let values = [0.0_f32, 1.0, f32::NAN]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let graph = graph_config(Metric::InnerProduct, 2); + let pool = pool(1); + let config = PiPNNConfig { + c_max: 2, + c_min: 1, + p_samp: 1.0, + fanout: vec![2], + leaf_k: 1, + replicas: 1, + }; + let context = PiPNNBuildContext::new(config, &graph, Metric::InnerProduct, &pool).unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + assert_graph_invariants(&actual, 3, 2); + assert_eq!(rows(actual), [vec![1], vec![0], vec![]]); + } + #[test] fn prunes_overfull_single_leaf_candidates_to_the_graph_degree() { let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; From 46e84562b016249e3443bf0a8438af47b68d018b Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 19 Aug 2026 05:45:11 +0000 Subject: [PATCH 56/58] perf(pipnn): deduplicate direct leaf edges once --- diskann/src/graph/pipnn/leaf_build.rs | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 20d276d180..18b952d329 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -63,7 +63,7 @@ pub(crate) enum LeafBuildError { struct LeafBuffers { point_values: Vec, neighbors: Vec, - local_adjacency: Vec>, + local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, } @@ -99,11 +99,10 @@ impl LeafBuffers { } fn prepare_local_adjacency(&mut self, point_count: usize) { - self.local_adjacency - .resize_with(point_count, AdjacencyList::new); + self.local_adjacency.resize_with(point_count, Vec::new); self.local_adjacency[..point_count] .iter_mut() - .for_each(AdjacencyList::clear); + .for_each(Vec::clear); } } @@ -124,7 +123,7 @@ impl DirectCandidates { Self { lists } } - fn add_leaf(&self, point_ids: &[u32], local_adjacency: &[AdjacencyList]) { + fn add_leaf(&self, point_ids: &[u32], local_adjacency: &[Vec]) { for (&source, additions) in point_ids.iter().zip(local_adjacency) { // `add_direct_leaf_candidates` checks every point ID before this append. self.lists[source as usize] @@ -261,7 +260,7 @@ fn add_symmetric_neighbors( point_ids: &[u32], leaf_k: usize, neighbors: &[LeafNeighbor], - local_adjacency: &mut [AdjacencyList], + local_adjacency: &mut [Vec], ) { for (source, source_neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { for neighbor in source_neighbors { @@ -625,7 +624,7 @@ mod tests { #[test] fn skips_duplicate_global_ids_without_self_edges() { - let mut graph = vec![crate::graph::AdjacencyList::new(); 2]; + let mut graph = vec![Vec::new(); 2]; add_symmetric_neighbors( &[7, 7], 1, @@ -641,13 +640,7 @@ mod tests { #[test] fn direct_candidate_accumulator_keeps_unique_sorted_lists() { let candidates = DirectCandidates::new(2); - candidates.add_leaf( - &[0, 1], - &[ - crate::graph::AdjacencyList::from_iter_untrusted([1, 1]), - crate::graph::AdjacencyList::from_iter_untrusted([0]), - ], - ); + candidates.add_leaf(&[0, 1], &[vec![1, 1], vec![0]]); assert_eq!(adjacency_lists(candidates.into_lists()), [vec![1], vec![0]]); } } From 74f8ac3783dc578c52afcfecbab62ba029765f67 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:09:29 +0000 Subject: [PATCH 57/58] refactor(pipnn): use the SIMD schema in stages --- diskann/src/graph/pipnn/leaf_build.rs | 23 ++++----------- diskann/src/graph/pipnn/mod.rs | 20 ++++--------- diskann/src/graph/pipnn/partitioning.rs | 37 ++++++------------------- 3 files changed, 20 insertions(+), 60 deletions(-) diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 18b952d329..bdc4fc06dd 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -21,12 +21,12 @@ use parking_lot::Mutex; use crate::{graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; -use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rayon::prelude::*; use super::{ kernel_metric::LeafMetric, leaf_kernel::{LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, select_leaf_neighbors}, + simd::PiPNNSIMDSchema, }; /// Failure while converting leaves into direct graph candidates. @@ -156,10 +156,7 @@ pub(super) fn build_leaf_candidates( requested_k: usize, ) -> Result>, LeafBuildError> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: LeafMetric, T: VectorRepr + 'static, { @@ -196,10 +193,7 @@ fn add_direct_leaf_candidates( candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: LeafMetric, T: VectorRepr + 'static, { @@ -288,13 +282,11 @@ fn grow(values: &mut Vec, len: usize, value: T) { mod tests { use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; - use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, - }; + use diskann_wide::arch::{self, Target1}; use half::f16; use std::collections::BTreeSet; + use super::super::simd::PiPNNSIMDSchema; use super::{ DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, build_leaf_candidates, @@ -326,10 +318,7 @@ mod tests { LeafBuildCall<'_, T>, > for DispatchLeafBuild where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, T: crate::utils::VectorRepr + 'static, { fn run( diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 661a08bcd8..5ed62698cf 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -47,14 +47,12 @@ use crate::{ }; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; -use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, -}; +use diskann_wide::arch::{self, Target1}; use rayon::ThreadPool; -use self::kernel_metric::{ - Cosine, CosineNormalized, InnerProduct, L2, LeafMetric, PartitionMetric, +use self::{ + kernel_metric::{Cosine, CosineNormalized, InnerProduct, L2, LeafMetric, PartitionMetric}, + simd::PiPNNSIMDSchema, }; /// PiPNN partition and leaf-selection policy. @@ -213,10 +211,7 @@ struct RunBuildGraph; impl Target1>>, BuildGraphCall<'_, '_, '_, T>> for RunBuildGraph where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, T: VectorRepr + Send + Sync + 'static, { fn run( @@ -256,10 +251,7 @@ fn build_graph_for( metric: Metric, ) -> ANNResult>> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: LeafMetric + PartitionMetric, T: VectorRepr + Send + Sync + 'static, { diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 9d13826d83..7325630f67 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -23,7 +23,6 @@ use diskann_utils::{ object_pool::{AsPooled, ObjectPool}, views::{MatrixView, MutMatrixView}, }; -use diskann_wide::{Architecture, SIMDMask, SIMDSelect, SIMDVector}; use rand::{SeedableRng, prelude::IndexedRandom}; use rayon::prelude::*; @@ -33,6 +32,7 @@ use super::{ partition_kernel::{ PartitionKernelWorkspace, PreparedLeaders, UNASSIGNED_LEADER, assign_leaders, }, + simd::PiPNNSIMDSchema, }; // These constants control internal batching and deterministic seed generation. @@ -106,10 +106,7 @@ pub(super) fn partition( config: &PiPNNConfig, ) -> ANNResult>> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { @@ -136,10 +133,7 @@ fn partition_replica( stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { @@ -197,10 +191,7 @@ fn partition_work_item( stripe_buffers: &StripeBufferPool, ) -> ANNResult<(Vec, Vec>)> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { @@ -277,10 +268,7 @@ fn assign_to_leaders( stripe_buffers: &StripeBufferPool, ) -> ANNResult>> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr + Send + Sync, { @@ -351,10 +339,7 @@ fn assign_point_stripe( assignments: &mut [u32], ) -> ANNResult<()> where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, M: PartitionMetric, T: VectorRepr, { @@ -555,10 +540,7 @@ fn assignment_stripe_point_count(leader_count: usize) -> usize { mod tests { use diskann_utils::views::{Matrix, MatrixView}; use diskann_vector::{Half, distance::Metric}; - use diskann_wide::{ - Architecture, SIMDMask, SIMDSelect, SIMDVector, - arch::{self, Target1}, - }; + use diskann_wide::arch::{self, Target1}; use super::*; @@ -571,10 +553,7 @@ mod tests { impl Target1>>, PartitionCall<'_, T>> for DispatchPartition where - A: Architecture, - A::f32x16: std::ops::Div, - ::Mask: SIMDSelect, - u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + A: PiPNNSIMDSchema, T: VectorRepr + Send + Sync, { fn run(self, arch: A, call: PartitionCall<'_, T>) -> ANNResult>> { From f9351bb2dbb22bd1226989b2b86e7f44196623d6 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:08:41 +0000 Subject: [PATCH 58/58] test(pipnn): clarify core graph contracts Keep exact graph-output, stale-state, boundary, and deterministic behavior checks while dropping allocation-only assertions. --- diskann/src/graph/pipnn/finalization.rs | 39 ++++++--- diskann/src/graph/pipnn/leaf_build.rs | 54 ++++++------- diskann/src/graph/pipnn/partitioning.rs | 100 +++++++++++------------- 3 files changed, 99 insertions(+), 94 deletions(-) diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs index 6f01c3dd25..754865bc9c 100644 --- a/diskann/src/graph/pipnn/finalization.rs +++ b/diskann/src/graph/pipnn/finalization.rs @@ -166,7 +166,8 @@ mod tests { } #[test] - fn preserves_lists_within_the_degree_bound() { + fn row_within_the_degree_bound_is_only_canonicalized() { + // Given let data = [0.0_f32, 1.0, 2.0, 3.0]; let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); let candidates = vec![ @@ -175,14 +176,18 @@ mod tests { candidate_list([]), candidate_list([]), ]; + let expected_canonical_row = [1, 3]; - let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + // When + let actual_rows = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - assert_eq!(&*actual[0], &[1, 3]); + // Then + assert_eq!(&*actual_rows[0], &expected_canonical_row); } #[test] - fn prunes_an_overfull_list_with_the_vamana_kernel() { + fn overfull_row_keeps_the_nearest_unoccluded_neighbors() { + // Given let data = [0.0_f32, 1.0, 2.0, -3.0]; let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); let candidates = vec![ @@ -191,14 +196,18 @@ mod tests { candidate_list([]), candidate_list([]), ]; + let expected_pruned_row = [1, 3]; - let actual = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); + // When + let actual_rows = prune_overfull(data, candidates, &graph_config(2), Metric::L2).unwrap(); - assert_eq!(&*actual[0], &[1, 3]); + // Then + assert_eq!(&*actual_rows[0], &expected_pruned_row); } #[test] fn reused_workspace_matches_fresh_pruning() { + // Given let data = [0.0_f32, 1.0, 2.0, -3.0, 4.0]; let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); let first = [3, 2, 1]; @@ -217,18 +226,26 @@ mod tests { .num_threads(1) .build() .unwrap(); - let fresh_first = pool + // When + let expected_first_row_from_fresh_workspace = pool .install(|| prune_overfull(data, candidates(&first, &[]), &graph, Metric::L2)) .unwrap(); - let fresh_second = pool + let expected_second_row_from_fresh_workspace = pool .install(|| prune_overfull(data, candidates(&[], &second), &graph, Metric::L2)) .unwrap(); - let reused = pool + let actual_rows_from_reused_workspace = pool .install(|| prune_overfull(data, candidates(&first, &second), &graph, Metric::L2)) .unwrap(); - assert_eq!(&*reused[0], &*fresh_first[0]); - assert_eq!(&*reused[1], &*fresh_second[1]); + // Then + assert_eq!( + &*actual_rows_from_reused_workspace[0], + &*expected_first_row_from_fresh_workspace[0] + ); + assert_eq!( + &*actual_rows_from_reused_workspace[1], + &*expected_second_row_from_fresh_workspace[1] + ); } } diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index bdc4fc06dd..2390e0713c 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -410,6 +410,7 @@ mod tests { #[test] fn leaf_adjacency_matches_an_independent_all_pairs_reference() { + // Given let points = [ [0.0_f32, 0.0], [1.0, 0.2], @@ -419,8 +420,10 @@ mod tests { [6.7, -3.2], ]; let flat: Vec<_> = points.into_iter().flatten().collect(); + let expected_adjacency = brute_force_symmetric_l2(&points, 2); - let actual = adjacency_lists( + // When + let actual_adjacency = adjacency_lists( build( view(&flat, points.len(), 2), &[(0..points.len() as u32).collect()], @@ -430,33 +433,41 @@ mod tests { .unwrap(), ); - assert_eq!(actual, brute_force_symmetric_l2(&points, 2)); + // Then + assert_eq!(actual_adjacency, expected_adjacency); } #[test] fn non_rankable_neighbors_are_omitted() { + // Given let data = [0.0_f32, 1.0, f32::NAN]; + let expected_adjacency = [vec![1], vec![0], vec![]]; + // When let graph = build(view(&data, 3, 1), &[vec![0, 1, 2]], 2, Metric::InnerProduct).unwrap(); + let actual_adjacency = adjacency_lists(graph); - assert_eq!(adjacency_lists(graph), [vec![1], vec![0], vec![]]); + // Then + assert_eq!(actual_adjacency, expected_adjacency); } #[test] - fn retains_and_deduplicates_candidates_from_overlapping_leaves() { + fn overlapping_leaves_contribute_each_candidate_once() { + // Given let data = [0.0_f32, 1.0, 2.0, 3.0]; let leaves = vec![vec![0, 1, 2], vec![0, 2, 3], vec![0, 1, 2]]; + let expected_adjacency = [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]]; + // When let graph = build(view(&data, 4, 1), &leaves, 2, Metric::L2).unwrap(); + let actual_adjacency = adjacency_lists(graph); - assert_eq!( - adjacency_lists(graph), - [vec![1, 2, 3], vec![0, 2], vec![0, 1, 3], vec![0, 2]] - ); + // Then + assert_eq!(actual_adjacency, expected_adjacency); } #[test] - fn symmetric_knn_can_give_one_point_more_than_two_k_candidates() { + fn symmetric_edges_can_give_a_center_more_than_two_k_neighbors() { let dimensions = 9; let mut data = vec![0.0_f32; 10 * dimensions]; for source in 1..10 { @@ -561,10 +572,10 @@ mod tests { let leaves: Vec> = (0..32) .map(|offset| (0..16).map(|point| (point + offset) % 64).collect()) .collect(); - let expected = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + let expected_candidate_order = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); for _ in 0..8 { - let actual = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); - assert_eq!(actual, expected); + let actual_candidate_order = build(view(&data, 64, 1), &leaves, 2, Metric::L2).unwrap(); + assert_eq!(actual_candidate_order, expected_candidate_order); } } @@ -588,22 +599,7 @@ mod tests { } #[test] - fn reuses_worker_buffers_for_smaller_leaves() { - let mut buffers = LeafBuffers::default(); - buffers.prepare(0, 64, 128, 2).unwrap(); - let point_values = buffers.point_values.as_ptr(); - let neighbors = buffers.neighbors.as_ptr(); - - buffers.prepare(1, 8, 128, 2).unwrap(); - - assert_eq!(buffers.point_values.as_ptr(), point_values); - assert_eq!(buffers.neighbors.as_ptr(), neighbors); - assert_eq!(buffers.point_values.len(), 64 * 128); - assert_eq!(buffers.neighbors.len(), 64 * 2); - } - - #[test] - fn reports_shape_overflow_before_allocating() { + fn leaf_buffer_preparation_reports_shape_overflow_before_allocating() { let mut buffers = LeafBuffers::default(); assert!(matches!( buffers.prepare(7, usize::MAX, 2, 1), @@ -612,7 +608,7 @@ mod tests { } #[test] - fn skips_duplicate_global_ids_without_self_edges() { + fn symmetric_edge_mapping_skips_duplicate_ids_instead_of_adding_self_edges() { let mut graph = vec![Vec::new(); 2]; add_symmetric_neighbors( &[7, 7], diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs index 7325630f67..395f71ede1 100644 --- a/diskann/src/graph/pipnn/partitioning.rs +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -670,7 +670,7 @@ mod tests { } #[test] - fn returns_one_leaf_at_and_below_c_max() { + fn partition_returns_one_leaf_when_point_count_does_not_exceed_c_max() { for points in [7, 8] { let data = clustered_data(points, 3); let leaves = partition_with_runtime_metric( @@ -726,30 +726,42 @@ mod tests { #[test] fn global_merge_canonicalizes_small_leaf_membership() { + // Given let leaves = vec![vec![9, 3, 1], vec![3, 2], vec![8]]; + let expected_canonical_membership = vec![vec![1, 2, 3, 8, 9]]; - let merged = merge_undersized_leaves(leaves, 4, 8).unwrap(); + // When + let actual_leaves = merge_undersized_leaves(leaves, 4, 8).unwrap(); - assert_eq!(merged, vec![vec![1, 2, 3, 8, 9]]); + // Then + assert_eq!(actual_leaves, expected_canonical_membership); } #[test] fn global_merge_never_overfills_before_reaching_c_min() { + // Given let leaves = vec![vec![0, 1, 2, 3], vec![4, 5, 6, 7], vec![8, 9, 10, 11]]; + let expected_capacity_bounded_leaves = + vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]]; - let merged = merge_undersized_leaves(leaves, 11, 11).unwrap(); + // When + let actual_leaves = merge_undersized_leaves(leaves, 11, 11).unwrap(); - assert_eq!( - merged, - vec![vec![0, 1, 2, 3, 4, 5, 6, 7], vec![8, 9, 10, 11]] - ); + // Then + assert_eq!(actual_leaves, expected_capacity_bounded_leaves); } #[test] fn global_merge_fills_exact_capacity_before_flushing() { - let merged = merge_undersized_leaves(vec![vec![0, 1], vec![2, 3]], 4, 4).unwrap(); + // Given + let leaves = vec![vec![0, 1], vec![2, 3]]; + let expected_exact_capacity_leaf = vec![vec![0, 1, 2, 3]]; + + // When + let actual_leaves = merge_undersized_leaves(leaves, 4, 4).unwrap(); - assert_eq!(merged, vec![vec![0, 1, 2, 3]]); + // Then + assert_eq!(actual_leaves, expected_exact_capacity_leaf); } #[test] @@ -783,23 +795,23 @@ mod tests { let f32_data: Vec = raw.iter().map(|&value| value as f32).collect(); let converted: Vec = raw.iter().copied().map(&convert).collect(); let config = config(2, 16, vec![2, 1], 1); - let expected = partition_with_runtime_metric( + let expected_f32_partition = partition_with_runtime_metric( MatrixView::try_from(&f32_data, points, dimensions).unwrap(), &config, Metric::L2, ) .unwrap(); - let actual = partition_with_runtime_metric( + let actual_converted_partition = partition_with_runtime_metric( MatrixView::try_from(&converted, points, dimensions).unwrap(), &config, Metric::L2, ) .unwrap_or_else(|error| panic!("{label} dimensions={dimensions}: {error}")); - assert_valid_partition_with_runtime_metric(&actual, points, 16, 1); + assert_valid_partition_with_runtime_metric(&actual_converted_partition, points, 16, 1); assert_eq!( - sorted_memberships(&actual), - sorted_memberships(&expected), + sorted_memberships(&actual_converted_partition), + sorted_memberships(&expected_f32_partition), "{label} dimensions={dimensions}" ); } @@ -859,7 +871,7 @@ mod tests { } #[test] - fn all_metrics_produce_valid_partitions() { + fn every_metric_covers_all_points_without_exceeding_c_max() { let data = directional_data(64, 8); let config = config(2, 20, vec![2], 1); @@ -875,7 +887,7 @@ mod tests { } #[test] - fn leader_count_is_bounded() { + fn sampled_leader_count_is_at_least_one_and_never_exceeds_the_cap() { assert_eq!(sampled_leader_count(1, 1.0), 1); assert_eq!(sampled_leader_count(10, 0.01), 2); assert_eq!(sampled_leader_count(50_000, 1.0), LEADER_CAP); @@ -888,31 +900,7 @@ mod tests { } #[test] - fn assignment_stripes_use_power_of_two_point_counts() { - assert_eq!(assignment_stripe_point_count(1_000), 128); - assert_eq!(assignment_stripe_point_count(256), 512); - assert_eq!( - assignment_stripe_point_count(1), - MAX_ASSIGNMENT_STRIPE_POINTS - ); - } - - #[test] - fn stripe_buffer_pool_reuses_returned_capacity() { - let pool = StripeBufferPool::new((), 0, None); - let points = { - let mut buffers = pool.get_ref(()); - buffers.point_values.resize(16, 0.0); - buffers.point_values.as_ptr() - }; - - let buffers = pool.get_ref(()); - assert_eq!(buffers.point_values.as_ptr(), points); - assert_eq!(buffers.point_values.len(), 16); - } - - #[test] - fn leader_assignment_handles_multiple_stripes() { + fn leader_assignment_preserves_clusters_across_multiple_stripes() { let points = 2_048; let data: Vec = (0..points).map(|point| point as f32).collect(); let data = MatrixView::try_from(data.as_slice(), points, 1).unwrap(); @@ -934,28 +922,32 @@ mod tests { #[test] fn scatter_omits_unassigned_slots() { - assert_eq!( - scatter_serial( - &[10, 11], - &[0, UNASSIGNED_LEADER, UNASSIGNED_LEADER, 1], - 2, - 2 - ), - [vec![10], vec![11]] - ); + // Given + let points = [10, 11]; + let assignments = [0, UNASSIGNED_LEADER, UNASSIGNED_LEADER, 1]; + let expected_clusters = [vec![10], vec![11]]; + + // When + let actual_clusters = scatter_serial(&points, &assignments, 2, 2); + + // Then + assert_eq!(actual_clusters, expected_clusters); } #[test] fn parallel_scatter_matches_serial_order() { + // Given let points: Vec = (0..PARALLEL_SCATTER_MIN_POINTS as u32).collect(); let assignments: Vec = points .iter() .flat_map(|point| [point % 7, (point + 3) % 7]) .collect(); - let expected = scatter_serial(&points, &assignments, 2, 7); - let actual = scatter_assignments(&points, &assignments, 2, 7).unwrap(); + // When + let expected_serial_clusters = scatter_serial(&points, &assignments, 2, 7); + let actual_parallel_clusters = scatter_assignments(&points, &assignments, 2, 7).unwrap(); - assert_eq!(actual, expected); + // Then + assert_eq!(actual_parallel_clusters, expected_serial_clusters); } }