competitive/tools/
avx_helper.rs

1#[macro_export]
2macro_rules! avx_helper {
3    (@avx512 $(#[$meta:meta])* $vis:vis fn $name:ident$(<$($T:ident),+>)?($($i:ident: $t:ty),*) -> $ret:ty where [$($clauses:tt)*] $body:block) => {
4        $(#[$meta])*
5        $vis fn $name$(<$($T)*>)?($($i: $t),*) -> $ret
6        where
7            $($clauses)*
8        {
9            if is_x86_feature_detected!("avx512f")
10                && is_x86_feature_detected!("avx512dq")
11                && is_x86_feature_detected!("avx512cd")
12                && is_x86_feature_detected!("avx512bw")
13                && is_x86_feature_detected!("avx512vl")
14            {
15                $crate::avx_helper!(@def_avx512 fn avx512$(<$($T)*>)?($($i: $t),*) -> $ret where [$($clauses)*] $body);
16                unsafe { avx512$(::<$($T),*>)?($($i),*) }
17            } else if is_x86_feature_detected!("avx2") {
18                $crate::avx_helper!(@def_avx2 fn avx2$(<$($T)*>)?($($i: $t),*) -> $ret where [$($clauses)*] $body);
19                unsafe { avx2$(::<$($T),*>)?($($i),*) }
20            } else {
21                $body
22            }
23        }
24    };
25    (@avx2 $(#[$meta:meta])* $vis:vis fn $name:ident$(<$($T:ident),+>)?($($i:ident: $t:ty),*) -> $ret:ty where [$($clauses:tt)*] $body:block) => {
26        $(#[$meta])*
27        $vis fn $name$(<$($T)*>)?($($i: $t),*) -> $ret
28        where
29            $($clauses)*
30        {
31            if is_x86_feature_detected!("avx2") {
32                $crate::avx_helper!(@def_avx2 fn avx2$(<$($T)*>)?($($i: $t),*) -> $ret where [$($clauses)*] $body);
33                unsafe { avx2$(::<$($T),*>)?($($i),*) }
34            } else {
35                $body
36            }
37        }
38    };
39    (@def_avx512 fn $name:ident$(<$($T:ident),+>)?($($args:tt)*) -> $ret:ty where [$($clauses:tt)*] $body:block) => {
40        #[target_feature(enable = "avx512f,avx512dq,avx512cd,avx512bw,avx512vl")]
41        unsafe fn $name$(<$($T)*>)?($($args)*) -> $ret
42        where
43            $($clauses)*
44        $body
45    };
46    (@def_avx2 fn $name:ident$(<$($T:ident),+>)?($($args:tt)*) -> $ret:ty where [$($clauses:tt)*] $body:block) => {
47        #[target_feature(enable = "avx2")]
48        unsafe fn $name$(<$($T)*>)?($($args)*) -> $ret
49        where
50            $($clauses)*
51        $body
52    };
53    (@$tag:ident $(#[$meta:meta])* $vis:vis fn $name:ident$(<$($T:ident),+>)?($($args:tt)*) -> $ret:ty $body:block) => {
54        $crate::avx_helper!(@$tag $(#[$meta])* $vis fn $name$(<$($T)*>)?($($args)*) -> $ret where [] $body);
55    };
56    (@$tag:ident $(#[$meta:meta])* $vis:vis fn $name:ident$(<$($T:ident),+>)?($($args:tt)*) $($t:tt)*) => {
57        $crate::avx_helper!(@$tag $(#[$meta])* $vis fn $name$(<$($T)*>)?($($args)*) -> () $($t)*);
58    };
59    ($($t:tt)*) => {
60        ::std::compile_error!($($t)*);
61    }
62}
competitive/tools/avx_helper.rs

competitive/tools/
avx_helper.rs
