utils/
multiversion.rs

1//! Implementation of [`multiversion!`](crate::multiversion!) macro.
2
3use std::error::Error;
4use std::fmt::{self, Display, Formatter};
5use std::str::FromStr;
6use std::sync::OnceLock;
7
8/// Macro to generate multiversioned functions.
9///
10/// This macro generates multiple versions of the provided functions, each optimized for different
11/// hardware, as well as dynamic dispatch functions which select the best version to use at runtime
12/// based on feature support.
13///
14/// This allows using instruction set extensions such as
15/// [AVX2](https://en.wikipedia.org/wiki/Advanced_Vector_Extensions#Advanced_Vector_Extensions_2)
16/// if the binary is run on a compatible processor, without enabling it at compile time, which would
17/// create a binary which can't be run on processors without AVX2 support.
18///
19/// The first three rules are supported:
20///
21/// 1. The first rule matches a dynamic dispatch function, optionally followed by extra private
22///    helper functions. Dynamic dispatch is controlled at runtime by supplying a path to a
23///    [`LazyLock`](std::sync::LazyLock) containing a [`Version`] value in the `dyn_dispatch`
24///    attribute.
25///
26/// 2. The second rule matches library functions without a dynamic dispatch function.
27///
28/// 3. The third rule takes a name of a library function generated using the second rule and
29///    expands to a [`LazyLock`](std::sync::LazyLock) evaluating to the fastest version, unless a
30///    global override is set using [`Version::set_override`]. The return value is intended to be
31///    stored in a `static` item for use in a `dyn_dispatch` attribute.
32///
33/// The first two rules start with a use statement which can be used for wildcard imports of
34/// multiversioned libraries. Each path will be expanded to include the version's name - e.g. the
35/// `avx2` version will expand `use {utils::simd::*};` into `use {utils::simd::avx2::*};`. Other
36/// imports should be outside the macro invocation.
37///
38/// Supported function syntax is limited as this is a declarative macro. Dynamic dispatch functions
39/// are the most limited, supporting only basic arguments and optionally a return type.
40///
41/// Versions enabling additional features require the `unsafe` crate feature to be enabled, as
42/// features are enabled using the
43/// [`target_feature`](https://doc.rust-lang.org/reference/attributes/codegen.html#the-target_feature-attribute)
44/// attribute, which requires callers to use `unsafe` blocks. A safe fallback version will always
45/// be generated.
46///
47/// See [`crate::md5`] as an example. [`multiversion_test!`](crate::multiversion_test!) can be used
48/// for testing multiversioned libraries.
49#[macro_export]
50macro_rules! multiversion {
51    // One dynamic dispatch function, optionally with extra helper functions.
52    (
53        use {$($($path:ident::)+*),*};
54
55        #[dyn_dispatch = $dispatch:path]
56        $(#[$m:meta])* $v:vis fn $name:ident($($arg_name:ident: $arg_type:ty),*$(,)?) $(-> $ret:ty)? $body:block
57
58        $($tail:tt)*
59    ) => {
60        /// [`multiversion!`] dynamic dispatch implementations.
61        mod $name {
62            #[allow(clippy::allow_attributes, unused_imports)]
63            use {super::*, $crate::multiversion}; // multiversion import needed for rustdoc links
64
65            $crate::multiversion!{
66                use {$($($path::)+*),*};
67                $(#[$m])* pub fn $name($($arg_name: $arg_type),*) $(-> $ret)? $body
68                $($tail)*
69            }
70        }
71
72        /// [`multiversion!`] dynamic dispatch function.
73        #[inline]
74        $v fn $name($($arg_name: $arg_type),*) $(-> $ret)? {
75            use $crate::multiversion::Version::*;
76
77            match *$dispatch {
78                Scalar => $name::scalar::$name($($arg_name),*),
79                Array128 => $name::array128::$name($($arg_name),*),
80                Array256 => $name::array256::$name($($arg_name),*),
81                #[cfg(feature = "all-simd")]
82                Array4096 => $name::array4096::$name($($arg_name),*),
83                #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
84                AVX2 => unsafe { $name::avx2::$name($($arg_name),*) },
85                #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
86                AVX2x2 => unsafe { $name::avx2x2::$name($($arg_name),*) },
87                #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
88                AVX2x4 => unsafe { $name::avx2x4::$name($($arg_name),*) },
89                #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
90                AVX2x8 => unsafe { $name::avx2x8::$name($($arg_name),*) },
91                #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
92                AVX512 => unsafe { $name::avx512::$name($($arg_name),*) },
93                #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
94                AVX512x2 => unsafe { $name::avx512x2::$name($($arg_name),*) },
95                #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
96                AVX512x4 => unsafe { $name::avx512x4::$name($($arg_name),*) },
97                #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
98                AVX512x8 => unsafe { $name::avx512x8::$name($($arg_name),*) },
99                #[cfg(all(feature = "unsafe", target_arch = "aarch64"))]
100                Neon => unsafe { $name::neon::$name($($arg_name),*) },
101                #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
102                Neonx2 => unsafe { $name::neonx2::$name($($arg_name),*) },
103                #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
104                Neonx4 => unsafe { $name::neonx4::$name($($arg_name),*) },
105                #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
106                Neonx8 => unsafe { $name::neonx8::$name($($arg_name),*) },
107            }
108        }
109    };
110
111    // Library-only definition, without a dynamic dispatch function.
112    (
113        use {$($($path:ident::)+*),*};
114
115        $($tail:tt)+
116    ) => {
117        /// [`multiversion!`] scalar implementation.
118        pub mod scalar {
119            #![allow(
120                clippy::reversed_empty_ranges,
121                clippy::range_plus_one,
122                clippy::modulo_one,
123                clippy::trivially_copy_pass_by_ref
124            )]
125
126            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
127            use {super::*, $($($path::)+scalar::*),*};
128
129            $($tail)*
130        }
131
132        /// [`multiversion!`] array128 implementation.
133        pub mod array128 {
134            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
135            use {super::*, $($($path::)+array128::*),*};
136
137            $($tail)*
138        }
139
140        /// [`multiversion!`] array256 implementation.
141        pub mod array256 {
142            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
143            use {super::*, $($($path::)+array256::*),*};
144
145            $($tail)*
146        }
147
148        /// [`multiversion!`] array4096 implementation.
149        #[cfg(feature="all-simd")]
150        pub mod array4096 {
151            #![allow(clippy::large_types_passed_by_value)]
152
153            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
154            use {super::*, $($($path::)+array4096::*),*};
155
156            $($tail)*
157        }
158
159        /// [`multiversion!`] avx2 implementation.
160        #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
161        pub mod avx2 {
162            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
163            use {super::*, $($($path::)+avx2::*),*};
164
165            $crate::multiversion!{@enable target_feature(enable = "avx2") $($tail)*}
166        }
167
168        /// [`multiversion!`] avx2x2 implementation.
169        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
170        pub mod avx2x2 {
171            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
172            use {super::*, $($($path::)+avx2x2::*),*};
173
174            $crate::multiversion!{@enable target_feature(enable = "avx2") $($tail)*}
175        }
176
177        /// [`multiversion!`] avx2x4 implementation.
178        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
179        pub mod avx2x4 {
180            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
181            use {super::*, $($($path::)+avx2x4::*),*};
182
183            $crate::multiversion!{@enable target_feature(enable = "avx2") $($tail)*}
184        }
185
186        /// [`multiversion!`] avx2x8 implementation.
187        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
188        pub mod avx2x8 {
189            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
190            use {super::*, $($($path::)+avx2x8::*),*};
191
192            $crate::multiversion!{@enable target_feature(enable = "avx2") $($tail)*}
193        }
194
195        /// [`multiversion!`] avx512 implementation.
196        #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
197        pub mod avx512 {
198            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
199            use {super::*, $($($path::)+avx512::*),*};
200
201            $crate::multiversion!{@enable target_feature(enable = "avx512f") $($tail)*}
202        }
203
204        /// [`multiversion!`] avx512x2 implementation.
205        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
206        pub mod avx512x2 {
207            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
208            use {super::*, $($($path::)+avx512x2::*),*};
209
210            $crate::multiversion!{@enable target_feature(enable = "avx512f") $($tail)*}
211        }
212
213        /// [`multiversion!`] avx512x4 implementation.
214        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
215        pub mod avx512x4 {
216            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
217            use {super::*, $($($path::)+avx512x4::*),*};
218
219            $crate::multiversion!{@enable target_feature(enable = "avx512f") $($tail)*}
220        }
221
222        /// [`multiversion!`] avx512x8 implementation.
223        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
224        pub mod avx512x8 {
225            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
226            use {super::*, $($($path::)+avx512x8::*),*};
227
228            $crate::multiversion!{@enable target_feature(enable = "avx512f") $($tail)*}
229        }
230
231        /// [`multiversion!`] neon implementation.
232        #[cfg(all(feature = "unsafe", target_arch = "aarch64"))]
233        pub mod neon {
234            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
235            use {super::*, $($($path::)+neon::*),*};
236
237            $crate::multiversion!{@enable target_feature(enable = "neon") $($tail)*}
238        }
239
240        /// [`multiversion!`] neonx2 implementation.
241        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
242        pub mod neonx2 {
243            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
244            use {super::*, $($($path::)+neonx2::*),*};
245
246            $crate::multiversion!{@enable target_feature(enable = "neon") $($tail)*}
247        }
248
249        /// [`multiversion!`] neonx4 implementation.
250        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
251        pub mod neonx4 {
252            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
253            use {super::*, $($($path::)+neonx4::*),*};
254
255            $crate::multiversion!{@enable target_feature(enable = "neon") $($tail)*}
256        }
257
258        /// [`multiversion!`] neonx8 implementation.
259        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
260        pub mod neonx8 {
261            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
262            use {super::*, $($($path::)+neonx8::*),*};
263
264            $crate::multiversion!{@enable target_feature(enable = "neon") $($tail)*}
265        }
266    };
267
268    // Microbenchmark for dynamic dispatch
269    (fastest($name:ident())) => {
270        ::std::sync::LazyLock::new(#[cfg_attr(target_family = "wasm", expect(unreachable_code))] || {
271            use $crate::multiversion::Version::*;
272
273            // Instant::now() isn't implemented in WebAssembly, so hardcode implementations
274            #[cfg(all(target_family = "wasm", target_feature = "simd128"))]
275            return Array256;
276            #[cfg(all(target_family = "wasm"))]
277            return Array128;
278
279            if let Some(version) = $crate::multiversion::Version::get_override() {
280                return version;
281            }
282
283            $crate::multiversion::VERSIONS
284                .iter()
285                .map(|&x| {
286                    let start = ::std::time::Instant::now();
287                    ::std::hint::black_box(match x {
288                        Scalar => scalar::$name(),
289                        Array128 => array128::$name(),
290                        Array256 => array256::$name(),
291                        #[cfg(feature = "all-simd")]
292                        Array4096 => array4096::$name(),
293                        #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
294                        AVX2 => unsafe { avx2::$name() },
295                        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
296                        AVX2x2 => unsafe { avx2x2::$name() },
297                        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
298                        AVX2x4 => unsafe { avx2x4::$name() },
299                        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
300                        AVX2x8 => unsafe { avx2x8::$name() },
301                        #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
302                        AVX512 => unsafe { avx512::$name() },
303                        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
304                        AVX512x2 => unsafe { avx512x2::$name() },
305                        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
306                        AVX512x4 => unsafe { avx512x4::$name() },
307                        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
308                        AVX512x8 => unsafe { avx512x8::$name() },
309                        #[cfg(all(feature = "unsafe", target_arch = "aarch64"))]
310                        Neon => unsafe { neon::$name() },
311                        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
312                        Neonx2 => unsafe { neonx2::$name() },
313                        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
314                        Neonx4 => unsafe { neonx4::$name() },
315                        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
316                        Neonx8 => unsafe { neonx8::$name() },
317                    });
318                    (start.elapsed(), x)
319                })
320                // .inspect(|x| { dbg!(x); })
321                .min_by_key(|x| x.0)
322                .unwrap()
323                .1
324        })
325    };
326
327    // Helper rules to add target_feature to helper functions
328    (@enable $t:meta) => {};
329    (@enable $t:meta
330        $(#[$m:meta])* $v:vis const $n:ident: $ty:ty = $e:expr; $($tail:tt)*
331    ) => {
332        $(#[$m])* $v const $n: $ty = $e;
333        $crate::multiversion!{@enable $t $($tail)*}
334    };
335    // Replace #[inline(always)] which is incompatible with target_feature with normal #[inline]
336    (@enable $t:meta
337        #[inline(always)] $($tail:tt)*
338    ) => {
339        $crate::multiversion!{@enable $t
340            #[inline] $($tail)*
341        }
342    };
343    (@enable $t:meta
344        $(#[$m:meta])* $v:vis fn $($tail:tt)*
345    ) => {
346        $crate::multiversion!{@one_item $t
347            $(#[$m])*
348            #[$t]
349            #[allow(clippy::allow_attributes, clippy::missing_safety_doc)]
350            $v fn $($tail)*
351        }
352    };
353
354    // Helper rule to pop the first item out of the tt sequence
355    (@one_item $t:meta $item:item $($tail:tt)*) => {
356        $item
357        $crate::multiversion!{@enable $t $($tail)*}
358    };
359}
360
361/// Helper for testing and benchmarking [`multiversion!`] library functions.
362///
363/// The first rule is for testing and creates individual
364/// [`#[test]`](https://doc.rust-lang.org/reference/attributes/testing.html) functions for each
365/// implementation.
366///
367/// The second rule is more general, duplicating the same expression for each implementation and
368/// is useful for benchmarking.
369///
370/// `#[target_feature(...)]` isn't applied to the test code as the feature-specific code should
371/// be elsewhere, inside a [`multiversion!`] macro.
372#[macro_export]
373macro_rules! multiversion_test {
374    (
375        use {$($($path:ident::)+*),*};
376
377        #[test]
378        $(#[$m:meta])* $v:vis fn multiversion() $body:block
379    ) => {
380        #[test]
381        $(#[$m])*
382        fn scalar() {
383            #![allow(
384                clippy::reversed_empty_ranges,
385                clippy::range_plus_one,
386                clippy::modulo_one,
387                clippy::trivially_copy_pass_by_ref
388            )]
389
390            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
391            use {$($($path::)+scalar::*),*};
392
393            $body
394        }
395
396        #[test]
397        $(#[$m])*
398        fn array128() {
399            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
400            use {$($($path::)+array128::*),*};
401
402            $body
403        }
404
405        #[test]
406        $(#[$m])*
407        fn array256() {
408            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
409            use {$($($path::)+array256::*),*};
410
411            $body
412        }
413
414        #[test]
415        #[cfg(feature = "all-simd")]
416        $(#[$m])*
417        fn array4096() {
418            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
419            use {$($($path::)+array4096::*),*};
420
421            $body
422        }
423
424        #[test]
425        #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
426        $(#[$m])*
427        fn avx2() {
428            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
429            use {$($($path::)+avx2::*),*};
430
431            if !$crate::multiversion::Version::AVX2.supported() {
432                use std::io::{stdout, Write};
433                let _ = writeln!(&mut stdout(), "warning: skipping test in {}::avx2 due to missing avx2 support", module_path!());
434                return;
435            }
436
437            unsafe { $body }
438        }
439
440        #[test]
441        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
442        $(#[$m])*
443        fn avx2x2() {
444            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
445            use {$($($path::)+avx2x2::*),*};
446
447            if !$crate::multiversion::Version::AVX2x2.supported() {
448                use std::io::{stdout, Write};
449                let _ = writeln!(&mut stdout(), "warning: skipping test in {}::avx2x2 due to missing avx2 support", module_path!());
450                return;
451            }
452
453            unsafe { $body }
454        }
455
456        #[test]
457        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
458        $(#[$m])*
459        fn avx2x4() {
460            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
461            use {$($($path::)+avx2x4::*),*};
462
463            if !$crate::multiversion::Version::AVX2x4.supported() {
464                use std::io::{stdout, Write};
465                let _ = writeln!(&mut stdout(), "warning: skipping test in {}::avx2x4 due to missing avx2 support", module_path!());
466                return;
467            }
468
469            unsafe { $body }
470        }
471
472        #[test]
473        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
474        $(#[$m])*
475        fn avx2x8() {
476            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
477            use {$($($path::)+avx2x8::*),*};
478
479            if !$crate::multiversion::Version::AVX2x8.supported() {
480                use std::io::{stdout, Write};
481                let _ = writeln!(&mut stdout(), "warning: skipping test in {}::avx2x8 due to missing avx2 support", module_path!());
482                return;
483            }
484
485            unsafe { $body }
486        }
487
488        #[test]
489        #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
490        $(#[$m])*
491        fn avx512() {
492            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
493            use {$($($path::)+avx512::*),*};
494
495            if !$crate::multiversion::Version::AVX512.supported() {
496                use std::io::{stdout, Write};
497                let _ = writeln!(&mut stdout(), "warning: skipping test in {}::avx512 due to missing avx512 support", module_path!());
498                return;
499            }
500
501            unsafe { $body }
502        }
503
504        #[test]
505        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
506        $(#[$m])*
507        fn avx512x2() {
508            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
509            use {$($($path::)+avx512x2::*),*};
510
511            if !$crate::multiversion::Version::AVX512x2.supported() {
512                use std::io::{stdout, Write};
513                let _ = writeln!(&mut stdout(), "warning: skipping test in {}::avx512x2 due to missing avx512 support", module_path!());
514                return;
515            }
516
517            unsafe { $body }
518        }
519
520        #[test]
521        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
522        $(#[$m])*
523        fn avx512x4() {
524            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
525            use {$($($path::)+avx512x4::*),*};
526
527            if !$crate::multiversion::Version::AVX512x4.supported() {
528                use std::io::{stdout, Write};
529                let _ = writeln!(&mut stdout(), "warning: skipping test in {}::avx512x4 due to missing avx512 support", module_path!());
530                return;
531            }
532
533            unsafe { $body }
534        }
535
536        #[test]
537        #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
538        $(#[$m])*
539        fn avx512x8() {
540            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
541            use {$($($path::)+avx512x8::*),*};
542
543            if !$crate::multiversion::Version::AVX512x8.supported() {
544                use std::io::{stdout, Write};
545                let _ = writeln!(&mut stdout(), "warning: skipping test in {}::avx512x8 due to missing avx512 support", module_path!());
546                return;
547            }
548
549            unsafe { $body }
550        }
551
552        #[test]
553        #[cfg(all(feature = "unsafe", target_arch = "aarch64"))]
554        $(#[$m])*
555        fn neon() {
556            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
557            use {$($($path::)+neon::*),*};
558
559            unsafe { $body }
560        }
561
562        #[test]
563        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
564        $(#[$m])*
565        fn neonx2() {
566            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
567            use {$($($path::)+neonx2::*),*};
568
569            unsafe { $body }
570        }
571
572        #[test]
573        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
574        $(#[$m])*
575        fn neonx4() {
576            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
577            use {$($($path::)+neonx4::*),*};
578
579            unsafe { $body }
580        }
581
582        #[test]
583        #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
584        $(#[$m])*
585        fn neonx8() {
586            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
587            use {$($($path::)+neonx8::*),*};
588
589            unsafe { $body }
590        }
591    };
592
593    (
594        use {$($($path:ident::)+*),*};
595
596        { $($tail:tt)+ }
597    ) => {
598        #[allow(
599            clippy::reversed_empty_ranges,
600            clippy::range_plus_one,
601            clippy::modulo_one,
602            clippy::trivially_copy_pass_by_ref
603        )]
604        {
605            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
606            use {$($($path::)+scalar::*),*};
607
608            $crate::multiversion_test!(@expr { $($tail)+ });
609        }
610
611        {
612            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
613            use {$($($path::)+array128::*),*};
614
615            $crate::multiversion_test!(@expr { $($tail)+ });
616        }
617
618        {
619            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
620            use {$($($path::)+array256::*),*};
621
622            $crate::multiversion_test!(@expr { $($tail)+ });
623        }
624
625        #[cfg(feature = "all-simd")]
626        #[allow(clippy::large_types_passed_by_value)]
627        {
628            #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
629            use {$($($path::)+array4096::*),*};
630
631            $crate::multiversion_test!(@expr { $($tail)+ });
632        }
633
634        #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
635        if $crate::multiversion::Version::AVX2.supported() {
636            unsafe {
637                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
638                use {$($($path::)+avx2::*),*};
639
640                $crate::multiversion_test!(@expr { $($tail)+ });
641            }
642
643            #[cfg(feature = "all-simd")]
644            unsafe {
645                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
646                use {$($($path::)+avx2x2::*),*};
647
648                $crate::multiversion_test!(@expr { $($tail)+ });
649            }
650
651            #[cfg(feature = "all-simd")]
652            unsafe {
653                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
654                use {$($($path::)+avx2x4::*),*};
655
656                $crate::multiversion_test!(@expr { $($tail)+ });
657            }
658
659            #[cfg(feature = "all-simd")]
660            unsafe {
661                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
662                use {$($($path::)+avx2x8::*),*};
663
664                $crate::multiversion_test!(@expr { $($tail)+ });
665            }
666        }
667
668        #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
669        if $crate::multiversion::Version::AVX512.supported() {
670            unsafe {
671                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
672                use {$($($path::)+avx512::*),*};
673
674                $crate::multiversion_test!(@expr { $($tail)+ });
675            }
676
677            #[cfg(feature = "all-simd")]
678            unsafe {
679                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
680                use {$($($path::)+avx512x2::*),*};
681
682                $crate::multiversion_test!(@expr { $($tail)+ });
683            }
684
685            #[cfg(feature = "all-simd")]
686            unsafe {
687                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
688                use {$($($path::)+avx512x4::*),*};
689
690                $crate::multiversion_test!(@expr { $($tail)+ });
691            }
692
693            #[cfg(feature = "all-simd")]
694            unsafe {
695                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
696                use {$($($path::)+avx512x8::*),*};
697
698                $crate::multiversion_test!(@expr { $($tail)+ });
699            }
700        }
701
702        #[cfg(all(feature = "unsafe", target_arch = "aarch64"))]
703        {
704            unsafe {
705                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
706                use {$($($path::)+neon::*),*};
707
708                $crate::multiversion_test!(@expr { $($tail)+ });
709            }
710
711            #[cfg(feature = "all-simd")]
712            unsafe {
713                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
714                use {$($($path::)+neonx2::*),*};
715
716                $crate::multiversion_test!(@expr { $($tail)+ });
717            }
718
719            #[cfg(feature = "all-simd")]
720            unsafe {
721                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
722                use {$($($path::)+neonx4::*),*};
723
724                $crate::multiversion_test!(@expr { $($tail)+ });
725            }
726
727            #[cfg(feature = "all-simd")]
728            unsafe {
729                #[allow(clippy::allow_attributes, unused_imports, clippy::wildcard_imports)]
730                use {$($($path::)+neonx8::*),*};
731
732                $crate::multiversion_test!(@expr { $($tail)+ });
733            }
734        }
735    };
736    (@expr $e:expr) => { $e }
737}
738
739macro_rules! versions_impl {
740    ($($
741        (#[$m:meta])*
742        $name:ident $(if $supported:expr)?,
743    )+) => {
744        /// Versions generated by [`multiversion!`](crate::multiversion!) on this platform.
745        #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Default)]
746        pub enum Version {
747            #[default] // First item is the default
748            $(
749                $(#[$m])*
750                $(#[doc = concat!("Requires `", stringify!($supported), "`.")])?
751                $name,
752            )+
753        }
754
755        impl Version {
756            /// Check if this version is supported at runtime.
757            #[must_use]
758            pub fn supported(self) -> bool {
759                match self {
760                    $($(#[$m])* Version::$name => true $(&& $supported)?,)+
761                }
762            }
763        }
764
765        impl FromStr for Version {
766            type Err = UnknownVersion;
767
768            /// Implementation is case-insensitive.
769            fn from_str(x: &str) -> Result<Self, Self::Err> {
770                $(
771                    $(#[$m])*
772                    if stringify!($name).eq_ignore_ascii_case(x) { return Ok(Version::$name); }
773                )+
774                Err(UnknownVersion(x.to_string()))
775            }
776        }
777
778        /// Runtime generated list of supported versions.
779        pub static VERSIONS: std::sync::LazyLock<Vec<Version>> = std::sync::LazyLock::new(|| {
780            let mut vec = vec![$($(#[$m])* Version::$name,)+];
781            vec.retain(|i| i.supported());
782            vec
783        });
784    };
785}
786versions_impl! {
787    Scalar,
788    Array128,
789    Array256,
790    #[cfg(feature = "all-simd")]
791    Array4096,
792    #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
793    AVX2 if std::arch::is_x86_feature_detected!("avx2"),
794    #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
795    AVX2x2 if std::arch::is_x86_feature_detected!("avx2"),
796    #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
797    AVX2x4 if std::arch::is_x86_feature_detected!("avx2"),
798    #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
799    AVX2x8 if std::arch::is_x86_feature_detected!("avx2"),
800    #[cfg(all(feature = "unsafe", any(target_arch = "x86", target_arch = "x86_64")))]
801    AVX512 if std::arch::is_x86_feature_detected!("avx512f"),
802    #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
803    AVX512x2 if std::arch::is_x86_feature_detected!("avx512f"),
804    #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
805    AVX512x4 if std::arch::is_x86_feature_detected!("avx512f"),
806    #[cfg(all(feature = "unsafe", feature = "all-simd", any(target_arch = "x86", target_arch = "x86_64")))]
807    AVX512x8 if std::arch::is_x86_feature_detected!("avx512f"),
808    #[cfg(all(feature = "unsafe", target_arch = "aarch64"))]
809    Neon,
810    #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
811    Neonx2,
812    #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
813    Neonx4,
814    #[cfg(all(feature = "unsafe", feature = "all-simd", target_arch = "aarch64"))]
815    Neonx8,
816}
817
818static OVERRIDE: OnceLock<Option<Version>> = OnceLock::new();
819
820impl Version {
821    /// Get the global version override, if any.
822    ///
823    /// This function will return [`Some`] containing the version provided to
824    /// [`Version::set_override`] if it was previously called, or [`None`] otherwise.
825    ///
826    /// All dynamic dispatch implementations should respect this value.
827    pub fn get_override() -> Option<Version> {
828        *OVERRIDE.get_or_init(|| None)
829    }
830
831    /// Set the global version override.
832    ///
833    /// # Panics
834    ///
835    /// This function will panic if any of the following conditions are met:
836    ///
837    /// - The provided version is not [supported](Self::supported).
838    ///
839    /// - This function is called more than once.
840    ///
841    /// - This function is called after calling [`Version::get_override`] (or any multiversioned
842    ///   dynamic dispatch function that respects the global override).
843    pub fn set_override(version: Version) {
844        assert!(version.supported(), "{version:?} is not supported!");
845
846        if OVERRIDE.set(Some(version)).is_err() {
847            // Value returned in Err() is the value passed to set, not the existing value
848            if Self::get_override().is_none() {
849                panic!("Version::set_override must be called before get_override");
850            } else {
851                panic!("Version::set_override called more than once");
852            }
853        }
854    }
855}
856
857/// Error type returned when trying to convert an invalid string to a [`Version`].
858#[derive(Debug)]
859pub struct UnknownVersion(String);
860
861impl Display for UnknownVersion {
862    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
863        write!(f, "unknown function version: {:#}", self.0)
864    }
865}
866
867impl Error for UnknownVersion {}