diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index c93553d97b..c9df12e669 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -145,3 +145,12 @@ jobs: cargo +nightly miri test --locked -p diskann --features pipnn --lib -- \ graph::pipnn::topk::tests::selection_returns_distinct_candidates_when_distances_tie \ graph::pipnn::topk::tests::a_single_pair_updates_exactly_its_two_endpoints::case_1_finite + + - name: PiPNN HashPrune pointer boundaries + env: + MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance + run: | + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::lsh::tests + cargo +nightly miri test --locked -p diskann --features pipnn --lib \ + graph::pipnn::hash_prune::tests diff --git a/Cargo.lock b/Cargo.lock index 37a6c99f4f..d2dce0e381 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -445,10 +445,12 @@ dependencies = [ "futures-util", "half", "hashbrown 0.16.1", + "libc", "num-traits", "parking_lot", "pin-project", "rand", + "rand_distr", "rayon", "relative-path 2.0.1", "rstest", diff --git a/diskann-benchmark/src/index/build.rs b/diskann-benchmark/src/index/build.rs index c8caee9fa2..1461c9f664 100644 --- a/diskann-benchmark/src/index/build.rs +++ b/diskann-benchmark/src/index/build.rs @@ -143,12 +143,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 09f93fe002..6452734a2d 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -79,8 +79,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(index_configuration.config.pruned_degree().get())?; + } } let pq_storage = PQStorage::new( @@ -184,8 +188,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 19e3aa7b88..779146a068 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, @@ -56,16 +56,19 @@ where let data = read_bin::(&mut builder.storage_provider.open_reader(&data_path)?)?; // The same seed controls leader sampling and the medoid sample below. - let config = PiPNNConfig { + let config = diskann::graph::pipnn::PiPNNConfig { seed: builder.index_configuration.random_seed, - ..config + ..parameters.into() }; - let context = PiPNNBuildContext::new( + let mut context = PiPNNBuildContext::new( config, &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 @@ -121,6 +124,7 @@ mod tests { fanout: vec![nz(10), nz(3)], k: nz(2), replicas: nz(1), + hash_prune: Some(crate::HashPruneParameters::default()), } } @@ -204,7 +208,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())); } @@ -218,7 +222,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( @@ -279,4 +283,34 @@ mod tests { assert!(format!("{error:?}").contains("p_samp (0) must be in (0, 1]")); } + + #[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 5b559e7a3b..57d6d547d8 100644 --- a/diskann-disk/src/build/configuration/build_algorithm.rs +++ b/diskann-disk/src/build/configuration/build_algorithm.rs @@ -30,6 +30,43 @@ pub struct PiPNNParameters { pub k: NonZeroUsize, /// Number of independent partition passes. pub replicas: NonZeroUsize, + /// 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")] @@ -47,6 +84,7 @@ impl Default for PiPNNParameters { fanout: FANOUT.to_vec(), k: K, replicas: NonZeroUsize::MIN, + hash_prune: None, } } } @@ -124,6 +162,15 @@ mod tests { ); assert_eq!(config.k.get(), 3); assert_eq!(config.replicas.get(), 1); + assert_eq!(config.hash_prune, None); + + let explicit: BuildAlgorithm = + serde_json::from_str(r#"{"algorithm":"PiPNN","hash_prune":{}}"#).unwrap(); + let BuildAlgorithm::PiPNN(explicit) = explicit else { + panic!("expected PiPNN"); + }; + assert_eq!(explicit.hash_prune, Some(HashPruneParameters::default())); + assert!( serde_json::from_str::(r#"{"algorithm":"PiPNN","l_max":72}"#).is_err() ); diff --git a/diskann-disk/src/build/configuration/disk_index_build_parameter.rs b/diskann-disk/src/build/configuration/disk_index_build_parameter.rs index a8b4017529..8adc7a14f8 100644 --- a/diskann-disk/src/build/configuration/disk_index_build_parameter.rs +++ b/diskann-disk/src/build/configuration/disk_index_build_parameter.rs @@ -198,9 +198,9 @@ impl DiskIndexBuildParameters { } #[cfg(feature = "pipnn")] - pub(crate) fn pipnn_config(&self) -> Option { + pub(crate) fn pipnn_parameters(&self) -> Option<&PiPNNParameters> { match &self.build_algorithm { - BuildAlgorithm::PiPNN(config) => Some(config.into()), + BuildAlgorithm::PiPNN(config) => Some(config), BuildAlgorithm::Vamana => None, } } diff --git a/diskann-disk/src/build/configuration/mod.rs b/diskann-disk/src/build/configuration/mod.rs index 4d407e340b..85e067f4ab 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 160de514d3..067af48b0e 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, SphericalBits, }; +#[cfg(feature = "pipnn")] +pub use configuration::{HashPruneParameters, PiPNNParameters}; diff --git a/diskann-disk/src/lib.rs b/diskann-disk/src/lib.rs index 284ac121f6..8b10b7915b 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, SphericalBits, }; +#[cfg(feature = "pipnn")] +pub use build::{HashPruneParameters, PiPNNParameters}; pub mod data_model; pub mod search; diff --git a/diskann-vector/src/lib.rs b/diskann-vector/src/lib.rs index e88dc12c9c..009b8da006 100644 --- a/diskann-vector/src/lib.rs +++ b/diskann-vector/src/lib.rs @@ -38,14 +38,17 @@ pub mod distance; pub mod norm; cfg_if::cfg_if! { - if #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))] { + // x86-64 guarantees SSE2; `_mm_prefetch` needs only SSE. + if #[cfg(target_arch = "x86_64")] { const CACHE_LINE_SIZE: usize = 64; #[inline(always)] unsafe fn prefetch_exactly(ptr: *const i8) { use std::arch::x86_64::*; for i in 0..N { - _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0); + // SAFETY: the caller guarantees that all `N` computed addresses are + // inside the allocation. + unsafe { _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0) }; } } @@ -56,7 +59,8 @@ cfg_if::cfg_if! { if CACHE_LINE_SIZE * i >= bytes { break; } - _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0); + // SAFETY: the loop uses only offsets below `bytes`. + unsafe { _mm_prefetch(ptr.add(i * CACHE_LINE_SIZE), _MM_HINT_T0) }; } } @@ -66,32 +70,48 @@ cfg_if::cfg_if! { pub fn prefetch_hint_max(vec: &[T]) { let vecsize = std::mem::size_of_val(vec); if vecsize >= MAX_CACHE_LINES * 64 { - // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated. + // SAFETY: the slice contains every address passed to prefetch. unsafe { prefetch_exactly::(vec.as_ptr().cast()) } } else { - // SAFETY: Pointer is in-bounds and use of the intrinsic is cfg gated. + // SAFETY: the slice covers `vecsize` bytes. unsafe { prefetch_at_most::(vec.as_ptr().cast(), vecsize) } } } + /// Prefetch a raw byte range without creating a slice. + /// + /// # Safety + /// + /// `ptr` must identify an allocation of at least `bytes` bytes. The allocation + /// must remain live for this call. The function creates no Rust reference. + /// The caller controls concurrent mutation of the range. + #[inline] + pub unsafe fn prefetch_hint_all_raw(ptr: *const u8, bytes: usize) { + use std::arch::x86_64::*; + + for offset in (0..bytes).step_by(CACHE_LINE_SIZE) { + // SAFETY: the caller guarantees the byte range, and `offset < bytes`. + unsafe { _mm_prefetch(ptr.add(offset).cast(), _MM_HINT_T0) }; + } + } + /// Prefetch the given vector in chunks of 64 bytes, which is a cache line size. /// The entire vector will be prefetched. #[inline] pub fn prefetch_hint_all(vec: &[T]) { - use std::arch::x86_64::*; - - let vecsize = std::mem::size_of_val(vec); - let num_prefetch_blocks = vecsize.div_ceil(64); - let vec_ptr = vec.as_ptr() as *const i8; - for d in 0..num_prefetch_blocks { - // SAFETY: Pointer is in-bounds and use of the intrinsic is gated by the - // `cfg`-guard on this function. - unsafe { - std::arch::x86_64::_mm_prefetch(vec_ptr.add(d * CACHE_LINE_SIZE), _MM_HINT_T0); - } - } } + // SAFETY: the slice remains live and covers exactly `size_of_val(vec)` bytes. + unsafe { prefetch_hint_all_raw(vec.as_ptr().cast(), std::mem::size_of_val(vec)) } + } } else { pub fn prefetch_hint_max(_vec: &[T]) {} + + /// Accept a raw prefetch range and do nothing. + /// + /// # Safety + /// + /// The pointer contract is the same as the x86-64 implementation. + pub unsafe fn prefetch_hint_all_raw(_ptr: *const u8, _bytes: usize) {} + pub fn prefetch_hint_all(_vec: &[T]) {} } } diff --git a/diskann-wide/src/arch/x86_64/v3/i16x16_.rs b/diskann-wide/src/arch/x86_64/v3/i16x16_.rs index bf8a1a3690..7506f84e35 100644 --- a/diskann-wide/src/arch/x86_64/v3/i16x16_.rs +++ b/diskann-wide/src/arch/x86_64/v3/i16x16_.rs @@ -16,7 +16,7 @@ use crate::{ v3::i16x8, }, }, - bitmask::BitMask, + bitmask::{BitMask, FromInt}, constant::Const, emulated::Emulated, helpers, @@ -89,16 +89,19 @@ impl SIMDMulAdd for i16x16 { impl SIMDPartialEq for i16x16 { #[inline(always)] fn eq_simd(self, other: Self) -> Self::Mask { - self.emulated() - .eq_simd(other.emulated()) - .as_arch(self.arch()) + // SAFETY: V3 includes AVX2 and BMI2. Each equal i16 lane contributes + // two adjacent movemask bits; pext keeps one bit per lane. + let bits = unsafe { + let bytes = _mm256_movemask_epi8(_mm256_cmpeq_epi16(self.0, other.0)) as u32; + _pext_u32(bytes, 0x5555_5555) as u16 + }; + BitMask::from_int(self.arch(), bits) } #[inline(always)] fn ne_simd(self, other: Self) -> Self::Mask { - self.emulated() - .ne_simd(other.emulated()) - .as_arch(self.arch()) + let equal = self.eq_simd(other); + BitMask::from_int(self.arch(), !equal.0) } } diff --git a/diskann-wide/src/doubled.rs b/diskann-wide/src/doubled.rs index 30d08e6cb1..f62f63cf57 100644 --- a/diskann-wide/src/doubled.rs +++ b/diskann-wide/src/doubled.rs @@ -433,6 +433,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 30ec5df32f..44ebfbe0c6 100644 --- a/diskann/Cargo.toml +++ b/diskann/Cargo.toml @@ -23,6 +23,7 @@ hashbrown = { version = "0.16.0", default-features = false, features = ["default num-traits.workspace = true parking_lot = { version = "0.12.5", optional = true } rand.workspace = true +rand_distr = { workspace = true, optional = true } rayon = { workspace = true, optional = true } thiserror.workspace = true tokio = { workspace = true, features = ["rt", "rt-multi-thread"] } @@ -34,6 +35,9 @@ diskann-wide = { workspace = true } dashmap = { workspace = true, optional = true } diskann-linalg = { 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 @@ -62,7 +66,14 @@ panic = "warn" default = ["tracing"] # Enable PiPNN batch graph construction. -pipnn = ["dep:diskann-linalg", "dep:parking_lot", "dep:rayon", "tracing"] +pipnn = [ + "dep:diskann-linalg", + "dep:libc", + "dep:parking_lot", + "dep:rand_distr", + "dep:rayon", + "tracing", +] # Enable "tracing" diagnostics. tracing = ["dep:tracing"] diff --git a/diskann/src/graph/adjacencylist.rs b/diskann/src/graph/adjacencylist.rs index 063ace9b2d..f5eef8b6a7 100644 --- a/diskann/src/graph/adjacencylist.rs +++ b/diskann/src/graph/adjacencylist.rs @@ -135,6 +135,19 @@ where } } + /// Take ownership of a vector that contains unique items. + /// + /// The caller must supply unique items. Debug builds check this condition. + #[cfg(feature = "pipnn")] + pub(crate) fn from_vec_trusted(edges: Vec) -> Self + where + I: ContainsSimd, + { + let list = Self { edges }; + list.debug_check_uniqueness(); + list + } + /// Resize the underlying storage to `capacity` elements and return a guard allowing /// full mutable access to the resized span. /// @@ -663,6 +676,17 @@ mod tests { } } + #[cfg(feature = "pipnn")] + #[test] + fn test_from_vec_trusted_preserves_order_and_allocation() { + let edges = vec![3_u32, 1, 2]; + let pointer = edges.as_ptr(); + let list = AdjacencyList::from_vec_trusted(edges); + + assert_eq!(&*list, &[3, 1, 2]); + assert_eq!(list.as_ptr(), pointer); + } + #[test] fn test_from_iter_untrusted() { let x = AdjacencyList::::from_iter_untrusted([]); diff --git a/diskann/src/graph/pipnn/bf16.rs b/diskann/src/graph/pipnn/bf16.rs new file mode 100644 index 0000000000..87fb78271f --- /dev/null +++ b/diskann/src/graph/pipnn/bf16.rs @@ -0,0 +1,105 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Lossy conversion between `f32` and bf16 storage. +//! +//! A bf16 value contains the upper 16 bits of an IEEE-754 `f32`. It keeps the +//! exponent and seven mantissa bits. For non-negative values, its `u16` bit order +//! matches `f32` numeric order. HashPrune applies a separate ordered-key transform +//! to signed distances. +//! +//! Conversion truncates the lower 16 bits. It does not round. It preserves sign, +//! infinity, signed zero, and the upper NaN payload bits. + +/// Convert `f32` → bf16 by truncating the lower 16 mantissa bits. +#[inline(always)] +pub(super) fn f32_to_bf16(v: f32) -> u16 { + (v.to_bits() >> 16) as u16 +} + +/// Reconstruct `f32` from a bf16 for conversion tests. +#[cfg(test)] +#[inline(always)] +fn bf16_to_f32(v: u16) -> f32 { + f32::from_bits((v as u32) << 16) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + fn bf16_roundtrip_preserves_exactly_representable_values( + #[values(0.0_f32, 1.0, 2.0, 0.5, 0.25, 4.0, -1.0, -0.5)] value: f32, + ) { + // Given: bf16 has seven mantissa bits, and these values have zero low mantissa bits. + let expected_value = value; + + // When + let actual_value = bf16_to_f32(f32_to_bf16(value)); + + // Then + assert_eq!(actual_value, expected_value); + } + + #[rstest] + fn bf16_truncation_keeps_relative_error_below_one_percent( + #[values(1e-10_f32, 1e10, std::f32::consts::PI, std::f32::consts::E)] value: f32, + ) { + // Given: bf16 truncation has at most about 2^-7 relative error for finite values. + let maximum_relative_error = 0.01; + + // When + let truncated = bf16_to_f32(f32_to_bf16(value)); + let actual_relative_error = ((truncated - value) / value).abs(); + + // Then + assert!(actual_relative_error <= maximum_relative_error); + } + + #[rstest] + #[case::positive_zero(0.0)] + #[case::negative_zero(-0.0)] + #[case::positive_infinity(f32::INFINITY)] + #[case::negative_infinity(f32::NEG_INFINITY)] + #[case::positive_nan(f32::NAN)] + #[case::negative_nan(f32::from_bits(0xFFC1_2345))] + #[case::minimum_positive(f32::MIN_POSITIVE)] + #[case::negative_minimum_positive(-f32::MIN_POSITIVE)] + fn bf16_conversion_preserves_the_upper_bits_of_special_values(#[case] value: f32) { + // Given + let expected_upper_bits = value.to_bits() & 0xFFFF_0000; + + // When + let actual_bits = bf16_to_f32(f32_to_bf16(value)).to_bits(); + + // Then + assert_eq!(actual_bits, expected_upper_bits); + } + + #[rstest] + fn bf16_truncation_moves_finite_values_toward_zero( + #[values(f32::MIN_POSITIVE, 0.1, 1.1, std::f32::consts::PI, f32::MAX)] value: f32, + ) { + // When + let positive_truncation = bf16_to_f32(f32_to_bf16(value)); + let negative_truncation = bf16_to_f32(f32_to_bf16(-value)); + + // Then + assert!((0.0..=value).contains(&positive_truncation)); + assert!((-value..=0.0).contains(&negative_truncation)); + } + + #[test] + fn bf16_encoding_preserves_non_negative_value_order() { + // For non-negative f32, bf16 (as u16) preserves ordering. + let xs: [f32; 6] = [0.0, 1e-10, 0.1, 1.0, 10.0, 1e10]; + let bs: Vec = xs.iter().map(|&x| f32_to_bf16(x)).collect(); + for w in bs.windows(2) { + assert!(w[0] <= w[1], "bf16 ordering broken: {} > {}", w[0], w[1]); + } + } +} diff --git a/diskann/src/graph/pipnn/conversion.rs b/diskann/src/graph/pipnn/conversion.rs index b8cbe1beb6..26827c6448 100644 --- a/diskann/src/graph/pipnn/conversion.rs +++ b/diskann/src/graph/pipnn/conversion.rs @@ -32,6 +32,19 @@ use diskann_wide::{ }; use half::f16; +/// Convert one row to `f32`, with the FP16 instructions of the running CPU +/// for `f16` input. Errors are those of `VectorRepr::as_f32_into`. +#[inline] +pub(super) fn as_f32_into(source: &[T], output: &mut [f32]) -> Result<(), T::Error> { + if TypeId::of::() == TypeId::of::() && source.len() == output.len() { + let source: &[f16] = bytemuck::cast_slice(source); + arch::dispatch2(SliceCast::::new(), output, source); + Ok(()) + } else { + T::as_f32_into(source, output) + } +} + /// Convert rows `ids` of `data` into consecutive rows of `output`. /// /// `f16` input selects the CPU's instructions once for the whole batch. @@ -155,3 +168,68 @@ mod tests { assert_eq!(output[5], f32::INFINITY); } } + +#[cfg(test)] +mod as_f32_into_tests { + use super::*; + + #[test] + fn fp16_conversion_preserves_bits_at_slice_boundaries() { + // IEEE-754 encodings for signed zeros, subnormal and normal limits, + // finite fractions, infinities and a quiet NaN with a payload. + let half_bits = [ + 0x0000, 0x8000, 0x0001, 0x03ff, 0x0400, 0x3c00, 0xbc00, 0x3555, 0x7bff, 0x7c00, 0xfc00, + 0x7e01, + ]; + let float_bits = [ + 0x00000000, 0x80000000, 0x33800000, 0x387fc000, 0x38800000, 0x3f800000, 0xbf800000, + 0x3eaaa000, 0x477fe000, 0x7f800000, 0xff800000, 0x7fc02000, + ]; + // Lengths around the Neon (4), AVX2 (8) and AVX-512 (16) widths, and an + // unrolled tail. + for length in [0, 3, 4, 7, 8, 9, 16, 17, 65] { + let source: Vec<_> = half_bits + .into_iter() + .cycle() + .take(length) + .map(f16::from_bits) + .collect(); + let expected: Vec<_> = float_bits.into_iter().cycle().take(length).collect(); + let mut output = vec![123.0_f32; length]; + + as_f32_into(&source, &mut output).unwrap(); + + assert_eq!( + output.into_iter().map(f32::to_bits).collect::>(), + expected, + "length {length}" + ); + } + } + + #[test] + fn mismatched_fp16_lengths_return_the_vector_repr_error() { + let source = [f16::ONE; 8]; + let mut output = [0.0_f32; 7]; + let expected = f16::as_f32_into(&source, &mut output).unwrap_err(); + + let actual = as_f32_into(&source, &mut output).unwrap_err(); + + assert_eq!(actual, expected); + } + + #[test] + fn other_representations_keep_their_conversion() { + fn check(source: [T; 2], expected: [f32; 2]) { + let mut output = [f32::NAN; 2]; + + as_f32_into(&source, &mut output).unwrap(); + + assert_eq!(output, expected, "{}", std::any::type_name::()); + } + + check([-2.5_f32, 1.25], [-2.5, 1.25]); + check([-2_i8, 127], [-2.0, 127.0]); + check([0_u8, 255], [0.0, 255.0]); + } +} diff --git a/diskann/src/graph/pipnn/hash_prune.rs b/diskann/src/graph/pipnn/hash_prune.rs new file mode 100644 index 0000000000..ff29383e9f --- /dev/null +++ b/diskann/src/graph/pipnn/hash_prune.rs @@ -0,0 +1,1546 @@ +/* + * 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. Only a +//! strictly smaller bf16 distance replaces that direction or the farthest entry +//! of a full reservoir. Equal distances keep the existing candidate, so ties can +//! depend on insertion order. +//! +//! Each source owns one lock. The lock protects its reservoir metadata and its +//! rows in the hash, distance, and neighbor arrays. `l_max` sets the logical +//! reservoir length. The `u8` metadata limits this value to 255. + +use parking_lot::lock_api::RawMutex as RawMutexTrait; +use std::cell::UnsafeCell; + +use super::{ + bf16::f32_to_bf16, + lsh::{LshSketches, MAX_PLANES}, + simd::Simd, +}; +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, +}; +#[cfg(not(miri))] +use rayon::prelude::*; + +/// Owned zero-initialized slab from `mmap(MAP_PRIVATE | MAP_ANONYMOUS)`. +#[cfg(target_os = "linux")] +struct MmapSlab { + ptr: *mut T, + len: usize, +} + +#[cfg(target_os = "linux")] +// SAFETY: the slab uniquely owns its mmap region until `drop`. Moving the slab +// transfers that ownership. `T: Send` permits transfer of initialized values. +unsafe impl Send for MmapSlab {} +#[cfg(target_os = "linux")] +// SAFETY: shared access exposes only `*const T`. `T: Sync` permits shared access +// to initialized values. HashPrune uses `UnsafeCell` and a point lock for writes. +unsafe impl Sync for MmapSlab {} + +#[cfg(target_os = "linux")] +impl MmapSlab { + fn new_zeroed(len: usize) -> ANNResult { + if len == 0 { + return Ok(Self { + ptr: std::ptr::NonNull::::dangling().as_ptr(), + len: 0, + }); + } + let bytes = len.checked_mul(std::mem::size_of::()).ok_or_else(|| { + ANNError::message(format!("HashPrune 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(|| { + ANNError::message(format!("HashPrune slab size {len} overflows usize")) + })?; + // SAFETY: `MEM_RESERVE | MEM_COMMIT` returns zero-initialized memory. + // `PAGE_READWRITE` permits all accesses used by this slab. + unsafe { + let ptr = winmem::VirtualAlloc( + std::ptr::null_mut(), + bytes, + winmem::MEM_RESERVE | winmem::MEM_COMMIT, + winmem::PAGE_READWRITE, + ); + if ptr.is_null() { + return Err( + ANNError::from(std::io::Error::last_os_error()).context(format!( + "VirtualAlloc failed for {bytes} HashPrune slab bytes" + )), + ); + } + Ok(Self { + ptr: ptr as *mut T, + len, + }) + } + } + + #[inline] + fn as_ptr(&self) -> *const T { + self.ptr + } +} + +#[cfg(windows)] +impl Drop for MmapSlab { + fn drop(&mut self) { + if self.len > 0 { + // SAFETY: this slab still uniquely owns the VirtualAlloc base pointer; + // MEM_RELEASE requires and receives `dwSize = 0`. + unsafe { + winmem::VirtualFree(self.ptr as winmem::Lpvoid, 0, winmem::MEM_RELEASE); + } + } + } +} + +/// Owned zero-initialized slab for platforms without `mmap` or `VirtualAlloc`. +#[cfg(not(any(target_os = "linux", windows)))] +struct MmapSlab(Vec); + +#[cfg(not(any(target_os = "linux", windows)))] +impl MmapSlab { + fn new_zeroed(len: usize) -> ANNResult { + Ok(Self(vec![T::default(); len])) + } + #[inline] + fn as_ptr(&self) -> *const T { + self.0.as_ptr() + } +} + +/// Largest reservoir length that fits in `ReservoirState`. +/// +/// `ReservoirState.len` and `ReservoirState.farthest_idx` are `u8`. Runtime `l_max` selects the +/// actual length. Values above this bound are invalid. +pub(crate) const MAX_RESERVOIR_LEN: usize = u8::MAX as usize; + +#[repr(C)] +struct ReservoirState { + len: u8, + farthest_idx: u8, + farthest_dist: u16, + _pad: [u8; 10], +} + +#[repr(C, align(16))] +struct LockedReservoirState { + lock: parking_lot::RawMutex, + state: UnsafeCell, +} + +impl LockedReservoirState { + fn new() -> Self { + Self { + lock: ::INIT, + state: UnsafeCell::new(ReservoirState::new_empty()), + } + } + + fn state_ptr(&self) -> *mut ReservoirState { + self.state.get() + } + + fn with_locked_state(&self, f: impl FnOnce(&mut ReservoirState) -> R) -> R { + struct UnlockOnDrop<'a>(&'a parking_lot::RawMutex); + impl Drop for UnlockOnDrop<'_> { + fn drop(&mut self) { + // SAFETY: the guard is created only after acquiring this mutex. + unsafe { self.0.unlock() }; + } + } + + self.lock.lock(); + let _guard = UnlockOnDrop(&self.lock); + // SAFETY: the mutex is separate from `state`, so contending threads may + // access the lock while this exclusive state reference is live. + f(unsafe { &mut *self.state.get() }) + } +} + +// SAFETY: `lock` guards every mutable access to `state`. Read-only extraction +// happens only after HashPrune is consumed, when no mutation can remain. +unsafe impl Sync for LockedReservoirState {} + +impl ReservoirState { + const fn new_empty() -> Self { + Self { + len: 0, + farthest_idx: 0, + farthest_dist: 0, + _pad: [0; 10], + } + } +} + +const _: [(); 16] = [(); std::mem::size_of::()]; + +// These pointers name one source reservoir's hash, distance, and neighbor rows. +// The caller must hold that source lock before it writes through a pointer. + +#[derive(Clone, Copy)] +struct ReservoirRows { + hashes: *mut u16, + distances: *mut u16, + neighbors: *mut u32, + row_stride: usize, +} + +#[derive(Clone, Copy)] +struct FindHashArgs { + hashes: *const u16, + row_stride: usize, + len: u8, + target: u16, +} + +#[derive(Clone, Copy)] +struct RelativeHashArgs { + src: *const f32, + dst: *const f32, + len: usize, +} + +type FindHash = Dispatched1, As>; +type RelativeHash = Dispatched1>; + +struct FindHashKernel; +struct RelativeHashKernel; +struct SelectFindHash; +struct SelectRelativeHash; + +impl Target for SelectFindHash +where + A: Architecture, + FindHashKernel: FTarget1, FindHashArgs>, +{ + fn run(self, arch: A) -> FindHash { + arch.dispatch1::, As>() + } +} + +impl Target for SelectRelativeHash +where + A: Architecture, + RelativeHashKernel: FTarget1, +{ + fn run(self, arch: A) -> RelativeHash { + arch.dispatch1::>() + } +} + +impl FTarget1, FindHashArgs> for FindHashKernel +where + A: Architecture, + A::i16x32: SIMDPartialEq, +{ + fn run(arch: A, args: FindHashArgs) -> Option { + find_hash_simd::(arch, args) + } +} + +impl FTarget1 for RelativeHashKernel +where + A: Simd, +{ + 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: A, args: RelativeHashArgs) -> u16 +where + A: Simd, +{ + if A::Vector::LANES >= MAX_PLANES { + // SAFETY: `src` and `dst` each contain `len <= MAX_PLANES <= A::Vector::LANES` + // values. The masked loads do not read inactive lanes. + let dst = unsafe { A::Vector::load_simd_first(arch, args.dst, args.len) }; + // SAFETY: `src` has the same checked length as `dst`. + let src = unsafe { A::Vector::load_simd_first(arch, args.src, args.len) }; + let active = (1_u64 << args.len) - 1; + return A::active_lanes((dst - src).ge_simd(A::Vector::splat(arch, 0.0))) as u16 + & active as u16; + } + + let mut bits = 0u16; + let mut offset = 0usize; + while offset < args.len { + let chunk_len = (args.len - offset).min(A::Vector::LANES); + // SAFETY: `offset + chunk_len <= args.len`. The masked loads do not read + // inactive lanes. + let dst = unsafe { A::Vector::load_simd_first(arch, args.dst.add(offset), chunk_len) }; + // SAFETY: `src` has the same checked length as `dst`. + let src = unsafe { A::Vector::load_simd_first(arch, args.src.add(offset), chunk_len) }; + let active = (1_u64 << chunk_len) - 1; + let chunk_bits = A::active_lanes((dst - src).ge_simd(A::Vector::splat(arch, 0.0))) & active; + bits |= (chunk_bits as u16) << offset; + offset += chunk_len; + } + bits +} + +/// Convert a bf16 distance to an order-preserving `u16` key. +/// +/// Raw bf16 bits are monotonic only for non-negative values. Inner-product +/// distance can be negative. This transform preserves the total numeric order +/// for both signs. +#[inline(always)] +fn ordered_distance_key(distance: f32) -> u16 { + let b = f32_to_bf16(distance); + if b & 0x8000 != 0 { !b } else { b | 0x8000 } +} + +/// Update the cached farthest entry for one reservoir. +/// +/// # Safety +/// +/// The caller holds the source lock. `state.len <= rows.row_stride`. The first +/// `state.len` entries of all three reservoir rows are initialized. +#[inline] +unsafe fn update_farthest(state: &mut ReservoirState, rows: ReservoirRows) { + if state.len == 0 { + state.farthest_dist = 0; + state.farthest_idx = 0; + return; + } + let mut max_idx: u8 = 0; + // SAFETY: `state.len > 0` and all active slots are initialized. + let mut max_distance = unsafe { *rows.distances }; + for i in 1..state.len as usize { + // SAFETY: `i < state.len <= rows.row_stride`, and all active entries are + // initialized. + let distance = unsafe { *rows.distances.add(i) }; + if distance > max_distance { + max_distance = distance; + max_idx = i as u8; + } + } + state.farthest_dist = max_distance; + state.farthest_idx = max_idx; +} + +/// Insert one edge into a locked reservoir. +/// +/// Only a strictly smaller bf16 distance replaces a matching hash or the +/// farthest entry of a full reservoir. Equal distances keep the existing entry. +/// +/// # 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 && dist_key >= state.farthest_dist { + 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_distance = unsafe { *rows.distances.add(idx) }; + if dist_key < current_distance { + 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; + // 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 dist_key > state.farthest_dist { + state.farthest_dist = dist_key; + state.farthest_idx = new_idx as u8; + } + return true; + } + + // The early rejection proved that the incoming distance is strictly smaller. + let idx = state.farthest_idx as usize; + // SAFETY: `idx < state.len <= rows.row_stride`; the caller holds the lock. + unsafe { + *rows.hashes.add(idx) = hash; + *rows.distances.add(idx) = dist_key; + *rows.neighbors.add(idx) = neighbor; + update_farthest(state, rows); + } + true +} + +/// Return at most `cap` neighbor IDs in distance order. +/// +/// `scratch` belongs to one Rayon extraction job and is reused for its rows. +/// +/// # Safety +/// +/// The caller must exclude mutation with the source lock or unique ownership. +/// `distances` and `neighbors` each point to `state.len` initialized entries. +unsafe fn collect_nearest_ids( + state: &ReservoirState, + distances: *const u16, + neighbors: *const u32, + cap: usize, + scratch: &mut Vec<(u32, u16)>, +) -> Vec { + let n = state.len as usize; + scratch.clear(); + scratch.reserve(n); + for i in 0..n { + // SAFETY: `i < n == state.len`, and both arrays have an initialized entry + // at `i`. + scratch.push(unsafe { (*neighbors.add(i), *distances.add(i)) }); + } + scratch.sort_unstable_by_key(|&(id, distance)| (distance, id)); + scratch[..n.min(cap)].iter().map(|&(id, _)| id).collect() +} + +/// Return at most `cap` neighbor IDs without sorting them. +/// +/// The caller does not depend on reservoir order. This function reads only the +/// neighbor row. +/// +/// # Safety +/// +/// The caller must exclude mutation with the source lock or unique ownership. +/// `neighbors` points to `state.len` initialized entries. +#[inline] +unsafe fn collect_neighbor_ids( + state: &ReservoirState, + neighbors: *const u32, + cap: usize, +) -> Vec { + let out_len = (state.len as usize).min(cap); + let mut out = Vec::with_capacity(out_len); + for i in 0..out_len { + // SAFETY: `i < out_len <= state.len`, and the neighbor entry is initialized. + out.push(unsafe { *neighbors.add(i) }); + } + out +} + +/// Bounded point reservoirs shared by parallel leaf workers. +/// +/// Source point `i` owns `states[i]` and row `i` in each reservoir array. +/// Its lock protects the metadata and all three rows. A worker holds at most one +/// source lock. Extraction consumes `HashPrune`, so no writer can remain. +pub(crate) struct HashPrune { + states: Vec, + hash_rows: UnsafeCell>, + distance_rows: UnsafeCell>, + neighbor_rows: UnsafeCell>, + row_stride: usize, + sketches: LshSketches, + l_max: usize, + find_hash: FindHash, + relative_hash: RelativeHash, +} + +// SAFETY: each mutable reservoir row is inside `UnsafeCell` and guarded by the +// matching source lock. Different source locks protect disjoint rows. Consuming +// extraction proves that no writer remains. +unsafe impl Send for HashPrune {} +// SAFETY: the same per-point lock protects mutation through shared HashPrune +// references; immutable sketches are safe to share. +unsafe impl Sync for HashPrune {} + +impl HashPrune { + /// Create one empty direction reservoir and LSH sketch for each dataset point. + pub(crate) fn new( + data: MatrixView<'_, T>, + num_planes: usize, + l_max: usize, + seed: u64, + ) -> ANNResult { + let npoints = data.nrows(); + let t0 = std::time::Instant::now(); + let sketches = LshSketches::try_new(data, num_planes, seed)?; + tracing::debug!( + elapsed_secs = t0.elapsed().as_secs_f64(), + "sketch computation" + ); + let t1 = std::time::Instant::now(); + let row_stride = l_max.next_multiple_of(32).max(32); + + let states: Vec<_> = (0..npoints).map(|_| LockedReservoirState::new()).collect(); + + // Each reservoir array has one `row_stride` row for each source point. + let total = npoints.checked_mul(row_stride).ok_or_else(|| { + ANNError::message(format!( + "HashPrune slab shape {npoints} x {row_stride} overflows usize" + )) + })?; + let hash_rows = MmapSlab::::new_zeroed(total)?; + let distance_rows = MmapSlab::::new_zeroed(total)?; + let neighbor_rows = MmapSlab::::new_zeroed(total)?; + + #[cfg(target_os = "linux")] + { + let state_bytes = states.len() * std::mem::size_of::(); + // SAFETY: each pointer names a contiguous allocation of `bytes`. + // `madvise` does not read or write the allocation. + unsafe { + for (ptr, bytes) in [ + (states.as_ptr() as *mut libc::c_void, state_bytes), + (hash_rows.as_ptr() as *mut libc::c_void, hash_rows.bytes()), + ( + distance_rows.as_ptr() as *mut libc::c_void, + distance_rows.bytes(), + ), + ( + neighbor_rows.as_ptr() as *mut libc::c_void, + neighbor_rows.bytes(), + ), + ] { + if bytes > 2 * 1024 * 1024 { + libc::madvise(ptr, bytes, libc::MADV_HUGEPAGE); + } + } + } + } + + tracing::debug!( + elapsed_secs = t1.elapsed().as_secs_f64(), + row_stride, + "reservoir allocation" + ); + + Ok(Self { + states, + hash_rows: UnsafeCell::new(hash_rows), + distance_rows: UnsafeCell::new(distance_rows), + neighbor_rows: UnsafeCell::new(neighbor_rows), + row_stride, + sketches, + l_max, + find_hash: arch::dispatch(SelectFindHash), + relative_hash: arch::dispatch(SelectRelativeHash), + }) + } + + /// Lock one source point's reservoir while `f` reads or updates it. + /// + /// RAII unlocks the point when the closure exits. + #[inline(always)] + fn with_locked_reservoir( + &self, + idx: usize, + f: impl FnOnce(&mut ReservoirState, ReservoirRows) -> R, + ) -> R { + let slot = &self.states[idx]; + let off = idx * self.row_stride; + // SAFETY: `idx` is in bounds. Each array has + // `states.len() * row_stride` elements. `UnsafeCell` permits these writes, + // and `with_locked_state` holds the source lock for the closure. + let rows = unsafe { + ReservoirRows { + hashes: (*self.hash_rows.get()).as_ptr().cast_mut().add(off), + distances: (*self.distance_rows.get()).as_ptr().cast_mut().add(off), + neighbors: (*self.neighbor_rows.get()).as_ptr().cast_mut().add(off), + row_stride: self.row_stride, + } + }; + slot.with_locked_state(|state| f(state, rows)) + } + + /// Add one leaf's weighted CSR edges to the point reservoirs. + /// + /// `point_ids` maps leaf-local positions to dataset IDs. `edge_offsets` and + /// `edges` form a CSR matrix with leaf-local targets. `sketch_scratch` stores + /// the gathered sketches for this leaf. + /// + /// The leaf builder supplies sorted unique dataset IDs and one monotonic + /// offset range per point. Each edge target is a valid leaf-local position. + /// Construction checks the full sketch and reservoir allocation shapes. + pub(crate) fn add_leaf_edges( + &self, + point_ids: &[u32], + edge_offsets: &[u32], + edges: &[(u32, f32)], + sketch_scratch: &mut Vec, + ) { + if edges.is_empty() { + return; + } + + let n = point_ids.len(); + let m = self.sketches.num_planes(); + let l_max = self.l_max as u8; + let sketch_len = n * m; + if sketch_scratch.len() < sketch_len { + sketch_scratch.resize(sketch_len, 0.0); + } + self.gather_sketches(point_ids, &mut sketch_scratch[..sketch_len]); + + for local_src in 0..n { + let start = edge_offsets[local_src] as usize; + let end = edge_offsets[local_src + 1] as usize; + if start == end { + continue; + } + let global_src = point_ids[local_src] as usize; + + if let Some(next) = (local_src + 1..n) + .find(|&i| edge_offsets[i] != edge_offsets[i + 1]) + .map(|i| point_ids[i] as usize) + { + let off = next * self.row_stride; + prefetch_hint_all(std::slice::from_ref(&self.states[next])); + // SAFETY: `next` is a dataset point ID. This range is the + // complete padded hash row for that point. Raw prefetch avoids + // a shared slice while another worker mutates the row. + unsafe { + let hashes = (*self.hash_rows.get()).as_ptr().add(off); + prefetch_hint_all_raw( + hashes.cast(), + self.row_stride * std::mem::size_of::(), + ); + } + } + + let src_sketch = &sketch_scratch[local_src * m..(local_src + 1) * m]; + self.with_locked_reservoir(global_src, |state, rows| { + for &(dst_local, dist) in &edges[start..end] { + let dst_index = dst_local as usize; + let global_dst = point_ids[dst_index]; + let dst_sketch = &sketch_scratch[dst_index * m..(dst_index + 1) * m]; + let hash = self.relative_hash.call(RelativeHashArgs { + src: src_sketch.as_ptr(), + dst: dst_sketch.as_ptr(), + len: m, + }); + // SAFETY: The source lock gives exclusive access to these rows. + // Each row has at least `l_max` entries. Insertion maintains + // `state.len <= l_max` and reads only initialized entries. + unsafe { + insert_reservoir_edge( + state, + rows, + hash, + global_dst, + dist, + l_max, + self.find_hash, + ) + }; + } + }); + } + } + + fn gather_sketches(&self, indices: &[u32], out: &mut [f32]) { + let dimensions = self.sketches.num_planes(); + let sketches = self.sketches.sketches(); + for (destination, &point) in out.chunks_exact_mut(dimensions).zip(indices) { + let first = point as usize * dimensions; + destination.copy_from_slice(&sketches[first..first + dimensions]); + } + } + + /// Consume the reservoirs and return at most `max_degree` nearest IDs per point. + pub(crate) fn into_nearest_lists(self, max_degree: usize) -> Vec> { + let row_stride = self.row_stride; + drop(self.sketches); + let HashPrune { + states, + hash_rows, + distance_rows, + neighbor_rows, + .. + } = self; + let hash_rows = hash_rows.into_inner(); + let distance_rows = distance_rows.into_inner(); + let neighbor_rows = neighbor_rows.into_inner(); + drop(hash_rows); + let extract = |scratch: &mut Vec<(u32, u16)>, i: usize| { + let off = i * row_stride; + // SAFETY: indexing proves that `i` names a live slot. This method + // consumes `self`, so no writer can overlap this state reference. + let state = unsafe { &*states[i].state_ptr() }; + // SAFETY: construction allocated `npoints * row_stride` entries; + // this loop keeps `i < npoints`, and insertion maintains + // `state.len <= l_max <= row_stride` initialized entries. + let ids = unsafe { + collect_nearest_ids( + state, + distance_rows.as_ptr().wrapping_add(off), + neighbor_rows.as_ptr().wrapping_add(off), + max_degree, + scratch, + ) + }; + // A neighbor always has the same relative hash for this source; + // insertion replaces an existing hash slot instead of appending. + AdjacencyList::from_vec_trusted(ids) + }; + + #[cfg(miri)] + { + let mut scratch = Vec::new(); + (0..states.len()) + .map(|i| extract(&mut scratch, i)) + .collect() + } + #[cfg(not(miri))] + { + (0..states.len()) + .into_par_iter() + .map_init(Vec::new, extract) + .collect() + } + } + + /// Consume the reservoirs and return all retained IDs without sorting them. + pub(crate) fn into_candidate_lists(self) -> Vec> { + let cap = self.l_max; + let row_stride = self.row_stride; + drop(self.sketches); + let HashPrune { + states, + hash_rows, + distance_rows, + neighbor_rows, + .. + } = self; + let hash_rows = hash_rows.into_inner(); + let distance_rows = distance_rows.into_inner(); + let neighbor_rows = neighbor_rows.into_inner(); + // Extraction reads only neighbor IDs. Drop the hash and distance arrays + // before the code creates the output lists. + drop(hash_rows); + drop(distance_rows); + let extract = |i: usize| { + let neighbors = neighbor_rows.as_ptr().wrapping_add(i * row_stride); + // SAFETY: indexing proves that `i` names a live slot. This method + // consumes `self`, so no writer can overlap this state reference. + let state = unsafe { &*states[i].state_ptr() }; + // SAFETY: construction allocated `npoints * row_stride` entries; + // this loop keeps `i < npoints`, and insertion maintains + // `state.len <= l_max <= row_stride` initialized entries. + let ids = unsafe { collect_neighbor_ids(state, neighbors, cap) }; + // Reservoir slots have unique hashes, and one neighbor cannot + // produce two hashes for the same source. + AdjacencyList::from_vec_trusted(ids) + }; + + #[cfg(miri)] + { + (0..states.len()).map(extract).collect() + } + #[cfg(not(miri))] + { + (0..states.len()).into_par_iter().map(extract).collect() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + fn build_hash_prune( + data: &[T], + points: usize, + dimensions: usize, + planes: usize, + l_max: usize, + ) -> ANNResult { + HashPrune::new( + MatrixView::try_from(data, points, dimensions).unwrap(), + planes, + l_max, + 42, + ) + } + + struct TestReservoir { + state: ReservoirState, + hashes: Vec, + distances: Vec, + neighbors: Vec, + row_stride: usize, + l_max: u8, + } + + impl TestReservoir { + fn new(l_max: usize) -> Self { + assert!(l_max <= MAX_RESERVOIR_LEN); + let row_stride = l_max.next_multiple_of(32).max(32); + Self { + state: ReservoirState::new_empty(), + hashes: vec![0; row_stride], + distances: vec![0; row_stride], + neighbors: vec![0; row_stride], + row_stride, + l_max: l_max as u8, + } + } + + fn rows(&self) -> ReservoirRows { + ReservoirRows { + hashes: self.hashes.as_ptr() as *mut u16, + distances: self.distances.as_ptr() as *mut u16, + neighbors: self.neighbors.as_ptr() as *mut u32, + row_stride: self.row_stride, + } + } + + fn insert(&mut self, hash: u16, neighbor: u32, distance: f32) -> bool { + let rows = self.rows(); + // SAFETY: the test owns the reservoir and holds its only mutable reference. + unsafe { + insert_reservoir_edge( + &mut self.state, + rows, + hash, + neighbor, + distance, + self.l_max, + arch::dispatch(SelectFindHash), + ) + } + } + + fn neighbors(&self) -> Vec<(u32, f32)> { + let mut entries: Vec<_> = self + .neighbors + .iter() + .copied() + .zip(self.distances.iter().copied()) + .take(self.len()) + .collect(); + entries.sort_unstable_by_key(|&(id, distance)| (distance, id)); + entries + .into_iter() + .map(|(id, key)| { + let bits = if key & 0x8000 != 0 { + key & 0x7fff + } else { + !key + }; + (id, f32::from_bits((bits as u32) << 16)) + }) + .collect() + } + + fn len(&self) -> usize { + self.state.len as usize + } + + fn is_empty(&self) -> bool { + self.state.len == 0 + } + } + + #[test] + fn reservoir_lock_serializes_state_mutation() { + let slot = LockedReservoirState::new(); + let start = std::sync::Barrier::new(3); + + std::thread::scope(|scope| { + for _ in 0..2 { + let slot = &slot; + let start = &start; + scope.spawn(move || { + start.wait(); + for _ in 0..16 { + slot.with_locked_state(|state| state.farthest_dist += 1); + } + }); + } + start.wait(); + }); + + assert_eq!(slot.with_locked_state(|state| state.farthest_dist), 32); + } + + fn add_edge(hp: &HashPrune, src: usize, dst: usize, distance: f32) { + let m = hp.sketches.num_planes(); + let sketches = hp.sketches.sketches(); + let hash = hp.relative_hash.call(RelativeHashArgs { + src: sketches[src * m..(src + 1) * m].as_ptr(), + dst: sketches[dst * m..(dst + 1) * m].as_ptr(), + len: m, + }); + let l_max = hp.l_max as u8; + hp.with_locked_reservoir(src, |state, rows| { + // SAFETY: `with_locked_reservoir` holds the source lock and supplies valid reservoir rows. + unsafe { + insert_reservoir_edge(state, rows, hash, dst as u32, distance, l_max, hp.find_hash) + }; + }); + } + + fn assert_sketch_source_type_matches_f32( + label: &str, + convert: impl Fn(u8) -> T, + reference: impl Fn(u8) -> f32, + ) where + T: VectorRepr + Send + Sync, + { + #[cfg(not(miri))] + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let points = 5; + // Miri keeps conversion and boundary coverage but omits duplicate sizes. + let dimensions: &[usize] = if cfg!(miri) { + &[1, 15, 16, 17, 33] + } else { + &[1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] + }; + let plane_counts: &[usize] = if cfg!(miri) { &[1, 16] } else { &[1, 8, 16] }; + for &dimensions in dimensions { + let raw: Vec = (0..points * dimensions) + .map(|index| { + let point = index / dimensions; + let dimension = index % dimensions; + (point + dimension) as u8 + }) + .collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + let f32_data: Vec = raw.iter().copied().map(&reference).collect(); + for &planes in plane_counts { + let build = || { + ( + LshSketches::try_new( + MatrixView::try_from(converted.as_slice(), points, dimensions).unwrap(), + planes, + 42, + ) + .unwrap(), + LshSketches::try_new( + MatrixView::try_from(f32_data.as_slice(), points, dimensions).unwrap(), + planes, + 42, + ) + .unwrap(), + ) + }; + #[cfg(miri)] + let (actual, expected) = build(); + #[cfg(not(miri))] + let (actual, expected) = pool.install(build); + assert_eq!( + actual.sketches(), + expected.sketches(), + "{label} dimensions={dimensions} planes={planes}" + ); + } + } + } + + // Source conversion. + + #[test] + fn f16_sketch_conversion_matches_f32_across_dimensions_and_planes() { + assert_sketch_source_type_matches_f32( + "f16", + |value| half::f16::from_f32_const(value as f32), + |value| value as f32, + ); + } + + #[test] + fn u8_sketch_conversion_matches_f32_across_dimensions_and_planes() { + assert_sketch_source_type_matches_f32("u8", |value| value, |value| value as f32); + } + + #[test] + fn i8_sketch_conversion_matches_f32_across_dimensions_and_planes() { + assert_sketch_source_type_matches_f32( + "i8", + |value| value as i8 - 11, + |value| (value as i8 - 11) as f32, + ); + } + + // Dispatched hash primitives. + + #[test] + fn relative_hash_matches_numeric_reference() { + let dispatched = arch::dispatch(SelectRelativeHash); + + let src = [ + 1.0, -2.0, 0.0, 7.5, -0.0, 3.25, -9.0, 4.0, 8.0, -1.5, 2.0, 0.0, 6.0, -3.0, 5.5, -7.25, + ]; + let dst = [ + 1.0, -3.0, 0.5, 7.0, 0.0, 3.25, -8.0, -4.0, 9.0, -1.5, -2.0, -0.0, 5.0, -2.0, 5.5, -8.0, + ]; + + for m in 0..=16 { + let mut expected_numeric_hash = 0u16; + for j in 0..m { + let diff: f32 = dst[j] - src[j]; + expected_numeric_hash |= ((diff >= 0.0) as u16) << j; + } + + let actual_dispatched_hash = dispatched.call(RelativeHashArgs { + src: src.as_ptr(), + dst: dst.as_ptr(), + len: m, + }); + assert_eq!(actual_dispatched_hash, expected_numeric_hash, "m={m}"); + } + } + + #[test] + fn relative_hash_defines_signed_zero_and_nan_buckets() { + let src = [0.0; 4]; + let dst = [ + 0.0, + -0.0, + f32::from_bits(0x7FC0_0000), + f32::from_bits(0xFFC0_0000), + ]; + + assert_eq!( + arch::dispatch(SelectRelativeHash).call(RelativeHashArgs { + src: src.as_ptr(), + dst: dst.as_ptr(), + len: dst.len(), + }), + 0b0011 + ); + } + + #[test] + fn find_hash_ignores_padding_and_returns_the_matching_active_index() { + let dispatched = arch::dispatch(SelectFindHash); + + for target in [0, 0xF00D] { + for len in [0usize, 1, 15, 16, 17, 31, 32, 33, 63, 64, 65, 254, 255] { + let row_stride = len.max(1).next_multiple_of(32); + let mut hashes = vec![target; row_stride]; + hashes[..len].fill(0x8001); + let args = |hashes: &[u16]| FindHashArgs { + hashes: hashes.as_ptr(), + row_stride, + len: len as u8, + target, + }; + + assert_eq!(dispatched.call(args(&hashes)), None, "len={len}"); + for index in [0, len / 2, len.saturating_sub(1)] { + if index < len { + hashes[index] = target; + assert_eq!(dispatched.call(args(&hashes)), Some(index), "len={len}"); + hashes[index] = 0x8001; + } + } + } + } + } + + // Storage and configuration. + + #[test] + fn slab_is_zeroed() { + let slab = MmapSlab::::new_zeroed(4).unwrap(); + assert!(!slab.as_ptr().is_null()); + // SAFETY: this test uniquely owns a live four-element slab. + let values = unsafe { std::slice::from_raw_parts(slab.as_ptr(), 4) }; + assert_eq!(values, &[0; 4]); + } + + #[test] + fn ordered_distance_key_preserves_bf16_order_for_all_signs() { + let values = [ + f32::NEG_INFINITY, + -100.0, + -0.0, + 0.0, + 0.25, + 100.0, + f32::INFINITY, + ]; + let keys: Vec<_> = values.iter().copied().map(ordered_distance_key).collect(); + assert!(keys.windows(2).all(|pair| pair[0] <= pair[1])); + } + + // Leaf ingestion and scratch reuse. + + #[test] + fn batched_leaf_edges_match_single_edge_reference() { + let data = [0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + let batched = build_hash_prune(&data, 4, 2, 8, 8).unwrap(); + let reference = build_hash_prune(&data, 4, 2, 8, 8).unwrap(); + let point_ids = [0, 1, 2, 3]; + let offsets = [0, 3, 6, 9, 12]; + let edges = [ + (1, 1.0), + (2, 1.0), + (3, 2.0), + (0, 1.0), + (2, 2.0), + (3, 1.0), + (0, 1.0), + (1, 2.0), + (3, 1.0), + (0, 2.0), + (1, 1.0), + (2, 1.0), + ]; + let mut scratch = Vec::new(); + + batched.add_leaf_edges(&point_ids, &offsets, &edges, &mut scratch); + for source in 0..point_ids.len() { + for &(target, distance) in + &edges[offsets[source] as usize..offsets[source + 1] as usize] + { + add_edge(&reference, source, target as usize, distance); + } + } + let canonicalize = |lists: Vec>| { + lists + .into_iter() + .map(|candidates| { + let mut ids = candidates.to_vec(); + ids.sort_unstable(); + ids + }) + .collect::>() + }; + let actual_batched_candidates = canonicalize(batched.into_candidate_lists()); + let expected_single_edge_candidates = canonicalize(reference.into_candidate_lists()); + + assert_eq!(actual_batched_candidates, expected_single_edge_candidates); + assert!( + actual_batched_candidates + .iter() + .all(|candidates| !candidates.is_empty()) + ); + } + + #[test] + fn reused_sketch_scratch_does_not_leak_candidates_between_leaves() { + // Given + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let hp = build_hash_prune(&data, 4, 1, 8, 4).unwrap(); + let expected_candidates = [vec![1], vec![0], vec![3], vec![2]]; + let mut sketch_scratch_with_stale_value = vec![99.0; 1]; + + // When + hp.add_leaf_edges( + &[0, 1], + &[0, 1, 2], + &[(1, 1.0), (0, 1.0)], + &mut sketch_scratch_with_stale_value, + ); + hp.add_leaf_edges( + &[2, 3], + &[0, 1, 2], + &[(1, 1.0), (0, 1.0)], + &mut sketch_scratch_with_stale_value, + ); + hp.add_leaf_edges( + &[0, 1], + &[0, 0, 0], + &[], + &mut sketch_scratch_with_stale_value, + ); + let actual_candidates: Vec<_> = hp + .into_candidate_lists() + .into_iter() + .map(Vec::from) + .collect(); + + // Then + assert_eq!(actual_candidates, expected_candidates); + } + + // TestReservoir replacement and ordering policy. + + #[test] + fn full_reservoir_evicts_the_farthest_candidate() { + // Given + let expected_capacity = 3; + let expected_neighbors = [(4, 0.5), (1, 1.0), (2, 2.0)]; + let mut reservoir = TestReservoir::new(expected_capacity); + assert!(reservoir.is_empty()); + reservoir.insert(0, 1, 1.0); + reservoir.insert(1, 2, 2.0); + reservoir.insert(2, 3, 3.0); + + // When + assert!(reservoir.insert(3, 4, 0.5)); + + // Then + assert_eq!(reservoir.len(), expected_capacity); + assert_eq!(reservoir.neighbors(), expected_neighbors); + } + + #[test] + fn same_hash_keeps_only_the_closest_candidate() { + // Given + let expected_candidate_count = 1; + let expected_closest_candidate = [(3, 1.0)]; + let mut reservoir = TestReservoir::new(5); + reservoir.insert(0, 1, 3.0); + + // When + reservoir.insert(0, 2, 2.0); + reservoir.insert(0, 3, 1.0); + assert!(!reservoir.insert(0, 4, 5.0)); + + // Then + assert_eq!(reservoir.len(), expected_candidate_count); + assert_eq!(reservoir.neighbors(), expected_closest_candidate); + } + + #[test] + fn appending_an_equal_distance_keeps_the_cached_farthest_slot() { + // Given: the first slot remains farthest when the second distance ties. + // Different hashes and IDs must not break that tie. + let mut reservoir = TestReservoir::new(2); + reservoir.insert(1, 30, 1.0); + reservoir.insert(9, 10, 1.0); + + // When: a closer edge replaces the first slot, leaving the second intact. + reservoir.insert(2, 20, 0.5); + + // Then + assert_eq!(reservoir.neighbors(), [(20, 0.5), (10, 1.0)]); + } + + #[rstest] + #[case::full(1)] + #[case::room_available(2)] + fn same_hash_bf16_tie_keeps_the_existing_candidate(#[case] capacity: usize) { + // Given: 1 + 2^-8 and 1 both truncate to bf16 1.0. + let mut reservoir = TestReservoir::new(capacity); + reservoir.insert(7, 20, 1.0 + 1.0 / 256.0); + + // When: a smaller ID and unrounded distance still cannot win a bf16 tie. + let inserted = reservoir.insert(7, 10, 1.0); + + // Then + assert!(!inserted); + assert_eq!(reservoir.neighbors(), [(20, 1.0)]); + } + + #[test] + fn full_reservoir_rejects_a_new_hash_tied_with_the_farthest_distance() { + // Given: 1 + 2^-8 truncates to bf16 1.0, above the other distance 0.5. + let mut reservoir = TestReservoir::new(2); + reservoir.insert(5, 30, 1.0 + 1.0 / 256.0); + reservoir.insert(6, 40, 0.5); + + // When: the incoming hash and ID are smaller, but its stored distance ties. + let inserted = reservoir.insert(1, 10, 1.0); + + // Then + assert!(!inserted); + assert_eq!(reservoir.neighbors(), [(40, 0.5), (30, 1.0)]); + } + + // Concurrency and consuming extraction. + + #[test] + fn parallel_unequal_distance_insertions_match_serial_neighbor_lists() { + // Given: all sketches share a hash. At the shared source 25, distances + // 1 and 2 make the nearer edge win regardless of which worker locks first. + let data = vec![0.0f32; 100 * 4]; + let parallel = build_hash_prune(&data, 100, 4, 4, 10).unwrap(); + let serial = build_hash_prune(&data, 100, 4, 4, 10).unwrap(); + + std::thread::scope(|scope| { + for sources in [0..25, 25..50] { + let parallel = ∥ + scope.spawn(move || { + for source in sources { + add_edge(parallel, source, source + 1, 1.0); + add_edge(parallel, source + 1, source, 2.0); + } + }); + } + }); + for source in 0..50 { + add_edge(&serial, source, source + 1, 1.0); + add_edge(&serial, source + 1, source, 2.0); + } + + assert_eq!(parallel.into_nearest_lists(5), serial.into_nearest_lists(5)); + } + + fn hash_prune_with_ranked_source_zero_candidates() -> HashPrune { + #[rustfmt::skip] + let point_vectors = [ + 0.0, 0.0, + 1.0, 0.0, + 0.0, 1.0, + -1.0, 0.0, + 0.0, -1.0, + 1.0, 1.0, + -1.0, 1.0, + 1.0, -1.0, + ]; + let hash_prune = build_hash_prune(&point_vectors, 8, 2, 16, 10).unwrap(); + // Use the neighbor ID as distance so the extraction order is explicit. + for target in 1..8 { + add_edge(&hash_prune, 0, target, target as f32); + } + hash_prune + } + + #[test] + fn candidate_extraction_returns_every_retained_neighbor() { + // Given + let hash_prune = hash_prune_with_ranked_source_zero_candidates(); + let expected_neighbor_ids: Vec<_> = (1..8).collect(); + + // When + let mut actual_neighbor_ids = hash_prune.into_candidate_lists()[0].to_vec(); + actual_neighbor_ids.sort_unstable(); + + // Then + assert_eq!(actual_neighbor_ids, expected_neighbor_ids); + } + + #[test] + fn nearest_extraction_keeps_the_requested_closest_neighbors() { + // Given + let hash_prune = hash_prune_with_ranked_source_zero_candidates(); + let expected_nearest_ids = [1, 2]; + + // When + let actual_nearest_lists = hash_prune.into_nearest_lists(2); + + // Then + assert_eq!(&*actual_nearest_lists[0], &expected_nearest_ids); + } + + #[test] + fn farthest_cache_updates_after_repeated_evictions() { + // Given + let expected_neighbors = [(14, 1.0), (13, 2.0), (12, 3.0)]; + let mut reservoir = TestReservoir::new(3); + reservoir.insert(0, 10, 5.0); + reservoir.insert(1, 11, 4.0); + reservoir.insert(2, 12, 3.0); + + // When + assert!(reservoir.insert(3, 13, 2.0)); + assert!(reservoir.insert(4, 14, 1.0)); + + // Then + assert_eq!(reservoir.neighbors(), expected_neighbors); + } + + #[test] + fn replacing_the_farthest_hash_refreshes_the_distance_threshold() { + // Given: hash 1 is farthest, then its closer replacement makes hash 2 farthest. + let mut reservoir = TestReservoir::new(2); + reservoir.insert(1, 10, 3.0); + reservoir.insert(2, 20, 2.0); + reservoir.insert(1, 30, 1.0); + + // When: 2.5 must not pass the refreshed threshold of 2.0. + let inserted = reservoir.insert(3, 40, 2.5); + + // Then + assert!(!inserted); + assert_eq!(reservoir.neighbors(), [(30, 1.0), (20, 2.0)]); + } + + #[test] + fn farthest_rescan_keeps_the_first_slot_among_equal_distances() { + // Given: after replacing the distance 2 entry, slots 1 and 2 tie at 1. + let mut reservoir = TestReservoir::new(3); + reservoir.insert(1, 10, 2.0); + reservoir.insert(2, 20, 1.0); + reservoir.insert(9, 30, 1.0); + reservoir.insert(3, 40, 0.5); + + // When: the next closer edge evicts slot 1, not the larger hash/ID in slot 2. + reservoir.insert(4, 50, 0.25); + + // Then + assert_eq!(reservoir.neighbors(), [(50, 0.25), (40, 0.5), (30, 1.0)]); + } + + #[test] + fn extraction_sorts_neighbors_when_the_farthest_slot_is_not_last() { + let mut reservoir = TestReservoir::new(4); + reservoir.insert(5, 1, 1.0); + reservoir.insert(10, 2, 3.0); + reservoir.insert(15, 3, 2.0); + reservoir.insert(3, 4, 0.5); + + assert_eq!( + reservoir.neighbors(), + [(4, 0.5), (1, 1.0), (3, 2.0), (2, 3.0)] + ); + } +} diff --git a/diskann/src/graph/pipnn/leaf_build.rs b/diskann/src/graph/pipnn/leaf_build.rs index 3b61957e55..af2d29c5cd 100644 --- a/diskann/src/graph/pipnn/leaf_build.rs +++ b/diskann/src/graph/pipnn/leaf_build.rs @@ -11,15 +11,20 @@ //! selected pair in both directions. Leaves overlap and run in parallel, so each //! data point has a locked candidate list. A job groups its edges by point first //! and then locks each list once. +//! +//! The HashPrune merge groups the edges of a leaf the same way. Each edge also +//! keeps its distance. The job adds the edges of each source point to the locked +//! reservoir of that point. use parking_lot::Mutex; -use crate::{ANNResult, graph::AdjacencyList, utils::VectorRepr}; +use crate::{ANNError, ANNResult, graph::AdjacencyList, utils::VectorRepr}; use diskann_utils::views::{MatrixView, MutMatrixView}; use rayon::prelude::*; use super::{ conversion::gather_as_f32, + hash_prune::HashPrune, leaf_kernel::{LeafKernelWorkspace, select_leaf_neighbors}, leaf_metric::LeafMetric, simd::Simd, @@ -29,12 +34,19 @@ use super::{ /// Scratch for one Rayon leaf job. /// /// Each buffer keeps the largest size that the job needed. A leaf uses a prefix. +/// The direct merge uses `local_adjacency`. The HashPrune merge uses the edge +/// buffers and `sketch_scratch`. #[derive(Default)] struct LeafBuffers { point_values: Vec, neighbors: Vec, local_adjacency: Vec>, kernel_workspace: LeafKernelWorkspace, + seen_pairs: Vec, + edge_offsets: Vec, + edges: Vec<(u32, f32)>, + edge_cursor: Vec, + sketch_scratch: Vec, } impl LeafBuffers { @@ -67,6 +79,14 @@ impl LeafBuffers { .iter_mut() .for_each(Vec::clear); } + + /// Grow the `source x target` flags of a leaf of `point_count` points. + /// + /// A leaf has at most `u32::MAX` distinct points, so the product fits in a + /// 64-bit `usize`. + fn prepare_seen_pairs(&mut self, point_count: usize) { + grow(&mut self.seen_pairs, point_count * point_count, false); + } } /// Build one candidate list per data point from the nearest neighbors in all @@ -98,10 +118,53 @@ where Ok(candidates.into_iter().map(Mutex::into_inner).collect()) } -/// Add the selected pairs of one leaf to the candidate lists. +/// Add the selected pairs of all leaves, with their distances, to the HashPrune +/// reservoirs. /// -/// The IDs of a leaf are distinct dataset rows, so `points x dimensions` is not -/// larger than the dataset. +/// Each selected pair adds both directions. +pub(super) fn add_hash_prune_candidates( + arch: A, + data: MatrixView<'_, T>, + leaves: Vec>, + requested_k: usize, + reservoirs: &HashPrune, +) -> ANNResult<()> +where + A: Simd, + M: LeafMetric, + T: VectorRepr, +{ + leaves + .par_iter() + .try_for_each_init(LeafBuffers::default, |buffers, point_ids| { + let leaf_k = + gather_leaf_neighbors::(arch, data, point_ids, requested_k, buffers)?; + if leaf_k == 0 { + return Ok(()); + } + let point_count = point_ids.len(); + buffers.prepare_seen_pairs(point_count); + let edge_count = build_symmetric_edge_csr( + 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, + }, + )?; + reservoirs.add_leaf_edges( + point_ids, + &buffers.edge_offsets[..point_count + 1], + &buffers.edges[..edge_count], + &mut buffers.sketch_scratch, + ); + Ok(()) + }) +} + +/// Add the selected pairs of one leaf to the candidate lists. fn add_leaf_candidates( arch: A, data: MatrixView<'_, T>, @@ -115,28 +178,17 @@ where M: LeafMetric, T: VectorRepr, { - let point_count = point_ids.len(); - let leaf_k = buffers.prepare(point_count, data.ncols(), requested_k); + let leaf_k = gather_leaf_neighbors::(arch, data, point_ids, requested_k, buffers)?; if leaf_k == 0 { return Ok(()); } - let point_values = &mut buffers.point_values[..point_count * data.ncols()]; - gather_as_f32(data, point_ids, point_values)?; - let points = MatrixView::try_from(&*point_values, point_count, data.ncols())?; - let neighbor_count = point_count * leaf_k; - let output = MutMatrixView::try_from( - &mut buffers.neighbors[..neighbor_count], - point_count, - leaf_k, - )?; - select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace)?; - + let point_count = point_ids.len(); buffers.prepare_local_adjacency(point_count); add_symmetric_neighbors( point_ids, leaf_k, - &buffers.neighbors[..neighbor_count], + &buffers.neighbors[..point_count * leaf_k], &mut buffers.local_adjacency[..point_count], ); for (&point_id, additions) in point_ids.iter().zip(&buffers.local_adjacency) { @@ -147,6 +199,41 @@ where Ok(()) } +/// Select the `k` nearest leaf points of each point of one leaf into +/// `buffers.neighbors`, and return the effective `k` of the leaf. +/// +/// The IDs of a leaf are distinct dataset rows, so `points x dimensions` is not +/// larger than the dataset. +fn gather_leaf_neighbors( + arch: A, + data: MatrixView<'_, T>, + point_ids: &[u32], + requested_k: usize, + buffers: &mut LeafBuffers, +) -> ANNResult +where + A: Simd, + M: LeafMetric, + T: VectorRepr, +{ + let point_count = point_ids.len(); + let leaf_k = buffers.prepare(point_count, data.ncols(), requested_k); + if leaf_k == 0 { + return Ok(0); + } + + let point_values = &mut buffers.point_values[..point_count * data.ncols()]; + gather_as_f32(data, point_ids, point_values)?; + let points = MatrixView::try_from(&*point_values, point_count, data.ncols())?; + let output = MutMatrixView::try_from( + &mut buffers.neighbors[..point_count * leaf_k], + point_count, + leaf_k, + )?; + select_leaf_neighbors::(arch, points, output, &mut buffers.kernel_workspace)?; + Ok(leaf_k) +} + /// Convert each selected pair from leaf positions to global IDs and add it to /// the edge lists of both points. fn add_symmetric_neighbors( @@ -171,6 +258,110 @@ fn add_symmetric_neighbors( } } +/// Reusable storage of [`build_symmetric_edge_csr`]. +struct EdgeBuffers<'a> { + /// `source x target` flags of the leaf. All flags are clear between calls. + seen: &'a mut [bool], + offsets: &'a mut Vec, + edges: &'a mut Vec<(u32, f32)>, + cursor: &'a mut Vec, +} + +/// Write the directed edges of one leaf as a CSR matrix and return the edge count. +/// +/// Each selected pair adds both directions, and each direction appears once. +/// Row `source` lists `(target, distance)` in leaf positions, in the order of +/// the first selection. `offsets[..=point_count]` holds the row bounds. +fn build_symmetric_edge_csr( + leaf_k: usize, + neighbors: &[Candidate], + buffers: EdgeBuffers<'_>, +) -> ANNResult { + let EdgeBuffers { + seen, + offsets, + edges, + cursor, + } = buffers; + let point_count = neighbors.len() / leaf_k; + grow(offsets, point_count + 1, 0); + offsets[..point_count + 1].fill(0); + + // The count pass sets the flag of each new direction and counts it in + // `offsets[source + 1]`. A source has at most `point_count - 1` targets, and + // a leaf has at most `u32::MAX` points, so a count fits in `u32`. + for (source, neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for neighbor in neighbors.iter().filter(|neighbor| neighbor.is_assigned()) { + let target = neighbor.local_idx as usize; + count_directed_edge(point_count, source, target, seen, offsets); + count_directed_edge(point_count, target, source, seen, offsets); + } + } + for point in 1..=point_count { + offsets[point] = offsets[point] + .checked_add(offsets[point - 1]) + .ok_or_else(|| ANNError::message("a leaf has more than u32::MAX directed edges"))?; + } + + let edge_count = offsets[point_count] as usize; + grow(edges, edge_count, (0, 0.0)); + grow(cursor, point_count, 0); + cursor[..point_count].copy_from_slice(&offsets[..point_count]); + let edges = &mut edges[..edge_count]; + let cursor = &mut cursor[..point_count]; + + // The write pass visits the directions in the same order. The first visit + // of a direction writes its edge and clears its flag for the next leaf. + for (source, neighbors) in neighbors.chunks_exact(leaf_k).enumerate() { + for neighbor in neighbors.iter().filter(|neighbor| neighbor.is_assigned()) { + let target = neighbor.local_idx as usize; + let distance = neighbor.distance; + write_counted_directed_edge(point_count, source, target, distance, seen, edges, cursor); + write_counted_directed_edge(point_count, target, source, distance, seen, edges, cursor); + } + } + Ok(edge_count) +} + +fn count_directed_edge( + point_count: usize, + source: usize, + target: usize, + seen: &mut [bool], + offsets: &mut [u32], +) { + let seen_entry = &mut seen[source * point_count + target]; + if !*seen_entry { + *seen_entry = true; + offsets[source + 1] += 1; + } +} + +fn write_counted_directed_edge( + point_count: usize, + source: usize, + target: usize, + distance: f32, + seen: &mut [bool], + edges: &mut [(u32, f32)], + cursor: &mut [u32], +) { + let seen_entry = &mut seen[source * point_count + target]; + if *seen_entry { + *seen_entry = false; + edges[cursor[source] as usize] = (target as u32, distance); + cursor[source] += 1; + } +} + +/// Grow `values` to at least `len` elements. The buffers never shrink, so a +/// smaller leaf uses a prefix. +fn grow(values: &mut Vec, len: usize, value: T) { + if values.len() < len { + values.resize(len, value); + } +} + #[cfg(test)] mod tests { use super::*; @@ -407,4 +598,77 @@ mod tests { ); } } + + #[test] + fn edge_csr_lists_each_direction_once_in_first_selection_order() { + // Each case gives the kernel output of a leaf with `k` columns. + // `Candidate::default()` is an empty slot. + for (case, k, neighbors, expected_offsets, expected_edges) in [ + ( + // The pair 1-2 is selected by both points. Its first selection, + // from point 1, sets the distance of both directions. + "three points", + 1, + vec![ + Candidate::new(1, 1.0), + Candidate::new(2, 2.0), + Candidate::new(1, 1.5), + ], + vec![0, 1, 3, 4], + vec![(1, 1.0), (0, 1.0), (2, 2.0), (1, 2.0)], + ), + ( + "empty slot", + 1, + vec![Candidate::new(1, 1.0), Candidate::default()], + vec![0, 1, 2], + vec![(1, 1.0), (0, 1.0)], + ), + ( + "pair selected from both ends", + 1, + vec![Candidate::new(1, 1.0), Candidate::new(0, 1.0)], + vec![0, 1, 2], + vec![(1, 1.0), (0, 1.0)], + ), + ( + "two neighbors per point", + 2, + vec![ + Candidate::new(1, 1.0), + Candidate::new(2, 4.0), + Candidate::new(0, 1.0), + Candidate::new(2, 1.0), + Candidate::new(1, 1.0), + Candidate::new(0, 4.0), + ], + vec![0, 2, 4, 6], + vec![(1, 1.0), (2, 4.0), (0, 1.0), (2, 1.0), (0, 4.0), (1, 1.0)], + ), + ] { + let point_count = neighbors.len() / k; + let mut seen = vec![false; point_count * point_count]; + // Stale values from a larger earlier leaf must not leak into this one. + let mut offsets = vec![9; point_count + 3]; + let mut edges = vec![(99, 99.0); 8]; + let mut cursor = vec![9; point_count + 3]; + + let edge_count = build_symmetric_edge_csr( + k, + &neighbors, + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap(); + + assert_eq!(edge_count, expected_edges.len(), "{case}"); + assert_eq!(offsets[..=point_count], expected_offsets, "{case}"); + assert_eq!(edges[..edge_count], expected_edges, "{case}"); + assert!(seen.iter().all(|&flag| !flag), "{case}: flags stay set"); + } + } } diff --git a/diskann/src/graph/pipnn/lsh.rs b/diskann/src/graph/pipnn/lsh.rs new file mode 100644 index 0000000000..6b40b8cd75 --- /dev/null +++ b/diskann/src/graph/pipnn/lsh.rs @@ -0,0 +1,229 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Random-hyperplane locality-sensitive hashing for dataset vectors. +//! +//! For each point `v`, the module computes +//! `Sketch(v) = [v · H_i for i in 0..num_planes]`. A seeded random generator +//! samples each hyperplane component from a standard normal distribution. +//! HashPrune compares two sketches to make a relative hash. +//! +//! `LshSketches` stores a row-major `npoints × num_planes` matrix. Each Rayon job +//! uses one `f32` conversion buffer for its source rows. `num_planes` cannot +//! exceed 16 because each relative hash is a `u16`. + +use crate::{ANNError, ANNResult, utils::VectorRepr}; +use diskann_utils::views::MatrixView; +use rand::SeedableRng; +use rand_distr::{Distribution, StandardNormal}; +#[cfg(not(miri))] +use rayon::prelude::*; + +/// Maximum number of hyperplanes (the hash output is `u16`). +pub(super) const MAX_PLANES: usize = 16; + +/// Precomputed LSH sketches for `npoints` vectors. +#[derive(Debug)] +pub(super) struct LshSketches { + num_planes: usize, + /// Row-major `npoints × num_planes`: `sketches[i*m + j] = dot(point_i, plane_j)`. + sketches: Vec, +} + +impl LshSketches { + /// Compute random-hyperplane projections for every point in `data`. + /// + /// Each worker converts one source row into reusable `f32` storage. Parallel + /// sketch work uses the currently installed Rayon pool. + pub(super) fn try_new( + data: MatrixView<'_, T>, + num_planes: usize, + seed: u64, + ) -> ANNResult { + let npoints = data.nrows(); + let ndims = data.ncols(); + let hyperplane_len = num_planes.checked_mul(ndims).ok_or_else(|| { + ANNError::message(format!( + "LSH matrix shape {num_planes} x {ndims} overflows usize" + )) + })?; + let sketch_len = npoints.checked_mul(num_planes).ok_or_else(|| { + ANNError::message(format!( + "LSH matrix shape {npoints} x {num_planes} overflows usize" + )) + })?; + + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let hyperplanes: Vec = (0..hyperplane_len) + .map(|_| StandardNormal.sample(&mut rng)) + .collect(); + + let mut sketches = vec![0.0f32; sketch_len]; + + #[cfg(miri)] + { + // Miri cannot execute Rayon's crossbeam dependency with strict provenance. + let mut buffer = Vec::new(); + for (point, sketch_row) in sketches.chunks_mut(num_planes).enumerate() { + fill_sketch_row(data, point, &hyperplanes, &mut buffer, sketch_row)?; + } + } + + #[cfg(not(miri))] + { + sketches + .par_chunks_mut(num_planes) + .enumerate() + .try_for_each_init(Vec::new, |buffer, (point, sketch_row)| { + fill_sketch_row(data, point, &hyperplanes, buffer, sketch_row) + })?; + } + + Ok(Self { + num_planes, + sketches, + }) + } + + /// Number of hyperplanes (also the number of bits in the hash). + #[inline] + pub(super) fn num_planes(&self) -> usize { + self.num_planes + } + + /// Return the row-major `npoints × num_planes` sketch buffer. + #[inline] + pub(super) fn sketches(&self) -> &[f32] { + &self.sketches + } +} + +fn fill_sketch_row( + data: MatrixView<'_, T>, + point: usize, + hyperplanes: &[f32], + buffer: &mut Vec, + sketch_row: &mut [f32], +) -> ANNResult<()> { + let dimensions = data.ncols(); + buffer.resize(dimensions, 0.0); + super::conversion::as_f32_into(data.row(point), &mut buffer[..dimensions]) + .map_err(Into::::into) + .map_err(|error| error.context(format!("converting LSH point {point}")))?; + for (plane_index, destination) in sketch_row.iter_mut().enumerate() { + let plane = &hyperplanes[plane_index * dimensions..(plane_index + 1) * dimensions]; + let mut dot = 0.0f32; + for dimension in 0..dimensions { + // SAFETY: both slices have exactly `dimensions` elements. + unsafe { + dot += *buffer.get_unchecked(dimension) * *plane.get_unchecked(dimension); + } + } + *destination = dot; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + fn run_with_threads(threads: usize, operation: impl FnOnce() -> R + Send) -> R { + #[cfg(miri)] + { + let _ = threads; + operation() + } + #[cfg(not(miri))] + { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + .install(operation) + } + } + + fn matrix_view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { + MatrixView::try_from(data, rows, columns).unwrap() + } + + #[test] + fn sketch_shape_has_one_value_per_point_and_plane() { + // Given + let point_vectors = [[1.0_f32, 0.0], [0.0, 1.0], [-1.0, 0.0]]; + let point_count = point_vectors.len(); + let dimensions = point_vectors[0].len(); + let plane_count = 4; + let expected_sketch_value_count = point_count * plane_count; + let data: Vec<_> = point_vectors.into_iter().flatten().collect(); + + // When + let sketches = run_with_threads(2, || { + LshSketches::try_new(matrix_view(&data, point_count, dimensions), plane_count, 42) + }) + .unwrap(); + + // Then + assert_eq!(sketches.num_planes(), plane_count); + assert_eq!(sketches.sketches().len(), expected_sketch_value_count); + } + + #[rstest] + #[case::first_seed(42)] + #[case::second_seed(99)] + fn sketches_match_seeded_serial_hyperplane_reference(#[case] seed: u64) { + // Given + let point_vectors = [ + [-3.0_f32, -2.0, -1.0, 0.0], + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + ]; + let npoints = point_vectors.len(); + let ndims = point_vectors[0].len(); + let planes = 5; + let data: Vec<_> = point_vectors.into_iter().flatten().collect(); + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let hyperplanes: Vec = (0..planes * ndims) + .map(|_| StandardNormal.sample(&mut rng)) + .collect(); + let expected_serial_sketch_values: Vec = data + .chunks_exact(ndims) + .flat_map(|point| { + hyperplanes + .chunks_exact(ndims) + .map(|plane| point.iter().zip(plane).map(|(x, h)| x * h).sum()) + }) + .collect(); + + // When + let actual_sketches = run_with_threads(2, || { + LshSketches::try_new(matrix_view(&data, npoints, ndims), planes, seed) + }) + .unwrap(); + + // Then + assert_eq!(actual_sketches.sketches(), expected_serial_sketch_values); + } + + #[rstest] + #[case::point_count_times_plane_count(usize::MAX, 0)] + #[case::dimension_times_plane_count(0, usize::MAX)] + fn sketch_construction_rejects_shape_overflow( + #[case] point_count: usize, + #[case] dimensions: usize, + ) { + // Given + let empty_data = matrix_view(&[] as &[f32], point_count, dimensions); + + // When + let error = + LshSketches::try_new(empty_data, 2, 42).expect_err("overflowing LSH shape must fail"); + + // Then + assert!(error.to_string().contains("overflows")); + } +} diff --git a/diskann/src/graph/pipnn/mod.rs b/diskann/src/graph/pipnn/mod.rs index fd9aab13c1..01f074db83 100644 --- a/diskann/src/graph/pipnn/mod.rs +++ b/diskann/src/graph/pipnn/mod.rs @@ -11,9 +11,13 @@ //! //! 1. `partitioning` splits the dataset into overlapping leaves. //! 2. `leaf_build` selects the `leaf_k` nearest neighbors of each point inside -//! each leaf and merges them into one candidate list per point. +//! each leaf. The direct merge adds them to one candidate list per point. +//! With [`HashPruneConfig`], `hash_prune` adds them to one bounded reservoir +//! per point instead. A reservoir keeps at most `l_max` candidates, at most +//! one for each direction hash. //! 3. `finalization` prunes each list that is longer than the graph degree with -//! Vamana RobustPrune. +//! Vamana RobustPrune. Without `final_prune`, a HashPrune build keeps the +//! nearest reservoir entries instead. //! //! [`build_graph`] selects the SIMD architecture and the metric once, so the //! steps compile for one architecture and one metric. It borrows one contiguous @@ -23,11 +27,14 @@ //! The build does not load providers or select start points. It also does not //! quantize, serialize, or search the graph. +mod bf16; mod conversion; mod finalization; +mod hash_prune; mod leaf_build; mod leaf_kernel; mod leaf_metric; +mod lsh; mod partition_kernel; mod partition_metric; mod partitioning; @@ -186,6 +193,19 @@ pub enum PiPNNConfigError { SamplingFraction(f64), #[error("fanout must not be empty")] EmptyFanout, + #[error("num_hash_planes ({0}) must be in [1, {max}]", max = lsh::MAX_PLANES)] + HashPlanes(usize), + #[error("l_max ({0}) must be in [1, {max}]", max = hash_prune::MAX_RESERVOIR_LEN)] + ReservoirLength(usize), + #[error( + "HashPrune capacity min(l_max={l_max}, hash buckets={hash_buckets}) must be at least \ + the graph degree ({degree})" + )] + HashPruneCapacity { + l_max: usize, + hash_buckets: usize, + degree: usize, + }, } crate::convert_error!(PiPNNConfigError); @@ -242,10 +262,52 @@ impl PiPNNConfig { } } +/// HashPrune parameters. +/// +/// Each edge `source -> target` gets a direction hash with one bit per random +/// hyperplane. The reservoir of `source` keeps the nearest target for each hash, +/// up to `l_max` targets. +#[derive(Clone, Debug, PartialEq)] +pub struct HashPruneConfig { + /// Number of random hyperplanes, one hash bit each. + pub num_hash_planes: usize, + /// Maximum number of candidates in the reservoir of a point. + pub l_max: usize, + /// Prune the reservoir candidates with Vamana RobustPrune. Without this + /// step, each point keeps its nearest candidates up to the graph degree. + pub final_prune: bool, +} + +impl HashPruneConfig { + /// Check the parameter ranges, and check that a reservoir can hold `degree` + /// candidates. + /// + /// A reservoir holds at most one candidate per hash, so it holds at most + /// `2^num_hash_planes` candidates. + pub fn validate(&self, degree: usize) -> Result<(), PiPNNConfigError> { + if !(1..=lsh::MAX_PLANES).contains(&self.num_hash_planes) { + return Err(PiPNNConfigError::HashPlanes(self.num_hash_planes)); + } + if !(1..=hash_prune::MAX_RESERVOIR_LEN).contains(&self.l_max) { + return Err(PiPNNConfigError::ReservoirLength(self.l_max)); + } + let hash_buckets = 1 << self.num_hash_planes; + if self.l_max.min(hash_buckets) < degree { + return Err(PiPNNConfigError::HashPruneCapacity { + l_max: self.l_max, + hash_buckets, + degree, + }); + } + Ok(()) + } +} + /// PiPNN parameters, graph policy, and the Rayon pool for one graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { config: PiPNNConfig, + hash_prune: Option, graph: &'a Config, metric: Metric, pool: &'a ThreadPool, @@ -265,11 +327,19 @@ impl<'a> PiPNNBuildContext<'a> { config.validate()?; Ok(Self { config, + hash_prune: None, graph, metric, pool, }) } + + /// Merge the leaf candidates with HashPrune reservoirs. + pub fn with_hash_prune(mut self, config: HashPruneConfig) -> Result { + config.validate(self.graph.pruned_degree().get())?; + self.hash_prune = Some(config); + Ok(self) + } } /// Build one adjacency list for each point of `data`. @@ -361,8 +431,28 @@ where // The leaf step takes the leaves by value, so their memory is free before // final pruning starts. let leaf_k = context.config.leaf_k.get(); - let candidates = tracing::info_span!("pipnn.leaf_build") - .in_scope(|| leaf_build::build_leaf_candidates::(arch, data, leaves, leaf_k))?; + let candidates = match &context.hash_prune { + None => tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::build_leaf_candidates::(arch, data, leaves, leaf_k) + })?, + Some(config) => { + let reservoirs = + hash_prune::HashPrune::new(data, config.num_hash_planes, config.l_max, 42)?; + tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::add_hash_prune_candidates::( + arch, + data, + leaves, + leaf_k, + &reservoirs, + ) + })?; + if !config.final_prune { + return Ok(reservoirs.into_nearest_lists(context.graph.pruned_degree().get())); + } + reservoirs.into_candidate_lists() + } + }; // Final pruning cuts each candidate list in place. Ok(tracing::info_span!("pipnn.finalization") .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric))) @@ -929,4 +1019,92 @@ mod construction_tests { ); } } + + #[test] + fn hash_prune_parameters_must_hold_the_graph_degree() { + let pool = thread_pool(1); + for (num_hash_planes, l_max, degree, expected) in [ + (0, 64, 2, Err(PiPNNConfigError::HashPlanes(0))), + (17, 64, 2, Err(PiPNNConfigError::HashPlanes(17))), + (8, 0, 2, Err(PiPNNConfigError::ReservoirLength(0))), + (8, 256, 2, Err(PiPNNConfigError::ReservoirLength(256))), + (16, 255, 2, Ok(())), + (8, 64, 64, Ok(())), + ( + 8, + 63, + 64, + Err(PiPNNConfigError::HashPruneCapacity { + l_max: 63, + hash_buckets: 256, + degree: 64, + }), + ), + // One plane gives two hashes, so a reservoir holds two candidates. + (1, 64, 2, Ok(())), + ( + 1, + 64, + 3, + Err(PiPNNConfigError::HashPruneCapacity { + l_max: 64, + hash_buckets: 2, + degree: 3, + }), + ), + ] { + let graph = graph_policy(degree, Metric::L2).unwrap(); + let context = + PiPNNBuildContext::new(partition_policy(), &graph, Metric::L2, &pool).unwrap(); + let config = HashPruneConfig { + num_hash_planes, + l_max, + final_prune: true, + }; + + let actual = context.with_hash_prune(config).map(|_| ()); + + assert_eq!( + actual, expected, + "planes={num_hash_planes}, l_max={l_max}, degree={degree}" + ); + } + } + + #[test] + fn hash_prune_keeps_the_nearest_neighbor_in_each_direction() { + // Two replicas add the same leaf in parallel. On a line, the two + // directions from a point have complementary hashes, so each reservoir + // keeps the nearest point on each side. Degree two keeps both. + let values = [-3.0_f32, 0.0, 1.0]; + let data = MatrixView::column_vector(&values[..]); + let graph = graph_policy(2, Metric::L2).unwrap(); + let config = PiPNNConfig { + c_max: nz(3), + fanout: vec![nz(1)], + leaf_k: nz(2), + replicas: nz(2), + ..partition_policy() + }; + for (threads, final_prune) in [(1, true), (4, true), (4, false)] { + let pool = thread_pool(threads); + let context = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(HashPruneConfig { + num_hash_planes: 8, + l_max: 16, + final_prune, + }) + .unwrap(); + + let actual = build_graph(data, &context).unwrap(); + + let actual: Vec> = actual.into_iter().map(Vec::from).collect(); + assert_eq!( + sorted_members_per_row(&actual), + [vec![1], vec![0, 2], vec![1]], + "{threads} threads, final_prune={final_prune}" + ); + } + } }