diff --git a/.cargo/mutants.toml b/.cargo/mutants.toml index c3a7e7a40..2d7811d4c 100644 --- a/.cargo/mutants.toml +++ b/.cargo/mutants.toml @@ -1,11 +1,12 @@ -# Kernel mutations that cannot be distinguished by native correctness tests. +# Mutations that cannot be distinguished by correctness tests. # V4 and Neon wrappers are exercised by SDE and AArch64 CI respectively. # The remaining transformations preserve outputs and only switch between -# equivalent scalar-tail or zero-clamp expressions. +# equivalent fast paths, scalar tails, zero clamps, or empty prune rounds. exclude_re = [ "diskann-pipnn/src/(leaf|partition)_kernel\\.rs:.*Target<.*V4", "diskann-pipnn/src/(leaf|partition)_kernel\\.rs:.*Target<.*Neon", "diskann-pipnn/src/leaf_kernel\\.rs:.*replace < with <= in pair_distance", "diskann-pipnn/src/partition_kernel\\.rs:.*replace \\* with / in process_(unary|binary)", "diskann-pipnn/src/leaf_kernel\\.rs:.*replace > with >= in .*run_simd", + "diskann/src/graph/prune\\.rs:[0-9]+:17: replace < with <= in robust_prune", ] diff --git a/Cargo.lock b/Cargo.lock index 0c9926781..8ea252f47 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -437,6 +437,7 @@ version = "0.55.0" dependencies = [ "anyhow", "bytemuck", + "criterion", "dashmap", "diskann-utils", "diskann-vector", diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index 14f6ba074..759e6769e 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -32,6 +32,7 @@ diskann-wide = { workspace = true } dashmap = { workspace = true, optional = true } [dev-dependencies] +criterion.workspace = true futures-util = { workspace = true, default-features = false } pin-project.workspace = true rand.workspace = true @@ -41,6 +42,11 @@ serde_json = { workspace = true } tokio = { workspace = true, features = ["macros", "sync"] } dashmap = { workspace = true } +[[bench]] +name = "robust_prune" +harness = false +required-features = ["testing"] + # Some 'cfg's in the source tree will be flagged by `cargo clippy -j 2 --workspace --no-deps --all-targets -- -D warnings` [lints.rust] unexpected_cfgs = { level = "warn", check-cfg = ['cfg(coverage)'] } diff --git a/diskann/benches/robust_prune.rs b/diskann/benches/robust_prune.rs new file mode 100644 index 000000000..d4a5cc36a --- /dev/null +++ b/diskann/benches/robust_prune.rs @@ -0,0 +1,140 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{hint::black_box, iter, time::Duration}; + +use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use diskann::{ + graph::{ + self, AdjacencyList, DiskANNIndex, + config::{MaxDegree, PruneKind}, + test::provider::{self as test_provider, Provider, StartPoint}, + }, + provider::NeighborAccessor, +}; +use diskann_vector::distance::Metric; + +const DEGREE: usize = 64; +const CANDIDATE_COUNTS: [usize; 3] = [64, 128, 750]; + +struct PruneCase { + index: DiskANNIndex, + strategy: test_provider::Strategy, +} + +impl PruneCase { + #[allow(clippy::unwrap_used)] // Deterministic benchmark fixture construction. + fn new(count: usize, kind: PruneKind, saturate: bool) -> Self { + let source = 0_u32; + let start_id = count as u32 + 1; + let mut source_neighbors = AdjacencyList::new(); + for candidate in 1..=count as u32 { + source_neighbors.push(candidate); + } + + let provider_config = test_provider::Config::new( + Metric::L2, + count.max(DEGREE), + StartPoint::new(start_id, vec![0.0]), + ) + .unwrap(); + let points = (0..=count as u32).map(|id| { + let neighbors = if id == source { + source_neighbors.clone() + } else { + AdjacencyList::new() + }; + (id, vec![id as f32], neighbors) + }); + let provider = Provider::new_from( + provider_config, + iter::once((start_id, AdjacencyList::new())), + points, + ) + .unwrap(); + let config = graph::config::Builder::new_with( + DEGREE, + MaxDegree::new(count.max(DEGREE)), + count.max(DEGREE), + kind, + |builder| { + builder + .alpha(1.2) + .saturate_after_prune(saturate) + .max_occlusion_size(count); + }, + ) + .build() + .unwrap(); + + Self { + index: DiskANNIndex::new(config, provider, None), + strategy: test_provider::Strategy::new(), + } + } + + #[allow(clippy::unwrap_used)] // A benchmark fixture failure must abort the sample. + async fn run(self) -> AdjacencyList { + self.index + .prune_range( + &self.strategy, + &test_provider::Context::default(), + iter::once(0), + ) + .await + .unwrap(); + + let mut neighbors = AdjacencyList::new(); + self.index + .provider() + .neighbors() + .get_neighbors(0, &mut neighbors) + .await + .unwrap(); + neighbors + } +} + +#[allow(clippy::unwrap_used)] // Runtime construction is benchmark harness setup. +fn benchmark_robust_prune(c: &mut Criterion) { + let runtime = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + let mut group = c.benchmark_group("vamana/robust-prune"); + + for count in CANDIDATE_COUNTS { + for (kind_name, kind) in [ + ("triangle", PruneKind::TriangleInequality), + ("occluding", PruneKind::Occluding), + ] { + for saturate in [false, true] { + let name = format!( + "{kind_name}/{count}-to-{DEGREE}/{}", + if saturate { "saturated" } else { "pruned" } + ); + group.throughput(Throughput::Elements(count as u64)); + group.bench_function(BenchmarkId::from_parameter(name), |bencher| { + bencher.iter_batched( + || PruneCase::new(count, kind, saturate), + |case| black_box(runtime.block_on(case.run())), + BatchSize::SmallInput, + ); + }); + } + } + } + + group.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(30) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(3)); + targets = benchmark_robust_prune +} +criterion_main!(benches); diff --git a/diskann/src/graph/index.rs b/diskann/src/graph/index.rs index 1df48072f..e97eefc6c 100644 --- a/diskann/src/graph/index.rs +++ b/diskann/src/graph/index.rs @@ -28,7 +28,8 @@ use super::{ self, Batch, InplaceDeleteStrategy, InsertStrategy, MultiInsertStrategy, PruneAccessor, PruneStrategy, SearchAccessor, SearchPostProcess, SearchStrategy, }, - internal::{BackedgeBuffer, SortedNeighbors, prune}, + internal::{BackedgeBuffer, SortedNeighbors}, + prune, search::{ PagedSearch, record::{NoopSearchRecord, SearchRecord, VisitedSearchRecord}, @@ -48,10 +49,7 @@ use crate::{ NeighborAccessorMut, SetElement, }, tracked_debug, tracked_error, tracked_trace, - utils::{ - IntoUsize, - async_tools::{self, DynamicBalancer}, - }, + utils::async_tools::{self, DynamicBalancer}, }; use diskann_utils::object_pool::{ObjectPool, PooledRef}; @@ -2378,7 +2376,8 @@ where view, |id| id == location, options, - ); + ) + .map_err(ANNError::opaque)?; Ok(()) } @@ -2452,7 +2451,8 @@ where view, |id| id == location, options, - ); + ) + .map_err(ANNError::opaque)?; Ok(()) } @@ -2530,7 +2530,8 @@ where view, |id| id == internal_id, options, - ); + ) + .map_err(ANNError::opaque)?; Ok(()) } @@ -2574,211 +2575,30 @@ where map: M, exclude: F, options: prune::Options, - ) where + ) -> Result<(), prune::RobustPruneError> + where M: View, C: for<'a, 'b> DistanceFunction, M::ElementRef<'b>, f32>, F: Fn(DP::InternalId) -> bool, { - assert!( - context.pool.len() <= u16::MAX.into_usize(), - "this has an upper bound set by `diskann::graph::Config` and should not exceed `u16::MAX`" + let policy = prune::Policy::new( + self.config.pruned_degree().get(), + self.config.alpha(), + self.config.prune_kind(), + options.force_saturate + || (self.config.saturate_after_prune() && self.config.alpha() > 1.0), ); - - let prune::Context { - pool, - states, - neighbors, - } = context; - - if pool.is_empty() { - neighbors.clear(); - return; - } - - let alpha = self.config.alpha(); - let degree = self.config.pruned_degree().get(); - - states.clear(); - states.resize(pool.len(), prune::State::default()); - - let mut current_alpha = 1.0f32; - let increment_factor = alpha.min(1.2); - - // To avoid many hash lookups, we pull out just the candidates we're going to prune - // into an auxiliary vector which can be accessed linearly. - // - // During the pruning phase, we store results by their relative position in the - // cache, and only resolve to their `local_id` at the end. - let cache: Vec<(f32, Option<_>)> = pool - .iter() - .map(|neighbor| { - // Filter out self loops. - let id = neighbor.id(); - if exclude(*id) { - (neighbor.distance(), None) - } else { - (neighbor.distance(), map.get(*id)) - } - }) - .collect(); - - // For an alpha value `A`, a candidate `i` is promoted to a neighbor if for all - // ``` - // max{j < i | j is a neighbor}(occlude_factor(i, j)) - // ``` - // This process happens with multiple values of `A`. - // - // We can compute this efficiently using the following rules: - // - // 1. For a candidate `i`, start scanning `j < i`, computing occlude factors. - // 2. If we find an occlude factor greater than `A`, record that `i` has visited - // `j`, stop computing occlude factors, and move on to `i + 1`. - // 3. If we reach `j == i - 1` with the maximum occlude factor less than `A`, then - // `i` gets promoted to a neighbor. - // - // On the implementation side, we use `states` in the following way: - // - // * `states[n].neighbor` is the **index** in `pool` of the `n`th **neighbor**. - // Note that a "neighbor" is a candidate that passes pruning. - // - // Very important: to get the index `j` in the above description, we need to - // check `pool[states[n].neighbor]`. - // - // This indexing naturally skips candidates `j` that have not been promoted to - // neighbors. - // - // * `states[i].occlude_factor` is the maximum occlude factor found for a candidate - // `i`. This gets set to `f32::MAX` when `i` is promoted to a neighbor which - // excludes it from future consideration. - // - // * `states[i].last_checked` is the highest value of `n` against which the - // occlude factor for `j = pool[states[n].neighbor]` has been checked. - // - // The maximum value this should reach is `i`. - // - // Note that we use `states` for both "candidate" and "neighbor" tracking. - let mut found = 0; - while found < degree { - for (i, (neighbor_distance, neighbor)) in cache.iter().enumerate() { - if found >= degree { - break; - } - - // The tracking states for candidate `i`. - let prune::State { - mut occlude_factor, - mut last_checked, - .. - } = states[i]; - - // If the occlusion factor for this neighbor is too high, skip it. - if occlude_factor > current_alpha { - continue; - } - - // Retrieval from the cache might not be perfect. - // - // This neighbor did not end up in the cache, then just skip it. - let neighbor = match neighbor { - Some(n) => n, - None => { - debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); - // SAFETY: We've already checked `states[i]`. - unsafe { states.get_unchecked_mut(i) }.occlude_factor = f32::MAX; - continue; - } - }; - - // Increment `position` until we've compared with all current entries in - // `result`. - // - // When the list is empty, the loop is skipped allowing the first undeleted - // element to be added. - while last_checked as usize != found { - let result_position = states[last_checked as usize].neighbor.into_usize(); - last_checked += 1; - - // If the position of this result in `pool` is greater than or equal - // the current working position, then skip this candidate. - if result_position >= i { - debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); - // SAFETY: We've already checked `states[i]`. - unsafe { states.get_unchecked_mut(i) }.last_checked = last_checked; - continue; - } - - // Otherwise, compute the distance between the result and this neighbor - // and update the occlude factor. - let distance = match &cache[result_position] { - (_, Some(v)) => { - computer.evaluate_similarity((*neighbor).reborrow(), v.reborrow()) - } - (_, None) => f32::MAX, - }; - - // Update occlude factor - occlude_factor = self.config.prune_kind().update_occlude_factor( - *neighbor_distance, - distance, - occlude_factor, - current_alpha, - ); - - // Check if the most recent update to the occlusion factor removes this - // neighbor from consideration. - if occlude_factor > current_alpha { - break; - } - } - - debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); - // SAFETY: We've already checked `states[i]`. - let state = unsafe { states.get_unchecked_mut(i) }; - - state.last_checked = last_checked; - if occlude_factor > current_alpha { - state.occlude_factor = occlude_factor; - continue; - } - - // This neighbor has passed all the requirements of being a candidate. - state.occlude_factor = f32::MAX; - - // This conversion should always succeed. - states[found].neighbor = i as u16; - found += 1; - } - - // Exit if we completed the final iteration. - if current_alpha == alpha { - break; - } - // Update current alpha for the next iteration. - current_alpha = (current_alpha * increment_factor).min(alpha); - } - - let mut guard = neighbors.resize(found); - std::iter::zip(guard.iter_mut(), states.iter()).for_each(|(d, s)| { - *d = *pool[s.neighbor.into_usize()].id(); - }); - guard.finish(found); - - debug_assert!(neighbors.len() <= degree, "max degree bound violated"); - - // Post processing saturation if enabled. - if options.force_saturate || (self.config.saturate_after_prune() && alpha > 1.0f32) { - for neighbor in context.pool.iter() { - if neighbors.len() >= degree { - break; - } - - if !exclude(*neighbor.id()) { - // `AdjacencyList` filters out duplicates. No need to explicitly - // check. - neighbors.push(*neighbor.id()); - } - } - } + let mut cache = Vec::new(); + prune::robust_prune( + context, + policy, + &mut cache, + |id| map.get(id), + |neighbor, selected| { + Ok(computer.evaluate_similarity((*neighbor).reborrow(), selected.reborrow())) + }, + exclude, + ) } /// Prune all nodes in `ids`. diff --git a/diskann/src/graph/internal/mod.rs b/diskann/src/graph/internal/mod.rs index ddaa6960e..3fa835fc1 100644 --- a/diskann/src/graph/internal/mod.rs +++ b/diskann/src/graph/internal/mod.rs @@ -8,5 +8,3 @@ pub(super) use sorted_neighbors::SortedNeighbors; mod backedge; pub(super) use backedge::BackedgeBuffer; - -pub(super) mod prune; diff --git a/diskann/src/graph/internal/prune.rs b/diskann/src/graph/internal/prune.rs deleted file mode 100644 index f5148be96..000000000 --- a/diskann/src/graph/internal/prune.rs +++ /dev/null @@ -1,176 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -use thiserror::Error; - -use super::SortedNeighbors; - -use crate::{ - ANNError, ANNErrorKind, error, graph::AdjacencyList, neighbor::Neighbor, utils::VectorId, -}; - -/// Options provided to prune. See the field-level documentation for more details. -/// -/// This struct should be kept cheap to construct. -#[derive(Debug, Clone, Copy)] -pub(crate) struct Options { - /// Force adjacency list saturation. - /// - /// Adjacency list saturation expands the post-pruning candidate list up to the - /// maximum degree by greedily adding skipped neighbors from the original candidate - /// pool. - pub(in crate::graph) force_saturate: bool, -} - -/// An aggregate of scratch space used by the pruning algorithm for allocation. -/// -/// The actual object passed to the pruning algorithms is [`Context`], which allows -/// sub-fields to be over-written as needed with local state if that is available instead. -#[derive(Debug)] -pub(crate) struct Scratch -where - I: VectorId, -{ - pub(in crate::graph) pool: Vec>, - pub(in crate::graph) states: Vec, - pub(in crate::graph) neighbors: AdjacencyList, -} - -impl Scratch -where - I: VectorId, -{ - /// Create a new empty scratch space. - /// - /// This function should not allocate. - pub(in crate::graph) fn new() -> Self { - Self { - pool: Vec::new(), - states: Vec::new(), - neighbors: AdjacencyList::new(), - } - } - - /// Convert `self` into a `Context`, truncating the internal `pool` list to a length of - /// `max_candidates`. - pub(in crate::graph) fn as_context(&mut self, max_candidates: usize) -> Context<'_, I> { - Context { - pool: SortedNeighbors::new(&mut self.pool, max_candidates), - states: &mut self.states, - neighbors: &mut self.neighbors, - } - } -} - -/// Arguments passed to the lowest-level pruning algorithm. -#[derive(Debug)] -pub(crate) struct Context<'ctx, I> -where - I: VectorId, -{ - /// Input: The list of candidates to prune. - pub(in crate::graph) pool: SortedNeighbors<'ctx, I>, - /// Scratch: State tracking for prune. - pub(in crate::graph) states: &'ctx mut Vec, - /// Output: The pruned candidates list. - pub(in crate::graph) neighbors: &'ctx mut AdjacencyList, -} - -/// Position-wise state tracking. -/// -/// Refer to the inline documentation in [`DiskANNIndex::occlude_list`] for documentation -/// on the use of these fields. -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct State { - /// The occlude factor for the pool item at the corresponding index. - pub(in crate::graph) occlude_factor: f32, - /// The index of the last checked neighbor. - pub(in crate::graph) last_checked: u16, - /// The candidate index of this neighbor. - pub(in crate::graph) neighbor: u16, -} - -#[derive(Debug, Clone, Copy, Error)] -#[error("retrieval of main vector id {} failed during prune aggregation", self.0)] -pub(crate) struct FailedVectorRetrieval(I) -where - I: VectorId; - -impl error::TransientError for FailedVectorRetrieval -where - I: VectorId, -{ - fn acknowledge(self, _why: D) - where - D: std::fmt::Display, - { - } - - #[track_caller] - #[inline(never)] - fn escalate(self, why: D) -> ANNError - where - D: std::fmt::Display, - { - ANNError::new(ANNErrorKind::IndexError, self).context(why.to_string()) - } -} - -/// Failure condition for [`DiskANNIndex::robust_prune_list`]. -/// -/// It's currently possible for retrieval of the id being pruned to fail due to a transient -/// error. We do not always want to escalate this as a hard error, and thus provide an -/// option for transient error handling. -#[derive(Debug)] -pub(crate) enum ListError -where - I: VectorId, -{ - /// A potentially transient error. - FailedVectorRetrieval(FailedVectorRetrieval), - /// A critical error. - Other(ANNError), -} - -impl ListError -where - I: VectorId, -{ - pub(in crate::graph) fn failed_retrieval(id: I) -> Self { - Self::FailedVectorRetrieval(FailedVectorRetrieval(id)) - } -} - -impl From for ListError -where - I: VectorId, -{ - fn from(err: ANNError) -> Self { - Self::Other(err) - } -} - -impl error::ToRanked for ListError -where - I: VectorId, -{ - type Transient = FailedVectorRetrieval; - type Error = ANNError; - - fn to_ranked(self) -> error::RankedError { - match self { - Self::FailedVectorRetrieval(err) => error::RankedError::Transient(err), - Self::Other(err) => error::RankedError::Error(err), - } - } - - fn from_transient(transient: Self::Transient) -> Self { - Self::FailedVectorRetrieval(transient) - } - - fn from_error(error: Self::Error) -> Self { - Self::Other(error) - } -} diff --git a/diskann/src/graph/mod.rs b/diskann/src/graph/mod.rs index 6bf4be7dd..dcbbeb42a 100644 --- a/diskann/src/graph/mod.rs +++ b/diskann/src/graph/mod.rs @@ -17,6 +17,8 @@ pub use config::Config; pub mod index; pub use index::DiskANNIndex; +pub mod prune; + mod start_point; pub use start_point::{SampleableForStart, StartPointStrategy}; diff --git a/diskann/src/graph/prune.rs b/diskann/src/graph/prune.rs new file mode 100644 index 000000000..49762d772 --- /dev/null +++ b/diskann/src/graph/prune.rs @@ -0,0 +1,404 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use thiserror::Error; + +use super::{config::PruneKind, internal::SortedNeighbors}; + +use crate::{ + ANNError, ANNErrorKind, error, + graph::AdjacencyList, + neighbor::Neighbor, + utils::{IntoUsize, VectorId}, +}; + +/// Options provided to prune. See the field-level documentation for more details. +/// +/// This struct should be kept cheap to construct. +#[derive(Debug, Clone, Copy)] +pub(crate) struct Options { + /// Force adjacency list saturation. + /// + /// Adjacency list saturation expands the post-pruning candidate list up to the + /// maximum degree by greedily adding skipped neighbors from the original candidate + /// pool. + pub(in crate::graph) force_saturate: bool, +} + +/// An aggregate of scratch space used by the pruning algorithm for allocation. +/// +/// The actual object passed to the pruning algorithms is [`Context`], which allows +/// sub-fields to be over-written as needed with local state if that is available instead. +#[derive(Debug)] +pub struct Scratch +where + I: VectorId, +{ + pub(in crate::graph) pool: Vec>, + pub(in crate::graph) states: Vec, + pub(in crate::graph) neighbors: AdjacencyList, +} + +impl Scratch +where + I: VectorId, +{ + /// Create a new empty scratch space. + /// + /// This function should not allocate. + pub fn new() -> Self { + Self { + pool: Vec::new(), + states: Vec::new(), + neighbors: AdjacencyList::new(), + } + } + + /// Convert `self` into a `Context`, truncating the internal `pool` list to a length of + /// `max_candidates`. + pub fn as_context(&mut self, max_candidates: usize) -> Context<'_, I> { + Context { + pool: SortedNeighbors::new(&mut self.pool, max_candidates), + states: &mut self.states, + neighbors: &mut self.neighbors, + } + } + + /// Candidate buffer used by callers before pruning. + pub fn candidates_mut(&mut self) -> &mut Vec> { + &mut self.pool + } + + /// The most recent pruned adjacency list. + pub fn neighbors(&self) -> &AdjacencyList { + &self.neighbors + } +} + +impl Default for Scratch +where + I: VectorId, +{ + fn default() -> Self { + Self::new() + } +} + +/// Arguments passed to the lowest-level pruning algorithm. +#[derive(Debug)] +pub struct Context<'ctx, I> +where + I: VectorId, +{ + /// Input: The list of candidates to prune. + pub(in crate::graph) pool: SortedNeighbors<'ctx, I>, + /// Scratch: State tracking for prune. + pub(in crate::graph) states: &'ctx mut Vec, + /// Output: The pruned candidates list. + pub(in crate::graph) neighbors: &'ctx mut AdjacencyList, +} + +/// Position-wise state tracking. +/// +/// Refer to the inline documentation in [`DiskANNIndex::occlude_list`] for documentation +/// on the use of these fields. +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct State { + /// The occlude factor for the pool item at the corresponding index. + pub(in crate::graph) occlude_factor: f32, + /// The index of the last checked neighbor. + pub(in crate::graph) last_checked: u16, + /// The candidate index of this neighbor. + pub(in crate::graph) neighbor: u16, +} + +/// Provider-independent policy for the Vamana robust-prune algorithm. +#[derive(Debug, Clone, Copy)] +pub struct Policy { + degree: usize, + alpha: f32, + prune_kind: PruneKind, + saturate: bool, +} + +impl Policy { + pub fn new(degree: usize, alpha: f32, prune_kind: PruneKind, saturate: bool) -> Self { + Self { + degree, + alpha, + prune_kind, + saturate, + } + } +} + +/// Failure returned by [`robust_prune`]. +#[derive(Debug, Error)] +pub enum RobustPruneError { + #[error("robust prune alpha must be finite and >= 1.0, got {0}")] + InvalidAlpha(f32), + #[error("robust prune supports at most {max} candidates, got {actual}")] + TooManyCandidates { actual: usize, max: usize }, + #[error("failed to reserve robust-prune workspace: {0}")] + Allocation(#[source] std::collections::TryReserveError), + #[error("distance computation failed: {0}")] + Distance(E), +} + +/// Run Vamana's robust-prune state machine over an already sorted candidate pool. +/// +/// Data access and distance evaluation stay with the caller so providers may fetch +/// asynchronously before entering this synchronous kernel, while in-memory builders +/// can use contiguous vector storage and specialized distance functions. +pub fn robust_prune( + context: &mut Context<'_, I>, + policy: Policy, + cache: &mut Vec<(f32, Option)>, + mut lookup: L, + mut distance: D, + exclude: X, +) -> Result<(), RobustPruneError> +where + I: VectorId, + L: FnMut(I) -> Option, + D: FnMut(&V, &V) -> Result, + X: Fn(I) -> bool, +{ + if !policy.alpha.is_finite() || policy.alpha < 1.0 { + return Err(RobustPruneError::InvalidAlpha(policy.alpha)); + } + if context.pool.len() > u16::MAX as usize { + return Err(RobustPruneError::TooManyCandidates { + actual: context.pool.len(), + max: u16::MAX as usize, + }); + } + + let Context { + pool, + states, + neighbors, + } = context; + + cache.clear(); + if pool.is_empty() { + neighbors.clear(); + return Ok(()); + } + + states + .try_reserve(pool.len().saturating_sub(states.len())) + .map_err(RobustPruneError::Allocation)?; + cache + .try_reserve(pool.len().saturating_sub(cache.len())) + .map_err(RobustPruneError::Allocation)?; + states.clear(); + states.resize(pool.len(), State::default()); + + let mut current_alpha = 1.0f32; + let increment_factor = policy.alpha.min(1.2); + + cache.extend(pool.iter().map(|neighbor| { + if exclude(*neighbor.id()) { + (neighbor.distance(), None) + } else { + (neighbor.distance(), lookup(*neighbor.id())) + } + })); + + // This is Vamana's existing lazy occlusion state machine. Keep the candidate and + // neighbor positions in one `State` array so retries at larger alpha values resume + // instead of recomputing comparisons. + let mut found = 0; + while found < policy.degree { + for (i, (neighbor_distance, neighbor)) in cache.iter().enumerate() { + if found >= policy.degree { + break; + } + + let State { + mut occlude_factor, + mut last_checked, + .. + } = states[i]; + + if occlude_factor > current_alpha { + continue; + } + + let neighbor = match neighbor { + Some(n) => n, + None => { + debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); + // SAFETY: `i` comes from iterating `cache`, which has the same length + // as `states`. + unsafe { states.get_unchecked_mut(i) }.occlude_factor = f32::MAX; + continue; + } + }; + + while last_checked as usize != found { + let result_position = states[last_checked as usize].neighbor.into_usize(); + last_checked += 1; + + if result_position >= i { + debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); + // SAFETY: `i` comes from iterating `cache`, which has the same length + // as `states`. + unsafe { states.get_unchecked_mut(i) }.last_checked = last_checked; + continue; + } + + let pair_distance = match &cache[result_position] { + (_, Some(value)) => { + distance(neighbor, value).map_err(RobustPruneError::Distance)? + } + (_, None) => f32::MAX, + }; + + occlude_factor = policy.prune_kind.update_occlude_factor( + *neighbor_distance, + pair_distance, + occlude_factor, + current_alpha, + ); + + if occlude_factor > current_alpha { + break; + } + } + + debug_assert!(states.get(i).is_some(), "index {i} is out of bounds"); + // SAFETY: `i` comes from iterating `cache`, which has the same length as + // `states`. + let state = unsafe { states.get_unchecked_mut(i) }; + + state.last_checked = last_checked; + if occlude_factor > current_alpha { + state.occlude_factor = occlude_factor; + continue; + } + + state.occlude_factor = f32::MAX; + states[found].neighbor = i as u16; + found += 1; + } + + if current_alpha == policy.alpha { + break; + } + current_alpha = (current_alpha * increment_factor).min(policy.alpha); + } + + let mut guard = neighbors.resize(found); + std::iter::zip(guard.iter_mut(), states.iter()).for_each(|(destination, state)| { + *destination = *pool[state.neighbor.into_usize()].id(); + }); + guard.finish(found); + + debug_assert!( + neighbors.len() <= policy.degree, + "max degree bound violated" + ); + + if policy.saturate { + for neighbor in pool.iter() { + if neighbors.len() >= policy.degree { + break; + } + if !exclude(*neighbor.id()) { + neighbors.push(*neighbor.id()); + } + } + } + + Ok(()) +} + +#[derive(Debug, Clone, Copy, Error)] +#[error("retrieval of main vector id {} failed during prune aggregation", self.0)] +pub(crate) struct FailedVectorRetrieval(I) +where + I: VectorId; + +impl error::TransientError for FailedVectorRetrieval +where + I: VectorId, +{ + fn acknowledge(self, _why: D) + where + D: std::fmt::Display, + { + } + + #[track_caller] + #[inline(never)] + fn escalate(self, why: D) -> ANNError + where + D: std::fmt::Display, + { + ANNError::new(ANNErrorKind::IndexError, self).context(why.to_string()) + } +} + +/// Failure condition for [`DiskANNIndex::robust_prune_list`]. +/// +/// It's currently possible for retrieval of the id being pruned to fail due to a transient +/// error. We do not always want to escalate this as a hard error, and thus provide an +/// option for transient error handling. +#[derive(Debug)] +pub(crate) enum ListError +where + I: VectorId, +{ + /// A potentially transient error. + FailedVectorRetrieval(FailedVectorRetrieval), + /// A critical error. + Other(ANNError), +} + +impl ListError +where + I: VectorId, +{ + pub(in crate::graph) fn failed_retrieval(id: I) -> Self { + Self::FailedVectorRetrieval(FailedVectorRetrieval(id)) + } +} + +impl From for ListError +where + I: VectorId, +{ + fn from(err: ANNError) -> Self { + Self::Other(err) + } +} + +impl error::ToRanked for ListError +where + I: VectorId, +{ + type Transient = FailedVectorRetrieval; + type Error = ANNError; + + fn to_ranked(self) -> error::RankedError { + match self { + Self::FailedVectorRetrieval(err) => error::RankedError::Transient(err), + Self::Other(err) => error::RankedError::Error(err), + } + } + + fn from_transient(transient: Self::Transient) -> Self { + Self::FailedVectorRetrieval(transient) + } + + fn from_error(error: Self::Error) -> Self { + Self::Other(error) + } +} + +#[cfg(test)] +mod tests; diff --git a/diskann/src/graph/prune/tests.rs b/diskann/src/graph/prune/tests.rs new file mode 100644 index 000000000..38c416813 --- /dev/null +++ b/diskann/src/graph/prune/tests.rs @@ -0,0 +1,205 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; +use crate::{ + error::{RankedError, ToRanked, TransientError}, + neighbor::Neighbor, +}; + +#[derive(Debug, PartialEq)] +struct DistanceFailure; + +fn run( + scratch: &mut Scratch, + policy: Policy, + lookup: impl FnMut(I) -> Option, + distance: impl FnMut(&V, &V) -> Result, +) -> Result<(), RobustPruneError> +where + I: VectorId, +{ + let max_candidates = scratch.candidates_mut().len(); + let mut context = scratch.as_context(max_candidates); + robust_prune( + &mut context, + policy, + &mut Vec::new(), + lookup, + distance, + |_| false, + ) +} + +#[test] +fn robust_prune_propagates_distance_failure() { + let mut scratch = Scratch::new(); + scratch + .candidates_mut() + .extend([Neighbor::new(1_u32, 1.0), Neighbor::new(2_u32, 2.0)]); + + let error = run( + &mut scratch, + Policy::new(2, 1.2, PruneKind::TriangleInequality, false), + Some, + |_, _| Err(DistanceFailure), + ) + .unwrap_err(); + + assert!(matches!(error, RobustPruneError::Distance(DistanceFailure))); +} + +#[test] +fn robust_prune_validates_alpha() { + for alpha in [f32::NAN, f32::INFINITY, 0.999] { + let mut scratch = Scratch::::new(); + let error = run( + &mut scratch, + Policy::new(1, alpha, PruneKind::TriangleInequality, false), + Some, + |_, _| Ok::<_, std::convert::Infallible>(0.0), + ) + .unwrap_err(); + assert!( + matches!(error, RobustPruneError::InvalidAlpha(value) if value.to_bits() == alpha.to_bits()) + ); + } + + let mut scratch = Scratch::::new(); + run( + &mut scratch, + Policy::new(1, 1.0, PruneKind::TriangleInequality, false), + Some, + |_, _| Ok::<_, std::convert::Infallible>(0.0), + ) + .unwrap(); +} + +#[test] +fn robust_prune_rejects_more_than_u16_candidates() { + let mut scratch = Scratch::new(); + scratch + .candidates_mut() + .extend((0..u16::MAX as u32).map(|id| Neighbor::new(id, id as f32))); + run( + &mut scratch, + Policy::new(0, 1.2, PruneKind::TriangleInequality, false), + Some, + |_, _| Ok::<_, std::convert::Infallible>(0.0), + ) + .unwrap(); + + scratch + .candidates_mut() + .push(Neighbor::new(u16::MAX as u32, u16::MAX as f32)); + let error = run( + &mut scratch, + Policy::new(1, 1.2, PruneKind::TriangleInequality, false), + Some, + |_, _| Ok::<_, std::convert::Infallible>(0.0), + ) + .unwrap_err(); + + assert!(matches!( + error, + RobustPruneError::TooManyCandidates { + actual, + max + } if actual == u16::MAX as usize + 1 && max == u16::MAX as usize + )); +} + +#[test] +fn scratch_exposes_pruned_neighbors() { + let mut scratch = Scratch::new(); + scratch + .candidates_mut() + .extend([Neighbor::new(1_u32, 1.0), Neighbor::new(2_u32, 2.0)]); + + run( + &mut scratch, + Policy::new(1, 1.2, PruneKind::TriangleInequality, false), + Some, + |left, right| Ok::<_, std::convert::Infallible>(left.abs_diff(*right) as f32), + ) + .unwrap(); + + assert_eq!(&**scratch.neighbors(), &[1]); +} + +#[test] +fn robust_prune_excludes_candidates_before_lookup() { + let mut scratch = Scratch::default(); + scratch + .candidates_mut() + .extend([Neighbor::new(1_u32, 1.0), Neighbor::new(2_u32, 2.0)]); + let mut context = scratch.as_context(2); + robust_prune( + &mut context, + Policy::new(1, 1.2, PruneKind::TriangleInequality, false), + &mut Vec::new(), + Some, + |left, right| Ok::<_, std::convert::Infallible>(left.abs_diff(*right) as f32), + |id| id == 1, + ) + .unwrap(); + + assert_eq!(&**scratch.neighbors(), &[2]); +} + +#[test] +fn list_errors_preserve_transient_and_fatal_rank() { + let transient = FailedVectorRetrieval(7_u32); + let escalated = transient.escalate("test escalation"); + assert!(escalated.to_string().contains("test escalation")); + + assert!(matches!( + ListError::failed_retrieval(8_u32).to_ranked(), + RankedError::Transient(FailedVectorRetrieval(8)) + )); + assert!(matches!( + as ToRanked>::from_transient(FailedVectorRetrieval(9)), + ListError::FailedVectorRetrieval(FailedVectorRetrieval(9)) + )); + + let fatal = ANNError::new( + ANNErrorKind::IndexError, + std::io::Error::other("fatal prune test error"), + ); + assert!(matches!(ListError::::from(fatal), ListError::Other(_))); + let fatal = ANNError::new( + ANNErrorKind::IndexError, + std::io::Error::other("fatal prune test error"), + ); + assert!(matches!( + as ToRanked>::from_error(fatal).to_ranked(), + RankedError::Error(_) + )); +} + +#[test] +fn robust_prune_clears_output_for_an_empty_pool() { + let mut scratch = Scratch::new(); + scratch.candidates_mut().push(Neighbor::new(1_u32, 1.0)); + run( + &mut scratch, + Policy::new(1, 1.2, PruneKind::TriangleInequality, false), + Some, + |_, _| Ok::<_, std::convert::Infallible>(0.0), + ) + .unwrap(); + assert_eq!(&**scratch.neighbors(), &[1]); + + scratch.candidates_mut().clear(); + run( + &mut scratch, + Policy::new(usize::MAX, 1.2, PruneKind::TriangleInequality, false), + Some, + |_, _| Ok::<_, std::convert::Infallible>(0.0), + ) + .unwrap(); + + assert!(scratch.neighbors().is_empty()); +} diff --git a/diskann/src/graph/test/cases/mod.rs b/diskann/src/graph/test/cases/mod.rs index 69f502d1f..f2cd6c26a 100644 --- a/diskann/src/graph/test/cases/mod.rs +++ b/diskann/src/graph/test/cases/mod.rs @@ -12,6 +12,7 @@ mod inline; mod inplace_delete; mod multihop; mod paged_search; +mod prune; mod range_search; /// Set to `true` and recompile to include full adjacency list state in participating diff --git a/diskann/src/graph/test/cases/prune.rs b/diskann/src/graph/test/cases/prune.rs new file mode 100644 index 000000000..0c5eb318d --- /dev/null +++ b/diskann/src/graph/test/cases/prune.rs @@ -0,0 +1,309 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::iter; + +use diskann_vector::distance::Metric; + +use crate::{ + graph::{ + self, AdjacencyList, DiskANNIndex, + config::PruneKind, + test::provider::{self as test_provider, Provider, StartPoint}, + }, + provider::NeighborAccessor, +}; + +struct PruneCase { + index: DiskANNIndex, + source: u32, +} + +struct PruneConfig { + metric: Metric, + source: u32, + degree: usize, + alpha: f32, + prune_kind: PruneKind, + saturate: bool, + max_occlusion_size: usize, +} + +impl PruneCase { + fn new( + vectors: Vec>, + candidates: impl IntoIterator, + config: PruneConfig, + ) -> Self { + let PruneConfig { + metric, + source, + degree, + alpha, + prune_kind, + saturate, + max_occlusion_size, + } = config; + let dimensions = vectors.first().expect("a source vector is required").len(); + assert!(vectors.iter().all(|vector| vector.len() == dimensions)); + assert!((source as usize) < vectors.len()); + + let mut source_neighbors = AdjacencyList::new(); + for candidate in candidates { + source_neighbors.push(candidate); + } + + let provider_max_degree = source_neighbors.len().max(degree); + let start_id = vectors.len() as u32; + let provider_config = test_provider::Config::new( + metric, + provider_max_degree, + StartPoint::new(start_id, vec![0.0; dimensions]), + ) + .unwrap(); + let points = vectors.into_iter().enumerate().map(|(id, vector)| { + let neighbors = if id as u32 == source { + source_neighbors.clone() + } else { + AdjacencyList::new() + }; + (id as u32, vector, neighbors) + }); + let provider = Provider::new_from( + provider_config, + iter::once((start_id, AdjacencyList::new())), + points, + ) + .unwrap(); + + let config = graph::config::Builder::new_with( + degree, + graph::config::MaxDegree::new(provider_max_degree), + 10, + prune_kind, + |builder| { + builder + .alpha(alpha) + .saturate_after_prune(saturate) + .max_occlusion_size(max_occlusion_size); + }, + ) + .build() + .unwrap(); + + Self { + index: DiskANNIndex::new(config, provider, None), + source, + } + } + + async fn run(self, strategy: &test_provider::Strategy) -> AdjacencyList { + self.index + .prune_range( + strategy, + &test_provider::Context::default(), + iter::once(self.source), + ) + .await + .unwrap(); + + let mut neighbors = AdjacencyList::new(); + self.index + .provider() + .neighbors() + .get_neighbors(self.source, &mut neighbors) + .await + .unwrap(); + neighbors + } +} + +fn l2_case( + positions: &[f32], + candidates: impl IntoIterator, + degree: usize, + alpha: f32, + saturate: bool, + max_occlusion_size: usize, +) -> PruneCase { + PruneCase::new( + positions.iter().map(|position| vec![*position]).collect(), + candidates, + PruneConfig { + metric: Metric::L2, + source: 0, + degree, + alpha, + prune_kind: PruneKind::TriangleInequality, + saturate, + max_occlusion_size, + }, + ) +} + +#[tokio::test(flavor = "current_thread")] +async fn rows_at_or_below_degree_are_unchanged() { + for (candidates, degree) in [(vec![], 2), (vec![2], 2), (vec![2, 1], 2)] { + let expected = candidates.clone(); + let actual = l2_case(&[0.0, 1.0, -1.0], candidates, degree, 1.2, false, 10) + .run(&test_provider::Strategy::new()) + .await; + assert_eq!(&*actual, expected); + } +} + +#[tokio::test(flavor = "current_thread")] +async fn equal_distances_keep_current_sorted_neighbor_order() { + let case = PruneCase::new( + vec![ + vec![0.0, 0.0, 0.0], + vec![1.0, 0.0, 0.0], + vec![0.0, 1.0, 0.0], + vec![0.0, 0.0, 1.0], + ], + [3, 1, 2], + PruneConfig { + metric: Metric::L2, + source: 0, + degree: 2, + alpha: 1.2, + prune_kind: PruneKind::TriangleInequality, + saturate: false, + max_occlusion_size: 10, + }, + ); + + assert_eq!(&*case.run(&test_provider::Strategy::new()).await, &[2, 1]); +} + +#[tokio::test(flavor = "current_thread")] +async fn triangle_prune_revisits_candidates_across_alpha_rounds() { + let once = l2_case( + &[0.0, 1.0, 8.0, 12.0, 16.0], + [1, 2, 3, 4], + 3, + 1.2, + false, + 10, + ) + .run(&test_provider::Strategy::new()) + .await; + let multiple = l2_case( + &[0.0, 1.0, 8.0, 12.0, 16.0], + [1, 2, 3, 4], + 3, + 1.44, + false, + 10, + ) + .run(&test_provider::Strategy::new()) + .await; + + assert_eq!(&*once, &[1, 3]); + assert_eq!(&*multiple, &[1, 3, 2]); +} + +#[tokio::test(flavor = "current_thread")] +async fn inner_product_uses_occluding_prune() { + let case = PruneCase::new( + vec![ + vec![1.0, 0.0], + vec![3.0, 0.0], + vec![2.0, 0.0], + vec![1.0, 1.0], + ], + [1, 2, 3], + PruneConfig { + metric: Metric::InnerProduct, + source: 0, + degree: 2, + alpha: 1.2, + prune_kind: PruneKind::Occluding, + saturate: false, + max_occlusion_size: 10, + }, + ); + + assert_eq!(&*case.run(&test_provider::Strategy::new()).await, &[1, 2]); +} + +#[tokio::test(flavor = "current_thread")] +async fn saturation_appends_candidates_in_pool_order() { + let unsaturated = l2_case(&[0.0, 1.0, 2.0, 3.0, 4.0], 1..=4, 3, 1.2, false, 10) + .run(&test_provider::Strategy::new()) + .await; + let saturated = l2_case(&[0.0, 1.0, 2.0, 3.0, 4.0], 1..=4, 3, 1.2, true, 10) + .run(&test_provider::Strategy::new()) + .await; + + assert_eq!(&*unsaturated, &[1]); + assert_eq!(&*saturated, &[1, 2, 3]); +} + +#[tokio::test(flavor = "current_thread")] +async fn configured_saturation_requires_alpha_above_one() { + let neighbors = l2_case(&[0.0, 1.0, 2.0, 3.0, 4.0], 1..=4, 3, 1.0, true, 10) + .run(&test_provider::Strategy::new()) + .await; + + assert_eq!(&*neighbors, &[1]); +} + +#[tokio::test(flavor = "current_thread")] +async fn self_and_unavailable_candidates_are_excluded() { + let case = l2_case(&[0.0, -1.0, 2.0, 3.0], [0, 1, 2, 3], 2, 1.2, false, 10); + let strategy = test_provider::Strategy::with_transient(true, [2]); + + assert_eq!(&*case.run(&strategy).await, &[1, 3]); +} + +#[tokio::test(flavor = "current_thread")] +async fn max_occlusion_size_truncates_to_nearest_candidates() { + let case = PruneCase::new( + vec![ + vec![0.0, 0.0], + vec![1.0, 0.0], + vec![0.0, 2.0], + vec![-3.0, 0.0], + vec![0.0, -4.0], + ], + [4, 3, 2, 1], + PruneConfig { + metric: Metric::L2, + source: 0, + degree: 3, + alpha: 1.2, + prune_kind: PruneKind::TriangleInequality, + saturate: false, + max_occlusion_size: 2, + }, + ); + + assert_eq!(&*case.run(&test_provider::Strategy::new()).await, &[1, 2]); +} + +#[tokio::test(flavor = "current_thread")] +async fn maximum_u16_candidate_pool_is_supported() { + let num_candidates = u16::MAX as usize; + let vectors = (0..=num_candidates) + .map(|position| vec![position as f32]) + .collect(); + let case = PruneCase::new( + vectors, + 1..=u16::MAX as u32, + PruneConfig { + metric: Metric::L2, + source: 0, + degree: 1, + alpha: 1.2, + prune_kind: PruneKind::TriangleInequality, + saturate: false, + max_occlusion_size: num_candidates, + }, + ); + + let strategy = test_provider::Strategy::with_transient(true, 1..u16::MAX as u32); + assert_eq!(&*case.run(&strategy).await, &[u16::MAX as u32]); +}