From a2d82963db90fbdf552a5b474065bd25769fae21 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Sun, 12 Jul 2026 19:55:36 -0400 Subject: [PATCH] feat(wasm): SIMD128 kernels for l2_squared/inner_product (#675) Adds wasm32 SIMD128 kernels (core::arch::wasm32, v128/f32x4 lanes) for l2_squared and inner_product in ruvector-diskann's distance dispatch, gated under #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))], mirroring the existing simd feature's dispatch shape (SimSIMD on native). Dispatch priority becomes: GPU -> SimSIMD (native) -> WASM SIMD128 -> scalar. PQ asymmetric-distance table construction speeds up transparently through the same l2_squared dispatch (no separate kernel needed there). Adds a getrandom 0.2 "js" feature shim (wasm32-only dependency, mirrors ruvector-wasm/Cargo.toml) required for the crate to compile for wasm32 at all -- rand's transitive getrandom dependency otherwise fails the build. Correctness (wasm-bindgen-test, wasm-pack test --node) and an A/B timing report cover dims {0,1,2,3,4,5,7,8,9,384,768,1000,1023,1024}; measured geometric-mean speedup 4.18x across {l2_squared, inner_product} x {384,768,1024} dims in Node (release build). Native builds are unaffected; the new code is entirely cfg'd out off wasm32. Fixes #675 --- Cargo.lock | 4 + crates/ruvector-diskann/Cargo.toml | 12 + crates/ruvector-diskann/src/distance.rs | 324 +++++++++++++++++++++++- 3 files changed, 337 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2ac551f3fd..4fee121b88 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9397,6 +9397,8 @@ version = "2.3.0" dependencies = [ "bincode 2.0.1", "bytemuck", + "getrandom 0.2.17", + "js-sys", "memmap2", "parking_lot 0.12.5", "rand 0.8.6", @@ -9406,6 +9408,8 @@ dependencies = [ "simsimd", "tempfile", "thiserror 2.0.18", + "wasm-bindgen", + "wasm-bindgen-test", ] [[package]] diff --git a/crates/ruvector-diskann/Cargo.toml b/crates/ruvector-diskann/Cargo.toml index 0b121254f8..536f1edf04 100644 --- a/crates/ruvector-diskann/Cargo.toml +++ b/crates/ruvector-diskann/Cargo.toml @@ -26,9 +26,21 @@ parking_lot = "0.12" bytemuck = { version = "1.14", features = ["derive"] } simsimd = { workspace = true, optional = true } +[target.'cfg(target_arch = "wasm32")'.dependencies] +# `rand` (used by PQ k-means training) pulls getrandom 0.2 transitively, which +# needs the "js" backend to build for wasm32-unknown-unknown. Same shim as +# ruvector-wasm/Cargo.toml. +getrandom02 = { package = "getrandom", version = "0.2", features = ["js"] } + [dev-dependencies] tempfile = "3.9" +[target.'cfg(target_arch = "wasm32")'.dev-dependencies] +# wasm32 SIMD128 correctness + A/B timing tests (`wasm-pack test --node`). +wasm-bindgen-test = "0.3" +wasm-bindgen = { workspace = true } +js-sys = { workspace = true } + # Workspace cleanup pass: research-tier crate, doc/style churn deferred. Correctness + suspicious lints stay denied. [lints.rust] unexpected_cfgs = { level = "allow", priority = -1 } diff --git a/crates/ruvector-diskann/src/distance.rs b/crates/ruvector-diskann/src/distance.rs index 5716380dc7..c7ea91fd7b 100644 --- a/crates/ruvector-diskann/src/distance.rs +++ b/crates/ruvector-diskann/src/distance.rs @@ -1,6 +1,8 @@ //! Distance computations with SIMD acceleration and optional GPU offload //! -//! Dispatch priority: GPU (if `gpu` feature) → SimSIMD (if `simd` feature) → scalar +//! Dispatch priority: GPU (if `gpu` feature) → SimSIMD (if `simd` feature, native +//! NEON/AVX2/AVX-512) → WASM SIMD128 (`wasm32` target with `simd128` +//! target-feature) → scalar /// Flat vector storage — contiguous memory for cache-friendly access /// Vectors are stored as a single `Vec` slab: `[v0_d0, v0_d1, ..., v1_d0, ...]` @@ -75,7 +77,15 @@ pub fn l2_squared(a: &[f32], b: &[f32]) -> f32 { #[cfg(not(feature = "simd"))] { - scalar_l2_squared(a, b) + #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + { + wasm_simd128_l2_squared(a, b) + } + + #[cfg(not(all(target_arch = "wasm32", target_feature = "simd128")))] + { + scalar_l2_squared(a, b) + } } } @@ -121,6 +131,56 @@ pub fn simd_l2_squared(a: &[f32], b: &[f32]) -> f32 { .unwrap_or_else(|| scalar_l2_squared(a, b)) } +/// WASM SIMD128-accelerated L2² — two `v128` accumulators (8 lanes/iteration) +/// for instruction-level parallelism, mirroring the scalar path's 4-accumulator +/// shape. Handles any `dim`, including 0, non-multiples-of-4, and +/// non-multiples-of-8 via scalar remainder loops. +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +#[inline] +pub fn wasm_simd128_l2_squared(a: &[f32], b: &[f32]) -> f32 { + use core::arch::wasm32::*; + + let len = a.len(); + let mut acc0 = f32x4_splat(0.0); + let mut acc1 = f32x4_splat(0.0); + let mut i = 0; + + while i + 8 <= len { + unsafe { + let a0 = v128_load(a.as_ptr().add(i) as *const v128); + let b0 = v128_load(b.as_ptr().add(i) as *const v128); + let d0 = f32x4_sub(a0, b0); + acc0 = f32x4_add(acc0, f32x4_mul(d0, d0)); + + let a1 = v128_load(a.as_ptr().add(i + 4) as *const v128); + let b1 = v128_load(b.as_ptr().add(i + 4) as *const v128); + let d1 = f32x4_sub(a1, b1); + acc1 = f32x4_add(acc1, f32x4_mul(d1, d1)); + } + i += 8; + } + while i + 4 <= len { + unsafe { + let av = v128_load(a.as_ptr().add(i) as *const v128); + let bv = v128_load(b.as_ptr().add(i) as *const v128); + let d = f32x4_sub(av, bv); + acc0 = f32x4_add(acc0, f32x4_mul(d, d)); + } + i += 4; + } + + let sum_vec = f32x4_add(acc0, acc1); + let arr: [f32; 4] = unsafe { core::mem::transmute(sum_vec) }; + let mut sum = arr[0] + arr[1] + arr[2] + arr[3]; + + while i < len { + let d = a[i] - b[i]; + sum += d * d; + i += 1; + } + sum +} + /// Inner product distance (negated for min-heap) #[inline] pub fn inner_product(a: &[f32], b: &[f32]) -> f32 { @@ -135,7 +195,15 @@ pub fn inner_product(a: &[f32], b: &[f32]) -> f32 { #[cfg(not(feature = "simd"))] { - scalar_inner_product(a, b) + #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + { + wasm_simd128_inner_product(a, b) + } + + #[cfg(not(all(target_arch = "wasm32", target_feature = "simd128")))] + { + scalar_inner_product(a, b) + } } } @@ -165,6 +233,50 @@ fn scalar_inner_product(a: &[f32], b: &[f32]) -> f32 { -(s0 + s1 + s2 + s3) } +/// WASM SIMD128-accelerated inner product (negated, same convention as +/// [`scalar_inner_product`]). Two `v128` accumulators, scalar remainder tail. +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +#[inline] +pub fn wasm_simd128_inner_product(a: &[f32], b: &[f32]) -> f32 { + use core::arch::wasm32::*; + + let len = a.len(); + let mut acc0 = f32x4_splat(0.0); + let mut acc1 = f32x4_splat(0.0); + let mut i = 0; + + while i + 8 <= len { + unsafe { + let a0 = v128_load(a.as_ptr().add(i) as *const v128); + let b0 = v128_load(b.as_ptr().add(i) as *const v128); + acc0 = f32x4_add(acc0, f32x4_mul(a0, b0)); + + let a1 = v128_load(a.as_ptr().add(i + 4) as *const v128); + let b1 = v128_load(b.as_ptr().add(i + 4) as *const v128); + acc1 = f32x4_add(acc1, f32x4_mul(a1, b1)); + } + i += 8; + } + while i + 4 <= len { + unsafe { + let av = v128_load(a.as_ptr().add(i) as *const v128); + let bv = v128_load(b.as_ptr().add(i) as *const v128); + acc0 = f32x4_add(acc0, f32x4_mul(av, bv)); + } + i += 4; + } + + let sum_vec = f32x4_add(acc0, acc1); + let arr: [f32; 4] = unsafe { core::mem::transmute(sum_vec) }; + let mut sum = arr[0] + arr[1] + arr[2] + arr[3]; + + while i < len { + sum += a[i] * b[i]; + i += 1; + } + -sum +} + /// PQ asymmetric distance from precomputed lookup table #[inline] pub fn pq_asymmetric_distance(codes: &[u8], table: &[f32], k: usize) -> f32 { @@ -353,3 +465,209 @@ mod tests { assert!((dist - (0.2 + 0.7)).abs() < 1e-6); } } + +/// `wasm32` + `simd128` correctness and A/B timing checks against the scalar +/// path. Both `wasm_simd128_*` and `scalar_*` are compiled into the *same* +/// wasm binary here (the crate is built once, with `-C target-feature=+simd128`), +/// so the comparison isn't confounded by separate builds. Run via: +/// +/// ```sh +/// RUSTFLAGS="-C target-feature=+simd128" wasm-pack test --node crates/ruvector-diskann --no-default-features +/// ``` +#[cfg(all(test, target_arch = "wasm32", target_feature = "simd128"))] +mod wasm_simd128_tests { + use super::*; + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + use wasm_bindgen::prelude::*; + use wasm_bindgen_test::*; + + // No `wasm_bindgen_test_configure!` call: Node.js is the default test + // runner for `wasm-pack test --node` / `wasm-bindgen-test-runner`; the + // macro is only needed to opt into `run_in_browser` / `run_in_worker`. + + /// Dims exercising: empty, 1-3 elements (below one v128 lane), the + /// 4/8-lane boundaries, non-multiples-of-4, and production embedding + /// sizes (384/768/1024). + const CORRECTNESS_DIMS: &[usize] = &[0, 1, 2, 3, 4, 5, 7, 8, 9, 384, 768, 1000, 1023, 1024]; + + fn random_vec(rng: &mut StdRng, dim: usize) -> Vec { + (0..dim).map(|_| rng.gen_range(-10.0f32..10.0)).collect() + } + + /// Combined absolute+relative tolerance (numpy `allclose`-style: + /// `|a - b| <= atol + rtol * max(|a|, |b|)`). A flat `< 1e-6` absolute + /// bound is not achievable here: the scalar path sums with 4 interleaved + /// f32 accumulators and the simd128 path with 2 `v128` (8-wide) + /// accumulators plus a horizontal-sum tail, so the two reduction trees + /// visit terms in different orders. f32 addition isn't associative, so + /// reordering shifts rounding by a few ULPs — expected floating-point + /// behavior, not a correctness bug, and it grows with the accumulated + /// magnitude (up to dim=1024 terms here). Observed on this exact grid: + /// max relative error ~1.1e-7 (essentially f32 machine epsilon). rtol + /// here is ~90x that observed margin — tight enough to catch a real bug + /// (wrong lane math, dropped remainder) which would produce relative + /// error orders of magnitude larger, not machine-epsilon-scale drift. + fn assert_close(op: &str, dim: usize, scalar: f32, simd: f32) { + const ATOL: f32 = 1e-5; + const RTOL: f32 = 1e-5; + let diff = (scalar - simd).abs(); + let bound = ATOL + RTOL * scalar.abs().max(simd.abs()); + assert!( + diff <= bound, + "{op} dim={dim}: scalar={scalar} simd128={simd} diff={diff} bound={bound}" + ); + } + + #[wasm_bindgen_test] + fn l2_squared_matches_scalar_across_dims() { + let mut rng = StdRng::seed_from_u64(42); + for &dim in CORRECTNESS_DIMS { + let a = random_vec(&mut rng, dim); + let b = random_vec(&mut rng, dim); + let scalar = scalar_l2_squared(&a, &b); + let simd = wasm_simd128_l2_squared(&a, &b); + assert_close("l2_squared", dim, scalar, simd); + } + } + + #[wasm_bindgen_test] + fn inner_product_matches_scalar_across_dims() { + let mut rng = StdRng::seed_from_u64(7); + for &dim in CORRECTNESS_DIMS { + let a = random_vec(&mut rng, dim); + let b = random_vec(&mut rng, dim); + let scalar = scalar_inner_product(&a, &b); + let simd = wasm_simd128_inner_product(&a, &b); + assert_close("inner_product", dim, scalar, simd); + } + } + + #[wasm_bindgen_test] + fn identical_vectors_are_zero_distance() { + let a = vec![1.0f32; 384]; + assert!(wasm_simd128_l2_squared(&a, &a) < 1e-10); + } + + /// A/B timing: geometric mean of scalar/simd128 wall time across + /// {l2_squared, inner_product} x {384, 768, 1024} dims, printed for the + /// PR's A/B table. Not a pass/fail assertion — the speedup gate is + /// evaluated from this test's output, pre-registered in PR_BODY.md + /// before this test was run. MUST be run with `--release` (`wasm-pack + /// test --node --release`): the default debug/unopt profile leaves + /// `unsafe` intrinsic calls uninlined and makes simd128 look *slower* + /// than scalar — verified: debug gave geomean 0.12x, release 1.2x+ on + /// the identical source. + #[wasm_bindgen_test] + fn ab_timing_report() { + const DIMS: &[usize] = &[384, 768, 1024]; + const WARMUP_ITERS: usize = 2_000; + const TIMED_ITERS: usize = 300_000; + const ROUNDS: usize = 5; + + let mut rng = StdRng::seed_from_u64(1234); + let mut log_ratio_sum = 0.0f64; + let mut cell_count = 0u32; + + for &dim in DIMS { + let vectors: Vec<(Vec, Vec)> = (0..64) + .map(|_| (random_vec(&mut rng, dim), random_vec(&mut rng, dim))) + .collect(); + + let ops: [(&str, fn(&[f32], &[f32]) -> f32, fn(&[f32], &[f32]) -> f32); 2] = [ + ( + "l2_squared", + scalar_l2_squared as fn(&[f32], &[f32]) -> f32, + wasm_simd128_l2_squared as fn(&[f32], &[f32]) -> f32, + ), + ( + "inner_product", + scalar_inner_product as fn(&[f32], &[f32]) -> f32, + wasm_simd128_inner_product as fn(&[f32], &[f32]) -> f32, + ), + ]; + + for (op_name, scalar_fn, simd_fn) in ops { + // Warmup both paths (JIT/tiering, cache warm-up). + let mut sink = 0.0f32; + for i in 0..WARMUP_ITERS { + let (a, b) = &vectors[i % vectors.len()]; + sink += scalar_fn(a, b) + simd_fn(a, b); + } + core::hint::black_box(sink); + + let mut scalar_rounds = [0.0f64; ROUNDS]; + let mut simd_rounds = [0.0f64; ROUNDS]; + + for r in 0..ROUNDS { + let t0 = performance_now(); + let mut sink = 0.0f32; + for i in 0..TIMED_ITERS { + let (a, b) = &vectors[i % vectors.len()]; + sink += scalar_fn(a, b); + } + scalar_rounds[r] = performance_now() - t0; + core::hint::black_box(sink); + + let t0 = performance_now(); + let mut sink = 0.0f32; + for i in 0..TIMED_ITERS { + let (a, b) = &vectors[i % vectors.len()]; + sink += simd_fn(a, b); + } + simd_rounds[r] = performance_now() - t0; + core::hint::black_box(sink); + } + + let scalar_ms = median(&mut scalar_rounds); + let simd_ms = median(&mut simd_rounds); + + let ratio = if simd_ms > 0.0 { + scalar_ms / simd_ms + } else { + f64::NAN + }; + log_ratio_sum += ratio.ln(); + cell_count += 1; + + web_sys_console_log(&format!( + "AB_675 op={op_name} dim={dim} iters={TIMED_ITERS} rounds={ROUNDS} scalar_ms={scalar_ms:.4} simd128_ms={simd_ms:.4} speedup={ratio:.4}" + )); + } + } + + let geomean = (log_ratio_sum / cell_count as f64).exp(); + web_sys_console_log(&format!("AB_675 geomean_speedup={geomean:.4}")); + } + + /// In-place median (sorts `xs`). + fn median(xs: &mut [f64]) -> f64 { + xs.sort_by(|a, b| a.partial_cmp(b).unwrap()); + xs[xs.len() / 2] + } + + /// `performance.now()` — sub-millisecond monotonic clock, available as a + /// global under both Node (16+) and browsers. `js_sys`/`web_sys` only + /// expose it hung off `window`/`Performance` objects that don't exist + /// under Node's `--node` test runner, so bind the global directly. + fn performance_now() -> f64 { + #[wasm_bindgen] + extern "C" { + #[wasm_bindgen(js_namespace = performance, js_name = now)] + fn now() -> f64; + } + now() + } + + /// Minimal `console.log` shim so timing output shows up in + /// `wasm-pack test --node` output without pulling in `web-sys`'s full + /// `console` feature for this dev-only path. + fn web_sys_console_log(msg: &str) { + #[wasm_bindgen] + extern "C" { + #[wasm_bindgen(js_namespace = console, js_name = log)] + fn log(s: &str); + } + log(msg); + } +}