numrs2/simd_optimize/
simd_select.rs

1//! SIMD implementation selection based on CPU features
2//!
3//! This module provides functionality for selecting the most efficient
4//! SIMD implementation based on the detected CPU features.
5
6use crate::simd_optimize::feature_detect::CpuFeatures;
7
8/// Represents the available SIMD implementations
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum SimdImplementation {
11    /// Generic scalar implementation (no SIMD)
12    Scalar,
13    /// SSE implementation (x86_64)
14    SSE,
15    /// AVX implementation (x86_64)
16    AVX,
17    /// AVX2 implementation (x86_64)
18    AVX2,
19    /// AVX-512 implementation (x86_64)
20    AVX512,
21    /// NEON implementation (aarch64)
22    NEON,
23    /// SVE implementation (aarch64)
24    SVE,
25}
26
27impl SimdImplementation {
28    /// Get the name of the SIMD implementation
29    pub fn name(&self) -> &'static str {
30        match self {
31            SimdImplementation::Scalar => "Scalar",
32            SimdImplementation::SSE => "SSE",
33            SimdImplementation::AVX => "AVX",
34            SimdImplementation::AVX2 => "AVX2",
35            SimdImplementation::AVX512 => "AVX512",
36            SimdImplementation::NEON => "NEON",
37            SimdImplementation::SVE => "SVE",
38        }
39    }
40
41    /// Check if the implementation is AVX2 or better
42    pub fn is_avx2_or_better(&self) -> bool {
43        matches!(self, SimdImplementation::AVX2 | SimdImplementation::AVX512)
44    }
45
46    /// Check if the implementation supports FMA operations
47    pub fn supports_fma(&self, features: &CpuFeatures) -> bool {
48        match self {
49            SimdImplementation::AVX2 | SimdImplementation::AVX512 => features.fma,
50            _ => false,
51        }
52    }
53
54    /// Check if the implementation is NEON or better
55    pub fn is_neon_or_better(&self) -> bool {
56        matches!(self, SimdImplementation::NEON | SimdImplementation::SVE)
57    }
58
59    /// Get the vector width in bits for this implementation
60    pub fn vector_width(&self) -> usize {
61        match self {
62            SimdImplementation::Scalar => 0,
63            SimdImplementation::SSE => 128,
64            SimdImplementation::AVX => 256,
65            SimdImplementation::AVX2 => 256,
66            SimdImplementation::AVX512 => 512,
67            SimdImplementation::NEON => 128,
68            SimdImplementation::SVE => 128, // Simplified - SVE can be variable
69        }
70    }
71}
72
73/// Select the most efficient SIMD implementation based on CPU features
74///
75/// # Arguments
76///
77/// * `features` - The detected CPU features
78///
79/// # Returns
80///
81/// The selected SIMD implementation
82pub fn select_simd_implementation(features: &CpuFeatures) -> SimdImplementation {
83    // Check for x86_64 features in order of preference
84    if features.avx512f {
85        return SimdImplementation::AVX512;
86    }
87
88    if features.avx2 {
89        return SimdImplementation::AVX2;
90    }
91
92    if features.avx {
93        return SimdImplementation::AVX;
94    }
95
96    if features.sse2 {
97        return SimdImplementation::SSE;
98    }
99
100    // Check for aarch64 features
101    if features.sve {
102        return SimdImplementation::SVE;
103    }
104
105    if features.neon {
106        return SimdImplementation::NEON;
107    }
108
109    // Fall back to scalar implementation
110    SimdImplementation::Scalar
111}
112
113/// Apply a specific SIMD strategy based on CPU features
114///
115/// This is a helper function to facilitate selecting different implementations
116/// based on the available CPU features.
117///
118/// # Arguments
119///
120/// * `features` - The detected CPU features
121/// * `scalar` - The scalar implementation to use if SIMD is not available
122/// * `sse` - The SSE implementation to use if SSE is available
123/// * `avx` - The AVX implementation to use if AVX is available
124/// * `avx2` - The AVX2 implementation to use if AVX2 is available
125/// * `avx512` - The AVX-512 implementation to use if AVX-512 is available
126/// * `neon` - The NEON implementation to use if NEON is available
127/// * `sve` - The SVE implementation to use if SVE is available
128///
129/// # Returns
130///
131/// The result of the selected implementation
132#[allow(clippy::too_many_arguments)]
133pub fn apply_simd_strategy<T, S, SSE, AVX, AVX2, AVX512, NEON, SVE>(
134    features: &CpuFeatures,
135    scalar: S,
136    sse: SSE,
137    avx: AVX,
138    avx2: AVX2,
139    avx512: AVX512,
140    neon: NEON,
141    sve: SVE,
142) -> T
143where
144    S: FnOnce() -> T,
145    SSE: FnOnce() -> T,
146    AVX: FnOnce() -> T,
147    AVX2: FnOnce() -> T,
148    AVX512: FnOnce() -> T,
149    NEON: FnOnce() -> T,
150    SVE: FnOnce() -> T,
151{
152    let implementation = select_simd_implementation(features);
153
154    match implementation {
155        SimdImplementation::AVX512 => avx512(),
156        SimdImplementation::AVX2 => avx2(),
157        SimdImplementation::AVX => avx(),
158        SimdImplementation::SSE => sse(),
159        SimdImplementation::NEON => neon(),
160        SimdImplementation::SVE => sve(),
161        SimdImplementation::Scalar => scalar(),
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    #[test]
170    fn test_simd_selection() {
171        // Test with no features
172        let features = CpuFeatures::default();
173        assert_eq!(
174            select_simd_implementation(&features),
175            SimdImplementation::Scalar
176        );
177
178        // Test with SSE2 only
179        let features = CpuFeatures {
180            sse2: true,
181            ..Default::default()
182        };
183        assert_eq!(
184            select_simd_implementation(&features),
185            SimdImplementation::SSE
186        );
187
188        // Test with AVX
189        let features = CpuFeatures {
190            sse2: true,
191            avx: true,
192            ..Default::default()
193        };
194        assert_eq!(
195            select_simd_implementation(&features),
196            SimdImplementation::AVX
197        );
198
199        // Test with AVX2
200        let features = CpuFeatures {
201            sse2: true,
202            avx: true,
203            avx2: true,
204            ..Default::default()
205        };
206        assert_eq!(
207            select_simd_implementation(&features),
208            SimdImplementation::AVX2
209        );
210
211        // Test with all x86_64 features
212        let features = CpuFeatures {
213            sse2: true,
214            avx: true,
215            avx2: true,
216            avx512f: true,
217            ..Default::default()
218        };
219        assert_eq!(
220            select_simd_implementation(&features),
221            SimdImplementation::AVX512
222        );
223
224        // Test with NEON
225        let features = CpuFeatures {
226            neon: true,
227            ..Default::default()
228        };
229        assert_eq!(
230            select_simd_implementation(&features),
231            SimdImplementation::NEON
232        );
233
234        // Test with SVE
235        let features = CpuFeatures {
236            neon: true,
237            sve: true,
238            ..Default::default()
239        };
240        assert_eq!(
241            select_simd_implementation(&features),
242            SimdImplementation::SVE
243        );
244    }
245
246    #[test]
247    fn test_simd_strategy() {
248        // Create a test function that returns a string based on the implementation
249        let apply_test = |features: &CpuFeatures| {
250            apply_simd_strategy(
251                features,
252                || "scalar",
253                || "sse",
254                || "avx",
255                || "avx2",
256                || "avx512",
257                || "neon",
258                || "sve",
259            )
260        };
261
262        // Test with no features
263        let features = CpuFeatures::default();
264        assert_eq!(apply_test(&features), "scalar");
265
266        // Test with SSE2 only
267        let features = CpuFeatures {
268            sse2: true,
269            ..Default::default()
270        };
271        assert_eq!(apply_test(&features), "sse");
272
273        // Test with AVX
274        let features = CpuFeatures {
275            sse2: true,
276            avx: true,
277            ..Default::default()
278        };
279        assert_eq!(apply_test(&features), "avx");
280
281        // Test with AVX2
282        let features = CpuFeatures {
283            sse2: true,
284            avx: true,
285            avx2: true,
286            ..Default::default()
287        };
288        assert_eq!(apply_test(&features), "avx2");
289    }
290}
numrs2/simd_optimize/simd_select.rs

numrs2/simd_optimize/
simd_select.rs
