utils/simd/
avx512.rs

1//! AVX512 vector implementations.
2//!
3//! Currently only requires AVX-512F.
4
5use std::array::from_fn;
6use std::ops::{Add, BitAnd, BitOr, BitXor, Not};
7
8#[cfg(target_arch = "x86_64")]
9#[allow(clippy::allow_attributes, clippy::wildcard_imports)]
10use std::arch::x86_64::*;
11
12#[cfg(target_arch = "x86")]
13#[allow(clippy::allow_attributes, clippy::wildcard_imports)]
14use std::arch::x86::*;
15
16/// AVX512 [u32] vector implementation.
17#[derive(Clone, Copy)]
18#[repr(transparent)]
19pub struct U32Vector<const V: usize, const L: usize>([__m512i; V]);
20
21impl<const V: usize, const L: usize> From<[u32; L]> for U32Vector<V, L> {
22    #[inline]
23    fn from(value: [u32; L]) -> Self {
24        Self(from_fn(|i| unsafe {
25            #[expect(
26                clippy::cast_ptr_alignment,
27                reason = "_mm512_loadu_si512 is an unaligned load which requires no alignment"
28            )]
29            _mm512_loadu_si512(value[i * 16..].as_ptr().cast::<__m512i>())
30        }))
31    }
32}
33
34impl<const V: usize, const L: usize> From<U32Vector<V, L>> for [u32; L] {
35    #[inline]
36    fn from(value: U32Vector<V, L>) -> Self {
37        let mut result = [0; L];
38        let (chunks, []) = result.as_chunks_mut::<16>() else {
39            unreachable!("L is a multiple of 16");
40        };
41        for (&v, r) in value.0.iter().zip(chunks) {
42            unsafe {
43                #[expect(
44                    clippy::cast_ptr_alignment,
45                    reason = "_mm512_storeu_si512 is an unaligned store which requires no alignment"
46                )]
47                _mm512_storeu_si512(r.as_mut_ptr().cast::<__m512i>(), v);
48            }
49        }
50        result
51    }
52}
53
54impl<const V: usize, const L: usize> Add for U32Vector<V, L> {
55    type Output = Self;
56
57    #[inline]
58    fn add(self, rhs: Self) -> Self::Output {
59        Self(from_fn(|i| unsafe {
60            _mm512_add_epi32(self.0[i], rhs.0[i])
61        }))
62    }
63}
64
65impl<const V: usize, const L: usize> BitAnd for U32Vector<V, L> {
66    type Output = Self;
67
68    #[inline]
69    fn bitand(self, rhs: Self) -> Self::Output {
70        Self(from_fn(|i| unsafe {
71            _mm512_and_si512(self.0[i], rhs.0[i])
72        }))
73    }
74}
75
76impl<const V: usize, const L: usize> BitOr for U32Vector<V, L> {
77    type Output = Self;
78
79    #[inline]
80    fn bitor(self, rhs: Self) -> Self::Output {
81        Self(from_fn(|i| unsafe { _mm512_or_si512(self.0[i], rhs.0[i]) }))
82    }
83}
84
85impl<const V: usize, const L: usize> BitXor for U32Vector<V, L> {
86    type Output = Self;
87
88    #[inline]
89    fn bitxor(self, rhs: Self) -> Self::Output {
90        Self(from_fn(|i| unsafe {
91            _mm512_xor_si512(self.0[i], rhs.0[i])
92        }))
93    }
94}
95
96impl<const V: usize, const L: usize> Not for U32Vector<V, L> {
97    type Output = Self;
98
99    #[inline]
100    fn not(self) -> Self::Output {
101        Self(from_fn(|i| unsafe {
102            _mm512_xor_si512(self.0[i], _mm512_set1_epi8(!0))
103        }))
104    }
105}
106
107impl<const V: usize, const L: usize> U32Vector<V, L> {
108    pub const LANES: usize = {
109        assert!(V * 16 == L);
110        L
111    };
112
113    #[inline]
114    #[must_use]
115    #[target_feature(enable = "avx512f")]
116    pub fn andnot(self, rhs: Self) -> Self {
117        Self(from_fn(|i| _mm512_andnot_si512(rhs.0[i], self.0[i])))
118    }
119
120    #[inline]
121    #[must_use]
122    #[target_feature(enable = "avx512f")]
123    pub fn splat(v: u32) -> Self {
124        Self(
125            #[expect(clippy::cast_possible_wrap)]
126            [_mm512_set1_epi32(v as i32); V],
127        )
128    }
129
130    #[inline]
131    #[must_use]
132    #[target_feature(enable = "avx512f")]
133    pub fn rotate_left(self, n: u32) -> Self {
134        Self(from_fn(|i| {
135            #[expect(clippy::cast_possible_wrap)]
136            _mm512_or_si512(
137                _mm512_sll_epi32(self.0[i], _mm_cvtsi32_si128(n as i32)),
138                _mm512_srl_epi32(self.0[i], _mm_cvtsi32_si128(32 - n as i32)),
139            )
140        }))
141    }
142}
143
144/// Vector implementations using a single AVX512 vector.
145pub mod avx512 {
146    /// The name of this backend.
147    pub const SIMD_BACKEND: &str = "avx512";
148
149    /// AVX512 vector with sixteen [u32] lanes.
150    pub type U32Vector = super::U32Vector<1, 16>;
151}
152
153/// Vector implementations using two AVX512 vectors.
154#[cfg(feature = "all-simd")]
155pub mod avx512x2 {
156    /// The name of this backend.
157    pub const SIMD_BACKEND: &str = "avx512x2";
158
159    /// Two AVX512 vectors with thirty-two total [u32] lanes.
160    pub type U32Vector = super::U32Vector<2, 32>;
161}
162
163/// Vector implementations using four AVX512 vectors.
164#[cfg(feature = "all-simd")]
165pub mod avx512x4 {
166    /// The name of this backend.
167    pub const SIMD_BACKEND: &str = "avx512x4";
168
169    /// Four AVX512 vectors with sixty-four total [u32] lanes.
170    pub type U32Vector = super::U32Vector<4, 64>;
171}
172
173/// Vector implementations using eight AVX512 vectors.
174#[cfg(feature = "all-simd")]
175pub mod avx512x8 {
176    /// The name of this backend.
177    pub const SIMD_BACKEND: &str = "avx512x8";
178
179    /// Eight AVX512 vectors with 128 total [u32] lanes.
180    pub type U32Vector = super::U32Vector<8, 128>;
181}
utils/simd/avx512.rs

utils/simd/
avx512.rs
