rssn_advanced/simd/
kernel.rs

1//! Runtime-dispatched SIMD kernel abstraction.
2//!
3//! [`SimdKernel`] is a trait that abstracts batch floating-point arithmetic.
4//! Concrete implementations cover the scalar fallback, AVX2, and NEON paths.
5//!
6//! [`global_kernel`] selects the best available kernel once per process via
7//! `OnceLock` — subsequent calls are a single pointer load.
8
9use std::sync::OnceLock;
10
11// =========================================================================
12// Trait
13// =========================================================================
14
15/// Batch arithmetic operations, abstracted over SIMD width and ISA.
16///
17/// Implementations must be `Send + Sync` so the global singleton is safe
18/// across thread boundaries. Each method operates element-wise over equal-
19/// length slices; callers guarantee slice lengths match.
20pub trait SimdKernel: Send + Sync {
21    /// Element-wise addition: `out[i] = a[i] + b[i]`.
22    fn batch_add(&self, a: &[f64], b: &[f64], out: &mut [f64]);
23    /// Element-wise subtraction: `out[i] = a[i] - b[i]`.
24    fn batch_sub(&self, a: &[f64], b: &[f64], out: &mut [f64]);
25    /// Element-wise multiplication: `out[i] = a[i] * b[i]`.
26    fn batch_mul(&self, a: &[f64], b: &[f64], out: &mut [f64]);
27    /// Element-wise division: `out[i] = a[i] / b[i]`.
28    fn batch_div(&self, a: &[f64], b: &[f64], out: &mut [f64]);
29    /// Fused multiply-add: `out[i] = a[i] * b[i] + c[i]`.
30    fn batch_fma(&self, a: &[f64], b: &[f64], c: &[f64], out: &mut [f64]);
31    /// Element-wise square root: `out[i] = sqrt(a[i])`.
32    fn batch_sqrt(&self, a: &[f64], out: &mut [f64]);
33    /// Element-wise negation: `out[i] = -a[i]`.
34    fn batch_neg(&self, a: &[f64], out: &mut [f64]);
35    /// Element-wise absolute value: `out[i] = |a[i]|`.
36    fn batch_abs(&self, a: &[f64], out: &mut [f64]);
37    /// Short name for logging/introspection.
38    fn name(&self) -> &'static str;
39}
40
41// =========================================================================
42// ScalarKernel — pure Rust fallback, always available
43// =========================================================================
44
45/// Pure-scalar fallback kernel. Works on every platform.
46pub struct ScalarKernel;
47
48impl SimdKernel for ScalarKernel {
49    fn batch_add(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
50        for i in 0..out.len() {
51            out[i] = a[i] + b[i];
52        }
53    }
54
55    fn batch_sub(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
56        for i in 0..out.len() {
57            out[i] = a[i] - b[i];
58        }
59    }
60
61    fn batch_mul(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
62        for i in 0..out.len() {
63            out[i] = a[i] * b[i];
64        }
65    }
66
67    fn batch_div(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
68        for i in 0..out.len() {
69            out[i] = a[i] / b[i];
70        }
71    }
72
73    fn batch_fma(&self, a: &[f64], b: &[f64], c: &[f64], out: &mut [f64]) {
74        for i in 0..out.len() {
75            out[i] = a[i].mul_add(b[i], c[i]);
76        }
77    }
78
79    fn batch_sqrt(&self, a: &[f64], out: &mut [f64]) {
80        for i in 0..out.len() {
81            out[i] = a[i].sqrt();
82        }
83    }
84
85    fn batch_neg(&self, a: &[f64], out: &mut [f64]) {
86        for i in 0..out.len() {
87            out[i] = -a[i];
88        }
89    }
90
91    fn batch_abs(&self, a: &[f64], out: &mut [f64]) {
92        for i in 0..out.len() {
93            out[i] = a[i].abs();
94        }
95    }
96
97    fn name(&self) -> &'static str {
98        "scalar"
99    }
100}
101
102// =========================================================================
103// Avx2Kernel — x86_64 AVX2 + FMA path
104// =========================================================================
105
106/// AVX2/FMA-accelerated kernel.
107///
108/// Guarded at construction time by a runtime `has_avx2()` check;
109/// never instantiated on non-x86_64 targets.
110#[cfg(target_arch = "x86_64")]
111pub struct Avx2Kernel;
112
113#[cfg(target_arch = "x86_64")]
114impl SimdKernel for Avx2Kernel {
115    fn batch_add(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
116        use crate::asm_presets::add_f64x4;
117        const LANES: usize = 4;
118        let n = out.len();
119        let mut i = 0;
120        while i + LANES <= n {
121            add_f64x4::apply(&a[i..i + LANES], &b[i..i + LANES], &mut out[i..i + LANES]);
122            i += LANES;
123        }
124        while i < n {
125            out[i] = a[i] + b[i];
126            i += 1;
127        }
128    }
129
130    fn batch_sub(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
131        use crate::asm_presets::sub_f64x4;
132        const LANES: usize = 4;
133        let n = out.len();
134        let mut i = 0;
135        while i + LANES <= n {
136            sub_f64x4::apply(&a[i..i + LANES], &b[i..i + LANES], &mut out[i..i + LANES]);
137            i += LANES;
138        }
139        while i < n {
140            out[i] = a[i] - b[i];
141            i += 1;
142        }
143    }
144
145    fn batch_mul(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
146        use crate::asm_presets::mul_f64x4;
147        const LANES: usize = 4;
148        let n = out.len();
149        let mut i = 0;
150        while i + LANES <= n {
151            mul_f64x4::apply(&a[i..i + LANES], &b[i..i + LANES], &mut out[i..i + LANES]);
152            i += LANES;
153        }
154        while i < n {
155            out[i] = a[i] * b[i];
156            i += 1;
157        }
158    }
159
160    fn batch_div(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
161        use crate::asm_presets::div_f64x4;
162        const LANES: usize = 4;
163        let n = out.len();
164        let mut i = 0;
165        while i + LANES <= n {
166            div_f64x4::apply(&a[i..i + LANES], &b[i..i + LANES], &mut out[i..i + LANES]);
167            i += LANES;
168        }
169        while i < n {
170            out[i] = a[i] / b[i];
171            i += 1;
172        }
173    }
174
175    fn batch_fma(&self, a: &[f64], b: &[f64], c: &[f64], out: &mut [f64]) {
176        use crate::asm_presets::fma_f64x4;
177        const LANES: usize = 4;
178        let len = out.len();
179        let mut idx = 0;
180        while idx + LANES <= len {
181            fma_f64x4::apply(
182                &a[idx..idx + LANES],
183                &b[idx..idx + LANES],
184                &c[idx..idx + LANES],
185                &mut out[idx..idx + LANES],
186            );
187            idx += LANES;
188        }
189        while idx < len {
190            out[idx] = a[idx].mul_add(b[idx], c[idx]);
191            idx += 1;
192        }
193    }
194
195    fn batch_sqrt(&self, a: &[f64], out: &mut [f64]) {
196        use crate::asm_presets::sqrt_f64x4;
197        const LANES: usize = 4;
198        let n = out.len();
199        let mut i = 0;
200        while i + LANES <= n {
201            sqrt_f64x4::apply(&a[i..i + LANES], &mut out[i..i + LANES]);
202            i += LANES;
203        }
204        while i < n {
205            out[i] = a[i].sqrt();
206            i += 1;
207        }
208    }
209
210    fn batch_neg(&self, a: &[f64], out: &mut [f64]) {
211        use crate::asm_presets::neg_f64x4;
212        const LANES: usize = 4;
213        let n = out.len();
214        let mut i = 0;
215        while i + LANES <= n {
216            neg_f64x4::apply(&a[i..i + LANES], &mut out[i..i + LANES]);
217            i += LANES;
218        }
219        while i < n {
220            out[i] = -a[i];
221            i += 1;
222        }
223    }
224
225    fn batch_abs(&self, a: &[f64], out: &mut [f64]) {
226        use crate::asm_presets::abs_f64x4;
227        const LANES: usize = 4;
228        let n = out.len();
229        let mut i = 0;
230        while i + LANES <= n {
231            abs_f64x4::apply(&a[i..i + LANES], &mut out[i..i + LANES]);
232            i += LANES;
233        }
234        while i < n {
235            out[i] = a[i].abs();
236            i += 1;
237        }
238    }
239
240    fn name(&self) -> &'static str {
241        "avx2"
242    }
243}
244
245// =========================================================================
246// NeonKernel — aarch64 NEON path
247// =========================================================================
248
249/// NEON-accelerated kernel for AArch64.
250#[cfg(target_arch = "aarch64")]
251pub struct NeonKernel;
252
253#[cfg(target_arch = "aarch64")]
254impl SimdKernel for NeonKernel {
255    fn batch_add(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
256        use crate::asm_presets::add_f64x2;
257        const LANES: usize = 2;
258        let n = out.len();
259        let mut i = 0;
260        while i + LANES <= n {
261            let a2: &[f64; 2] = (&a[i..i + LANES]).try_into().unwrap();
262            let b2: &[f64; 2] = (&b[i..i + LANES]).try_into().unwrap();
263            let o2: &mut [f64; 2] = (&mut out[i..i + LANES]).try_into().unwrap();
264            add_f64x2::apply(a2, b2, o2);
265            i += LANES;
266        }
267        while i < n {
268            out[i] = a[i] + b[i];
269            i += 1;
270        }
271    }
272
273    fn batch_sub(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
274        use crate::asm_presets::sub_f64x2;
275        const LANES: usize = 2;
276        let n = out.len();
277        let mut i = 0;
278        while i + LANES <= n {
279            let a2: &[f64; 2] = (&a[i..i + LANES]).try_into().unwrap();
280            let b2: &[f64; 2] = (&b[i..i + LANES]).try_into().unwrap();
281            let o2: &mut [f64; 2] = (&mut out[i..i + LANES]).try_into().unwrap();
282            sub_f64x2::apply(a2, b2, o2);
283            i += LANES;
284        }
285        while i < n {
286            out[i] = a[i] - b[i];
287            i += 1;
288        }
289    }
290
291    fn batch_mul(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
292        use crate::asm_presets::mul_f64x2;
293        const LANES: usize = 2;
294        let n = out.len();
295        let mut i = 0;
296        while i + LANES <= n {
297            let a2: &[f64; 2] = (&a[i..i + LANES]).try_into().unwrap();
298            let b2: &[f64; 2] = (&b[i..i + LANES]).try_into().unwrap();
299            let o2: &mut [f64; 2] = (&mut out[i..i + LANES]).try_into().unwrap();
300            mul_f64x2::apply(a2, b2, o2);
301            i += LANES;
302        }
303        while i < n {
304            out[i] = a[i] * b[i];
305            i += 1;
306        }
307    }
308
309    fn batch_div(&self, a: &[f64], b: &[f64], out: &mut [f64]) {
310        use crate::asm_presets::div_f64x2;
311        const LANES: usize = 2;
312        let n = out.len();
313        let mut i = 0;
314        while i + LANES <= n {
315            let a2: &[f64; 2] = (&a[i..i + LANES]).try_into().unwrap();
316            let b2: &[f64; 2] = (&b[i..i + LANES]).try_into().unwrap();
317            let o2: &mut [f64; 2] = (&mut out[i..i + LANES]).try_into().unwrap();
318            div_f64x2::apply(a2, b2, o2);
319            i += LANES;
320        }
321        while i < n {
322            out[i] = a[i] / b[i];
323            i += 1;
324        }
325    }
326
327    fn batch_fma(&self, a: &[f64], b: &[f64], c: &[f64], out: &mut [f64]) {
328        // NEON FMA falls back to scalar mul_add (vfma intrinsic not exposed for f64x2 in stable).
329        for i in 0..out.len() {
330            out[i] = a[i].mul_add(b[i], c[i]);
331        }
332    }
333
334    fn batch_sqrt(&self, a: &[f64], out: &mut [f64]) {
335        use crate::asm_presets::sqrt_f64x2;
336        const LANES: usize = 2;
337        let n = out.len();
338        let mut i = 0;
339        while i + LANES <= n {
340            let a2: &[f64; 2] = (&a[i..i + LANES]).try_into().unwrap();
341            let o2: &mut [f64; 2] = (&mut out[i..i + LANES]).try_into().unwrap();
342            sqrt_f64x2::apply(a2, o2);
343            i += LANES;
344        }
345        while i < n {
346            out[i] = a[i].sqrt();
347            i += 1;
348        }
349    }
350
351    fn batch_neg(&self, a: &[f64], out: &mut [f64]) {
352        use crate::asm_presets::neg_f64x2;
353        const LANES: usize = 2;
354        let n = out.len();
355        let mut i = 0;
356        while i + LANES <= n {
357            let a2: &[f64; 2] = (&a[i..i + LANES]).try_into().unwrap();
358            let o2: &mut [f64; 2] = (&mut out[i..i + LANES]).try_into().unwrap();
359            neg_f64x2::apply(a2, o2);
360            i += LANES;
361        }
362        while i < n {
363            out[i] = -a[i];
364            i += 1;
365        }
366    }
367
368    fn batch_abs(&self, a: &[f64], out: &mut [f64]) {
369        use crate::asm_presets::abs_f64x2;
370        const LANES: usize = 2;
371        let n = out.len();
372        let mut i = 0;
373        while i + LANES <= n {
374            let a2: &[f64; 2] = (&a[i..i + LANES]).try_into().unwrap();
375            let o2: &mut [f64; 2] = (&mut out[i..i + LANES]).try_into().unwrap();
376            abs_f64x2::apply(a2, o2);
377            i += LANES;
378        }
379        while i < n {
380            out[i] = a[i].abs();
381            i += 1;
382        }
383    }
384
385    fn name(&self) -> &'static str {
386        "neon"
387    }
388}
389
390// =========================================================================
391// Global kernel selection
392// =========================================================================
393
394static GLOBAL_KERNEL: OnceLock<Box<dyn SimdKernel>> = OnceLock::new();
395
396/// Returns the best available [`SimdKernel`] for this process.
397///
398/// The first call probes CPU feature flags and selects the optimal
399/// implementation; all subsequent calls return the cached pointer with a
400/// single Acquire load.
401///
402/// Priority: AVX2 (`x86_64`) > NEON (aarch64) > scalar.
403pub fn global_kernel() -> &'static dyn SimdKernel {
404    GLOBAL_KERNEL.get_or_init(select_best_kernel).as_ref()
405}
406
407fn select_best_kernel() -> Box<dyn SimdKernel> {
408    #[cfg(target_arch = "x86_64")]
409    if crate::simd::detect::has_avx2() {
410        return Box::new(Avx2Kernel);
411    }
412
413    #[cfg(target_arch = "aarch64")]
414    if crate::simd::detect::has_neon() {
415        return Box::new(NeonKernel);
416    }
417
418    Box::new(ScalarKernel)
419}
420
421#[cfg(test)]
422mod tests {
423    use super::*;
424
425    #[test]
426    fn scalar_kernel_add() {
427        let a = [1.0_f64, 2.0, 3.0, 4.0];
428        let b = [10.0_f64; 4];
429        let mut out = [0.0_f64; 4];
430        ScalarKernel.batch_add(&a, &b, &mut out);
431        assert_eq!(out, [11.0, 12.0, 13.0, 14.0]);
432    }
433
434    #[test]
435    fn scalar_kernel_fma() {
436        let a = [2.0_f64, 3.0];
437        let b = [4.0_f64, 5.0];
438        let c = [1.0_f64, 1.0];
439        let mut out = [0.0_f64; 2];
440        ScalarKernel.batch_fma(&a, &b, &c, &mut out);
441        assert_eq!(out, [9.0, 16.0]);
442    }
443
444    #[test]
445    fn global_kernel_returns_valid_name() {
446        let name = global_kernel().name();
447        assert!(!name.is_empty());
448    }
449
450    #[test]
451    fn global_kernel_add_matches_scalar() {
452        let a: Vec<f64> = (0..16).map(f64::from).collect();
453        let b = vec![1.0_f64; 16];
454        let mut out_global = vec![0.0_f64; 16];
455        let mut out_scalar = vec![0.0_f64; 16];
456        global_kernel().batch_add(&a, &b, &mut out_global);
457        ScalarKernel.batch_add(&a, &b, &mut out_scalar);
458        for i in 0..16 {
459            assert!(
460                (out_global[i] - out_scalar[i]).abs() < 1e-12,
461                "mismatch at {i}"
462            );
463        }
464    }
465}
rssn_advanced/simd/kernel.rs

rssn_advanced/simd/
kernel.rs
