diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 1df300d528..2cd7cbc322 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -147,3 +147,12 @@ jobs: graph::pipnn::leaf_kernel::tests::rank_leaf_dots_tests cargo +nightly miri test --locked -p diskann --features pipnn --lib \ graph::pipnn::partition_kernel::tests::rank_leader_dots_tests + + - name: PiPNN HashPrune pointer boundaries + env: + MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance + run: | + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::lsh::tests + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::hash_prune::tests diff --git a/Cargo.lock b/Cargo.lock index f756d1dc63..d441451904 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -445,10 +445,12 @@ dependencies = [ "futures-util", "half", "hashbrown 0.16.1", + "libc", "num-traits", "parking_lot", "pin-project", "rand", + "rand_distr", "rayon", "relative-path 2.0.1", "rstest", diff --git a/diskann-benchmark/src/index/build.rs b/diskann-benchmark/src/index/build.rs index f8b2de462b..55e7d50c2d 100644 --- a/diskann-benchmark/src/index/build.rs +++ b/diskann-benchmark/src/index/build.rs @@ -142,12 +142,15 @@ where let started = std::time::Instant::now(); let adjacency = { - let context = diskann::graph::pipnn::PiPNNBuildContext::new( + let mut context = diskann::graph::pipnn::PiPNNBuildContext::new( parameters.into(), &graph, metric, &pool, )?; + if let Some(hash_prune) = ¶meters.hash_prune { + context = context.with_hash_prune(hash_prune.into())?; + } diskann::graph::pipnn::build_graph(data.as_view(), &context)? }; let start_points = input diff --git a/diskann-disk/src/build/builder/build.rs b/diskann-disk/src/build/builder/build.rs index 7b7e13e551..cb05020768 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -77,8 +77,12 @@ where index_writer: DiskIndexWriter, ) -> ANNResult { #[cfg(feature = "pipnn")] - if let Some(config) = disk_build_param.pipnn_config() { - config.validate()?; + if let Some(parameters) = disk_build_param.pipnn_parameters() { + diskann::graph::pipnn::PiPNNConfig::from(parameters).validate()?; + if let Some(hash_prune) = ¶meters.hash_prune { + diskann::graph::pipnn::HashPruneConfig::from(hash_prune) + .validate_for_degree(index_configuration.config.pruned_degree().get())?; + } } let pq_storage = PQStorage::new( @@ -182,8 +186,8 @@ where async fn build_graph(&mut self, pool: RayonThreadPoolRef<'_>) -> ANNResult<()> { #[cfg(feature = "pipnn")] - if let Some(config) = self.disk_build_param.pipnn_config() { - return pipnn::build_graph(self, pool, config); + if let Some(parameters) = self.disk_build_param.pipnn_parameters().cloned() { + return pipnn::build_graph(self, pool, ¶meters); } match determine_build_strategy::( diff --git a/diskann-disk/src/build/builder/build/pipnn.rs b/diskann-disk/src/build/builder/build/pipnn.rs index 488ea8ae7e..bc3f3f4d67 100644 --- a/diskann-disk/src/build/builder/build/pipnn.rs +++ b/diskann-disk/src/build/builder/build/pipnn.rs @@ -11,7 +11,7 @@ //! //! PiPNN and Vamana use the same disk graph format. -use diskann::graph::pipnn::{PiPNNBuildContext, PiPNNConfig}; +use diskann::graph::pipnn::PiPNNBuildContext; use diskann::{utils::VectorRepr, ANNError, ANNResult}; use diskann_providers::{ storage::{save_adjacency_graph, StorageReadProvider, StorageWriteProvider}, @@ -20,13 +20,13 @@ use diskann_providers::{ use diskann_utils::io::{read_bin, Metadata}; use super::{u32_try_from, DiskIndexBuilder}; -use crate::data_model::GraphDataType; +use crate::{data_model::GraphDataType, PiPNNParameters}; /// Build PiPNN adjacency and persist it through the canonical disk graph writer. pub(super) fn build_graph( builder: &DiskIndexBuilder<'_, Data, StorageProvider>, pool: RayonThreadPoolRef<'_>, - config: PiPNNConfig, + parameters: &PiPNNParameters, ) -> ANNResult<()> where Data: GraphDataType, @@ -55,12 +55,15 @@ where // supplied Rayon pool. let data = read_bin::(&mut builder.storage_provider.open_reader(&data_path)?)?; - let context = PiPNNBuildContext::new( - config, + let mut context = PiPNNBuildContext::new( + parameters.into(), &builder.index_configuration.config, builder.index_configuration.dist_metric, pool.as_rayon(), )?; + if let Some(hash_prune) = ¶meters.hash_prune { + context = context.with_hash_prune(hash_prune.into())?; + } let adjacency = diskann::graph::pipnn::build_graph(data.as_view(), &context)?; // The disk header requires a start point. Use the same sampled medoid policy @@ -116,6 +119,7 @@ mod tests { fanout: vec![10, 3], k: 2, replicas: 1, + hash_prune: Some(crate::HashPruneParameters::default()), } } @@ -199,7 +203,7 @@ mod tests { let builder = builder(&storage, 3, 8, 1.0, 1.2, parameters.clone()); let pool = create_thread_pool(1).unwrap(); - let error = super::build_graph(&builder, pool.as_ref(), (¶meters).into()).unwrap_err(); + let error = super::build_graph(&builder, pool.as_ref(), ¶meters).unwrap_err(); assert!(format!("{error:?}").contains("configured point count 3")); assert!(!storage.exists(&builder.index_writer.get_mem_index_file())); } @@ -213,7 +217,7 @@ mod tests { let builder = builder(&storage, points, dimensions, 1.0, 1.2, parameters.clone()); let pool = create_thread_pool(1).unwrap(); - super::build_graph(&builder, pool.as_ref(), (¶meters).into()).unwrap(); + super::build_graph(&builder, pool.as_ref(), ¶meters).unwrap(); let mut header = [0_u8; 24]; std::io::Read::read_exact( @@ -274,4 +278,34 @@ mod tests { assert!(format!("{error:?}").contains("c_max must be greater than zero")); } + + #[test] + fn builder_rejects_hash_prune_capacity_before_quantizer_artifacts() { + let storage = VirtualStorageProvider::new_memory(); + let parameters = PiPNNParameters { + hash_prune: Some(crate::HashPruneParameters { + num_hash_planes: 12, + l_max: 16, + final_prune: true, + }), + ..PiPNNParameters::default() + }; + let params = DiskIndexBuildParameters::new_pipnn( + MemoryBudget::try_from_gb(1.0).unwrap(), + NumPQChunks::new_with(1, 1).unwrap(), + parameters, + ); + let config = IndexConfiguration::new(Metric::L2, 1, 1, ONE, 1, graph_config(32, 1.2)); + let writer = + DiskIndexWriter::new("/data.fbin".into(), "/index".into(), None, 4096).unwrap(); + + let error = match DiskIndexBuilder::, _>::new(&storage, params, config, writer) { + Ok(_) => panic!("HashPrune capacity below graph degree must be rejected"), + Err(error) => error, + }; + + assert!(format!("{error:?}").contains("must be at least the graph degree (32)")); + assert!(!storage.exists("/index_pq_pivots.bin")); + assert!(!storage.exists("/index_pq_compressed.bin")); + } } diff --git a/diskann-disk/src/build/configuration/build_algorithm.rs b/diskann-disk/src/build/configuration/build_algorithm.rs index d33e6b0b89..787c8f3e3d 100644 --- a/diskann-disk/src/build/configuration/build_algorithm.rs +++ b/diskann-disk/src/build/configuration/build_algorithm.rs @@ -29,6 +29,43 @@ pub struct PiPNNParameters { pub k: usize, /// Number of independent partition passes. pub replicas: usize, + /// HashPrune policy. `None` keeps all unique direct candidates. + pub hash_prune: Option, +} + +/// HashPrune parameters in the JSON build configuration. +#[cfg(feature = "pipnn")] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct HashPruneParameters { + /// Number of random-hyperplane sketch dimensions. + pub num_hash_planes: usize, + /// Maximum number of candidates retained per point. + pub l_max: usize, + /// Apply Vamana RobustPrune after reservoir extraction. + pub final_prune: bool, +} + +#[cfg(feature = "pipnn")] +impl Default for HashPruneParameters { + fn default() -> Self { + Self { + num_hash_planes: 12, + l_max: 64, + final_prune: true, + } + } +} + +#[cfg(feature = "pipnn")] +impl From<&HashPruneParameters> for diskann::graph::pipnn::HashPruneConfig { + fn from(config: &HashPruneParameters) -> Self { + Self { + num_hash_planes: config.num_hash_planes, + l_max: config.l_max, + final_prune: config.final_prune, + } + } } #[cfg(feature = "pipnn")] @@ -41,6 +78,7 @@ impl Default for PiPNNParameters { fanout: vec![8, 3], k: 2, replicas: 1, + hash_prune: None, } } } @@ -116,6 +154,15 @@ mod tests { assert_eq!(config.fanout, [10, 3]); assert_eq!(config.k, 3); assert_eq!(config.replicas, 1); + assert_eq!(config.hash_prune, None); + + let explicit: BuildAlgorithm = + serde_json::from_str(r#"{"algorithm":"PiPNN","hash_prune":{}}"#).unwrap(); + let BuildAlgorithm::PiPNN(explicit) = explicit else { + panic!("expected PiPNN"); + }; + assert_eq!(explicit.hash_prune, Some(HashPruneParameters::default())); + assert!( serde_json::from_str::(r#"{"algorithm":"PiPNN","l_max":72}"#).is_err() ); diff --git a/diskann-disk/src/build/configuration/disk_index_build_parameter.rs b/diskann-disk/src/build/configuration/disk_index_build_parameter.rs index a8b4017529..8adc7a14f8 100644 --- a/diskann-disk/src/build/configuration/disk_index_build_parameter.rs +++ b/diskann-disk/src/build/configuration/disk_index_build_parameter.rs @@ -198,9 +198,9 @@ impl DiskIndexBuildParameters { } #[cfg(feature = "pipnn")] - pub(crate) fn pipnn_config(&self) -> Option { + pub(crate) fn pipnn_parameters(&self) -> Option<&PiPNNParameters> { match &self.build_algorithm { - BuildAlgorithm::PiPNN(config) => Some(config.into()), + BuildAlgorithm::PiPNN(config) => Some(config), BuildAlgorithm::Vamana => None, } } diff --git a/diskann-disk/src/build/configuration/mod.rs b/diskann-disk/src/build/configuration/mod.rs index a7e343fb57..d2a26bba41 100644 --- a/diskann-disk/src/build/configuration/mod.rs +++ b/diskann-disk/src/build/configuration/mod.rs @@ -5,7 +5,7 @@ pub mod build_algorithm; pub use build_algorithm::BuildAlgorithm; #[cfg(feature = "pipnn")] -pub use build_algorithm::PiPNNParameters; +pub use build_algorithm::{HashPruneParameters, PiPNNParameters}; pub mod disk_index_build_parameter; pub use disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks}; diff --git a/diskann-disk/src/build/mod.rs b/diskann-disk/src/build/mod.rs index 27f4c124aa..c3f304664c 100644 --- a/diskann-disk/src/build/mod.rs +++ b/diskann-disk/src/build/mod.rs @@ -12,9 +12,9 @@ pub mod builder; pub mod configuration; // Re-export key types for convenience -#[cfg(feature = "pipnn")] -pub use configuration::PiPNNParameters; pub use configuration::{ disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters, QuantizationType, }; +#[cfg(feature = "pipnn")] +pub use configuration::{HashPruneParameters, PiPNNParameters}; diff --git a/diskann-disk/src/lib.rs b/diskann-disk/src/lib.rs index 5d9e6c368d..1704f93cf8 100644 --- a/diskann-disk/src/lib.rs +++ b/diskann-disk/src/lib.rs @@ -14,12 +14,12 @@ pub(crate) mod test_utils; pub mod error; pub mod build; -#[cfg(feature = "pipnn")] -pub use build::PiPNNParameters; pub use build::{ disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters, QuantizationType, }; +#[cfg(feature = "pipnn")] +pub use build::{HashPruneParameters, PiPNNParameters}; pub mod data_model; pub mod search; diff --git a/diskann-vector/src/lib.rs b/diskann-vector/src/lib.rs index e88dc12c9c..009b8da006 100644 --- a/diskann-vector/src/lib.rs +++ b/diskann-vector/src/lib.rs @@ -38,14 +38,17 @@ pub mod distance; pub mod norm; cfg_if::cfg_if! { - if #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] { + // x86-64 guarantees SSE2; `_mm_prefetch` needs only SSE. + if #[cfg(target_arch = "x86_64")] { const CACHE_LINE_SIZE: usize = 64; #[inline(always)] unsafe fn prefetch_exactly(ptr: *const i8) { use std::arch::x86_64::*; for i in 0..N { - _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0); + // SAFETY: the caller guarantees that all `N` computed addresses are + // inside the allocation. + unsafe { _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0) }; } } @@ -56,7 +59,8 @@ cfg_if::cfg_if! { if CACHE_LINE_SIZE * i >= bytes { break; } - _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0); + // SAFETY: the loop uses only offsets below `bytes`. + unsafe { _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0) }; } } @@ -66,32 +70,48 @@ cfg_if::cfg_if! { pub fn prefetch_hint_max(vec: &[T]) { let vecsize = std::mem::size_of_val(vec); if vecsize >= MAX_CACHE_LINES * 64 { - // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated. + // SAFETY: the slice contains every address passed to prefetch. unsafe { prefetch_exactly::(vec.as_ptr().cast()) } } else { - // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated. + // SAFETY: the slice covers `vecsize` bytes. unsafe { prefetch_at_most::(vec.as_ptr().cast(), vecsize) } } } + /// Prefetch a raw byte range without creating a slice. + /// + /// # Safety + /// + /// `ptr` must identify an allocation of at least `bytes` bytes. The allocation + /// must remain live for this call. The function creates no Rust reference. + /// The caller controls concurrent mutation of the range. + #[inline] + pub unsafe fn prefetch_hint_all_raw(ptr: *const u8, bytes: usize) { + use std::arch::x86_64::*; + + for offset in (0..bytes).step_by(CACHE_LINE_SIZE) { + // SAFETY: the caller guarantees the byte range, and `offset < bytes`. + unsafe { _mm_prefetch(ptr.add(offset).cast(), _MM_HINT_T0) }; + } + } + /// Prefetch the given vector in chunks of 64 bytes, which is a cache line size. /// The entire vector will be prefetched. #[inline] pub fn prefetch_hint_all(vec: &[T]) { - use std::arch::x86_64::*; - - let vecsize = std::mem::size_of_val(vec); - let num_prefetch_blocks = vecsize.div_ceil(64); - let vec_ptr = vec.as_ptr() as *const i8; - for d in 0..num_prefetch_blocks { - // SAFETY: Pointer is in-bounds and use of the intrinsic is gated by the - // `cfg`-guard on this function. - unsafe { - std::arch::x86_64::_mm_prefetch(vec_ptr.add(d * CACHE_LINE_SIZE), _MM_HINT_T0); - } - } } + // SAFETY: the slice remains live and covers exactly `size_of_val(vec)` bytes. + unsafe { prefetch_hint_all_raw(vec.as_ptr().cast(), std::mem::size_of_val(vec)) } + } } else { pub fn prefetch_hint_max(_vec: &[T]) {} + + /// Accept a raw prefetch range and do nothing. + /// + /// # Safety + /// + /// The pointer contract is the same as the x86-64 implementation. + pub unsafe fn prefetch_hint_all_raw(_ptr: *const u8, _bytes: usize) {} + pub fn prefetch_hint_all(_vec: &[T]) {} } } diff --git a/diskann-wide/src/arch/x86_64/v3/i16x16_.rs b/diskann-wide/src/arch/x86_64/v3/i16x16_.rs index bf8a1a3690..7506f84e35 100644 --- a/diskann-wide/src/arch/x86_64/v3/i16x16_.rs +++ b/diskann-wide/src/arch/x86_64/v3/i16x16_.rs @@ -16,7 +16,7 @@ use crate::{ v3::i16x8, }, }, - bitmask::BitMask, + bitmask::{BitMask, FromInt}, constant::Const, emulated::Emulated, helpers, @@ -89,16 +89,19 @@ impl SIMDMulAdd for i16x16 { impl SIMDPartialEq for i16x16 { #[inline(always)] fn eq_simd(self, other: Self) -> Self::Mask { - self.emulated() - .eq_simd(other.emulated()) - .as_arch(self.arch()) + // SAFETY: V3 includes AVX2 and BMI2. Each equal i16 lane contributes + // two adjacent movemask bits; pext keeps one bit per lane. + let bits = unsafe { + let bytes = _mm256_movemask_epi8(_mm256_cmpeq_epi16(self.0, other.0)) as u32; + _pext_u32(bytes, 0x5555_5555) as u16 + }; + BitMask::from_int(self.arch(), bits) } #[inline(always)] fn ne_simd(self, other: Self) -> Self::Mask { - self.emulated() - .ne_simd(other.emulated()) - .as_arch(self.arch()) + let equal = self.eq_simd(other); + BitMask::from_int(self.arch(), !equal.0) } } diff --git a/diskann-wide/src/doubled.rs b/diskann-wide/src/doubled.rs index d6adcb7b13..1c4cbb3f52 100644 --- a/diskann-wide/src/doubled.rs +++ b/diskann-wide/src/doubled.rs @@ -442,6 +442,13 @@ macro_rules! double_mask { let hi = <$repr>::keep_first(arch, i.saturating_sub({ $N / 2 })); Self(lo, hi) } + + #[inline(always)] + fn first(&self) -> Option { + self.0 + .first() + .or_else(|| self.1.first().map(|index| index + { $N / 2 })) + } } impl From<$crate::doubled::Doubled<$repr>> diff --git a/diskann/Cargo.toml b/diskann/Cargo.toml index d8dcf211b8..589bae1ee5 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -24,6 +24,7 @@ hashbrown = { version = "0.16.0", default-features = false, features = ["default num-traits.workspace = true parking_lot = { version = "0.12.5", optional = true } rand.workspace = true +rand_distr = { workspace = true, optional = true } rayon = { workspace = true, optional = true } thiserror.workspace = true tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } @@ -34,6 +35,9 @@ diskann-wide = { workspace = true } # Optional Dependencies dashmap = { workspace = true, optional = true } +[target.'cfg(target_os = "linux")'.dependencies] +libc = { version = "0.2", optional = true } + [dev-dependencies] futures-util = { workspace = true, default-features = false } pin-project.workspace = true @@ -61,7 +65,14 @@ panic = "warn" default = ["tracing"] # Enable PiPNN batch graph construction. -pipnn = ["dep:diskann-linalg", "dep:parking_lot", "dep:rayon", "tracing"] +pipnn = [ + "dep:diskann-linalg", + "dep:libc", + "dep:parking_lot", + "dep:rand_distr", + "dep:rayon", + "tracing", +] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/adjacencylist.rs b/diskann/src/graph/adjacencylist.rs index 2923621870..85db7962f8 100644 --- a/diskann/src/graph/adjacencylist.rs +++ b/diskann/src/graph/adjacencylist.rs @@ -135,6 +135,19 @@ where } } + /// Take ownership of a vector that contains unique items. + /// + /// The caller must supply unique items. Debug builds check this condition. + #[cfg(feature = "pipnn")] + pub(crate) fn from_vec_trusted(edges: Vec) -> Self + where + I: ContainsSimd, + { + let list = Self { edges }; + list.debug_check_uniqueness(); + list + } + /// Resize the underlying storage to `capacity` elements and return a guard allowing /// full mutable access to the resized span. /// @@ -667,6 +680,17 @@ mod tests { } } + #[cfg(feature = "pipnn")] + #[test] + fn test_from_vec_trusted_preserves_order_and_allocation() { + let edges = vec![3_u32, 1, 2]; + let pointer = edges.as_ptr(); + let list = AdjacencyList::from_vec_trusted(edges); + + assert_eq!(&*list, &[3, 1, 2]); + assert_eq!(list.as_ptr(), pointer); + } + #[test] fn test_from_iter_untrusted() { let x = AdjacencyList::::from_iter_untrusted([]); diff --git a/diskann/src/graph/pipnn/bf16.rs b/diskann/src/graph/pipnn/bf16.rs new file mode 100644 index 0000000000..87fb78271f --- /dev/null +++ b/diskann/src/graph/pipnn/bf16.rs @@ -0,0 +1,105 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Lossy conversion between `f32` and bf16 storage. +//! +//! A bf16 value contains the upper 16 bits of an IEEE-754 `f32`. It keeps the +//! exponent and seven mantissa bits. For non-negative values, its `u16` bit order +//! matches `f32` numeric order. HashPrune applies a separate ordered-key transform +//! to signed distances. +//! +//! Conversion truncates the lower 16 bits. It does not round. It preserves sign, +//! infinity, signed zero, and the upper NaN payload bits. + +/// Convert `f32` → bf16 by truncating the lower 16 mantissa bits. +#[inline(always)] +pub(super) fn f32_to_bf16(v: f32) -> u16 { + (v.to_bits() >> 16) as u16 +} + +/// Reconstruct `f32` from a bf16 for conversion tests. +#[cfg(test)] +#[inline(always)] +fn bf16_to_f32(v: u16) -> f32 { + f32::from_bits((v as u32) << 16) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + fn bf16_roundtrip_preserves_exactly_representable_values( + #[values(0.0_f32, 1.0, 2.0, 0.5, 0.25, 4.0, -1.0, -0.5)] value: f32, + ) { + // Given: bf16 has seven mantissa bits, and these values have zero low mantissa bits. + let expected_value = value; + + // When + let actual_value = bf16_to_f32(f32_to_bf16(value)); + + // Then + assert_eq!(actual_value, expected_value); + } + + #[rstest] + fn bf16_truncation_keeps_relative_error_below_one_percent( + #[values(1e-10_f32, 1e10, std::f32::consts::PI, std::f32::consts::E)] value: f32, + ) { + // Given: bf16 truncation has at most about 2^-7 relative error for finite values. + let maximum_relative_error = 0.01; + + // When + let truncated = bf16_to_f32(f32_to_bf16(value)); + let actual_relative_error = ((truncated - value) / value).abs(); + + // Then + assert!(actual_relative_error <= maximum_relative_error); + } + + #[rstest] + #[case::positive_zero(0.0)] + #[case::negative_zero(-0.0)] + #[case::positive_infinity(f32::INFINITY)] + #[case::negative_infinity(f32::NEG_INFINITY)] + #[case::positive_nan(f32::NAN)] + #[case::negative_nan(f32::from_bits(0xFFC1_2345))] + #[case::minimum_positive(f32::MIN_POSITIVE)] + #[case::negative_minimum_positive(-f32::MIN_POSITIVE)] + fn bf16_conversion_preserves_the_upper_bits_of_special_values(#[case] value: f32) { + // Given + let expected_upper_bits = value.to_bits() & 0xFFFF_0000; + + // When + let actual_bits = bf16_to_f32(f32_to_bf16(value)).to_bits(); + + // Then + assert_eq!(actual_bits, expected_upper_bits); + } + + #[rstest] + fn bf16_truncation_moves_finite_values_toward_zero( + #[values(f32::MIN_POSITIVE, 0.1, 1.1, std::f32::consts::PI, f32::MAX)] value: f32, + ) { + // When + let positive_truncation = bf16_to_f32(f32_to_bf16(value)); + let negative_truncation = bf16_to_f32(f32_to_bf16(-value)); + + // Then + assert!((0.0..=value).contains(&positive_truncation)); + assert!((-value..=0.0).contains(&negative_truncation)); + } + + #[test] + fn bf16_encoding_preserves_non_negative_value_order() { + // For non-negative f32, bf16 (as u16) preserves ordering. + let xs: [f32; 6] = [0.0, 1e-10, 0.1, 1.0, 10.0, 1e10]; + let bs: Vec = xs.iter().map(|&x| f32_to_bf16(x)).collect(); + for w in bs.windows(2) { + assert!(w[0] <= w[1], "bf16 ordering broken: {} > {}", w[0], w[1]); + } + } +} diff --git a/diskann/src/graph/pipnn/hash_prune.rs b/diskann/src/graph/pipnn/hash_prune.rs new file mode 100644 index 0000000000..db4f12a5d6 --- /dev/null +++ b/diskann/src/graph/pipnn/hash_prune.rs @@ -0,0 +1,1553 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Merge leaf edges into bounded per-point reservoirs. +//! +//! For edge `source → target`, relative-hash bit `j` records whether the target's +//! projection on hyperplane `j` is at least the source's projection. The hash +//! groups edges with similar residual directions. +//! +//! A source reservoir keeps at most one neighbor for each relative hash. A closer +//! edge replaces the edge for that direction. A full reservoir accepts only an +//! edge below its farthest total key. +//! +//! Each source owns one lock. The lock protects its reservoir metadata and its +//! rows in the hash, distance, and neighbor arrays. `l_max` sets the logical +//! reservoir length. The `u8` metadata limits this value to 255. + +use parking_lot::lock_api::RawMutex as RawMutexTrait; +use std::cell::UnsafeCell; + +use super::{ + bf16::f32_to_bf16, + lsh::{LshSketches, MAX_PLANES}, + simd::{PiPNNSIMDSchema, PiPNNSIMDVector}, +}; +use crate::{ANNError, ANNResult, graph::AdjacencyList, utils::VectorRepr}; +use bytemuck::Pod; +use diskann_utils::views::MatrixView; +use diskann_vector::{prefetch_hint_all, prefetch_hint_all_raw}; +use diskann_wide::{ + Architecture, SIMDMask, SIMDPartialEq, SIMDVector, + arch::{self, Dispatched1, FTarget1, Target}, + lifetime::As, +}; +#[cfg(not(miri))] +use rayon::prelude::*; + +/// Owned zero-initialized slab from `mmap(MAP_PRIVATE | MAP_ANONYMOUS)`. +#[cfg(target_os = "linux")] +struct MmapSlab { + ptr: *mut T, + len: usize, +} + +#[cfg(target_os = "linux")] +// SAFETY: the slab uniquely owns its mmap region until `drop`. Moving the slab +// transfers that ownership. `T: Send` permits transfer of initialized values. +unsafe impl Send for MmapSlab {} +#[cfg(target_os = "linux")] +// SAFETY: shared access exposes only `*const T`. `T: Sync` permits shared access +// to initialized values. HashPrune uses `UnsafeCell` and a point lock for writes. +unsafe impl Sync for MmapSlab {} + +#[cfg(target_os = "linux")] +impl MmapSlab { + fn new_zeroed(len: usize) -> ANNResult { + if len == 0 { + return Ok(Self { + ptr: std::ptr::NonNull::::dangling().as_ptr(), + len: 0, + }); + } + let bytes = len + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| super::config_error(format!("slab size {len} overflows usize")))?; + // SAFETY: `MAP_ANONYMOUS` returns zero-initialized memory. + // `PROT_READ | PROT_WRITE` permits all accesses used by this slab. + unsafe { + let ptr = libc::mmap( + std::ptr::null_mut(), + bytes, + libc::PROT_READ | libc::PROT_WRITE, + libc::MAP_PRIVATE | libc::MAP_ANONYMOUS, + -1, + 0, + ); + if ptr == libc::MAP_FAILED { + return Err(ANNError::from(std::io::Error::last_os_error()) + .context(format!("mmap failed for {bytes} HashPrune slab bytes"))); + } + Ok(Self { + ptr: ptr as *mut T, + len, + }) + } + } + + #[inline] + fn as_ptr(&self) -> *const T { + self.ptr + } + + #[inline] + fn bytes(&self) -> usize { + self.len * std::mem::size_of::() + } +} + +#[cfg(target_os = "linux")] +impl Drop for MmapSlab { + fn drop(&mut self) { + if self.len > 0 { + // SAFETY: this slab still uniquely owns the mmap base pointer and exact + // byte count established by `new_zeroed`; `self.len > 0` excludes the + // dangling zero-length representation. + unsafe { + libc::munmap(self.ptr as *mut libc::c_void, self.bytes()); + } + } + } +} + +/// Owned zero-initialized slab from `VirtualAlloc`. +#[cfg(windows)] +mod winmem { + pub(super) type Lpvoid = *mut core::ffi::c_void; + pub(super) const MEM_COMMIT: u32 = 0x0000_1000; + pub(super) const MEM_RESERVE: u32 = 0x0000_2000; + pub(super) const MEM_RELEASE: u32 = 0x0000_8000; + pub(super) const PAGE_READWRITE: u32 = 0x04; + + unsafe extern "system" { + pub(super) fn VirtualAlloc( + lpAddress: Lpvoid, + dwSize: usize, + flAllocationType: u32, + flProtect: u32, + ) -> Lpvoid; + pub(super) fn VirtualFree(lpAddress: Lpvoid, dwSize: usize, dwFreeType: u32) -> i32; + } +} + +#[cfg(windows)] +struct MmapSlab { + ptr: *mut T, + len: usize, +} + +#[cfg(windows)] +// SAFETY: the slab uniquely owns its `VirtualAlloc` region until `drop`. Moving +// the slab transfers that ownership. `T: Send` permits transfer of initialized values. +unsafe impl Send for MmapSlab {} +#[cfg(windows)] +// SAFETY: shared access exposes only `*const T`. `T: Sync` permits shared access +// to initialized values. HashPrune uses `UnsafeCell` and a point lock for writes. +unsafe impl Sync for MmapSlab {} + +#[cfg(windows)] +impl MmapSlab { + fn new_zeroed(len: usize) -> ANNResult { + if len == 0 { + return Ok(Self { + ptr: std::ptr::NonNull::::dangling().as_ptr(), + len: 0, + }); + } + let bytes = len + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| super::config_error(format!("slab size {len} overflows usize")))?; + // SAFETY: `MEM_RESERVE | MEM_COMMIT` returns zero-initialized memory. + // `PAGE_READWRITE` permits all accesses used by this slab. + unsafe { + let ptr = winmem::VirtualAlloc( + std::ptr::null_mut(), + bytes, + winmem::MEM_RESERVE | winmem::MEM_COMMIT, + winmem::PAGE_READWRITE, + ); + if ptr.is_null() { + return Err( + ANNError::from(std::io::Error::last_os_error()).context(format!( + "VirtualAlloc failed for {bytes} HashPrune slab bytes" + )), + ); + } + Ok(Self { + ptr: ptr as *mut T, + len, + }) + } + } + + #[inline] + fn as_ptr(&self) -> *const T { + self.ptr + } +} + +#[cfg(windows)] +impl Drop for MmapSlab { + fn drop(&mut self) { + if self.len > 0 { + // SAFETY: this slab still uniquely owns the VirtualAlloc base pointer; + // MEM_RELEASE requires and receives `dwSize = 0`. + unsafe { + winmem::VirtualFree(self.ptr as winmem::Lpvoid, 0, winmem::MEM_RELEASE); + } + } + } +} + +/// Owned zero-initialized slab for platforms without `mmap` or `VirtualAlloc`. +#[cfg(not(any(target_os = "linux", windows)))] +struct MmapSlab(Vec); + +#[cfg(not(any(target_os = "linux", windows)))] +impl MmapSlab { + fn new_zeroed(len: usize) -> ANNResult { + Ok(Self(vec![T::default(); len])) + } + #[inline] + fn as_ptr(&self) -> *const T { + self.0.as_ptr() + } +} + +/// Largest reservoir length that fits in `ReservoirState`. +/// +/// `ReservoirState.len` and `ReservoirState.farthest_idx` are `u8`. Runtime `l_max` selects the +/// actual length. Values above this bound are invalid. +pub(crate) const MAX_RESERVOIR_LEN: usize = u8::MAX as usize; + +#[repr(C)] +struct ReservoirState { + len: u8, + farthest_idx: u8, + farthest_dist: u16, + _pad: [u8; 10], +} + +#[repr(C, align(16))] +struct LockedReservoirState { + lock: parking_lot::RawMutex, + state: UnsafeCell, +} + +impl LockedReservoirState { + fn new() -> Self { + Self { + lock: ::INIT, + state: UnsafeCell::new(ReservoirState::new_empty()), + } + } + + fn state_ptr(&self) -> *mut ReservoirState { + self.state.get() + } + + fn with_locked_state(&self, f: impl FnOnce(&mut ReservoirState) -> R) -> R { + struct UnlockOnDrop<'a>(&'a parking_lot::RawMutex); + impl Drop for UnlockOnDrop<'_> { + fn drop(&mut self) { + // SAFETY: the guard is created only after acquiring this mutex. + unsafe { self.0.unlock() }; + } + } + + self.lock.lock(); + let _guard = UnlockOnDrop(&self.lock); + // SAFETY: the mutex is separate from `state`, so contending threads may + // access the lock while this exclusive state reference is live. + f(unsafe { &mut *self.state.get() }) + } +} + +// SAFETY: `lock` guards every mutable access to `state`. Read-only extraction +// happens only after HashPrune is consumed, when no mutation can remain. +unsafe impl Sync for LockedReservoirState {} + +impl ReservoirState { + const fn new_empty() -> Self { + Self { + len: 0, + farthest_idx: 0, + farthest_dist: 0, + _pad: [0; 10], + } + } +} + +const _: [(); 16] = [(); std::mem::size_of::()]; + +// These pointers name one source reservoir's hash, distance, and neighbor rows. +// The caller must hold that source lock before it writes through a pointer. + +#[derive(Clone, Copy)] +struct ReservoirRows { + hashes: *mut u16, + distances: *mut u16, + neighbors: *mut u32, + row_stride: usize, +} + +#[derive(Clone, Copy)] +struct FindHashArgs { + hashes: *const u16, + row_stride: usize, + len: u8, + target: u16, +} + +#[derive(Clone, Copy)] +struct RelativeHashArgs { + src: *const f32, + dst: *const f32, + len: usize, +} + +type FindHash = Dispatched1, As>; +type RelativeHash = Dispatched1>; + +struct FindHashKernel; +struct RelativeHashKernel; +struct SelectFindHash; +struct SelectRelativeHash; + +impl Target for SelectFindHash +where + A: Architecture, + FindHashKernel: FTarget1, FindHashArgs>, +{ + fn run(self, arch: A) -> FindHash { + arch.dispatch1::, As>() + } +} + +impl Target for SelectRelativeHash +where + A: Architecture, + RelativeHashKernel: FTarget1, +{ + fn run(self, arch: A) -> RelativeHash { + arch.dispatch1::>() + } +} + +impl FTarget1, FindHashArgs> for FindHashKernel +where + A: Architecture, + A::i16x32: SIMDPartialEq, +{ + fn run(arch: A, args: FindHashArgs) -> Option { + find_hash_simd::(arch, args) + } +} + +impl FTarget1 for RelativeHashKernel +where + A: PiPNNSIMDSchema, +{ + fn run(arch: A, args: RelativeHashArgs) -> u16 { + relative_hash_simd::(arch, args) + } +} + +/// Find an existing relative-direction bucket in one source reservoir. +/// +/// The SIMD backend has no `u16` vector. An `i16` load keeps each hash bit +/// pattern, so equality gives the same result. +fn find_hash_simd(arch: F::Arch, args: FindHashArgs) -> Option +where + F: SIMDVector + SIMDPartialEq, +{ + let len = args.len as usize; + let target = F::splat(arch, args.target as i16); + let chunks = len.div_ceil(F::LANES).min(args.row_stride / F::LANES); + for chunk in 0..chunks { + // SAFETY: `insert_reservoir_edge` supplies a hash row with `row_stride` elements. + // `chunks <= row_stride / F::LANES`, so this full load stays in the row. + let values = unsafe { F::load_simd(arch, args.hashes.add(chunk * F::LANES).cast::()) }; + if let Some(offset) = values.eq_simd(target).first() { + let lane = chunk * F::LANES + offset; + if lane < len { + return Some(lane); + } + } + } + None +} + +/// Return the relative hash for two sketches. +/// +/// Bit `j` is one when `dst[j] - src[j] >= 0.0`. Equality and signed zero set +/// the bit on every architecture. +fn relative_hash_simd(arch: F::Arch, args: RelativeHashArgs) -> u16 +where + F: PiPNNSIMDVector, +{ + if F::LANES >= MAX_PLANES { + // SAFETY: `src` and `dst` each contain `len <= MAX_PLANES <= F::LANES` + // values. The masked loads do not read inactive lanes. + let dst = unsafe { F::load_simd_first(arch, args.dst, args.len) }; + // SAFETY: `src` has the same checked length as `dst`. + let src = unsafe { F::load_simd_first(arch, args.src, args.len) }; + let active = (1_u64 << args.len) - 1; + return F::active_lanes((dst - src).ge_simd(F::splat(arch, 0.0))) as u16 & active as u16; + } + + let mut bits = 0u16; + let mut offset = 0usize; + while offset < args.len { + let chunk_len = (args.len - offset).min(F::LANES); + // SAFETY: `offset + chunk_len <= args.len`. The masked loads do not read + // inactive lanes. + let dst = unsafe { F::load_simd_first(arch, args.dst.add(offset), chunk_len) }; + // SAFETY: `src` has the same checked length as `dst`. + let src = unsafe { F::load_simd_first(arch, args.src.add(offset), chunk_len) }; + let active = (1_u64 << chunk_len) - 1; + let chunk_bits = F::active_lanes((dst - src).ge_simd(F::splat(arch, 0.0))) & active; + bits |= (chunk_bits as u16) << offset; + offset += chunk_len; + } + bits +} + +/// Convert a bf16 distance to an order-preserving `u16` key. +/// +/// Raw bf16 bits are monotonic only for non-negative values. Inner-product +/// distance can be negative. This transform preserves the total numeric order +/// for both signs. +#[inline(always)] +fn ordered_distance_key(distance: f32) -> u16 { + let b = f32_to_bf16(distance); + if b & 0x8000 != 0 { !b } else { b | 0x8000 } +} + +/// Update the cached farthest entry for one reservoir. +/// +/// # Safety +/// +/// The caller holds the source lock. `state.len <= rows.row_stride`. The first +/// `state.len` entries of all three reservoir rows are initialized. +#[inline] +unsafe fn update_farthest(state: &mut ReservoirState, rows: ReservoirRows) { + if state.len == 0 { + state.farthest_dist = 0; + state.farthest_idx = 0; + return; + } + // The total key is `(distance, residual hash, neighbor ID)`. The residual + // hash resolves equal bf16 distances. The ID resolves the remaining ties. + let mut max_idx: u8 = 0; + // SAFETY: `state.len > 0` and all active slots are initialized. + let mut max_key = unsafe { (*rows.distances, *rows.hashes, *rows.neighbors) }; + for i in 1..state.len as usize { + // SAFETY: `i < state.len <= rows.row_stride`, and all active entries are + // initialized. + let key = unsafe { + ( + *rows.distances.add(i), + *rows.hashes.add(i), + *rows.neighbors.add(i), + ) + }; + if key > max_key { + max_key = key; + max_idx = i as u8; + } + } + state.farthest_dist = max_key.0; + state.farthest_idx = max_idx; +} + +/// Insert one edge into a locked reservoir. +/// +/// The function replaces a matching hash only when the new edge has a smaller +/// total key. A full reservoir accepts only a key below its farthest key. +/// +/// # Safety +/// +/// The caller holds the source lock. Each pointer in `rows` is valid for +/// `row_stride` elements. `state.len <= l_max <= row_stride`. The first +/// `state.len` entries of all three rows are initialized. +#[inline(always)] +unsafe fn insert_reservoir_edge( + state: &mut ReservoirState, + rows: ReservoirRows, + hash: u16, + neighbor: u32, + distance: f32, + l_max: u8, + find_hash: FindHash, +) -> bool { + let dist_key = ordered_distance_key(distance); + + if state.len >= l_max { + let farthest = state.farthest_idx as usize; + // SAFETY: a full reservoir has `farthest < state.len` initialized slots. + let farthest_key = unsafe { + ( + state.farthest_dist, + *rows.hashes.add(farthest), + *rows.neighbors.add(farthest), + ) + }; + if (dist_key, hash, neighbor) >= farthest_key { + return false; + } + } + + if let Some(idx) = find_hash.call(FindHashArgs { + hashes: rows.hashes, + row_stride: rows.row_stride, + len: state.len, + target: hash, + }) { + // SAFETY: `idx < state.len <= rows.row_stride`. + let current_key = unsafe { (*rows.distances.add(idx), *rows.neighbors.add(idx)) }; + if (dist_key, neighbor) < current_key { + let was_farthest = idx == state.farthest_idx as usize; + // SAFETY: `idx < state.len <= rows.row_stride`. The entry is + // initialized, and the caller holds the source lock. + unsafe { + *rows.neighbors.add(idx) = neighbor; + *rows.distances.add(idx) = dist_key; + } + if was_farthest { + // SAFETY: the caller still holds the source lock. The reservoir rows + // and initialized prefix are unchanged. + unsafe { update_farthest(state, rows) }; + } + return true; + } + return false; + } + + if state.len < l_max { + let new_idx = state.len as usize; + let becomes_farthest = if state.len == 0 { + true + } else { + let farthest = state.farthest_idx as usize; + // SAFETY: `farthest < state.len` identifies an initialized slot. + let farthest_key = unsafe { + ( + state.farthest_dist, + *rows.hashes.add(farthest), + *rows.neighbors.add(farthest), + ) + }; + (dist_key, hash, neighbor) > farthest_key + }; + // SAFETY: `new_idx < l_max <= rows.row_stride`; the caller holds the lock. + unsafe { + *rows.hashes.add(new_idx) = hash; + *rows.distances.add(new_idx) = dist_key; + *rows.neighbors.add(new_idx) = neighbor; + } + state.len += 1; + if becomes_farthest { + state.farthest_dist = dist_key; + state.farthest_idx = new_idx as u8; + } + return true; + } + + // The full-reservoir early rejection above proved that the incoming + // `(distance, residual hash, ID)` key is better than the cached farthest key. + let idx = state.farthest_idx as usize; + // SAFETY: `idx < state.len <= rows.row_stride`; the caller holds the lock. + unsafe { + *rows.hashes.add(idx) = hash; + *rows.distances.add(idx) = dist_key; + *rows.neighbors.add(idx) = neighbor; + update_farthest(state, rows); + } + true +} + +/// Return at most `cap` neighbor IDs in distance order. +/// +/// `scratch` belongs to one Rayon extraction job and is reused for its rows. +/// +/// # Safety +/// +/// The caller must exclude mutation with the source lock or unique ownership. +/// `distances` and `neighbors` each point to `state.len` initialized entries. +unsafe fn collect_nearest_ids( + state: &ReservoirState, + distances: *const u16, + neighbors: *const u32, + cap: usize, + scratch: &mut Vec<(u32, u16)>, +) -> Vec { + let n = state.len as usize; + scratch.clear(); + scratch.reserve(n); + for i in 0..n { + // SAFETY: `i < n == state.len`, and both arrays have an initialized entry + // at `i`. + scratch.push(unsafe { (*neighbors.add(i), *distances.add(i)) }); + } + scratch.sort_unstable_by_key(|&(id, distance)| (distance, id)); + scratch[..n.min(cap)].iter().map(|&(id, _)| id).collect() +} + +/// Return at most `cap` neighbor IDs without sorting them. +/// +/// The caller does not depend on reservoir order. This function reads only the +/// neighbor row. +/// +/// # Safety +/// +/// The caller must exclude mutation with the source lock or unique ownership. +/// `neighbors` points to `state.len` initialized entries. +#[inline] +unsafe fn collect_neighbor_ids( + state: &ReservoirState, + neighbors: *const u32, + cap: usize, +) -> Vec { + let out_len = (state.len as usize).min(cap); + let mut out = Vec::with_capacity(out_len); + for i in 0..out_len { + // SAFETY: `i < out_len <= state.len`, and the neighbor entry is initialized. + out.push(unsafe { *neighbors.add(i) }); + } + out +} + +/// Bounded point reservoirs shared by parallel leaf workers. +/// +/// Source point `i` owns `states[i]` and row `i` in each reservoir array. +/// Its lock protects the metadata and all three rows. A worker holds at most one +/// source lock. Extraction consumes `HashPrune`, so no writer can remain. +pub(crate) struct HashPrune { + states: Vec, + hash_rows: UnsafeCell>, + distance_rows: UnsafeCell>, + neighbor_rows: UnsafeCell>, + row_stride: usize, + sketches: LshSketches, + l_max: usize, + find_hash: FindHash, + relative_hash: RelativeHash, +} + +// SAFETY: each mutable reservoir row is inside `UnsafeCell` and guarded by the +// matching source lock. Different source locks protect disjoint rows. Consuming +// extraction proves that no writer remains. +unsafe impl Send for HashPrune {} +// SAFETY: the same per-point lock protects mutation through shared HashPrune +// references; immutable sketches are safe to share. +unsafe impl Sync for HashPrune {} + +impl HashPrune { + /// Create one empty direction reservoir and LSH sketch for each dataset point. + pub(crate) fn new( + data: MatrixView<'_, T>, + num_planes: usize, + l_max: usize, + seed: u64, + ) -> ANNResult { + let npoints = data.nrows(); + let t0 = std::time::Instant::now(); + let sketches = LshSketches::try_new(data, num_planes, seed)?; + tracing::debug!( + elapsed_secs = t0.elapsed().as_secs_f64(), + "sketch computation" + ); + let t1 = std::time::Instant::now(); + let row_stride = l_max.next_multiple_of(32).max(32); + + let states: Vec<_> = (0..npoints).map(|_| LockedReservoirState::new()).collect(); + + // Each reservoir array has one `row_stride` row for each source point. + let total = npoints.checked_mul(row_stride).ok_or_else(|| { + super::config_error(format!( + "HashPrune slab shape {npoints} x {row_stride} overflows usize" + )) + })?; + let hash_rows = MmapSlab::::new_zeroed(total)?; + let distance_rows = MmapSlab::::new_zeroed(total)?; + let neighbor_rows = MmapSlab::::new_zeroed(total)?; + + #[cfg(target_os = "linux")] + { + let state_bytes = states.len() * std::mem::size_of::(); + // SAFETY: each pointer names a contiguous allocation of `bytes`. + // `madvise` does not read or write the allocation. + unsafe { + for (ptr, bytes) in [ + (states.as_ptr() as *mut libc::c_void, state_bytes), + (hash_rows.as_ptr() as *mut libc::c_void, hash_rows.bytes()), + ( + distance_rows.as_ptr() as *mut libc::c_void, + distance_rows.bytes(), + ), + ( + neighbor_rows.as_ptr() as *mut libc::c_void, + neighbor_rows.bytes(), + ), + ] { + if bytes > 2 * 1024 * 1024 { + libc::madvise(ptr, bytes, libc::MADV_HUGEPAGE); + } + } + } + } + + tracing::debug!( + elapsed_secs = t1.elapsed().as_secs_f64(), + row_stride, + "reservoir allocation" + ); + + Ok(Self { + states, + hash_rows: UnsafeCell::new(hash_rows), + distance_rows: UnsafeCell::new(distance_rows), + neighbor_rows: UnsafeCell::new(neighbor_rows), + row_stride, + sketches, + l_max, + find_hash: arch::dispatch(SelectFindHash), + relative_hash: arch::dispatch(SelectRelativeHash), + }) + } + + /// Lock one source point's reservoir while `f` reads or updates it. + /// + /// RAII unlocks the point when the closure exits. + #[inline(always)] + fn with_locked_reservoir( + &self, + idx: usize, + f: impl FnOnce(&mut ReservoirState, ReservoirRows) -> R, + ) -> R { + let slot = &self.states[idx]; + let off = idx * self.row_stride; + // SAFETY: `idx` is in bounds. Each array has + // `states.len() * row_stride` elements. `UnsafeCell` permits these writes, + // and `with_locked_state` holds the source lock for the closure. + let rows = unsafe { + ReservoirRows { + hashes: (*self.hash_rows.get()).as_ptr().cast_mut().add(off), + distances: (*self.distance_rows.get()).as_ptr().cast_mut().add(off), + neighbors: (*self.neighbor_rows.get()).as_ptr().cast_mut().add(off), + row_stride: self.row_stride, + } + }; + slot.with_locked_state(|state| f(state, rows)) + } + + /// Add one leaf's weighted CSR edges to the point reservoirs. + /// + /// `point_ids` maps leaf-local positions to dataset IDs. `edge_offsets` and + /// `edges` form a CSR matrix with leaf-local targets. `sketch_scratch` stores + /// the gathered sketches for this leaf. + /// + /// The leaf builder supplies sorted unique dataset IDs and one monotonic + /// offset range per point. Each edge target is a valid leaf-local position. + /// Construction checks the full sketch and reservoir allocation shapes. + pub(crate) fn add_leaf_edges( + &self, + point_ids: &[u32], + edge_offsets: &[u32], + edges: &[(u32, f32)], + sketch_scratch: &mut Vec, + ) { + if edges.is_empty() { + return; + } + + let n = point_ids.len(); + let m = self.sketches.num_planes(); + let l_max = self.l_max as u8; + let sketch_len = n * m; + if sketch_scratch.len() < sketch_len { + sketch_scratch.resize(sketch_len, 0.0); + } + self.gather_sketches(point_ids, &mut sketch_scratch[..sketch_len]); + + for local_src in 0..n { + let start = edge_offsets[local_src] as usize; + let end = edge_offsets[local_src + 1] as usize; + if start == end { + continue; + } + let global_src = point_ids[local_src] as usize; + + if let Some(next) = (local_src + 1..n) + .find(|&i| edge_offsets[i] != edge_offsets[i + 1]) + .map(|i| point_ids[i] as usize) + { + let off = next * self.row_stride; + prefetch_hint_all(std::slice::from_ref(&self.states[next])); + // SAFETY: `next` is a dataset point ID. This range is the + // complete padded hash row for that point. Raw prefetch avoids + // a shared slice while another worker mutates the row. + unsafe { + let hashes = (*self.hash_rows.get()).as_ptr().add(off); + prefetch_hint_all_raw( + hashes.cast(), + self.row_stride * std::mem::size_of::(), + ); + } + } + + let src_sketch = &sketch_scratch[local_src * m..(local_src + 1) * m]; + self.with_locked_reservoir(global_src, |state, rows| { + for &(dst_local, dist) in &edges[start..end] { + let dst_index = dst_local as usize; + let global_dst = point_ids[dst_index]; + let dst_sketch = &sketch_scratch[dst_index * m..(dst_index + 1) * m]; + let hash = self.relative_hash.call(RelativeHashArgs { + src: src_sketch.as_ptr(), + dst: dst_sketch.as_ptr(), + len: m, + }); + // SAFETY: The source lock gives exclusive access to these rows. + // Each row has at least `l_max` entries. Insertion maintains + // `state.len <= l_max` and reads only initialized entries. + unsafe { + insert_reservoir_edge( + state, + rows, + hash, + global_dst, + dist, + l_max, + self.find_hash, + ) + }; + } + }); + } + } + + fn gather_sketches(&self, indices: &[u32], out: &mut [f32]) { + let dimensions = self.sketches.num_planes(); + let sketches = self.sketches.sketches(); + for (destination, &point) in out.chunks_exact_mut(dimensions).zip(indices) { + let first = point as usize * dimensions; + destination.copy_from_slice(&sketches[first..first + dimensions]); + } + } + + /// Consume the reservoirs and return at most `max_degree` nearest IDs per point. + #[allow(clippy::disallowed_methods)] // build_graph installs the caller-owned pool. + pub(crate) fn into_nearest_lists(self, max_degree: usize) -> Vec> { + let row_stride = self.row_stride; + drop(self.sketches); + let HashPrune { + states, + hash_rows, + distance_rows, + neighbor_rows, + .. + } = self; + let hash_rows = hash_rows.into_inner(); + let distance_rows = distance_rows.into_inner(); + let neighbor_rows = neighbor_rows.into_inner(); + drop(hash_rows); + let extract = |scratch: &mut Vec<(u32, u16)>, i: usize| { + let off = i * row_stride; + // SAFETY: indexing proves that `i` names a live slot. This method + // consumes `self`, so no writer can overlap this state reference. + let state = unsafe { &*states[i].state_ptr() }; + // SAFETY: construction allocated `npoints * row_stride` entries; + // this loop keeps `i < npoints`, and insertion maintains + // `state.len <= l_max <= row_stride` initialized entries. + let ids = unsafe { + collect_nearest_ids( + state, + distance_rows.as_ptr().wrapping_add(off), + neighbor_rows.as_ptr().wrapping_add(off), + max_degree, + scratch, + ) + }; + // A neighbor always has the same relative hash for this source; + // insertion replaces an existing hash slot instead of appending. + AdjacencyList::from_vec_trusted(ids) + }; + + #[cfg(miri)] + { + let mut scratch = Vec::new(); + (0..states.len()) + .map(|i| extract(&mut scratch, i)) + .collect() + } + #[cfg(not(miri))] + { + (0..states.len()) + .into_par_iter() + .map_init(Vec::new, extract) + .collect() + } + } + + /// Consume the reservoirs and return all retained IDs without sorting them. + #[allow(clippy::disallowed_methods)] // build_graph installs the caller-owned pool. + pub(crate) fn into_candidate_lists(self) -> Vec> { + let cap = self.l_max; + let row_stride = self.row_stride; + drop(self.sketches); + let HashPrune { + states, + hash_rows, + distance_rows, + neighbor_rows, + .. + } = self; + let hash_rows = hash_rows.into_inner(); + let distance_rows = distance_rows.into_inner(); + let neighbor_rows = neighbor_rows.into_inner(); + // Extraction reads only neighbor IDs. Drop the hash and distance arrays + // before the code creates the output lists. + drop(hash_rows); + drop(distance_rows); + let extract = |i: usize| { + let neighbors = neighbor_rows.as_ptr().wrapping_add(i * row_stride); + // SAFETY: indexing proves that `i` names a live slot. This method + // consumes `self`, so no writer can overlap this state reference. + let state = unsafe { &*states[i].state_ptr() }; + // SAFETY: construction allocated `npoints * row_stride` entries; + // this loop keeps `i < npoints`, and insertion maintains + // `state.len <= l_max <= row_stride` initialized entries. + let ids = unsafe { collect_neighbor_ids(state, neighbors, cap) }; + // Reservoir slots have unique hashes, and one neighbor cannot + // produce two hashes for the same source. + AdjacencyList::from_vec_trusted(ids) + }; + + #[cfg(miri)] + { + (0..states.len()).map(extract).collect() + } + #[cfg(not(miri))] + { + (0..states.len()).into_par_iter().map(extract).collect() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn build_hash_prune( + data: &[T], + points: usize, + dimensions: usize, + planes: usize, + l_max: usize, + ) -> ANNResult { + HashPrune::new( + MatrixView::try_from(data, points, dimensions).unwrap(), + planes, + l_max, + 42, + ) + } + + struct TestReservoir { + state: ReservoirState, + hashes: Vec, + distances: Vec, + neighbors: Vec, + row_stride: usize, + l_max: u8, + } + + impl TestReservoir { + fn new(l_max: usize) -> Self { + assert!(l_max <= MAX_RESERVOIR_LEN); + let row_stride = l_max.next_multiple_of(32).max(32); + Self { + state: ReservoirState::new_empty(), + hashes: vec![0; row_stride], + distances: vec![0; row_stride], + neighbors: vec![0; row_stride], + row_stride, + l_max: l_max as u8, + } + } + + fn rows(&self) -> ReservoirRows { + ReservoirRows { + hashes: self.hashes.as_ptr() as *mut u16, + distances: self.distances.as_ptr() as *mut u16, + neighbors: self.neighbors.as_ptr() as *mut u32, + row_stride: self.row_stride, + } + } + + fn insert(&mut self, hash: u16, neighbor: u32, distance: f32) -> bool { + let rows = self.rows(); + // SAFETY: the test owns the reservoir and holds its only mutable reference. + unsafe { + insert_reservoir_edge( + &mut self.state, + rows, + hash, + neighbor, + distance, + self.l_max, + arch::dispatch(SelectFindHash), + ) + } + } + + fn neighbors(&self) -> Vec<(u32, f32)> { + let mut entries: Vec<_> = self + .neighbors + .iter() + .copied() + .zip(self.distances.iter().copied()) + .take(self.len()) + .collect(); + entries.sort_unstable_by_key(|&(id, distance)| (distance, id)); + entries + .into_iter() + .map(|(id, key)| { + let bits = if key & 0x8000 != 0 { + key & 0x7fff + } else { + !key + }; + (id, f32::from_bits((bits as u32) << 16)) + }) + .collect() + } + + fn len(&self) -> usize { + self.state.len as usize + } + + fn is_empty(&self) -> bool { + self.state.len == 0 + } + } + + #[test] + fn reservoir_lock_serializes_state_mutation() { + let slot = LockedReservoirState::new(); + let start = std::sync::Barrier::new(3); + + std::thread::scope(|scope| { + for _ in 0..2 { + let slot = &slot; + let start = &start; + scope.spawn(move || { + start.wait(); + for _ in 0..16 { + slot.with_locked_state(|state| state.farthest_dist += 1); + } + }); + } + start.wait(); + }); + + assert_eq!(slot.with_locked_state(|state| state.farthest_dist), 32); + } + + fn add_edge(hp: &HashPrune, src: usize, dst: usize, distance: f32) { + let m = hp.sketches.num_planes(); + let sketches = hp.sketches.sketches(); + let hash = hp.relative_hash.call(RelativeHashArgs { + src: sketches[src * m..(src + 1) * m].as_ptr(), + dst: sketches[dst * m..(dst + 1) * m].as_ptr(), + len: m, + }); + let l_max = hp.l_max as u8; + hp.with_locked_reservoir(src, |state, rows| { + // SAFETY: `with_locked_reservoir` holds the source lock and supplies valid reservoir rows. + unsafe { + insert_reservoir_edge(state, rows, hash, dst as u32, distance, l_max, hp.find_hash) + }; + }); + } + + fn assert_sketch_source_type_matches_f32( + label: &str, + convert: impl Fn(u8) -> T, + reference: impl Fn(u8) -> f32, + ) where + T: VectorRepr + Send + Sync, + { + #[cfg(not(miri))] + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let points = 5; + // Miri keeps conversion and boundary coverage but omits duplicate sizes. + let dimensions: &[usize] = if cfg!(miri) { + &[1, 15, 16, 17, 33] + } else { + &[1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] + }; + let plane_counts: &[usize] = if cfg!(miri) { &[1, 16] } else { &[1, 8, 16] }; + for &dimensions in dimensions { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let point = index / dimensions; + let dimension = index % dimensions; + (point + dimension) as u8 + }) + .collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + let f32_data: Vec = raw.iter().copied().map(&reference).collect(); + for &planes in plane_counts { + let build = || { + ( + LshSketches::try_new( + MatrixView::try_from(converted.as_slice(), points, dimensions).unwrap(), + planes, + 42, + ) + .unwrap(), + LshSketches::try_new( + MatrixView::try_from(f32_data.as_slice(), points, dimensions).unwrap(), + planes, + 42, + ) + .unwrap(), + ) + }; + #[cfg(miri)] + let (actual, expected) = build(); + #[cfg(not(miri))] + let (actual, expected) = pool.install(build); + assert_eq!( + actual.sketches(), + expected.sketches(), + "{label} dimensions={dimensions} planes={planes}" + ); + } + } + } + + // Source conversion. + + #[test] + fn f16_sketch_conversion_matches_f32_across_dimensions_and_planes() { + assert_sketch_source_type_matches_f32( + "f16", + |value| half::f16::from_f32_const(value as f32), + |value| value as f32, + ); + } + + #[test] + fn u8_sketch_conversion_matches_f32_across_dimensions_and_planes() { + assert_sketch_source_type_matches_f32("u8", |value| value, |value| value as f32); + } + + #[test] + fn i8_sketch_conversion_matches_f32_across_dimensions_and_planes() { + assert_sketch_source_type_matches_f32( + "i8", + |value| value as i8 - 11, + |value| (value as i8 - 11) as f32, + ); + } + + // Dispatched hash primitives. + + #[test] + fn relative_hash_matches_numeric_reference() { + let dispatched = arch::dispatch(SelectRelativeHash); + + let src = [ + 1.0, -2.0, 0.0, 7.5, -0.0, 3.25, -9.0, 4.0, 8.0, -1.5, 2.0, 0.0, 6.0, -3.0, 5.5, -7.25, + ]; + let dst = [ + 1.0, -3.0, 0.5, 7.0, 0.0, 3.25, -8.0, -4.0, 9.0, -1.5, -2.0, -0.0, 5.0, -2.0, 5.5, -8.0, + ]; + + for m in 0..=16 { + let mut expected_numeric_hash = 0u16; + for j in 0..m { + let diff: f32 = dst[j] - src[j]; + expected_numeric_hash |= ((diff >= 0.0) as u16) << j; + } + + let actual_dispatched_hash = dispatched.call(RelativeHashArgs { + src: src.as_ptr(), + dst: dst.as_ptr(), + len: m, + }); + assert_eq!(actual_dispatched_hash, expected_numeric_hash, "m={m}"); + } + } + + #[test] + fn relative_hash_defines_signed_zero_and_nan_buckets() { + let src = [0.0; 4]; + let dst = [ + 0.0, + -0.0, + f32::from_bits(0x7FC0_0000), + f32::from_bits(0xFFC0_0000), + ]; + + assert_eq!( + arch::dispatch(SelectRelativeHash).call(RelativeHashArgs { + src: src.as_ptr(), + dst: dst.as_ptr(), + len: dst.len(), + }), + 0b0011 + ); + } + + #[test] + fn find_hash_ignores_padding_and_returns_the_matching_active_index() { + let dispatched = arch::dispatch(SelectFindHash); + + for target in [0, 0xF00D] { + for len in [0usize, 1, 15, 16, 17, 31, 32, 33, 63, 64, 65, 254, 255] { + let row_stride = len.max(1).next_multiple_of(32); + let mut hashes = vec![target; row_stride]; + hashes[..len].fill(0x8001); + let args = |hashes: &[u16]| FindHashArgs { + hashes: hashes.as_ptr(), + row_stride, + len: len as u8, + target, + }; + + assert_eq!(dispatched.call(args(&hashes)), None, "len={len}"); + for index in [0, len / 2, len.saturating_sub(1)] { + if index < len { + hashes[index] = target; + assert_eq!(dispatched.call(args(&hashes)), Some(index), "len={len}"); + hashes[index] = 0x8001; + } + } + } + } + } + + // Storage and configuration. + + #[test] + fn slab_is_zeroed() { + let slab = MmapSlab::::new_zeroed(4).unwrap(); + assert!(!slab.as_ptr().is_null()); + // SAFETY: this test uniquely owns a live four-element slab. + let values = unsafe { std::slice::from_raw_parts(slab.as_ptr(), 4) }; + assert_eq!(values, &[0; 4]); + } + + #[test] + fn ordered_distance_key_preserves_bf16_order_for_all_signs() { + let values = [ + f32::NEG_INFINITY, + -100.0, + -0.0, + 0.0, + 0.25, + 100.0, + f32::INFINITY, + ]; + let keys: Vec<_> = values.iter().copied().map(ordered_distance_key).collect(); + assert!(keys.windows(2).all(|pair| pair[0] <= pair[1])); + } + + // Leaf ingestion and scratch reuse. + + #[test] + fn batched_leaf_edges_match_single_edge_reference() { + let data = [0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + let batched = build_hash_prune(&data, 4, 2, 8, 8).unwrap(); + let reference = build_hash_prune(&data, 4, 2, 8, 8).unwrap(); + let point_ids = [0, 1, 2, 3]; + let offsets = [0, 3, 6, 9, 12]; + let edges = [ + (1, 1.0), + (2, 1.0), + (3, 2.0), + (0, 1.0), + (2, 2.0), + (3, 1.0), + (0, 1.0), + (1, 2.0), + (3, 1.0), + (0, 2.0), + (1, 1.0), + (2, 1.0), + ]; + let mut scratch = Vec::new(); + + batched.add_leaf_edges(&point_ids, &offsets, &edges, &mut scratch); + for source in 0..point_ids.len() { + for &(target, distance) in + &edges[offsets[source] as usize..offsets[source + 1] as usize] + { + add_edge(&reference, source, target as usize, distance); + } + } + let canonicalize = |lists: Vec>| { + lists + .into_iter() + .map(|candidates| { + let mut ids = candidates.to_vec(); + ids.sort_unstable(); + ids + }) + .collect::>() + }; + let actual_batched_candidates = canonicalize(batched.into_candidate_lists()); + let expected_single_edge_candidates = canonicalize(reference.into_candidate_lists()); + + assert_eq!(actual_batched_candidates, expected_single_edge_candidates); + assert!( + actual_batched_candidates + .iter() + .all(|candidates| !candidates.is_empty()) + ); + } + + #[test] + fn reused_sketch_scratch_does_not_leak_candidates_between_leaves() { + // Given + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let hp = build_hash_prune(&data, 4, 1, 8, 4).unwrap(); + let expected_candidates = [vec![1], vec![0], vec![3], vec![2]]; + let mut sketch_scratch_with_stale_value = vec![99.0; 1]; + + // When + hp.add_leaf_edges( + &[0, 1], + &[0, 1, 2], + &[(1, 1.0), (0, 1.0)], + &mut sketch_scratch_with_stale_value, + ); + hp.add_leaf_edges( + &[2, 3], + &[0, 1, 2], + &[(1, 1.0), (0, 1.0)], + &mut sketch_scratch_with_stale_value, + ); + hp.add_leaf_edges( + &[0, 1], + &[0, 0, 0], + &[], + &mut sketch_scratch_with_stale_value, + ); + let actual_candidates: Vec<_> = hp + .into_candidate_lists() + .into_iter() + .map(Vec::from) + .collect(); + + // Then + assert_eq!(actual_candidates, expected_candidates); + } + + // TestReservoir replacement and ordering policy. + + #[test] + fn full_reservoir_evicts_the_farthest_candidate() { + // Given + let expected_capacity = 3; + let expected_neighbors = [(4, 0.5), (1, 1.0), (2, 2.0)]; + let mut reservoir = TestReservoir::new(expected_capacity); + assert!(reservoir.is_empty()); + reservoir.insert(0, 1, 1.0); + reservoir.insert(1, 2, 2.0); + reservoir.insert(2, 3, 3.0); + + // When + assert!(reservoir.insert(3, 4, 0.5)); + + // Then + assert_eq!(reservoir.len(), expected_capacity); + assert_eq!(reservoir.neighbors(), expected_neighbors); + } + + #[test] + fn same_hash_keeps_only_the_closest_candidate() { + // Given + let expected_candidate_count = 1; + let expected_closest_candidate = [(3, 1.0)]; + let mut reservoir = TestReservoir::new(5); + reservoir.insert(0, 1, 3.0); + + // When + reservoir.insert(0, 2, 2.0); + reservoir.insert(0, 3, 1.0); + assert!(!reservoir.insert(0, 4, 5.0)); + + // Then + assert_eq!(reservoir.len(), expected_candidate_count); + assert_eq!(reservoir.neighbors(), expected_closest_candidate); + } + + #[test] + fn equal_distances_are_ordered_by_neighbor_id() { + // Given + let expected_candidate_count = 3; + let expected_neighbor_id_order = [(1, 1.0), (2, 1.0), (3, 1.0)]; + let mut reservoir = TestReservoir::new(5); + reservoir.insert(0, 1, 1.0); + reservoir.insert(1, 2, 1.0); + reservoir.insert(2, 3, 1.0); + + // When + let actual_neighbors = reservoir.neighbors(); + + // Then + assert_eq!(reservoir.len(), expected_candidate_count); + assert_eq!(actual_neighbors, expected_neighbor_id_order); + } + + #[test] + fn same_hash_bf16_ties_are_history_independent() { + for order in [[0, 1], [1, 0]] { + let candidates = [(7, 20, 1.0), (7, 10, 1.0)]; + let mut reservoir = TestReservoir::new(2); + for index in order { + let (hash, neighbor, distance) = candidates[index]; + reservoir.insert(hash, neighbor, distance); + } + assert_eq!(reservoir.neighbors(), [(10, 1.0)], "order={order:?}"); + } + } + + #[test] + fn full_reservoir_bf16_ties_are_history_independent() { + let permutations = [ + [0, 1, 2], + [0, 2, 1], + [1, 0, 2], + [1, 2, 0], + [2, 0, 1], + [2, 1, 0], + ]; + let candidates = [(1, 30, 1.0), (2, 10, 1.0), (3, 20, 1.0)]; + for order in permutations { + let mut reservoir = TestReservoir::new(2); + for index in order { + let (hash, neighbor, distance) = candidates[index]; + reservoir.insert(hash, neighbor, distance); + } + let mut actual = reservoir.neighbors(); + actual.sort_unstable_by_key(|&(neighbor, _)| neighbor); + assert_eq!(actual, [(10, 1.0), (30, 1.0)], "order={order:?}"); + } + } + + // Concurrency and consuming extraction. + + #[test] + fn parallel_insertion_matches_serial_neighbor_lists() { + let data = vec![0.0f32; 100 * 4]; + let parallel = build_hash_prune(&data, 100, 4, 4, 10).unwrap(); + let serial = build_hash_prune(&data, 100, 4, 4, 10).unwrap(); + + std::thread::scope(|scope| { + for sources in [0..25, 25..50] { + let parallel = ∥ + scope.spawn(move || { + for source in sources { + add_edge(parallel, source, source + 1, 1.0); + add_edge(parallel, source + 1, source, 1.0); + } + }); + } + }); + for source in 0..50 { + add_edge(&serial, source, source + 1, 1.0); + add_edge(&serial, source + 1, source, 1.0); + } + + assert_eq!(parallel.into_nearest_lists(5), serial.into_nearest_lists(5)); + } + + fn hash_prune_with_ranked_source_zero_candidates() -> HashPrune { + #[rustfmt::skip] + let point_vectors = [ + 0.0, 0.0, + 1.0, 0.0, + 0.0, 1.0, + -1.0, 0.0, + 0.0, -1.0, + 1.0, 1.0, + -1.0, 1.0, + 1.0, -1.0, + ]; + let hash_prune = build_hash_prune(&point_vectors, 8, 2, 16, 10).unwrap(); + // Use the neighbor ID as distance so the extraction order is explicit. + for target in 1..8 { + add_edge(&hash_prune, 0, target, target as f32); + } + hash_prune + } + + #[test] + fn candidate_extraction_returns_every_retained_neighbor() { + // Given + let hash_prune = hash_prune_with_ranked_source_zero_candidates(); + let expected_neighbor_ids: Vec<_> = (1..8).collect(); + + // When + let mut actual_neighbor_ids = hash_prune.into_candidate_lists()[0].to_vec(); + actual_neighbor_ids.sort_unstable(); + + // Then + assert_eq!(actual_neighbor_ids, expected_neighbor_ids); + } + + #[test] + fn nearest_extraction_keeps_the_requested_closest_neighbors() { + // Given + let hash_prune = hash_prune_with_ranked_source_zero_candidates(); + let expected_nearest_ids = [1, 2]; + + // When + let actual_nearest_lists = hash_prune.into_nearest_lists(2); + + // Then + assert_eq!(&*actual_nearest_lists[0], &expected_nearest_ids); + } + + #[test] + fn farthest_cache_updates_after_repeated_evictions() { + // Given + let expected_neighbors = [(14, 1.0), (13, 2.0), (12, 3.0)]; + let mut reservoir = TestReservoir::new(3); + reservoir.insert(0, 10, 5.0); + reservoir.insert(1, 11, 4.0); + reservoir.insert(2, 12, 3.0); + + // When + assert!(reservoir.insert(3, 13, 2.0)); + assert!(reservoir.insert(4, 14, 1.0)); + + // Then + assert_eq!(reservoir.neighbors(), expected_neighbors); + } + + #[test] + fn extraction_sorts_neighbors_when_the_farthest_slot_is_not_last() { + let mut reservoir = TestReservoir::new(4); + reservoir.insert(5, 1, 1.0); + reservoir.insert(10, 2, 3.0); + reservoir.insert(15, 3, 2.0); + reservoir.insert(3, 4, 0.5); + + assert_eq!( + reservoir.neighbors(), + [(4, 0.5), (1, 1.0), (3, 2.0), (2, 3.0)] + ); + } +} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 6c1f2c822e..f9b3e01227 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -11,11 +11,12 @@ //! 1. Gather each ID and convert its vector to reusable `f32` storage. //! 2. Call the leaf kernel for Gram construction, norms, and local ranking. //! 3. Convert local positions to global point IDs. -//! 4. Add both edge directions to global candidate lists. +//! 4. Add both edge directions to direct candidates or HashPrune reservoirs. //! -//! Overlapping leaves run concurrently. A worker locks one destination list only -//! while it adds one leaf's IDs. Reusable buffers keep their largest allocation. -//! Each operation uses an explicit active prefix. +//! Overlapping leaves run concurrently. The direct path locks one destination +//! list while it adds IDs. The HashPrune path locks one source reservoir while it +//! adds weighted edges. Reusable buffers keep their largest allocation. Each +//! operation uses an explicit active prefix. use parking_lot::Mutex; @@ -29,7 +30,7 @@ use super::{ simd::PiPNNSIMDSchema, }; -/// Failure while converting leaves into direct graph candidates. +/// Failure while converting leaves into graph candidates. #[derive(Debug, thiserror::Error)] pub(crate) enum LeafBuildError { #[error("leaf {leaf} shape {rows} x {columns} overflows usize")] @@ -53,18 +54,26 @@ pub(crate) enum LeafBuildError { #[source] source: crate::ANNError, }, + #[error("leaf {leaf} produced too many directed edges")] + TooManyEdges { leaf: usize }, } /// Reusable buffers for one Rayon leaf job. /// -/// The numerical vectors keep the largest leaf shape that this job observed. -/// The job creates local adjacency lists only when the effective `k` is not zero. +/// The buffers keep the largest leaf shape that this job observed. The direct +/// path uses `local_adjacency`. The HashPrune path uses the CSR and sketch +/// buffers. #[derive(Default)] struct LeafBuffers { point_values: Vec, neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, + seen_pairs: Vec, + edge_offsets: Vec, + edges: Vec<(u32, f32)>, + edge_cursor: Vec, + sketch_scratch: Vec, } impl LeafBuffers { @@ -83,6 +92,13 @@ impl LeafBuffers { rows: point_count, columns: dimension_count, })?; + point_count + .checked_mul(point_count) + .ok_or(LeafBuildError::ShapeOverflow { + leaf, + rows: point_count, + columns: point_count, + })?; let leaf_k = leaf_neighbor_count(point_count, requested_k); let neighbor_count = point_count @@ -104,6 +120,11 @@ impl LeafBuffers { .iter_mut() .for_each(Vec::clear); } + + fn prepare_seen_pairs(&mut self, point_count: usize) { + // `prepare` checked this product for the same leaf shape. + grow(&mut self.seen_pairs, point_count * point_count, false); + } } /// Concurrent candidate lists indexed by global point ID. @@ -178,6 +199,56 @@ where Ok(candidates.into_lists()) } +/// Add weighted symmetric leaf edges to HashPrune reservoirs. +#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. +pub(super) fn add_hash_prune_candidates( + arch: A, + data: MatrixView<'_, T>, + leaves: Vec>, + requested_k: usize, + hash_prune: &super::hash_prune::HashPrune, +) -> Result<(), LeafBuildError> +where + A: PiPNNSIMDSchema, + M: LeafMetric, + T: VectorRepr + 'static, +{ + leaves.par_iter().enumerate().try_for_each_init( + LeafBuffers::default, + |buffers, (leaf, point_ids)| { + let leaf_k = gather_leaf_neighbors::( + arch, + data, + leaf, + point_ids, + requested_k, + buffers, + )?; + let point_count = point_ids.len(); + buffers.prepare_seen_pairs(point_count); + let edge_count = build_symmetric_edge_csr( + leaf, + point_ids, + leaf_k, + &buffers.neighbors[..point_count * leaf_k], + EdgeBuffers { + seen: &mut buffers.seen_pairs[..point_count * point_count], + offsets: &mut buffers.edge_offsets, + edges: &mut buffers.edges, + cursor: &mut buffers.edge_cursor, + }, + )?; + hash_prune.add_leaf_edges( + point_ids, + &buffers.edge_offsets[..point_count + 1], + &buffers.edges[..edge_count], + &mut buffers.sketch_scratch, + ); + Ok(()) + }, + ) +} + /// Add one leaf's symmetric neighbors to the direct candidate lists. /// /// Reusable buffers can be longer than this leaf, so all accesses use the current @@ -192,6 +263,40 @@ fn add_direct_leaf_candidates( buffers: &mut LeafBuffers, candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> +where + A: PiPNNSIMDSchema, + M: LeafMetric, + T: VectorRepr + 'static, +{ + let leaf_k = + gather_leaf_neighbors::(arch, data, leaf, point_ids, requested_k, buffers)?; + if leaf_k == 0 { + return Ok(()); + } + buffers.prepare_local_adjacency(point_ids.len()); + add_symmetric_neighbors( + point_ids, + leaf_k, + &buffers.neighbors[..point_ids.len() * leaf_k], + &mut buffers.local_adjacency[..point_ids.len()], + ); + candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]); + Ok(()) +} + +/// Select local nearest neighbors for one leaf. +/// +/// The function gathers leaf IDs into a packed `f32` matrix. The leaf kernel +/// owns Gram construction, norm preparation, and local ranking. This function +/// returns the effective neighbor count for graph-edge mapping. +fn gather_leaf_neighbors( + arch: A, + data: MatrixView<'_, T>, + leaf: usize, + point_ids: &[u32], + requested_k: usize, + buffers: &mut LeafBuffers, +) -> Result where A: PiPNNSIMDSchema, M: LeafMetric, @@ -200,7 +305,7 @@ where let (leaf_k, neighbor_value_count) = buffers.prepare(leaf, point_ids.len(), data.ncols(), requested_k)?; if leaf_k == 0 { - return Ok(()); + return Ok(0); } let point_value_count = point_ids.len() * data.ncols(); @@ -238,18 +343,12 @@ where })?; select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - - buffers.prepare_local_adjacency(point_ids.len()); - add_symmetric_neighbors( - point_ids, - leaf_k, - &buffers.neighbors[..neighbor_value_count], - &mut buffers.local_adjacency[..point_ids.len()], - ); - candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]); - Ok(()) + Ok(leaf_k) } +/// Add symmetric dataset IDs from one leaf-kernel result. +/// +/// The leaf kernel returns only leaf-local positions in `point_ids`. fn add_symmetric_neighbors( point_ids: &[u32], leaf_k: usize, @@ -272,6 +371,132 @@ fn add_symmetric_neighbors( } } +struct EdgeBuffers<'a> { + seen: &'a mut [bool], + offsets: &'a mut Vec, + edges: &'a mut Vec<(u32, f32)>, + cursor: &'a mut Vec, +} + +/// Create directed leaf edges for HashPrune ingestion. +/// +/// Each selected neighbor pair contributes both directions. Duplicate directions +/// appear once. Each target is a position in `point_ids`. +/// Build weighted CSR edges from one leaf-kernel result. +/// +/// The leaf kernel returns only leaf-local positions in `point_ids`. +fn build_symmetric_edge_csr( + leaf: usize, + point_ids: &[u32], + leaf_k: usize, + neighbors: &[LeafNeighbor], + buffers: EdgeBuffers<'_>, +) -> Result { + let EdgeBuffers { + seen, + offsets, + edges, + cursor, + } = buffers; + let point_count = point_ids.len(); + grow(offsets, point_count + 1, 0); + offsets[..point_count + 1].fill(0); + if leaf_k == 0 { + return Ok(0); + } + + // The prior successful write pass left the active `seen` area clear. This + // count pass marks each unique directed edge. + for (source, neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for neighbor in neighbors { + if !neighbor.is_assigned() { + continue; + } + let target = neighbor.target as usize; + count_directed_edge(leaf, point_count, source, target, seen, offsets)?; + count_directed_edge(leaf, point_count, target, source, seen, offsets)?; + } + } + for point in 1..=point_count { + offsets[point] = offsets[point] + .checked_add(offsets[point - 1]) + .ok_or(LeafBuildError::TooManyEdges { leaf })?; + } + + let edge_count = offsets[point_count] as usize; + grow(edges, edge_count, (0, 0.0)); + grow(cursor, point_count, 0); + cursor[..point_count].copy_from_slice(&offsets[..point_count]); + let edges = &mut edges[..edge_count]; + let cursor = &mut cursor[..point_count]; + + // This pass visits the same directions as the count pass. The first + // occurrence writes its edge and clears its mark for the next leaf. + for (source, neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for neighbor in neighbors { + if !neighbor.is_assigned() { + continue; + } + let target = neighbor.target as usize; + write_counted_directed_edge( + point_count, + source, + target, + neighbor.distance, + seen, + edges, + cursor, + ); + write_counted_directed_edge( + point_count, + target, + source, + neighbor.distance, + seen, + edges, + cursor, + ); + } + } + Ok(edge_count) +} + +fn count_directed_edge( + leaf: usize, + point_count: usize, + source: usize, + target: usize, + seen: &mut [bool], + offsets: &mut [u32], +) -> Result<(), LeafBuildError> { + let seen_entry = &mut seen[source * point_count + target]; + if !*seen_entry { + *seen_entry = true; + offsets[source + 1] = offsets[source + 1] + .checked_add(1) + .ok_or(LeafBuildError::TooManyEdges { leaf })?; + } + Ok(()) +} + +fn write_counted_directed_edge( + point_count: usize, + source: usize, + target: usize, + distance: f32, + seen: &mut [bool], + edges: &mut [(u32, f32)], + cursor: &mut [u32], +) { + let seen_entry = &mut seen[source * point_count + target]; + if *seen_entry { + *seen_entry = false; + let edge_slot = cursor[source] as usize; + edges[edge_slot] = (target as u32, distance); + cursor[source] += 1; + } +} + fn grow(values: &mut Vec, len: usize, value: T) { if values.len() < len { values.resize(len, value); @@ -287,10 +512,10 @@ mod tests { use rstest::rstest; use std::collections::BTreeSet; - use super::super::simd::PiPNNSIMDSchema; + use super::super::{leaf_kernel::LeafNeighbor, simd::PiPNNSIMDSchema}; use super::{ - DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, - build_leaf_candidates, + DirectCandidates, EdgeBuffers, LeafBuffers, LeafBuildError, add_symmetric_neighbors, + build_leaf_candidates, build_symmetric_edge_csr, }; fn matrix_view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { @@ -693,4 +918,126 @@ mod tests { candidates.add_leaf(&[0, 1], &[vec![1, 1], vec![0]]); assert_eq!(adjacency_lists(candidates.into_lists()), [vec![1], vec![0]]); } + + #[test] + fn symmetric_edge_csr_contains_both_directions_in_source_order() { + // Given + let point_ids = [10, 20, 30]; + let neighbors = [ + LeafNeighbor::new(1, 1.0), + LeafNeighbor::new(2, 2.0), + LeafNeighbor::new(1, 1.5), + ]; + let expected_edge_count = 4; + let expected_offsets = [0, 1, 3, 4]; + let expected_edges = [(1, 1.0), (0, 1.0), (2, 2.0), (1, 2.0)]; + let mut seen = vec![false; 9]; + let mut offsets = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = Vec::new(); + + // When + let actual_edge_count = build_symmetric_edge_csr( + 0, + &point_ids, + 1, + &neighbors, + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap(); + + // Then + assert_eq!(actual_edge_count, expected_edge_count); + assert_eq!(offsets, expected_offsets); + assert_eq!(edges, expected_edges); + } + + #[test] + fn symmetric_edge_csr_omits_unassigned_neighbors() { + let point_ids = [10, 20]; + let neighbors = [LeafNeighbor::new(1, 1.0), LeafNeighbor::default()]; + let mut seen = vec![false; 4]; + let mut offsets = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = Vec::new(); + + let count = build_symmetric_edge_csr( + 0, + &point_ids, + 1, + &neighbors, + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap(); + + assert_eq!(count, 2); + assert_eq!(offsets, [0, 1, 2]); + assert_eq!(edges, [(1, 1.0), (0, 1.0)]); + } + + #[test] + fn symmetric_edge_csr_deduplicates_edges_seen_from_both_endpoints() { + let point_ids = [10, 20]; + let neighbors = [LeafNeighbor::new(1, 1.0), LeafNeighbor::new(0, 1.0)]; + let mut seen = vec![false; 4]; + let mut offsets = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = Vec::new(); + + let count = build_symmetric_edge_csr( + 0, + &point_ids, + 1, + &neighbors, + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap(); + + assert_eq!(count, 2); + assert_eq!(offsets, [0, 1, 2]); + assert_eq!(edges, [(1, 1.0), (0, 1.0)]); + assert!(seen.iter().all(|&entry| !entry)); + } + #[test] + fn zero_k_edge_csr_has_empty_adjacency() { + let point_ids = [10, 20, 30]; + let zero_k = 0; + let mut seen = vec![false; 9]; + let mut offsets = Vec::new(); + let mut edges = vec![(99, 99.0)]; + let mut cursor = Vec::new(); + + let count = build_symmetric_edge_csr( + 0, + &point_ids, + zero_k, + &[], + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap(); + + assert_eq!(count, 0); + assert_eq!(offsets, [0, 0, 0, 0]); + assert_eq!(edges, [(99, 99.0)]); + } } diff --git a/diskann/src/graph/pipnn/lsh.rs b/diskann/src/graph/pipnn/lsh.rs new file mode 100644 index 0000000000..7bcd532527 --- /dev/null +++ b/diskann/src/graph/pipnn/lsh.rs @@ -0,0 +1,274 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Random-hyperplane locality-sensitive hashing for dataset vectors. +//! +//! For each point `v`, the module computes +//! `Sketch(v) = [v · H_i for i in 0..num_planes]`. A seeded random generator +//! samples each hyperplane component from a standard normal distribution. +//! HashPrune compares two sketches to make a relative hash. +//! +//! `LshSketches` stores a row-major `npoints × num_planes` matrix. Each Rayon job +//! uses one `f32` conversion buffer for its source rows. `num_planes` cannot +//! exceed 16 because each relative hash is a `u16`. + +use crate::{ANNError, ANNResult, utils::VectorRepr}; +use diskann_utils::views::MatrixView; +use rand::SeedableRng; +use rand_distr::{Distribution, StandardNormal}; +#[cfg(not(miri))] +use rayon::prelude::*; + +/// Maximum number of hyperplanes (the hash output is `u16`). +pub(super) const MAX_PLANES: usize = 16; + +/// Precomputed LSH sketches for `npoints` vectors. +#[derive(Debug)] +pub(super) struct LshSketches { + num_planes: usize, + /// Row-major `npoints × num_planes`: `sketches[i*m + j] = dot(point_i, plane_j)`. + sketches: Vec, +} + +impl LshSketches { + /// Compute random-hyperplane projections for every point in `data`. + /// + /// Each worker converts one source row into reusable `f32` storage. Parallel + /// sketch work uses the currently installed Rayon pool. + pub(super) fn try_new( + data: MatrixView<'_, T>, + num_planes: usize, + seed: u64, + ) -> ANNResult { + let npoints = data.nrows(); + let ndims = data.ncols(); + let hyperplane_len = num_planes.checked_mul(ndims).ok_or_else(|| { + ANNError::message(format!( + "LSH matrix shape {num_planes} x {ndims} overflows usize" + )) + })?; + let sketch_len = npoints.checked_mul(num_planes).ok_or_else(|| { + ANNError::message(format!( + "LSH matrix shape {npoints} x {num_planes} overflows usize" + )) + })?; + + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let hyperplanes: Vec = (0..hyperplane_len) + .map(|_| StandardNormal.sample(&mut rng)) + .collect(); + + let mut sketches = vec![0.0f32; sketch_len]; + + #[cfg(miri)] + { + // Miri cannot execute Rayon's crossbeam dependency with strict provenance. + let mut buffer = Vec::new(); + for (point, sketch_row) in sketches.chunks_mut(num_planes).enumerate() { + fill_sketch_row(data, point, &hyperplanes, &mut buffer, sketch_row)?; + } + } + + #[cfg(not(miri))] + { + #[allow(clippy::disallowed_methods)] // caller installs the complete build in its pool. + sketches + .par_chunks_mut(num_planes) + .enumerate() + .try_for_each_init(Vec::new, |buffer, (point, sketch_row)| { + fill_sketch_row(data, point, &hyperplanes, buffer, sketch_row) + })?; + } + + Ok(Self { + num_planes, + sketches, + }) + } + + /// Number of hyperplanes (also the number of bits in the hash). + #[inline] + pub(super) fn num_planes(&self) -> usize { + self.num_planes + } + + /// Return the row-major `npoints × num_planes` sketch buffer. + #[inline] + pub(super) fn sketches(&self) -> &[f32] { + &self.sketches + } +} + +fn fill_sketch_row( + data: MatrixView<'_, T>, + point: usize, + hyperplanes: &[f32], + buffer: &mut Vec, + sketch_row: &mut [f32], +) -> ANNResult<()> { + let dimensions = data.ncols(); + buffer.resize(dimensions, 0.0); + T::as_f32_into(data.row(point), &mut buffer[..dimensions]) + .map_err(Into::::into) + .map_err(|error| error.context(format!("converting LSH point {point}")))?; + for (plane_index, destination) in sketch_row.iter_mut().enumerate() { + let plane = &hyperplanes[plane_index * dimensions..(plane_index + 1) * dimensions]; + let mut dot = 0.0f32; + for dimension in 0..dimensions { + // SAFETY: both slices have exactly `dimensions` elements. + unsafe { + dot += *buffer.get_unchecked(dimension) * *plane.get_unchecked(dimension); + } + } + *destination = dot; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + fn run_with_threads(threads: usize, operation: impl FnOnce() -> R + Send) -> R { + #[cfg(miri)] + { + let _ = threads; + operation() + } + #[cfg(not(miri))] + { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + .install(operation) + } + } + + fn matrix_view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + MatrixView::try_from(data, rows, columns).unwrap() + } + + #[test] + fn sketch_shape_has_one_value_per_point_and_plane() { + // Given + let point_vectors = [[1.0_f32, 0.0], [0.0, 1.0], [-1.0, 0.0]]; + let point_count = point_vectors.len(); + let dimensions = point_vectors[0].len(); + let plane_count = 4; + let expected_sketch_value_count = point_count * plane_count; + let data: Vec<_> = point_vectors.into_iter().flatten().collect(); + + // When + let sketches = run_with_threads(2, || { + LshSketches::try_new(matrix_view(&data, point_count, dimensions), plane_count, 42) + }) + .unwrap(); + + // Then + assert_eq!(sketches.num_planes(), plane_count); + assert_eq!(sketches.sketches().len(), expected_sketch_value_count); + } + + #[rstest] + #[case::first_seed(42)] + #[case::second_seed(99)] + fn sketches_match_seeded_serial_hyperplane_reference(#[case] seed: u64) { + // Given + let point_vectors = [ + [-3.0_f32, -2.0, -1.0, 0.0], + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + ]; + let npoints = point_vectors.len(); + let ndims = point_vectors[0].len(); + let planes = 5; + let data: Vec<_> = point_vectors.into_iter().flatten().collect(); + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let hyperplanes: Vec = (0..planes * ndims) + .map(|_| StandardNormal.sample(&mut rng)) + .collect(); + let expected_serial_sketch_values: Vec = data + .chunks_exact(ndims) + .flat_map(|point| { + hyperplanes + .chunks_exact(ndims) + .map(|plane| point.iter().zip(plane).map(|(x, h)| x * h).sum()) + }) + .collect(); + + // When + let actual_sketches = run_with_threads(2, || { + LshSketches::try_new(matrix_view(&data, npoints, ndims), planes, seed) + }) + .unwrap(); + + // Then + assert_eq!(actual_sketches.sketches(), expected_serial_sketch_values); + } + + #[test] + fn zero_points_produce_an_empty_sketch() { + // Given + let zero_point_count = 0; + let dimensions = 7; + let plane_count = 4; + + // When + let sketches = run_with_threads(2, || { + LshSketches::try_new( + matrix_view(&[] as &[f32], zero_point_count, dimensions), + plane_count, + 42, + ) + }) + .unwrap(); + + // Then + assert_eq!(sketches.num_planes(), plane_count); + assert!(sketches.sketches().is_empty()); + } + + #[test] + fn zero_dimensions_produce_zero_dot_products() { + // Given + let point_count = 3; + let zero_dimensions = 0; + let plane_count = 2; + let expected_zero_dot_products = vec![0.0; point_count * plane_count]; + + // When + let sketches = run_with_threads(2, || { + LshSketches::try_new( + matrix_view(&[] as &[f32], point_count, zero_dimensions), + plane_count, + 42, + ) + }) + .unwrap(); + + // Then + assert_eq!(sketches.sketches(), expected_zero_dot_products); + } + + #[rstest] + #[case::point_count_times_plane_count(usize::MAX, 0)] + #[case::dimension_times_plane_count(0, usize::MAX)] + fn sketch_construction_rejects_shape_overflow( + #[case] point_count: usize, + #[case] dimensions: usize, + ) { + // Given + let empty_data = matrix_view(&[] as &[f32], point_count, dimensions); + + // When + let error = + LshSketches::try_new(empty_data, 2, 42).expect_err("overflowing LSH shape must fail"); + + // Then + assert!(error.to_string().contains("overflows")); + } +} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 5ed62698cf..5c946aa00f 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -12,9 +12,10 @@ //! 1. `partitioning` samples leaders and makes overlapping leaves. Each leaf has //! at most `c_max` points. //! 2. `leaf_build` computes a lower-triangular Gram matrix for each leaf. It -//! selects local neighbors and merges their global point IDs. -//! 3. `finalization` applies Vamana RobustPrune to each candidate list that is -//! longer than the graph degree. +//! selects local neighbors. The direct path merges their global point IDs. +//! The HashPrune path sends weighted edges to bounded point reservoirs. +//! 3. `finalization` applies Vamana RobustPrune to direct candidates. It also +//! prunes HashPrune candidates when `final_prune` is true. //! //! `diskann-wide` selects architecture `A`. One match selects metric marker `M`. //! The build passes both concrete types through all replicas, recursive @@ -34,9 +35,12 @@ mod kernel_metric; mod simd; +mod bf16; mod finalization; +mod hash_prune; mod leaf_build; mod leaf_kernel; +mod lsh; mod partition_kernel; mod partitioning; @@ -112,6 +116,55 @@ impl PiPNNConfig { } } +/// HashPrune policy for bounded candidate reservoirs. +#[derive(Clone, Debug, PartialEq)] +pub struct HashPruneConfig { + /// Number of random-hyperplane bits in each relative-direction hash. + pub num_hash_planes: usize, + /// Maximum number of direction buckets retained for each source point. + pub l_max: usize, + /// Apply Vamana RobustPrune after reservoir extraction. + pub final_prune: bool, +} + +impl HashPruneConfig { + /// Check the structural HashPrune limits. + pub fn validate(&self) -> ANNResult<()> { + if !(1..=lsh::MAX_PLANES).contains(&self.num_hash_planes) { + return Err(config_error(format!( + "num_hash_planes ({}) must be in [1, {}]", + self.num_hash_planes, + lsh::MAX_PLANES + ))); + } + if !(1..=hash_prune::MAX_RESERVOIR_LEN).contains(&self.l_max) { + return Err(config_error(format!( + "l_max ({}) must be in [1, {}]", + self.l_max, + hash_prune::MAX_RESERVOIR_LEN + ))); + } + Ok(()) + } + + /// Check that the reservoir and hash space can hold `degree` neighbors. + pub fn validate_for_degree(&self, degree: usize) -> ANNResult<()> { + self.validate()?; + let hash_capacity = 1usize + .checked_shl(self.num_hash_planes as u32) + .unwrap_or(usize::MAX); + let candidate_capacity = self.l_max.min(hash_capacity); + if candidate_capacity < degree { + return Err(config_error(format!( + "HashPrune capacity min(l_max={}, hash buckets={hash_capacity}) must be at least \ + the graph degree ({degree})", + self.l_max + ))); + } + Ok(()) + } +} + /// PiPNN policy and borrowed execution resources for one graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { @@ -119,6 +172,7 @@ pub struct PiPNNBuildContext<'a> { pub(crate) graph: &'a Config, pub(crate) metric: Metric, pub(crate) pool: &'a ThreadPool, + hash_prune: Option, } impl<'a> PiPNNBuildContext<'a> { @@ -142,8 +196,16 @@ impl<'a> PiPNNBuildContext<'a> { graph, metric, pool, + hash_prune: None, }) } + + /// Enable HashPrune candidate merging for this build. + pub fn with_hash_prune(mut self, config: HashPruneConfig) -> ANNResult { + config.validate_for_degree(self.graph.pruned_degree().get())?; + self.hash_prune = Some(config); + Ok(self) + } } /// Build one PiPNN adjacency list for each point in `data`. @@ -242,8 +304,9 @@ where /// Run the PiPNN graph pipeline for one selected metric implementation. /// -/// The function builds overlapping leaves, merges direct candidates, and applies -/// final graph-degree pruning. +/// The function builds overlapping leaves and runs the configured candidate +/// merge. It prunes direct candidates to graph degree. It prunes HashPrune +/// candidates when `final_prune` is true. fn build_graph_for( arch: A, data: MatrixView<'_, T>, @@ -257,16 +320,49 @@ where { let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition::(arch, data, &context.config))?; - // Leaf jobs borrow individual ID lists. This call consumes the leaf vector, - // so its complete allocation drops when leaf construction returns. - let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates::(arch, data, leaves, context.config.leaf_k) - .map_err(ANNError::new) - })?; - // Finalization consumes each candidate list. It reuses that list's allocation - // for the final adjacency when the graph policy permits it. - tracing::info_span!("pipnn.finalization") - .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) + match &context.hash_prune { + None => { + // Leaf jobs borrow individual ID lists. This call consumes the leaf + // vector, so its allocation drops when leaf construction returns. + let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::build_leaf_candidates::( + arch, + data, + leaves, + context.config.leaf_k, + ) + .map_err(ANNError::new) + })?; + tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) + } + Some(config) => { + // `HashPrune` lives until all leaf jobs finish. A leaf job locks only + // one source reservoir at a time. + let hash_prune = + hash_prune::HashPrune::new(data, config.num_hash_planes, config.l_max, 42)?; + // This call consumes the leaves. Each weighted CSR list exists only + // during its leaf job. The reservoirs retain the selected edges. + tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::add_hash_prune_candidates::( + arch, + data, + leaves, + context.config.leaf_k, + &hash_prune, + ) + .map_err(ANNError::new) + })?; + if config.final_prune { + let candidates = hash_prune.into_candidate_lists(); + tracing::info_span!("pipnn.finalization").in_scope(|| { + finalization::prune_overfull(data, candidates, context.graph, metric) + }) + } else { + Ok(hash_prune.into_nearest_lists(context.graph.pruned_degree().get())) + } + } + } } fn effective_metric(metric: Metric) -> Metric { @@ -290,25 +386,40 @@ fn config_error(message: impl std::fmt::Display) -> ANNError { mod tests { use super::*; use half::f16; + use rstest::rstest; #[test] - fn integer_normalized_cosine_uses_unnormalized_cosine() { - for metric in [ + fn integer_vectors_use_cosine_when_normalized_cosine_is_requested() { + assert_eq!( + effective_metric::(Metric::CosineNormalized), + Metric::Cosine + ); + assert_eq!( + effective_metric::(Metric::CosineNormalized), + Metric::Cosine + ); + } + + #[rstest] + fn metric_selection_is_unchanged_for_float_vectors( + #[values( Metric::L2, Metric::Cosine, Metric::CosineNormalized, - Metric::InnerProduct, - ] { - let expected = if metric == Metric::CosineNormalized { - Metric::Cosine - } else { - metric - }; - assert_eq!(effective_metric::(metric), expected); - assert_eq!(effective_metric::(metric), expected); - assert_eq!(effective_metric::(metric), metric); - assert_eq!(effective_metric::(metric), metric); - } + Metric::InnerProduct + )] + metric: Metric, + ) { + assert_eq!(effective_metric::(metric), metric); + assert_eq!(effective_metric::(metric), metric); + } + + #[rstest] + fn integer_vectors_keep_non_normalized_metric_selection( + #[values(Metric::L2, Metric::Cosine, Metric::InnerProduct)] metric: Metric, + ) { + assert_eq!(effective_metric::(metric), metric); + assert_eq!(effective_metric::(metric), metric); } } #[cfg(test)] @@ -318,12 +429,13 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod build_graph_tests { - use super::{PiPNNBuildContext, PiPNNConfig, build_graph}; + use super::{HashPruneConfig, PiPNNBuildContext, PiPNNConfig, build_graph}; use crate::graph::config::{self, MaxDegree}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; use half::f16; use rand::{Rng, SeedableRng, rngs::StdRng}; + use rstest::rstest; fn pipnn_config() -> PiPNNConfig { PiPNNConfig { @@ -344,17 +456,26 @@ mod build_graph_tests { .unwrap() } - fn pool(threads: usize) -> rayon::ThreadPool { + fn thread_pool(threads: usize) -> rayon::ThreadPool { rayon::ThreadPoolBuilder::new() .num_threads(threads) .build() .unwrap() } - fn rows(graph: Vec>) -> Vec> { + fn adjacency_rows(graph: Vec>) -> Vec> { graph.into_iter().map(Vec::from).collect() } + fn deterministic_point_values(points: usize, dimensions: usize) -> Vec { + (0..points) + .flat_map(|point| { + (0..dimensions) + .map(move |dimension| point as f32 + dimension as f32 / dimensions as f32) + }) + .collect() + } + fn assert_graph_invariants( graph: &[crate::graph::AdjacencyList], points: usize, @@ -375,34 +496,45 @@ mod build_graph_tests { } #[test] - fn builds_a_single_leaf_graph_for_real_dataset_ids() { - let data = [0.0_f32, 1.0, 2.0, 3.0]; - let data = MatrixView::try_from(&data[..], 4, 1).unwrap(); + fn single_leaf_build_maps_local_neighbors_to_dataset_ids() { + // Given + let point_values = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&point_values[..], 4, 1).unwrap(); let graph = graph_config(Metric::L2, 2); - let pool = pool(2); + let pool = thread_pool(2); let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + let expected_adjacency = [vec![1], vec![0, 2], vec![1, 3], vec![2]]; - let actual = build_graph(data, &context).unwrap(); + // When + let actual_adjacency = adjacency_rows(build_graph(data, &context).unwrap()); - assert_eq!(rows(actual), [vec![1], vec![0, 2], vec![1, 3], vec![2]]); + // Then + assert_eq!(actual_adjacency, expected_adjacency); + } + #[test] + fn degree_one_pruning_keeps_adjacent_neighbors_on_a_line() { + // Given + let point_values = [0.0_f32, 1.0, 2.0, 3.0]; + let data = MatrixView::try_from(&point_values[..], 4, 1).unwrap(); let graph = graph_config(Metric::L2, 1); + let pool = thread_pool(2); let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + let expected_adjacency = [vec![1], vec![2], vec![3], vec![2]]; - let pruned = build_graph(data, &context).unwrap(); + // When + let actual_adjacency = adjacency_rows(build_graph(data, &context).unwrap()); - assert_graph_invariants(&pruned, 4, 1); - for (source, neighbors) in pruned.iter().enumerate() { - assert_eq!(source.abs_diff(neighbors[0] as usize), 1); - } + // Then + assert_eq!(actual_adjacency, expected_adjacency); } #[test] - fn omits_non_rankable_candidates_without_invalid_ids() { + fn non_rankable_points_leave_empty_adjacency_without_sentinel_ids() { let values = [0.0_f32, 1.0, f32::NAN]; let data = MatrixView::try_from(&values[..], 3, 1).unwrap(); let graph = graph_config(Metric::InnerProduct, 2); - let pool = pool(1); + let pool = thread_pool(1); let config = PiPNNConfig { c_max: 2, c_min: 1, @@ -412,19 +544,20 @@ mod build_graph_tests { replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, Metric::InnerProduct, &pool).unwrap(); + let expected_rankable_adjacency = [vec![1], vec![0], vec![]]; - let actual = build_graph(data, &context).unwrap(); + let actual_graph = build_graph(data, &context).unwrap(); - assert_graph_invariants(&actual, 3, 2); - assert_eq!(rows(actual), [vec![1], vec![0], vec![]]); + assert_graph_invariants(&actual_graph, 3, 2); + assert_eq!(adjacency_rows(actual_graph), expected_rankable_adjacency); } #[test] - fn prunes_overfull_single_leaf_candidates_to_the_graph_degree() { + fn single_leaf_adjacency_is_bounded_by_the_graph_degree() { let data = [0.0_f32, 1.0, 2.0, 3.0, 4.0]; let data = MatrixView::try_from(&data[..], 5, 1).unwrap(); let graph = graph_config(Metric::L2, 1); - let pool = pool(2); + let pool = thread_pool(2); let config = PiPNNConfig { c_max: 5, c_min: 1, @@ -435,60 +568,88 @@ mod build_graph_tests { }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - let actual = build_graph(data, &context).unwrap(); + let actual_graph = build_graph(data, &context).unwrap(); - assert_graph_invariants(&actual, 5, 1); - assert!(actual.iter().all(|row| row.len() == 1)); + assert_graph_invariants(&actual_graph, 5, 1); + assert!(actual_graph.iter().all(|row| row.len() == 1)); } - #[test] - fn rejects_empty_dataset_dimensions_at_the_public_boundary() { + #[rstest] + #[case::zero_points(0, 4)] + #[case::zero_dimensions(4, 0)] + fn empty_dataset_dimension_is_rejected(#[case] point_count: usize, #[case] dimensions: usize) { + // Given let graph = graph_config(Metric::L2, 2); - let pool = pool(1); + let pool = thread_pool(1); let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + let empty_data = MatrixView::try_from(&[] as &[f32], point_count, dimensions).unwrap(); - let no_rows = MatrixView::try_from(&[] as &[f32], 0, 4).unwrap(); - let no_columns = MatrixView::try_from(&[] as &[f32], 4, 0).unwrap(); + // When + let result = build_graph(empty_data, &context); - assert!(build_graph(no_rows, &context).is_err()); - assert!(build_graph(no_columns, &context).is_err()); + // Then + assert!(result.is_err()); } - #[test] - fn supports_every_source_type_and_metric() { - fn build( - values: &[T], - metric: Metric, - ) { - let data = MatrixView::try_from(values, 6, 2).unwrap(); - let graph = graph_config(metric, 2); - let pool = pool(2); - let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); - let actual = build_graph(data, &context).unwrap(); - assert_graph_invariants(&actual, 6, 2); - } + fn assert_graph_build_succeeds( + values: &[T], + metric: Metric, + ) { + let data = MatrixView::try_from(values, 6, 2).unwrap(); + let graph = graph_config(metric, 2); + let pool = thread_pool(2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, metric, &pool).unwrap(); + let actual_graph = build_graph(data, &context).unwrap(); + assert_graph_invariants(&actual_graph, 6, 2); + } - let values = [ - 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, 0.5, 0.5, -0.5, -0.5, - ]; - for metric in [ + #[rstest] + fn f32_graph_build_succeeds_with_each_metric( + #[values( Metric::L2, Metric::Cosine, Metric::CosineNormalized, - Metric::InnerProduct, - ] { - build(&values, metric); - } - build(&values.map(f16::from_f32), Metric::L2); - build(&[1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2], Metric::L2); - build(&[1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1], Metric::L2); + Metric::InnerProduct + )] + metric: Metric, + ) { + let diagonal = std::f32::consts::FRAC_1_SQRT_2; + let unit_vectors = [ + 1.0_f32, 0.0, 0.0, 1.0, -1.0, 0.0, 0.0, -1.0, diagonal, diagonal, -diagonal, -diagonal, + ]; + + assert_graph_build_succeeds(&unit_vectors, metric); + } + + #[test] + fn f16_graph_build_succeeds_with_l2() { + let values = [ + 1.0_f32, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0, 1.0, 1.0, 2.0, 2.0, + ]; + assert_graph_build_succeeds(&values.map(f16::from_f32), Metric::L2); } #[test] - fn integer_normalized_cosine_matches_cosine() { - fn assert_match(values: &[T]) { + fn u8_graph_build_succeeds_with_l2() { + let values = [1_u8, 0, 0, 1, 2, 0, 0, 2, 1, 1, 2, 2]; + assert_graph_build_succeeds(&values, Metric::L2); + } + + #[test] + fn i8_graph_build_succeeds_with_l2() { + let values = [1_i8, 0, 0, 1, -1, 0, 0, -1, 1, 1, -1, -1]; + assert_graph_build_succeeds(&values, Metric::L2); + } + + #[test] + fn integer_vector_graphs_match_cosine_when_normalized_cosine_is_requested() { + fn assert_integer_graphs_match_cosine< + T: crate::utils::VectorRepr + Send + Sync + 'static, + >( + values: &[T], + ) { let data = MatrixView::try_from(values, 8, 2).unwrap(); - let pool = pool(2); + let pool = thread_pool(2); let build = |metric| { let graph = graph_config(metric, 2); let config = PiPNNConfig { @@ -500,23 +661,23 @@ mod build_graph_tests { replicas: 1, }; let context = PiPNNBuildContext::new(config, &graph, metric, &pool).unwrap(); - rows(build_graph(data, &context).unwrap()) + adjacency_rows(build_graph(data, &context).unwrap()) }; assert_eq!(build(Metric::CosineNormalized), build(Metric::Cosine)); } - assert_match(&[1_u8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 200, 2, 2, 1, 1, 2]); - assert_match(&[1_i8, 0, 100, 1, 2, 0, 0, 1, 1, 1, 120, 2, 2, 1, 1, 2]); + assert_integer_graphs_match_cosine(&[1_u8, 0, 2, 0, 0, 1, 0, 2, 1, 1, 2, 1, 1, 2, 2, 2]); + assert_integer_graphs_match_cosine(&[ + 1_i8, 0, -1, 0, 0, 1, 0, -1, 1, 1, -1, -1, 1, -1, -1, 1, + ]); } #[test] - fn is_deterministic_for_a_fixed_pool_size() { - let data: Vec = (0..96 * 4) - .map(|value| ((value * 17 + 3) % 101) as f32) - .collect(); + fn graph_build_is_deterministic_for_a_fixed_pool_size() { + let data = deterministic_point_values(96, 4); let data = MatrixView::try_from(&data[..], 96, 4).unwrap(); let graph = graph_config(Metric::L2, 8); - let pool = pool(4); + let pool = thread_pool(4); let config = PiPNNConfig { c_max: 16, c_min: 4, @@ -535,7 +696,7 @@ mod build_graph_tests { } #[test] - fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { + fn graph_build_preserves_invariants_across_fixed_seed_inputs() { let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); for case in 0..24 { let points = rng.random_range(4..=32); @@ -548,7 +709,7 @@ mod build_graph_tests { .collect(); let data = MatrixView::try_from(&values[..], points, dimensions).unwrap(); let graph = graph_config(Metric::L2, degree); - let pool = pool(2); + let pool = thread_pool(2); let config = PiPNNConfig { c_max, c_min, @@ -559,11 +720,60 @@ mod build_graph_tests { }; let context = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool).unwrap(); - let actual = build_graph(data, &context) + let actual_graph = build_graph(data, &context) .unwrap_or_else(|error| panic!("randomized case {case} failed: {error}")); - assert_graph_invariants(&actual, points, degree); + assert_graph_invariants(&actual_graph, points, degree); } } + + #[test] + fn parallel_hash_prune_build_is_set_invariant() { + let points = 64; + let dimensions = 4; + let values = deterministic_point_values(points, dimensions); + let data = MatrixView::try_from(values.as_slice(), points, dimensions).unwrap(); + let graph = graph_config(Metric::L2, 8); + let pool = thread_pool(4); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.25, + fanout: vec![3, 2], + leaf_k: 3, + replicas: 2, + }; + let hash_prune = HashPruneConfig { + num_hash_planes: 8, + l_max: 16, + final_prune: true, + }; + let build = || { + let context = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(hash_prune.clone()) + .unwrap(); + build_graph(data, &context).unwrap() + }; + + let first = build(); + let second = build(); + let canonicalize = |graph: &[crate::graph::AdjacencyList]| { + graph + .iter() + .map(|row| { + let mut ids = row.to_vec(); + ids.sort_unstable(); + ids + }) + .collect::>() + }; + + // Parallel finalization can order equal candidates differently. Compare + // the retained neighbor sets. + assert_eq!(canonicalize(&first), canonicalize(&second)); + assert_graph_invariants(&first, points, 8); + assert!(first.iter().any(|row| !row.is_empty())); + } } #[cfg(test)] #[allow( @@ -572,9 +782,10 @@ mod build_graph_tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod config_tests { - use super::{PiPNNBuildContext, PiPNNConfig}; + use super::{HashPruneConfig, PiPNNBuildContext, PiPNNConfig}; use crate::graph::config::{self, MaxDegree}; use diskann_vector::distance::Metric; + use rstest::rstest; fn pipnn_config() -> PiPNNConfig { PiPNNConfig { @@ -588,93 +799,147 @@ mod config_tests { } fn graph_config(metric: Metric, alpha: f32) -> crate::graph::Config { - config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { + graph_config_with_degree(metric, alpha, 64) + } + + fn graph_config_with_degree(metric: Metric, alpha: f32, degree: usize) -> crate::graph::Config { + config::Builder::new_with(degree, MaxDegree::same(), 72, metric.into(), |builder| { builder.alpha(alpha); }) .build() .unwrap() } - fn pool() -> rayon::ThreadPool { + fn two_thread_pool() -> rayon::ThreadPool { rayon::ThreadPoolBuilder::new() .num_threads(2) .build() .unwrap() } - #[test] - fn rejects_each_invalid_algorithm_parameter() { + #[rstest] + #[case::zero_c_max(PiPNNConfig { c_max: 0, ..pipnn_config() })] + #[case::zero_c_min(PiPNNConfig { c_min: 0, ..pipnn_config() })] + #[case::c_min_above_c_max(PiPNNConfig { c_min: 513, ..pipnn_config() })] + #[case::zero_sampling_probability(PiPNNConfig { p_samp: 0.0, ..pipnn_config() })] + #[case::negative_sampling_probability(PiPNNConfig { p_samp: -0.01, ..pipnn_config() })] + #[case::sampling_probability_above_one(PiPNNConfig { p_samp: 1.01, ..pipnn_config() })] + #[case::nan_sampling_probability(PiPNNConfig { p_samp: f64::NAN, ..pipnn_config() })] + #[case::empty_fanout(PiPNNConfig { fanout: Vec::new(), ..pipnn_config() })] + #[case::zero_later_fanout(PiPNNConfig { fanout: vec![1, 0], ..pipnn_config() })] + #[case::zero_leaf_k(PiPNNConfig { leaf_k: 0, ..pipnn_config() })] + #[case::zero_replicas(PiPNNConfig { replicas: 0, ..pipnn_config() })] + fn invalid_algorithm_parameter_is_rejected(#[case] invalid_config: PiPNNConfig) { let graph = graph_config(Metric::L2, 1.2); - let pool = pool(); - let mut cases = [ - PiPNNConfig { - c_max: 0, - ..pipnn_config() - }, - PiPNNConfig { - c_min: 0, - ..pipnn_config() - }, - PiPNNConfig { - c_min: 513, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: 0.0, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: -0.01, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: 1.01, - ..pipnn_config() - }, - PiPNNConfig { - p_samp: f64::NAN, - ..pipnn_config() - }, - PiPNNConfig { - fanout: Vec::new(), - ..pipnn_config() - }, - PiPNNConfig { - fanout: vec![1, 0], - ..pipnn_config() - }, - PiPNNConfig { - leaf_k: 0, - ..pipnn_config() - }, - PiPNNConfig { - replicas: 0, - ..pipnn_config() - }, - ]; + let pool = two_thread_pool(); - for config in &mut cases { - PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) - .expect_err("invalid PiPNN config must be rejected"); - } + PiPNNBuildContext::new(invalid_config, &graph, Metric::L2, &pool) + .expect_err("invalid PiPNN config must be rejected"); } #[test] - fn rejects_graph_policy_for_a_different_metric() { + fn graph_policy_for_a_different_metric_is_rejected() { let graph = graph_config(Metric::InnerProduct, 1.2); - let pool = pool(); + let pool = two_thread_pool(); let error = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap_err(); assert!(error.to_string().contains("prune kind")); } + #[rstest] + #[case::below_one(0.9)] + #[case::nan(f32::NAN)] + #[case::infinity(f32::INFINITY)] + fn build_context_accepts_alpha_allowed_by_graph_config(#[case] alpha: f32) { + let pool = two_thread_pool(); + let graph = graph_config(Metric::L2, alpha); + + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + } + + #[rstest] + #[case::zero_hash_planes(HashPruneConfig { + num_hash_planes: 0, + l_max: 64, + final_prune: true, + })] + #[case::too_many_hash_planes(HashPruneConfig { + num_hash_planes: 17, + l_max: 64, + final_prune: true, + })] + #[case::zero_l_max(HashPruneConfig { + num_hash_planes: 8, + l_max: 0, + final_prune: true, + })] + #[case::l_max_above_storage_limit(HashPruneConfig { + num_hash_planes: 8, + l_max: 256, + final_prune: true, + })] + fn invalid_hash_prune_parameter_is_rejected(#[case] invalid_config: HashPruneConfig) { + let graph = graph_config(Metric::L2, 1.2); + let pool = two_thread_pool(); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + + assert!(context.with_hash_prune(invalid_config).is_err()); + } + #[test] - fn does_not_add_alpha_validation_beyond_graph_config() { - let pool = pool(); - for alpha in [0.9, f32::NAN, f32::INFINITY] { - let graph = graph_config(Metric::L2, alpha); - PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); - } + fn candidate_capacity_below_graph_degree_is_rejected() { + let graph = graph_config(Metric::L2, 1.2); + let pool = two_thread_pool(); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + let below_degree_capacity = HashPruneConfig { + num_hash_planes: 8, + l_max: 63, + final_prune: true, + }; + + assert!(context.with_hash_prune(below_degree_capacity).is_err()); + } + + #[test] + fn candidate_capacity_equal_to_graph_degree_is_accepted() { + let graph = graph_config(Metric::L2, 1.2); + let pool = two_thread_pool(); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + let equal_degree_capacity = HashPruneConfig { + num_hash_planes: 8, + l_max: 64, + final_prune: true, + }; + + context.with_hash_prune(equal_degree_capacity).unwrap(); + } + + #[test] + fn hash_bucket_capacity_equal_to_graph_degree_is_accepted() { + let pool = two_thread_pool(); + let graph = graph_config_with_degree(Metric::L2, 1.2, 2); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + let two_hash_buckets = HashPruneConfig { + num_hash_planes: 1, + l_max: 64, + final_prune: true, + }; + + context.with_hash_prune(two_hash_buckets).unwrap(); + } + + #[test] + fn hash_bucket_capacity_below_graph_degree_is_rejected() { + let pool = two_thread_pool(); + let graph = graph_config_with_degree(Metric::L2, 1.2, 3); + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + let two_hash_buckets = HashPruneConfig { + num_hash_planes: 1, + l_max: 64, + final_prune: true, + }; + + assert!(context.with_hash_prune(two_hash_buckets).is_err()); } }