diff --git a/Cargo.lock b/Cargo.lock index 26286aeda3..3e687081cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -446,8 +446,10 @@ dependencies = [ "half", "hashbrown 0.16.1", "num-traits", + "parking_lot", "pin-project", "rand", + "rayon", "relative-path 2.0.1", "rstest", "serde", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index f36533dcbd..20745a3f69 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -21,7 +21,9 @@ 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 tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } tracing = { workspace = true, optional = true } @@ -59,8 +61,8 @@ panic = "warn" [features] default = ["tracing"] -# Enable PiPNN numerical kernels. -pipnn = ["dep:diskann-linalg"] +# Enable PiPNN batch graph construction. +pipnn = ["dep:diskann-linalg", "dep:parking_lot", "dep:rayon", "diskann-utils/rayon", "tracing"] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/pipnn/conversion.rs b/diskann/src/graph/pipnn/conversion.rs new file mode 100644 index 0000000000..e7860db6d5 --- /dev/null +++ b/diskann/src/graph/pipnn/conversion.rs @@ -0,0 +1,174 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Gather dataset rows into packed `f32` storage. +//! +//! `f16` rows use the FP16 instructions of the running CPU. The generic +//! `VectorRepr::as_f32_into` path uses the instructions of the build target +//! instead. Crates that depend on `diskann` do not inherit this repository's +//! `x86-64-v3` target flag, so for them that path converts `f16` in software, +//! which is many times slower. +//! +//! Other element types use `VectorRepr::as_f32_into` for every target. `f32` +//! rows are copied. The baseline x86-64 target has SSE2 instructions that +//! convert `u8` and `i8` to `f32`, so the compiler vectorizes that loop without +//! the `x86-64-v3` flag. The baseline target has no `f16` conversion +//! instruction; F16C starts at `x86-64-v3`. + +use std::any::TypeId; + +use crate::{ANNResult, utils::VectorRepr}; +use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_vector::conversion::SliceCast; +use diskann_wide::{ + Architecture, + arch::{self, Target2, Target3}, +}; +use half::f16; + +/// Return the length of a row of `data` after conversion to `f32`. +/// +/// Compressed element types store a row in a different number of elements than +/// the `f32` vector it converts to. +pub(super) fn f32_dimensions(data: MatrixView<'_, T>) -> ANNResult { + T::full_dimension(data.row(0)).map_err(Into::into) +} + +/// Convert rows `ids` of `data` into the rows of `output`, one row per ID. +/// +/// `output` has [`f32_dimensions`] columns. `f16` input selects the CPU's +/// instructions once for the whole batch. +pub(super) fn gather_as_f32( + data: MatrixView<'_, T>, + ids: &[u32], + mut output: MutMatrixView<'_, f32>, +) -> ANNResult<()> { + debug_assert_eq!(output.nrows(), ids.len()); + if TypeId::of::() == TypeId::of::() { + arch::dispatch3(GatherFp16, data, ids, output); + } else { + for (&id, destination) in ids.iter().zip(output.row_iter_mut()) { + T::as_f32_into(data.row(id as usize), destination).map_err(Into::into)?; + } + } + Ok(()) +} + +struct GatherFp16; + +impl Target3, &[u32], MutMatrixView<'_, f32>> for GatherFp16 +where + A: Architecture, + T: VectorRepr, + for<'a, 'b> SliceCast: Target2, +{ + #[inline(always)] + fn run( + self, + arch: A, + data: MatrixView<'_, T>, + ids: &[u32], + mut output: MutMatrixView<'_, f32>, + ) { + for (&id, destination) in ids.iter().zip(output.row_iter_mut()) { + let source = bytemuck::cast_slice(data.row(id as usize)); + SliceCast::::new().run(arch, destination, source); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gathered_rows_follow_requested_ids() { + // Row IDs are out of order and contain a repeat. + fn check(values: [T; 6], expected: [f32; 6]) { + let data = MatrixView::try_from(&values[..], 3, 2).unwrap(); + let mut output = [f32::NAN; 6]; + + gather_as_f32( + data, + &[2, 0, 2], + MutMatrixView::try_from(output.as_mut_slice(), 3, 2).unwrap(), + ) + .unwrap(); + + assert_eq!(output, expected, "{}", std::any::type_name::()); + } + + check( + [1.5_f32, -2.0, 7.0, 4.0, 0.5, 9.0], + [0.5, 9.0, 1.5, -2.0, 0.5, 9.0], + ); + check([1_i8, -2, 7, 4, -5, 9], [-5.0, 9.0, 1.0, -2.0, -5.0, 9.0]); + check([1_u8, 2, 7, 4, 5, 9], [5.0, 9.0, 1.0, 2.0, 5.0, 9.0]); + check( + [1.5_f32, -2.0, 7.0, 4.0, 0.5, 9.0].map(f16::from_f32), + [0.5, 9.0, 1.5, -2.0, 0.5, 9.0], + ); + } + + #[test] + fn fp16_gather_converts_every_coordinate() { + // x86 converts eight values at a time. The dimensions cover a short row, + // full vectors with and without a tail, and an embedding with a tail. + for dimensions in [7, 8, 9, 32, 33, 1537] { + let mut values: Vec<_> = (0..3 * dimensions) + .map(|index| f16::from_f32((index % 31) as f32 - 15.0)) + .collect(); + for (row, values) in values.chunks_exact_mut(dimensions).enumerate() { + values[dimensions - 1] = f16::from_f32(100.0 + row as f32); + } + let ids = [2, 0, 2]; + let expected: Vec<_> = ids + .iter() + .flat_map(|&row| { + values[row as usize * dimensions..(row as usize + 1) * dimensions] + .iter() + .map(|value| value.to_f32()) + }) + .collect(); + let mut output = vec![f32::NAN; expected.len()]; + + gather_as_f32( + MatrixView::try_from(values.as_slice(), 3, dimensions).unwrap(), + &ids, + MutMatrixView::try_from(output.as_mut_slice(), 3, dimensions).unwrap(), + ) + .unwrap(); + + assert_eq!(output, expected, "{dimensions} dimensions"); + } + } + + #[test] + fn fp16_gather_preserves_special_values_and_zero_signs() { + let values = [ + f16::ZERO, + f16::NEG_ZERO, + f16::INFINITY, + f16::NEG_INFINITY, + f16::from_bits(1), + f16::NAN, + ]; + let mut output = [42.0; 6]; + + gather_as_f32( + MatrixView::try_from(&values[..], 2, 3).unwrap(), + &[1, 0], + MutMatrixView::try_from(output.as_mut_slice(), 2, 3).unwrap(), + ) + .unwrap(); + + assert_eq!(output[0], f32::NEG_INFINITY); + assert_eq!(output[1], 2.0_f32.powi(-24)); + assert!(output[2].is_nan()); + assert_eq!(output[3].to_bits(), 0.0_f32.to_bits()); + assert_eq!(output[4].to_bits(), (-0.0_f32).to_bits()); + assert_eq!(output[5], f32::INFINITY); + } +} diff --git a/diskann/src/graph/pipnn/finalization.rs b/diskann/src/graph/pipnn/finalization.rs new file mode 100644 index 0000000000..a36ec7ac43 --- /dev/null +++ b/diskann/src/graph/pipnn/finalization.rs @@ -0,0 +1,299 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Final pruning: cut each candidate list to the graph degree with RobustPrune. +//! +//! A list within the degree passes through unchanged. A longer list is sorted by +//! distance to its source point and pruned with the shared Vamana RobustPrune, +//! which owns occlusion and the alpha rounds. The selected IDs overwrite the list +//! in place. + +use crate::{ + graph::{ + AdjacencyList, Config, + internal::{SortedNeighbors, prune}, + }, + neighbor::Neighbor, + utils::VectorRepr, +}; +use diskann_utils::views::MatrixView; +use diskann_vector::{DistanceFunction, distance::Metric}; +use rayon::prelude::*; + +/// RobustPrune scratch for one Rayon job. +#[derive(Default)] +struct PruneWorkspace { + candidates: Vec>, + mapped_candidates: Vec>>, + states: Vec, +} + +impl PruneWorkspace { + /// Prune the candidates of point `source` to the graph degree. + fn prune( + &mut self, + data: MatrixView<'_, T>, + source: usize, + mut source_candidates: AdjacencyList, + graph: &Config, + distance: &T::Distance, + ) -> AdjacencyList { + let degree = graph.pruned_degree().get(); + if source_candidates.len() <= degree { + return source_candidates; + } + + let source_vector = data.row(source); + self.candidates.clear(); + self.candidates + .extend(source_candidates.iter().copied().map(|candidate| { + Neighbor::new( + candidate, + distance.evaluate_similarity(source_vector, data.row(candidate as usize)), + ) + })); + // Like Vamana, prune only the `max_occlusion_size` nearest candidates. The + // graph configuration keeps this size within the `u16` positions of + // RobustPrune. + // + // Sort before mapping: changing the element type can change the unstable + // sort's order for equal distances, and therefore RobustPrune's choices. + let sorted = SortedNeighbors::new(&mut self.candidates, graph.max_occlusion_size().get()); + let mapped = sorted.map_in(&mut self.mapped_candidates, |&candidate| { + (candidate as usize != source).then_some(candidate) + }); + self.states.clear(); + self.states.resize(sorted.len(), prune::State::default()); + let selected = prune::robust_prune( + mapped, + &mut self.states, + degree, + graph.alpha(), + graph.prune_kind(), + |left, right| { + distance.evaluate_similarity(data.row(*left as usize), data.row(*right as usize)) + }, + ); + + let mut output = source_candidates.resize(selected); + for (destination, state) in output.iter_mut().zip(&self.states) { + *destination = *sorted[state.neighbor as usize].id(); + } + output.finish(selected); + source_candidates + } +} + +/// Prune each candidate list that is longer than the graph degree. +/// +/// `candidates` holds one list of dataset IDs for each data row. +pub(crate) fn prune_overfull( + data: MatrixView<'_, T>, + candidates: Vec>, + graph: &Config, + metric: Metric, +) -> Vec> { + let distance = T::distance(metric, Some(data.ncols())); + + // This runs in the pool of the build context (see `build_graph`). + candidates + .into_par_iter() + .enumerate() + .map_init( + PruneWorkspace::default, + |workspace, (source, candidates)| { + workspace.prune(data, source, candidates, graph, &distance) + }, + ) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::graph::{config, pipnn::test_support::thread_pool}; + + fn pruning_config( + degree: usize, + metric: Metric, + alpha: f32, + ) -> Result { + config::Builder::new_with(degree, config::MaxDegree::same(), 16, metric.into(), |b| { + b.alpha(alpha); + }) + .build() + } + + fn candidates_in_order(ids: &[u32]) -> AdjacencyList { + let mut candidates = AdjacencyList::new(); + candidates.extend_from_slice(ids); + candidates + } + + /// Prune the candidates of source 0 in one-dimensional L2 data. + fn prune_l2(values: &[f32], ids: &[u32], degree: usize, alpha: f32) -> Vec { + let data = MatrixView::try_from(values, values.len(), 1).unwrap(); + let graph = pruning_config(degree, Metric::L2, alpha).unwrap(); + let distance = f32::distance(Metric::L2, Some(1)); + let pruned = + PruneWorkspace::default().prune(data, 0, candidates_in_order(ids), &graph, &distance); + pruned.to_vec() + } + + #[test] + fn candidates_within_degree_keep_their_order() { + let values = [0.0_f32, 1.0, 3.0, -2.0]; + + for ids in [&[][..], &[2][..], &[3, 1][..]] { + assert_eq!(prune_l2(&values, ids, 2, 1.0), ids); + } + } + + #[test] + fn overfull_candidates_follow_the_metric_pruning_rule() { + // Source 0 has candidates [3, 2, 1]; ID 1 is nearest for every metric. + // L2 occludes ID 2: 9 / 4 > 1. + // Cosine retains ID 2: equal directions give an occlusion ratio of 1. + // Normalized cosine retains ID 3: 1 / 1.6 < 1. + // Inner product occludes ID 2: its dot with ID 1 is 6, versus 2 with the source. + for (metric, values, expected) in [ + ( + Metric::L2, + [0.0, 0.0, 1.0, 0.0, 3.0, 0.0, -2.0, 0.0], + [1, 3], + ), + ( + Metric::Cosine, + [1.0, 0.0, 2.0, 0.0, 1.0, 1.0, -1.0, 0.0], + [1, 2], + ), + ( + Metric::CosineNormalized, + [1.0, 0.0, 0.8, 0.6, -1.0, 0.0, 0.0, -1.0], + [1, 3], + ), + ( + Metric::InnerProduct, + [1.0, 0.0, 3.0, 0.0, 2.0, 1.0, -1.0, 4.0], + [1, 3], + ), + ] { + let data = MatrixView::try_from(&values[..], 4, 2).unwrap(); + let graph = pruning_config(2, metric, 1.0).unwrap(); + let distance = f32::distance(metric, Some(2)); + + let actual = PruneWorkspace::default().prune( + data, + 0, + candidates_in_order(&[3, 2, 1]), + &graph, + &distance, + ); + + assert_eq!(&*actual, expected, "{metric:?}"); + } + } + + #[test] + fn pruning_excludes_the_source_without_shifting_selected_ids() { + assert_eq!( + prune_l2(&[0.0, 1.0, 3.0, -2.0], &[3, 0, 2, 1], 2, 1.0), + [1, 3] + ); + } + + #[test] + fn graph_alpha_controls_which_occluded_neighbors_return() { + // After selecting x=1, x=3 has ratio 9/4 and x=4 has ratio 16/9. + // Alpha 2 admits x=4 while x=3 remains occluded. + let values = [0.0_f32, 1.0, 3.0, 4.0]; + + for (alpha, expected) in [(1.0, &[1][..]), (2.0, &[1, 3][..])] { + assert_eq!( + prune_l2(&values, &[3, 2, 1], 2, alpha), + expected, + "alpha {alpha}" + ); + } + } + + #[test] + fn a_reused_workspace_does_not_carry_state_between_sources() { + let values = [0.0_f32, 1.0, 3.0, -2.0, 9.0]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + let graph = pruning_config(2, Metric::L2, 1.0).unwrap(); + let distance = f32::distance(Metric::L2, Some(1)); + let mut workspace = PruneWorkspace::default(); + + // One workspace serves a large list, a smaller list, then a large list again. + for (source, ids, expected) in [ + (0, &[4, 3, 2, 1][..], &[1, 3][..]), + (4, &[3, 2, 1][..], &[2][..]), + (1, &[4, 3, 2, 0][..], &[0, 2][..]), + ] { + let actual = workspace.prune(data, source, candidates_in_order(ids), &graph, &distance); + + assert_eq!(&*actual, expected, "source {source}"); + } + } + + #[test] + fn pruning_reads_only_the_max_occlusion_size_nearest_candidates() { + // Source 0 is at x=0. Candidates 2 and 3 lie on one side and candidate 1 + // on the other, so 2 occludes 3 but not 1. A pool of two candidates + // drops 1, the farthest. + let values = [0.0_f32, -3.0, 1.0, 2.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let distance = f32::distance(Metric::L2, Some(1)); + for (max_occlusion_size, expected) in [(2, &[2][..]), (3, &[2, 1][..])] { + let graph = config::Builder::new_with( + 2, + config::MaxDegree::same(), + 16, + Metric::L2.into(), + |b| { + b.alpha(1.0).max_occlusion_size(max_occlusion_size); + }, + ) + .build() + .unwrap(); + + let actual = PruneWorkspace::default().prune( + data, + 0, + candidates_in_order(&[1, 2, 3]), + &graph, + &distance, + ); + + assert_eq!( + &*actual, expected, + "max_occlusion_size {max_occlusion_size}" + ); + } + } + + #[test] + fn parallel_pruning_keeps_each_result_with_its_source() { + let values = [0.0_f32, 1.0, 4.0, 9.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let graph = pruning_config(1, Metric::L2, 1.0).unwrap(); + + for workers in [1, 3] { + let candidates = [[3, 2, 1], [3, 2, 0], [3, 1, 0], [2, 1, 0]] + .map(|ids| candidates_in_order(&ids)) + .to_vec(); + + let actual = thread_pool(workers) + .install(|| prune_overfull(data, candidates, &graph, Metric::L2)); + + assert_eq!( + actual.into_iter().map(Vec::from).collect::>(), + [vec![1], vec![0], vec![1], vec![2]], + "{workers} workers" + ); + } + } +} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs new file mode 100644 index 0000000000..c2d112432f --- /dev/null +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -0,0 +1,426 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Leaf construction: select the nearest neighbors inside each leaf and merge +//! them into one candidate list per data point. +//! +//! A leaf job gathers the vectors of its points as `f32`. The leaf kernel then +//! selects the `k` nearest leaf points of each point, and the job adds each +//! selected pair in both directions. Leaves overlap and run in parallel, so each +//! data point has a locked candidate list. A job groups its edges by point first +//! and then locks each list once. + +use parking_lot::Mutex; + +use crate::{ANNResult, graph::AdjacencyList, utils::VectorRepr}; +use diskann_utils::views::{MatrixView, MutMatrixView}; +use rayon::prelude::*; + +use super::{ + conversion::{f32_dimensions, gather_as_f32}, + leaf_kernel::{LeafKernelWorkspace, select_leaf_neighbors}, + leaf_metric::LeafMetric, + simd::Simd, + topk::Candidate, +}; + +/// Scratch for one Rayon leaf job. +/// +/// Each buffer keeps the largest size that the job needed. A leaf uses a prefix. +#[derive(Default)] +struct LeafBuffers { + point_values: Vec, + neighbors: Vec, + local_adjacency: Vec>, + kernel_workspace: LeafKernelWorkspace, +} + +impl LeafBuffers { + /// Grow the buffers for a leaf of `point_count` points and return the + /// effective `k` of the leaf. + /// + /// The buffers never shrink. A smaller leaf uses a prefix, so a later large + /// leaf does not fill the same slots again. + fn prepare(&mut self, point_count: usize, dimension_count: usize, requested_k: usize) -> usize { + // A point has at most `point_count - 1` other points in its leaf. Wider rows + // would hold only empty slots. + let leaf_k = requested_k.min(point_count.saturating_sub(1)); + let value_count = point_count * dimension_count; + if self.point_values.len() < value_count { + self.point_values.resize(value_count, 0.0); + } + let neighbor_count = point_count * leaf_k; + if self.neighbors.len() < neighbor_count { + self.neighbors.resize(neighbor_count, Candidate::default()); + } + leaf_k + } + + /// Clear the edge lists of the first `point_count` leaf points. + fn prepare_local_adjacency(&mut self, point_count: usize) { + if self.local_adjacency.len() < point_count { + self.local_adjacency.resize_with(point_count, Vec::new); + } + self.local_adjacency[..point_count] + .iter_mut() + .for_each(Vec::clear); + } +} + +/// Build one candidate list per data point from the nearest neighbors in all +/// leaves. +/// +/// Each selected pair adds both directions. A list holds global IDs without +/// duplicates, in the order in which the parallel jobs add them. The lists stay +/// unsorted: final pruning ranks candidates by distance, and a graph search reads +/// all neighbors of a point, so the order only breaks distance ties. +pub(super) fn build_leaf_candidates( + arch: A, + data: MatrixView<'_, T>, + leaves: Vec>, + requested_k: usize, +) -> ANNResult>> +where + A: Simd, + M: LeafMetric, + T: VectorRepr, +{ + let candidates: Vec<_> = (0..data.nrows()) + .map(|_| Mutex::new(AdjacencyList::new())) + .collect(); + let dimensions = f32_dimensions(data)?; + leaves + .par_iter() + .try_for_each_init(LeafBuffers::default, |buffers, point_ids| { + add_leaf_candidates::( + arch, + data, + dimensions, + point_ids, + requested_k, + buffers, + &candidates, + ) + })?; + Ok(candidates.into_iter().map(Mutex::into_inner).collect()) +} + +/// Add the selected pairs of one leaf to the candidate lists. +/// +/// The IDs of a leaf are distinct dataset rows, so `points x dimensions` is not +/// larger than the dataset. +fn add_leaf_candidates( + arch: A, + data: MatrixView<'_, T>, + dimensions: usize, + point_ids: &[u32], + requested_k: usize, + buffers: &mut LeafBuffers, + candidates: &[Mutex>], +) -> ANNResult<()> +where + A: Simd, + M: LeafMetric, + T: VectorRepr, +{ + let point_count = point_ids.len(); + let leaf_k = buffers.prepare(point_count, dimensions, requested_k); + if leaf_k == 0 { + return Ok(()); + } + + buffers.prepare_local_adjacency(point_count); + let mut points = MutMatrixView::try_from( + &mut buffers.point_values[..point_count * dimensions], + point_count, + dimensions, + )?; + gather_as_f32(data, point_ids, points.as_mut_view())?; + let mut output = MutMatrixView::try_from( + &mut buffers.neighbors[..point_count * leaf_k], + point_count, + leaf_k, + )?; + select_leaf_neighbors::( + arch, + points.as_view(), + output.as_mut_view(), + &mut buffers.kernel_workspace, + )?; + + add_symmetric_neighbors( + point_ids, + output.as_view(), + &mut buffers.local_adjacency[..point_count], + ); + for (&point_id, additions) in point_ids.iter().zip(&buffers.local_adjacency) { + candidates[point_id as usize] + .lock() + .extend_from_slice(additions); + } + Ok(()) +} + +/// Convert each selected pair from leaf positions to global IDs and add it to +/// the edge lists of both points. +fn add_symmetric_neighbors( + point_ids: &[u32], + neighbors: MatrixView<'_, Candidate>, + local_adjacency: &mut [Vec], +) { + for (source, source_neighbors) in neighbors.row_iter().enumerate() { + for neighbor in source_neighbors { + // Skip the empty slots that the kernel leaves for NaN and infinite + // distances. + if !neighbor.is_assigned() { + continue; + } + let target = neighbor.local_idx as usize; + // The kernel never selects a point for itself, and the IDs of a leaf + // are unique, so no list gets a self edge. + local_adjacency[source].push(point_ids[target]); + local_adjacency[target].push(point_ids[source]); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::graph::pipnn::{ + L2, + test_support::{dense_points, sorted_members_per_row, thread_pool}, + }; + use diskann_wide::ARCH; + use half::f16; + use std::collections::BTreeSet; + + fn build(values: &[f32], leaves: Vec>, k: usize, workers: usize) -> Vec> { + let data = MatrixView::try_from(values, values.len(), 1).unwrap(); + let candidates = thread_pool(workers) + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves, k)) + .unwrap(); + sorted_members_per_row(&candidates.into_iter().map(Vec::from).collect::>()) + } + + #[test] + fn selected_pairs_add_global_ids_in_both_directions() { + // The leaf lists IDs [5, 1, 3], at coordinates [9, 0, 2]. + // The directed choices are 1 -> 3, 3 -> 1 and 5 -> 3. + let values = [100.0_f32, 0.0, -100.0, 2.0, 200.0, 9.0]; + + let actual = build(&values, vec![vec![5, 1, 3]], 1, 1); + + assert_eq!( + actual, + [vec![], vec![3], vec![], vec![1, 5], vec![], vec![3]] + ); + } + + #[test] + fn overlapping_leaves_merge_each_neighbor_once() { + let values = [0.0_f32, 1.0, 4.0, 9.0, 16.0]; + let leaves = vec![vec![0, 1, 2], vec![1, 2, 3], vec![2, 3, 4], vec![0, 1, 2]]; + + for workers in [1, 3] { + let actual = build(&values, leaves.clone(), 1, workers); + + assert_eq!( + actual, + [vec![1], vec![0, 2], vec![1, 3], vec![2, 4], vec![3]], + "{workers} workers" + ); + } + } + + #[test] + fn leaves_without_pairs_add_no_candidates() { + let values = [0.0_f32, 1.0, 4.0]; + + for (leaves, k) in [ + (vec![], 2), + (vec![vec![]], 2), + (vec![vec![2]], 2), + (vec![vec![0, 1, 2]], 0), + ] { + let case = format!("leaves {leaves:?}, k={k}"); + + let actual = build(&values, leaves, k, 1); + + assert_eq!(actual, [Vec::::new(), vec![], vec![]], "{case}"); + } + } + + #[test] + fn k_above_the_leaf_size_selects_every_other_point() { + let values = [100.0_f32, 0.0, -100.0, 2.0, 200.0, 9.0]; + + let actual = build(&values, vec![vec![1, 3, 5]], 99, 1); + + assert_eq!( + actual, + [vec![], vec![3, 5], vec![], vec![1, 5], vec![], vec![1, 3]] + ); + } + + #[test] + fn unrankable_pairs_add_no_candidates() { + let values = [0.0_f32, 3.0, f32::NAN]; + + let actual = build(&values, vec![vec![0, 1, 2]], 2, 1); + + assert_eq!(actual, [vec![1], vec![0], vec![]]); + } + + #[test] + fn candidates_match_the_leaf_kernel_on_each_leaf() { + // The kernel tests cover metrics, dimensions, widths and architectures. + // This test covers what leaf build adds: the gather of each element type, + // k cut to the leaf size, buffers reused across leaf shapes, and the merge + // of overlapping leaves on four workers. + let dimensions = 24; + let values = dense_points(500, dimensions, 1290); + let half: Vec<_> = values.iter().map(|&x| f16::from_f32(x)).collect(); + // Dense points are multiples of 1/8 in [-2, 14], so eight times them fits + // `i8`, and an offset of 16 makes them fit `u8`. + let signed: Vec<_> = values.iter().map(|&x| (8.0 * x) as i8).collect(); + let unsigned: Vec<_> = values.iter().map(|&x| (8.0 * x + 16.0) as u8).collect(); + // Leaves of 1 to 40 points put each k above, equal to and below the other + // points of a leaf. + for k in 1..=10 { + check_against_kernel(&values, dimensions, k); + check_against_kernel(&half, dimensions, k); + check_against_kernel(&signed, dimensions, k); + check_against_kernel(&unsigned, dimensions, k); + } + } + + /// Compare `build_leaf_candidates` with the leaf kernel run on each leaf + /// alone. The kernel sees the same rows in the same order, so ties break the + /// same way and the check is exact. + fn check_against_kernel(values: &[T], dimensions: usize, k: usize) { + let point_count = values.len() / dimensions; + let data = MatrixView::try_from(values, point_count, dimensions).unwrap(); + // Sixty overlapping leaves of 1 to 40 points. + let leaves: Vec> = (0..60u32) + .map(|leaf| { + let ids: BTreeSet<_> = (0..1 + leaf % 40) + .map(|i| (leaf * 7 + i * 13) % point_count as u32) + .collect(); + ids.into_iter().collect() + }) + .collect(); + + let mut expected = vec![BTreeSet::new(); point_count]; + for leaf in &leaves { + let width = k.min(leaf.len() - 1); + if width == 0 { + continue; + } + // Convert with `VectorRepr`, not with the gather under test. + let mut rows = vec![0.0; leaf.len() * dimensions]; + for (&id, row) in leaf.iter().zip(rows.chunks_exact_mut(dimensions)) { + T::as_f32_into(data.row(id as usize), row).unwrap(); + } + let mut output = vec![Candidate::default(); leaf.len() * width]; + select_leaf_neighbors::<_, L2>( + ARCH, + MatrixView::try_from(rows.as_slice(), leaf.len(), dimensions).unwrap(), + MutMatrixView::try_from(output.as_mut_slice(), leaf.len(), width).unwrap(), + &mut LeafKernelWorkspace::default(), + ) + .unwrap(); + for (source, row) in output.chunks_exact(width).enumerate() { + for neighbor in row.iter().filter(|neighbor| neighbor.is_assigned()) { + let (source, target) = (leaf[source], leaf[neighbor.local_idx as usize]); + expected[source as usize].insert(target); + expected[target as usize].insert(source); + } + } + } + let expected: Vec> = expected + .into_iter() + .map(|ids| ids.into_iter().collect()) + .collect(); + + let actual = thread_pool(4) + .install(|| build_leaf_candidates::<_, L2, _>(ARCH, data, leaves, k)) + .unwrap(); + + let actual: Vec> = actual.into_iter().map(Vec::from).collect(); + assert_eq!( + sorted_members_per_row(&actual), + expected, + "{}, k={k}", + std::any::type_name::() + ); + } + + #[test] + fn reused_buffers_do_not_carry_edges_between_leaves() { + let values = [0.0_f32, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0]; + let data = MatrixView::try_from(&values[..], 7, 1).unwrap(); + let mut buffers = LeafBuffers::default(); + + // The same worker processes a leaf, a smaller leaf, then more neighbors + // per point. Every call must use only its current IDs and active rows. + for (ids, requested_k, expected) in [ + ( + vec![0, 1, 2, 3], + 1, + vec![ + vec![1], + vec![0, 2], + vec![1, 3], + vec![2], + vec![], + vec![], + vec![], + ], + ), + ( + vec![4, 6], + 99, + vec![vec![], vec![], vec![], vec![], vec![6], vec![], vec![4]], + ), + ( + vec![1, 3, 5], + 2, + vec![ + vec![], + vec![3, 5], + vec![], + vec![1, 5], + vec![], + vec![1, 3], + vec![], + ], + ), + ] { + let candidates: Vec<_> = (0..7).map(|_| Mutex::new(AdjacencyList::new())).collect(); + + add_leaf_candidates::<_, L2, _>( + ARCH, + data, + 1, + &ids, + requested_k, + &mut buffers, + &candidates, + ) + .unwrap(); + + let actual: Vec<_> = candidates + .into_iter() + .map(|row| Vec::from(row.into_inner())) + .collect(); + assert_eq!( + sorted_members_per_row(&actual), + expected, + "leaf {ids:?}, k={requested_k}" + ); + } + } +} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 882788dc79..cd1669016b 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -3,35 +3,50 @@ * Licensed under the MIT license. */ -//! Numerical kernels for PiPNN graph construction. +//! Provider-independent [PiPNN](https://arxiv.org/html/2602.21247v1) graph construction. //! -//! Each kernel computes a distance matrix with metric `M`, then selects the nearest -//! candidates from it with SIMD architecture `A`: +//! PiPNN builds a graph without graph search, in three steps. A leaf is a small +//! cluster of at most `c_max` points; the build selects neighbors only inside +//! leaves. //! -//! - `partition_kernel` assigns each point to its nearest partition leaders. It -//! uses `partition_metric`. -//! - `leaf_kernel` finds the nearest neighbors of each point in a leaf. It uses -//! `leaf_metric`. -//! - `topk` holds the selection code of both kernels. `simd` sets its vector width. +//! 1. `partitioning` splits the dataset into overlapping leaves. +//! 2. `leaf_build` selects the `leaf_k` nearest neighbors of each point inside +//! each leaf and merges them into one candidate list per point. +//! 3. `finalization` prunes each list that is longer than the graph degree with +//! Vamana RobustPrune. //! -//! The metric code does not use `A`, because the GEMM and norm routines select -//! their own SIMD code. Callers choose the metric once per graph build, so each -//! kernel compiles for one metric and has no run-time metric check. - -#![expect( - dead_code, - reason = "graph construction integrates these kernels in the next PR" -)] +//! [`build_graph`] selects the SIMD architecture and the metric once, so the +//! steps compile for one architecture and one metric. It borrows one contiguous +//! [`MatrixView`] and returns one adjacency list for each point. Each step takes +//! ownership of the previous output and frees it when it returns. +//! +//! The build does not load providers or select start points. It also does not +//! quantize, serialize, or search the graph. +mod conversion; +mod finalization; +mod leaf_build; mod leaf_kernel; mod leaf_metric; mod partition_kernel; mod partition_metric; +mod partitioning; mod simd; mod topk; -use crate::{ANNError, ANNResult}; -use diskann_utils::views::MutMatrixView; +use std::num::NonZeroUsize; + +use crate::{ + ANNError, ANNResult, + graph::{AdjacencyList, Config}, + utils::VectorRepr, +}; +use diskann_utils::views::{MatrixView, MutMatrixView}; +use diskann_vector::distance::Metric; +use diskann_wide::arch::{self, Target2}; +use rayon::ThreadPool; + +use self::{leaf_metric::LeafMetric, partition_metric::PartitionMetric, simd::Simd}; /// Squared Euclidean distance. pub(super) struct L2; @@ -164,11 +179,248 @@ mod cosine_distance_tests { } } +/// Error for invalid PiPNN parameters. +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum PiPNNConfigError { + #[error("p_samp ({0}) must be in (0, 1]")] + SamplingFraction(f64), + #[error("fanout must not be empty")] + EmptyFanout, + #[error("max_leaders ({0}) must be at least 2")] + MaxLeaders(usize), +} + +crate::convert_error!(PiPNNConfigError); + +/// PiPNN partition and leaf parameters. +/// +/// The DiskANN graph [`Config`] supplies the degree, alpha, and prune kind. +/// +/// The default configuration is `c_max = 512`, `p_samp = 0.005`, +/// `max_leaders = 1000`, `fanout = [10, 3]`, `leaf_k = 2` and `replicas = 1`. +/// Each point then joins about 30 leaves, the product of the fanout values. +/// Larger `c_max`, fanout values, `leaf_k` and `replicas` all give each point +/// more or nearer candidates, which raises recall and makes the build slower. +#[derive(Clone, Debug, PartialEq)] +pub struct PiPNNConfig { + /// Maximum number of points in a leaf. A larger leaf holds more of the true + /// neighbors of its points, so the selected neighbors are nearer. The leaf + /// distance work grows with the square of the leaf size. + pub c_max: NonZeroUsize, + /// Fraction of the points of a cluster that a split samples as leaders, in + /// `(0, 1]`. More leaders give smaller children and fewer split levels, but + /// each split costs `points x leaders x dimensions` multiply-adds. A split + /// samples at least 2 and at most `max_leaders` leaders. + pub p_samp: f64, + /// Maximum number of leaders in one split, at least 2. A split costs + /// `points x leaders x dimensions` multiply-adds, so the cap bounds the + /// first splits of a large dataset, where `p_samp` alone samples many + /// leaders: 10M points at `p_samp = 0.005` would give 50,000. The PiPNN + /// paper uses about 1,000. + pub max_leaders: NonZeroUsize, + /// Number of nearest leaders that each point joins, for each split level. + /// Levels after this schedule use one leader. The schedule must not be empty. + /// Larger values make leaves overlap more, so each point gets candidates + /// from more leaves. The number of leaves and the leaf build time grow with + /// the product of the values. + pub fanout: Vec, + /// Number of nearest neighbors that each point selects inside a leaf. Each + /// selected pair adds a candidate in both directions. A larger value gives + /// more candidates and makes final pruning slower. + pub leaf_k: NonZeroUsize, + /// Number of independent partitionings of the dataset. Each replica repeats + /// the partition and leaf work with new leaders, which adds candidates and + /// build time. + pub replicas: NonZeroUsize, + /// Seed of leader sampling. The same seed gives the same leaves. `None` + /// draws a seed from the operating system, so each build samples new leaves. + pub seed: Option, +} + +impl PiPNNConfig { + /// Check the parameters that the field types do not constrain. + pub fn validate(&self) -> Result<(), PiPNNConfigError> { + if !(0.0 < self.p_samp && self.p_samp <= 1.0) { + return Err(PiPNNConfigError::SamplingFraction(self.p_samp)); + } + if self.fanout.is_empty() { + return Err(PiPNNConfigError::EmptyFanout); + } + if self.max_leaders.get() < 2 { + return Err(PiPNNConfigError::MaxLeaders(self.max_leaders.get())); + } + Ok(()) + } +} + +/// PiPNN parameters, graph policy, and the Rayon pool for one graph build. +#[derive(Debug)] +pub struct PiPNNBuildContext<'a> { + config: PiPNNConfig, + graph: &'a Config, + metric: Metric, + pool: &'a ThreadPool, +} + +impl<'a> PiPNNBuildContext<'a> { + /// Check `config` and combine it with the graph policy. + /// + /// Final pruning uses the prune kind of `graph`, as Vamana does. Build + /// `graph` with the prune kind of `metric` (`metric.into()`). + pub fn new( + config: PiPNNConfig, + graph: &'a Config, + metric: Metric, + pool: &'a ThreadPool, + ) -> Result { + config.validate()?; + Ok(Self { + config, + graph, + metric, + pool, + }) + } +} + +/// Build one adjacency list for each point of `data`. +/// +/// Each step converts the rows it reads to `f32` with +/// [`VectorRepr::as_f32_into`], so any element type that converts to `f32` +/// works, including compressed ones. +/// +/// The graph links dataset points only; the caller selects start points and +/// serializes the index. +pub fn build_graph( + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, +) -> ANNResult>> +where + T: VectorRepr, +{ + 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, +{ + 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(ANNError::message(format!( + "PiPNN dataset point count ({}) exceeds the u32 graph ID limit", + data.nrows() + ))); + } + arch::dispatch2_no_features(BuildGraph, data, context) +} + +/// Runs the build for the architecture that `diskann-wide` selects at run time. +struct BuildGraph; + +impl Target2>>, MatrixView<'_, T>, &PiPNNBuildContext<'_>> + for BuildGraph +where + A: Simd, + T: VectorRepr, +{ + fn run( + self, + arch: A, + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, + ) -> ANNResult>> { + let metric = context.metric; + match metric { + Metric::L2 => build_graph_for::(arch, data, context, metric), + Metric::Cosine => build_graph_for::(arch, data, context, metric), + Metric::CosineNormalized => { + build_graph_for::(arch, data, context, metric) + } + Metric::InnerProduct => { + build_graph_for::(arch, data, context, metric) + } + } + } +} + +/// Run the three build steps for architecture `A` and metric `M`. +fn build_graph_for( + arch: A, + data: MatrixView<'_, T>, + context: &PiPNNBuildContext<'_>, + metric: Metric, +) -> ANNResult>> +where + A: Simd, + M: LeafMetric + PartitionMetric, + T: VectorRepr, +{ + let leaves = tracing::info_span!("pipnn.partition") + .in_scope(|| partitioning::partition::(arch, data, &context.config))?; + // The leaf step takes the leaves by value, so their memory is free before + // final pruning starts. + let leaf_k = context.config.leaf_k.get(); + let candidates = tracing::info_span!("pipnn.leaf_build") + .in_scope(|| leaf_build::build_leaf_candidates::(arch, data, leaves, leaf_k))?; + // Final pruning cuts each candidate list in place. + Ok(tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric))) +} + #[cfg(test)] mod test_support { use super::simd::Simd; use diskann_vector::distance::Metric; + // Sort the members of each row but keep the row order. A row position is a + // leader column or a graph source, so sorting the rows would hide a wrong + // assignment. Repeated members stay visible. + pub(super) fn sorted_members_per_row(rows: &[Vec]) -> Vec> { + rows.iter() + .map(|row| { + let mut members = row.clone(); + members.sort_unstable(); + members + }) + .collect() + } + + pub(super) fn nz(value: usize) -> std::num::NonZeroUsize { + std::num::NonZeroUsize::new(value).unwrap() + } + + /// Build a Rayon pool with `threads` workers. + pub(super) fn thread_pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + } + + #[test] + fn membership_comparison_preserves_group_positions_and_duplicate_counts() { + let rows = [vec![9, 3, 9], vec![], vec![2, 1], vec![2, 1]]; + + assert_eq!( + sorted_members_per_row(&rows), + [vec![3, 9, 9], vec![], vec![1, 2], vec![1, 2]] + ); + } + /// A test body that runs once for each architecture. pub(super) trait ArchCheck { fn check(&self, arch: A); @@ -320,3 +572,345 @@ mod test_support { ); } } + +#[cfg(test)] +mod construction_tests { + use super::*; + use crate::graph::config; + use half::f16; + use test_support::{nz, sorted_members_per_row, thread_pool}; + + fn partition_policy() -> PiPNNConfig { + PiPNNConfig { + c_max: nz(4), + p_samp: 1.0, + max_leaders: nz(1000), + fanout: vec![nz(2)], + leaf_k: nz(1), + replicas: nz(1), + seed: Some(1290), + } + } + + fn graph_policy(degree: usize, metric: Metric) -> Result { + config::Builder::new_with(degree, config::MaxDegree::same(), 16, metric.into(), |b| { + b.alpha(1.0); + }) + .build() + } + + /// Build the graph of `values` with `threads` workers. + fn build( + values: &[T], + dimensions: usize, + config: PiPNNConfig, + degree: usize, + metric: Metric, + threads: usize, + ) -> Vec> { + let data = MatrixView::try_from(values, values.len() / dimensions, dimensions).unwrap(); + let graph = graph_policy(degree, metric).unwrap(); + let pool = thread_pool(threads); + let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); + let graph = build_graph(data, &context).unwrap(); + graph.into_iter().map(Vec::from).collect() + } + + #[test] + fn invalid_sampling_fractions_and_an_empty_fanout_are_rejected() { + for p_samp in [ + 0.0, + -0.5, + 1.0 + f64::EPSILON, + f64::NAN, + f64::INFINITY, + f64::NEG_INFINITY, + ] { + let config = PiPNNConfig { + p_samp, + ..partition_policy() + }; + + let error = config.validate().unwrap_err(); + + assert!( + matches!(error, PiPNNConfigError::SamplingFraction(value) if value.to_bits() == p_samp.to_bits()), + "p_samp {p_samp}: {error}" + ); + } + let config = PiPNNConfig { + fanout: vec![], + ..partition_policy() + }; + assert_eq!(config.validate(), Err(PiPNNConfigError::EmptyFanout)); + } + + #[test] + fn sampling_fraction_endpoints_are_accepted() { + for p_samp in [f64::from_bits(1), 1.0] { + let config = PiPNNConfig { + p_samp, + ..partition_policy() + }; + + assert_eq!(config.validate(), Ok(()), "p_samp {p_samp}"); + } + } + + #[test] + fn a_leader_cap_below_two_is_rejected() { + let config = PiPNNConfig { + max_leaders: nz(1), + ..partition_policy() + }; + + assert_eq!(config.validate(), Err(PiPNNConfigError::MaxLeaders(1))); + } + + #[test] + fn a_build_context_checks_the_configuration() { + let config = PiPNNConfig { + fanout: vec![], + ..partition_policy() + }; + let graph = graph_policy(2, Metric::L2).unwrap(); + let pool = thread_pool(1); + + let error = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap_err(); + + assert_eq!(error, PiPNNConfigError::EmptyFanout); + } + + #[test] + fn an_empty_dataset_axis_is_rejected_before_building() { + let graph = graph_policy(2, Metric::L2).unwrap(); + let pool = thread_pool(1); + let context = + PiPNNBuildContext::new(partition_policy(), &graph, Metric::L2, &pool).unwrap(); + + for (rows, columns, expected_message) in [ + (0, 2, "at least one data point"), + (2, 0, "at least one data dimension"), + ] { + let data = MatrixView::try_from(&[] as &[f32], rows, columns).unwrap(); + + let error = build_graph(data, &context).unwrap_err(); + + assert!(error.to_string().contains(expected_message), "{error}"); + } + } + + #[test] + fn native_vector_types_build_the_expected_neighbor_graph() { + // Nearest choices 0 -> 1, 1 -> 0 and 2 -> 1 become a symmetric chain. + fn check(values: [T; 3]) { + let actual = build(&values, 1, partition_policy(), 2, Metric::L2, 2); + + assert_eq!( + sorted_members_per_row(&actual), + [vec![1], vec![0, 2], vec![1]], + "{}", + std::any::type_name::() + ); + } + + check([0.0_f32, 1.0, 4.0]); + check([0.0_f32, 1.0, 4.0].map(f16::from_f32)); + check([0_i8, 1, 4]); + check([0_u8, 1, 4]); + } + + #[test] + fn the_requested_metric_determines_leaf_neighbors() { + // L2's nearest choices are [2, 2, 0, 2]; cosine pairs similar directions + // 0 <-> 1 and 2 <-> 3; dot products give [1, 3, 3, 2]. + // The second coordinate sits in the last dimension of an embedding. + let dimensions = 1537; + for (metric, expected) in [ + (Metric::L2, vec![vec![2], vec![2], vec![0, 1, 3], vec![2]]), + (Metric::Cosine, vec![vec![1], vec![0], vec![3], vec![2]]), + ( + Metric::CosineNormalized, + vec![vec![1], vec![0], vec![3], vec![2]], + ), + ( + Metric::InnerProduct, + vec![vec![1], vec![0, 3], vec![3], vec![1, 2]], + ), + ] { + let values = test_support::packed_points( + &[[1.0, 0.0], [5.0, 2.0], [1.0, 3.0], [0.0, 9.0]], + dimensions, + metric == Metric::CosineNormalized, + ); + + let actual = build(&values, dimensions, partition_policy(), 3, metric, 2); + + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "{metric:?}" + ); + } + } + + #[test] + fn splitting_uses_the_requested_metric_to_group_points() { + // Four points force splitting at c_max=2. Every point is sampled, so + // leader order cannot change memberships. With L2, leader 0 gets {0,2} + // and leader 2 gets all four points, which then split into singletons. + // Cosine pairs directions {0,1} and {2,3}. Inner product picks leaders + // {1,2} for points 0/1 and {2,3} for points 2/3; the oversized leader-2 + // cluster then splits into those same pairs using fanout one. + // Every two-point leaf contributes its only pair; degree three retains + // every edge. The graph therefore exposes partition membership directly. + let config = PiPNNConfig { + c_max: nz(2), + ..partition_policy() + }; + for (metric, expected) in [ + (Metric::L2, vec![vec![2], vec![], vec![0], vec![]]), + (Metric::Cosine, vec![vec![1], vec![0], vec![3], vec![2]]), + ( + Metric::CosineNormalized, + vec![vec![1], vec![0], vec![3], vec![2]], + ), + ( + Metric::InnerProduct, + vec![vec![1], vec![0], vec![3], vec![2]], + ), + ] { + let values = test_support::packed_points( + &[[1.0, 0.0], [6.0, 1.0], [2.0, 4.0], [0.0, 9.0]], + 2, + metric == Metric::CosineNormalized, + ); + + let actual = build(&values, 2, config.clone(), 3, metric, 2); + + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "{metric:?}" + ); + } + } + + #[test] + fn final_pruning_cuts_merged_candidates_to_the_degree() { + let config = PiPNNConfig { + leaf_k: nz(2), + ..partition_policy() + }; + + let actual = build(&[0.0_f32, 1.0, 4.0], 1, config, 1, Metric::L2, 2); + + assert_eq!(actual, [vec![1], vec![0], vec![1]]); + } + + #[test] + fn overlapping_partitions_build_the_same_deduplicated_graph() { + // Sampling every point with fanout 2 makes two overlapping copies + // of each close pair. Neither replicas nor worker count may duplicate edges. + for (workers, replicas) in [(1, 1), (3, 2)] { + let config = PiPNNConfig { + c_max: nz(2), + replicas: nz(replicas), + ..partition_policy() + }; + + let actual = build( + &[0.0_f32, 1.0, 10.0, 11.0], + 1, + config, + 2, + Metric::L2, + workers, + ); + + assert_eq!( + actual, + [vec![1], vec![0], vec![3], vec![2]], + "{workers} workers, {replicas} replicas" + ); + } + } + + #[test] + fn a_single_point_has_one_empty_adjacency_list() { + let actual = build(&[3.0_f32], 1, partition_policy(), 2, Metric::L2, 1); + + assert_eq!(actual, [Vec::::new()]); + } + + #[test] + fn multi_level_builds_give_valid_rows_for_every_metric() { + // With 1,000 points and c_max 32, the build splits over several levels, + // makes overlapping leaves, merges their candidates on four workers and + // prunes lists above the degree. Each row must be non-empty, within the + // degree, and free of duplicates and self edges. The f16 build checks + // that every step reads f16 rows, which the smaller builds never split. + fn check(values: &[T], metric: Metric) { + let (dimensions, degree) = (16, 8); + let config = PiPNNConfig { + c_max: nz(32), + p_samp: 0.05, + max_leaders: nz(1000), + fanout: vec![nz(4), nz(2)], + leaf_k: nz(2), + replicas: nz(1), + seed: Some(1290), + }; + + let actual = build(values, dimensions, config, degree, metric, 4); + + for (point, neighbors) in actual.iter().enumerate() { + let mut distinct = neighbors.clone(); + distinct.sort_unstable(); + distinct.dedup(); + assert!( + !neighbors.is_empty() + && neighbors.len() <= degree + && distinct.len() == neighbors.len() + && !neighbors.contains(&(point as u32)), + "{metric:?}, {}, point {point}: {neighbors:?}", + std::any::type_name::() + ); + } + } + + for metric in [ + Metric::L2, + Metric::Cosine, + Metric::CosineNormalized, + Metric::InnerProduct, + ] { + let mut values = test_support::dense_points(1000, 16, 1290); + if metric == Metric::CosineNormalized { + test_support::normalize(&mut values, 16); + } + check(&values, metric); + check( + &values.iter().map(|&x| f16::from_f32(x)).collect::>(), + metric, + ); + } + } + + #[test] + fn copies_of_one_vector_build_leaf_local_edges() { + // Every leader is a copy, so each point ties at every split. The build + // splits the six copies into two leaves of three points. + let actual = build(&[1.0_f32; 6], 1, partition_policy(), 2, Metric::L2, 2); + + for (point, neighbors) in actual.iter().enumerate() { + assert!(!neighbors.is_empty(), "point {point}"); + assert!( + neighbors + .iter() + .all(|&neighbor| neighbor / 3 == point as u32 / 3), + "point {point}: {neighbors:?}" + ); + } + } +} diff --git a/diskann/src/graph/pipnn/partitioning.rs b/diskann/src/graph/pipnn/partitioning.rs new file mode 100644 index 0000000000..0950871920 --- /dev/null +++ b/diskann/src/graph/pipnn/partitioning.rs @@ -0,0 +1,970 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Randomized ball carving: recursive, overlapping partitioning of the dataset +//! into leaves of at most `c_max` points. +//! +//! One split samples leaders from a cluster and assigns each point to its +//! `fanout` nearest leaders. The points of one leader form a child cluster, so +//! children overlap when `fanout > 1`. A child with at most `c_max` points is a +//! leaf; a larger child is split again at the next level. `fanout[level]` gives +//! the fanout of each level, and levels after the schedule use fanout 1. Each +//! replica repeats the process with a different seed. +//! +//! A fanout-1 split in which one child gets every point makes no progress, for +//! example when all sampled leaders are copies of one vector. The split cuts +//! that child into equal leaves in point order, and the build logs a warning +//! with the number of points it cut. These leaves group points arbitrarily, +//! which can lower recall. +//! +//! The partition kernel computes point-to-leader distances and selects the +//! nearest leaders. This module samples leaders, gathers vectors, and groups the +//! assignments into children. +//! +//! A split assigns its points in parallel stripes. An assignment stripe is a +//! batch of up to `ASSIGNMENT_STRIPE_POINTS` points that one task gathers as +//! `f32` and ranks against all leaders of the split. +//! +//! Buffer sizes use plain multiplication. The IDs of one split are distinct +//! dataset rows, so a gathered `rows x dimensions` buffer is not larger than the +//! dataset. An assignment buffer holds `points x fanout` IDs. The fanout of a +//! split is below its leader count, which is at most its point count, so with at +//! most `u32::MAX` points this count fits in a 64-bit `usize`. + +use crate::{ANNResult, utils::VectorRepr}; +use diskann_utils::{ + object_pool::{AsPooled, ObjectPool}, + views::{Matrix, MatrixView, MutMatrixView}, +}; +use rand::{Rng, SeedableRng, prelude::IndexedRandom, rngs::StdRng}; +use rayon::prelude::*; + +use super::{ + PiPNNConfig, + conversion::{f32_dimensions, gather_as_f32}, + partition_kernel::{PartitionKernelWorkspace, assign_leaders}, + partition_metric::PartitionMetric, + simd::Simd, + topk::UNASSIGNED, +}; + +/// Number of points in one assignment stripe. Short stripes repeat the fixed +/// cost of a kernel call, and tall stripes overflow the last-level cache once +/// every worker holds one. The height does not follow the detected cache size, +/// because a virtual machine reports a last-level cache that other tenants +/// share. +const ASSIGNMENT_STRIPE_POINTS: usize = 256; +/// Clusters with at least this many points group their points by leader in +/// parallel. Such clusters occur at the first levels, where few splits run at +/// the same time, so a serial pass would leave most workers idle. Smaller +/// clusters group serially, because the parallel pass costs more than it saves. +const PARALLEL_GROUPING_MIN_POINTS: usize = 100_000; + +/// A cluster with more than `c_max` points that waits for its split. +struct Cluster { + point_ids: Vec, + /// Position in the fanout schedule. + level: usize, + /// Seed of the random stream of this cluster's split. The parent split + /// draws it, so every split samples its leaders from its own stream. + seed: u64, +} + +/// The result of one split: leaves, and clusters for the next level. +struct ClusterSplit { + next_level: Vec, + leaves: Vec>, + /// Number of points that the split cut into leaves in point order. + cut_points: usize, +} + +/// Scratch for one assignment stripe: the gathered `f32` points and the kernel +/// workspace. +#[derive(Default)] +struct StripeBuffers { + point_values: Vec, + kernel_workspace: PartitionKernelWorkspace, +} + +impl AsPooled<()> for StripeBuffers { + fn create(_: ()) -> Self { + Self::default() + } + + fn modify(&mut self, _: ()) { + // Keep the largest allocation across leases. `assign_point_stripe` sets + // the active prefix before each read. + } +} + +/// Stripe scratch shared by all splits of a build. +/// +/// `ObjectPool` locks only to give or take back a lease. Numerical work holds the +/// lease, not the pool lock. +type StripeBufferPool = ObjectPool; + +/// Build the leaves of all replicas. +/// +/// A point whose distances to all leaders of a split are not rankable, for +/// example a NaN vector, joins no child of that split. +pub(super) fn partition( + arch: A, + data: MatrixView<'_, T>, + config: &PiPNNConfig, +) -> ANNResult>> +where + A: Simd, + M: PartitionMetric, + T: VectorRepr + Send + Sync, +{ + let mut leaves = Vec::new(); + let mut cut_points = 0; + let stripe_buffers = StripeBufferPool::new((), 0, None); + // Replica `r` uses seed `base + r`. `StdRng::seed_from_u64` gives + // independent streams for adjacent seeds. + let base_seed = config.seed.unwrap_or_else(rand::random); + for replica in 0..config.replicas.get() { + let seed = base_seed.wrapping_add(replica as u64); + let (mut replica_leaves, replica_cut_points) = + partition_replica::(arch, data, config, seed, &stripe_buffers)?; + leaves.append(&mut replica_leaves); + cut_points += replica_cut_points; + } + if cut_points > 0 { + tracing::warn!( + cut_points, + "PiPNN cut {cut_points} points into leaves in point order because no split \ + could shrink their clusters; these leaves group points arbitrarily and can \ + lower recall" + ); + } + Ok(leaves) +} + +/// Split one replica, one level at a time, until every cluster is a leaf. +/// +/// The loop ends. The fanout schedule is finite, and each split sends its +/// clusters one level deeper. After the schedule, fanout is 1, so each point +/// joins one child. A child is then smaller than its parent, or it holds every +/// point and [`split_cluster`] cuts it into leaves. +fn partition_replica( + arch: A, + data: MatrixView<'_, T>, + config: &PiPNNConfig, + seed: u64, + stripe_buffers: &StripeBufferPool, +) -> ANNResult<(Vec>, usize)> +where + A: Simd, + M: PartitionMetric, + T: VectorRepr + Send + Sync, +{ + let initial_point_ids = (0..data.nrows() as u32).collect(); + let c_max = config.c_max.get(); + if data.nrows() <= c_max { + return Ok((vec![initial_point_ids], 0)); + } + + let mut leaves = Vec::new(); + let mut cut_points = 0; + let mut pending = vec![Cluster { + point_ids: initial_point_ids, + level: 0, + seed, + }]; + while !pending.is_empty() { + // Indexed parallel collection keeps the parent order. + let splits: ANNResult> = pending + .into_par_iter() + .map(|cluster| split_cluster::(arch, data, config, cluster, stripe_buffers)) + .collect(); + + let mut next_level = Vec::new(); + for mut split in splits? { + next_level.append(&mut split.next_level); + leaves.append(&mut split.leaves); + cut_points += split.cut_points; + } + pending = next_level; + } + Ok((leaves, cut_points)) +} + +/// Split one cluster into leaves and clusters for the next level. +fn split_cluster( + arch: A, + data: MatrixView<'_, T>, + config: &PiPNNConfig, + cluster: Cluster, + stripe_buffers: &StripeBufferPool, +) -> ANNResult +where + A: Simd, + M: PartitionMetric, + T: VectorRepr + Send + Sync, +{ + let mut rng = StdRng::seed_from_u64(cluster.seed); + let fanout = config + .fanout + .get(cluster.level) + .map_or(1, |fanout| fanout.get()); + let c_max = config.c_max.get(); + let leader_count = sampled_leader_count( + cluster.point_ids.len(), + config.p_samp, + config.max_leaders.get(), + ); + if fanout > 1 && fanout >= leader_count { + // Every point would join every leader, and each child would repeat this + // cluster. Split the cluster at the next level instead. + return Ok(ClusterSplit { + next_level: vec![Cluster { + point_ids: cluster.point_ids, + level: cluster.level + 1, + seed: rng.random(), + }], + leaves: Vec::new(), + cut_points: 0, + }); + } + + let leaders = sample_leaders(&cluster.point_ids, leader_count, &mut rng); + let children = assign_to_leaders::( + arch, + data, + &cluster.point_ids, + &leaders, + fanout, + stripe_buffers, + )?; + + let mut next_level = Vec::new(); + let mut leaves = Vec::new(); + let mut cut_points = 0; + for child in children { + if child.is_empty() { + continue; + } + if child.len() <= c_max { + leaves.push(child); + } else if fanout == 1 && child.len() == cluster.point_ids.len() { + // Every point chose this leader, so the split made no progress. With + // L2 or cosine, each leader is nearest to itself, so this needs all + // sampled leaders to be copies of one vector. With inner product, a + // leader with a large norm can also win every point. Another split can + // stall the same way, so cut the child into equal leaves in point + // order. The cut groups points arbitrarily, which can lower recall but + // keeps the graph valid. + let leaf_len = child.len().div_ceil(child.len().div_ceil(c_max)); + leaves.extend(child.chunks(leaf_len).map(<[u32]>::to_vec)); + cut_points += child.len(); + } else { + next_level.push(Cluster { + point_ids: child, + level: cluster.level + 1, + seed: rng.random(), + }); + } + } + Ok(ClusterSplit { + next_level, + leaves, + cut_points, + }) +} + +/// Sample `count` distinct leaders from the points of a cluster. +fn sample_leaders(points: &[u32], count: usize, rng: &mut StdRng) -> Vec { + points.choose_multiple(rng, count).copied().collect() +} + +/// Return the number of leaders for a cluster of `points` points. +/// +/// A split needs two leaders. `max_leaders` bounds the cost of one split. +fn sampled_leader_count(points: usize, sampling_fraction: f64, max_leaders: usize) -> usize { + ((points as f64 * sampling_fraction).ceil() as usize).clamp(2, max_leaders) +} + +/// Assign each point to its `fanout` nearest leaders and group the points by +/// leader. +/// +/// Stripes of points run in parallel. Each child keeps the input order of its +/// points, so the next split samples the same leaders for any thread count. +fn assign_to_leaders( + arch: A, + data: MatrixView<'_, T>, + point_ids: &[u32], + leader_ids: &[u32], + fanout: usize, + stripe_buffers: &StripeBufferPool, +) -> ANNResult>> +where + A: Simd, + M: PartitionMetric, + T: VectorRepr + Send + Sync, +{ + let dimensions = f32_dimensions(data)?; + let mut leader_values = Matrix::new(0.0f32, leader_ids.len(), dimensions); + gather_as_f32(data, leader_ids, leader_values.as_mut_view())?; + + let leaders = M::create_leaders(leader_values.as_view()); + let leader_count = M::leader_count(&leaders); + + let mut assignments = Matrix::new(0u32, point_ids.len(), fanout); + // This runs in the pool of the build context (see `build_graph`). A lease + // for each stripe takes one lock, which costs far less than the GEMM and + // ranking of the stripe. + assignments + .par_window_iter_mut(ASSIGNMENT_STRIPE_POINTS) + .zip(point_ids.par_chunks(ASSIGNMENT_STRIPE_POINTS)) + .try_for_each(|(stripe_assignments, stripe_ids)| { + assign_point_stripe::( + arch, + data, + stripe_ids, + dimensions, + &leaders, + &mut stripe_buffers.get_ref(()), + stripe_assignments, + ) + })?; + + Ok(group_by_leader( + point_ids, + assignments.as_view(), + leader_count, + )) +} + +/// Gather one stripe of points as `f32` and write the leader columns of the +/// nearest leaders of each point to its row of `assignments`, one per column. +#[inline] +fn assign_point_stripe( + arch: A, + data: MatrixView<'_, T>, + point_ids: &[u32], + dimensions: usize, + leaders: &M::Leaders<'_>, + buffers: &mut StripeBuffers, + assignments: MutMatrixView<'_, u32>, +) -> ANNResult<()> +where + A: Simd, + M: PartitionMetric, + T: VectorRepr, +{ + let point_count = point_ids.len(); + let point_values_len = point_count * dimensions; + // The buffer keeps its largest length. This stripe uses a prefix of it. + if buffers.point_values.len() < point_values_len { + buffers.point_values.resize(point_values_len, 0.0); + } + let StripeBuffers { + point_values, + kernel_workspace, + } = buffers; + let mut points = MutMatrixView::try_from( + &mut point_values[..point_values_len], + point_count, + dimensions, + )?; + gather_as_f32(data, point_ids, points.as_mut_view())?; + assign_leaders::( + arch, + points.as_view(), + leaders, + assignments, + kernel_workspace, + ) +} + +/// Group the points by assigned leader. Unassigned slots join no child. +/// +/// Both paths keep the input order of the points inside each child, because the +/// next split samples its leaders by position. +fn group_by_leader( + points: &[u32], + assignments: MatrixView<'_, u32>, + leaders: usize, +) -> Vec> { + if points.len() < PARALLEL_GROUPING_MIN_POINTS { + return group_by_leader_serial(points, assignments, leaders); + } + + let chunk_points = points.len().div_ceil(rayon::current_num_threads()); + // Group each chunk of points on its own. Indexed collection keeps the chunk + // order. + let locals: Vec<_> = points + .par_chunks(chunk_points) + .zip(assignments.par_window_iter(chunk_points)) + .map(|(points, assignments)| group_by_leader_serial(points, assignments, leaders)) + .collect(); + + let mut sizes = vec![0usize; leaders]; + for local in &locals { + for (size, cluster) in sizes.iter_mut().zip(local) { + *size += cluster.len(); + } + } + + // Join the chunk parts of each child in chunk order. + sizes + .into_par_iter() + .enumerate() + .map(|(leader, size)| { + let mut cluster = Vec::with_capacity(size); + for local in &locals { + cluster.extend_from_slice(&local[leader]); + } + cluster + }) + .collect() +} + +fn group_by_leader_serial( + points: &[u32], + assignments: MatrixView<'_, u32>, + leaders: usize, +) -> Vec> { + let mut sizes = vec![0usize; leaders]; + for &leader in assignments.as_slice() { + if leader != UNASSIGNED { + sizes[leader as usize] += 1; + } + } + let mut clusters: Vec> = sizes.into_iter().map(Vec::with_capacity).collect(); + for (&point, point_assignments) in points.iter().zip(assignments.row_iter()) { + for &leader in point_assignments { + if leader != UNASSIGNED { + clusters[leader as usize].push(point); + } + } + } + clusters +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::graph::pipnn::{ + L2, + test_support::{self, dense_points, nz, sorted_members_per_row, thread_pool}, + }; + use diskann_vector::distance::Metric; + use diskann_wide::ARCH; + use std::collections::HashSet; + + // Neither leaf order nor point order within a leaf is part of the result. + // Preserve repeated leaves and IDs so comparison still detects duplicates. + fn sorted_leaf_memberships(leaves: &[Vec]) -> Vec> { + let mut leaves = sorted_members_per_row(leaves); + leaves.sort_unstable(); + leaves + } + + fn splitting_config() -> PiPNNConfig { + PiPNNConfig { + c_max: nz(2), + p_samp: 1.0, + max_leaders: nz(1000), + fanout: vec![nz(2)], + leaf_k: nz(1), + replicas: nz(1), + seed: Some(1290), + } + } + + fn split( + data: MatrixView<'_, f32>, + config: &PiPNNConfig, + parent: Cluster, + ) -> ClusterSplit { + let buffers = StripeBufferPool::new((), 0, None); + thread_pool(2) + .install(|| split_cluster::<_, M, _>(ARCH, data, config, parent, &buffers)) + .unwrap() + } + + #[test] + fn sampled_leaders_are_a_repeatable_subset_without_replacement() { + // (points, sampling fraction, expected leaders) + for (point_count, fraction, expected_count) in [ + (7, 0.01, 2), // A split needs two leaders. + (7, 0.31, 3), // 2.17 rounds up. + (7, 1.0, 7), + (999, 1.0, 999), + (1000, 1.0, 1000), + (1001, 1.0, 1000), + ] { + let case = format!("{point_count} points, fraction {fraction}"); + let ids: Vec<_> = (0..point_count).map(|id| 10 + 3 * id).collect(); + + let count = sampled_leader_count(ids.len(), fraction, 1000); + let mut leaders = sample_leaders(&ids, count, &mut StdRng::seed_from_u64(1290)); + + assert_eq!(leaders.len(), expected_count, "{case}"); + let distinct: HashSet<_> = leaders.iter().collect(); + assert_eq!(distinct.len(), expected_count, "{case}"); + assert!(leaders.iter().all(|id| ids.contains(id)), "{case}"); + let mut repeated = sample_leaders(&ids, count, &mut StdRng::seed_from_u64(1290)); + leaders.sort_unstable(); + repeated.sort_unstable(); + assert_eq!(leaders, repeated, "{case}"); + } + } + + #[test] + fn points_join_their_nearest_leaders() { + // Leader columns are global IDs [2, 0], at x=10 and x=0. + // The requested points [5, 4, 1] are at x=8, x=2 and x=100. + let values = [0.0_f32, 100.0, 10.0, -100.0, 2.0, 8.0]; + let data = MatrixView::try_from(&values[..], 6, 1).unwrap(); + let buffers = StripeBufferPool::new((), 0, None); + let pool = thread_pool(2); + + for (fanout, expected) in [ + (1, [vec![5, 1], vec![4]]), + (2, [vec![5, 4, 1], vec![5, 4, 1]]), + ] { + let actual = pool + .install(|| { + assign_to_leaders::<_, L2, _>(ARCH, data, &[5, 4, 1], &[2, 0], fanout, &buffers) + }) + .unwrap(); + + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "fanout {fanout}" + ); + } + } + + #[test] + fn multiple_stripes_and_their_partial_tail_preserve_all_assignments() { + // Two full stripes, with three points left over. + // Each dense row selects two of the three centers at 10, 0 and 4. + // Every row selects 4; its other center is 10 for x=10 and 0 otherwise. + let point_count = 2 * ASSIGNMENT_STRIPE_POINTS + 3; + let dimensions = 384; + let values: Vec<_> = (0..point_count) + .flat_map(|point| std::iter::repeat_n([0.0_f32, 4.0, 10.0][point % 3], dimensions)) + .collect(); + let data = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let ids: Vec<_> = (0..point_count as u32).rev().collect(); + let expected = [ + ids.iter().copied().filter(|id| id % 3 == 2).collect(), + ids.iter().copied().filter(|id| id % 3 != 2).collect(), + ids.clone(), + ]; + let buffers = StripeBufferPool::new((), 0, None); + + for workers in [1, 3] { + let actual = thread_pool(workers) + .install(|| { + assign_to_leaders::<_, L2, _>(ARCH, data, &ids, &[2, 0, 1], 2, &buffers) + }) + .unwrap(); + + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "{workers} workers" + ); + } + } + + #[test] + fn stripes_assign_each_point_to_its_nearest_leaders() { + // 3,000 points fill 11 stripes and a partial tail on several workers. + // 300 leaders fill whole SIMD leader groups. + // Dense fixtures keep the f32 scores exact, so the check needs no + // tolerance. It accepts any order of tied leaders. + let (point_count, dimensions, fanout) = (3000, 32, 3); + let values = dense_points(point_count, dimensions, 1290); + let data = MatrixView::try_from(values.as_slice(), point_count, dimensions).unwrap(); + let row = |id: u32| &values[id as usize * dimensions..][..dimensions]; + let ids: Vec = (0..point_count as u32).rev().collect(); + let leaders: Vec = (0..point_count as u32).step_by(10).collect(); + let buffers = StripeBufferPool::new((), 0, None); + + let clusters = thread_pool(4) + .install(|| assign_to_leaders::<_, L2, _>(ARCH, data, &ids, &leaders, fanout, &buffers)) + .unwrap(); + + let mut chosen = vec![Vec::new(); point_count]; + for (column, cluster) in clusters.iter().enumerate() { + // `ids` is descending, and grouping keeps the input order. + assert!(cluster.is_sorted_by(|a, b| a > b), "leader {column}"); + for &id in cluster { + chosen[id as usize].push(column); + } + } + for (id, columns) in chosen.iter().enumerate() { + let distance = |column: usize| { + test_support::distance(Metric::L2, row(id as u32), row(leaders[column])) + }; + let farthest_chosen = columns + .iter() + .map(|&column| distance(column)) + .fold(f64::MIN, f64::max); + let nearest_other = (0..leaders.len()) + .filter(|column| !columns.contains(column)) + .map(distance) + .fold(f64::INFINITY, f64::min); + assert_eq!(columns.len(), fanout, "point {id}"); + assert!( + farthest_chosen <= nearest_other, + "point {id}: chosen {columns:?} at {farthest_chosen}, other at {nearest_other}" + ); + } + } + + #[test] + fn reused_stripe_buffers_do_not_carry_points() { + // Points and leaders lie on a line, so the nearest leaders follow from + // the coordinates. The first stripe grows the buffer. The second stripe is + // smaller and reads only a prefix of it. + let values = [100.0_f32, 0.0, 8.0, 10.0, 2.0]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + let leader_values = [0.0_f32, 10.0, 100.0]; + let leaders = L2::create_leaders(MatrixView::try_from(&leader_values[..], 3, 1).unwrap()); + let fanout = 2; + let mut buffers = StripeBuffers::default(); + + for ids in [&[4, 2, 0, 3, 1][..], &[2, 1][..]] { + let mut assignments = Matrix::new(0, ids.len(), fanout); + + assign_point_stripe::<_, L2, _>( + ARCH, + data, + ids, + 1, + &leaders, + &mut buffers, + assignments.as_mut_view(), + ) + .unwrap(); + + // The `fanout` leader columns nearest to each point. In one dimension, + // |x - leader| orders the leaders the same way as squared L2. No two + // leaders are at the same distance from a point. + let distance = + |id: u32, column: u32| (values[id as usize] - leader_values[column as usize]).abs(); + let expected: Vec> = ids + .iter() + .map(|&id| { + let mut columns: Vec = (0..leader_values.len() as u32).collect(); + columns.sort_by(|&a, &b| distance(id, a).total_cmp(&distance(id, b))); + columns.truncate(fanout); + columns + }) + .collect(); + let actual: Vec<_> = assignments.row_iter().map(<[u32]>::to_vec).collect(); + assert_eq!(actual, expected, "points {ids:?}"); + } + } + + #[test] + fn grouping_by_leader_drops_unassigned_slots() { + let pattern = [ + [2, 0], + [UNASSIGNED, 1], + [0, UNASSIGNED], + [UNASSIGNED, UNASSIGNED], + ]; + let pool = thread_pool(3); + + // Counts around the parallel threshold cover both grouping paths and + // parallel chunks of unequal length. + for point_count in [ + 7, + PARALLEL_GROUPING_MIN_POINTS - 1, + PARALLEL_GROUPING_MIN_POINTS, + PARALLEL_GROUPING_MIN_POINTS + 1, + ] { + let ids: Vec<_> = (0..point_count).map(|i| 1_000_000 - i as u32).collect(); + let assignments: Vec<_> = (0..point_count).flat_map(|i| pattern[i % 4]).collect(); + let assignments = MatrixView::try_from(assignments.as_slice(), point_count, 2).unwrap(); + let expected = [ + ids.iter().copied().step_by(2).collect(), + ids.iter().copied().skip(1).step_by(4).collect(), + ids.iter().copied().step_by(4).collect(), + vec![], + ]; + + let actual = pool.install(|| group_by_leader(&ids, assignments, 4)); + + assert_eq!( + sorted_members_per_row(&actual), + sorted_members_per_row(&expected), + "{point_count} points" + ); + } + } + + #[test] + fn a_split_uses_the_configured_fanout_then_falls_back_to_one_leader() { + // Every point is a leader, so membership is fixed regardless of output order. + let values = [0.0_f32, 1.0, 10.0, 11.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + + for (level, expected) in [ + (0, vec![vec![0, 1], vec![0, 1], vec![2, 3], vec![2, 3]]), + (1, vec![vec![0], vec![1], vec![2], vec![3]]), + ] { + let parent = Cluster { + point_ids: vec![0, 1, 2, 3], + level, + seed: 1290, + }; + + let actual = split::(data, &splitting_config(), parent); + + assert_eq!( + sorted_leaf_memberships(&actual.leaves), + sorted_leaf_memberships(&expected), + "level {level}" + ); + assert!(actual.next_level.is_empty(), "level {level}"); + assert_eq!(actual.cut_points, 0, "level {level}"); + } + } + + #[test] + fn children_above_cmax_are_queued_at_the_next_level() { + let values = [0.0_f32, 1.0, 10.0, 11.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let config = PiPNNConfig { + c_max: nz(1), + ..splitting_config() + }; + let parent = Cluster { + point_ids: vec![0, 1, 2, 3], + level: 0, + seed: 1290, + }; + + let actual = split::(data, &config, parent); + + assert!(actual.leaves.is_empty()); + assert!(actual.next_level.iter().all(|child| child.level == 1)); + let ids: Vec<_> = actual + .next_level + .into_iter() + .map(|child| child.point_ids) + .collect(); + assert_eq!( + sorted_leaf_memberships(&ids), + [vec![0, 1], vec![0, 1], vec![2, 3], vec![2, 3]] + ); + } + + #[test] + fn a_fanout_one_split_without_progress_cuts_equal_leaves() { + // Copies of one vector tie for every leader, so all points choose the + // same leader and the child equals its parent. + let values = [1.0_f32; 5]; + let data = MatrixView::try_from(&values[..], 5, 1).unwrap(); + let config = PiPNNConfig { + fanout: vec![nz(1)], + ..splitting_config() + }; + let parent = Cluster { + point_ids: (0..5).collect(), + level: 0, + seed: 1290, + }; + + let actual = split::(data, &config, parent); + + assert_eq!( + sorted_leaf_memberships(&actual.leaves), + [vec![0, 1], vec![2, 3], vec![4]] + ); + assert!(actual.next_level.is_empty()); + assert_eq!(actual.cut_points, 5); + } + + #[test] + fn a_child_that_repeats_its_parent_at_fanout_two_is_split_again() { + // Points at x=0, x=10 and x=4 are all leaders. Each point joins itself and + // x=4, so the child of x=4 holds all points. With fanout above one this + // is overlap, not a stalled split. + let values = [0.0_f32, 10.0, 4.0]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let parent = Cluster { + point_ids: vec![0, 1, 2], + level: 0, + seed: 1290, + }; + + let actual = split::(data, &splitting_config(), parent); + + assert_eq!( + sorted_leaf_memberships(&actual.leaves), + [vec![0, 2], vec![1]] + ); + let [child] = &actual.next_level[..] else { + panic!("expected one pending cluster"); + }; + assert_eq!( + (child.point_ids.as_slice(), child.level), + (&[0, 1, 2][..], 1) + ); + } + + #[test] + fn a_level_without_more_leaders_than_its_fanout_is_skipped() { + // Four points at p_samp 0.01 give two leaders. Fanout 2 would put every + // point into both children, so the cluster moves to the next level. + let values = [0.0_f32, 1.0, 10.0, 11.0]; + let data = MatrixView::try_from(&values[..], 4, 1).unwrap(); + let config = PiPNNConfig { + p_samp: 0.01, + ..splitting_config() + }; + let parent = Cluster { + point_ids: vec![0, 1, 2, 3], + level: 0, + seed: 1290, + }; + + let actual = split::(data, &config, parent); + + assert!(actual.leaves.is_empty()); + let [child] = &actual.next_level[..] else { + panic!("expected one pending cluster"); + }; + assert_eq!( + (child.point_ids.as_slice(), child.level), + (&[0, 1, 2, 3][..], 1) + ); + } + + #[test] + fn data_within_cmax_forms_one_complete_leaf_per_replica() { + let values = [0.0_f32, 1.0, 10.0]; + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + let pool = thread_pool(1); + + for c_max in [3, 4] { + let config = PiPNNConfig { + c_max: nz(c_max), + replicas: nz(2), + ..splitting_config() + }; + + let actual = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + + assert_eq!( + sorted_leaf_memberships(&actual), + [vec![0, 1, 2], vec![0, 1, 2]], + "c_max {c_max}" + ); + } + } + + #[test] + fn recursive_partitioning_preserves_leaf_membership_across_worker_counts() { + let point_count = 129; + let values = dense_points(point_count, 17, 1290); + let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); + let config = PiPNNConfig { + c_max: nz(16), + p_samp: 0.2, + fanout: vec![nz(2), nz(2)], + replicas: nz(2), + ..splitting_config() + }; + + let first = thread_pool(1) + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + let second = thread_pool(3) + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + + assert_eq!( + sorted_leaf_memberships(&first), + sorted_leaf_memberships(&second) + ); + let mut ids: Vec<_> = first.iter().flatten().copied().collect(); + ids.sort_unstable(); + ids.dedup(); + assert_eq!(ids, (0..point_count as u32).collect::>()); + for leaf in first { + assert!( + !leaf.is_empty() && leaf.len() <= config.c_max.get(), + "leaf {leaf:?}" + ); + assert_eq!( + leaf.iter().copied().collect::>().len(), + leaf.len(), + "duplicate IDs in leaf {leaf:?}" + ); + } + } + + #[test] + fn replicas_use_different_seeds() { + // With one seed for every replica, the second replica would repeat the + // leaves of the first. + let point_count = 33; + let values = dense_points(point_count, 17, 1290); + let data = MatrixView::try_from(values.as_slice(), point_count, 17).unwrap(); + let config = PiPNNConfig { + c_max: nz(8), + p_samp: 0.25, + fanout: vec![nz(1)], + ..splitting_config() + }; + let pool = thread_pool(2); + let leaves = |replicas| { + let config = PiPNNConfig { + replicas: nz(replicas), + ..config.clone() + }; + let leaves = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + sorted_leaf_memberships(&leaves) + }; + + let one = leaves(1); + let two = leaves(2); + + assert!(two.iter().any(|leaf| !one.contains(leaf))); + } + + #[test] + fn unrankable_points_join_no_leaf() { + let config = PiPNNConfig { + c_max: nz(1), + fanout: vec![nz(1)], + ..splitting_config() + }; + let pool = thread_pool(1); + + for (values, expected) in [ + ([0.0, 3.0, f32::NAN], vec![vec![0], vec![1]]), + ([f32::NAN; 3], vec![]), + ] { + let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); + + let actual = pool + .install(|| partition::<_, L2, _>(ARCH, data, &config)) + .unwrap(); + + assert_eq!( + sorted_leaf_memberships(&actual), + sorted_leaf_memberships(&expected), + "values {values:?}" + ); + } + } +}