utils/simd/
avx2.rs

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

utils/simd/
avx2.rs
