nodedb_vector/distance/simd/
runtime.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Runtime SIMD detection and dispatch table.
4
5use super::bbq;
6use super::hamming::fast_hamming;
7use super::scalar::{scalar_cosine, scalar_ip, scalar_l2};
8use crate::distance::typed_scalar;
9
10#[cfg(target_arch = "x86_64")]
11use super::{avx2, avx512};
12
13// The NEON tier reinterprets the little-endian payload, so a big-endian
14// aarch64 build falls through to the scalar kernel instead.
15#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
16use super::neon;
17
18#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
19use super::wasm_simd128;
20
21/// Function pointer type for half-precision byte-level distance kernels.
22type HalfFn = fn(&[u8], &[u8], usize) -> f32;
23
24/// Fused BBQ decode-and-distance: centred query bytes, packed sign
25/// bits, and the candidate's residual scale. No intermediate Vec<f32>.
26pub type BbqFn = fn(&[u8], &[u8], f32, usize) -> f32;
27
28/// Selected SIMD runtime — function pointers to the best available kernels.
29pub struct SimdRuntime {
30    pub l2_squared: fn(&[f32], &[f32]) -> f32,
31    pub cosine_distance: fn(&[f32], &[f32]) -> f32,
32    pub neg_inner_product: fn(&[f32], &[f32]) -> f32,
33    pub hamming: fn(&[u8], &[u8]) -> u32,
34    pub name: &'static str,
35    /// F16 fused decode-and-compute kernels (no intermediate Vec<f32>).
36    pub l2_squared_f16: HalfFn,
37    pub cosine_distance_f16: HalfFn,
38    pub neg_inner_product_f16: HalfFn,
39    /// BF16 fused decode-and-compute kernels.
40    pub l2_squared_bf16: HalfFn,
41    pub cosine_distance_bf16: HalfFn,
42    pub neg_inner_product_bf16: HalfFn,
43    /// Fused BBQ decode-and-distance kernel (see `BbqFn`).
44    pub l2_bbq: BbqFn,
45}
46
47impl SimdRuntime {
48    /// Detect CPU features and select the best kernels.
49    pub fn detect() -> Self {
50        #[cfg(target_arch = "x86_64")]
51        {
52            let has_vpopcntdq = std::is_x86_feature_detected!("avx512vpopcntdq");
53
54            if std::is_x86_feature_detected!("avx512f") {
55                let name = if has_vpopcntdq {
56                    "avx512+vpopcntdq"
57                } else {
58                    "avx512"
59                };
60                let rt = Self {
61                    l2_squared: avx512::l2_squared,
62                    cosine_distance: avx512::cosine_distance,
63                    neg_inner_product: avx512::neg_inner_product,
64                    hamming: fast_hamming,
65                    name,
66                    l2_squared_f16: typed_scalar::l2_squared_f16,
67                    cosine_distance_f16: typed_scalar::cosine_f16,
68                    neg_inner_product_f16: typed_scalar::neg_inner_product_f16,
69                    l2_squared_bf16: typed_scalar::l2_squared_bf16,
70                    cosine_distance_bf16: typed_scalar::cosine_bf16,
71                    neg_inner_product_bf16: typed_scalar::neg_inner_product_bf16,
72                    l2_bbq: avx512::l2_bbq,
73                };
74                tracing::info!(kernel = rt.name, "vector SIMD kernel selected");
75                debug_assert!(
76                    !has_vpopcntdq || rt.name == "avx512+vpopcntdq",
77                    "AVX-512 VPOPCNTDQ available but kernel did not select it"
78                );
79                return rt;
80            }
81            if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
82                let rt = Self {
83                    l2_squared: avx2::l2_squared,
84                    cosine_distance: avx2::cosine_distance,
85                    neg_inner_product: avx2::neg_inner_product,
86                    hamming: fast_hamming,
87                    name: "avx2+fma",
88                    l2_squared_f16: typed_scalar::l2_squared_f16,
89                    cosine_distance_f16: typed_scalar::cosine_f16,
90                    neg_inner_product_f16: typed_scalar::neg_inner_product_f16,
91                    l2_squared_bf16: typed_scalar::l2_squared_bf16,
92                    cosine_distance_bf16: typed_scalar::cosine_bf16,
93                    neg_inner_product_bf16: typed_scalar::neg_inner_product_bf16,
94                    l2_bbq: avx2::l2_bbq,
95                };
96                tracing::info!(kernel = rt.name, "vector SIMD kernel selected");
97                return rt;
98            }
99        }
100        #[cfg(all(target_arch = "aarch64", target_endian = "little"))]
101        {
102            let rt = Self {
103                l2_squared: neon::l2_squared,
104                cosine_distance: neon::cosine_distance,
105                neg_inner_product: neon::neg_inner_product,
106                hamming: fast_hamming,
107                name: "neon",
108                l2_squared_f16: typed_scalar::l2_squared_f16,
109                cosine_distance_f16: typed_scalar::cosine_f16,
110                neg_inner_product_f16: typed_scalar::neg_inner_product_f16,
111                l2_squared_bf16: typed_scalar::l2_squared_bf16,
112                cosine_distance_bf16: typed_scalar::cosine_bf16,
113                neg_inner_product_bf16: typed_scalar::neg_inner_product_bf16,
114                l2_bbq: neon::l2_bbq,
115            };
116            tracing::info!(kernel = rt.name, "vector SIMD kernel selected");
117            return rt;
118        }
119        // The tier module is itself compile-time gated on `target_feature =
120        // "simd128"`, so this arm's guard must match that gate exactly: no
121        // runtime probe exists for wasm features.
122        #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
123        {
124            let rt = Self {
125                l2_squared: wasm_simd128::l2_squared,
126                cosine_distance: wasm_simd128::cosine_distance,
127                neg_inner_product: wasm_simd128::neg_inner_product,
128                hamming: fast_hamming,
129                name: "wasm-simd128",
130                l2_squared_f16: typed_scalar::l2_squared_f16,
131                cosine_distance_f16: typed_scalar::cosine_f16,
132                neg_inner_product_f16: typed_scalar::neg_inner_product_f16,
133                l2_squared_bf16: typed_scalar::l2_squared_bf16,
134                cosine_distance_bf16: typed_scalar::cosine_bf16,
135                neg_inner_product_bf16: typed_scalar::neg_inner_product_bf16,
136                l2_bbq: wasm_simd128::l2_bbq,
137            };
138            tracing::info!(kernel = rt.name, "vector SIMD kernel selected");
139            return rt;
140        }
141        #[allow(unreachable_code)]
142        {
143            let rt = Self {
144                l2_squared: scalar_l2,
145                cosine_distance: scalar_cosine,
146                neg_inner_product: scalar_ip,
147                hamming: fast_hamming,
148                name: "scalar",
149                l2_squared_f16: typed_scalar::l2_squared_f16,
150                cosine_distance_f16: typed_scalar::cosine_f16,
151                neg_inner_product_f16: typed_scalar::neg_inner_product_f16,
152                l2_squared_bf16: typed_scalar::l2_squared_bf16,
153                cosine_distance_bf16: typed_scalar::cosine_bf16,
154                neg_inner_product_bf16: typed_scalar::neg_inner_product_bf16,
155                l2_bbq: bbq::l2_bbq,
156            };
157            tracing::info!(kernel = rt.name, "vector SIMD kernel selected");
158            rt
159        }
160    }
161}
162
163/// Global SIMD runtime — initialized once, used everywhere.
164static RUNTIME: std::sync::OnceLock<SimdRuntime> = std::sync::OnceLock::new();
165
166/// Get the global SIMD runtime (auto-detects on first call).
167pub fn runtime() -> &'static SimdRuntime {
168    RUNTIME.get_or_init(SimdRuntime::detect)
169}
170
171#[cfg(test)]
172mod tests {
173    use super::super::hamming::fast_hamming;
174    use super::super::scalar::{scalar_cosine, scalar_ip, scalar_l2};
175    use super::*;
176
177    #[test]
178    fn runtime_detects_features() {
179        let rt = SimdRuntime::detect();
180        assert!(!rt.name.is_empty());
181        tracing::info!("SIMD runtime: {}", rt.name);
182    }
183
184    #[test]
185    fn l2_matches_scalar() {
186        let rt = runtime();
187        let a: Vec<f32> = (0..768).map(|i| (i as f32) * 0.01).collect();
188        let b: Vec<f32> = (0..768).map(|i| (i as f32) * 0.01 + 0.001).collect();
189
190        let simd_result = (rt.l2_squared)(&a, &b);
191        let scalar_result = scalar_l2(&a, &b);
192        assert!(
193            (simd_result - scalar_result).abs() < 0.01,
194            "simd={simd_result}, scalar={scalar_result}"
195        );
196    }
197
198    #[test]
199    fn cosine_matches_scalar() {
200        let rt = runtime();
201        let a: Vec<f32> = (0..768).map(|i| (i as f32).sin()).collect();
202        let b: Vec<f32> = (0..768).map(|i| (i as f32).cos()).collect();
203
204        let simd_result = (rt.cosine_distance)(&a, &b);
205        let scalar_result = scalar_cosine(&a, &b);
206        assert!(
207            (simd_result - scalar_result).abs() < 0.001,
208            "simd={simd_result}, scalar={scalar_result}"
209        );
210    }
211
212    #[test]
213    fn ip_matches_scalar() {
214        let rt = runtime();
215        let a: Vec<f32> = (0..128).map(|i| (i as f32) * 0.1).collect();
216        let b: Vec<f32> = (0..128).map(|i| (i as f32) * 0.2).collect();
217
218        let simd_result = (rt.neg_inner_product)(&a, &b);
219        let scalar_result = scalar_ip(&a, &b);
220        assert!(
221            (simd_result - scalar_result).abs() < 0.1,
222            "simd={simd_result}, scalar={scalar_result}"
223        );
224    }
225
226    #[test]
227    fn hamming_matches() {
228        let a = vec![0b10101010u8; 16];
229        let b = vec![0b01010101u8; 16];
230        assert_eq!(fast_hamming(&a, &b), 128);
231    }
232
233    #[test]
234    fn small_vectors() {
235        let rt = runtime();
236        let a = [1.0f32, 2.0, 3.0];
237        let b = [4.0f32, 5.0, 6.0];
238        let l2 = (rt.l2_squared)(&a, &b);
239        assert!((l2 - 27.0).abs() < 0.01);
240    }
241}
nodedb_vector/distance/simd/runtime.rs

nodedb_vector/distance/simd/
runtime.rs
