From 0bcf39bca0c797f54782ee671a757843e126e1a0 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 06:00:26 +0000 Subject: [PATCH 01/16] pipnn: add HashPrune candidate merging --- Cargo.lock | 3 + diskann-pipnn/Cargo.toml | 3 + diskann-pipnn/src/bf16.rs | 60 ++ diskann-pipnn/src/cpu_dispatch.rs | 73 ++ diskann-pipnn/src/hash_prune.rs | 1002 +++++++++++++++++++++++++ diskann-pipnn/src/hash_prune/tests.rs | 338 +++++++++ diskann-pipnn/src/leaf_build.rs | 210 +++++- diskann-pipnn/src/lib.rs | 112 ++- diskann-pipnn/src/lsh.rs | 273 +++++++ diskann-pipnn/tests/build_graph.rs | 40 +- diskann-pipnn/tests/config.rs | 42 +- 11 files changed, 2134 insertions(+), 22 deletions(-) create mode 100644 diskann-pipnn/src/bf16.rs create mode 100644 diskann-pipnn/src/cpu_dispatch.rs create mode 100644 diskann-pipnn/src/hash_prune.rs create mode 100644 diskann-pipnn/src/hash_prune/tests.rs create mode 100644 diskann-pipnn/src/lsh.rs diff --git a/Cargo.lock b/Cargo.lock index 41d8cef11..161c06d62 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -695,7 +695,10 @@ dependencies = [ "diskann-vector", "diskann-wide", "half", + "libc", + "parking_lot", "rand", + "rand_distr", "rayon", "thiserror 2.0.17", "tracing", diff --git a/diskann-pipnn/Cargo.toml b/diskann-pipnn/Cargo.toml index 3f706b02d..8405eb5b7 100644 --- a/diskann-pipnn/Cargo.toml +++ b/diskann-pipnn/Cargo.toml @@ -16,7 +16,10 @@ diskann-linalg.workspace = true diskann-utils.workspace = true diskann-vector.workspace = true diskann-wide.workspace = true +libc = "0.2" +parking_lot = "0.12" rand.workspace = true +rand_distr.workspace = true rayon.workspace = true thiserror.workspace = true tracing.workspace = true diff --git a/diskann-pipnn/src/bf16.rs b/diskann-pipnn/src/bf16.rs new file mode 100644 index 000000000..68d55280b --- /dev/null +++ b/diskann-pipnn/src/bf16.rs @@ -0,0 +1,60 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Lossy `f32` ↔ `u16` (bf16) conversion for compact distance storage. +//! +//! `bf16` is the upper 16 bits of an IEEE-754 `f32`: same exponent, mantissa +//! truncated from 23 to 7 bits. For non-negative values, `bf16` bit ordering +//! matches `f32` ordering, so a packed `bf16` can be compared as `u16` to give +//! the correct distance order — useful for keeping a top-k tracker as compact +//! 8-byte entries. + +/// Convert `f32` → bf16 by truncating the lower 16 mantissa bits. +#[inline(always)] +pub fn f32_to_bf16(v: f32) -> u16 { + (v.to_bits() >> 16) as u16 +} + +/// Reconstruct `f32` from a bf16, zero-filling the lower mantissa bits. +#[inline(always)] +pub fn bf16_to_f32(v: u16) -> f32 { + f32::from_bits((v as u32) << 16) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn roundtrip_bf16_exact_values() { + // bf16 has 7 mantissa bits. Pick f32 values whose lower 16 mantissa + // bits are zero so the truncation is lossless. + for &x in &[0.0_f32, 1.0, 2.0, 0.5, 0.25, 4.0, -1.0, -0.5] { + let back = bf16_to_f32(f32_to_bf16(x)); + assert_eq!(back, x, "exact bf16 roundtrip failed for {}", x); + } + } + + #[test] + fn truncation_is_within_relative_tolerance() { + // Values that aren't bf16-exact still round to within ~2^-7 relative. + use std::f32::consts::{E, PI}; + for &x in &[1e-10_f32, 1e10, PI, E] { + let back = bf16_to_f32(f32_to_bf16(x)); + let rel = ((back - x) / x).abs(); + assert!(rel <= 0.01, "{} → {}: rel error {} > 1%", x, back, rel); + } + } + + #[test] + fn ordering_preserved_for_non_negative() { + // For non-negative f32, bf16 (as u16) preserves ordering. + let xs: [f32; 6] = [0.0, 1e-10, 0.1, 1.0, 10.0, 1e10]; + let bs: Vec = xs.iter().map(|&x| f32_to_bf16(x)).collect(); + for w in bs.windows(2) { + assert!(w[0] <= w[1], "bf16 ordering broken: {} > {}", w[0], w[1]); + } + } +} diff --git a/diskann-pipnn/src/cpu_dispatch.rs b/diskann-pipnn/src/cpu_dispatch.rs new file mode 100644 index 000000000..cfe7e2b30 --- /dev/null +++ b/diskann-pipnn/src/cpu_dispatch.rs @@ -0,0 +1,73 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Cached CPU capabilities used by PiPNN's private SIMD kernels. +//! +//! Callers ask for the vector width required by a kernel family instead of +//! matching ISA names themselves. This keeps feature detection in one place +//! while allowing each family to request only the features it actually uses. + +use std::sync::OnceLock; + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub(crate) enum VectorWidth { + Wide, + Narrow, + Scalar, +} + +#[derive(Copy, Clone, Debug)] +struct CpuFeatures { + avx512f: bool, + avx512bw: bool, + avx2: bool, +} + +static FEATURES: OnceLock = OnceLock::new(); + +fn features() -> CpuFeatures { + *FEATURES.get_or_init(|| { + #[cfg(target_arch = "x86_64")] + let detected = CpuFeatures { + avx512f: std::is_x86_feature_detected!("avx512f"), + avx512bw: std::is_x86_feature_detected!("avx512bw"), + avx2: std::is_x86_feature_detected!("avx2"), + }; + + #[cfg(not(target_arch = "x86_64"))] + let detected = CpuFeatures { + avx512f: false, + avx512bw: false, + avx2: false, + }; + + tracing::info!(?detected, "PiPNN CPU features detected"); + detected + }) +} + +/// Width for kernels that use only packed `f32` arithmetic. +pub(crate) fn f32_width() -> VectorWidth { + let features = features(); + if features.avx512f { + VectorWidth::Wide + } else if features.avx2 { + VectorWidth::Narrow + } else { + VectorWidth::Scalar + } +} + +/// Width for packed 16-bit integer kernels. +pub(crate) fn u16_width() -> VectorWidth { + let features = features(); + if features.avx512f && features.avx512bw { + VectorWidth::Wide + } else if features.avx2 { + VectorWidth::Narrow + } else { + VectorWidth::Scalar + } +} diff --git a/diskann-pipnn/src/hash_prune.rs b/diskann-pipnn/src/hash_prune.rs new file mode 100644 index 000000000..7df3bd516 --- /dev/null +++ b/diskann-pipnn/src/hash_prune.rs @@ -0,0 +1,1002 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! HashPrune: LSH-based online pruning for merging edges from overlapping partitions. +//! +//! Storage is AoSoA hot/cold split: +//! - `hot: Vec` — one 16-byte slot per point (mutex + len + farthest). +//! Early rejection and lock acquisition only touch this slab. +//! - Three cold slabs (`cold_hashes`, `cold_distances`, `cold_neighbors`), each a +//! single `MmapSlab` of `npoints * scan_lanes` elements. Splitting hashes / +//! distances / neighbors into three contiguous arrays lets `find_hash` walk +//! pure u16 hashes (32 per cache line) instead of 8-byte mixed-AoS entries. +//! +//! Each slab is one contiguous allocation, so `madvise(HUGEPAGE)` is effective +//! when the kernel actually backs THP. +//! +//! `l_max` is dynamic user input; the cold slab stride (`scan_lanes`) and the +//! hash scan width both scale with it at runtime. The only fixed +//! bound is `MAX_RESERVOIR_LEN = 255`, the structural limit of the `u8` +//! `HotSlot.len` / `farthest_idx` fields; a larger `l_max` is rejected at +//! construction time. + +use parking_lot::lock_api::RawMutex as RawMutexTrait; +use std::cell::UnsafeCell; + +use crate::bf16::{bf16_to_f32, f32_to_bf16}; +use crate::lsh::{LshSketchError, LshSketches}; +use bytemuck::Pod; +use diskann::{utils::VectorRepr, ANNError, ANNResult}; +use rayon::prelude::*; + +/// Owned slab allocated via direct `mmap(MAP_PRIVATE | MAP_ANONYMOUS)`. The +/// kernel backs the range with its zero-page until first write, so we get +/// true lazy faulting for the AoSoA cold slabs rather than eagerly committing +/// the full reservoir allocation. +#[cfg(target_os = "linux")] +struct MmapSlab { + ptr: *mut T, + len: usize, +} + +#[cfg(target_os = "linux")] +// SAFETY: the allocation contains only `Pod` values and ownership transfers +// with the slab; mutation is synchronized by HashPrune's per-row locks. +unsafe impl Send for MmapSlab {} +#[cfg(target_os = "linux")] +// SAFETY: shared access exposes only immutable pointers; HashPrune synchronizes +// every mutation to a row. +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(|| crate::config_error(format!("slab size {len} overflows usize")))?; + // SAFETY: MAP_ANONYMOUS gives a zero-backed VA region; PROT_RW makes + // it readable/writable. Pages allocate on first write only. + 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: ptr was returned by mmap with this byte count. + unsafe { + libc::munmap(self.ptr as *mut libc::c_void, self.bytes()); + } + } + } +} + +#[cfg(target_os = "linux")] +impl std::ops::Deref for MmapSlab { + type Target = [T]; + fn deref(&self) -> &[T] { + // SAFETY: ptr+len describe a valid initialized slice (zero-init). + unsafe { std::slice::from_raw_parts(self.ptr, self.len) } + } +} + +/// Windows counterpart of the Linux mmap slab. `VirtualAlloc(MEM_RESERVE | +/// MEM_COMMIT, PAGE_READWRITE)` reserves a zero-backed anonymous range whose +/// pages fault in on first write — the same lazy-commit behavior as +/// `mmap(MAP_ANONYMOUS)`, so Windows gets the same peak-RSS win instead of the +/// eager-fault `Vec` fallback. +#[cfg(windows)] +mod winmem { + // Minimal FFI to the Win32 memory API — avoids pulling in the `windows` + // crate for four extern declarations. + 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; + + 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: see the Linux implementation; Windows provides the same zeroed, +// process-owned allocation semantics. +unsafe impl Send for MmapSlab {} +#[cfg(windows)] +// SAFETY: shared mutation is synchronized by HashPrune's per-row locks. +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(|| crate::config_error(format!("slab size {len} overflows usize")))?; + // SAFETY: MEM_RESERVE|MEM_COMMIT + PAGE_READWRITE returns a zero-backed + // RW region; physical pages fault in on first write only. Windows + // zero-fills committed pages, matching mmap's MAP_ANONYMOUS contract. + unsafe { + let ptr = winmem::VirtualAlloc( + std::ptr::null_mut(), + bytes, + winmem::MEM_RESERVE | winmem::MEM_COMMIT, + winmem::PAGE_READWRITE, + ); + if ptr.is_null() { + return Err( + ANNError::from(std::io::Error::last_os_error()).context(format!( + "VirtualAlloc failed for {bytes} HashPrune slab bytes" + )), + ); + } + Ok(Self { + ptr: ptr as *mut T, + len, + }) + } + } + + #[inline] + fn as_ptr(&self) -> *const T { + self.ptr + } + + #[inline] + #[allow(dead_code)] // parity with the Linux slab; madvise (Linux-only) is the sole caller + fn bytes(&self) -> usize { + self.len * std::mem::size_of::() + } +} + +#[cfg(windows)] +impl Drop for MmapSlab { + fn drop(&mut self) { + if self.len > 0 { + // SAFETY: ptr came from VirtualAlloc; MEM_RELEASE requires dwSize=0. + unsafe { + winmem::VirtualFree(self.ptr as winmem::LPVOID, 0, winmem::MEM_RELEASE); + } + } + } +} + +#[cfg(windows)] +impl std::ops::Deref for MmapSlab { + type Target = [T]; + fn deref(&self) -> &[T] { + // SAFETY: ptr+len describe a valid initialized slice (zero-init). + unsafe { std::slice::from_raw_parts(self.ptr, self.len) } + } +} + +/// Fallback slab for platforms that are neither Linux nor Windows: regular Vec. +/// Eager-fault behavior tracks the host allocator. +#[cfg(not(any(target_os = "linux", windows)))] +struct MmapSlab(Vec); + +#[cfg(not(any(target_os = "linux", windows)))] +impl MmapSlab { + fn new_zeroed(len: usize) -> ANNResult { + let mut values = Vec::new(); + values + .try_reserve_exact(len) + .map_err(ANNError::opaque) + .map_err(|error| error.context(format!("reserving {len} HashPrune slab elements")))?; + values.resize_with(len, T::default); + Ok(Self(values)) + } + #[inline] + fn as_ptr(&self) -> *const T { + self.0.as_ptr() + } + #[inline] + fn bytes(&self) -> usize { + self.0.len() * std::mem::size_of::() + } +} + +#[cfg(not(any(target_os = "linux", windows)))] +impl std::ops::Deref for MmapSlab { + type Target = [T]; + fn deref(&self) -> &[T] { + &self.0 + } +} + +/// Structural upper bound on per-node reservoir length: `HotSlot.len` and +/// `farthest_idx` are `u8`, so a reservoir can hold at most 255 entries. This +/// is an overflow guard, NOT the reservoir size — the cold slab stride +/// (`scan_lanes`) is sized to the runtime `l_max`, so the list scales with the +/// user's `l_max` up to this bound. `find_hash_simd` scans `scan_lanes / 32` +/// chunks, also runtime-sized. +pub(crate) const MAX_RESERVOIR_LEN: usize = u8::MAX as usize; + +/// Compute LSH sketches over `data` (row-major `npoints × ndims` of `T`). +fn sketches_from_data( + data: &[T], + npoints: usize, + ndims: usize, + num_planes: usize, + seed: u64, +) -> ANNResult { + LshSketches::try_new(npoints, ndims, num_planes, seed, |i, out| { + T::as_f32_into(&data[i * ndims..(i + 1) * ndims], out) + }) + .map_err(|error| match error { + LshSketchError::InvalidPlaneCount { actual, max } => { + crate::config_error(format!("num_hash_planes ({actual}) must be in 1..={max}")) + } + LshSketchError::ShapeOverflow { rows, columns } => ANNError::log_index_error(format!( + "LSH matrix shape {rows} x {columns} overflows usize" + )), + LshSketchError::Allocation(error) => ANNError::opaque(error), + LshSketchError::Fill(error) => error.into(), + }) +} + +// ─── HotSlot: 16-byte per-point mutex + cached fields ───────────────────────── + +#[repr(C, align(16))] +struct HotSlot { + lock: parking_lot::RawMutex, + len: u8, + farthest_idx: u8, + _pad0: u8, + farthest_dist: u16, + _pad1: [u8; 10], +} + +#[repr(transparent)] +struct LockedHotSlot(UnsafeCell); + +impl LockedHotSlot { + fn new() -> Self { + Self(UnsafeCell::new(HotSlot::new_empty())) + } + + fn get(&self) -> *mut HotSlot { + self.0.get() + } +} + +// SAFETY: every mutable access to the contained HotSlot is guarded by its +// embedded RawMutex. Read-only extraction happens only after HashPrune is +// consumed, when no mutation can remain. +unsafe impl Sync for LockedHotSlot {} + +impl HotSlot { + const fn new_empty() -> Self { + Self { + lock: ::INIT, + len: 0, + farthest_idx: 0, + _pad0: 0, + farthest_dist: 0, + _pad1: [0; 10], + } + } +} + +const _: () = assert!(std::mem::size_of::() == 16); + +// ─── Cold slabs ─────────────────────────────────────────────────────────────── +// +// Each per-point reservoir lives at index `idx` across three runtime-sized +// slabs: hashes, distances, neighbors. The stride is `scan_lanes` (l_max +// rounded up to a multiple of 32 so the AVX-512 / AVX-2 find_hash scan can +// stay aligned). At l_max=64 the stride is 64 and the per-point cold cost is +// 64 * 8 = 512 B; at l_max=128 the stride is 128 and the per-point cost is +// 1024 B. No fixed-size padding — the stride is the runtime l_max. +// +// `ColdSlotPtrs` is the lightweight view passed into `insert_locked` and the +// scan/update helpers — three raw pointers + the stride. Mutation safety is +// established by the caller via `HotSlot.lock`. + +#[derive(Clone, Copy)] +struct ColdSlotPtrs { + hashes: *mut u16, + distances: *mut u16, + neighbors: *mut u32, + scan_lanes: usize, +} + +#[inline(always)] +fn round_up_to_32(n: usize) -> usize { + n.div_ceil(32) * 32 +} + +// ─── find_hash SIMD: 32-way u16 compare ─────────────────────────────────────── + +type FindHash = unsafe fn(*const u16, usize, u8, u16) -> Option; +type RelativeHash = unsafe fn(*const f32, *const f32, usize) -> u16; + +fn select_find_hash() -> FindHash { + match crate::cpu_dispatch::u16_width() { + #[cfg(target_arch = "x86_64")] + crate::cpu_dispatch::VectorWidth::Wide => find_hash_wide, + #[cfg(target_arch = "x86_64")] + crate::cpu_dispatch::VectorWidth::Narrow => find_hash_narrow, + _ => find_hash_scalar, + } +} + +fn select_relative_hash() -> RelativeHash { + match crate::cpu_dispatch::f32_width() { + #[cfg(target_arch = "x86_64")] + crate::cpu_dispatch::VectorWidth::Wide => relative_hash_local_wide, + #[cfg(target_arch = "x86_64")] + crate::cpu_dispatch::VectorWidth::Narrow => relative_hash_local_narrow, + _ => relative_hash_local_scalar, + } +} + +/// SAFETY: caller guarantees AVX-512F + AVX-512BW at runtime and `hashes` +/// spans `scan_lanes` elements. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f", enable = "avx512bw")] +unsafe fn find_hash_wide( + hashes: *const u16, + scan_lanes: usize, + len: u8, + target: u16, +) -> Option { + use std::arch::x86_64::*; + let len = len as usize; + let t = _mm512_set1_epi16(target as i16); + // Scan only the chunks that cover valid entries (len), not the full + // scan_lanes capacity. At avg_deg ~60 (l_max 128) this halves the scan. + let chunks = len.div_ceil(32).min(scan_lanes / 32); + for chunk in 0..chunks { + let v = _mm512_loadu_si512(hashes.add(chunk * 32) as *const __m512i); + let mask = _mm512_cmpeq_epi16_mask(v, t); + if mask != 0 { + let lane = chunk * 32 + mask.trailing_zeros() as usize; + if lane < len { + return Some(lane); + } + } + } + None +} + +/// SAFETY: caller guarantees AVX2 at runtime. AVX2 is the workspace baseline +/// (`x86-64-v3`) so this path covers every supported deployment. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +unsafe fn find_hash_narrow( + hashes: *const u16, + scan_lanes: usize, + len: u8, + target: u16, +) -> Option { + use std::arch::x86_64::*; + let len = len as usize; + let t = _mm256_set1_epi16(target as i16); + // Scan only the chunks covering valid entries (len), not full scan_lanes. + let chunks = len.div_ceil(16).min(scan_lanes / 16); + for chunk in 0..chunks { + let v = _mm256_loadu_si256(hashes.add(chunk * 16) as *const __m256i); + let m = _mm256_cmpeq_epi16(v, t); + let bits = _mm256_movemask_epi8(m) as u32; + if bits != 0 { + let lane = chunk * 16 + (bits.trailing_zeros() as usize) / 2; + if lane < len { + return Some(lane); + } + } + } + None +} + +/// SAFETY: `hashes` must point at `len` valid `u16` slots. +#[inline(always)] +unsafe fn find_hash_scalar( + hashes: *const u16, + _scan_lanes: usize, + len: u8, + target: u16, +) -> Option { + (0..len as usize).find(|&i| *hashes.add(i) == target) +} + +// ─── relative_hash_local: AoSoA-friendly per-pair sketch comparison ─────────── +// +// `relative_hash_local(src, dst, m)` returns the `m`-bit pattern formed by +// `sign(dst[j] - src[j])` for j in 0..m. Used as the in-leaf LSH bucket for HP +// insertion (the "local sketch" cache hits L1). +// +// The implementation is selected once when HashPrune is constructed. + +/// SAFETY: `src` / `dst` must point at `m` valid `f32` slots, `m <= 16`. +unsafe fn relative_hash_local_scalar(src: *const f32, dst: *const f32, m: usize) -> u16 { + let mut h: u16 = 0; + for j in 0..m { + let diff = *dst.add(j) - *src.add(j); + let bit = ((!diff.is_sign_negative()) as u16) << j; + h |= bit; + } + h +} + +/// AVX2-vectorized `relative_hash_local`. Matches the scalar kernel's +/// `!is_sign_negative` semantics exactly (sign bit of `dst-src`), so the +/// produced graph is bit-identical to the scalar path it replaces on AVX2. +/// +/// SAFETY: caller guarantees AVX2 at runtime. `m <= 16` enforced upstream. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2")] +unsafe fn relative_hash_local_narrow(src: *const f32, dst: *const f32, m: usize) -> u16 { + use std::arch::x86_64::*; + let mut h: u16 = 0; + let mut j = 0usize; + while j + 8 <= m { + let d = _mm256_sub_ps(_mm256_loadu_ps(dst.add(j)), _mm256_loadu_ps(src.add(j))); + // movemask = sign bits of (dst-src); we want the NON-negative lanes + // (sign clear), i.e. `!is_sign_negative`, matching the scalar kernel. + let signs = _mm256_movemask_ps(d) as u16; + h |= (!signs & 0xFF) << j; + j += 8; + } + while j < m { + let diff = *dst.add(j) - *src.add(j); + h |= ((!diff.is_sign_negative()) as u16) << j; + j += 1; + } + h +} + +/// SAFETY: caller guarantees AVX-512F at runtime. `m <= 16` enforced upstream. +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx512f")] +unsafe fn relative_hash_local_wide(src: *const f32, dst: *const f32, m: usize) -> u16 { + use std::arch::x86_64::*; + let kmask: u16 = if m >= 16 { 0xFFFF } else { (1u16 << m) - 1 }; + let dst_v = _mm512_maskz_loadu_ps(kmask, dst); + let src_v = _mm512_maskz_loadu_ps(kmask, src); + let diff = _mm512_sub_ps(dst_v, src_v); + let mask = _mm512_cmp_ps_mask::<_CMP_GE_OQ>(diff, _mm512_setzero_ps()); + mask & kmask +} + +// ─── Per-reservoir mutation helpers (caller holds lock) ─────────────────────── + +/// Map a bf16-rounded `f32` to an order-preserving `u16` so raw integer +/// comparison matches float ordering for ALL signs. Raw bf16-bit compares are +/// monotonic only for non-negative values; InnerProduct distances are `-dot` +/// (negative), which otherwise sort inverted and make the reservoir evict its +/// best edges. For non-negative inputs this only sets the top bit on every +/// value, so L2/Cosine orderings — and the resulting graphs — are unchanged. +/// Inverse: [`key_to_bf16`]. +#[inline(always)] +fn ordered_key(distance: f32) -> u16 { + let b = f32_to_bf16(distance); + if b & 0x8000 != 0 { + !b + } else { + b | 0x8000 + } +} + +/// Inverse of [`ordered_key`]: recover the bf16 bits for distance readback. +#[inline(always)] +fn key_to_bf16(key: u16) -> u16 { + if key & 0x8000 != 0 { + key & 0x7FFF + } else { + !key + } +} + +/// SAFETY: caller holds the slot lock; pointers in `cold` are valid for +/// `scan_lanes` elements each. +#[inline] +unsafe fn update_farthest(hot: &mut HotSlot, cold: ColdSlotPtrs) { + if hot.len == 0 { + hot.farthest_dist = 0; + hot.farthest_idx = 0; + return; + } + let mut max_dist: u16 = 0; + let mut max_idx: u8 = 0; + for i in 0..hot.len as usize { + let d = *cold.distances.add(i); + if d > max_dist { + max_dist = d; + max_idx = i as u8; + } + } + hot.farthest_dist = max_dist; + hot.farthest_idx = max_idx; +} + +/// SAFETY: caller holds the slot lock; pointers in `cold` are valid for +/// `scan_lanes` elements each, and `l_max <= scan_lanes`. +#[inline(always)] +unsafe fn insert_locked( + hot: &mut HotSlot, + cold: ColdSlotPtrs, + hash: u16, + neighbor: u32, + distance: f32, + l_max: u8, + find_hash: FindHash, +) -> bool { + let dist_key = ordered_key(distance); + + if hot.len >= l_max && dist_key >= hot.farthest_dist { + return false; + } + + if let Some(idx) = find_hash(cold.hashes, cold.scan_lanes, hot.len, hash) { + if dist_key < *cold.distances.add(idx) { + let was_farthest = idx == hot.farthest_idx as usize; + *cold.neighbors.add(idx) = neighbor; + *cold.distances.add(idx) = dist_key; + if was_farthest { + update_farthest(hot, cold); + } + return true; + } + return false; + } + + if hot.len < l_max { + let new_idx = hot.len as usize; + *cold.hashes.add(new_idx) = hash; + *cold.distances.add(new_idx) = dist_key; + *cold.neighbors.add(new_idx) = neighbor; + hot.len += 1; + if dist_key >= hot.farthest_dist { + hot.farthest_dist = dist_key; + hot.farthest_idx = new_idx as u8; + } + return true; + } + + if dist_key < hot.farthest_dist { + let idx = hot.farthest_idx as usize; + *cold.hashes.add(idx) = hash; + *cold.distances.add(idx) = dist_key; + *cold.neighbors.add(idx) = neighbor; + update_farthest(hot, cold); + return true; + } + false +} + +/// Collect the reservoir's entries sorted by distance, truncated to `cap`. +/// A thread-local scratch `Vec`, sized to the reservoir's runtime fill and +/// reused across calls, avoids per-row allocation during extraction. +/// +/// SAFETY: caller holds the slot lock; `distances` and `neighbors` are valid +/// for `hot.len` elements. +unsafe fn collect_sorted_neighbors( + hot: &HotSlot, + distances: *const u16, + neighbors: *const u32, + cap: usize, +) -> Vec<(u32, f32)> { + thread_local! { + static SCRATCH: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; + } + let n = hot.len as usize; + SCRATCH.with(|cell| { + let mut scratch = cell.borrow_mut(); + scratch.clear(); + scratch.reserve(n); + for i in 0..n { + scratch.push((*neighbors.add(i), *distances.add(i))); + } + scratch.sort_unstable_by_key(|&(_, d)| d); + let out_len = n.min(cap); + let mut out = Vec::with_capacity(out_len); + for &(id, d) in &scratch[..out_len] { + out.push((id, bf16_to_f32(key_to_bf16(d)))); + } + out + }) +} + +/// Collect the reservoir's neighbor ids, truncated to `cap`, WITHOUT sorting. +/// Reservoir order is intentionally not preserved. Reading only `neighbors` +/// lets the caller drop the hashes and distances slabs before extraction; any +/// ordering required by a later graph-finalization policy belongs to that caller. +/// +/// SAFETY: caller holds the slot lock (or owns the reservoir); `neighbors` is +/// valid for `hot.len` elements. +#[inline] +unsafe fn collect_neighbor_ids(hot: &HotSlot, neighbors: *const u32, cap: usize) -> Vec { + let out_len = (hot.len as usize).min(cap); + let mut out = Vec::with_capacity(out_len); + for i in 0..out_len { + out.push(*neighbors.add(i)); + } + out +} + +// ─── HashPrune ──────────────────────────────────────────────────────────────── + +pub(crate) struct HashPrune { + hot: Vec, + /// AoSoA hashes slab: `npoints * scan_lanes` u16. + cold_hashes: MmapSlab, + /// AoSoA distances slab (bf16 in u16): `npoints * scan_lanes`. + cold_distances: MmapSlab, + /// AoSoA neighbors slab: `npoints * scan_lanes` u32. + cold_neighbors: MmapSlab, + /// Per-slot stride. Equals `round_up_to_32(l_max).max(32)`. Always a + /// multiple of 32 so the AVX-512 / AVX-2 find_hash scan stays aligned. + scan_lanes: usize, + sketches: LshSketches, + l_max: usize, + find_hash: FindHash, + relative_hash: RelativeHash, +} + +// SAFETY: HotSlot has interior mutability via RawMutex. The cold slabs are +// plain bit-pattern arrays; each per-point slot is guarded by HotSlot[i].lock. +// Disjoint-index parallel access is safe. +unsafe impl Send for HashPrune {} +// SAFETY: the same per-row lock protects all mutation through shared +// references, and immutable sketches are safe to share. +unsafe impl Sync for HashPrune {} + +impl HashPrune { + pub(crate) fn new( + data: &[T], + npoints: usize, + ndims: usize, + num_planes: usize, + l_max: usize, + seed: u64, + ) -> ANNResult { + if !(1..=MAX_RESERVOIR_LEN).contains(&l_max) { + return Err(crate::config_error(format!( + "HashPrune l_max ({l_max}) must be in 1..={MAX_RESERVOIR_LEN}" + ))); + } + + let t0 = std::time::Instant::now(); + let sketches = sketches_from_data(data, npoints, ndims, num_planes, seed)?; + tracing::debug!( + elapsed_secs = t0.elapsed().as_secs_f64(), + "sketch computation" + ); + let t1 = std::time::Instant::now(); + let scan_lanes = round_up_to_32(l_max).max(32); + + // Hot slab: one HotSlot per point, contiguous. + let mut hot: Vec = Vec::new(); + hot.try_reserve_exact(npoints) + .map_err(ANNError::opaque) + .map_err(|error| error.context(format!("reserving {npoints} HashPrune rows")))?; + for _ in 0..npoints { + hot.push(LockedHotSlot::new()); + } + + // Three cold slabs, each `npoints * scan_lanes` elements, allocated + // via mmap so the kernel keeps them zero-backed (no physical pages + // until first write). At scan_lanes = 64 the per-point cold cost is + // 64 * 8 = 512 B; at scan_lanes = 128 it is 1024 B. Reservoirs that + // never fill past the avg fill don't touch the high pages. + let total = npoints.checked_mul(scan_lanes).ok_or_else(|| { + crate::config_error(format!( + "HashPrune slab shape {npoints} x {scan_lanes} overflows usize" + )) + })?; + let cold_hashes = MmapSlab::::new_zeroed(total)?; + let cold_distances = MmapSlab::::new_zeroed(total)?; + let cold_neighbors = MmapSlab::::new_zeroed(total)?; + + // Hint hugepages on slabs > 2 MB so DTLB pressure scales with 2 MB + // pages instead of 4 KB. Non-fatal on failure; no-op on kernels + // without THP. Linux-only: the Windows MEM_LARGE_PAGES equivalent must + // be requested at VirtualAlloc time AND needs SeLockMemoryPrivilege + // (off by default), so a failure would abort the slab rather than + // silently fall back — not worth it for a DTLB hint. + #[cfg(target_os = "linux")] + { + let hot_bytes = hot.len() * std::mem::size_of::(); + // SAFETY: each slab backs a contiguous allocation of the indicated + // byte length. madvise is non-fatal on failure. + unsafe { + for (ptr, bytes) in [ + (hot.as_ptr() as *mut libc::c_void, hot_bytes), + ( + cold_hashes.as_ptr() as *mut libc::c_void, + cold_hashes.bytes(), + ), + ( + cold_distances.as_ptr() as *mut libc::c_void, + cold_distances.bytes(), + ), + ( + cold_neighbors.as_ptr() as *mut libc::c_void, + cold_neighbors.bytes(), + ), + ] { + if bytes > 2 * 1024 * 1024 { + libc::madvise(ptr, bytes, libc::MADV_HUGEPAGE); + } + } + } + } + + tracing::debug!( + elapsed_secs = t1.elapsed().as_secs_f64(), + scan_lanes, + "reservoir allocation" + ); + + Ok(Self { + hot, + cold_hashes, + cold_distances, + cold_neighbors, + scan_lanes, + sketches, + l_max, + find_hash: select_find_hash(), + relative_hash: select_relative_hash(), + }) + } + + /// Locks the per-slot mutex at `idx`, runs `f` with mutable access to the + /// hot slot and its cold-slab pointers, and unlocks on return or panic. + #[inline(always)] + fn with_locked(&self, idx: usize, f: impl FnOnce(&mut HotSlot, ColdSlotPtrs) -> R) -> R { + struct UnlockOnDrop { + hot_ptr: *mut HotSlot, + } + impl Drop for UnlockOnDrop { + fn drop(&mut self) { + // SAFETY: only constructed after locking this slot's mutex. + unsafe { (*self.hot_ptr).lock.unlock() }; + } + } + assert!(idx < self.hot.len(), "HashPrune row index out of bounds"); + // SAFETY: idx bounds-checked above, so `idx * scan_lanes + scan_lanes` + // is within each cold slab's capacity. UnlockOnDrop unlocks on panic. + unsafe { + let hot_ptr = self.hot[idx].get(); + let off = idx * self.scan_lanes; + let cold = ColdSlotPtrs { + hashes: (self.cold_hashes.as_ptr() as *mut u16).add(off), + distances: (self.cold_distances.as_ptr() as *mut u16).add(off), + neighbors: (self.cold_neighbors.as_ptr() as *mut u32).add(off), + scan_lanes: self.scan_lanes, + }; + (*hot_ptr).lock.lock(); + let _guard = UnlockOnDrop { hot_ptr }; + f(&mut *hot_ptr, cold) + } + } + + /// Merge one leaf's CSR edge list into the global reservoirs. + /// + /// Sketch layout and gathering are HashPrune implementation details; the + /// caller only lends a reusable buffer to avoid per-leaf allocation. + 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; + } + self.add_leaf_edges_with_scratch(point_ids, edge_offsets, edges, sketch_scratch); + } + + fn add_leaf_edges_with_scratch( + &self, + point_ids: &[u32], + edge_offsets: &[u32], + edges: &[(u32, f32)], + sketch_scratch: &mut Vec, + ) { + let n = point_ids.len(); + let m = self.sketches.num_planes(); + let l_max = self.l_max as u8; + debug_assert_eq!(edge_offsets.len(), n + 1); + let sketch_len = n * m; + 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; + + // Prefetch the next non-empty source's hot and cold slots. + #[cfg(all(target_arch = "x86_64", target_feature = "sse"))] + { + use std::arch::x86_64::*; + if local_src + 1 < n { + let mut nxt = local_src + 1; + while nxt < n && edge_offsets[nxt] == edge_offsets[nxt + 1] { + nxt += 1; + } + if nxt < n { + let nxt_global = point_ids[nxt] as usize; + let off = nxt_global * self.scan_lanes; + // SAFETY: nxt_global < npoints; off + scan_lanes + // bounded by slab capacity. Prefetch is non-fatal. + unsafe { + let hot_p = self.hot[nxt_global].get().cast::(); + _mm_prefetch::<{ _MM_HINT_T0 }>(hot_p); + // hashes array covers `scan_lanes` u16 = scan_lanes*2 bytes; + // prefetch in 64-byte cache-line strides. + let hashes_p = self.cold_hashes.as_ptr().add(off) as *const i8; + let hashes_bytes = self.scan_lanes * std::mem::size_of::(); + let mut b = 0usize; + while b < hashes_bytes { + _mm_prefetch::<{ _MM_HINT_T0 }>(hashes_p.add(b)); + b += 64; + } + } + } + } + } + + let src_sketch = &sketch_scratch[local_src * m..(local_src + 1) * m]; + self.with_locked(global_src, |hot, cold| { + for &(dst_local, dist) in &edges[start..end] { + let global_dst = point_ids[dst_local as usize]; + let dst_sketch = + &sketch_scratch[dst_local as usize * m..(dst_local as usize + 1) * m]; + debug_assert!(m <= 16, "num_planes <= 16 enforced by validate"); + // SAFETY: m <= 16, sketches are m-elt slices. + let hash = unsafe { + (self.relative_hash)(src_sketch.as_ptr(), dst_sketch.as_ptr(), m) + }; + // SAFETY: cold ptrs from with_locked are valid for scan_lanes. + unsafe { + insert_locked(hot, cold, hash, global_dst, dist, l_max, self.find_hash) + }; + } + }); + } + } + + fn gather_sketches(&self, indices: &[u32], out: &mut [f32]) { + let m = self.sketches.num_planes(); + let src = self.sketches.sketches(); + debug_assert_eq!(out.len(), indices.len() * m); + for (i, &idx) in indices.iter().enumerate() { + let g = idx as usize; + out[i * m..(i + 1) * m].copy_from_slice(&src[g * m..(g + 1) * m]); + } + } + + /// Extract the nearest `max_degree` candidates retained by HashPrune. + #[allow(clippy::disallowed_methods)] // build_graph installs the caller-owned pool. + pub(crate) fn into_nearest_lists(self, max_degree: usize) -> Vec> { + let scan_lanes = self.scan_lanes; + drop(self.sketches); + let HashPrune { + hot, + cold_hashes, + cold_distances, + cold_neighbors, + .. + } = self; + drop(cold_hashes); + (0..hot.len()) + .into_par_iter() + .map(|i| { + let off = i * scan_lanes; + // SAFETY: extraction owns every row, so no mutation remains. + let hot = unsafe { &*hot[i].get() }; + // SAFETY: i < hot.len() == npoints, so off + scan_lanes is + // within both cold slabs. + let nbrs = unsafe { + collect_sorted_neighbors( + hot, + cold_distances.as_ptr().wrapping_add(off), + cold_neighbors.as_ptr().wrapping_add(off), + max_degree, + ) + }; + nbrs.into_iter().map(|(id, _)| id).collect() + }) + .collect() + } + + /// Extract each point's full reservoir as candidate IDs. Drops the hashes + /// and distances slabs (2/3 of the reservoir) before materializing the copy, + /// so only the neighbors slab overlaps it. + #[allow(clippy::disallowed_methods)] // build_graph installs the caller-owned pool. + pub(crate) fn into_candidate_lists(self) -> Vec> { + let cap = self.l_max; + let scan_lanes = self.scan_lanes; + drop(self.sketches); + let HashPrune { + hot, + cold_hashes, + cold_distances, + cold_neighbors, + .. + } = self; + // Neither the hashes (LSH dedup index) nor the distances (bf16 + // keep-closer key) are read again — free them before the copy so the + // reservoir+copy overlap is just the neighbors slab. + drop(cold_hashes); + drop(cold_distances); + (0..hot.len()) + .into_par_iter() + .map(|i| { + let neighbors = cold_neighbors.as_ptr().wrapping_add(i * scan_lanes); + // SAFETY: extraction owns every row; no mutation remains. + let hot = unsafe { &*hot[i].get() }; + // SAFETY: i < hot.len() == npoints; the row spans scan_lanes + // slots and hot.len <= l_max <= scan_lanes. + unsafe { collect_neighbor_ids(hot, neighbors, cap) } + }) + .collect() + } +} + +#[cfg(test)] +mod tests; diff --git a/diskann-pipnn/src/hash_prune/tests.rs b/diskann-pipnn/src/hash_prune/tests.rs new file mode 100644 index 000000000..a14ae3b9c --- /dev/null +++ b/diskann-pipnn/src/hash_prune/tests.rs @@ -0,0 +1,338 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::*; + +struct Reservoir { + hot: HotSlot, + hashes: Vec, + distances: Vec, + neighbors: Vec, + scan_lanes: usize, + l_max: u8, +} + +impl Reservoir { + fn new(l_max: usize) -> Self { + assert!(l_max <= MAX_RESERVOIR_LEN); + let scan_lanes = round_up_to_32(l_max).max(32); + Self { + hot: HotSlot::new_empty(), + hashes: vec![0; scan_lanes], + distances: vec![0; scan_lanes], + neighbors: vec![0; scan_lanes], + scan_lanes, + l_max: l_max as u8, + } + } + + fn cold(&self) -> ColdSlotPtrs { + ColdSlotPtrs { + hashes: self.hashes.as_ptr() as *mut u16, + distances: self.distances.as_ptr() as *mut u16, + neighbors: self.neighbors.as_ptr() as *mut u32, + scan_lanes: self.scan_lanes, + } + } + + fn insert(&mut self, hash: u16, neighbor: u32, distance: f32) -> bool { + let cold = self.cold(); + // SAFETY: the test owns the reservoir and holds its only mutable reference. + unsafe { + insert_locked( + &mut self.hot, + cold, + hash, + neighbor, + distance, + self.l_max, + select_find_hash(), + ) + } + } + + fn neighbors(&self) -> Vec<(u32, f32)> { + let cold = self.cold(); + // SAFETY: the test owns the reservoir; all cold slabs span scan_lanes entries. + unsafe { collect_sorted_neighbors(&self.hot, cold.distances, cold.neighbors, usize::MAX) } + } + + fn len(&self) -> usize { + self.hot.len as usize + } + + fn is_empty(&self) -> bool { + self.hot.len == 0 + } +} + +fn add_edge(hp: &HashPrune, src: usize, dst: usize, distance: f32) { + let m = hp.sketches.num_planes(); + let sketches = hp.sketches.sketches(); + // SAFETY: both offsets select complete `m`-element sketch rows. + let hash = unsafe { + (hp.relative_hash)( + sketches.as_ptr().add(src * m), + sketches.as_ptr().add(dst * m), + m, + ) + }; + let l_max = hp.l_max as u8; + hp.with_locked(src, |hot, cold| { + // SAFETY: with_locked guards the row and supplies valid cold-slab pointers. + unsafe { insert_locked(hot, cold, hash, dst as u32, distance, l_max, hp.find_hash) }; + }); +} + +#[test] +fn test_relative_hash_local_narrow_matches_scalar_semantics() { + #[cfg(target_arch = "x86_64")] + { + if !is_x86_feature_detected!("avx2") { + return; + } + + let src = [ + 1.0, -2.0, 0.0, 7.5, -0.0, 3.25, -9.0, 4.0, 8.0, -1.5, 2.0, 0.0, 6.0, -3.0, 5.5, -7.25, + ]; + let dst = [ + 1.0, -3.0, 0.5, 7.0, 0.0, 3.25, -8.0, -4.0, 9.0, -1.5, -2.0, -0.0, 5.0, -2.0, 5.5, -8.0, + ]; + + for m in 0..=16 { + let mut expected = 0u16; + for j in 0..m { + let diff: f32 = dst[j] - src[j]; + expected |= ((!diff.is_sign_negative()) as u16) << j; + } + + // SAFETY: the feature guard above checked AVX2 and both arrays + // contain 16 elements, so every tested prefix is in bounds. + let actual = unsafe { relative_hash_local_narrow(src.as_ptr(), dst.as_ptr(), m) }; + assert_eq!(actual, expected, "m={m}"); + } + } +} + +#[test] +fn test_reservoir_basic() { + let mut reservoir = Reservoir::new(3); + assert!(reservoir.is_empty()); + + assert!(reservoir.insert(0, 1, 1.0)); + assert!(reservoir.insert(1, 2, 2.0)); + assert!(reservoir.insert(2, 3, 3.0)); + assert_eq!(reservoir.len(), 3); + + assert!(reservoir.insert(3, 4, 0.5)); + assert_eq!(reservoir.len(), 3); + + let neighbors = reservoir.neighbors(); + assert!(!neighbors.iter().any(|(id, _)| *id == 3)); + assert!(neighbors.iter().any(|(id, _)| *id == 4)); +} + +#[test] +fn test_reservoir_same_hash_keeps_closer() { + let mut reservoir = Reservoir::new(10); + + assert!(reservoir.insert(0, 1, 2.0)); + assert_eq!(reservoir.len(), 1); + + assert!(reservoir.insert(0, 2, 1.0)); + assert_eq!(reservoir.len(), 1); + + let neighbors = reservoir.neighbors(); + assert_eq!(neighbors[0].0, 2); + assert_eq!(neighbors[0].1, 1.0); + + assert!(!reservoir.insert(0, 3, 5.0)); + assert_eq!(reservoir.len(), 1); +} + +#[test] +fn test_hash_prune_end_to_end() { + let data = vec![0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + + let hp = HashPrune::new(&data, 4, 2, 4, 10, 42).unwrap(); + + add_edge(&hp, 0, 1, 1.0); + add_edge(&hp, 0, 2, 1.0); + add_edge(&hp, 0, 3, 1.414); + add_edge(&hp, 1, 0, 1.0); + add_edge(&hp, 1, 3, 1.0); + add_edge(&hp, 2, 0, 1.0); + add_edge(&hp, 2, 3, 1.0); + add_edge(&hp, 3, 1, 1.0); + add_edge(&hp, 3, 2, 1.0); + + let graph = hp.into_nearest_lists(3); + assert_eq!(graph.len(), 4); + + for (i, neighbors) in graph.iter().enumerate() { + assert!(!neighbors.is_empty(), "point {} has no neighbors", i); + } +} + +#[test] +fn test_reservoir_lazy_allocation() { + let mut res = Reservoir::new(5); + assert!(res.is_empty()); + assert!(res.insert(0, 1, 1.0)); + assert_eq!(res.len(), 1); +} + +#[test] +fn test_reservoir_insert_then_evict_cycle() { + let mut res = Reservoir::new(3); + res.insert(0, 10, 3.0); + res.insert(1, 11, 2.0); + res.insert(2, 12, 1.0); + assert_eq!(res.len(), 3); + assert!(res.insert(3, 13, 0.5)); + assert_eq!(res.len(), 3); + let neighbors = res.neighbors(); + assert!(neighbors.iter().all(|&(_, d)| d <= 2.0)); +} + +#[test] +fn test_reservoir_all_same_hash() { + let mut res = Reservoir::new(5); + res.insert(0, 1, 3.0); + res.insert(0, 2, 2.0); + res.insert(0, 3, 1.0); + assert_eq!(res.len(), 1); + let neighbors = res.neighbors(); + assert_eq!(neighbors[0].0, 3); + assert_eq!(neighbors[0].1, 1.0); +} + +#[test] +fn test_reservoir_all_same_distance() { + let mut res = Reservoir::new(5); + res.insert(0, 1, 1.0); + res.insert(1, 2, 1.0); + res.insert(2, 3, 1.0); + assert_eq!(res.len(), 3); +} + +#[test] +#[allow(clippy::disallowed_methods)] +fn test_hash_prune_parallel_safety() { + use rayon::prelude::*; + let data = vec![0.0f32; 100 * 4]; + let hp = HashPrune::new(&data, 100, 4, 4, 10, 42).unwrap(); + (0..50).into_par_iter().for_each(|i| { + add_edge(&hp, i, (i + 1) % 100, 1.0); + add_edge(&hp, (i + 1) % 100, i, 1.0); + }); + let graph = hp.into_nearest_lists(5); + assert_eq!(graph.len(), 100); +} + +#[test] +fn test_hash_prune_high_degree_limit() { + let data = vec![0.0f32; 10 * 2]; + let hp = HashPrune::new(&data, 10, 2, 4, 10, 42).unwrap(); + for i in 0..10 { + for j in 0..10 { + if i != j { + add_edge(&hp, i, j, (i as f32 - j as f32).abs()); + } + } + } + let graph = hp.into_nearest_lists(1); + for neighbors in &graph { + assert!( + neighbors.len() <= 1, + "max_degree=1 should limit to 1 neighbor" + ); + } +} + +#[test] +fn test_hash_prune_extract_sorted() { + let data = vec![0.0f32; 4 * 2]; + let hp = HashPrune::new(&data, 4, 2, 4, 10, 42).unwrap(); + add_edge(&hp, 0, 1, 3.0); + add_edge(&hp, 0, 2, 1.0); + add_edge(&hp, 0, 3, 2.0); + let graph = hp.into_nearest_lists(3); + assert!(!graph[0].is_empty()); +} + +#[test] +fn test_into_candidate_lists_returns_full_reservoir() { + let data = vec![0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + let hp = HashPrune::new(&data, 4, 2, 4, 10, 42).unwrap(); + add_edge(&hp, 0, 1, 1.0); + add_edge(&hp, 0, 2, 1.0); + add_edge(&hp, 0, 3, 1.414); + add_edge(&hp, 1, 0, 1.0); + add_edge(&hp, 2, 0, 1.0); + add_edge(&hp, 3, 0, 1.414); + + let full = hp.into_candidate_lists(); + assert_eq!(full.len(), 4); + assert!(!full[0].is_empty(), "node 0 should have neighbors"); + // ids-only, unsorted: every id is one of node 0's inserted neighbors + // {1,2,3} with no duplicates (the LSH bucket may keep-closer-collapse a + // colliding pair on this tiny 4-plane sketch, so we don't assert all 3). + let mut n0 = full[0].clone(); + n0.sort_unstable(); + let deduped = { + let mut d = n0.clone(); + d.dedup(); + d + }; + assert_eq!(n0, deduped, "no duplicate ids in a reservoir row"); + assert!( + n0.iter().all(|&id| (1..=3).contains(&id)), + "node 0 ids must be a subset of its inserted neighbors {{1,2,3}}, got {:?}", + n0 + ); +} + +#[test] +fn test_into_nearest_lists_truncates_to_max_degree() { + let data = vec![0.0f32; 4 * 2]; + let hp = HashPrune::new(&data, 4, 2, 4, 10, 42).unwrap(); + add_edge(&hp, 0, 1, 1.0); + add_edge(&hp, 0, 2, 2.0); + add_edge(&hp, 0, 3, 3.0); + + let graph = hp.into_nearest_lists(2); + assert!( + graph[0].len() <= 2, + "bounded graph extraction should truncate to max_degree" + ); +} + +#[test] +fn test_reservoir_farthest_cache_after_eviction() { + let mut res = Reservoir::new(3); + res.insert(0, 10, 5.0); + res.insert(1, 11, 4.0); + res.insert(2, 12, 3.0); + assert!(res.insert(3, 13, 2.0)); + assert!(res.insert(4, 14, 1.0)); + let neighbors = res.neighbors(); + assert_eq!(neighbors.len(), 3); + for &(_, d) in &neighbors { + assert!(d <= 3.1, "expected dist <= 3.0, got {}", d); + } +} + +#[test] +fn test_reservoir_farthest_insert_before_farthest_idx() { + let mut res = Reservoir::new(4); + res.insert(5, 1, 1.0); + res.insert(10, 2, 3.0); + res.insert(15, 3, 2.0); + res.insert(3, 4, 0.5); + let neighbors = res.neighbors(); + assert_eq!(neighbors.len(), 4); + assert_eq!(neighbors[0].0, 4); +} diff --git a/diskann-pipnn/src/leaf_build.rs b/diskann-pipnn/src/leaf_build.rs index 2e48dc224..103415098 100644 --- a/diskann-pipnn/src/leaf_build.rs +++ b/diskann-pipnn/src/leaf_build.rs @@ -84,6 +84,8 @@ pub(crate) enum LeafBuildError { InvalidLocalPosition { position: u32, points: usize }, #[error("candidate row {point} is poisoned")] PoisonedCandidateRow { point: u32 }, + #[error("leaf {leaf} produced too many directed edges")] + TooManyEdges { leaf: usize }, } /// Scratch leased to one Rayon job and reused for successive leaves. @@ -100,6 +102,11 @@ struct LeafBuffers { local_graph: Vec>, top_k: LeafTopKWorkspace, seen_ids: HashSet, + seen_pairs: Vec, + edge_offsets: Vec, + edges: Vec<(u32, f32)>, + edge_cursor: Vec, + sketch_scratch: Vec, } impl LeafBuffers { @@ -141,6 +148,7 @@ impl LeafBuffers { nearest_values, LeafNeighbor::default(), )?; + grow("leaf seen pairs", &mut self.seen_pairs, dot_values, false)?; Ok(actual_k) } @@ -232,6 +240,53 @@ where candidates.into_rows() } +/// Build leaf-local edges and stream them into HashPrune reservoirs. +#[allow(clippy::disallowed_methods)] // The supplied pool owns this terminal operation. +pub(crate) fn add_hash_prune_candidates( + data: MatrixView<'_, T>, + leaves: Vec>, + k: usize, + metric: Metric, + hash_prune: &crate::hash_prune::HashPrune, +) -> Result<(), LeafBuildError> +where + T: VectorRepr + 'static, +{ + if data.ncols() == 0 { + return Err(LeafBuildError::EmptyDimensions); + } + if data.nrows() > u32::MAX as usize { + return Err(LeafBuildError::TooManyPoints(data.nrows())); + } + + leaves.par_iter().enumerate().try_for_each_init( + LeafBuffers::default, + |buffers, (leaf, point_ids)| { + let actual_k = compute_leaf(data, leaf, point_ids, k, metric, buffers)?; + let points = point_ids.len(); + let edge_count = build_symmetric_edge_csr( + leaf, + point_ids, + actual_k, + &buffers.nearest[..points * actual_k], + EdgeBuffers { + seen: &mut buffers.seen_pairs[..points * points], + offsets: &mut buffers.edge_offsets, + edges: &mut buffers.edges, + cursor: &mut buffers.edge_cursor, + }, + )?; + hash_prune.add_leaf_edges( + point_ids, + &buffers.edge_offsets, + &buffers.edges[..edge_count], + &mut buffers.sketch_scratch, + ); + Ok(()) + }, + ) +} + /// Build and publish one leaf's symmetric nearest-neighbor rows. /// /// Validation precedes all dataset indexing. Sorted partition output takes the @@ -247,6 +302,31 @@ fn build_leaf( buffers: &mut LeafBuffers, candidates: &DirectCandidates, ) -> Result<(), LeafBuildError> +where + T: VectorRepr + 'static, +{ + let actual_k = compute_leaf(data, leaf, point_ids, k, metric, buffers)?; + if actual_k == 0 { + return Ok(()); + } + buffers.prepare_local_graph(point_ids.len())?; + add_symmetric_edges( + point_ids, + actual_k, + &buffers.nearest[..point_ids.len() * actual_k], + &mut buffers.local_graph[..point_ids.len()], + )?; + candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) +} + +fn compute_leaf( + data: MatrixView<'_, T>, + leaf: usize, + point_ids: &[u32], + k: usize, + metric: Metric, + buffers: &mut LeafBuffers, +) -> Result where T: VectorRepr + 'static, { @@ -283,7 +363,7 @@ where } let actual_k = buffers.prepare(leaf, point_ids.len(), data.ncols(), k)?; if actual_k == 0 { - return Ok(()); + return Ok(0); } let point_values = point_ids.len() * data.ncols(); @@ -320,15 +400,7 @@ where &mut buffers.top_k, ) .map_err(|source| LeafBuildError::Kernel { leaf, source })?; - - buffers.prepare_local_graph(point_ids.len())?; - add_symmetric_edges( - point_ids, - actual_k, - &buffers.nearest[..nearest_values], - &mut buffers.local_graph[..point_ids.len()], - )?; - candidates.add_leaf(point_ids, &buffers.local_graph[..point_ids.len()]) + Ok(actual_k) } fn add_symmetric_edges( @@ -356,6 +428,124 @@ fn add_symmetric_edges( Ok(()) } +struct EdgeBuffers<'a> { + seen: &'a mut [bool], + offsets: &'a mut Vec, + edges: &'a mut Vec<(u32, f32)>, + cursor: &'a mut Vec, +} + +fn build_symmetric_edge_csr( + leaf: usize, + point_ids: &[u32], + k: usize, + nearest: &[LeafNeighbor], + buffers: EdgeBuffers<'_>, +) -> Result { + let EdgeBuffers { + seen, + offsets, + edges, + cursor, + } = buffers; + let points = point_ids.len(); + if k == 0 { + resize("leaf edge offsets", offsets, points + 1, 0)?; + offsets.fill(0); + edges.clear(); + return Ok(0); + } + seen.fill(false); + resize("leaf edge offsets", offsets, points + 1, 0)?; + offsets.fill(0); + + for (source, neighbors) in nearest.chunks_exact(k).enumerate() { + for neighbor in neighbors { + let target = neighbor.position as usize; + if target >= points { + return Err(LeafBuildError::InvalidLocalPosition { + position: neighbor.position, + points, + }); + } + count_directed_edge(leaf, points, source, target, seen, offsets)?; + count_directed_edge(leaf, points, target, source, seen, offsets)?; + } + } + for index in 1..=points { + offsets[index] = offsets[index] + .checked_add(offsets[index - 1]) + .ok_or(LeafBuildError::TooManyEdges { leaf })?; + } + + let edge_count = offsets[points] as usize; + resize("leaf edges", edges, edge_count, (0, 0.0))?; + resize("leaf edge cursor", cursor, points, 0)?; + cursor.copy_from_slice(&offsets[..points]); + seen.fill(false); + + for (source, neighbors) in nearest.chunks_exact(k).enumerate() { + for neighbor in neighbors { + let target = neighbor.position as usize; + write_directed_edge( + points, + source, + target, + neighbor.distance, + seen, + edges, + cursor, + ); + write_directed_edge( + points, + target, + source, + neighbor.distance, + seen, + edges, + cursor, + ); + } + } + Ok(edge_count) +} + +fn count_directed_edge( + leaf: usize, + points: usize, + source: usize, + target: usize, + seen: &mut [bool], + offsets: &mut [u32], +) -> Result<(), LeafBuildError> { + let seen = &mut seen[source * points + target]; + if !*seen { + *seen = true; + offsets[source + 1] = offsets[source + 1] + .checked_add(1) + .ok_or(LeafBuildError::TooManyEdges { leaf })?; + } + Ok(()) +} + +fn write_directed_edge( + points: usize, + source: usize, + target: usize, + distance: f32, + seen: &mut [bool], + edges: &mut [(u32, f32)], + cursor: &mut [u32], +) { + let seen = &mut seen[source * points + target]; + if !*seen { + *seen = true; + let position = cursor[source] as usize; + edges[position] = (target as u32, distance); + cursor[source] += 1; + } +} + fn grow( buffer: &'static str, values: &mut Vec, diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index 7684871db..f325f14cd 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -16,9 +16,13 @@ //! lifetime boundary and releases each stage's large scratch before the next //! outer pipeline allocation. +mod bf16; +pub(crate) mod cpu_dispatch; mod finalization; +mod hash_prune; mod leaf_build; pub mod leaf_kernel; +mod lsh; pub mod partition_kernel; mod partitioning; @@ -95,6 +99,44 @@ impl PiPNNConfig { } } +/// HashPrune reservoir policy. +#[derive(Clone, Debug, PartialEq)] +pub struct HashPruneConfig { + /// Number of random-hyperplane sketch dimensions. + pub num_hash_planes: usize, + /// Maximum retained candidates per point. + pub l_max: usize, + /// Apply shared Vamana RobustPrune after HashPrune extraction. + pub final_prune: bool, +} + +impl HashPruneConfig { + /// Validate HashPrune's structural bounds. + pub fn validate(&self) -> ANNResult<()> { + if !(1..=lsh::MAX_PLANES).contains(&self.num_hash_planes) { + return Err(config_error(format!( + "num_hash_planes ({}) must be in [1, {}]", + self.num_hash_planes, + lsh::MAX_PLANES + ))); + } + if !(1..=hash_prune::MAX_RESERVOIR_LEN).contains(&self.l_max) { + return Err(config_error(format!( + "l_max ({}) must be in [1, {}]", + self.l_max, + hash_prune::MAX_RESERVOIR_LEN + ))); + } + Ok(()) + } +} + +#[derive(Clone, Debug)] +enum CandidateMerge { + Direct, + HashPrune(HashPruneConfig), +} + /// Validated, borrowed policy and execution context for one PiPNN graph build. #[derive(Debug)] pub struct PiPNNBuildContext<'a> { @@ -102,6 +144,7 @@ pub struct PiPNNBuildContext<'a> { pub(crate) graph: &'a Config, pub(crate) metric: Metric, pub(crate) pool: &'a ThreadPool, + candidate_merge: CandidateMerge, } impl<'a> PiPNNBuildContext<'a> { @@ -131,8 +174,16 @@ impl<'a> PiPNNBuildContext<'a> { graph, metric, pool, + candidate_merge: CandidateMerge::Direct, }) } + + /// Enable HashPrune candidate merging for this build. + pub fn with_hash_prune(mut self, config: HashPruneConfig) -> ANNResult { + config.validate()?; + self.candidate_merge = CandidateMerge::HashPrune(config); + Ok(self) + } } /// Build PiPNN adjacency for real rows in `data`. @@ -188,16 +239,57 @@ where let partition = partitioning::PartitionConfig::from(&context.config); let leaves = tracing::info_span!("pipnn.partition") .in_scope(|| partitioning::partition(data, partition, metric))?; - // `leaves` is consumed here. Workers borrow individual ID rows during the - // parallel pass, and the complete partition allocation drops on return. - let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { - leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) - .map_err(ANNError::opaque) - })?; - // Finalization consumes candidate rows and reuses their allocations for the - // resulting adjacency where possible. - tracing::info_span!("pipnn.finalization") - .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) + match &context.candidate_merge { + CandidateMerge::Direct => { + // `leaves` is consumed here. Workers borrow individual ID rows + // during the parallel pass, and the partition allocation drops on + // return. + let candidates = tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::build_leaf_candidates(data, leaves, context.config.k, metric) + .map_err(ANNError::opaque) + })?; + // Finalization consumes candidate rows and reuses their allocations + // for the resulting adjacency where possible. + tracing::info_span!("pipnn.finalization") + .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) + } + CandidateMerge::HashPrune(config) => { + let hash_prune = hash_prune::HashPrune::new( + data.as_slice(), + data.nrows(), + data.ncols(), + config.num_hash_planes, + config.l_max, + 42, + )?; + tracing::info_span!("pipnn.leaf_build").in_scope(|| { + leaf_build::add_hash_prune_candidates( + data, + leaves, + context.config.k, + metric, + &hash_prune, + ) + .map_err(ANNError::opaque) + })?; + if config.final_prune { + let candidates = rows_to_adjacency(hash_prune.into_candidate_lists()); + tracing::info_span!("pipnn.finalization").in_scope(|| { + finalization::prune_overfull(data, candidates, context.graph, metric) + }) + } else { + Ok(rows_to_adjacency( + hash_prune.into_nearest_lists(context.graph.pruned_degree().get()), + )) + } + } + } +} + +fn rows_to_adjacency(rows: Vec>) -> Vec> { + rows.into_iter() + .map(AdjacencyList::from_iter_untrusted) + .collect() } fn effective_metric(metric: Metric) -> Metric { diff --git a/diskann-pipnn/src/lsh.rs b/diskann-pipnn/src/lsh.rs new file mode 100644 index 000000000..7b344f3c9 --- /dev/null +++ b/diskann-pipnn/src/lsh.rs @@ -0,0 +1,273 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Random-hyperplane LSH (Locality-Sensitive Hashing) for `f32` vectors. +//! +//! Computes `Sketch(v) = [v · H_i for i in 0..num_planes]` where `H_i` are +//! random unit-Gaussian hyperplanes. Callers use these sketches to derive +//! application-specific hashes. +//! +//! Sketches are computed in parallel via rayon, with caller-provided per-point +//! `f32` conversion (so f16, u8, etc. don't need a full upfront f32 copy). +//! `num_planes ≤ 16` so the result fits in `u16`. + +use rand::SeedableRng; +use rand_distr::{Distribution, StandardNormal}; +use rayon::prelude::*; + +/// Maximum number of hyperplanes (the hash output is `u16`). +pub const MAX_PLANES: usize = 16; + +/// Failure while constructing LSH sketches. +#[derive(Debug)] +pub enum LshSketchError { + /// The relative hash must use between one and 16 bits. + InvalidPlaneCount { actual: usize, max: usize }, + /// A matrix shape overflowed `usize`. + ShapeOverflow { rows: usize, columns: usize }, + /// A sketch buffer could not be allocated. + Allocation(std::collections::TryReserveError), + /// The caller could not materialize a point in `f32`. + Fill(E), +} + +impl std::fmt::Display for LshSketchError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InvalidPlaneCount { actual, max } => { + write!(f, "num_planes ({actual}) must be in 1..={max}") + } + Self::ShapeOverflow { rows, columns } => { + write!(f, "LSH matrix shape {rows} x {columns} overflows usize") + } + Self::Allocation(error) => write!(f, "LSH allocation failed: {error}"), + Self::Fill(error) => write!(f, "point conversion failed: {error}"), + } + } +} + +impl std::error::Error for LshSketchError +where + E: std::error::Error + 'static, +{ + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::InvalidPlaneCount { .. } => None, + Self::ShapeOverflow { .. } => None, + Self::Allocation(error) => Some(error), + Self::Fill(error) => Some(error), + } + } +} + +/// Precomputed LSH sketches for `npoints` vectors. +pub struct LshSketches { + num_planes: usize, + /// Row-major `npoints × num_planes`: `sketches[i*m + j] = dot(point_i, plane_j)`. + sketches: Vec, +} + +impl LshSketches { + /// Compute LSH sketches for `npoints` points of dimension `ndims`. + /// + /// `fill_point` is called once per point with the point's destination + /// buffer of length `ndims` (the buffer is allocated per worker thread + /// and reused across calls). This avoids requiring a contiguous `&[f32]` + /// for the whole dataset when the caller stores points as f16, u8, etc. + /// + /// Caller MUST be inside a `rayon::ThreadPool::install(...)` scope — + /// parallel work runs on the current pool. + /// + /// Returns [`LshSketchError::InvalidPlaneCount`] unless `num_planes` fits + /// the non-empty 16-bit relative-hash representation. + pub fn try_new( + npoints: usize, + ndims: usize, + num_planes: usize, + seed: u64, + fill_point: F, + ) -> Result> + where + F: Fn(usize, &mut [f32]) -> Result<(), E> + Send + Sync, + E: Send, + { + if !(1..=MAX_PLANES).contains(&num_planes) { + return Err(LshSketchError::InvalidPlaneCount { + actual: num_planes, + max: MAX_PLANES, + }); + } + + let hyperplane_len = + num_planes + .checked_mul(ndims) + .ok_or(LshSketchError::ShapeOverflow { + rows: num_planes, + columns: ndims, + })?; + let sketch_len = npoints + .checked_mul(num_planes) + .ok_or(LshSketchError::ShapeOverflow { + rows: npoints, + columns: num_planes, + })?; + + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + // Random unit-Gaussian hyperplanes, row-major `num_planes × ndims`. + let mut hyperplanes: Vec = Vec::new(); + hyperplanes + .try_reserve_exact(hyperplane_len) + .map_err(LshSketchError::Allocation)?; + hyperplanes.resize_with(hyperplane_len, || StandardNormal.sample(&mut rng)); + + let mut sketches = Vec::new(); + sketches + .try_reserve_exact(sketch_len) + .map_err(LshSketchError::Allocation)?; + sketches.resize(sketch_len, 0.0f32); + + #[allow(clippy::disallowed_methods)] // see module docstring; caller is in pool.install(). + sketches + .par_chunks_mut(num_planes) + .enumerate() + .try_for_each_init(Vec::new, |buffer, (i, sketch_row)| { + if buffer.len() < ndims { + let additional = ndims - buffer.len(); + buffer + .try_reserve(additional) + .map_err(LshSketchError::Allocation)?; + } + buffer.resize(ndims, 0.0); + fill_point(i, &mut buffer[..ndims]).map_err(LshSketchError::Fill)?; + for j in 0..num_planes { + let plane = &hyperplanes[j * ndims..(j + 1) * ndims]; + let mut dot = 0.0f32; + for d in 0..ndims { + // SAFETY: `d` is in 0..ndims; both buffers have length `ndims`. + unsafe { + dot += *buffer.get_unchecked(d) * *plane.get_unchecked(d); + } + } + sketch_row[j] = dot; + } + Ok(()) + })?; + + Ok(Self { + num_planes, + sketches, + }) + } + + /// Number of hyperplanes (also the number of bits in the hash). + #[inline] + pub fn num_planes(&self) -> usize { + self.num_planes + } + + /// Raw access to the row-major `npoints × num_planes` sketch buffer. + /// Callers can scatter-gather a small per-leaf cache of sketches to avoid + /// touching the multi-hundred-MB global buffer in tight inner loops. + #[inline] + pub fn sketches(&self) -> &[f32] { + &self.sketches + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn build_pool(threads: usize) -> rayon::ThreadPool { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + } + + #[test] + fn computes_expected_sketch_shape() { + let pool = build_pool(2); + let data: Vec = vec![1.0, 0.0, 0.0, 1.0, -1.0, 0.0]; + let sk = pool + .install(|| { + LshSketches::try_new(3, 2, 4, 42, |i, out| { + out.copy_from_slice(&data[i * 2..(i + 1) * 2]); + Ok::<_, std::convert::Infallible>(()) + }) + }) + .unwrap(); + assert_eq!(sk.num_planes(), 4); + assert_eq!(sk.sketches().len(), 12); + } + + #[test] + fn different_seeds_change_hashes() { + let pool = build_pool(2); + let data: Vec = (0..32 * 4).map(|i| (i as f32).sin()).collect(); + let build = |seed| { + pool.install(|| { + LshSketches::try_new(32, 4, 12, seed, |i, out| { + out.copy_from_slice(&data[i * 4..(i + 1) * 4]); + Ok::<_, std::convert::Infallible>(()) + }) + }) + .unwrap() + }; + let sk1 = build(42); + let sk2 = build(99); + assert_ne!(sk1.sketches(), sk2.sketches()); + } + + #[test] + fn try_new_propagates_fill_errors() { + #[derive(Debug, PartialEq)] + struct FillError; + + let pool = build_pool(1); + let error = match pool.install(|| LshSketches::try_new(1, 2, 4, 42, |_, _| Err(FillError))) + { + Ok(_) => panic!("fill error should be propagated"), + Err(error) => error, + }; + + assert!(matches!(error, LshSketchError::Fill(FillError))); + } + + #[test] + fn try_new_reports_too_many_planes() { + let pool = build_pool(1); + let error = match pool.install(|| { + LshSketches::try_new(1, 2, 17, 42, |_, _| Ok::<_, std::convert::Infallible>(())) + }) { + Ok(_) => panic!("too many planes should be rejected"), + Err(error) => error, + }; + + assert!(matches!( + error, + LshSketchError::InvalidPlaneCount { + actual: 17, + max: 16 + } + )); + } + + #[test] + fn try_new_rejects_zero_planes() { + let pool = build_pool(1); + let error = match pool.install(|| { + LshSketches::try_new(1, 2, 0, 42, |_, _| Ok::<_, std::convert::Infallible>(())) + }) { + Ok(_) => panic!("zero planes should be rejected"), + Err(error) => error, + }; + + assert!(matches!( + error, + LshSketchError::InvalidPlaneCount { actual: 0, max: 16 } + )); + } +} diff --git a/diskann-pipnn/tests/build_graph.rs b/diskann-pipnn/tests/build_graph.rs index 04086be8e..5fd3bf36d 100644 --- a/diskann-pipnn/tests/build_graph.rs +++ b/diskann-pipnn/tests/build_graph.rs @@ -4,7 +4,7 @@ */ use diskann::graph::config::{self, MaxDegree}; -use diskann_pipnn::{build_graph, PiPNNBuildContext, PiPNNConfig}; +use diskann_pipnn::{build_graph, HashPruneConfig, PiPNNBuildContext, PiPNNConfig}; use diskann_utils::views::MatrixView; use diskann_vector::distance::Metric; use half::f16; @@ -193,6 +193,44 @@ fn is_deterministic_for_a_fixed_pool_size() { assert_graph_invariants(&first, 96, 8); } +#[test] +fn hash_prune_build_is_deterministic_and_degree_bounded() { + let points = 64; + let dimensions = 4; + let values: Vec = (0..points * dimensions) + .map(|value| ((value * 17 + 3) % 101) as f32) + .collect(); + let data = MatrixView::try_from(values.as_slice(), points, dimensions).unwrap(); + let graph = graph_config(Metric::L2, 8); + let pool = pool(4); + let config = PiPNNConfig { + c_max: 16, + c_min: 4, + p_samp: 0.25, + fanout: vec![3, 2], + k: 3, + replicas: 2, + }; + let hash_prune = HashPruneConfig { + num_hash_planes: 8, + l_max: 16, + final_prune: true, + }; + let build = || { + let context = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(hash_prune.clone()) + .unwrap(); + build_graph(data, &context).unwrap() + }; + + let first = build(); + let second = build(); + + assert_eq!(first, second); + assert_graph_invariants(&first, points, 8); +} + #[test] fn fixed_seed_randomized_sweeps_preserve_graph_invariants() { let mut rng = StdRng::seed_from_u64(0x857a_d38b_44c2_0f11); diff --git a/diskann-pipnn/tests/config.rs b/diskann-pipnn/tests/config.rs index 923faee9e..9e684a2b5 100644 --- a/diskann-pipnn/tests/config.rs +++ b/diskann-pipnn/tests/config.rs @@ -4,7 +4,7 @@ */ use diskann::graph::config::{self, MaxDegree}; -use diskann_pipnn::{PiPNNBuildContext, PiPNNConfig}; +use diskann_pipnn::{HashPruneConfig, PiPNNBuildContext, PiPNNConfig}; use diskann_vector::distance::Metric; fn pipnn_config() -> PiPNNConfig { @@ -103,6 +103,46 @@ fn rejects_each_invalid_algorithm_parameter() { } } +#[test] +fn validates_hash_prune_parameters() { + let graph = graph_config(Metric::L2, 1.2); + let pool = pool(); + for config in [ + HashPruneConfig { + num_hash_planes: 0, + l_max: 64, + final_prune: true, + }, + HashPruneConfig { + num_hash_planes: 17, + l_max: 64, + final_prune: true, + }, + HashPruneConfig { + num_hash_planes: 8, + l_max: 0, + final_prune: true, + }, + HashPruneConfig { + num_hash_planes: 8, + l_max: 256, + final_prune: true, + }, + ] { + let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); + assert!(context.with_hash_prune(config).is_err()); + } + + PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(HashPruneConfig { + num_hash_planes: 8, + l_max: 64, + final_prune: true, + }) + .unwrap(); +} + #[test] fn rejects_graph_policy_for_a_different_metric() { let graph = graph_config(Metric::InnerProduct, 1.2); From bd9b08bcf572cf01d7411180a1cae35c0c770d89 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 06:00:26 +0000 Subject: [PATCH 02/16] build: configure HashPrune integration --- diskann-benchmark/src/index/build.rs | 5 ++- diskann-disk/src/build/builder/build.rs | 11 ++++-- diskann-disk/src/build/builder/build/pipnn.rs | 13 ++++--- .../src/build/builder/build/pipnn/tests.rs | 3 +- .../build/configuration/build_algorithm.rs | 39 +++++++++++++++++++ .../disk_index_build_parameter.rs | 9 ++--- diskann-disk/src/build/configuration/mod.rs | 2 +- diskann-disk/src/build/mod.rs | 4 +- diskann-disk/src/lib.rs | 4 +- 9 files changed, 69 insertions(+), 21 deletions(-) diff --git a/diskann-benchmark/src/index/build.rs b/diskann-benchmark/src/index/build.rs index 326cf5dda..19387a4ba 100644 --- a/diskann-benchmark/src/index/build.rs +++ b/diskann-benchmark/src/index/build.rs @@ -161,8 +161,11 @@ where let started = std::time::Instant::now(); let adjacency = { - let context = + let mut context = diskann_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_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 20a01949d..7d519f62f 100644 --- a/diskann-disk/src/build/builder/build.rs +++ b/diskann-disk/src/build/builder/build.rs @@ -76,8 +76,11 @@ 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_pipnn::PiPNNConfig::from(parameters).validate()?; + if let Some(hash_prune) = ¶meters.hash_prune { + diskann_pipnn::HashPruneConfig::from(hash_prune).validate()?; + } } let pq_storage = PQStorage::new( @@ -181,8 +184,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 09ad6a760..538787eb2 100644 --- a/diskann-disk/src/build/builder/build/pipnn.rs +++ b/diskann-disk/src/build/builder/build/pipnn.rs @@ -26,7 +26,7 @@ //! search and loading cannot distinguish which builder produced the graph. use diskann::{utils::VectorRepr, ANNError, ANNResult}; -use diskann_pipnn::{PiPNNBuildContext, PiPNNConfig}; +use diskann_pipnn::PiPNNBuildContext; use diskann_providers::{ storage::{save_adjacency_graph, StorageReadProvider, StorageWriteProvider}, utils::{find_medoid_with_sampling, RayonThreadPoolRef, MAX_MEDOID_SAMPLE_SIZE}, @@ -34,13 +34,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, @@ -70,12 +70,15 @@ where // before the outer disk pipeline continues. let data = read_bin::(&mut builder.storage_provider.open_reader(&data_path)?)?; - let context = PiPNNBuildContext::new( - config, + let mut context = PiPNNBuildContext::new( + parameters.into(), &builder.index_configuration.config, builder.index_configuration.dist_metric, pool.as_rayon(), )?; + if let Some(hash_prune) = ¶meters.hash_prune { + context = context.with_hash_prune(hash_prune.into())?; + } let adjacency = diskann_pipnn::build_graph(data.as_view(), &context)?; // Start-node policy belongs to the persisted index, not the core graph diff --git a/diskann-disk/src/build/builder/build/pipnn/tests.rs b/diskann-disk/src/build/builder/build/pipnn/tests.rs index 8edb57f84..c1edd4401 100644 --- a/diskann-disk/src/build/builder/build/pipnn/tests.rs +++ b/diskann-disk/src/build/builder/build/pipnn/tests.rs @@ -33,6 +33,7 @@ fn pipnn() -> PiPNNParameters { fanout: vec![10, 3], k: 2, replicas: 1, + hash_prune: Some(crate::HashPruneParameters::default()), } } @@ -141,7 +142,7 @@ fn pipnn_graph_adapter_writes_real_point_header() { 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( diff --git a/diskann-disk/src/build/configuration/build_algorithm.rs b/diskann-disk/src/build/configuration/build_algorithm.rs index f798138a1..bfab5aed8 100644 --- a/diskann-disk/src/build/configuration/build_algorithm.rs +++ b/diskann-disk/src/build/configuration/build_algorithm.rs @@ -29,6 +29,43 @@ pub struct PiPNNParameters { pub k: usize, /// Number of independent partition passes. pub replicas: usize, + /// HashPrune policy, or `None` for exact direct-candidate accumulation. + pub hash_prune: Option, +} + +/// JSON-facing HashPrune parameters. +#[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 candidates retained per point. + pub l_max: usize, + /// Apply shared 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_pipnn::HashPruneConfig { + fn from(config: &HashPruneParameters) -> Self { + Self { + num_hash_planes: config.num_hash_planes, + l_max: config.l_max, + final_prune: config.final_prune, + } + } } #[cfg(feature = "pipnn")] @@ -41,6 +78,7 @@ impl Default for PiPNNParameters { fanout: vec![8, 3], k: 2, replicas: 1, + hash_prune: Some(HashPruneParameters::default()), } } } @@ -116,6 +154,7 @@ mod tests { assert_eq!(config.fanout, [10, 3]); assert_eq!(config.k, 3); assert_eq!(config.replicas, 1); + assert_eq!(config.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 f09af881d..8db135d3c 100644 --- a/diskann-disk/src/build/configuration/disk_index_build_parameter.rs +++ b/diskann-disk/src/build/configuration/disk_index_build_parameter.rs @@ -196,9 +196,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, } } @@ -290,11 +290,10 @@ mod dataset_test { fn new_pipnn_uses_the_common_disk_pipeline_parameters() { let pq = NumPQChunks::new_with(1, 128).unwrap(); let parameters = PiPNNParameters::default(); - let config = diskann_pipnn::PiPNNConfig::from(¶meters); let budget = MemoryBudget::try_from_gb(2.0).unwrap(); - let params = DiskIndexBuildParameters::new_pipnn(budget, pq, parameters); + let params = DiskIndexBuildParameters::new_pipnn(budget, pq, parameters.clone()); - assert_eq!(params.pipnn_config(), Some(config)); + assert_eq!(params.pipnn_parameters(), Some(¶meters)); assert_eq!(params.build_memory_limit(), budget); assert_eq!(params.search_pq_chunks(), pq); assert_eq!( diff --git a/diskann-disk/src/build/configuration/mod.rs b/diskann-disk/src/build/configuration/mod.rs index a7e343fb5..d2a26bba4 100644 --- a/diskann-disk/src/build/configuration/mod.rs +++ b/diskann-disk/src/build/configuration/mod.rs @@ -5,7 +5,7 @@ pub mod build_algorithm; pub use build_algorithm::BuildAlgorithm; #[cfg(feature = "pipnn")] -pub use build_algorithm::PiPNNParameters; +pub use build_algorithm::{HashPruneParameters, PiPNNParameters}; pub mod disk_index_build_parameter; pub use disk_index_build_parameter::{DiskIndexBuildParameters, MemoryBudget, NumPQChunks}; diff --git a/diskann-disk/src/build/mod.rs b/diskann-disk/src/build/mod.rs index 27f4c124a..c3f304664 100644 --- a/diskann-disk/src/build/mod.rs +++ b/diskann-disk/src/build/mod.rs @@ -12,9 +12,9 @@ pub mod builder; pub mod configuration; // Re-export key types for convenience -#[cfg(feature = "pipnn")] -pub use configuration::PiPNNParameters; pub use configuration::{ disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters, QuantizationType, }; +#[cfg(feature = "pipnn")] +pub use configuration::{HashPruneParameters, PiPNNParameters}; diff --git a/diskann-disk/src/lib.rs b/diskann-disk/src/lib.rs index a7d3d29a5..13e1c2f4b 100644 --- a/diskann-disk/src/lib.rs +++ b/diskann-disk/src/lib.rs @@ -12,12 +12,12 @@ pub(crate) mod test_utils; pub mod build; -#[cfg(feature = "pipnn")] -pub use build::PiPNNParameters; pub use build::{ disk_index_build_parameter, filter_parameter, BuildAlgorithm, DiskIndexBuildParameters, QuantizationType, }; +#[cfg(feature = "pipnn")] +pub use build::{HashPruneParameters, PiPNNParameters}; pub mod data_model; pub mod search; From 3a10d35523c222b0a67f9545c5571ef30ee432c4 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 06:02:51 +0000 Subject: [PATCH 03/16] bench(pipnn): compare candidate merge modes --- diskann-pipnn/Cargo.toml | 4 ++ diskann-pipnn/benches/hash_prune.rs | 87 +++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+) create mode 100644 diskann-pipnn/benches/hash_prune.rs diff --git a/diskann-pipnn/Cargo.toml b/diskann-pipnn/Cargo.toml index 8405eb5b7..42408ca0d 100644 --- a/diskann-pipnn/Cargo.toml +++ b/diskann-pipnn/Cargo.toml @@ -36,5 +36,9 @@ harness = false name = "core" harness = false +[[bench]] +name = "hash_prune" +harness = false + [lints] workspace = true diff --git a/diskann-pipnn/benches/hash_prune.rs b/diskann-pipnn/benches/hash_prune.rs new file mode 100644 index 000000000..1038af418 --- /dev/null +++ b/diskann-pipnn/benches/hash_prune.rs @@ -0,0 +1,87 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{hint::black_box, time::Duration}; + +use criterion::{criterion_group, criterion_main, Criterion, Throughput}; +use diskann::graph::config::{self, MaxDegree}; +use diskann_pipnn::{build_graph, HashPruneConfig, PiPNNBuildContext, PiPNNConfig}; +use diskann_utils::views::MatrixView; +use diskann_vector::distance::Metric; + +const POINTS: usize = 2_048; +const DIMENSIONS: usize = 128; + +fn data() -> Vec { + (0..POINTS * DIMENSIONS) + .map(|index| { + ((index.wrapping_mul(1_664_525).wrapping_add(1_013_904_223) % 2_003) as f32 - 1_001.0) + / 1_001.0 + }) + .collect() +} + +fn benchmark_candidate_merge(c: &mut Criterion) { + let values = data(); + let view = MatrixView::try_from(values.as_slice(), POINTS, DIMENSIONS).unwrap(); + let graph = + config::Builder::new_with(64, MaxDegree::same(), 72, Metric::L2.into(), |builder| { + builder.alpha(1.2); + }) + .build() + .unwrap(); + let config = PiPNNConfig { + c_max: 256, + c_min: 64, + p_samp: 0.01, + fanout: vec![4, 2], + k: 3, + replicas: 1, + }; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap(); + let direct = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool).unwrap(); + let pruned = PiPNNBuildContext::new(config.clone(), &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(HashPruneConfig { + num_hash_planes: 12, + l_max: 64, + final_prune: true, + }) + .unwrap(); + let reservoir_only = PiPNNBuildContext::new(config, &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(HashPruneConfig { + num_hash_planes: 12, + l_max: 64, + final_prune: false, + }) + .unwrap(); + + let mut group = c.benchmark_group("pipnn/candidate-merge"); + group.throughput(Throughput::Elements(POINTS as u64)); + for (name, context) in [ + ("direct", direct), + ("hash-prune+robust-prune", pruned), + ("hash-prune-only", reservoir_only), + ] { + group.bench_function(name, |bencher| { + bencher.iter(|| black_box(build_graph(view, &context).unwrap())); + }); + } + group.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(10) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(5)); + targets = benchmark_candidate_merge +} +criterion_main!(benches); From b589fd8faac99bf1b76c5d2eeacd2f6fb3b17a42 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 09:50:43 +0000 Subject: [PATCH 04/16] refactor(pipnn): dispatch HashPrune through diskann-wide --- diskann-pipnn/src/cpu_dispatch.rs | 73 ----- diskann-pipnn/src/hash_prune.rs | 324 ++++++++++----------- diskann-pipnn/src/hash_prune/tests.rs | 176 +++++++++-- diskann-pipnn/src/leaf_build/tests.rs | 64 +++- diskann-pipnn/src/lib.rs | 1 - diskann-pipnn/src/lsh.rs | 57 +++- diskann-pipnn/tests/build_graph.rs | 13 +- diskann-vector/src/lib.rs | 3 +- diskann-wide/src/arch/x86_64/v3/i16x16_.rs | 17 +- diskann-wide/src/doubled.rs | 7 + 10 files changed, 439 insertions(+), 296 deletions(-) delete mode 100644 diskann-pipnn/src/cpu_dispatch.rs diff --git a/diskann-pipnn/src/cpu_dispatch.rs b/diskann-pipnn/src/cpu_dispatch.rs deleted file mode 100644 index cfe7e2b30..000000000 --- a/diskann-pipnn/src/cpu_dispatch.rs +++ /dev/null @@ -1,73 +0,0 @@ -/* - * Copyright (c) Microsoft Corporation. - * Licensed under the MIT license. - */ - -//! Cached CPU capabilities used by PiPNN's private SIMD kernels. -//! -//! Callers ask for the vector width required by a kernel family instead of -//! matching ISA names themselves. This keeps feature detection in one place -//! while allowing each family to request only the features it actually uses. - -use std::sync::OnceLock; - -#[derive(Copy, Clone, Debug, PartialEq, Eq)] -pub(crate) enum VectorWidth { - Wide, - Narrow, - Scalar, -} - -#[derive(Copy, Clone, Debug)] -struct CpuFeatures { - avx512f: bool, - avx512bw: bool, - avx2: bool, -} - -static FEATURES: OnceLock = OnceLock::new(); - -fn features() -> CpuFeatures { - *FEATURES.get_or_init(|| { - #[cfg(target_arch = "x86_64")] - let detected = CpuFeatures { - avx512f: std::is_x86_feature_detected!("avx512f"), - avx512bw: std::is_x86_feature_detected!("avx512bw"), - avx2: std::is_x86_feature_detected!("avx2"), - }; - - #[cfg(not(target_arch = "x86_64"))] - let detected = CpuFeatures { - avx512f: false, - avx512bw: false, - avx2: false, - }; - - tracing::info!(?detected, "PiPNN CPU features detected"); - detected - }) -} - -/// Width for kernels that use only packed `f32` arithmetic. -pub(crate) fn f32_width() -> VectorWidth { - let features = features(); - if features.avx512f { - VectorWidth::Wide - } else if features.avx2 { - VectorWidth::Narrow - } else { - VectorWidth::Scalar - } -} - -/// Width for packed 16-bit integer kernels. -pub(crate) fn u16_width() -> VectorWidth { - let features = features(); - if features.avx512f && features.avx512bw { - VectorWidth::Wide - } else if features.avx2 { - VectorWidth::Narrow - } else { - VectorWidth::Scalar - } -} diff --git a/diskann-pipnn/src/hash_prune.rs b/diskann-pipnn/src/hash_prune.rs index 7df3bd516..8d5db68df 100644 --- a/diskann-pipnn/src/hash_prune.rs +++ b/diskann-pipnn/src/hash_prune.rs @@ -29,6 +29,12 @@ use crate::bf16::{bf16_to_f32, f32_to_bf16}; use crate::lsh::{LshSketchError, LshSketches}; use bytemuck::Pod; use diskann::{utils::VectorRepr, ANNError, ANNResult}; +use diskann_vector::prefetch_hint_all; +use diskann_wide::{ + arch::{self, Dispatched1, FTarget1, Target}, + lifetime::As, + Architecture, SIMDMask, SIMDPartialEq, SIMDPartialOrd, SIMDVector, +}; use rayon::prelude::*; /// Owned slab allocated via direct `mmap(MAP_PRIVATE | MAP_ANONYMOUS)`. The @@ -363,79 +369,93 @@ fn round_up_to_32(n: usize) -> usize { // ─── find_hash SIMD: 32-way u16 compare ─────────────────────────────────────── -type FindHash = unsafe fn(*const u16, usize, u8, u16) -> Option; -type RelativeHash = unsafe fn(*const f32, *const f32, usize) -> u16; +#[derive(Clone, Copy)] +struct FindHashArgs { + hashes: *const u16, + scan_lanes: usize, + len: u8, + target: u16, +} -fn select_find_hash() -> FindHash { - match crate::cpu_dispatch::u16_width() { - #[cfg(target_arch = "x86_64")] - crate::cpu_dispatch::VectorWidth::Wide => find_hash_wide, - #[cfg(target_arch = "x86_64")] - crate::cpu_dispatch::VectorWidth::Narrow => find_hash_narrow, - _ => find_hash_scalar, +#[derive(Clone, Copy)] +struct RelativeHashArgs { + src: *const f32, + dst: *const f32, + len: usize, +} + +type FindHash = Dispatched1, As>; +type RelativeHash = Dispatched1>; + +struct FindHashKernel; +struct RelativeHashKernel; +struct SelectFindHash; +struct SelectRelativeHash; + +impl Target for SelectFindHash +where + A: Architecture, + FindHashKernel: FTarget1, FindHashArgs>, +{ + fn run(self, arch: A) -> FindHash { + arch.dispatch1::, As>() + } +} + +impl Target for SelectRelativeHash +where + A: Architecture, + RelativeHashKernel: FTarget1, +{ + fn run(self, arch: A) -> RelativeHash { + arch.dispatch1::>() } } +fn select_find_hash() -> FindHash { + arch::dispatch(SelectFindHash) +} + fn select_relative_hash() -> RelativeHash { - match crate::cpu_dispatch::f32_width() { - #[cfg(target_arch = "x86_64")] - crate::cpu_dispatch::VectorWidth::Wide => relative_hash_local_wide, - #[cfg(target_arch = "x86_64")] - crate::cpu_dispatch::VectorWidth::Narrow => relative_hash_local_narrow, - _ => relative_hash_local_scalar, + arch::dispatch(SelectRelativeHash) +} + +impl FTarget1, FindHashArgs> for FindHashKernel +where + A: Architecture, + A::i16x32: SIMDPartialEq, +{ + fn run(arch: A, args: FindHashArgs) -> Option { + find_hash_simd::(arch, args) } } -/// SAFETY: caller guarantees AVX-512F + AVX-512BW at runtime and `hashes` -/// spans `scan_lanes` elements. -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx512f", enable = "avx512bw")] -unsafe fn find_hash_wide( - hashes: *const u16, - scan_lanes: usize, - len: u8, - target: u16, -) -> Option { - use std::arch::x86_64::*; - let len = len as usize; - let t = _mm512_set1_epi16(target as i16); - // Scan only the chunks that cover valid entries (len), not the full - // scan_lanes capacity. At avg_deg ~60 (l_max 128) this halves the scan. - let chunks = len.div_ceil(32).min(scan_lanes / 32); - for chunk in 0..chunks { - let v = _mm512_loadu_si512(hashes.add(chunk * 32) as *const __m512i); - let mask = _mm512_cmpeq_epi16_mask(v, t); - if mask != 0 { - let lane = chunk * 32 + mask.trailing_zeros() as usize; - if lane < len { - return Some(lane); - } - } +impl FTarget1 for RelativeHashKernel +where + A: Architecture, + A::f32x16: SIMDPartialOrd + std::ops::Sub, + u64: From<<<::Mask as SIMDMask>::BitMask as SIMDMask>::Underlying>, +{ + fn run(arch: A, args: RelativeHashArgs) -> u16 { + relative_hash_simd::(arch, args) } - None } -/// SAFETY: caller guarantees AVX2 at runtime. AVX2 is the workspace baseline -/// (`x86-64-v3`) so this path covers every supported deployment. -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx2")] -unsafe fn find_hash_narrow( - hashes: *const u16, - scan_lanes: usize, - len: u8, - target: u16, -) -> Option { - use std::arch::x86_64::*; - let len = len as usize; - let t = _mm256_set1_epi16(target as i16); - // Scan only the chunks covering valid entries (len), not full scan_lanes. - let chunks = len.div_ceil(16).min(scan_lanes / 16); +/// Hashes are compared for equality only, so the lanes are loaded as `i16` +/// (diskann-wide has no `u16` vector type) — the bit patterns, and therefore +/// the equality result, are identical. +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.scan_lanes / F::LANES); for chunk in 0..chunks { - let v = _mm256_loadu_si256(hashes.add(chunk * 16) as *const __m256i); - let m = _mm256_cmpeq_epi16(v, t); - let bits = _mm256_movemask_epi8(m) as u32; - if bits != 0 { - let lane = chunk * 16 + (bits.trailing_zeros() as usize) / 2; + // SAFETY: every full load stays inside the padded `scan_lanes` 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); } @@ -444,74 +464,29 @@ unsafe fn find_hash_narrow( None } -/// SAFETY: `hashes` must point at `len` valid `u16` slots. -#[inline(always)] -unsafe fn find_hash_scalar( - hashes: *const u16, - _scan_lanes: usize, - len: u8, - target: u16, -) -> Option { - (0..len as usize).find(|&i| *hashes.add(i) == target) -} - -// ─── relative_hash_local: AoSoA-friendly per-pair sketch comparison ─────────── -// -// `relative_hash_local(src, dst, m)` returns the `m`-bit pattern formed by -// `sign(dst[j] - src[j])` for j in 0..m. Used as the in-leaf LSH bucket for HP -// insertion (the "local sketch" cache hits L1). -// -// The implementation is selected once when HashPrune is constructed. - -/// SAFETY: `src` / `dst` must point at `m` valid `f32` slots, `m <= 16`. -unsafe fn relative_hash_local_scalar(src: *const f32, dst: *const f32, m: usize) -> u16 { - let mut h: u16 = 0; - for j in 0..m { - let diff = *dst.add(j) - *src.add(j); - let bit = ((!diff.is_sign_negative()) as u16) << j; - h |= bit; - } - h -} - -/// AVX2-vectorized `relative_hash_local`. Matches the scalar kernel's -/// `!is_sign_negative` semantics exactly (sign bit of `dst-src`), so the -/// produced graph is bit-identical to the scalar path it replaces on AVX2. -/// -/// SAFETY: caller guarantees AVX2 at runtime. `m <= 16` enforced upstream. -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx2")] -unsafe fn relative_hash_local_narrow(src: *const f32, dst: *const f32, m: usize) -> u16 { - use std::arch::x86_64::*; - let mut h: u16 = 0; - let mut j = 0usize; - while j + 8 <= m { - let d = _mm256_sub_ps(_mm256_loadu_ps(dst.add(j)), _mm256_loadu_ps(src.add(j))); - // movemask = sign bits of (dst-src); we want the NON-negative lanes - // (sign clear), i.e. `!is_sign_negative`, matching the scalar kernel. - let signs = _mm256_movemask_ps(d) as u16; - h |= (!signs & 0xFF) << j; - j += 8; - } - while j < m { - let diff = *dst.add(j) - *src.add(j); - h |= ((!diff.is_sign_negative()) as u16) << j; - j += 1; - } - h -} - -/// SAFETY: caller guarantees AVX-512F at runtime. `m <= 16` enforced upstream. -#[cfg(target_arch = "x86_64")] -#[target_feature(enable = "avx512f")] -unsafe fn relative_hash_local_wide(src: *const f32, dst: *const f32, m: usize) -> u16 { - use std::arch::x86_64::*; - let kmask: u16 = if m >= 16 { 0xFFFF } else { (1u16 << m) - 1 }; - let dst_v = _mm512_maskz_loadu_ps(kmask, dst); - let src_v = _mm512_maskz_loadu_ps(kmask, src); - let diff = _mm512_sub_ps(dst_v, src_v); - let mask = _mm512_cmp_ps_mask::<_CMP_GE_OQ>(diff, _mm512_setzero_ps()); - mask & kmask +/// Bit `j` of the returned hash is `dst[j] - src[j] >= 0.0`; equality, +/// including signed zero, hashes as non-negative on every backend. +fn relative_hash_simd(arch: F::Arch, args: RelativeHashArgs) -> u16 +where + F: SIMDVector + SIMDPartialOrd + std::ops::Sub, + u64: From<<::BitMask as SIMDMask>::Underlying>, +{ + debug_assert!(args.len <= F::LANES); + debug_assert!(F::LANES <= u16::BITS as usize); + + // SAFETY: the first `len` lanes are valid and masked loads do not access + // inactive lanes. HashPrune limits sketches to 16 planes. + let dst = unsafe { F::load_simd_first(arch, args.dst, args.len) }; + // SAFETY: as above. + let src = unsafe { F::load_simd_first(arch, args.src, args.len) }; + let bits = u64::from( + (dst - src) + .ge_simd(F::splat(arch, 0.0)) + .bitmask() + .to_underlying(), + ); + let active = ((1_u32 << args.len) - 1) as u16; + bits as u16 & active } // ─── Per-reservoir mutation helpers (caller holds lock) ─────────────────────── @@ -555,7 +530,8 @@ unsafe fn update_farthest(hot: &mut HotSlot, cold: ColdSlotPtrs) { let mut max_dist: u16 = 0; let mut max_idx: u8 = 0; for i in 0..hot.len as usize { - let d = *cold.distances.add(i); + // SAFETY: guaranteed by this function's contract. + let d = unsafe { *cold.distances.add(i) }; if d > max_dist { max_dist = d; max_idx = i as u8; @@ -583,13 +559,23 @@ unsafe fn insert_locked( return false; } - if let Some(idx) = find_hash(cold.hashes, cold.scan_lanes, hot.len, hash) { - if dist_key < *cold.distances.add(idx) { + if let Some(idx) = find_hash.call(FindHashArgs { + hashes: cold.hashes, + scan_lanes: cold.scan_lanes, + len: hot.len, + target: hash, + }) { + // SAFETY: `idx < hot.len <= cold.scan_lanes`. + if dist_key < unsafe { *cold.distances.add(idx) } { let was_farthest = idx == hot.farthest_idx as usize; - *cold.neighbors.add(idx) = neighbor; - *cold.distances.add(idx) = dist_key; + // SAFETY: as above; the caller holds the slot lock. + unsafe { + *cold.neighbors.add(idx) = neighbor; + *cold.distances.add(idx) = dist_key; + } if was_farthest { - update_farthest(hot, cold); + // SAFETY: this function's contract provides the same invariants. + unsafe { update_farthest(hot, cold) }; } return true; } @@ -598,9 +584,12 @@ unsafe fn insert_locked( if hot.len < l_max { let new_idx = hot.len as usize; - *cold.hashes.add(new_idx) = hash; - *cold.distances.add(new_idx) = dist_key; - *cold.neighbors.add(new_idx) = neighbor; + // SAFETY: `new_idx < l_max <= cold.scan_lanes`; the caller holds the lock. + unsafe { + *cold.hashes.add(new_idx) = hash; + *cold.distances.add(new_idx) = dist_key; + *cold.neighbors.add(new_idx) = neighbor; + } hot.len += 1; if dist_key >= hot.farthest_dist { hot.farthest_dist = dist_key; @@ -611,10 +600,13 @@ unsafe fn insert_locked( if dist_key < hot.farthest_dist { let idx = hot.farthest_idx as usize; - *cold.hashes.add(idx) = hash; - *cold.distances.add(idx) = dist_key; - *cold.neighbors.add(idx) = neighbor; - update_farthest(hot, cold); + // SAFETY: `idx < hot.len <= cold.scan_lanes`; the caller holds the lock. + unsafe { + *cold.hashes.add(idx) = hash; + *cold.distances.add(idx) = dist_key; + *cold.neighbors.add(idx) = neighbor; + update_farthest(hot, cold); + } return true; } false @@ -642,7 +634,8 @@ unsafe fn collect_sorted_neighbors( scratch.clear(); scratch.reserve(n); for i in 0..n { - scratch.push((*neighbors.add(i), *distances.add(i))); + // SAFETY: guaranteed by this function's contract. + scratch.push(unsafe { (*neighbors.add(i), *distances.add(i)) }); } scratch.sort_unstable_by_key(|&(_, d)| d); let out_len = n.min(cap); @@ -666,7 +659,8 @@ unsafe fn collect_neighbor_ids(hot: &HotSlot, neighbors: *const u32, cap: usize) let out_len = (hot.len as usize).min(cap); let mut out = Vec::with_capacity(out_len); for i in 0..out_len { - out.push(*neighbors.add(i)); + // SAFETY: guaranteed by this function's contract. + out.push(unsafe { *neighbors.add(i) }); } out } @@ -870,34 +864,13 @@ impl HashPrune { let global_src = point_ids[local_src] as usize; // Prefetch the next non-empty source's hot and cold slots. - #[cfg(all(target_arch = "x86_64", target_feature = "sse"))] + if let Some(next) = (local_src + 1..n) + .find(|&i| edge_offsets[i] != edge_offsets[i + 1]) + .map(|i| point_ids[i] as usize) { - use std::arch::x86_64::*; - if local_src + 1 < n { - let mut nxt = local_src + 1; - while nxt < n && edge_offsets[nxt] == edge_offsets[nxt + 1] { - nxt += 1; - } - if nxt < n { - let nxt_global = point_ids[nxt] as usize; - let off = nxt_global * self.scan_lanes; - // SAFETY: nxt_global < npoints; off + scan_lanes - // bounded by slab capacity. Prefetch is non-fatal. - unsafe { - let hot_p = self.hot[nxt_global].get().cast::(); - _mm_prefetch::<{ _MM_HINT_T0 }>(hot_p); - // hashes array covers `scan_lanes` u16 = scan_lanes*2 bytes; - // prefetch in 64-byte cache-line strides. - let hashes_p = self.cold_hashes.as_ptr().add(off) as *const i8; - let hashes_bytes = self.scan_lanes * std::mem::size_of::(); - let mut b = 0usize; - while b < hashes_bytes { - _mm_prefetch::<{ _MM_HINT_T0 }>(hashes_p.add(b)); - b += 64; - } - } - } - } + let off = next * self.scan_lanes; + prefetch_hint_all(std::slice::from_ref(&self.hot[next])); + prefetch_hint_all(&self.cold_hashes[off..off + self.scan_lanes]); } let src_sketch = &sketch_scratch[local_src * m..(local_src + 1) * m]; @@ -907,10 +880,11 @@ impl HashPrune { let dst_sketch = &sketch_scratch[dst_local as usize * m..(dst_local as usize + 1) * m]; debug_assert!(m <= 16, "num_planes <= 16 enforced by validate"); - // SAFETY: m <= 16, sketches are m-elt slices. - let hash = unsafe { - (self.relative_hash)(src_sketch.as_ptr(), dst_sketch.as_ptr(), m) - }; + let hash = self.relative_hash.call(RelativeHashArgs { + src: src_sketch.as_ptr(), + dst: dst_sketch.as_ptr(), + len: m, + }); // SAFETY: cold ptrs from with_locked are valid for scan_lanes. unsafe { insert_locked(hot, cold, hash, global_dst, dist, l_max, self.find_hash) diff --git a/diskann-pipnn/src/hash_prune/tests.rs b/diskann-pipnn/src/hash_prune/tests.rs index a14ae3b9c..01340211a 100644 --- a/diskann-pipnn/src/hash_prune/tests.rs +++ b/diskann-pipnn/src/hash_prune/tests.rs @@ -71,14 +71,11 @@ impl Reservoir { fn add_edge(hp: &HashPrune, src: usize, dst: usize, distance: f32) { let m = hp.sketches.num_planes(); let sketches = hp.sketches.sketches(); - // SAFETY: both offsets select complete `m`-element sketch rows. - let hash = unsafe { - (hp.relative_hash)( - sketches.as_ptr().add(src * m), - sketches.as_ptr().add(dst * m), - m, - ) - }; + 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(src, |hot, cold| { // SAFETY: with_locked guards the row and supplies valid cold-slab pointers. @@ -87,33 +84,154 @@ fn add_edge(hp: &HashPrune, src: usize, dst: usize, distance: f32) { } #[test] -fn test_relative_hash_local_narrow_matches_scalar_semantics() { - #[cfg(target_arch = "x86_64")] - { - if !is_x86_feature_detected!("avx2") { - return; +fn test_relative_hash_matches_numeric_reference() { + let dispatched = select_relative_hash(); + + let src = [ + 1.0, -2.0, 0.0, 7.5, -0.0, 3.25, -9.0, 4.0, 8.0, -1.5, 2.0, 0.0, 6.0, -3.0, 5.5, -7.25, + ]; + let dst = [ + 1.0, -3.0, 0.5, 7.0, 0.0, 3.25, -8.0, -4.0, 9.0, -1.5, -2.0, -0.0, 5.0, -2.0, 5.5, -8.0, + ]; + + for m in 0..=16 { + let mut expected = 0u16; + for j in 0..m { + let diff: f32 = dst[j] - src[j]; + expected |= ((diff >= 0.0) as u16) << j; } - let src = [ - 1.0, -2.0, 0.0, 7.5, -0.0, 3.25, -9.0, 4.0, 8.0, -1.5, 2.0, 0.0, 6.0, -3.0, 5.5, -7.25, - ]; - let dst = [ - 1.0, -3.0, 0.5, 7.0, 0.0, 3.25, -8.0, -4.0, 9.0, -1.5, -2.0, -0.0, 5.0, -2.0, 5.5, -8.0, - ]; - - for m in 0..=16 { - let mut expected = 0u16; - for j in 0..m { - let diff: f32 = dst[j] - src[j]; - expected |= ((!diff.is_sign_negative()) as u16) << j; + let actual = dispatched.call(RelativeHashArgs { + src: src.as_ptr(), + dst: dst.as_ptr(), + len: m, + }); + assert_eq!(actual, expected, "m={m}"); + } +} + +#[test] +fn test_relative_hash_defines_signed_zero_and_nan_buckets() { + let src = [0.0; 4]; + let dst = [ + 0.0, + -0.0, + f32::from_bits(0x7FC0_0000), + f32::from_bits(0xFFC0_0000), + ]; + + assert_eq!( + select_relative_hash().call(RelativeHashArgs { + src: src.as_ptr(), + dst: dst.as_ptr(), + len: dst.len(), + }), + 0b0011 + ); +} + +#[test] +fn test_find_hash_handles_padded_boundaries_and_all_bit_patterns() { + let dispatched = select_find_hash(); + + for target in [0, 0xF00D] { + for len in [0, 1, 15, 16, 17, 31, 32, 33, 63, 64, 65, 254, 255] { + let scan_lanes = round_up_to_32(len.max(1)); + let mut hashes = vec![target; scan_lanes]; + hashes[..len].fill(0x8001); + let args = |hashes: &[u16]| FindHashArgs { + hashes: hashes.as_ptr(), + scan_lanes, + 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; + } } + } + } +} - // SAFETY: the feature guard above checked AVX2 and both arrays - // contain 16 elements, so every tested prefix is in bounds. - let actual = unsafe { relative_hash_local_narrow(src.as_ptr(), dst.as_ptr(), m) }; - assert_eq!(actual, expected, "m={m}"); +#[test] +fn test_slab_is_zeroed_and_reports_its_bytes() { + let slab = MmapSlab::::new_zeroed(4).unwrap(); + assert_eq!(slab.bytes(), 4 * std::mem::size_of::()); + assert_eq!(slab.len(), 4); + assert!(!slab.as_ptr().is_null()); + assert_eq!(&*slab, &[0; 4]); +} + +#[test] +fn test_round_up_to_32_boundaries() { + assert_eq!(round_up_to_32(0), 0); + assert_eq!(round_up_to_32(1), 32); + assert_eq!(round_up_to_32(32), 32); + assert_eq!(round_up_to_32(33), 64); +} + +#[test] +fn test_ordered_key_roundtrips_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_key).collect(); + assert!(keys.windows(2).all(|pair| pair[0] <= pair[1])); + for (value, key) in values.into_iter().zip(keys) { + assert_eq!( + bf16_to_f32(key_to_bf16(key)), + bf16_to_f32(f32_to_bf16(value)) + ); + } +} + +#[test] +fn test_add_leaf_edges_matches_single_edge_reference() { + let data = [0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + let batched = HashPrune::new(&data, 4, 2, 8, 8, 42).unwrap(); + let reference = HashPrune::new(&data, 4, 2, 8, 8, 42).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 mut actual = batched.into_candidate_lists(); + let mut expected = reference.into_candidate_lists(); + actual.iter_mut().for_each(|row| row.sort_unstable()); + expected.iter_mut().for_each(|row| row.sort_unstable()); + + assert_eq!(actual, expected); + assert!(actual.iter().all(|row| !row.is_empty())); } #[test] diff --git a/diskann-pipnn/src/leaf_build/tests.rs b/diskann-pipnn/src/leaf_build/tests.rs index 0236370de..9d9276cbc 100644 --- a/diskann-pipnn/src/leaf_build/tests.rs +++ b/diskann-pipnn/src/leaf_build/tests.rs @@ -9,9 +9,10 @@ use half::f16; use std::collections::BTreeSet; use super::{ - add_symmetric_edges, allocation_error, build_leaf_candidates, DirectCandidates, LeafBuffers, - LeafBuildError, + add_symmetric_edges, allocation_error, build_leaf_candidates, build_symmetric_edge_csr, + DirectCandidates, EdgeBuffers, LeafBuffers, LeafBuildError, }; +use crate::leaf_kernel::LeafNeighbor; fn view(data: &[T], rows: usize, columns: usize) -> MatrixView<'_, T> { MatrixView::try_from(data, rows, columns).unwrap() @@ -356,6 +357,65 @@ fn skips_duplicate_global_ids_without_self_edges() { assert!(graph.iter().all(|row| row.is_empty())); } +#[test] +fn symmetric_edge_csr_matches_expected_rows() { + let point_ids = [10, 20, 30]; + let nearest = [ + LeafNeighbor::new(1, 1.0), + LeafNeighbor::new(2, 2.0), + LeafNeighbor::new(1, 1.5), + ]; + let mut seen = vec![false; 9]; + let mut offsets = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = Vec::new(); + + let count = build_symmetric_edge_csr( + 0, + &point_ids, + 1, + &nearest, + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap(); + + assert_eq!(count, 4); + assert_eq!(offsets, [0, 1, 3, 4]); + assert_eq!(edges, [(1, 1.0), (0, 1.0), (2, 2.0), (1, 2.0)]); +} + +#[test] +fn zero_k_edge_csr_has_empty_rows() { + let point_ids = [10, 20, 30]; + let mut seen = vec![false; 9]; + let mut offsets = Vec::new(); + let mut edges = vec![(99, 99.0)]; + let mut cursor = Vec::new(); + + let count = build_symmetric_edge_csr( + 0, + &point_ids, + 0, + &[], + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap(); + + assert_eq!(count, 0); + assert_eq!(offsets, [0, 0, 0, 0]); + assert!(edges.is_empty()); +} + #[test] fn poisoned_candidate_rows_return_errors() { let candidates = DirectCandidates::new(1).unwrap(); diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index f325f14cd..d024ac144 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -17,7 +17,6 @@ //! outer pipeline allocation. mod bf16; -pub(crate) mod cpu_dispatch; mod finalization; mod hash_prune; mod leaf_build; diff --git a/diskann-pipnn/src/lsh.rs b/diskann-pipnn/src/lsh.rs index 7b344f3c9..e933c338f 100644 --- a/diskann-pipnn/src/lsh.rs +++ b/diskann-pipnn/src/lsh.rs @@ -203,6 +203,57 @@ mod tests { assert_eq!(sk.sketches().len(), 12); } + #[test] + fn sketches_match_serial_hyperplane_reference() { + let npoints = 3; + let ndims = 4; + let planes = 5; + let seed = 42; + let data: Vec = (0..npoints * ndims) + .map(|value| value as f32 - 3.0) + .collect(); + let actual = build_pool(2) + .install(|| { + LshSketches::try_new(npoints, ndims, planes, seed, |i, out| { + out.copy_from_slice(&data[i * ndims..(i + 1) * ndims]); + Ok::<_, std::convert::Infallible>(()) + }) + }) + .unwrap(); + + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let hyperplanes: Vec = (0..planes * ndims) + .map(|_| StandardNormal.sample(&mut rng)) + .collect(); + let mut expected = Vec::::new(); + for point in data.chunks_exact(ndims) { + for plane in hyperplanes.chunks_exact(ndims) { + expected.push(point.iter().zip(plane).map(|(x, h)| x * h).sum()); + } + } + assert_eq!(actual.sketches(), expected); + } + + #[test] + fn reports_shape_overflow_and_formats_errors() { + let hyperplanes = match LshSketches::try_new(1, usize::MAX, 2, 42, |_, _| { + Ok::<_, std::convert::Infallible>(()) + }) { + Ok(_) => panic!("hyperplane shape overflow must fail"), + Err(error) => error, + }; + assert!(matches!(hyperplanes, LshSketchError::ShapeOverflow { .. })); + assert!(hyperplanes.to_string().contains("overflows")); + + let sketches = match LshSketches::try_new(usize::MAX, 1, 2, 42, |_, _| { + Ok::<_, std::convert::Infallible>(()) + }) { + Ok(_) => panic!("sketch shape overflow must fail"), + Err(error) => error, + }; + assert!(matches!(sketches, LshSketchError::ShapeOverflow { .. })); + } + #[test] fn different_seeds_change_hashes() { let pool = build_pool(2); @@ -223,7 +274,8 @@ mod tests { #[test] fn try_new_propagates_fill_errors() { - #[derive(Debug, PartialEq)] + #[derive(Debug, PartialEq, thiserror::Error)] + #[error("fill failed")] struct FillError; let pool = build_pool(1); @@ -233,7 +285,8 @@ mod tests { Err(error) => error, }; - assert!(matches!(error, LshSketchError::Fill(FillError))); + assert!(matches!(&error, LshSketchError::Fill(FillError))); + assert!(std::error::Error::source(&error).is_some()); } #[test] diff --git a/diskann-pipnn/tests/build_graph.rs b/diskann-pipnn/tests/build_graph.rs index 5fd3bf36d..e0c5f9eac 100644 --- a/diskann-pipnn/tests/build_graph.rs +++ b/diskann-pipnn/tests/build_graph.rs @@ -194,7 +194,7 @@ fn is_deterministic_for_a_fixed_pool_size() { } #[test] -fn hash_prune_build_is_deterministic_and_degree_bounded() { +fn parallel_hash_prune_builds_preserve_graph_invariants() { let points = 64; let dimensions = 4; let values: Vec = (0..points * dimensions) @@ -224,11 +224,12 @@ fn hash_prune_build_is_deterministic_and_degree_bounded() { build_graph(data, &context).unwrap() }; - let first = build(); - let second = build(); - - assert_eq!(first, second); - assert_graph_invariants(&first, points, 8); + // Reservoir updates commute only until a full bucket must evict an edge; + // concurrent leaf order can therefore change ties without changing the contract. + for graph in [build(), build()] { + assert_graph_invariants(&graph, points, 8); + assert!(graph.iter().any(|row| !row.is_empty())); + } } #[test] diff --git a/diskann-vector/src/lib.rs b/diskann-vector/src/lib.rs index e88dc12c9..8774432da 100644 --- a/diskann-vector/src/lib.rs +++ b/diskann-vector/src/lib.rs @@ -38,7 +38,8 @@ 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)] diff --git a/diskann-wide/src/arch/x86_64/v3/i16x16_.rs b/diskann-wide/src/arch/x86_64/v3/i16x16_.rs index bf8a1a369..7506f84e3 100644 --- a/diskann-wide/src/arch/x86_64/v3/i16x16_.rs +++ b/diskann-wide/src/arch/x86_64/v3/i16x16_.rs @@ -16,7 +16,7 @@ use crate::{ v3::i16x8, }, }, - bitmask::BitMask, + bitmask::{BitMask, FromInt}, constant::Const, emulated::Emulated, helpers, @@ -89,16 +89,19 @@ impl SIMDMulAdd for i16x16 { impl SIMDPartialEq for i16x16 { #[inline(always)] fn eq_simd(self, other: Self) -> Self::Mask { - self.emulated() - .eq_simd(other.emulated()) - .as_arch(self.arch()) + // SAFETY: V3 includes AVX2 and BMI2. Each equal i16 lane contributes + // two adjacent movemask bits; pext keeps one bit per lane. + let bits = unsafe { + let bytes = _mm256_movemask_epi8(_mm256_cmpeq_epi16(self.0, other.0)) as u32; + _pext_u32(bytes, 0x5555_5555) as u16 + }; + BitMask::from_int(self.arch(), bits) } #[inline(always)] fn ne_simd(self, other: Self) -> Self::Mask { - self.emulated() - .ne_simd(other.emulated()) - .as_arch(self.arch()) + let equal = self.eq_simd(other); + BitMask::from_int(self.arch(), !equal.0) } } diff --git a/diskann-wide/src/doubled.rs b/diskann-wide/src/doubled.rs index 0a1ca5015..58f7ba176 100644 --- a/diskann-wide/src/doubled.rs +++ b/diskann-wide/src/doubled.rs @@ -435,6 +435,13 @@ macro_rules! double_mask { let hi = <$repr>::keep_first(arch, i.saturating_sub({ $N / 2 })); Self(lo, hi) } + + #[inline(always)] + fn first(&self) -> Option { + self.0 + .first() + .or_else(|| self.1.first().map(|index| index + { $N / 2 })) + } } impl From<$crate::doubled::Doubled<$repr>> From 408b1d60514e964f1c5756c9bf8969a762caf888 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 11:33:46 +0000 Subject: [PATCH 05/16] perf(pipnn): move unique HashPrune rows into adjacency --- diskann-pipnn/src/hash_prune.rs | 16 +++++++++++----- diskann-pipnn/src/hash_prune/tests.rs | 17 ++++++++++++----- diskann-pipnn/src/lib.rs | 12 ++---------- diskann/src/graph/adjacencylist.rs | 23 +++++++++++++++++++++++ 4 files changed, 48 insertions(+), 20 deletions(-) diff --git a/diskann-pipnn/src/hash_prune.rs b/diskann-pipnn/src/hash_prune.rs index 8d5db68df..df1d3f9e6 100644 --- a/diskann-pipnn/src/hash_prune.rs +++ b/diskann-pipnn/src/hash_prune.rs @@ -28,7 +28,7 @@ use std::cell::UnsafeCell; use crate::bf16::{bf16_to_f32, f32_to_bf16}; use crate::lsh::{LshSketchError, LshSketches}; use bytemuck::Pod; -use diskann::{utils::VectorRepr, ANNError, ANNResult}; +use diskann::{graph::AdjacencyList, utils::VectorRepr, ANNError, ANNResult}; use diskann_vector::prefetch_hint_all; use diskann_wide::{ arch::{self, Dispatched1, FTarget1, Target}, @@ -906,7 +906,7 @@ impl HashPrune { /// Extract the nearest `max_degree` candidates retained by HashPrune. #[allow(clippy::disallowed_methods)] // build_graph installs the caller-owned pool. - pub(crate) fn into_nearest_lists(self, max_degree: usize) -> Vec> { + pub(crate) fn into_nearest_lists(self, max_degree: usize) -> Vec> { let scan_lanes = self.scan_lanes; drop(self.sketches); let HashPrune { @@ -933,7 +933,10 @@ impl HashPrune { max_degree, ) }; - nbrs.into_iter().map(|(id, _)| id).collect() + let ids = nbrs.into_iter().map(|(id, _)| id).collect(); + // A neighbor always has the same relative hash for this source; + // insertion replaces an existing hash slot instead of appending. + AdjacencyList::from_vec_trusted(ids) }) .collect() } @@ -942,7 +945,7 @@ impl HashPrune { /// and distances slabs (2/3 of the reservoir) before materializing the copy, /// so only the neighbors slab overlaps it. #[allow(clippy::disallowed_methods)] // build_graph installs the caller-owned pool. - pub(crate) fn into_candidate_lists(self) -> Vec> { + pub(crate) fn into_candidate_lists(self) -> Vec> { let cap = self.l_max; let scan_lanes = self.scan_lanes; drop(self.sketches); @@ -966,7 +969,10 @@ impl HashPrune { let hot = unsafe { &*hot[i].get() }; // SAFETY: i < hot.len() == npoints; the row spans scan_lanes // slots and hot.len <= l_max <= scan_lanes. - unsafe { collect_neighbor_ids(hot, neighbors, cap) } + let ids = unsafe { collect_neighbor_ids(hot, neighbors, cap) }; + // Reservoir slots have unique hashes, and one neighbor cannot + // produce two hashes for the same source. + AdjacencyList::from_vec_trusted(ids) }) .collect() } diff --git a/diskann-pipnn/src/hash_prune/tests.rs b/diskann-pipnn/src/hash_prune/tests.rs index 01340211a..fcfc29c49 100644 --- a/diskann-pipnn/src/hash_prune/tests.rs +++ b/diskann-pipnn/src/hash_prune/tests.rs @@ -225,10 +225,17 @@ fn test_add_leaf_edges_matches_single_edge_reference() { add_edge(&reference, source, target as usize, distance); } } - let mut actual = batched.into_candidate_lists(); - let mut expected = reference.into_candidate_lists(); - actual.iter_mut().for_each(|row| row.sort_unstable()); - expected.iter_mut().for_each(|row| row.sort_unstable()); + let canonicalize = |rows: Vec>| { + rows.into_iter() + .map(|row| { + let mut ids = row.to_vec(); + ids.sort_unstable(); + ids + }) + .collect::>() + }; + let actual = canonicalize(batched.into_candidate_lists()); + let expected = canonicalize(reference.into_candidate_lists()); assert_eq!(actual, expected); assert!(actual.iter().all(|row| !row.is_empty())); @@ -398,7 +405,7 @@ fn test_into_candidate_lists_returns_full_reservoir() { // ids-only, unsorted: every id is one of node 0's inserted neighbors // {1,2,3} with no duplicates (the LSH bucket may keep-closer-collapse a // colliding pair on this tiny 4-plane sketch, so we don't assert all 3). - let mut n0 = full[0].clone(); + let mut n0 = full[0].to_vec(); n0.sort_unstable(); let deduped = { let mut d = n0.clone(); diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index d024ac144..1c5b95ffb 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -272,25 +272,17 @@ where .map_err(ANNError::opaque) })?; if config.final_prune { - let candidates = rows_to_adjacency(hash_prune.into_candidate_lists()); + let candidates = hash_prune.into_candidate_lists(); tracing::info_span!("pipnn.finalization").in_scope(|| { finalization::prune_overfull(data, candidates, context.graph, metric) }) } else { - Ok(rows_to_adjacency( - hash_prune.into_nearest_lists(context.graph.pruned_degree().get()), - )) + Ok(hash_prune.into_nearest_lists(context.graph.pruned_degree().get())) } } } } -fn rows_to_adjacency(rows: Vec>) -> Vec> { - rows.into_iter() - .map(AdjacencyList::from_iter_untrusted) - .collect() -} - fn effective_metric(metric: Metric) -> Metric { use std::any::TypeId; diff --git a/diskann/src/graph/adjacencylist.rs b/diskann/src/graph/adjacencylist.rs index 8071001ac..6874ae315 100644 --- a/diskann/src/graph/adjacencylist.rs +++ b/diskann/src/graph/adjacencylist.rs @@ -135,6 +135,19 @@ where } } + /// Construct an adjacency list by taking ownership of a vector whose items + /// are known to be unique. + /// + /// Uniqueness is checked in debug builds but not in release builds. + pub fn from_vec_trusted(edges: Vec) -> Self + where + I: ContainsSimd, + { + let list = Self { edges }; + list.debug_check_uniqueness(); + list + } + /// Resize the underlying storage to `capacity` elements and return a guard allowing /// full mutable access to the resized span. /// @@ -667,6 +680,16 @@ mod tests { } } + #[test] + fn test_from_vec_trusted_preserves_order_and_allocation() { + let edges = vec![3_u32, 1, 2]; + let pointer = edges.as_ptr(); + let list = AdjacencyList::from_vec_trusted(edges); + + assert_eq!(&*list, &[3, 1, 2]); + assert_eq!(list.as_ptr(), pointer); + } + #[test] fn test_from_iter_untrusted() { let x = AdjacencyList::::from_iter_untrusted([]); From e517242ce0d788da390db8bf09bd2cf2abc4c156 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:01:34 +0000 Subject: [PATCH 06/16] fix(pipnn): borrow consumed leaf IDs during edge build --- diskann-pipnn/src/leaf_build.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-pipnn/src/leaf_build.rs b/diskann-pipnn/src/leaf_build.rs index 103415098..7e56d296a 100644 --- a/diskann-pipnn/src/leaf_build.rs +++ b/diskann-pipnn/src/leaf_build.rs @@ -266,7 +266,7 @@ where let points = point_ids.len(); let edge_count = build_symmetric_edge_csr( leaf, - point_ids, + &point_ids, actual_k, &buffers.nearest[..points * actual_k], EdgeBuffers { From 27cd5abace236aec51850ed033c6adacce091255 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 13:18:26 +0000 Subject: [PATCH 07/16] fix(pipnn): validate HashPrune capacity against graph degree --- diskann-pipnn/src/lib.rs | 7 +++++++ diskann-pipnn/tests/config.rs | 5 +++++ 2 files changed, 12 insertions(+) diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index 1c5b95ffb..d6ac48f5d 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -180,6 +180,13 @@ impl<'a> PiPNNBuildContext<'a> { /// Enable HashPrune candidate merging for this build. pub fn with_hash_prune(mut self, config: HashPruneConfig) -> ANNResult { config.validate()?; + let degree = self.graph.pruned_degree().get(); + if config.l_max < degree { + return Err(config_error(format!( + "HashPrune l_max ({}) must be at least the graph degree ({degree})", + config.l_max + ))); + } self.candidate_merge = CandidateMerge::HashPrune(config); Ok(self) } diff --git a/diskann-pipnn/tests/config.rs b/diskann-pipnn/tests/config.rs index 9e684a2b5..71f6551f7 100644 --- a/diskann-pipnn/tests/config.rs +++ b/diskann-pipnn/tests/config.rs @@ -128,6 +128,11 @@ fn validates_hash_prune_parameters() { l_max: 256, final_prune: true, }, + HashPruneConfig { + num_hash_planes: 8, + l_max: 63, + final_prune: true, + }, ] { let context = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool).unwrap(); assert!(context.with_hash_prune(config).is_err()); From fbc2547b4fb449b72559d756f93c154bf1d14f52 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:01:17 +0000 Subject: [PATCH 08/16] ci: run Miri on HashPrune pointer kernels --- .github/workflows/nightly.yml | 46 +++++++++++++++++++++++++++++++---- 1 file changed, 41 insertions(+), 5 deletions(-) diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index fab755b6a..597f72774 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -17,7 +17,20 @@ env: RUST_CONFIG: 'build.rustflags=["-Dwarnings"]' RUST_BACKTRACE: 1 CARGO_TERM_COLOR: always - DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2,integration-test" + DISKANN_FEATURES: >- + virtual_storage, + spherical-quantization, + product-quantization, + tracing, + experimental_diversity_search, + disk-index, + flatbuffers, + linalg, + codegen, + multi-vector, + bftree, + inmem2, + integration-test defaults: run: @@ -36,7 +49,12 @@ jobs: run: rustup show && rustup component add clippy - uses: Swatinem/rust-cache@v2 - name: "clippy --workspace --all-targets" - run: cargo clippy --locked --workspace --all-targets --no-deps --config "$RUST_CONFIG" -- -Dwarnings + run: | + cargo clippy --locked --workspace \ + --all-targets \ + --no-deps \ + --config "$RUST_CONFIG" \ + -- -Dwarnings clippy-features: name: clippy-features (macos) @@ -52,7 +70,7 @@ jobs: cargo clippy --locked --workspace \ --all-targets \ --no-deps \ - --features ${{ env.DISKANN_FEATURES }} \ + --features "${{ env.DISKANN_FEATURES }}" \ --config "$RUST_CONFIG" \ -- -Dwarnings @@ -109,7 +127,7 @@ jobs: set -euxo pipefail cargo nextest run --locked --workspace \ --config "$RUST_CONFIG" \ - --features ${{ env.DISKANN_FEATURES }} + --features "${{ env.DISKANN_FEATURES }}" cargo test --locked --doc --workspace --config "$RUST_CONFIG" @@ -133,6 +151,24 @@ jobs: - uses: Swatinem/rust-cache@v2 - name: miri - run: cargo +nightly miri nextest run --locked --package diskann-quantization + run: | + cargo +nightly miri nextest run --locked \ + --package diskann-quantization + env: + MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance + + # Full PiPNN tests use Rayon, whose crossbeam dependency is not strict- + # provenance clean under Miri. Exercise the private raw-pointer dispatch + # kernels directly without weakening MIRIFLAGS. + - name: miri PiPNN HashPrune kernels + run: | + find_hash_test="hash_prune::tests::" + find_hash_test+="test_find_hash_handles_padded_" + find_hash_test+="boundaries_and_all_bit_patterns" + cargo +nightly miri test --locked --package diskann-pipnn \ + "$find_hash_test" --lib -- --exact + cargo +nightly miri test --locked --package diskann-pipnn \ + hash_prune::tests::test_relative_hash_matches_numeric_reference \ + --lib -- --exact env: MIRIFLAGS: -Zmiri-disable-isolation -Zmiri-strict-provenance From cca75ddd8f5825091d8c873c28ef8859b4573a89 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:19:55 +0000 Subject: [PATCH 09/16] refactor(pipnn): scope extraction scratch to Rayon jobs --- diskann-pipnn/src/hash_prune.rs | 41 ++++++++++++--------------- diskann-pipnn/src/hash_prune/tests.rs | 11 ++++++- 2 files changed, 28 insertions(+), 24 deletions(-) diff --git a/diskann-pipnn/src/hash_prune.rs b/diskann-pipnn/src/hash_prune.rs index df1d3f9e6..0bb8dcd86 100644 --- a/diskann-pipnn/src/hash_prune.rs +++ b/diskann-pipnn/src/hash_prune.rs @@ -613,8 +613,8 @@ unsafe fn insert_locked( } /// Collect the reservoir's entries sorted by distance, truncated to `cap`. -/// A thread-local scratch `Vec`, sized to the reservoir's runtime fill and -/// reused across calls, avoids per-row allocation during extraction. +/// A Rayon-owned scratch `Vec`, sized to the reservoir's runtime fill and +/// reused within one extraction job, avoids per-row allocation. /// /// SAFETY: caller holds the slot lock; `distances` and `neighbors` are valid /// for `hot.len` elements. @@ -623,28 +623,22 @@ unsafe fn collect_sorted_neighbors( distances: *const u16, neighbors: *const u32, cap: usize, + scratch: &mut Vec<(u32, u16)>, ) -> Vec<(u32, f32)> { - thread_local! { - static SCRATCH: std::cell::RefCell> = - const { std::cell::RefCell::new(Vec::new()) }; - } let n = hot.len as usize; - SCRATCH.with(|cell| { - let mut scratch = cell.borrow_mut(); - scratch.clear(); - scratch.reserve(n); - for i in 0..n { - // SAFETY: guaranteed by this function's contract. - scratch.push(unsafe { (*neighbors.add(i), *distances.add(i)) }); - } - scratch.sort_unstable_by_key(|&(_, d)| d); - let out_len = n.min(cap); - let mut out = Vec::with_capacity(out_len); - for &(id, d) in &scratch[..out_len] { - out.push((id, bf16_to_f32(key_to_bf16(d)))); - } - out - }) + scratch.clear(); + scratch.reserve(n); + for i in 0..n { + // SAFETY: guaranteed by this function's contract. + scratch.push(unsafe { (*neighbors.add(i), *distances.add(i)) }); + } + scratch.sort_unstable_by_key(|&(_, d)| d); + let out_len = n.min(cap); + let mut out = Vec::with_capacity(out_len); + for &(id, d) in &scratch[..out_len] { + out.push((id, bf16_to_f32(key_to_bf16(d)))); + } + out } /// Collect the reservoir's neighbor ids, truncated to `cap`, WITHOUT sorting. @@ -919,7 +913,7 @@ impl HashPrune { drop(cold_hashes); (0..hot.len()) .into_par_iter() - .map(|i| { + .map_init(Vec::new, |scratch, i| { let off = i * scan_lanes; // SAFETY: extraction owns every row, so no mutation remains. let hot = unsafe { &*hot[i].get() }; @@ -931,6 +925,7 @@ impl HashPrune { cold_distances.as_ptr().wrapping_add(off), cold_neighbors.as_ptr().wrapping_add(off), max_degree, + scratch, ) }; let ids = nbrs.into_iter().map(|(id, _)| id).collect(); diff --git a/diskann-pipnn/src/hash_prune/tests.rs b/diskann-pipnn/src/hash_prune/tests.rs index fcfc29c49..de2c70e88 100644 --- a/diskann-pipnn/src/hash_prune/tests.rs +++ b/diskann-pipnn/src/hash_prune/tests.rs @@ -55,8 +55,17 @@ impl Reservoir { fn neighbors(&self) -> Vec<(u32, f32)> { let cold = self.cold(); + let mut scratch = Vec::new(); // SAFETY: the test owns the reservoir; all cold slabs span scan_lanes entries. - unsafe { collect_sorted_neighbors(&self.hot, cold.distances, cold.neighbors, usize::MAX) } + unsafe { + collect_sorted_neighbors( + &self.hot, + cold.distances, + cold.neighbors, + usize::MAX, + &mut scratch, + ) + } } fn len(&self) -> usize { From 2c30cf8eeb1887354dab97445ce7f60b5058a00f Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 15:44:56 +0000 Subject: [PATCH 10/16] perf(pipnn): defer HashPrune leaf release to stage exit --- diskann-pipnn/src/leaf_build.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-pipnn/src/leaf_build.rs b/diskann-pipnn/src/leaf_build.rs index 7e56d296a..103415098 100644 --- a/diskann-pipnn/src/leaf_build.rs +++ b/diskann-pipnn/src/leaf_build.rs @@ -266,7 +266,7 @@ where let points = point_ids.len(); let edge_count = build_symmetric_edge_csr( leaf, - &point_ids, + point_ids, actual_k, &buffers.nearest[..points * actual_k], EdgeBuffers { From 3eba8859ad47f225dbd289020d527bcc8bfc1eb5 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:37:29 +0000 Subject: [PATCH 11/16] fix(pipnn): validate HashPrune hash-space capacity --- diskann-pipnn/src/lib.rs | 9 +++++++-- diskann-pipnn/tests/config.rs | 22 +++++++++++++++++++++- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index d6ac48f5d..af4d05d79 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -181,9 +181,14 @@ impl<'a> PiPNNBuildContext<'a> { pub fn with_hash_prune(mut self, config: HashPruneConfig) -> ANNResult { config.validate()?; let degree = self.graph.pruned_degree().get(); - if config.l_max < degree { + let hash_capacity = 1usize + .checked_shl(config.num_hash_planes as u32) + .unwrap_or(usize::MAX); + let candidate_capacity = config.l_max.min(hash_capacity); + if candidate_capacity < degree { return Err(config_error(format!( - "HashPrune l_max ({}) must be at least the graph degree ({degree})", + "HashPrune capacity min(l_max={}, hash buckets={hash_capacity}) must be at least \ + the graph degree ({degree})", config.l_max ))); } diff --git a/diskann-pipnn/tests/config.rs b/diskann-pipnn/tests/config.rs index 71f6551f7..84eb14054 100644 --- a/diskann-pipnn/tests/config.rs +++ b/diskann-pipnn/tests/config.rs @@ -19,7 +19,11 @@ fn pipnn_config() -> PiPNNConfig { } fn graph_config(metric: Metric, alpha: f32) -> diskann::graph::Config { - config::Builder::new_with(64, MaxDegree::same(), 72, metric.into(), |builder| { + graph_config_with_degree(metric, alpha, 64) +} + +fn graph_config_with_degree(metric: Metric, alpha: f32, degree: usize) -> diskann::graph::Config { + config::Builder::new_with(degree, MaxDegree::same(), 72, metric.into(), |builder| { builder.alpha(alpha); }) .build() @@ -148,6 +152,22 @@ fn validates_hash_prune_parameters() { .unwrap(); } +#[test] +fn validates_hash_bucket_capacity_against_graph_degree() { + let pool = pool(); + for (degree, accepted) in [(2, true), (3, false)] { + let graph = graph_config_with_degree(Metric::L2, 1.2, degree); + let result = PiPNNBuildContext::new(pipnn_config(), &graph, Metric::L2, &pool) + .unwrap() + .with_hash_prune(HashPruneConfig { + num_hash_planes: 1, + l_max: 64, + final_prune: true, + }); + assert_eq!(result.is_ok(), accepted, "degree={degree}"); + } +} + #[test] fn rejects_graph_policy_for_a_different_metric() { let graph = graph_config(Metric::InnerProduct, 1.2); From eac8b60ec40bb9f8f98112946927ff201df31536 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 30 Jul 2026 07:34:09 +0000 Subject: [PATCH 12/16] perf(pipnn): retain HashPrune sketch scratch capacity --- diskann-pipnn/src/hash_prune.rs | 4 +++- diskann-pipnn/src/hash_prune/tests.rs | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/diskann-pipnn/src/hash_prune.rs b/diskann-pipnn/src/hash_prune.rs index 0bb8dcd86..1f8a803e6 100644 --- a/diskann-pipnn/src/hash_prune.rs +++ b/diskann-pipnn/src/hash_prune.rs @@ -846,7 +846,9 @@ impl HashPrune { let l_max = self.l_max as u8; debug_assert_eq!(edge_offsets.len(), n + 1); let sketch_len = n * m; - sketch_scratch.resize(sketch_len, 0.0); + 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 { diff --git a/diskann-pipnn/src/hash_prune/tests.rs b/diskann-pipnn/src/hash_prune/tests.rs index de2c70e88..3c5e0259a 100644 --- a/diskann-pipnn/src/hash_prune/tests.rs +++ b/diskann-pipnn/src/hash_prune/tests.rs @@ -229,6 +229,10 @@ fn test_add_leaf_edges_matches_single_edge_reference() { let mut scratch = Vec::new(); batched.add_leaf_edges(&point_ids, &offsets, &edges, &mut scratch); + let scratch_len = scratch.len(); + batched.add_leaf_edges(&point_ids[..2], &[0, 0, 0], &[], &mut scratch); + assert_eq!(scratch.len(), scratch_len); + 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); From 10576044c2fddbb8e566a24bf3302f6717308624 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:03:25 +0000 Subject: [PATCH 13/16] fix(pipnn): declare HashPrune bytemuck dependency --- Cargo.lock | 1 + diskann-pipnn/Cargo.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 161c06d62..1e0c78578 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -688,6 +688,7 @@ dependencies = [ name = "diskann-pipnn" version = "0.55.0" dependencies = [ + "bytemuck", "criterion", "diskann", "diskann-linalg", diff --git a/diskann-pipnn/Cargo.toml b/diskann-pipnn/Cargo.toml index 42408ca0d..6c9255199 100644 --- a/diskann-pipnn/Cargo.toml +++ b/diskann-pipnn/Cargo.toml @@ -11,6 +11,7 @@ license.workspace = true edition.workspace = true [dependencies] +bytemuck.workspace = true diskann.workspace = true diskann-linalg.workspace = true diskann-utils.workspace = true From 0c3beaf2fcaf336cf2a7a450d8f3a36931209e98 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 03:34:48 +0000 Subject: [PATCH 14/16] docs(pipnn): map HashPrune reservoir flow --- diskann-pipnn/src/bf16.rs | 40 ++++++++++++- diskann-pipnn/src/hash_prune.rs | 39 ++++++++++++- diskann-pipnn/src/hash_prune/tests.rs | 81 +++++++++++++++++++++++++++ diskann-pipnn/src/leaf_build.rs | 32 +++++++++-- diskann-pipnn/src/leaf_build/tests.rs | 57 +++++++++++++++++++ diskann-pipnn/src/lib.rs | 42 +++++++++++--- diskann-pipnn/src/lsh.rs | 61 +++++++++++++++++--- 7 files changed, 329 insertions(+), 23 deletions(-) diff --git a/diskann-pipnn/src/bf16.rs b/diskann-pipnn/src/bf16.rs index 68d55280b..be61703b2 100644 --- a/diskann-pipnn/src/bf16.rs +++ b/diskann-pipnn/src/bf16.rs @@ -8,8 +8,12 @@ //! `bf16` is the upper 16 bits of an IEEE-754 `f32`: same exponent, mantissa //! truncated from 23 to 7 bits. For non-negative values, `bf16` bit ordering //! matches `f32` ordering, so a packed `bf16` can be compared as `u16` to give -//! the correct distance order — useful for keeping a top-k tracker as compact -//! 8-byte entries. +//! the correct distance order. HashPrune stores signed distances through a +//! separate order-preserving key transform before comparing these bits. +//! +//! Conversion is truncation, not round-to-nearest. It preserves the sign bit, +//! infinities, signed zero, and the upper NaN payload bits exactly; lower payload +//! bits are discarded with the lower mantissa. /// Convert `f32` → bf16 by truncating the lower 16 mantissa bits. #[inline(always)] @@ -48,6 +52,38 @@ mod tests { } } + #[test] + fn special_values_preserve_their_upper_bits() { + for value in [ + 0.0_f32, + -0.0, + f32::INFINITY, + f32::NEG_INFINITY, + f32::NAN, + f32::from_bits(0xFFC1_2345), + f32::MIN_POSITIVE, + -f32::MIN_POSITIVE, + ] { + let expected = value.to_bits() & 0xFFFF_0000; + assert_eq!( + bf16_to_f32(f32_to_bf16(value)).to_bits(), + expected, + "value_bits={:08x}", + value.to_bits() + ); + } + } + + #[test] + fn finite_truncation_moves_toward_zero() { + for value in [f32::MIN_POSITIVE, 0.1, 1.1, std::f32::consts::PI, f32::MAX] { + let positive = bf16_to_f32(f32_to_bf16(value)); + let negative = bf16_to_f32(f32_to_bf16(-value)); + assert!((0.0..=value).contains(&positive), "value={value}"); + assert!((-value..=0.0).contains(&negative), "value={value}"); + } + } + #[test] fn ordering_preserved_for_non_negative() { // For non-negative f32, bf16 (as u16) preserves ordering. diff --git a/diskann-pipnn/src/hash_prune.rs b/diskann-pipnn/src/hash_prune.rs index 1f8a803e6..64e32d0c6 100644 --- a/diskann-pipnn/src/hash_prune.rs +++ b/diskann-pipnn/src/hash_prune.rs @@ -5,6 +5,34 @@ //! HashPrune: LSH-based online pruning for merging edges from overlapping partitions. //! +//! ```text +//! dataset ──> random-hyperplane sketches (one row per point) +//! │ +//! leaf-local symmetric CSR ──> gather leaf sketches +//! │ +//! v +//! relative hash(source, target) +//! │ +//! v +//! lock source row ──> update bounded reservoir +//! │ +//! consume HashPrune at stage exit +//! ┌──────────┴──────────┐ +//! v v +//! nearest rows unsorted candidate rows +//! (for RobustPrune) +//! ``` +//! +//! Each source row retains at most one neighbor per relative hash bucket: +//! +//! | Existing state | Incoming edge | Action | +//! | --- | --- | --- | +//! | matching hash | closer | replace that bucket | +//! | matching hash | not closer | reject | +//! | free bucket | any distance | append | +//! | full | closer than cached farthest | replace farthest | +//! | full | not closer | reject before scanning hashes | +//! //! Storage is AoSoA hot/cold split: //! - `hot: Vec` — one 16-byte slot per point (mutex + len + farthest). //! Early rejection and lock acquisition only touch this slab. @@ -125,8 +153,7 @@ impl std::ops::Deref for MmapSlab { /// Windows counterpart of the Linux mmap slab. `VirtualAlloc(MEM_RESERVE | /// MEM_COMMIT, PAGE_READWRITE)` reserves a zero-backed anonymous range whose /// pages fault in on first write — the same lazy-commit behavior as -/// `mmap(MAP_ANONYMOUS)`, so Windows gets the same peak-RSS win instead of the -/// eager-fault `Vec` fallback. +/// `mmap(MAP_ANONYMOUS)`, rather than the fallback `Vec`'s eager initialization. #[cfg(windows)] mod winmem { // Minimal FFI to the Win32 memory API — avoids pulling in the `windows` @@ -661,6 +688,12 @@ unsafe fn collect_neighbor_ids(hot: &HotSlot, neighbors: *const u32, cap: usize) // ─── HashPrune ──────────────────────────────────────────────────────────────── +/// Global bounded reservoirs shared by parallel leaf workers. +/// +/// Row `i` consists of `hot[i]` and the `i * scan_lanes` range in each cold +/// slab. The embedded row mutex is the sole synchronization boundary: callers +/// never hold two row locks, and sketch computation/extraction happens without +/// a lock. Consuming extraction proves that no writer can remain. pub(crate) struct HashPrune { hot: Vec, /// AoSoA hashes slab: `npoints * scan_lanes` u16. @@ -687,6 +720,8 @@ unsafe impl Send for HashPrune {} unsafe impl Sync for HashPrune {} impl HashPrune { + /// Precompute immutable sketches and allocate one lazy-backed reservoir row + /// per dataset point. pub(crate) fn new( data: &[T], npoints: usize, diff --git a/diskann-pipnn/src/hash_prune/tests.rs b/diskann-pipnn/src/hash_prune/tests.rs index 3c5e0259a..2766560a0 100644 --- a/diskann-pipnn/src/hash_prune/tests.rs +++ b/diskann-pipnn/src/hash_prune/tests.rs @@ -92,6 +92,63 @@ fn add_edge(hp: &HashPrune, src: usize, dst: usize, distance: f32) { }); } +fn assert_sketch_source_type_matches_f32( + label: &str, + convert: impl Fn(u8) -> T, + reference: impl Fn(u8) -> f32, +) where + T: VectorRepr + Send + Sync, +{ + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .unwrap(); + let points = 5; + for dimensions in [1, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33] { + let raw: Vec = (0..points * dimensions) + .map(|index| ((index * 7 + index / dimensions * 3) % 23) as u8) + .collect(); + let converted: Vec = raw.iter().copied().map(&convert).collect(); + let f32_data: Vec = raw.iter().copied().map(&reference).collect(); + for planes in [1, 8, 16] { + let (actual, expected) = pool.install(|| { + ( + sketches_from_data(&converted, points, dimensions, planes, 42).unwrap(), + sketches_from_data(&f32_data, points, dimensions, planes, 42).unwrap(), + ) + }); + assert_eq!( + actual.sketches(), + expected.sketches(), + "{label} dimensions={dimensions} planes={planes}" + ); + } + } +} + +#[test] +fn test_f16_sketch_conversion_matches_f32_across_dimensions_and_planes() { + assert_sketch_source_type_matches_f32( + "f16", + |value| half::f16::from_f32(value as f32), + |value| value as f32, + ); +} + +#[test] +fn test_u8_sketch_conversion_matches_f32_across_dimensions_and_planes() { + assert_sketch_source_type_matches_f32("u8", |value| value, |value| value as f32); +} + +#[test] +fn test_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, + ); +} + #[test] fn test_relative_hash_matches_numeric_reference() { let dispatched = select_relative_hash(); @@ -184,6 +241,30 @@ fn test_round_up_to_32_boundaries() { assert_eq!(round_up_to_32(33), 64); } +#[test] +fn test_hash_prune_accepts_structural_l_max_boundaries() { + let data = [0.0_f32]; + let low = HashPrune::new(&data, 1, 1, 1, 1, 42).unwrap(); + assert_eq!(low.l_max, 1); + assert_eq!(low.scan_lanes, 32); + + let high = HashPrune::new(&data, 1, 1, 1, MAX_RESERVOIR_LEN, 42).unwrap(); + assert_eq!(high.l_max, MAX_RESERVOIR_LEN); + assert_eq!(high.scan_lanes, 256); +} + +#[test] +fn test_hash_prune_rejects_l_max_outside_structural_boundaries() { + for l_max in [0, MAX_RESERVOIR_LEN + 1] { + let result = HashPrune::new(&[0.0_f32], 1, 1, 1, l_max, 42); + let error = match result { + Ok(_) => panic!("l_max={l_max} must be rejected"), + Err(error) => error, + }; + assert!(format!("{error:?}").contains(&format!("l_max ({l_max})"))); + } +} + #[test] fn test_ordered_key_roundtrips_bf16_order_for_all_signs() { let values = [ diff --git a/diskann-pipnn/src/leaf_build.rs b/diskann-pipnn/src/leaf_build.rs index 103415098..c422d6fa7 100644 --- a/diskann-pipnn/src/leaf_build.rs +++ b/diskann-pipnn/src/leaf_build.rs @@ -12,10 +12,22 @@ //! 3. runs the dual-endpoint leaf top-k kernel; and //! 4. translates leaf-local positions back to dataset IDs. //! -//! The final step merges symmetric adjacency rows under per-point locks because -//! overlapping leaves are processed concurrently. Numeric buffers retain their -//! high-water length; every consumer therefore receives an explicit active -//! prefix rather than treating `Vec::len()` as the current leaf shape. +//! ```text +//! leaf IDs ─> gathered f32 rows ─> lower A·Aᵀ ─> local top-k +//! │ +//! ┌─────────────────────────┴─────────────────────┐ +//! v v +//! symmetric ID rows symmetric weighted CSR +//! │ │ +//! per-dataset-row mutex HashPrune rows +//! ``` +//! +//! Direct merging uses per-point locks because overlapping leaves are processed +//! concurrently. HashPrune instead builds one deduplicated weighted CSR row per +//! leaf source; the global reservoir owns its own row lock. Numeric and CSR +//! buffers retain their high-water length, so every consumer receives an +//! explicit active prefix rather than treating `Vec::len()` as the current leaf +//! shape. use std::{ collections::{HashSet, TryReserveError}, @@ -435,6 +447,18 @@ struct EdgeBuffers<'a> { cursor: &'a mut Vec, } +/// Materialize the symmetric local top-k relation as deduplicated weighted CSR. +/// +/// | Pass | `seen[source, target]` | Output | +/// | --- | --- | --- | +/// | count | first occurrence increments row count | counts in `offsets[1..]` | +/// | prefix | not used | exclusive row boundaries in `offsets` | +/// | write | first occurrence reserves one cursor slot | `(local_target, distance)` | +/// +/// Both directions are visited for every top-k edge. The dense `seen` matrix is +/// no larger than the already-live dot-product matrix and is reused by the +/// worker. Targets remain leaf-local here; HashPrune translates through +/// `point_ids` while the leaf mapping is still available. fn build_symmetric_edge_csr( leaf: usize, point_ids: &[u32], diff --git a/diskann-pipnn/src/leaf_build/tests.rs b/diskann-pipnn/src/leaf_build/tests.rs index 9d9276cbc..36891e531 100644 --- a/diskann-pipnn/src/leaf_build/tests.rs +++ b/diskann-pipnn/src/leaf_build/tests.rs @@ -389,6 +389,63 @@ fn symmetric_edge_csr_matches_expected_rows() { assert_eq!(edges, [(1, 1.0), (0, 1.0), (2, 2.0), (1, 2.0)]); } +#[test] +fn symmetric_edge_csr_deduplicates_edges_seen_from_both_endpoints() { + let point_ids = [10, 20]; + let nearest = [LeafNeighbor::new(1, 1.0), LeafNeighbor::new(0, 1.0)]; + let mut seen = vec![false; 4]; + let mut offsets = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = Vec::new(); + + let count = build_symmetric_edge_csr( + 0, + &point_ids, + 1, + &nearest, + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap(); + + assert_eq!(count, 2); + assert_eq!(offsets, [0, 1, 2]); + assert_eq!(edges, [(1, 1.0), (0, 1.0)]); +} + +#[test] +fn symmetric_edge_csr_rejects_out_of_range_local_targets() { + let mut seen = vec![false; 4]; + let mut offsets = Vec::new(); + let mut edges = Vec::new(); + let mut cursor = Vec::new(); + let error = build_symmetric_edge_csr( + 7, + &[10, 20], + 1, + &[LeafNeighbor::new(2, 1.0), LeafNeighbor::new(0, 1.0)], + EdgeBuffers { + seen: &mut seen, + offsets: &mut offsets, + edges: &mut edges, + cursor: &mut cursor, + }, + ) + .unwrap_err(); + + assert!(matches!( + error, + LeafBuildError::InvalidLocalPosition { + position: 2, + points: 2 + } + )); +} + #[test] fn zero_k_edge_csr_has_empty_rows() { let point_ids = [10, 20, 30]; diff --git a/diskann-pipnn/src/lib.rs b/diskann-pipnn/src/lib.rs index af4d05d79..9440e8509 100644 --- a/diskann-pipnn/src/lib.rs +++ b/diskann-pipnn/src/lib.rs @@ -5,16 +5,35 @@ //! Provider-independent PiPNN graph construction. //! -//! The crate owns three synchronous stages: overlapping partition generation, -//! leaf-local candidate construction, and graph-degree finalization. The caller -//! supplies a contiguous `MatrixView`, DiskANN graph policy, and the Rayon pool -//! that must execute every parallel iterator. Providers, start/frozen points, -//! quantization, persistence, and search remain outside this boundary. +//! The crate owns overlapping partition generation, leaf-local nearest-neighbor +//! construction, candidate merging, and optional graph-degree finalization. The +//! caller supplies a contiguous `MatrixView`, DiskANN graph policy, and the Rayon +//! pool that must execute every parallel iterator. Providers, start/frozen +//! points, quantization, persistence, and search remain outside this boundary. +//! +//! ```text +//! MatrixView ──> overlapping partition ──> owned leaves ──> leaf top-k +//! │ +//! ┌─────────────────────────┴──────────────────┐ +//! v v +//! direct candidate rows HashPrune reservoirs +//! │ │ │ +//! v final_prune extract +//! RobustPrune │ │ +//! └───────────────────────────────┴───────────┘ +//! adjacency rows +//! ``` +//! +//! | Merge policy | Intermediate bound | Finalization | +//! | --- | --- | --- | +//! | direct | all unique leaf candidates | shared RobustPrune | +//! | HashPrune + `final_prune` | `l_max` hash buckets per point | shared RobustPrune | +//! | HashPrune only | `l_max` hash buckets per point | nearest `max_degree` extraction | //! //! Stage outputs are owned values. Partition leaves move into leaf construction; -//! candidate rows move into finalization. That ownership both documents the -//! lifetime boundary and releases each stage's large scratch before the next -//! outer pipeline allocation. +//! candidate rows or reservoirs move into extraction/finalization. That ownership +//! documents the lifetime boundary and releases each stage's large allocation +//! before the next outer pipeline allocation. mod bf16; mod finalization; @@ -265,6 +284,9 @@ where .in_scope(|| finalization::prune_overfull(data, candidates, context.graph, metric)) } CandidateMerge::HashPrune(config) => { + // Sketches and reservoir slabs must exist before leaf workers start; + // workers borrow this stage owner and mutate only one locked source + // row at a time. let hash_prune = hash_prune::HashPrune::new( data.as_slice(), data.nrows(), @@ -273,6 +295,8 @@ where config.l_max, 42, )?; + // Leaves are consumed at this boundary. Their local weighted CSR + // rows are temporary; only bounded reservoir state survives. tracing::info_span!("pipnn.leaf_build").in_scope(|| { leaf_build::add_hash_prune_candidates( data, @@ -284,6 +308,8 @@ where .map_err(ANNError::opaque) })?; if config.final_prune { + // Consuming extraction drops sketches and unused cold slabs + // before candidate adjacency is materialized. let candidates = hash_prune.into_candidate_lists(); tracing::info_span!("pipnn.finalization").in_scope(|| { finalization::prune_overfull(data, candidates, context.graph, metric) diff --git a/diskann-pipnn/src/lsh.rs b/diskann-pipnn/src/lsh.rs index e933c338f..5d58b4de0 100644 --- a/diskann-pipnn/src/lsh.rs +++ b/diskann-pipnn/src/lsh.rs @@ -5,13 +5,25 @@ //! Random-hyperplane LSH (Locality-Sensitive Hashing) for `f32` vectors. //! -//! Computes `Sketch(v) = [v · H_i for i in 0..num_planes]` where `H_i` are -//! random unit-Gaussian hyperplanes. Callers use these sketches to derive -//! application-specific hashes. +//! Computes `Sketch(v) = [v · H_i for i in 0..num_planes]` where each +//! hyperplane component is sampled from a standard normal distribution. Callers +//! use differences between two sketch rows to derive relative hash bits. //! -//! Sketches are computed in parallel via rayon, with caller-provided per-point -//! `f32` conversion (so f16, u8, etc. don't need a full upfront f32 copy). -//! `num_planes ≤ 16` so the result fits in `u16`. +//! ```text +//! seeded RNG ──> hyperplanes [planes × dimensions] (immutable) +//! │ +//! source row ──> worker f32 scratch ──> dot products ──> sketch row [planes] +//! ``` +//! +//! | Buffer | Shape | Lifetime | +//! | --- | --- | --- | +//! | hyperplanes | `num_planes × ndims` | construction call | +//! | conversion scratch | `ndims` per Rayon job | reused across point rows | +//! | sketches | `npoints × num_planes` | owned by `LshSketches` | +//! +//! Sketches are computed in parallel via Rayon, with caller-provided per-point +//! `f32` conversion so f16, u8, and i8 storage does not require a full upfront +//! f32 copy. `num_planes ≤ 16` keeps every relative hash in a `u16`. use rand::SeedableRng; use rand_distr::{Distribution, StandardNormal}; @@ -115,7 +127,7 @@ impl LshSketches { })?; let mut rng = rand::rngs::StdRng::seed_from_u64(seed); - // Random unit-Gaussian hyperplanes, row-major `num_planes × ndims`. + // Standard-normal hyperplanes, row-major `num_planes × ndims`. let mut hyperplanes: Vec = Vec::new(); hyperplanes .try_reserve_exact(hyperplane_len) @@ -234,6 +246,41 @@ mod tests { assert_eq!(actual.sketches(), expected); } + #[test] + fn zero_points_produces_an_empty_sketch_without_filling() { + let sketches = build_pool(2) + .install(|| { + LshSketches::try_new( + 0, + 7, + 4, + 42, + |_, _| -> Result<(), std::convert::Infallible> { + panic!("zero points must not invoke fill_point") + }, + ) + }) + .unwrap(); + assert_eq!(sketches.num_planes(), 4); + assert!(sketches.sketches().is_empty()); + } + + #[test] + fn zero_dimensions_produces_zero_dot_products() { + let calls = std::sync::atomic::AtomicUsize::new(0); + let sketches = build_pool(2) + .install(|| { + LshSketches::try_new(3, 0, 2, 42, |_, out| { + assert!(out.is_empty()); + calls.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Ok::<_, std::convert::Infallible>(()) + }) + }) + .unwrap(); + assert_eq!(calls.load(std::sync::atomic::Ordering::Relaxed), 3); + assert_eq!(sketches.sketches(), &[0.0; 6]); + } + #[test] fn reports_shape_overflow_and_formats_errors() { let hyperplanes = match LshSketches::try_new(1, usize::MAX, 2, 42, |_, _| { From 4e2830ddd9d5820ea655abe61ecadd20c773025a Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 04:16:18 +0000 Subject: [PATCH 15/16] test(pipnn): exercise sketch scratch reuse --- diskann-pipnn/src/hash_prune/tests.rs | 16 ++++++++++++++++ diskann-pipnn/src/lsh.rs | 17 +++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/diskann-pipnn/src/hash_prune/tests.rs b/diskann-pipnn/src/hash_prune/tests.rs index 2766560a0..fcdce4f5d 100644 --- a/diskann-pipnn/src/hash_prune/tests.rs +++ b/diskann-pipnn/src/hash_prune/tests.rs @@ -335,6 +335,22 @@ fn test_add_leaf_edges_matches_single_edge_reference() { assert!(actual.iter().all(|row| !row.is_empty())); } +#[test] +fn test_add_leaf_edges_grows_and_then_reuses_sketch_scratch() { + let data = [0.0_f32, 1.0, 2.0, 3.0]; + let hp = HashPrune::new(&data, 4, 1, 8, 4, 42).unwrap(); + let mut scratch = vec![99.0; 1]; + + hp.add_leaf_edges(&[0, 1], &[0, 1, 2], &[(1, 1.0), (0, 1.0)], &mut scratch); + assert_eq!(scratch.len(), 16); + let capacity = scratch.capacity(); + + hp.add_leaf_edges(&[2, 3], &[0, 1, 2], &[(1, 1.0), (0, 1.0)], &mut scratch); + assert_eq!(scratch.len(), 16); + assert_eq!(scratch.capacity(), capacity); + assert!(hp.into_candidate_lists().iter().all(|row| row.len() == 1)); +} + #[test] fn test_reservoir_basic() { let mut reservoir = Reservoir::new(3); diff --git a/diskann-pipnn/src/lsh.rs b/diskann-pipnn/src/lsh.rs index 5d58b4de0..d361b15ce 100644 --- a/diskann-pipnn/src/lsh.rs +++ b/diskann-pipnn/src/lsh.rs @@ -246,6 +246,23 @@ mod tests { assert_eq!(actual.sketches(), expected); } + #[test] + fn worker_conversion_scratch_grows_across_point_rows() { + let lengths = std::sync::Mutex::new(Vec::new()); + let sketches = build_pool(1) + .install(|| { + LshSketches::try_new(2, 5, 1, 42, |i, out| { + lengths.lock().unwrap().push(out.len()); + out.fill((i + 1) as f32); + Ok::<_, std::convert::Infallible>(()) + }) + }) + .unwrap(); + assert_eq!(*lengths.lock().unwrap(), [5, 5]); + assert_eq!(sketches.sketches().len(), 2); + assert_ne!(sketches.sketches()[0], sketches.sketches()[1]); + } + #[test] fn zero_points_produces_an_empty_sketch_without_filling() { let sketches = build_pool(2) From 25ab02bd6395b7e471457a85c4d01a0177265341 Mon Sep 17 00:00:00 2001 From: Weiyao Luo <9347182+SeliMeli@users.noreply.github.com> Date: Fri, 31 Jul 2026 04:21:21 +0000 Subject: [PATCH 16/16] test(disk): adapt PiPNN adapter mismatch case --- diskann-disk/src/build/builder/build/pipnn/tests.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diskann-disk/src/build/builder/build/pipnn/tests.rs b/diskann-disk/src/build/builder/build/pipnn/tests.rs index c1edd4401..6649b22fb 100644 --- a/diskann-disk/src/build/builder/build/pipnn/tests.rs +++ b/diskann-disk/src/build/builder/build/pipnn/tests.rs @@ -115,7 +115,7 @@ fn pipnn_graph_adapter_rejects_point_count_mismatch() { 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())); }