From cbf2f5f9776a9dad80645a3d43dfe0455c73a739 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Sun, 2 Aug 2026 13:45:15 -0400 Subject: [PATCH 1/4] feat(graph): optional lattice-embed kernels for the schema scan path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Routes the schema layer's dot product and Euclidean scoring through lattice-embed behind an off-by-default `lattice-simd` feature, with a passthrough feature on ruvector-graph-wasm. The dependency is optional and its kernels compile for wasm32, which is what keeps this layer WASM-safe and no-feature-build-safe — the two properties that kept simsimd out of it. Equal-length guards preserve the existing truncating behaviour for mismatched slices, which the kernels do not share. --- crates/ruvector-graph-wasm/Cargo.toml | 7 ++ crates/ruvector-graph/Cargo.toml | 11 +++ crates/ruvector-graph/src/schema.rs | 110 +++++++++++++++++++++++++- 3 files changed, 125 insertions(+), 3 deletions(-) diff --git a/crates/ruvector-graph-wasm/Cargo.toml b/crates/ruvector-graph-wasm/Cargo.toml index 38c51ad733..354a087e1e 100644 --- a/crates/ruvector-graph-wasm/Cargo.toml +++ b/crates/ruvector-graph-wasm/Cargo.toml @@ -62,7 +62,14 @@ wasm-bindgen-test = "0.3" [features] default = [] +# Note: on wasm32 this is effectively a no-op, because ruvector-core excludes +# simsimd on that target and falls back to scalar. simd = ["ruvector-core/simd"] +# Vectorized kernels that do reach wasm32, for both the core distance functions +# and the graph schema scan path. Build with +# `RUSTFLAGS="-C target-feature=+simd128"`; without that flag the kernels fall +# back to scalar and this feature is a no-op. +lattice-simd = ["ruvector-graph/lattice-simd"] # Ensure getrandom uses wasm_js/js features for WASM [target.'cfg(target_arch = "wasm32")'.dependencies] diff --git a/crates/ruvector-graph/Cargo.toml b/crates/ruvector-graph/Cargo.toml index b96d8a3e1e..92e470fafc 100644 --- a/crates/ruvector-graph/Cargo.toml +++ b/crates/ruvector-graph/Cargo.toml @@ -23,6 +23,10 @@ hnsw_rs = { workspace = true, optional = true } # SIMD and performance simsimd = { workspace = true, optional = true } +# Kernels only: `default-features = false` keeps `lattice-inference` and the +# model/tokenizer/download stack out of the tree. Unlike simsimd, these kernels +# build for wasm32, which is what lets the schema layer stay WASM-safe. +lattice-embed = { version = "0.7.0", optional = true, default-features = false } rayon = { workspace = true } crossbeam = { workspace = true } num_cpus = "1.16" @@ -112,6 +116,13 @@ full = ["simd", "storage", "async-runtime", "compression", "hnsw_rs", "ruvector- # SIMD optimizations simd = ["ruvector-core/simd", "simsimd"] +# Explicit SIMD kernels for the schema-layer scan path, via lattice-embed. +# Distinct from `simd` above in two ways: it reaches wasm32 (where simsimd is +# unavailable, so `simd` leaves the scan scalar), and it carries no non-optional +# dependency, so a no-feature build is unchanged. Raises the effective MSRV to +# 1.93 for anyone who enables it. Off by default. +lattice-simd = ["dep:lattice-embed"] + # Storage backends storage = ["redb", "memmap2"] diff --git a/crates/ruvector-graph/src/schema.rs b/crates/ruvector-graph/src/schema.rs index 65ee6c40fe..7032c0777e 100644 --- a/crates/ruvector-graph/src/schema.rs +++ b/crates/ruvector-graph/src/schema.rs @@ -94,6 +94,13 @@ impl DistanceMetric { // Single fused pass: accumulate q·c and c·c together so the // candidate slice is read once (half the memory traffic of two // separate `dot` calls). + // + // Deliberately left scalar even under `lattice-simd`. A kernel + // cosine takes only `(a, b)` and recomputes the query norm per + // candidate, which would discard the `query_norm` hoist this + // signature exists for, and would silently ignore a caller-supplied + // `query_norm` that differs from `‖query‖`. Vectorizing this arm + // needs a kernel that accepts a precomputed query norm. let n = query.len().min(candidate.len()); let mut qc = 0.0f32; let mut cc = 0.0f32; @@ -110,6 +117,15 @@ impl DistanceMetric { } } DistanceMetric::Euclidean => { + // Same equal-length guard as `dot`: the loop below truncates to + // the shorter slice, the kernel does not. + #[cfg(feature = "lattice-simd")] + { + if query.len() == candidate.len() { + return -lattice_embed::simd::euclidean_distance(query, candidate); + } + } + let n = query.len().min(candidate.len()); let mut sum = 0.0f32; for i in 0..n { @@ -156,10 +172,24 @@ pub fn score_property( #[inline] fn dot(a: &[f32], b: &[f32]) -> f32 { + // With `lattice-simd`, use explicit kernels rather than relying on + // autovectorization. This stays WASM-safe and no-feature-build-safe, the two + // properties that kept `simsimd` out of this layer: the dependency is + // optional, and its kernels compile to `simd128` on wasm32 where `simsimd` + // is unavailable. + // + // The equal-length guard is required, not defensive. The iterator form below + // truncates to the shorter slice, while the kernel returns 0.0 on a length + // mismatch, so unequal inputs must keep taking the scalar path to preserve + // this function's existing behaviour. + #[cfg(feature = "lattice-simd")] + { + if a.len() == b.len() { + return lattice_embed::simd::dot_product(a, b); + } + } + // Iterator form so LLVM auto-vectorizes (SSE/AVX/NEON) without bounds checks. - // SIMD via `simsimd`/ruvector-core is a follow-up (ADR-252 P5) but is - // deliberately not a hard dependency here so the schema layer stays WASM- and - // no-feature-build-safe. a.iter().zip(b.iter()).map(|(x, y)| x * y).sum() } @@ -688,4 +718,78 @@ mod tests { ); assert!(s.validate_node(&n).is_ok()); } + + /// Deterministic pseudo-random vectors, no dev-dependency needed. + fn vecs(dim: usize, seed: u32) -> (Vec, Vec) { + let mut s = seed.wrapping_mul(2_654_435_761).wrapping_add(1); + let mut next = || { + s ^= s << 13; + s ^= s >> 17; + s ^= s << 5; + (s as f32 / u32::MAX as f32) * 2.0 - 1.0 + }; + ( + (0..dim).map(|_| next()).collect(), + (0..dim).map(|_| next()).collect(), + ) + } + + /// Whichever `dot` backend is compiled in must agree with a naive scalar sum. + /// + /// This is what catches a wrong kernel or a wrong adapter: swapping + /// `dot_product` for a different kernel, or dropping the negation on the + /// Euclidean arm, still compiles and fails here. + #[test] + fn test_score_pre_matches_scalar_reference() { + // Dimensions straddling 4/8/16-lane widths and their remainders. + for dim in [1usize, 3, 4, 7, 8, 15, 16, 17, 31, 64, 127, 384, 768] { + for seed in 0..4u32 { + let (q, c) = vecs(dim, seed); + + let want_dot: f32 = q.iter().zip(&c).map(|(x, y)| x * y).sum(); + let got_dot = DistanceMetric::DotProduct.score_pre(&q, &c, 0.0); + assert!( + (got_dot - want_dot).abs() <= 1e-3 * want_dot.abs().max(1.0), + "dot mismatch dim={dim} seed={seed}: got {got_dot}, want {want_dot}" + ); + + let want_euc = + -q.iter().zip(&c).map(|(x, y)| (x - y) * (x - y)).sum::().sqrt(); + let got_euc = DistanceMetric::Euclidean.score_pre(&q, &c, 0.0); + assert!( + (got_euc - want_euc).abs() <= 1e-3 * want_euc.abs().max(1.0), + "euclidean mismatch dim={dim} seed={seed}: got {got_euc}, want {want_euc}" + ); + + let qn = DistanceMetric::Cosine.query_norm(&q); + let want_qn = q.iter().map(|x| x * x).sum::().sqrt(); + assert!( + (qn - want_qn).abs() <= 1e-3 * want_qn.abs().max(1.0), + "query_norm mismatch dim={dim} seed={seed}: got {qn}, want {want_qn}" + ); + } + } + } + + /// Unequal lengths must keep the truncating scalar behaviour. The kernels + /// return 0.0 on a length mismatch, so a missing guard would show up here as + /// a zero instead of the truncated dot product. + #[test] + fn test_unequal_lengths_truncate_not_zero() { + let q = vec![1.0f32, 2.0, 3.0, 4.0]; + let c = vec![1.0f32, 1.0, 1.0]; + + let got = DistanceMetric::DotProduct.score_pre(&q, &c, 0.0); + assert!( + (got - 6.0).abs() < 1e-5, + "expected truncated dot 6.0, got {got}" + ); + + let got = DistanceMetric::Euclidean.score_pre(&q, &c, 0.0); + let want = -(0.0f32 + 1.0 + 4.0).sqrt(); + assert!( + (got - want).abs() < 1e-5, + "expected truncated euclidean {want}, got {got}" + ); + } } From afe654ab688dea465710ae4184bc47706c12986e Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Sun, 2 Aug 2026 13:48:33 -0400 Subject: [PATCH 2/4] feat(graph): route the cosine scan arm through the precomputed-norm kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the schema scan path: the cosine arm previously stayed scalar because a two-argument kernel cosine recomputes the query norm per candidate, which discards the hoist `score_pre`'s signature exists for and ignores a caller-supplied norm. lattice-embed 0.7.1 adds a kernel that takes the precomputed norm and rescales by it, so the arm now routes with the same equal-length guard as the others. The pin moves to 0.7.1 because that function does not exist in 0.7.0, making it a build requirement rather than a tidy-up. Tests gain cosine score coverage (previously only its hoisted norm was checked), an assertion that a caller-supplied norm rescales rather than being ignored, and a truncation case for the cosine arm. Also corrects a comment on ruvector-graph-wasm's `simd` feature: it said ruvector-core excludes simsimd on wasm32. It does not — the crate still resolves into the wasm32 graph. What is gated is core's SimSIMD call sites, on `not(target_arch = "wasm32")`, so the feature takes the scalar arm there. --- Cargo.lock | 21 +++++++- crates/ruvector-graph-wasm/Cargo.toml | 6 ++- crates/ruvector-graph/Cargo.toml | 2 +- crates/ruvector-graph/src/schema.rs | 73 +++++++++++++++++++++++---- 4 files changed, 87 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 03d9322bc8..bf41af3d16 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", @@ -9636,6 +9654,7 @@ dependencies = [ "hnsw_rs", "hyper 1.10.1", "lalrpop-util", + "lattice-embed 0.7.1", "lru 0.16.4", "lz4", "memmap2", diff --git a/crates/ruvector-graph-wasm/Cargo.toml b/crates/ruvector-graph-wasm/Cargo.toml index 354a087e1e..8d21fd3fe3 100644 --- a/crates/ruvector-graph-wasm/Cargo.toml +++ b/crates/ruvector-graph-wasm/Cargo.toml @@ -62,8 +62,10 @@ wasm-bindgen-test = "0.3" [features] default = [] -# Note: on wasm32 this is effectively a no-op, because ruvector-core excludes -# simsimd on that target and falls back to scalar. +# On wasm32 this does not vectorize: ruvector-core gates its SimSIMD call sites +# on `not(target_arch = "wasm32")` and takes the scalar arm there. Note it is the +# call sites that are cfg'd out, not the crate — `simsimd` still resolves into +# the wasm32 dependency graph. simd = ["ruvector-core/simd"] # Vectorized kernels that do reach wasm32, for both the core distance functions # and the graph schema scan path. Build with diff --git a/crates/ruvector-graph/Cargo.toml b/crates/ruvector-graph/Cargo.toml index 92e470fafc..513b2a2be7 100644 --- a/crates/ruvector-graph/Cargo.toml +++ b/crates/ruvector-graph/Cargo.toml @@ -26,7 +26,7 @@ simsimd = { workspace = true, optional = true } # Kernels only: `default-features = false` keeps `lattice-inference` and the # model/tokenizer/download stack out of the tree. Unlike simsimd, these kernels # build for wasm32, which is what lets the schema layer stay WASM-safe. -lattice-embed = { version = "0.7.0", optional = true, default-features = false } +lattice-embed = { version = "0.7.1", optional = true, default-features = false } rayon = { workspace = true } crossbeam = { workspace = true } num_cpus = "1.16" diff --git a/crates/ruvector-graph/src/schema.rs b/crates/ruvector-graph/src/schema.rs index 7032c0777e..d7ccfce2c7 100644 --- a/crates/ruvector-graph/src/schema.rs +++ b/crates/ruvector-graph/src/schema.rs @@ -91,16 +91,25 @@ impl DistanceMetric { match self { DistanceMetric::DotProduct => dot(query, candidate), DistanceMetric::Cosine => { + // Takes the precomputed `query_norm` rather than recomputing + // `‖query‖` per candidate, so the hoist this signature exists for + // survives, and a caller-supplied norm that differs from `‖query‖` + // rescales the result exactly as the scalar arm below does. + // + // Equal-length guard for the same reason as `dot`: the fused loop + // truncates to the shorter slice, the kernel returns 0.0. + #[cfg(feature = "lattice-simd")] + { + if query.len() == candidate.len() { + return lattice_embed::simd::cosine_similarity_pre_normalized( + query, candidate, query_norm, + ); + } + } + // Single fused pass: accumulate q·c and c·c together so the // candidate slice is read once (half the memory traffic of two // separate `dot` calls). - // - // Deliberately left scalar even under `lattice-simd`. A kernel - // cosine takes only `(a, b)` and recomputes the query norm per - // candidate, which would discard the `query_norm` hoist this - // signature exists for, and would silently ignore a caller-supplied - // `query_norm` that differs from `‖query‖`. Vectorizing this arm - // needs a kernel that accepts a precomputed query norm. let n = query.len().min(candidate.len()); let mut qc = 0.0f32; let mut cc = 0.0f32; @@ -175,8 +184,8 @@ fn dot(a: &[f32], b: &[f32]) -> f32 { // With `lattice-simd`, use explicit kernels rather than relying on // autovectorization. This stays WASM-safe and no-feature-build-safe, the two // properties that kept `simsimd` out of this layer: the dependency is - // optional, and its kernels compile to `simd128` on wasm32 where `simsimd` - // is unavailable. + // optional, and its kernels compile to `simd128` on wasm32 — where the + // SimSIMD-backed paths are cfg'd out and fall back to scalar. // // The equal-length guard is required, not defensive. The iterator form below // truncates to the shorter slice, while the kernel returns 0.0 on a length @@ -753,8 +762,12 @@ mod tests { "dot mismatch dim={dim} seed={seed}: got {got_dot}, want {want_dot}" ); - let want_euc = - -q.iter().zip(&c).map(|(x, y)| (x - y) * (x - y)).sum::().sqrt(); + let want_euc = -q + .iter() + .zip(&c) + .map(|(x, y)| (x - y) * (x - y)) + .sum::() + .sqrt(); let got_euc = DistanceMetric::Euclidean.score_pre(&q, &c, 0.0); assert!( (got_euc - want_euc).abs() <= 1e-3 * want_euc.abs().max(1.0), @@ -767,6 +780,35 @@ mod tests { (qn - want_qn).abs() <= 1e-3 * want_qn.abs().max(1.0), "query_norm mismatch dim={dim} seed={seed}: got {qn}, want {want_qn}" ); + + // The cosine score itself, not just its hoisted norm. Reference is + // built from the naive sums rather than from `query_norm` above, so + // a norm that is wrong in the same direction on both sides cannot + // cancel out and pass. + let qc: f32 = q.iter().zip(&c).map(|(x, y)| x * y).sum(); + let cn = c.iter().map(|y| y * y).sum::().sqrt(); + let want_cos = if want_qn == 0.0 || cn == 0.0 { + 0.0 + } else { + qc / (want_qn * cn) + }; + let got_cos = DistanceMetric::Cosine.score_pre(&q, &c, qn); + assert!( + (got_cos - want_cos).abs() <= 1e-3 * want_cos.abs().max(1.0), + "cosine mismatch dim={dim} seed={seed}: got {got_cos}, want {want_cos}" + ); + + // A caller-supplied norm that is not `‖query‖` must rescale the + // result, not be ignored. This is the property that made a plain + // two-argument kernel unusable for this signature, so it is + // asserted rather than assumed. + let scaled = DistanceMetric::Cosine.score_pre(&q, &c, qn * 2.0); + assert!( + (scaled - want_cos / 2.0).abs() <= 1e-3 * (want_cos / 2.0).abs().max(1.0), + "cosine ignored the supplied query_norm dim={dim} seed={seed}: \ + got {scaled}, want {}", + want_cos / 2.0 + ); } } } @@ -791,5 +833,14 @@ mod tests { (got - want).abs() < 1e-5, "expected truncated euclidean {want}, got {got}" ); + + // Cosine truncates too: q·c and ‖c‖ both over the first 3 lanes. + let qn = (1.0f32 + 4.0 + 9.0 + 16.0).sqrt(); + let got = DistanceMetric::Cosine.score_pre(&q, &c, qn); + let want = 6.0f32 / (qn * 3.0f32.sqrt()); + assert!( + (got - want).abs() < 1e-5, + "expected truncated cosine {want}, got {got}" + ); } } From 9060c776462903246752527abe224a2d49469440 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:04:56 -0400 Subject: [PATCH 3/4] test(graph): pin the pre-normalized cosine backend selection Add a lattice-simd-gated regression test that fails if the equal-length cosine arm in DistanceMetric::score_pre stops calling lattice_embed::simd::cosine_similarity_pre_normalized and silently falls through to the scalar reference implementation. The existing parity tests only compare numerical output, which the scalar fallback also satisfies. Co-Authored-By: Claude Sonnet 5 --- crates/ruvector-graph/src/schema.rs | 34 +++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/crates/ruvector-graph/src/schema.rs b/crates/ruvector-graph/src/schema.rs index d7ccfce2c7..107b93c69d 100644 --- a/crates/ruvector-graph/src/schema.rs +++ b/crates/ruvector-graph/src/schema.rs @@ -19,6 +19,14 @@ use crate::types::PropertyValue; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +/// Test-only witness for which cosine backend `score_pre` actually took. Lets a +/// test assert *selection*, not just numerical agreement: a scalar reference +/// can match the kernel's output by construction while the kernel call itself +/// has been silently reverted to the fallback arm. +#[cfg(all(test, feature = "lattice-simd"))] +static COSINE_LATTICE_ROUTE_HIT: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + /// Declared type of a node/edge property. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum PropertyType { @@ -101,6 +109,8 @@ impl DistanceMetric { #[cfg(feature = "lattice-simd")] { if query.len() == candidate.len() { + #[cfg(test)] + COSINE_LATTICE_ROUTE_HIT.store(true, std::sync::atomic::Ordering::Relaxed); return lattice_embed::simd::cosine_similarity_pre_normalized( query, candidate, query_norm, ); @@ -843,4 +853,28 @@ mod tests { "expected truncated cosine {want}, got {got}" ); } + + /// Guards *selection*, not just numerical agreement: reverting the + /// `lattice-simd` cosine arm to the scalar fallback still produces a + /// correct score (that's the point of the fallback), so + /// `test_score_pre_matches_scalar_reference` alone would keep passing. + /// This test fails if the equal-length branch stops calling + /// `lattice_embed::simd::cosine_similarity_pre_normalized`. + #[cfg(feature = "lattice-simd")] + #[test] + fn test_cosine_equal_length_routes_through_lattice_backend() { + COSINE_LATTICE_ROUTE_HIT.store(false, std::sync::atomic::Ordering::Relaxed); + + let q = vec![1.0f32, 2.0, 3.0, 4.0]; + let c = vec![4.0f32, 3.0, 2.0, 1.0]; + let qn = DistanceMetric::Cosine.query_norm(&q); + let _ = DistanceMetric::Cosine.score_pre(&q, &c, qn); + + assert!( + COSINE_LATTICE_ROUTE_HIT.load(std::sync::atomic::Ordering::Relaxed), + "expected the equal-length cosine path to call \ + lattice_embed::simd::cosine_similarity_pre_normalized; \ + the scalar fallback ran instead" + ); + } } From d241ad2e6b3c41dcb3b832ed53720045b45184c4 Mon Sep 17 00:00:00 2001 From: OceanLi <122793010+ohdearquant@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:41:15 -0400 Subject: [PATCH 4/4] test(graph): witness the cosine backend after the call returns Move the lattice-cosine selection witness from a pre-call flag to a per-thread Cell holding the value the kernel actually returned, set only once the call has completed. The selection test cross-checks that value bit-for-bit against a fresh, independent call into lattice_embed::simd::cosine_similarity_pre_normalized, so a reversion that swaps in the scalar computation (even keeping the same store) still diverges. thread_local storage also means another test's cosine calls, running on their own thread, cannot satisfy this test's assertion. --- crates/ruvector-graph/src/schema.rs | 66 ++++++++++++++++++++++------- 1 file changed, 50 insertions(+), 16 deletions(-) diff --git a/crates/ruvector-graph/src/schema.rs b/crates/ruvector-graph/src/schema.rs index 107b93c69d..27fd473d79 100644 --- a/crates/ruvector-graph/src/schema.rs +++ b/crates/ruvector-graph/src/schema.rs @@ -23,9 +23,17 @@ use std::collections::HashMap; /// test assert *selection*, not just numerical agreement: a scalar reference /// can match the kernel's output by construction while the kernel call itself /// has been silently reverted to the fallback arm. +/// +/// Per-thread (rather than a shared global) so that `cargo test`'s default +/// one-thread-per-test execution can't let one test's cosine calls satisfy +/// another test's assertion. Holds the actual value the lattice call +/// returned (not just a flag), set only after that call returns, so the +/// witness proves what came back rather than merely that a guarded branch +/// was entered. #[cfg(all(test, feature = "lattice-simd"))] -static COSINE_LATTICE_ROUTE_HIT: std::sync::atomic::AtomicBool = - std::sync::atomic::AtomicBool::new(false); +thread_local! { + static COSINE_LATTICE_ROUTE_HIT: std::cell::Cell> = const { std::cell::Cell::new(None) }; +} /// Declared type of a node/edge property. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -109,11 +117,12 @@ impl DistanceMetric { #[cfg(feature = "lattice-simd")] { if query.len() == candidate.len() { - #[cfg(test)] - COSINE_LATTICE_ROUTE_HIT.store(true, std::sync::atomic::Ordering::Relaxed); - return lattice_embed::simd::cosine_similarity_pre_normalized( + let result = lattice_embed::simd::cosine_similarity_pre_normalized( query, candidate, query_norm, ); + #[cfg(test)] + COSINE_LATTICE_ROUTE_HIT.with(|hit| hit.set(Some(result))); + return result; } } @@ -858,23 +867,48 @@ mod tests { /// `lattice-simd` cosine arm to the scalar fallback still produces a /// correct score (that's the point of the fallback), so /// `test_score_pre_matches_scalar_reference` alone would keep passing. - /// This test fails if the equal-length branch stops calling - /// `lattice_embed::simd::cosine_similarity_pre_normalized`. + /// + /// The witness is written only after + /// `lattice_embed::simd::cosine_similarity_pre_normalized` returns, and + /// this test cross-checks the recorded value bit-for-bit against a + /// second, independent direct call to that same kernel function. A + /// reversion that swaps the call for an inline scalar computation (even + /// one bound to the same local and stored the same way) still shows up + /// here: the scalar sum's rounding practically never matches the + /// kernel's, so the two values diverge. Reverting the whole arm instead + /// leaves the witness unset, which the `.expect` below catches. #[cfg(feature = "lattice-simd")] #[test] fn test_cosine_equal_length_routes_through_lattice_backend() { - COSINE_LATTICE_ROUTE_HIT.store(false, std::sync::atomic::Ordering::Relaxed); + COSINE_LATTICE_ROUTE_HIT.with(|hit| hit.set(None)); - let q = vec![1.0f32, 2.0, 3.0, 4.0]; - let c = vec![4.0f32, 3.0, 2.0, 1.0]; + // dim=17 straddles the 16-lane width with a one-element remainder, + // so the kernel's reduction order can't coincidentally match a + // sequential scalar sum. + let (q, c) = vecs(17, 5); let qn = DistanceMetric::Cosine.query_norm(&q); - let _ = DistanceMetric::Cosine.score_pre(&q, &c, qn); + let got = DistanceMetric::Cosine.score_pre(&q, &c, qn); - assert!( - COSINE_LATTICE_ROUTE_HIT.load(std::sync::atomic::Ordering::Relaxed), - "expected the equal-length cosine path to call \ - lattice_embed::simd::cosine_similarity_pre_normalized; \ - the scalar fallback ran instead" + let recorded = COSINE_LATTICE_ROUTE_HIT.with(|hit| hit.get()).expect( + "expected the equal-length cosine path to record a post-call \ + witness; the scalar fallback ran instead (or the lattice \ + arm never returned through the witnessed path)", + ); + assert_eq!( + recorded.to_bits(), + got.to_bits(), + "witness value diverged from score_pre's own return value" + ); + + let direct = lattice_embed::simd::cosine_similarity_pre_normalized(&q, &c, qn); + assert_eq!( + recorded.to_bits(), + direct.to_bits(), + "expected the equal-length cosine path's witnessed value to \ + bit-match a fresh, independent call into \ + lattice_embed::simd::cosine_similarity_pre_normalized; got a \ + different value, so the scan path did not actually return the \ + kernel's result" ); } }