Skip to main content

nodedb_query/simd_filter/
runtime.rs

1// SPDX-License-Identifier: Apache-2.0
2
3// ---------------------------------------------------------------------------
4// Runtime dispatch
5// ---------------------------------------------------------------------------
6
7use super::scalar::{
8    CmpOp, scalar_cmp_f64, scalar_cmp_i64, scalar_eq_u32, scalar_ne_u32, scalar_range_i64,
9};
10
11/// SIMD runtime for filter-to-bitmask operations.
12pub struct FilterSimdRuntime {
13    /// `values[i] == target` → bit i set.
14    pub eq_u32: fn(&[u32], u32) -> Vec<u64>,
15    /// `values[i] != target` → bit i set.
16    pub ne_u32: fn(&[u32], u32) -> Vec<u64>,
17    /// `values[i] > threshold` → bit i set.
18    pub gt_f64: fn(&[f64], f64) -> Vec<u64>,
19    /// `values[i] >= threshold` → bit i set.
20    pub gte_f64: fn(&[f64], f64) -> Vec<u64>,
21    /// `values[i] < threshold` → bit i set.
22    pub lt_f64: fn(&[f64], f64) -> Vec<u64>,
23    /// `values[i] <= threshold` → bit i set.
24    pub lte_f64: fn(&[f64], f64) -> Vec<u64>,
25    /// `values[i] > threshold` → bit i set (i64).
26    pub gt_i64: fn(&[i64], i64) -> Vec<u64>,
27    /// `values[i] >= threshold` → bit i set (i64).
28    pub gte_i64: fn(&[i64], i64) -> Vec<u64>,
29    /// `values[i] < threshold` → bit i set (i64).
30    pub lt_i64: fn(&[i64], i64) -> Vec<u64>,
31    /// `values[i] <= threshold` → bit i set (i64).
32    pub lte_i64: fn(&[i64], i64) -> Vec<u64>,
33    /// `min <= values[i] <= max` → bit i set.
34    pub range_i64: fn(&[i64], i64, i64) -> Vec<u64>,
35    pub name: &'static str,
36}
37
38impl FilterSimdRuntime {
39    pub fn detect() -> Self {
40        #[cfg(target_arch = "x86_64")]
41        {
42            if std::is_x86_feature_detected!("avx512f") {
43                return Self {
44                    eq_u32: super::avx512::avx512_eq_u32,
45                    ne_u32: super::avx512::avx512_ne_u32,
46                    gt_f64: super::avx512::avx512_gt_f64,
47                    gte_f64: super::avx512::avx512_gte_f64,
48                    lt_f64: super::avx512::avx512_lt_f64,
49                    lte_f64: super::avx512::avx512_lte_f64,
50                    gt_i64: super::avx512::avx512_gt_i64,
51                    gte_i64: super::avx512::avx512_gte_i64,
52                    lt_i64: super::avx512::avx512_lt_i64,
53                    lte_i64: super::avx512::avx512_lte_i64,
54                    range_i64: super::avx512::avx512_range_i64,
55                    name: "avx512",
56                };
57            }
58            if std::is_x86_feature_detected!("avx2") {
59                return Self {
60                    eq_u32: super::avx2::avx2_eq_u32,
61                    ne_u32: super::avx2::avx2_ne_u32,
62                    gt_f64: super::avx2::avx2_gt_f64,
63                    gte_f64: super::avx2::avx2_gte_f64,
64                    lt_f64: super::avx2::avx2_lt_f64,
65                    lte_f64: super::avx2::avx2_lte_f64,
66                    gt_i64: super::avx2::avx2_gt_i64,
67                    gte_i64: super::avx2::avx2_gte_i64,
68                    lt_i64: super::avx2::avx2_lt_i64,
69                    lte_i64: super::avx2::avx2_lte_i64,
70                    range_i64: super::avx2::avx2_range_i64,
71                    name: "avx2",
72                };
73            }
74        }
75        #[cfg(target_arch = "aarch64")]
76        {
77            return Self {
78                eq_u32: super::neon::neon_eq_u32,
79                ne_u32: super::neon::neon_ne_u32,
80                gt_f64: super::neon::neon_gt_f64,
81                gte_f64: super::neon::neon_gte_f64,
82                lt_f64: super::neon::neon_lt_f64,
83                lte_f64: super::neon::neon_lte_f64,
84                gt_i64: super::neon::neon_gt_i64,
85                gte_i64: super::neon::neon_gte_i64,
86                lt_i64: super::neon::neon_lt_i64,
87                lte_i64: super::neon::neon_lte_i64,
88                range_i64: super::neon::neon_range_i64,
89                name: "neon",
90            };
91        }
92        #[cfg(target_arch = "wasm32")]
93        {
94            return Self {
95                eq_u32: super::wasm::wasm_eq_u32,
96                ne_u32: super::wasm::wasm_ne_u32,
97                gt_f64: |v, t| scalar_cmp_f64(v, t, CmpOp::Gt),
98                gte_f64: |v, t| scalar_cmp_f64(v, t, CmpOp::Gte),
99                lt_f64: |v, t| scalar_cmp_f64(v, t, CmpOp::Lt),
100                lte_f64: |v, t| scalar_cmp_f64(v, t, CmpOp::Lte),
101                gt_i64: |v, t| scalar_cmp_i64(v, t, CmpOp::Gt),
102                gte_i64: |v, t| scalar_cmp_i64(v, t, CmpOp::Gte),
103                lt_i64: |v, t| scalar_cmp_i64(v, t, CmpOp::Lt),
104                lte_i64: |v, t| scalar_cmp_i64(v, t, CmpOp::Lte),
105                range_i64: scalar_range_i64,
106                name: "wasm-simd128",
107            };
108        }
109        #[allow(unreachable_code)]
110        Self {
111            eq_u32: scalar_eq_u32,
112            ne_u32: scalar_ne_u32,
113            gt_f64: |v, t| scalar_cmp_f64(v, t, CmpOp::Gt),
114            gte_f64: |v, t| scalar_cmp_f64(v, t, CmpOp::Gte),
115            lt_f64: |v, t| scalar_cmp_f64(v, t, CmpOp::Lt),
116            lte_f64: |v, t| scalar_cmp_f64(v, t, CmpOp::Lte),
117            gt_i64: |v, t| scalar_cmp_i64(v, t, CmpOp::Gt),
118            gte_i64: |v, t| scalar_cmp_i64(v, t, CmpOp::Gte),
119            lt_i64: |v, t| scalar_cmp_i64(v, t, CmpOp::Lt),
120            lte_i64: |v, t| scalar_cmp_i64(v, t, CmpOp::Lte),
121            range_i64: scalar_range_i64,
122            name: "scalar",
123        }
124    }
125}
126
127static FILTER_RUNTIME: std::sync::OnceLock<FilterSimdRuntime> = std::sync::OnceLock::new();
128
129/// Get the global filter SIMD runtime.
130pub fn filter_runtime() -> &'static FilterSimdRuntime {
131    FILTER_RUNTIME.get_or_init(FilterSimdRuntime::detect)
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use crate::simd_filter::bitmask::{bitmask_to_indices, popcount};
138
139    #[test]
140    fn runtime_detects() {
141        let rt = filter_runtime();
142        assert!(!rt.name.is_empty());
143    }
144
145    #[test]
146    fn eq_u32_basic() {
147        let rt = filter_runtime();
148        let values: Vec<u32> = (0..100).collect();
149        let mask = (rt.eq_u32)(&values, 42);
150        assert_eq!(popcount(&mask), 1);
151        let indices = bitmask_to_indices(&mask);
152        assert_eq!(indices, vec![42]);
153    }
154
155    #[test]
156    fn ne_u32_basic() {
157        let rt = filter_runtime();
158        let values: Vec<u32> = (0..100).collect();
159        let mask = (rt.ne_u32)(&values, 42);
160        assert_eq!(popcount(&mask), 99);
161    }
162
163    #[test]
164    fn eq_u32_repeated() {
165        let rt = filter_runtime();
166        // 1000 values cycling through 0..8.
167        let values: Vec<u32> = (0..1000).map(|i| (i % 8) as u32).collect();
168        let mask = (rt.eq_u32)(&values, 3);
169        assert_eq!(popcount(&mask), 125); // 1000/8
170        let indices = bitmask_to_indices(&mask);
171        assert!(indices.iter().all(|&i| values[i as usize] == 3));
172    }
173
174    #[test]
175    fn gt_f64_basic() {
176        let rt = filter_runtime();
177        let values: Vec<f64> = (0..1000).map(|i| i as f64).collect();
178        let mask = (rt.gt_f64)(&values, 500.0);
179        assert_eq!(popcount(&mask), 499); // 501..999
180    }
181
182    #[test]
183    fn gte_f64_basic() {
184        let rt = filter_runtime();
185        let values: Vec<f64> = (0..1000).map(|i| i as f64).collect();
186        let mask = (rt.gte_f64)(&values, 500.0);
187        assert_eq!(popcount(&mask), 500); // 500..999
188    }
189
190    #[test]
191    fn range_i64_basic() {
192        let rt = filter_runtime();
193        let values: Vec<i64> = (0..1000).collect();
194        let mask = (rt.range_i64)(&values, 100, 200);
195        assert_eq!(popcount(&mask), 101); // 100..=200
196    }
197
198    #[test]
199    fn empty_input() {
200        let rt = filter_runtime();
201        assert!(popcount(&(rt.eq_u32)(&[], 0)) == 0);
202        assert!(popcount(&(rt.gt_f64)(&[], 0.0)) == 0);
203        assert!(popcount(&(rt.range_i64)(&[], 0, 100)) == 0);
204    }
205
206    #[test]
207    fn large_input_eq_u32() {
208        let rt = filter_runtime();
209        let n: u32 = 10_000;
210        let values: Vec<u32> = (0..n).map(|i| i % 256).collect();
211        let mask = (rt.eq_u32)(&values, 0);
212        // 0, 256, 512, ... → ceil(n/256) occurrences.
213        let expected = values.iter().filter(|&&v| v == 0).count() as u64;
214        assert_eq!(popcount(&mask), expected);
215        let indices = bitmask_to_indices(&mask);
216        assert!(indices.iter().all(|&i| values[i as usize] == 0));
217    }
218
219    #[test]
220    fn i64_comparisons() {
221        let rt = filter_runtime();
222        let values: Vec<i64> = (0..100).collect();
223
224        assert_eq!(popcount(&(rt.gt_i64)(&values, 50)), 49);
225        assert_eq!(popcount(&(rt.gte_i64)(&values, 50)), 50);
226        assert_eq!(popcount(&(rt.lt_i64)(&values, 50)), 50);
227        assert_eq!(popcount(&(rt.lte_i64)(&values, 50)), 51);
228    }
229}