diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 924bc53a4..4b63fabea 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -71,6 +71,7 @@ jobs: - test-workspace - test-workspace-features - coverage + - miri-pipnn - vectorset-clippy - vectorset-fmt - vectorset-build @@ -95,6 +96,35 @@ jobs: steps: - run: exit 0 + miri-pipnn: + needs: basics + name: strict-provenance PiPNN kernels + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install Rust nightly with Miri + run: rustup toolchain install nightly --component miri + + - uses: Swatinem/rust-cache@v2 + + - name: Run HashPrune pointer boundaries + env: + MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance + run: | + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::hash_prune::tests::find_hash_handles_padded_boundaries_and_all_bit_patterns \ + -- --exact + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::hash_prune::tests::relative_hash_matches_numeric_reference \ + -- --exact + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::hash_prune::tests::full_reservoir_evicts_the_farthest_candidate \ + -- --exact + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::hash_prune::tests::hot_slot_contention_serializes_state_mutation \ + -- --exact + fmt: name: format check runs-on: ubuntu-latest diff --git a/Cargo.lock b/Cargo.lock index dfa8a99ab..5ada406ab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -445,9 +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", "serde", diff --git a/diskann-benchmark/src/index/build.rs b/diskann-benchmark/src/index/build.rs index 39d63f67c..59c108b0e 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 7b7e13e55..cb0502076 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 488ea8ae7..bc3f3f4d6 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 f2ec56a0e..9f06f5fa0 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 a8b401752..8adc7a14f 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 a7e343fb5..d2a26bba4 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 27f4c124a..c3f304664 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 5d9e6c368..1704f93cf 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 e88dc12c9..009b8da00 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 bf8a1a369..7506f84e3 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 0a1ca5015..58f7ba176 100644 --- a/diskann-wide/src/doubled.rs +++ b/diskann-wide/src/doubled.rs @@ -435,6 +435,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 697cf1eb3..08bd0dd27 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -22,7 +22,9 @@ half = { workspace = true, features = ["bytemuck", "num-traits"] } # while other crates use default-features = true. Keeping version 0.16.0 consistent. hashbrown = { version = "0.16.0", default-features = false, features = ["default-hasher"] } num-traits.workspace = true +parking_lot = { version = "0.12", 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"] } @@ -33,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 @@ -59,7 +64,14 @@ panic = "warn" default = ["tracing"] # Enable PiPNN batch graph construction. -pipnn = ["dep:diskann-linalg", "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 8071001ac..1543afe97 100644 --- a/diskann/src/graph/adjacencylist.rs +++ b/diskann/src/graph/adjacencylist.rs @@ -135,6 +135,18 @@ where } } + /// Take ownership of a vector that contains unique items. + /// + /// The caller must supply unique items. Debug builds check this condition. + 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 +679,16 @@ mod tests { } } + #[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 000000000..c26c82e5c --- /dev/null +++ b/diskann/src/graph/pipnn/bf16.rs @@ -0,0 +1,95 @@ +/* + * 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::*; + + #[test] + fn roundtrip_bf16_exact_values() { + // bf16 has 7 mantissa bits. Pick f32 values whose lower 16 mantissa + // bits are zero so the truncation is lossless. + for &x in &[0.0_f32, 1.0, 2.0, 0.5, 0.25, 4.0, -1.0, -0.5] { + let back = bf16_to_f32(f32_to_bf16(x)); + assert_eq!(back, x, "exact bf16 roundtrip failed for {}", x); + } + } + + #[test] + fn truncation_is_within_relative_tolerance() { + // A non-exact finite value has at most about 2^-7 relative truncation error. + use std::f32::consts::{E, PI}; + for &x in &[1e-10_f32, 1e10, PI, E] { + let back = bf16_to_f32(f32_to_bf16(x)); + let rel = ((back - x) / x).abs(); + assert!(rel <= 0.01, "{} → {}: rel error {} > 1%", x, back, rel); + } + } + + #[test] + fn special_values_preserve_their_upper_bits() { + for value in [ + 0.0_f32, + -0.0, + f32::INFINITY, + f32::NEG_INFINITY, + f32::NAN, + f32::from_bits(0xFFC1_2345), + f32::MIN_POSITIVE, + -f32::MIN_POSITIVE, + ] { + let expected = value.to_bits() & 0xFFFF_0000; + assert_eq!( + bf16_to_f32(f32_to_bf16(value)).to_bits(), + expected, + "value_bits={:08x}", + value.to_bits() + ); + } + } + + #[test] + fn finite_truncation_moves_toward_zero() { + for value in [f32::MIN_POSITIVE, 0.1, 1.1, std::f32::consts::PI, f32::MAX] { + let positive = bf16_to_f32(f32_to_bf16(value)); + let negative = bf16_to_f32(f32_to_bf16(-value)); + assert!((0.0..=value).contains(&positive), "value={value}"); + assert!((-value..=0.0).contains(&negative), "value={value}"); + } + } + + #[test] + fn ordering_preserved_for_non_negative() { + // 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 000000000..cba691aa1 --- /dev/null +++ b/diskann/src/graph/pipnn/hash_prune.rs @@ -0,0 +1,1588 @@ +/* + * 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}; +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, SIMDPartialOrd, SIMDVector, + arch::{self, Dispatched1, FTarget1, Target}, + lifetime::As, +}; +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 + } + + #[inline] + #[allow(dead_code)] + fn bytes(&self) -> usize { + self.len * std::mem::size_of::() + } +} + +#[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 { + let mut values = Vec::new(); + values + .try_reserve_exact(len) + .map_err(ANNError::new) + .map_err(|error| error.context(format!("reserving {len} HashPrune slab elements")))?; + values.resize_with(len, T::default); + Ok(Self(values)) + } + #[inline] + fn as_ptr(&self) -> *const T { + self.0.as_ptr() + } + #[inline] + fn bytes(&self) -> usize { + self.0.len() * std::mem::size_of::() + } +} + +/// 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::>() + } +} + +fn select_find_hash() -> FindHash { + arch::dispatch(SelectFindHash) +} + +fn select_relative_hash() -> RelativeHash { + arch::dispatch(SelectRelativeHash) +} + +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: Architecture, + A::f32x16: SIMDPartialOrd + std::ops::Sub, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + 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: SIMDVector + SIMDPartialOrd + std::ops::Sub, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + // SAFETY: `src` and `dst` each contain `len` values. Construction checks + // `len <= 16 <= F::LANES`. The masked load does not read inactive lanes. + let dst = unsafe { F::load_simd_first(arch, args.dst, args.len) }; + // SAFETY: `dst` and `src` have the same checked length. + let src = unsafe { F::load_simd_first(arch, args.src, args.len) }; + let bits = u64::from( + (dst - src) + .ge_simd(F::splat(arch, 0.0)) + .bitmask() + .to_underlying(), + ); + let active = ((1_u32 << args.len) - 1) as u16; + bits as u16 & active +} + +/// 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 { + if !(1..=MAX_RESERVOIR_LEN).contains(&l_max) { + return Err(super::config_error(format!( + "HashPrune l_max ({l_max}) must be in 1..={MAX_RESERVOIR_LEN}" + ))); + } + + 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 mut states: Vec = Vec::new(); + states + .try_reserve_exact(npoints) + .map_err(ANNError::new) + .map_err(|error| error.context(format!("reserving {npoints} HashPrune reservoirs")))?; + for _ in 0..npoints { + states.push(LockedReservoirState::new()); + } + + // 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: select_find_hash(), + relative_hash: select_relative_hash(), + }) + } + + /// Lock one source point's reservoir while `f` reads or updates it. + /// + /// RAII unlocks the point when the closure exits. + /// + /// Returns an error when `idx` is outside the reservoir array or its row + /// offset overflows `usize`. + #[inline(always)] + fn with_locked_reservoir( + &self, + idx: usize, + f: impl FnOnce(&mut ReservoirState, ReservoirRows) -> R, + ) -> ANNResult { + let slot = self.states.get(idx).ok_or_else(|| { + ANNError::message(format!( + "HashPrune point ID {idx} is outside {} reservoirs", + self.states.len() + )) + })?; + let off = idx.checked_mul(self.row_stride).ok_or_else(|| { + ANNError::message(format!( + "HashPrune row offset {idx} x {} overflows usize", + 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, + } + }; + Ok(slot.with_locked_state(|state| f(state, rows))) + } + + /// Merge one leaf's CSR edges into 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. + /// + /// Returns an error for an invalid CSR shape, point ID, or local target. + pub(crate) fn add_leaf_edges( + &self, + point_ids: &[u32], + edge_offsets: &[u32], + edges: &[(u32, f32)], + sketch_scratch: &mut Vec, + ) -> ANNResult<()> { + let n = point_ids.len(); + let expected_offsets = n.checked_add(1).ok_or_else(|| { + ANNError::message(format!("HashPrune point count {n} overflows usize")) + })?; + if edge_offsets.len() != expected_offsets { + return Err(ANNError::message(format!( + "HashPrune expected {expected_offsets} edge offsets, got {}", + edge_offsets.len() + ))); + } + if edges.is_empty() { + return Ok(()); + } + + let m = self.sketches.num_planes(); + let l_max = self.l_max as u8; + let sketch_len = n.checked_mul(m).ok_or_else(|| { + ANNError::message(format!("HashPrune sketch shape {n} x {m} overflows usize")) + })?; + 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 || end > edges.len() { + return Err(ANNError::message(format!( + "HashPrune CSR range {start}..{end} is outside {} edges", + edges.len() + ))); + } + 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, so this raw range is the + // complete padded hash segment for that point. Raw prefetch avoids + // creating a shared slice while another worker mutates the segment. + 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| -> ANNResult<()> { + for &(dst_local, dist) in &edges[start..end] { + let dst_index = dst_local as usize; + let global_dst = *point_ids.get(dst_index).ok_or_else(|| { + ANNError::message(format!( + "HashPrune local target {dst_local} is outside {n} leaf points" + )) + })?; + 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: `with_locked_reservoir` holds this source's lock for the + // closure and supplies its three exact reservoir rows. + // `l_max` was validated at construction and `insert_reservoir_edge` + // maintains initialized entries through `state.len`. + unsafe { + insert_reservoir_edge( + state, + rows, + hash, + global_dst, + dist, + l_max, + self.find_hash, + ) + }; + } + Ok(()) + })??; + } + Ok(()) + } + + fn gather_sketches(&self, indices: &[u32], out: &mut [f32]) -> ANNResult<()> { + let m = self.sketches.num_planes(); + let expected = indices.len().checked_mul(m).ok_or_else(|| { + ANNError::message(format!( + "HashPrune sketch shape {} x {m} overflows usize", + indices.len() + )) + })?; + if out.len() != expected { + return Err(ANNError::message(format!( + "HashPrune expected {expected} gathered sketch values, got {}", + out.len() + ))); + } + let src = self.sketches.sketches(); + for (i, &idx) in indices.iter().enumerate() { + let start = (idx as usize).checked_mul(m).ok_or_else(|| { + ANNError::message(format!( + "HashPrune sketch offset {idx} x {m} overflows usize" + )) + })?; + let end = start.checked_add(m).ok_or_else(|| { + ANNError::message(format!("HashPrune sketch row {idx} overflows usize")) + })?; + let source = src.get(start..end).ok_or_else(|| { + ANNError::message(format!("HashPrune point ID {idx} has no sketch row")) + })?; + let output_start = i * m; + out[output_start..output_start + m].copy_from_slice(source); + } + Ok(()) + } + + /// 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); + (0..states.len()) + .into_par_iter() + .map_init(Vec::new, |scratch, i| { + 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) + }) + .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); + (0..states.len()) + .into_par_iter() + .map(|i| { + 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) + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn 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 Reservoir { + state: ReservoirState, + hashes: Vec, + distances: Vec, + neighbors: Vec, + row_stride: usize, + l_max: u8, + } + + impl Reservoir { + 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, + select_find_hash(), + ) + } + } + + 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) + }; + }) + .unwrap(); + } + + fn assert_sketch_source_type_matches_f32( + label: &str, + convert: impl Fn(u8) -> T, + reference: impl Fn(u8) -> f32, + ) where + T: VectorRepr + Send + Sync, + { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let points = 5; + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| ((index * 7 + index / dimensions * 3) % 23) 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 [1, 8, 16] { + let (actual, expected) = pool.install(|| { + ( + 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(), + ) + }); + 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(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 = select_relative_hash(); + + 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 = 0u16; + for j in 0..m { + let diff: f32 = dst[j] - src[j]; + expected |= ((diff >= 0.0) as u16) << j; + } + + let actual = dispatched.call(RelativeHashArgs { + src: src.as_ptr(), + dst: dst.as_ptr(), + len: m, + }); + assert_eq!(actual, expected, "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!( + select_relative_hash().call(RelativeHashArgs { + src: src.as_ptr(), + dst: dst.as_ptr(), + len: dst.len(), + }), + 0b0011 + ); + } + + #[test] + fn find_hash_handles_padded_boundaries_and_all_bit_patterns() { + let dispatched = select_find_hash(); + + 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_and_reports_its_bytes() { + let slab = MmapSlab::::new_zeroed(4).unwrap(); + assert_eq!(slab.bytes(), 4 * std::mem::size_of::()); + 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 accepts_structural_l_max_boundaries() { + let data = [0.0_f32]; + let low = hash_prune(&data, 1, 1, 1, 1).unwrap(); + assert_eq!(low.l_max, 1); + assert_eq!(low.row_stride, 32); + + let high = hash_prune(&data, 1, 1, 1, MAX_RESERVOIR_LEN).unwrap(); + assert_eq!(high.l_max, MAX_RESERVOIR_LEN); + assert_eq!(high.row_stride, 256); + } + + #[test] + fn rejects_l_max_outside_structural_boundaries() { + for l_max in [0, MAX_RESERVOIR_LEN + 1] { + let result = hash_prune(&[0.0_f32], 1, 1, 1, l_max); + let error = match result { + Ok(_) => panic!("l_max={l_max} must be rejected"), + Err(error) => error, + }; + assert!(format!("{error:?}").contains(&format!("l_max ({l_max})"))); + } + } + + #[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 = hash_prune(&data, 4, 2, 8, 8).unwrap(); + let reference = 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) + .unwrap(); + 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 = canonicalize(batched.into_candidate_lists()); + let expected = canonicalize(reference.into_candidate_lists()); + + assert_eq!(actual, expected); + assert!(actual.iter().all(|candidates| !candidates.is_empty())); + } + + #[test] + fn leaf_edges_reject_invalid_csr_and_point_ids() { + let data = [0.0_f32, 1.0]; + let hp = hash_prune(&data, 2, 1, 1, 2).unwrap(); + let mut scratch = Vec::new(); + + assert!(hp.add_leaf_edges(&[0], &[0], &[], &mut scratch).is_err()); + assert!( + hp.add_leaf_edges(&[2], &[0, 1], &[(0, 1.0)], &mut scratch) + .is_err() + ); + assert!( + hp.add_leaf_edges(&[0, 1], &[0, 1, 1], &[(2, 1.0)], &mut scratch) + .is_err() + ); + assert!( + hp.add_leaf_edges(&[0, 1], &[0, 1, 0], &[(1, 1.0)], &mut scratch) + .is_err() + ); + } + + #[test] + fn leaf_edges_grow_then_reuse_sketch_scratch() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let hp = hash_prune(&data, 4, 1, 8, 4).unwrap(); + let mut scratch = vec![99.0; 1]; + + hp.add_leaf_edges(&[0, 1], &[0, 1, 2], &[(1, 1.0), (0, 1.0)], &mut scratch) + .unwrap(); + assert_eq!(scratch.len(), 16); + let capacity = scratch.capacity(); + + hp.add_leaf_edges(&[2, 3], &[0, 1, 2], &[(1, 1.0), (0, 1.0)], &mut scratch) + .unwrap(); + assert_eq!(scratch.len(), 16); + assert_eq!(scratch.capacity(), capacity); + + hp.add_leaf_edges(&[0, 1], &[0, 0, 0], &[], &mut scratch) + .unwrap(); + assert_eq!(scratch.len(), 16); + assert_eq!(scratch.capacity(), capacity); + assert!( + hp.into_candidate_lists() + .iter() + .all(|candidates| candidates.len() == 1) + ); + } + + // Reservoir replacement and ordering policy. + + #[test] + fn full_reservoir_evicts_the_farthest_candidate() { + let mut reservoir = Reservoir::new(3); + assert!(reservoir.is_empty()); + + assert!(reservoir.insert(0, 1, 1.0)); + assert!(reservoir.insert(1, 2, 2.0)); + assert!(reservoir.insert(2, 3, 3.0)); + assert!(reservoir.insert(3, 4, 0.5)); + + assert_eq!(reservoir.len(), 3); + assert_eq!(reservoir.neighbors(), [(4, 0.5), (1, 1.0), (2, 2.0)]); + } + + #[test] + fn same_hash_keeps_only_the_closest_candidate() { + let mut reservoir = Reservoir::new(5); + + assert!(reservoir.insert(0, 1, 3.0)); + assert!(reservoir.insert(0, 2, 2.0)); + assert!(reservoir.insert(0, 3, 1.0)); + assert!(!reservoir.insert(0, 4, 5.0)); + + assert_eq!(reservoir.len(), 1); + assert_eq!(reservoir.neighbors(), [(3, 1.0)]); + } + + #[test] + fn equal_distances_are_ordered_by_neighbor_id() { + let mut res = Reservoir::new(5); + res.insert(0, 1, 1.0); + res.insert(1, 2, 1.0); + res.insert(2, 3, 1.0); + assert_eq!(res.len(), 3); + assert_eq!(res.neighbors(), [(1, 1.0), (2, 1.0), (3, 1.0)]); + } + + #[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 = Reservoir::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 = Reservoir::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] + #[allow(clippy::disallowed_methods)] + fn parallel_insertion_matches_serial_neighbor_lists() { + use rayon::prelude::*; + + let data = vec![0.0f32; 100 * 4]; + let parallel = hash_prune(&data, 100, 4, 4, 10).unwrap(); + let serial = hash_prune(&data, 100, 4, 4, 10).unwrap(); + + (0..50).into_par_iter().for_each(|source| { + add_edge(¶llel, source, (source + 1) % 100, 1.0); + add_edge(¶llel, (source + 1) % 100, source, 1.0); + }); + for source in 0..50 { + add_edge(&serial, source, (source + 1) % 100, 1.0); + add_edge(&serial, (source + 1) % 100, source, 1.0); + } + + assert_eq!(parallel.into_nearest_lists(5), serial.into_nearest_lists(5)); + } + + #[test] + fn extraction_returns_full_candidates_and_truncates_to_nearest_degree() { + #[rustfmt::skip] + let data = [ + 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 full = hash_prune(&data, 8, 2, 16, 10).unwrap(); + let nearest = hash_prune(&data, 8, 2, 16, 10).unwrap(); + for target in 1..8 { + add_edge(&full, 0, target, target as f32); + add_edge(&nearest, 0, target, target as f32); + } + + let mut full_ids = full.into_candidate_lists()[0].to_vec(); + full_ids.sort_unstable(); + assert_eq!(full_ids, (1..8).collect::>()); + assert_eq!(&*nearest.into_nearest_lists(2)[0], &[1, 2]); + } + + #[test] + fn farthest_cache_updates_after_repeated_evictions() { + let mut reservoir = Reservoir::new(3); + reservoir.insert(0, 10, 5.0); + reservoir.insert(1, 11, 4.0); + reservoir.insert(2, 12, 3.0); + assert!(reservoir.insert(3, 13, 2.0)); + assert!(reservoir.insert(4, 14, 1.0)); + + assert_eq!(reservoir.neighbors(), [(14, 1.0), (13, 2.0), (12, 3.0)]); + } + + #[test] + fn sorted_extraction_handles_an_early_farthest_slot() { + let mut reservoir = Reservoir::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 25d7df817..34c9a8de5 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -12,11 +12,12 @@ //! 2. Compute the lower triangle of `A · Aᵀ`. //! 3. Select local neighbors for both points of each pair. //! 4. Convert local positions to global point IDs. -//! 5. Add both edge directions to global candidate lists. +//! 5. 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 std::{collections::TryReserveError, sync::Mutex}; @@ -32,7 +33,7 @@ use super::{ }, }; -/// 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 build requires at least one dimension")] @@ -89,12 +90,21 @@ pub(crate) enum LeafBuildError { InvalidLocalTarget { target: u32, points: usize }, #[error("candidate list for point {point} is poisoned")] PoisonedCandidateList { point: u32 }, + #[error("leaf {leaf} produced too many directed edges")] + TooManyEdges { leaf: usize }, + #[error("HashPrune rejected the edge data for leaf {leaf}")] + HashPrune { + leaf: usize, + #[source] + source: crate::ANNError, + }, } /// 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, @@ -102,6 +112,11 @@ struct LeafBuffers { 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 { @@ -152,6 +167,7 @@ impl LeafBuffers { neighbor_count, LeafNeighbor::default(), )?; + grow("leaf seen pairs", &mut self.seen_pairs, dot_count, false)?; Ok((leaf_k, neighbor_count)) } @@ -264,6 +280,66 @@ where 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: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, + T: VectorRepr + 'static, +{ + if data.ncols() == 0 { + return Err(LeafBuildError::EmptyDimensions); + } + if data.nrows() > u32::MAX as usize { + return Err(LeafBuildError::TooManyPoints(data.nrows())); + } + + leaves.par_iter().enumerate().try_for_each_init( + LeafBuffers::default, + |buffers, (leaf, point_ids)| { + let leaf_k = select_leaf_neighbors::( + arch, + data, + leaf, + point_ids, + requested_k, + buffers, + )?; + let point_count = point_ids.len(); + 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, + &buffers.edges[..edge_count], + &mut buffers.sketch_scratch, + ) + .map_err(|source| LeafBuildError::HashPrune { leaf, source }) + }, + ) +} + /// Add one leaf's symmetric neighbors to the direct candidate lists. /// /// The function rejects empty, duplicate, unsorted, or out-of-range point IDs. @@ -278,6 +354,42 @@ fn add_direct_leaf_candidates( buffers: &mut LeafBuffers, candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> +where + A: Architecture, + A::f32x16: std::ops::Div, + ::Mask: SIMDSelect, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, + M: KernelMetric, + T: VectorRepr + 'static, +{ + let leaf_k = + select_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()]) +} + +/// Select local nearest neighbors for one leaf. +/// +/// The function checks point IDs, converts their vectors to `f32`, and computes +/// the lower Gram matrix. It writes leaf-local neighbors to `buffers.neighbors` +/// and returns the effective neighbor count. +fn select_leaf_neighbors( + arch: A, + data: MatrixView<'_, T>, + leaf: usize, + point_ids: &[u32], + requested_k: usize, + buffers: &mut LeafBuffers, +) -> Result where A: Architecture, A::f32x16: std::ops::Div, @@ -310,7 +422,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(); @@ -353,15 +465,7 @@ where })?; nearest_neighbors::(arch, dots, output, &mut buffers.kernel_workspace) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - - buffers.prepare_local_adjacency(point_ids.len())?; - add_symmetric_neighbors( - point_ids, - leaf_k, - &buffers.neighbors[..neighbor_value_count], - &mut buffers.local_adjacency[..point_ids.len()], - )?; - candidates.add_leaf(point_ids, &buffers.local_adjacency[..point_ids.len()]) + Ok(leaf_k) } fn add_symmetric_neighbors( @@ -389,6 +493,128 @@ fn add_symmetric_neighbors( Ok(()) } +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`. +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(); + if leaf_k == 0 { + resize("leaf edge offsets", offsets, point_count + 1, 0)?; + offsets.fill(0); + edges.clear(); + return Ok(0); + } + seen.fill(false); + resize("leaf edge offsets", offsets, point_count + 1, 0)?; + offsets.fill(0); + + for (source, neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for neighbor in neighbors { + let target = neighbor.target as usize; + if target >= point_count { + return Err(LeafBuildError::InvalidLocalTarget { + target: neighbor.target, + points: point_count, + }); + } + 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; + resize("leaf edges", edges, edge_count, (0, 0.0))?; + resize("leaf edge cursor", cursor, point_count, 0)?; + cursor.copy_from_slice(&offsets[..point_count]); + seen.fill(false); + + for (source, neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for neighbor in neighbors { + let target = neighbor.target as usize; + write_directed_edge( + point_count, + source, + target, + neighbor.distance, + seen, + edges, + cursor, + ); + write_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_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 = true; + let edge_slot = cursor[source] as usize; + edges[edge_slot] = (target as u32, distance); + cursor[source] += 1; + } +} + fn grow( buffer: &'static str, values: &mut Vec, @@ -442,9 +668,10 @@ mod tests { use half::f16; use std::collections::BTreeSet; + use super::super::leaf_kernel::LeafNeighbor; use super::{ - DirectCandidates, LeafBuffers, LeafBuildError, add_symmetric_neighbors, allocation_error, - build_leaf_candidates, + DirectCandidates, EdgeBuffers, LeafBuffers, LeafBuildError, add_symmetric_neighbors, + allocation_error, build_leaf_candidates, build_symmetric_edge_csr, }; fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { @@ -888,4 +1115,117 @@ mod tests { [vec![1], vec![0]] ); } + + #[test] + fn symmetric_edge_csr_matches_expected_adjacency() { + let point_ids = [10, 20, 30]; + let neighbors = [ + LeafNeighbor::new(1, 1.0), + LeafNeighbor::new(2, 2.0), + LeafNeighbor::new(1, 1.5), + ]; + let mut seen = vec![false; 9]; + 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, 4); + assert_eq!(offsets, [0, 1, 3, 4]); + assert_eq!(edges, [(1, 1.0), (0, 1.0), (2, 2.0), (1, 2.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)]); + } + #[test] + fn symmetric_edge_csr_rejects_out_of_range_local_targets() { + let mut seen = vec![false; 4]; + let mut offsets = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = Vec::new(); + let error = build_symmetric_edge_csr( + 7, + &[10, 20], + 1, + &[LeafNeighbor::new(2, 1.0), LeafNeighbor::new(0, 1.0)], + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap_err(); + + assert!(matches!( + error, + LeafBuildError::InvalidLocalTarget { + target: 2, + points: 2 + } + )); + } + #[test] + fn zero_k_edge_csr_has_empty_adjacency() { + let point_ids = [10, 20, 30]; + 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, + 0, + &[], + 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!(edges.is_empty()); + } } diff --git a/diskann/src/graph/pipnn/lsh.rs b/diskann/src/graph/pipnn/lsh.rs new file mode 100644 index 000000000..6f93876cb --- /dev/null +++ b/diskann/src/graph/pipnn/lsh.rs @@ -0,0 +1,217 @@ +/* + * 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}; +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 { + if !(1..=MAX_PLANES).contains(&num_planes) { + return Err(ANNError::message(format!( + "num_planes ({num_planes}) must be in 1..={MAX_PLANES}" + ))); + } + 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 mut hyperplanes: Vec = Vec::new(); + hyperplanes + .try_reserve_exact(hyperplane_len) + .map_err(ANNError::new)?; + hyperplanes.resize_with(hyperplane_len, || StandardNormal.sample(&mut rng)); + + let mut sketches = Vec::new(); + sketches + .try_reserve_exact(sketch_len) + .map_err(ANNError::new)?; + sketches.resize(sketch_len, 0.0f32); + + #[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)| { + if buffer.len() < ndims { + buffer + .try_reserve(ndims - buffer.len()) + .map_err(ANNError::new)?; + } + buffer.resize(ndims, 0.0); + T::as_f32_into(data.row(point), &mut buffer[..ndims]) + .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 * ndims..(plane_index + 1) * ndims]; + let mut dot = 0.0f32; + for dimension in 0..ndims { + // SAFETY: both slices have exactly `ndims` elements. + unsafe { + dot += + *buffer.get_unchecked(dimension) * *plane.get_unchecked(dimension); + } + } + *destination = dot; + } + Ok::<(), ANNError>(()) + })?; + + 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 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn build_pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + } + + fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + MatrixView::try_from(data, rows, columns).unwrap() + } + + #[test] + fn computes_expected_sketch_shape() { + let data = [1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; + let sketches = build_pool(2) + .install(|| LshSketches::try_new(view(&data, 3, 2), 4, 42)) + .unwrap(); + + assert_eq!(sketches.num_planes(), 4); + assert_eq!(sketches.sketches().len(), 12); + } + + #[test] + fn sketches_match_seeded_serial_hyperplane_reference() { + let npoints = 3; + let ndims = 4; + let planes = 5; + let data: Vec = (0..npoints * ndims) + .map(|value| value as f32 - 3.0) + .collect(); + + for seed in [42, 99] { + let actual = build_pool(2) + .install(|| LshSketches::try_new(view(&data, npoints, ndims), planes, seed)) + .unwrap(); + + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let hyperplanes: Vec = (0..planes * ndims) + .map(|_| StandardNormal.sample(&mut rng)) + .collect(); + let expected: 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(); + assert_eq!(actual.sketches(), expected, "seed={seed}"); + } + } + + #[test] + fn zero_points_produce_an_empty_sketch() { + let sketches = build_pool(2) + .install(|| LshSketches::try_new(view(&[] as &[f32], 0, 7), 4, 42)) + .unwrap(); + + assert_eq!(sketches.num_planes(), 4); + assert!(sketches.sketches().is_empty()); + } + + #[test] + fn zero_dimensions_produce_zero_dot_products() { + let sketches = build_pool(2) + .install(|| LshSketches::try_new(view(&[] as &[f32], 3, 0), 2, 42)) + .unwrap(); + + assert_eq!(sketches.sketches(), &[0.0; 6]); + } + + #[test] + fn rejects_shape_overflow() { + for data in [ + view(&[] as &[f32], 0, usize::MAX), + view(&[] as &[f32], usize::MAX, 0), + ] { + let error = + LshSketches::try_new(data, 2, 42).expect_err("overflowing LSH shape must fail"); + assert!(error.to_string().contains("overflows")); + } + } + + #[test] + fn rejects_plane_counts_outside_u16_capacity() { + for planes in [0, MAX_PLANES + 1] { + let error = LshSketches::try_new(view(&[0.0_f32], 1, 1), planes, 42).unwrap_err(); + assert!(error.to_string().contains("must be in 1..=16")); + } + } +} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index 4a6dabcae..b25b35811 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 @@ -33,9 +34,12 @@ mod kernel_metric; +mod bf16; mod finalization; +mod hash_prune; mod leaf_build; mod leaf_kernel; +mod lsh; mod partition_kernel; mod partitioning; @@ -114,6 +118,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> { @@ -121,6 +174,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> { @@ -144,8 +198,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`. @@ -239,8 +301,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>, @@ -256,16 +319,47 @@ 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.k) - .map_err(ANNError::new) - })?; - // Finalization consumes each candidate list. It reuses that list's allocation - // for the final adjacency when the graph policy permits it. - tracing::info_span!("pipnn.finalization") - .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, M::METRIC)) + 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.k) + .map_err(ANNError::new) + })?; + // Finalization consumes each candidate list. It reuses that list's + // allocation when the graph policy permits it. + tracing::info_span!("pipnn.finalization").in_scope(|| { + finalization::prune_overfull(data, candidates, context.graph, M::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.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, M::METRIC) + }) + } else { + Ok(hash_prune.into_nearest_lists(context.graph.pruned_degree().get())) + } + } + } } fn effective_metric(metric: Metric) -> Metric { @@ -317,7 +411,7 @@ mod tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod build_graph_tests { - use super::{PiPNNBuildContext, PiPNNConfig, build_graph, leaf_kernel}; + use super::{HashPruneConfig, PiPNNBuildContext, PiPNNConfig, build_graph, leaf_kernel}; use crate::graph::config::{self, MaxDegree}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; @@ -541,6 +635,57 @@ mod build_graph_tests { assert_graph_invariants(&actual, points, degree); } } + + #[test] + fn parallel_hash_prune_build_is_set_invariant() { + let points = 64; + let dimensions = 4; + let values: Vec = (0..points * dimensions) + .map(|value| ((value * 17 + 3) % 101) as f32) + .collect(); + let data = MatrixView::try_from(values.as_slice(), points, dimensions).unwrap(); + let graph = graph_config(Metric::L2, 8); + let pool = pool(4); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.25, + fanout: vec![3, 2], + k: 3, + replicas: 2, + }; + let 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( @@ -549,7 +694,7 @@ mod build_graph_tests { reason = "deterministic test fixture construction must abort on invalid setup" )] mod config_tests { - use super::{PiPNNBuildContext, PiPNNConfig, leaf_kernel}; + use super::{HashPruneConfig, PiPNNBuildContext, PiPNNConfig, leaf_kernel}; use crate::graph::config::{self, MaxDegree}; use diskann_vector::distance::Metric; @@ -565,7 +710,11 @@ 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() @@ -658,4 +807,65 @@ mod config_tests { PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); } } + + #[test] + fn rejects_invalid_hash_prune_parameters_and_accepts_valid_boundary() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + for config in [ + HashPruneConfig { + num_hash_planes: 0, + l_max: 64, + final_prune: true, + }, + HashPruneConfig { + num_hash_planes: 17, + l_max: 64, + final_prune: true, + }, + HashPruneConfig { + num_hash_planes: 8, + l_max: 0, + final_prune: true, + }, + HashPruneConfig { + num_hash_planes: 8, + l_max: 256, + final_prune: true, + }, + HashPruneConfig { + num_hash_planes: 8, + l_max: 63, + final_prune: true, + }, + ] { + let context = + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + assert!(context.with_hash_prune(config).is_err()); + } + + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(HashPruneConfig { + num_hash_planes: 8, + l_max: 64, + final_prune: true, + }) + .unwrap(); + } + #[test] + fn requires_hash_bucket_capacity_to_cover_graph_degree() { + let pool = pool(); + for (degree, accepted) in [(2, true), (3, false)] { + let graph = graph_config_with_degree(Metric::L2, 1.2, degree); + let result = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(HashPruneConfig { + num_hash_planes: 1, + l_max: 64, + final_prune: true, + }); + assert_eq!(result.is_ok(), accepted, "degree={degree}"); + } + } }