diff --git a/Cargo.lock b/Cargo.lock index 03d9322bc8..0c4a093577 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4938,6 +4938,24 @@ dependencies = [ "tracing", ] +[[package]] +name = "lattice-embed" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a8471670d8eb3dc5b52b7c977f1be449a4d39404ebd47bd084a1e922fa6002e" +dependencies = [ + "async-trait", + "blake3", + "chrono", + "lru 0.16.4", + "parking_lot 0.12.5", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tracing", +] + [[package]] name = "lattice-inference" version = "0.6.1" @@ -9203,7 +9221,7 @@ dependencies = [ "dashmap 6.2.1", "hf-hub", "hnsw_rs", - "lattice-embed", + "lattice-embed 0.6.1", "memmap2", "mockall", "ndarray 0.16.1", @@ -9407,6 +9425,7 @@ dependencies = [ "bincode 2.0.1", "bytemuck", "getrandom 0.2.17", + "lattice-embed 0.7.1", "memmap2", "parking_lot 0.12.5", "rand 0.8.6", diff --git a/crates/ruvector-diskann/Cargo.toml b/crates/ruvector-diskann/Cargo.toml index bf4adcb7f8..0ed9793389 100644 --- a/crates/ruvector-diskann/Cargo.toml +++ b/crates/ruvector-diskann/Cargo.toml @@ -11,6 +11,12 @@ description = "DiskANN/Vamana — SSD-friendly approximate nearest neighbor sear default = [] gpu = [] # Feature flag for GPU acceleration (CUDA/Metal stubs) simd = ["simsimd"] +# Route the f32 distance kernels through lattice-embed rather than SimSIMD. +# Takes precedence over `simd` and over the wasm32 SIMD128 path when both are +# enabled, so exactly one route is selected for any feature/target +# combination (the other backends' helper functions still compile; only the +# dispatch arm actually called changes). +lattice-simd = ["lattice-embed"] # BET 1 (ADR-200): fixed-topology reuse + periodic rebuild under metric drift. reuse-under-drift = [] @@ -24,6 +30,14 @@ thiserror = { workspace = true } rand = { workspace = true } parking_lot = "0.12" bytemuck = { version = "1.14", features = ["derive"] } +# Kernels only: `native` is the feature that pulls the model stack, so with +# default features off this is a pure-Rust SIMD dependency with no tokenizer, +# no model loader, and no download path. Unlike `simsimd` below it is not +# target-gated, because it compiles for wasm32 as well. The kernels themselves +# are pure Rust, but that does not make the feature's dependency tree +# C-toolchain-free: blake3 (a transitive dependency) needs a target C compiler +# on AArch64, where its build script compiles c/blake3_neon.c. +lattice-embed = { version = "0.7.1", optional = true, default-features = false } [target.'cfg(target_arch = "wasm32")'.dependencies] # `rand` (used by PQ k-means training) pulls getrandom 0.2 transitively, which diff --git a/crates/ruvector-diskann/src/distance.rs b/crates/ruvector-diskann/src/distance.rs index 3c7dc34c2b..34f4565c3a 100644 --- a/crates/ruvector-diskann/src/distance.rs +++ b/crates/ruvector-diskann/src/distance.rs @@ -1,8 +1,8 @@ //! Distance computations with SIMD acceleration and optional GPU offload //! -//! Dispatch priority: GPU (if `gpu` feature) → SimSIMD (if `simd` feature, native -//! NEON/AVX2/AVX-512) → WASM SIMD128 (`wasm32` target with `simd128` -//! target-feature) → scalar +//! Dispatch priority: GPU (if `gpu` feature) → lattice-embed (if `lattice-simd` +//! feature) → SimSIMD (if `simd` feature, native NEON/AVX2/AVX-512) → WASM +//! SIMD128 (`wasm32` target with `simd128` target-feature) → scalar use crate::error::{DiskAnnError, Result}; use memmap2::Mmap; @@ -189,24 +189,80 @@ impl FlatVectors { // Distance functions — auto-dispatch based on features // ============================================================================ +/// Set from inside [`l2_lattice`]/[`inner_lattice`] below, only after the +/// real kernel call returns — so a dispatch arm that swaps the wrapper call +/// for a scalar/native fallback bypasses the flag along with the kernel. +/// Thread-local (not a shared global) so a sibling test running on another +/// thread can't set the flag between this test's reset and its assert. +/// Compiled only for `lattice-simd` test builds — zero footprint elsewhere. +/// See `lattice_backend_is_actually_invoked` below. +#[cfg(all(test, feature = "lattice-simd"))] +thread_local! { + static LATTICE_L2_INVOKED: std::cell::Cell = const { std::cell::Cell::new(false) }; + static LATTICE_INNER_INVOKED: std::cell::Cell = const { std::cell::Cell::new(false) }; +} + +/// Sole caller of the `lattice-embed` L2 kernel — the dispatch arm in +/// [`l2_squared`] can only reach the kernel through this wrapper, so +/// replacing the arm's call expression with a fallback also removes the +/// witness store. +#[cfg(feature = "lattice-simd")] +#[inline] +fn l2_lattice(a: &[f32], b: &[f32]) -> f32 { + let result = lattice_embed::simd::squared_euclidean_distance(a, b); + #[cfg(test)] + LATTICE_L2_INVOKED.with(|invoked| invoked.set(true)); + result +} + +/// Sole caller of the `lattice-embed` dot-product kernel — see +/// [`l2_lattice`] for why the store lives here rather than inline in the +/// dispatch arm. Negated to match the other backends: this returns a +/// distance for a min-heap, not a similarity. `SpatialSimilarity::inner` is a +/// plain alias for `dot`, so both sides negate the same raw dot product. +#[cfg(feature = "lattice-simd")] +#[inline] +fn inner_lattice(a: &[f32], b: &[f32]) -> f32 { + let result = -lattice_embed::simd::dot_product(a, b); + #[cfg(test)] + LATTICE_INNER_INVOKED.with(|invoked| invoked.set(true)); + result +} + /// L2 squared distance — dispatches to best available implementation #[inline] pub fn l2_squared(a: &[f32], b: &[f32]) -> f32 { assert_eq!(a.len(), b.len(), "distance vectors must have equal lengths"); - #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + #[cfg(feature = "lattice-simd")] + { + l2_lattice(a, b) + } + + #[cfg(all( + not(feature = "lattice-simd"), + target_arch = "wasm32", + target_feature = "simd128" + ))] { wasm_simd128_l2_squared(a, b) } - #[cfg(all(not(target_arch = "wasm32"), feature = "simd"))] + #[cfg(all( + not(feature = "lattice-simd"), + not(target_arch = "wasm32"), + feature = "simd" + ))] { simd_l2_squared(a, b) } - #[cfg(any( - all(target_arch = "wasm32", not(target_feature = "simd128")), - all(not(target_arch = "wasm32"), not(feature = "simd")) + #[cfg(all( + not(feature = "lattice-simd"), + any( + all(target_arch = "wasm32", not(target_feature = "simd128")), + all(not(target_arch = "wasm32"), not(feature = "simd")) + ) ))] { scalar_l2_squared(a, b) @@ -313,27 +369,48 @@ pub fn wasm_simd128_l2_squared(a: &[f32], b: &[f32]) -> f32 { pub fn inner_product(a: &[f32], b: &[f32]) -> f32 { assert_eq!(a.len(), b.len(), "distance vectors must have equal lengths"); - #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + #[cfg(feature = "lattice-simd")] + { + inner_lattice(a, b) + } + + #[cfg(all( + not(feature = "lattice-simd"), + target_arch = "wasm32", + target_feature = "simd128" + ))] { wasm_simd128_inner_product(a, b) } - #[cfg(all(not(target_arch = "wasm32"), feature = "simd"))] + #[cfg(all( + not(feature = "lattice-simd"), + not(target_arch = "wasm32"), + feature = "simd" + ))] { simsimd::SpatialSimilarity::inner(a, b) .map(|d| -(d as f32)) .unwrap_or_else(|| scalar_inner_product(a, b)) } - #[cfg(any( - all(target_arch = "wasm32", not(target_feature = "simd128")), - all(not(target_arch = "wasm32"), not(feature = "simd")) + #[cfg(all( + not(feature = "lattice-simd"), + any( + all(target_arch = "wasm32", not(target_feature = "simd128")), + all(not(target_arch = "wasm32"), not(feature = "simd")) + ) ))] { scalar_inner_product(a, b) } } +/// Retained under `lattice-simd` as the fallback the `simd` backend still +/// calls if SimSIMD returns `None`. The parity test's reference is the local +/// `naive_inner_product`, not this function. `scalar_l2_squared` is `pub` and +/// so needs no such annotation. +#[cfg_attr(feature = "lattice-simd", allow(dead_code))] #[inline] fn scalar_inner_product(a: &[f32], b: &[f32]) -> f32 { let mut s0 = 0.0f32; @@ -579,6 +656,97 @@ mod tests { assert!((inner_product(&a, &b) - (-32.0)).abs() < 1e-6); } + /// Checks whichever backend compiled in against naive scalar references, + /// across dimensions chosen to straddle 4/8/16-lane widths **and their + /// remainders** so tail handling is exercised. The references are naive + /// single-pass loops rather than `scalar_l2_squared` / `scalar_inner_product`, + /// which are themselves 4-accumulator implementations, so a reduction-order + /// bug in the shared shape cannot hide. + /// + /// Backend-independent by construction: it covers the SimSIMD path, the + /// lattice path, and the plain scalar path. The cross-dimension parity + /// tests that existed before this were gated on + /// `all(target_arch = "wasm32", target_feature = "simd128")`, so no native + /// backend was checked at any dimension wider than 3. + #[test] + fn backend_matches_scalar_reference() { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + fn naive_l2_squared(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() + } + + fn naive_inner_product(a: &[f32], b: &[f32]) -> f32 { + -a.iter().zip(b).map(|(x, y)| x * y).sum::() + } + + const DIMS: &[usize] = &[ + 0, 1, 2, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 384, 768, 1000, 1023, 1024, + ]; + + // Combined absolute+relative tolerance, numpy `allclose`-style. A flat + // absolute bound is not achievable: each backend uses a different + // reduction tree, f32 addition is not associative, so reordering shifts + // rounding by a few ULPs and that drift grows with accumulated + // magnitude. A real bug (wrong lane math, dropped remainder) produces + // relative error orders of magnitude above machine epsilon. + const ATOL: f32 = 1e-5; + const RTOL: f32 = 1e-5; + + let mut rng = StdRng::seed_from_u64(42); + for &dim in DIMS { + let a: Vec = (0..dim).map(|_| rng.gen_range(-10.0f32..10.0)).collect(); + let b: Vec = (0..dim).map(|_| rng.gen_range(-10.0f32..10.0)).collect(); + + for (op, got, want) in [ + ("l2_squared", l2_squared(&a, &b), naive_l2_squared(&a, &b)), + ( + "inner_product", + inner_product(&a, &b), + naive_inner_product(&a, &b), + ), + ] { + let bound = ATOL + RTOL * got.abs().max(want.abs()); + assert!( + (got - want).abs() <= bound, + "{op} dim={dim}: got={got} want={want} bound={bound}" + ); + } + } + } + + /// `backend_matches_scalar_reference` above only checks numerical parity + /// against a naive oracle; scalar and lattice implement the same + /// arithmetic, so a `lattice-simd` build that silently fell back to the + /// scalar/native kernel would still pass it. This pins actual backend + /// *selection*: it fails if either dispatch arm's call to + /// [`l2_lattice`]/[`inner_lattice`] is replaced by a fallback route while + /// the feature stays declared, since the witness flags are only set + /// inside those wrappers, after the real kernel call returns. Thread-local + /// storage means a sibling test running concurrently on another thread + /// can't set these flags between this test's reset and its assert. + #[test] + #[cfg(feature = "lattice-simd")] + fn lattice_backend_is_actually_invoked() { + LATTICE_L2_INVOKED.with(|invoked| invoked.set(false)); + LATTICE_INNER_INVOKED.with(|invoked| invoked.set(false)); + + let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0]; + let b = vec![6.0f32, 7.0, 8.0, 9.0, 10.0]; + let _ = l2_squared(&a, &b); + let _ = inner_product(&a, &b); + + assert!( + LATTICE_L2_INVOKED.with(|invoked| invoked.get()), + "l2_squared did not route through the lattice-embed kernel" + ); + assert!( + LATTICE_INNER_INVOKED.with(|invoked| invoked.get()), + "inner_product did not route through the lattice-embed kernel" + ); + } + #[test] fn test_flat_vectors() { let mut fv = FlatVectors::new(3);