diff --git a/Cargo.lock b/Cargo.lock index b28f10fe4..0ff14ce18 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -438,6 +438,7 @@ dependencies = [ "anyhow", "bytemuck", "dashmap", + "diskann-linalg", "diskann-utils", "diskann-vector", "diskann-wide", @@ -447,6 +448,7 @@ dependencies = [ "num-traits", "pin-project", "rand", + "rayon", "relative-path 2.0.1", "serde", "serde_json", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 72911c910..697cf1eb3 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 } @@ -57,7 +59,7 @@ panic = "warn" default = ["tracing"] # Enable PiPNN batch graph construction. -pipnn = [] +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 000000000..55d86103e --- /dev/null +++ b/diskann/src/graph/pipnn/finalization.rs @@ -0,0 +1,293 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! 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. 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 +//! source vectors and metric distances. + +use crate::{ + ANNError, ANNResult, + graph::{ + AdjacencyList, Config, + internal::{SortedNeighbors, prune}, + }, + neighbor::Neighbor, + utils::VectorRepr, +}; +use diskann_utils::views::MatrixView; +use diskann_vector::{DistanceFunction, distance::Metric}; +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_candidates`. +#[derive(Default)] +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. +pub(crate) fn prune_overfull( + data: MatrixView<'_, T>, + candidates: Vec>, + graph: &Config, + metric: Metric, +) -> ANNResult>> +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())); + + // `build_graph` runs this Rayon operation in the pool from the build context. + #[allow(clippy::disallowed_methods)] + candidates + .into_par_iter() + .enumerate() + .map_init( + PruneWorkspace::default, + |workspace, (source, mut source_candidates)| { + // 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); + } + + 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| { + Neighbor::new( + candidate, + distance + .evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + + 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.candidate_slots.clear(); + workspace + .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.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 + .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()); + + let selected = prune::robust_prune( + &sorted, + &workspace.candidate_slots, + workspace.prune_states.as_mut_slice(), + degree, + graph.alpha(), + graph.prune_kind(), + |left, right| { + distance.evaluate_similarity( + data.row(*left as usize), + data.row(*right as usize), + ) + }, + ); + + let mut guard = source_candidates.resize(selected); + for (destination, state) in guard.iter_mut().zip(workspace.prune_states.iter()) { + *destination = *sorted[state.neighbor as usize].id(); + } + guard.finish(selected); + Ok(source_candidates) + }, + ) + .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::{ + 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 new file mode 100644 index 000000000..25d7df817 --- /dev/null +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -0,0 +1,891 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Leaf-local graph construction and candidate accumulation. +//! +//! 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. +//! +//! 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 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::KernelMetric, + leaf_kernel::{ + LeafKernelError, LeafKernelWorkspace, LeafNeighbor, leaf_neighbor_count, nearest_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, + 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, + additional: usize, + #[source] + source: TryReserveError, + }, + #[error("failed to convert point {point} in leaf {leaf}")] + Conversion { + leaf: usize, + point: u32, + #[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, + }, + #[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. +/// +/// 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, + dots: Vec, + neighbors: Vec, + local_adjacency: Vec>, + kernel_workspace: LeafKernelWorkspace, +} + +impl LeafBuffers { + fn prepare( + &mut self, + leaf: usize, + point_count: usize, + dimension_count: usize, + requested_k: usize, + ) -> Result<(usize, usize), LeafBuildError> { + 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 = + point_count + .checked_mul(leaf_k) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: point_count, + 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(), + )?; + 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))?; + self.local_adjacency + .resize_with(point_count, AdjacencyList::new); + self.local_adjacency[..point_count] + .iter_mut() + .for_each(AdjacencyList::clear); + Ok(()) + } +} + +/// Concurrent candidate lists indexed by global point ID. +/// +/// 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>>, +} + +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 add_leaf( + &self, + point_ids: &[u32], + local_adjacency: &[AdjacencyList], + ) -> Result<(), LeafBuildError> { + 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 + .lock() + .map_err(|_| poisoned_candidate_list(source))?; + candidates.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 = candidates + .into_inner() + .map_err(|_| poisoned_candidate_list(point as u32))?; + candidates.sort(); + output.push(candidates); + } + Ok(output) + } +} + +/// Build direct graph candidates from all overlapping leaves. +/// +/// 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, + 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); + } + 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)| { + add_direct_leaf_candidates::( + arch, + data, + leaf, + point_ids, + requested_k, + buffers, + &candidates, + ) + }, + )?; + 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. +fn add_direct_leaf_candidates( + arch: A, + data: MatrixView<'_, T>, + leaf: usize, + point_ids: &[u32], + requested_k: usize, + 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() { + 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 { + return Ok(()); + } + + let point_value_count = point_ids.len() * data.ncols(); + let dot_count = point_ids.len() * point_ids.len(); + + for (&point, point_output) in point_ids + .iter() + .zip(buffers.point_values[..point_value_count].chunks_exact_mut(data.ncols())) + { + 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( + 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", + })?; + let output = MutMatrixView::try_from( + &mut buffers.neighbors[..neighbor_value_count], + point_ids.len(), + leaf_k, + ) + .map_err(|_| LeafBuildError::InvalidView { + leaf, + buffer: "leaf output", + })?; + nearest_neighbors::(arch, dots, output, &mut buffers.kernel_workspace) + .map_err(|source| LeafBuildError::Kernel { leaf, source })?; + + 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()]) +} + +fn add_symmetric_neighbors( + point_ids: &[u32], + 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]; + 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> { + 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, + } +} + +fn poisoned_candidate_list(point: u32) -> LeafBuildError { + LeafBuildError::PoisonedCandidateList { point } +} + +#[cfg(test)] +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; + + 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() + } + + 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], + k: usize, + metric: Metric, + ) -> Result>, LeafBuildError> + where + T: crate::utils::VectorRepr + 'static, + { + pool().install(|| { + arch::dispatch1_no_features( + DispatchLeafBuild(metric), + LeafBuildCall { + data, + leaves: leaves.to_vec(), + k, + }, + ) + }) + } + + 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)) + })); + } + + 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 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] + 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]; + 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 1c51887d8..4a6dabcae 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,29 +3,659 @@ * Licensed under the MIT license. */ -//! Numerical kernels for PiPNN graph construction. +//! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! [`partition_kernel`] converts point-to-leader dot products into sorted leader -//! positions. The output width sets the fanout. One workspace stores the -//! runtime-sized tracker and reuses it for each point. +//! PiPNN builds graph candidates in three steps. A leader is a sampled dataset +//! point that acts as the center of one child partition. A leaf is a bounded +//! child partition used for local neighbor selection. //! -//! [`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. +//! 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. //! -//! `kernel_metric` defines the scalar and SIMD formulas. It also defines the -//! required norm units for each metric. +//! `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. //! -//! The graph builder selects architecture `A` and metric `M` once. It passes -//! these concrete types to both kernels. +//! [`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. //! -//! Each kernel checks all view and scale relationships before unchecked SIMD -//! access. The kernels borrow their matrices. They write only to caller-owned -//! output and workspace. -#[allow(dead_code)] +//! 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 output before it creates another graph representation. + mod kernel_metric; -#[allow(dead_code)] +mod finalization; +mod leaf_build; mod leaf_kernel; -#[allow(dead_code)] mod partition_kernel; +mod partitioning; + +use crate::{ + ANNError, ANNResult, + graph::{AdjacencyList, Config}, + utils::VectorRepr, +}; +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}; + +/// PiPNN partition and leaf-selection policy. +/// +/// 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 child-partition centers. + pub p_samp: f64, + /// 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, + /// Number of independent partition passes over the dataset. + pub replicas: usize, +} + +impl PiPNNConfig { + /// 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")); + } + 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 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!( + "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")); + } + Ok(()) + } +} + +/// PiPNN policy and borrowed execution resources for one 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> { + /// Check and combine PiPNN configuration with DiskANN graph policy. + pub fn new( + config: PiPNNConfig, + graph: &'a Config, + metric: Metric, + pool: &'a ThreadPool, + ) -> ANNResult { + config.validate()?; + 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 one PiPNN adjacency list for each point in `data`. +/// +/// 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`. +/// The build therefore uses norm-aware cosine for these two input types. +pub fn build_graph( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync + 'static, +{ + context + .pool + .install(|| validate_and_dispatch_build(data, context)) +} + +/// Check dataset bounds and select the architecture and metric implementation. +fn validate_and_dispatch_build( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr + Send + Sync + 'static, +{ + if data.nrows() == 0 { + return Err(ANNError::message("PiPNN requires at least one data point")); + } + if data.ncols() == 0 { + return Err(ANNError::message( + "PiPNN requires at least one data dimension", + )); + } + if data.nrows() > u32::MAX as usize { + return Err(config_error(format!( + "dataset point count ({}) exceeds the u32 graph ID limit", + data.nrows() + ))); + } + // 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, + 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) + } + } + } +} + +/// Run the PiPNN graph pipeline for one selected metric implementation. +/// +/// The function builds overlapping leaves, merges direct candidates, and applies +/// final graph-degree pruning. +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::(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.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. + tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, M::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::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, leaf_kernel}; + 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_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); + let pool = pool(2); + let config = PiPNNConfig { + c_max: 5, + c_min: 1, + p_samp: 0.5, + fanout: vec![2], + k: leaf_kernel::MAX_LEAF_NEIGHBORS, + 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)] +#[allow( + clippy::expect_used, + clippy::unwrap_used, + reason = "deterministic test fixture construction must abort on invalid setup" +)] +mod config_tests { + use super::{PiPNNBuildContext, PiPNNConfig, leaf_kernel}; + 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 { + k: 0, + ..pipnn_config() + }, + PiPNNConfig { + k: leaf_kernel::MAX_LEAF_NEIGHBORS + 1, + ..pipnn_config() + }, + PiPNNConfig { + replicas: 0, + ..pipnn_config() + }, + ]; + + for config in &mut cases { + PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .expect_err("invalid PiPNN config must be rejected"); + } + } + + #[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!(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 new file mode 100644 index 000000000..c988b38c3 --- /dev/null +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -0,0 +1,1248 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Deterministic overlapping partition construction for PiPNN. +//! +//! 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 +//! deterministic seed. + +use std::collections::HashSet; + +use crate::{ANNError, ANNResult, utils::VectorRepr}; +use diskann_linalg::Transpose; +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::KernelMetric, + partition_kernel::{ + PartitionInput, PartitionKernelWorkspace, PartitionScales, nearest_leaders, + }, +}; + +// 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; +const ASSIGNMENT_CACHE_TARGET_BYTES: usize = 524_288; +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; + +/// 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")] + 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, + 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("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, + point_scales: Vec, + kernel: PartitionKernelWorkspace, +} + +impl AsPooled<()> for StripeBuffers { + fn create(_: ()) -> Self { + Self::default() + } + + fn modify(&mut self, _: ()) { + // Keep the largest allocation across leases. `assign_point_stripe` defines the + // active prefix before each read. + } +} + +/// Reusable buffers for point-to-leader assignment. +/// +/// `ObjectPool` locks only when it gives or receives a lease. Numerical work +/// holds the lease, not the pool lock. +type StripeBufferPool = ObjectPool; + +/// Build overlapping bounded leaves for all configured replicas. +/// +/// 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>, + 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 { + 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) +} + +/// 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>, + config: &PiPNNConfig, + 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())?; + 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); + } + + let mut leaves = Vec::new(); + let mut work = Vec::new(); + work.try_reserve_exact(1).map_err(ANNError::new)?; + work.push(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. + #[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 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)?; + 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); + }; + Err(ANNError::new(PartitionError::IterationLimit { + size: largest.indices.len(), + level: largest.level, + limit: MAX_PARTITION_ITERATIONS, + })) +} + +/// Split one oversized cluster into child partitions. +/// +/// 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, + 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(); + 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::(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() { + 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)) +} + +/// Sample point IDs that act as centers for one partition split. +fn sample_leaders(points: &[u32], sampling_fraction: f64, seed: u64) -> ANNResult> { + 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) +} + +/// 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) +} + +fn replica_seed(replica: usize) -> u64 { + PARTITION_SEED.wrapping_add((replica as u64).wrapping_mul(REPLICA_SEED_STEP)) +} + +// 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 cluster point to its nearest sampled partition centers. +/// +/// 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>, + point_ids: &[u32], + leader_ids: &[u32], + fanout: usize, + 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(); + 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) { + filled_vec(leader_ids.len(), 0.0f32)? + } else { + Vec::new() + }; + for (scale, leader_vector) in leader_scales + .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(); + } + } + + 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_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 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) + .enumerate() + .try_for_each(|(worker, worker_assignments)| { + let mut buffers = stripe_buffers.get_ref(()); + 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_point_stripe::( + arch, + data, + &point_ids[first_point..first_point + stripe_point_count], + &leader_values, + &leader_scales, + fanout, + &mut buffers, + stripe_assignments, + )?; + } + Ok::<(), ANNError>(()) + })?; + + scatter_assignments(point_ids, &assignments, fanout, leader_ids.len()) +} + +/// 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. +#[inline] +#[allow(clippy::too_many_arguments)] +fn assign_point_stripe( + arch: A, + data: MatrixView<'_, T>, + point_ids: &[u32], + leader_values: &[f32], + leader_scales: &[f32], + 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(); + 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.points, point_values_len, 0.0)?; + grow_fallible(&mut buffers.dots, dots_len, 0.0)?; + let StripeBuffers { + 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]; + gather_vectors(data, point_ids, point_values)?; + diskann_linalg::sgemm( + Transpose::None, + Transpose::Ordinary, + point_count, + leader_count, + dimensions, + 1.0, + point_values, + leader_values, + None, + dots, + ) + .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 + .iter_mut() + .zip(point_values.chunks_exact(dimensions)) + { + *scale = FastL2NormSquared.evaluate(point_values); + } + &*point_scales + } else { + &[] + }; + let scales = match M::METRIC { + Metric::L2 => PartitionScales::L2 { + leader_squared_norms: leader_scales, + }, + 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::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, scales }, + output, + kernel_workspace, + ) + .map_err(ANNError::new) +} + +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)?; + } + Ok(()) +} + +/// Group assigned point IDs by child partition. +/// +/// 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], + fanout: usize, + leaders: usize, +) -> ANNResult>> { + if points.len() < PARALLEL_SCATTER_MIN_POINTS { + return 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. + #[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 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)?; + 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(), + }) + })?; + } + } + + // `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() + .enumerate() + .map(|(leader, size)| { + let mut cluster = Vec::new(); + cluster.try_reserve_exact(size).map_err(ANNError::new)?; + for local in &locals { + cluster.extend_from_slice(&local[leader]); + } + Ok(cluster) + }) + .collect() +} + +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::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, + }) + })?; + } + 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) +} + +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) +} + +/// 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 merge_undersized_leaves( + 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::new)?; + small_leaves + .try_reserve(leaves.len()) + .map_err(ANNError::new)?; + 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::new)?; + + 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(), + }) + })?; + if combined > c_max { + 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)?); + } + } + + if !small.is_empty() { + 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(), + }) + })?; + if combined <= c_max { + last.try_reserve(remainder.len()).map_err(ANNError::new)?; + last.append(&mut remainder); + last.sort_unstable(); + } + } + if !remainder.is_empty() { + merged.push(remainder); + } + } + + 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()); + 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) +} + +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 grow_fallible(values: &mut Vec, len: usize, value: T) -> ANNResult<()> { + if values.len() >= len { + return Ok(()); + } + 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 { + rows.checked_mul(cols) + .ok_or_else(|| ANNError::new(PartitionError::ShapeOverflow { buffer, rows, cols })) +} + +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 { + point_count.next_power_of_two() / 2 + }; + point_count.clamp(MIN_ASSIGNMENT_STRIPE_POINTS, MAX_ASSIGNMENT_STRIPE_POINTS) +} + +#[cfg(test)] +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, + 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_with_runtime_metric( + 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_with_runtime_metric( + 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_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_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_with_runtime_metric(data.as_view(), &config(2, 8, vec![2], 1), Metric::L2) + .unwrap(); + + 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_with_runtime_metric(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 = merge_undersized_leaves(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 = 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]] + ); + } + + #[test] + fn global_merge_fills_exact_capacity_before_flushing() { + let merged = merge_undersized_leaves(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_with_runtime_metric( + data.as_view(), + &config(3, 12, vec![3, 2], 3), + Metric::CosineNormalized, + ) + .unwrap(); + + assert_valid_partition_with_runtime_metric(&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 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| { + 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_with_runtime_metric( + MatrixView::try_from(&f32_data, points, dimensions).unwrap(), + &config, + Metric::L2, + ) + .unwrap(); + 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_with_runtime_metric(&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 + } + + // 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(); + 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::<_, super::super::kernel_metric::L2, _>( + diskann_wide::ARCH, + data, + &[2], + &[0, 1], + 1, + &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_with_runtime_metric(data.as_view(), &config, metric).unwrap(); + assert_valid_partition_with_runtime_metric(&leaves, 64, 20, 1); + } + } + + #[test] + fn leader_count_is_bounded() { + 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] + 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::<_, super::super::kernel_metric::L2, _>( + diskann_wide::ARCH, + data, + &point_ids, + &[0, 2_047], + 1, + &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_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_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, + } + ); + } +}