Skip to main content

rten_simd/arch/x86_64/
avx2.rs

1use std::arch::x86_64::{
2    __m128i, __m256, __m256i, _CMP_EQ_OQ, _CMP_GE_OQ, _CMP_GT_OQ, _CMP_LE_OQ, _CMP_LT_OQ,
3    _MM_FROUND_TO_NEAREST_INT, _MM_HINT_ET0, _MM_HINT_T0, _mm_add_ps, _mm_cvtss_f32, _mm_movehl_ps,
4    _mm_prefetch, _mm_setr_epi8, _mm_shuffle_epi8, _mm_shuffle_ps, _mm_unpacklo_epi64,
5    _mm256_add_epi8, _mm256_add_epi16, _mm256_add_epi32, _mm256_add_ps, _mm256_and_ps,
6    _mm256_and_si256, _mm256_andnot_ps, _mm256_andnot_si256, _mm256_blendv_epi8, _mm256_blendv_ps,
7    _mm256_castps256_ps128, _mm256_castsi256_si128, _mm256_cmp_ps, _mm256_cmpeq_epi8,
8    _mm256_cmpeq_epi16, _mm256_cmpeq_epi32, _mm256_cmpgt_epi8, _mm256_cmpgt_epi16,
9    _mm256_cmpgt_epi32, _mm256_cvtepi8_epi16, _mm256_cvtepi16_epi32, _mm256_cvtepi32_ps,
10    _mm256_cvtepu8_epi16, _mm256_cvtph_ps, _mm256_cvtps_epi32, _mm256_cvtps_ph,
11    _mm256_cvttps_epi32, _mm256_div_ps, _mm256_extractf128_ps, _mm256_extracti128_si256,
12    _mm256_fmadd_ps, _mm256_fnmadd_ps, _mm256_insertf128_si256, _mm256_loadu_ps,
13    _mm256_loadu_si256, _mm256_maskload_epi32, _mm256_maskload_ps, _mm256_maskstore_epi32,
14    _mm256_maskstore_ps, _mm256_max_ps, _mm256_min_ps, _mm256_movemask_epi8, _mm256_mul_ps,
15    _mm256_mullo_epi16, _mm256_mullo_epi32, _mm256_or_ps, _mm256_or_si256, _mm256_packs_epi32,
16    _mm256_packus_epi16, _mm256_permute2x128_si256, _mm256_permute4x64_epi64, _mm256_round_ps,
17    _mm256_set_m128i, _mm256_set1_epi8, _mm256_set1_epi16, _mm256_set1_epi32, _mm256_set1_ps,
18    _mm256_setr_m128i, _mm256_setzero_si256, _mm256_slli_epi16, _mm256_slli_epi32,
19    _mm256_srai_epi16, _mm256_srai_epi32, _mm256_srli_epi16, _mm256_storeu_ps, _mm256_storeu_si256,
20    _mm256_sub_epi8, _mm256_sub_epi16, _mm256_sub_epi32, _mm256_sub_ps, _mm256_unpackhi_epi8,
21    _mm256_unpackhi_epi16, _mm256_unpacklo_epi8, _mm256_unpacklo_epi16, _mm256_xor_ps,
22    _mm256_xor_si256,
23};
24use std::is_x86_feature_detected;
25use std::mem::transmute;
26
27use super::super::{lanes, simd_type};
28use crate::f16;
29use crate::ops::{
30    BitOps, Concat, Extend, FloatOps, IntOps, Interleave, MaskOps, Narrow, NarrowSaturate, NumOps,
31    SignedIntOps, ToFloat,
32};
33use crate::{Isa, Mask, Simd};
34
35simd_type!(F32x8, __m256, f32, M32, Avx2Isa);
36simd_type!(F16x16, __m256i, f16, M16, Avx2Isa);
37simd_type!(I32x8, __m256i, i32, M32, Avx2Isa);
38simd_type!(I16x16, __m256i, i16, M16, Avx2Isa);
39simd_type!(I8x32, __m256i, i8, M8, Avx2Isa);
40simd_type!(U8x32, __m256i, u8, M8, Avx2Isa);
41simd_type!(U16x16, __m256i, u16, M16, Avx2Isa);
42simd_type!(U32x8, __m256i, u32, M32, Avx2Isa);
43
44#[derive(Copy, Clone)]
45pub struct Avx2Isa {
46    _private: (),
47}
48
49impl Avx2Isa {
50    pub fn new() -> Option<Self> {
51        if is_x86_feature_detected!("avx2")
52            && is_x86_feature_detected!("fma")
53            && is_x86_feature_detected!("f16c")
54        {
55            Some(Avx2Isa { _private: () })
56        } else {
57            None
58        }
59    }
60}
61
62// Safety: AVX2 is supported as `Avx2Isa::new` checks this.
63unsafe impl Isa for Avx2Isa {
64    type M32 = M32;
65    type M16 = M16;
66    type M8 = M8;
67    type F32 = F32x8;
68    type I32 = I32x8;
69    type I16 = I16x16;
70    type I8 = I8x32;
71    type U8 = U8x32;
72    type U16 = U16x16;
73    type U32 = U32x8;
74    type F16 = F16x16;
75    type Bits = I32x8;
76
77    fn f32(
78        self,
79    ) -> impl FloatOps<f32, Simd = Self::F32, Int = Self::I32>
80    + NarrowSaturate<f32, f16, Output = Self::F16> {
81        self
82    }
83
84    fn f16(self) -> impl Extend<f16, Output = Self::F32, Simd = Self::F16> {
85        self
86    }
87
88    fn i32(
89        self,
90    ) -> impl SignedIntOps<i32, Simd = Self::I32>
91    + NarrowSaturate<i32, i16, Output = Self::I16>
92    + Concat<i32>
93    + ToFloat<i32, Output = Self::F32> {
94        self
95    }
96
97    fn i16(
98        self,
99    ) -> impl SignedIntOps<i16, Simd = Self::I16>
100    + NarrowSaturate<i16, u8, Output = Self::U8>
101    + Extend<i16, Output = Self::I32>
102    + Interleave<i16> {
103        self
104    }
105
106    fn i8(
107        self,
108    ) -> impl SignedIntOps<i8, Simd = Self::I8> + Extend<i8, Output = Self::I16> + Interleave<i8>
109    {
110        self
111    }
112
113    fn u8(
114        self,
115    ) -> impl IntOps<u8, Simd = Self::U8> + Extend<u8, Output = Self::U16> + Interleave<u8> {
116        self
117    }
118
119    fn u16(self) -> impl IntOps<u16, Simd = Self::U16> {
120        self
121    }
122
123    fn m32(self) -> impl MaskOps<Self::M32> {
124        self
125    }
126
127    fn m16(self) -> impl MaskOps<Self::M16> {
128        self
129    }
130
131    fn m8(self) -> impl MaskOps<Self::M8> {
132        self
133    }
134}
135
136macro_rules! simd_ops_common {
137    ($simd:ty, $mask:ty) => {
138        type Simd = $simd;
139
140        #[inline]
141        fn len(self) -> usize {
142            lanes::<$simd>()
143        }
144
145        #[inline]
146        fn prefetch(self, ptr: *const <$simd as Simd>::Elem) {
147            unsafe { _mm_prefetch(ptr as *const i8, _MM_HINT_T0) }
148        }
149
150        #[inline]
151        fn prefetch_write(self, ptr: *mut <$simd as Simd>::Elem) {
152            unsafe { _mm_prefetch(ptr as *const i8, _MM_HINT_ET0) }
153        }
154    };
155}
156
157macro_rules! simd_int_ops_common {
158    ($simd:ty) => {
159        #[inline]
160        fn and(self, x: $simd, y: $simd) -> $simd {
161            unsafe { _mm256_and_si256(x.0, y.0) }.into()
162        }
163
164        #[inline]
165        fn or(self, x: $simd, y: $simd) -> $simd {
166            unsafe { _mm256_or_si256(x.0, y.0) }.into()
167        }
168
169        #[inline]
170        fn xor(self, x: $simd, y: $simd) -> $simd {
171            unsafe { _mm256_xor_si256(x.0, y.0) }.into()
172        }
173
174        #[inline]
175        fn not(self, x: $simd) -> $simd {
176            unsafe { _mm256_andnot_si256(x.0, _mm256_set1_epi8(-1)) }.into()
177        }
178    };
179}
180
181unsafe impl BitOps<f32> for Avx2Isa {
182    simd_ops_common!(F32x8, M32);
183
184    #[inline]
185    fn first_n_mask(self, n: usize) -> M32 {
186        let mask: [i32; 8] = std::array::from_fn(|i| if i < n { -1 } else { 0 });
187        M32::from_float(unsafe { _mm256_loadu_ps(mask.as_ptr() as *const f32) })
188    }
189
190    #[inline]
191    fn and(self, x: F32x8, y: F32x8) -> F32x8 {
192        unsafe { _mm256_and_ps(x.0, y.0) }.into()
193    }
194
195    #[inline]
196    fn not(self, x: F32x8) -> F32x8 {
197        let all_ones: F32x8 = self.splat(f32::from_bits(0xFFFFFFFF));
198        unsafe { _mm256_andnot_ps(x.0, all_ones.0) }.into()
199    }
200
201    #[inline]
202    fn or(self, x: F32x8, y: F32x8) -> F32x8 {
203        unsafe { _mm256_or_ps(x.0, y.0) }.into()
204    }
205
206    #[inline]
207    fn xor(self, x: F32x8, y: F32x8) -> F32x8 {
208        unsafe { _mm256_xor_ps(x.0, y.0) }.into()
209    }
210
211    #[inline]
212    fn splat(self, x: f32) -> F32x8 {
213        unsafe { _mm256_set1_ps(x) }.into()
214    }
215
216    #[inline]
217    unsafe fn load_ptr(self, ptr: *const f32) -> F32x8 {
218        unsafe { _mm256_loadu_ps(ptr) }.into()
219    }
220
221    #[inline]
222    fn select(self, x: F32x8, y: F32x8, mask: M32) -> F32x8 {
223        unsafe { _mm256_blendv_ps(y.0, x.0, mask.as_float()) }.into()
224    }
225
226    #[inline]
227    unsafe fn load_ptr_mask(self, ptr: *const f32, mask: M32) -> F32x8 {
228        unsafe { _mm256_maskload_ps(ptr, mask.0) }.into()
229    }
230
231    #[inline]
232    unsafe fn store_ptr_mask(self, x: F32x8, ptr: *mut f32, mask: M32) {
233        unsafe { _mm256_maskstore_ps(ptr, mask.0, x.0) }
234    }
235
236    #[inline]
237    unsafe fn store_ptr(self, x: F32x8, ptr: *mut f32) {
238        unsafe { _mm256_storeu_ps(ptr, x.0) }
239    }
240}
241
242unsafe impl NumOps<f32> for Avx2Isa {
243    #[inline]
244    fn add(self, x: F32x8, y: F32x8) -> F32x8 {
245        unsafe { _mm256_add_ps(x.0, y.0) }.into()
246    }
247
248    #[inline]
249    fn sub(self, x: F32x8, y: F32x8) -> F32x8 {
250        unsafe { _mm256_sub_ps(x.0, y.0) }.into()
251    }
252
253    #[inline]
254    fn mul(self, x: F32x8, y: F32x8) -> F32x8 {
255        unsafe { _mm256_mul_ps(x.0, y.0) }.into()
256    }
257
258    #[inline]
259    fn mul_add(self, a: F32x8, b: F32x8, c: F32x8) -> F32x8 {
260        unsafe { _mm256_fmadd_ps(a.0, b.0, c.0) }.into()
261    }
262
263    #[inline]
264    fn lt(self, x: F32x8, y: F32x8) -> M32 {
265        M32::from_float(unsafe { _mm256_cmp_ps(x.0, y.0, _CMP_LT_OQ) })
266    }
267
268    #[inline]
269    fn le(self, x: F32x8, y: F32x8) -> M32 {
270        M32::from_float(unsafe { _mm256_cmp_ps(x.0, y.0, _CMP_LE_OQ) })
271    }
272
273    #[inline]
274    fn eq(self, x: F32x8, y: F32x8) -> M32 {
275        M32::from_float(unsafe { _mm256_cmp_ps(x.0, y.0, _CMP_EQ_OQ) })
276    }
277
278    #[inline]
279    fn ge(self, x: F32x8, y: F32x8) -> M32 {
280        M32::from_float(unsafe { _mm256_cmp_ps(x.0, y.0, _CMP_GE_OQ) })
281    }
282
283    #[inline]
284    fn gt(self, x: F32x8, y: F32x8) -> M32 {
285        M32::from_float(unsafe { _mm256_cmp_ps(x.0, y.0, _CMP_GT_OQ) })
286    }
287
288    #[inline]
289    fn min(self, x: F32x8, y: F32x8) -> F32x8 {
290        unsafe { _mm256_min_ps(x.0, y.0) }.into()
291    }
292
293    #[inline]
294    fn max(self, x: F32x8, y: F32x8) -> F32x8 {
295        unsafe { _mm256_max_ps(x.0, y.0) }.into()
296    }
297
298    #[inline]
299    fn sum(self, x: F32x8) -> f32 {
300        // See https://stackoverflow.com/a/13222410/434243
301        unsafe {
302            let hi_4 = _mm256_extractf128_ps(x.0, 1);
303            let lo_4 = _mm256_castps256_ps128(x.0);
304            let sum_4 = _mm_add_ps(lo_4, hi_4);
305            let lo_2 = sum_4;
306            let hi_2 = _mm_movehl_ps(sum_4, sum_4);
307            let sum_2 = _mm_add_ps(lo_2, hi_2);
308            let lo = sum_2;
309            let hi = _mm_shuffle_ps(sum_2, sum_2, 0x1);
310            let sum = _mm_add_ps(lo, hi);
311            _mm_cvtss_f32(sum)
312        }
313    }
314}
315
316impl FloatOps<f32> for Avx2Isa {
317    type Int = <Self as Isa>::I32;
318
319    #[inline]
320    fn div(self, x: F32x8, y: F32x8) -> F32x8 {
321        unsafe { _mm256_div_ps(x.0, y.0) }.into()
322    }
323
324    #[inline]
325    fn abs(self, x: F32x8) -> F32x8 {
326        unsafe { _mm256_andnot_ps(_mm256_set1_ps(-0.0), x.0) }.into()
327    }
328
329    #[inline]
330    fn neg(self, x: F32x8) -> F32x8 {
331        unsafe { _mm256_xor_ps(x.0, _mm256_set1_ps(-0.0)) }.into()
332    }
333
334    #[inline]
335    fn mul_sub_from(self, a: F32x8, b: F32x8, c: F32x8) -> F32x8 {
336        unsafe { _mm256_fnmadd_ps(a.0, b.0, c.0) }.into()
337    }
338
339    #[inline]
340    fn round_ties_even(self, x: F32x8) -> F32x8 {
341        unsafe { _mm256_round_ps(x.0, _MM_FROUND_TO_NEAREST_INT) }.into()
342    }
343
344    #[inline]
345    fn to_int_trunc(self, x: F32x8) -> Self::Int {
346        unsafe { _mm256_cvttps_epi32(x.0) }.into()
347    }
348
349    #[inline]
350    fn to_int_round(self, x: F32x8) -> Self::Int {
351        unsafe { _mm256_cvtps_epi32(x.0) }.into()
352    }
353}
354
355unsafe impl BitOps<i32> for Avx2Isa {
356    simd_ops_common!(I32x8, M32);
357    simd_int_ops_common!(I32x8);
358
359    #[inline]
360    fn first_n_mask(self, n: usize) -> M32 {
361        let mask: [i32; 8] = std::array::from_fn(|i| if i < n { -1 } else { 0 });
362        M32(unsafe { _mm256_loadu_si256(mask.as_ptr() as *const __m256i) })
363    }
364
365    #[inline]
366    fn splat(self, x: i32) -> I32x8 {
367        unsafe { _mm256_set1_epi32(x) }.into()
368    }
369
370    #[inline]
371    unsafe fn load_ptr(self, ptr: *const i32) -> I32x8 {
372        unsafe { _mm256_loadu_si256(ptr as *const __m256i) }.into()
373    }
374
375    #[inline]
376    fn select(self, x: I32x8, y: I32x8, mask: M32) -> I32x8 {
377        unsafe { _mm256_blendv_epi8(y.0, x.0, mask.0) }.into()
378    }
379
380    #[inline]
381    unsafe fn store_ptr(self, x: I32x8, ptr: *mut i32) {
382        unsafe { _mm256_storeu_si256(ptr as *mut __m256i, x.0) }
383    }
384
385    #[inline]
386    unsafe fn load_ptr_mask(self, ptr: *const i32, mask: M32) -> I32x8 {
387        unsafe { _mm256_maskload_epi32(ptr, mask.0) }.into()
388    }
389
390    #[inline]
391    unsafe fn store_ptr_mask(self, x: I32x8, ptr: *mut i32, mask: M32) {
392        unsafe { _mm256_maskstore_epi32(ptr, mask.0, x.0) }
393    }
394}
395
396unsafe impl NumOps<i32> for Avx2Isa {
397    #[inline]
398    fn add(self, x: I32x8, y: I32x8) -> I32x8 {
399        unsafe { _mm256_add_epi32(x.0, y.0) }.into()
400    }
401
402    #[inline]
403    fn sub(self, x: I32x8, y: I32x8) -> I32x8 {
404        unsafe { _mm256_sub_epi32(x.0, y.0) }.into()
405    }
406
407    #[inline]
408    fn mul(self, x: I32x8, y: I32x8) -> I32x8 {
409        unsafe { _mm256_mullo_epi32(x.0, y.0) }.into()
410    }
411
412    #[inline]
413    fn eq(self, x: I32x8, y: I32x8) -> M32 {
414        M32(unsafe { _mm256_cmpeq_epi32(x.0, y.0) })
415    }
416
417    #[inline]
418    fn ge(self, x: I32x8, y: I32x8) -> M32 {
419        M32(unsafe { _mm256_or_si256(_mm256_cmpgt_epi32(x.0, y.0), _mm256_cmpeq_epi32(x.0, y.0)) })
420    }
421
422    #[inline]
423    fn gt(self, x: I32x8, y: I32x8) -> M32 {
424        M32(unsafe { _mm256_cmpgt_epi32(x.0, y.0) })
425    }
426}
427
428impl IntOps<i32> for Avx2Isa {
429    #[inline]
430    fn shift_left<const SHIFT: i32>(self, x: I32x8) -> I32x8 {
431        unsafe { _mm256_slli_epi32(x.0, SHIFT) }.into()
432    }
433
434    #[inline]
435    fn shift_right<const SHIFT: i32>(self, x: I32x8) -> I32x8 {
436        unsafe { _mm256_srai_epi32(x.0, SHIFT) }.into()
437    }
438}
439
440impl SignedIntOps<i32> for Avx2Isa {
441    #[inline]
442    fn neg(self, x: I32x8) -> I32x8 {
443        unsafe { _mm256_sub_epi32(_mm256_setzero_si256(), x.0) }.into()
444    }
445}
446
447/// Copied from unstable `_MM_SHUFFLE` function in `core::arch::x86`.
448const fn _mm_shuffle(z: u32, y: u32, x: u32, w: u32) -> i32 {
449    ((z << 6) | (y << 4) | (x << 2) | w) as i32
450}
451
452impl NarrowSaturate<i32, i16> for Avx2Isa {
453    type Output = I16x16;
454
455    #[inline]
456    fn narrow_saturate(self, low: I32x8, high: I32x8) -> I16x16 {
457        unsafe {
458            // AVX2 pack functions treat each input as 2 128-bit lanes and
459            // interleave narrowed 64-bit blocks from each input. Shuffle the
460            // output to get narrowed lanes from `low` followed by lanes from
461            // high.
462            let packed = _mm256_packs_epi32(low.0, high.0);
463            _mm256_permute4x64_epi64(packed, _mm_shuffle(3, 1, 2, 0))
464        }
465        .into()
466    }
467}
468
469impl Concat<i32> for Avx2Isa {
470    #[inline]
471    fn concat_low(self, a: I32x8, b: I32x8) -> I32x8 {
472        unsafe {
473            let a_lo = _mm256_castsi256_si128(a.0);
474            let b_lo = _mm256_castsi256_si128(b.0);
475            _mm256_set_m128i(b_lo, a_lo)
476        }
477        .into()
478    }
479
480    #[inline]
481    fn concat_high(self, a: I32x8, b: I32x8) -> I32x8 {
482        unsafe {
483            let a_hi = _mm256_extracti128_si256(a.0, 1);
484            let b_hi = _mm256_extracti128_si256(b.0, 1);
485            _mm256_set_m128i(b_hi, a_hi)
486        }
487        .into()
488    }
489}
490
491impl ToFloat<i32> for Avx2Isa {
492    type Output = F32x8;
493
494    #[inline]
495    fn to_float(self, x: I32x8) -> F32x8 {
496        unsafe { _mm256_cvtepi32_ps(x.0) }.into()
497    }
498}
499
500unsafe impl BitOps<i16> for Avx2Isa {
501    simd_ops_common!(I16x16, M16);
502    simd_int_ops_common!(I16x16);
503
504    #[inline]
505    fn first_n_mask(self, n: usize) -> M16 {
506        let mask: [i16; 16] = std::array::from_fn(|i| if i < n { -1 } else { 0 });
507        M16(unsafe { _mm256_loadu_si256(mask.as_ptr() as *const __m256i) })
508    }
509
510    #[inline]
511    fn splat(self, x: i16) -> I16x16 {
512        unsafe { _mm256_set1_epi16(x) }.into()
513    }
514
515    #[inline]
516    unsafe fn load_ptr(self, ptr: *const i16) -> I16x16 {
517        unsafe { _mm256_loadu_si256(ptr as *const __m256i) }.into()
518    }
519
520    #[inline]
521    fn select(self, x: I16x16, y: I16x16, mask: M16) -> I16x16 {
522        unsafe { _mm256_blendv_epi8(y.0, x.0, mask.0) }.into()
523    }
524
525    #[inline]
526    unsafe fn store_ptr(self, x: I16x16, ptr: *mut i16) {
527        unsafe { _mm256_storeu_si256(ptr as *mut __m256i, x.0) }
528    }
529
530    #[inline]
531    unsafe fn load_ptr_mask(self, ptr: *const i16, mask: M16) -> I16x16 {
532        // There is no native masked-load instruction for i16, so fall back to
533        // scalar loads.
534        let mask = _mm256_movemask_epi8(mask.0) as u32;
535        let xs: [i16; 16] = std::array::from_fn(|i| {
536            let mask_bit = mask & (1 << (i * 2 + 1));
537            if mask_bit != 0 {
538                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
539                unsafe { *ptr.add(i) }
540            } else {
541                0
542            }
543        });
544        self.load_ptr(xs.as_ptr())
545    }
546
547    #[inline]
548    unsafe fn store_ptr_mask(self, x: I16x16, ptr: *mut i16, mask: M16) {
549        // There is no native masked-store instruction for i16, so fall back to
550        // scalar store.
551        let xs = Simd::to_array(x);
552        let mask = _mm256_movemask_epi8(mask.0) as u32;
553        for i in 0..16 {
554            let mask_bit = mask & (1 << (i * 2 + 1));
555            if mask_bit != 0 {
556                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
557                unsafe { *ptr.add(i) = xs[i] }
558            }
559        }
560    }
561}
562
563unsafe impl NumOps<i16> for Avx2Isa {
564    #[inline]
565    fn add(self, x: I16x16, y: I16x16) -> I16x16 {
566        unsafe { _mm256_add_epi16(x.0, y.0) }.into()
567    }
568
569    #[inline]
570    fn sub(self, x: I16x16, y: I16x16) -> I16x16 {
571        unsafe { _mm256_sub_epi16(x.0, y.0) }.into()
572    }
573
574    #[inline]
575    fn mul(self, x: I16x16, y: I16x16) -> I16x16 {
576        unsafe { _mm256_mullo_epi16(x.0, y.0) }.into()
577    }
578
579    #[inline]
580    fn eq(self, x: I16x16, y: I16x16) -> M16 {
581        M16(unsafe { _mm256_cmpeq_epi16(x.0, y.0) })
582    }
583
584    #[inline]
585    fn ge(self, x: I16x16, y: I16x16) -> M16 {
586        M16(unsafe { _mm256_or_si256(_mm256_cmpgt_epi16(x.0, y.0), _mm256_cmpeq_epi16(x.0, y.0)) })
587    }
588
589    #[inline]
590    fn gt(self, x: I16x16, y: I16x16) -> M16 {
591        M16(unsafe { _mm256_cmpgt_epi16(x.0, y.0) })
592    }
593}
594
595impl IntOps<i16> for Avx2Isa {
596    #[inline]
597    fn shift_left<const SHIFT: i32>(self, x: I16x16) -> I16x16 {
598        unsafe { _mm256_slli_epi16(x.0, SHIFT) }.into()
599    }
600
601    #[inline]
602    fn shift_right<const SHIFT: i32>(self, x: I16x16) -> I16x16 {
603        unsafe { _mm256_srai_epi16(x.0, SHIFT) }.into()
604    }
605}
606
607impl SignedIntOps<i16> for Avx2Isa {
608    #[inline]
609    fn neg(self, x: I16x16) -> I16x16 {
610        unsafe { _mm256_sub_epi16(_mm256_setzero_si256(), x.0) }.into()
611    }
612}
613
614impl NarrowSaturate<i16, u8> for Avx2Isa {
615    type Output = U8x32;
616
617    #[inline]
618    fn narrow_saturate(self, low: I16x16, high: I16x16) -> U8x32 {
619        unsafe {
620            // AVX2 pack functions treat each input as 2 128-bit lanes and
621            // interleave narrowed 64-bit blocks from each input. Shuffle the
622            // output to get narrowed lanes from `low` followed by lanes from
623            // high.
624            let packed = _mm256_packus_epi16(low.0, high.0);
625            _mm256_permute4x64_epi64(packed, _mm_shuffle(3, 1, 2, 0))
626        }
627        .into()
628    }
629}
630
631impl Interleave<i16> for Avx2Isa {
632    #[inline]
633    fn interleave_low(self, a: I16x16, b: I16x16) -> I16x16 {
634        unsafe {
635            // AB{N} = Interleaved Nth 64-bit block.
636            let lo = _mm256_unpacklo_epi16(a.0, b.0); // AB0 AB2
637            let hi = _mm256_unpackhi_epi16(a.0, b.0); // AB1 AB3
638            _mm256_insertf128_si256(lo, _mm256_castsi256_si128(hi), 1) // AB0 AB1
639        }
640        .into()
641    }
642
643    #[inline]
644    fn interleave_high(self, a: I16x16, b: I16x16) -> I16x16 {
645        unsafe {
646            // AB{N} = Interleaved Nth 64-bit block.
647            let lo = _mm256_unpacklo_epi16(a.0, b.0); // AB0 AB2
648            let hi = _mm256_unpackhi_epi16(a.0, b.0); // AB1 AB3
649            _mm256_permute2x128_si256(lo, hi, 0x31) // AB2 AB3
650        }
651        .into()
652    }
653}
654
655unsafe impl BitOps<i8> for Avx2Isa {
656    simd_ops_common!(I8x32, M8);
657    simd_int_ops_common!(I8x32);
658
659    #[inline]
660    fn first_n_mask(self, n: usize) -> M8 {
661        let mask: [i8; 32] = std::array::from_fn(|i| if i < n { -1 } else { 0 });
662        M8(unsafe { _mm256_loadu_si256(mask.as_ptr() as *const __m256i) })
663    }
664
665    #[inline]
666    fn splat(self, x: i8) -> I8x32 {
667        unsafe { _mm256_set1_epi8(x) }.into()
668    }
669
670    #[inline]
671    unsafe fn load_ptr(self, ptr: *const i8) -> I8x32 {
672        unsafe { _mm256_loadu_si256(ptr as *const __m256i) }.into()
673    }
674
675    #[inline]
676    fn select(self, x: I8x32, y: I8x32, mask: <I8x32 as Simd>::Mask) -> I8x32 {
677        unsafe { _mm256_blendv_epi8(y.0, x.0, mask.0) }.into()
678    }
679
680    #[inline]
681    unsafe fn store_ptr(self, x: I8x32, ptr: *mut i8) {
682        unsafe { _mm256_storeu_si256(ptr as *mut __m256i, x.0) }
683    }
684
685    #[inline]
686    unsafe fn load_ptr_mask(self, ptr: *const i8, mask: M8) -> I8x32 {
687        // There is no native masked-load instruction for i8, so fall back to
688        // scalar loads.
689        let mask = _mm256_movemask_epi8(mask.0) as u32;
690        let xs: [i8; 32] = std::array::from_fn(|i| {
691            let mask_bit = mask & (1 << i);
692            if mask_bit != 0 {
693                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
694                unsafe { *ptr.add(i) }
695            } else {
696                0
697            }
698        });
699        self.load_ptr(xs.as_ptr())
700    }
701
702    #[inline]
703    unsafe fn store_ptr_mask(self, x: I8x32, ptr: *mut i8, mask: M8) {
704        // There is no native masked-store instruction for i8, so fall back to
705        // scalar store.
706        let xs = Simd::to_array(x);
707        let mask = _mm256_movemask_epi8(mask.0) as u32;
708        for i in 0..32 {
709            let mask_bit = mask & (1 << i);
710            if mask_bit != 0 {
711                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
712                unsafe { *ptr.add(i) = xs[i] }
713            }
714        }
715    }
716}
717
718unsafe impl NumOps<i8> for Avx2Isa {
719    #[inline]
720    fn add(self, x: I8x32, y: I8x32) -> I8x32 {
721        unsafe { _mm256_add_epi8(x.0, y.0) }.into()
722    }
723
724    #[inline]
725    fn sub(self, x: I8x32, y: I8x32) -> I8x32 {
726        unsafe { _mm256_sub_epi8(x.0, y.0) }.into()
727    }
728
729    #[inline]
730    fn mul(self, x: I8x32, y: I8x32) -> I8x32 {
731        let x_lo = Extend::<i8>::extend_low(self, x);
732        let x_hi = Extend::<i8>::extend_high(self, x);
733        let y_lo = Extend::<i8>::extend_low(self, y);
734        let y_hi = Extend::<i8>::extend_high(self, y);
735
736        let i16_ops = self.i16();
737        let prod_lo = i16_ops.mul(x_lo, y_lo);
738        let prod_hi = i16_ops.mul(x_hi, y_hi);
739
740        self.narrow_truncate(prod_lo, prod_hi)
741    }
742
743    #[inline]
744    fn eq(self, x: I8x32, y: I8x32) -> M8 {
745        M8(unsafe { _mm256_cmpeq_epi8(x.0, y.0) })
746    }
747
748    #[inline]
749    fn ge(self, x: I8x32, y: I8x32) -> M8 {
750        M8(unsafe { _mm256_or_si256(_mm256_cmpgt_epi8(x.0, y.0), _mm256_cmpeq_epi8(x.0, y.0)) })
751    }
752
753    #[inline]
754    fn gt(self, x: I8x32, y: I8x32) -> M8 {
755        M8(unsafe { _mm256_cmpgt_epi8(x.0, y.0) })
756    }
757}
758
759impl IntOps<i8> for Avx2Isa {
760    #[inline]
761    fn shift_left<const SHIFT: i32>(self, x: I8x32) -> I8x32 {
762        let x_lo = Extend::<i8>::extend_low(self, x);
763        let x_hi = Extend::<i8>::extend_high(self, x);
764
765        let i16_ops = self.i16();
766        let y_lo = i16_ops.shift_left::<SHIFT>(x_lo);
767        let y_hi = i16_ops.shift_left::<SHIFT>(x_hi);
768
769        self.narrow_truncate(y_lo, y_hi)
770    }
771
772    #[inline]
773    fn shift_right<const SHIFT: i32>(self, x: I8x32) -> I8x32 {
774        let x_lo = Extend::<i8>::extend_low(self, x);
775        let x_hi = Extend::<i8>::extend_high(self, x);
776
777        let i16_ops = self.i16();
778        let y_lo = i16_ops.shift_right::<SHIFT>(x_lo);
779        let y_hi = i16_ops.shift_right::<SHIFT>(x_hi);
780
781        self.narrow_truncate(y_lo, y_hi)
782    }
783}
784
785impl SignedIntOps<i8> for Avx2Isa {
786    #[inline]
787    fn neg(self, x: I8x32) -> I8x32 {
788        unsafe { _mm256_sub_epi8(_mm256_setzero_si256(), x.0) }.into()
789    }
790}
791
792#[inline]
793fn interleave_low_x8(a: __m256i, b: __m256i) -> __m256i {
794    unsafe {
795        // AB{N} = Interleaved Nth 64-bit block.
796        let lo = _mm256_unpacklo_epi8(a, b); // AB0 AB2
797        let hi = _mm256_unpackhi_epi8(a, b); // AB1 AB3
798        _mm256_insertf128_si256(lo, _mm256_castsi256_si128(hi), 1) // AB0 AB1
799    }
800}
801
802#[inline]
803fn interleave_high_x8(a: __m256i, b: __m256i) -> __m256i {
804    unsafe {
805        // AB{N} = Interleaved Nth 64-bit block.
806        let lo = _mm256_unpacklo_epi8(a, b); // AB0 AB2
807        let hi = _mm256_unpackhi_epi8(a, b); // AB1 AB3
808        _mm256_permute2x128_si256(lo, hi, 0x31) // AB2 AB3
809    }
810}
811
812impl Interleave<i8> for Avx2Isa {
813    #[inline]
814    fn interleave_low(self, a: I8x32, b: I8x32) -> I8x32 {
815        interleave_low_x8(a.0, b.0).into()
816    }
817
818    #[inline]
819    fn interleave_high(self, a: I8x32, b: I8x32) -> I8x32 {
820        interleave_high_x8(a.0, b.0).into()
821    }
822}
823
824unsafe impl BitOps<u8> for Avx2Isa {
825    simd_ops_common!(U8x32, M8);
826    simd_int_ops_common!(U8x32);
827
828    #[inline]
829    fn first_n_mask(self, n: usize) -> M8 {
830        let mask: [i8; 32] = std::array::from_fn(|i| if i < n { -1 } else { 0 });
831        M8(unsafe { _mm256_loadu_si256(mask.as_ptr() as *const __m256i) })
832    }
833
834    #[inline]
835    fn splat(self, x: u8) -> U8x32 {
836        unsafe { _mm256_set1_epi8(x as i8) }.into()
837    }
838
839    #[inline]
840    unsafe fn load_ptr(self, ptr: *const u8) -> U8x32 {
841        unsafe { _mm256_loadu_si256(ptr as *const __m256i) }.into()
842    }
843
844    #[inline]
845    fn select(self, x: U8x32, y: U8x32, mask: M8) -> U8x32 {
846        unsafe { _mm256_blendv_epi8(y.0, x.0, mask.0) }.into()
847    }
848
849    #[inline]
850    unsafe fn store_ptr(self, x: U8x32, ptr: *mut u8) {
851        unsafe { _mm256_storeu_si256(ptr as *mut __m256i, x.0) }
852    }
853
854    #[inline]
855    unsafe fn load_ptr_mask(self, ptr: *const u8, mask: M8) -> U8x32 {
856        // There is no native masked-load instruction for u8, so fall back to
857        // scalar loads.
858        let mask = _mm256_movemask_epi8(mask.0) as u32;
859        let xs: [u8; 32] = std::array::from_fn(|i| {
860            let mask_bit = mask & (1 << i);
861            if mask_bit != 0 {
862                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
863                unsafe { *ptr.add(i) }
864            } else {
865                0
866            }
867        });
868        self.load_ptr(xs.as_ptr())
869    }
870
871    #[inline]
872    unsafe fn store_ptr_mask(self, x: U8x32, ptr: *mut u8, mask: M8) {
873        // There is no native masked-store instruction for u8, so fall back to
874        // scalar store.
875        let xs = Simd::to_array(x);
876        let mask = _mm256_movemask_epi8(mask.0) as u32;
877        for i in 0..32 {
878            let mask_bit = mask & (1 << i);
879            if mask_bit != 0 {
880                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
881                unsafe { *ptr.add(i) = xs[i] }
882            }
883        }
884    }
885}
886
887unsafe impl NumOps<u8> for Avx2Isa {
888    #[inline]
889    fn add(self, x: U8x32, y: U8x32) -> U8x32 {
890        unsafe { _mm256_add_epi8(x.0, y.0) }.into()
891    }
892
893    #[inline]
894    fn sub(self, x: U8x32, y: U8x32) -> U8x32 {
895        unsafe { _mm256_sub_epi8(x.0, y.0) }.into()
896    }
897
898    #[inline]
899    fn mul(self, x: U8x32, y: U8x32) -> U8x32 {
900        let x_lo = Extend::<u8>::extend_low(self, x);
901        let x_hi = Extend::<u8>::extend_high(self, x);
902        let y_lo = Extend::<u8>::extend_low(self, y);
903        let y_hi = Extend::<u8>::extend_high(self, y);
904
905        let u16_ops = self.u16();
906        let prod_lo = u16_ops.mul(x_lo, y_lo);
907        let prod_hi = u16_ops.mul(x_hi, y_hi);
908
909        self.narrow_truncate(prod_lo, prod_hi)
910    }
911
912    #[inline]
913    fn eq(self, x: U8x32, y: U8x32) -> M8 {
914        M8(unsafe { _mm256_cmpeq_epi8(x.0, y.0) })
915    }
916
917    #[inline]
918    fn ge(self, x: U8x32, y: U8x32) -> M8 {
919        let xy_eq = <Self as NumOps<u8>>::eq(self, x, y);
920        let xy_gt = <Self as NumOps<u8>>::gt(self, x, y);
921        M8(unsafe { _mm256_or_si256(xy_eq.0, xy_gt.0) })
922    }
923
924    #[inline]
925    fn gt(self, x: U8x32, y: U8x32) -> M8 {
926        // AVX2 lacks u8 comparison. Shift both values to i8 and use signed compare.
927        M8(unsafe {
928            let mask = _mm256_set1_epi8(0x80u8 as i8);
929            let x_i8 = _mm256_xor_si256(x.0, mask);
930            let y_i8 = _mm256_xor_si256(y.0, mask);
931            _mm256_cmpgt_epi8(x_i8, y_i8)
932        })
933    }
934}
935
936unsafe impl BitOps<u16> for Avx2Isa {
937    simd_ops_common!(U16x16, M16);
938    simd_int_ops_common!(U16x16);
939
940    #[inline]
941    fn first_n_mask(self, n: usize) -> M16 {
942        let mask: [i16; 16] = std::array::from_fn(|i| if i < n { -1 } else { 0 });
943        M16(unsafe { _mm256_loadu_si256(mask.as_ptr() as *const __m256i) })
944    }
945
946    #[inline]
947    fn splat(self, x: u16) -> U16x16 {
948        unsafe { _mm256_set1_epi16(x as i16) }.into()
949    }
950
951    #[inline]
952    unsafe fn load_ptr(self, ptr: *const u16) -> U16x16 {
953        unsafe { _mm256_loadu_si256(ptr as *const __m256i) }.into()
954    }
955
956    #[inline]
957    fn select(self, x: U16x16, y: U16x16, mask: M16) -> U16x16 {
958        unsafe { _mm256_blendv_epi8(y.0, x.0, mask.0) }.into()
959    }
960
961    #[inline]
962    unsafe fn store_ptr(self, x: U16x16, ptr: *mut u16) {
963        unsafe { _mm256_storeu_si256(ptr as *mut __m256i, x.0) }
964    }
965
966    #[inline]
967    unsafe fn load_ptr_mask(self, ptr: *const u16, mask: M16) -> U16x16 {
968        // There is no native masked-load instruction for i16, so fall back to
969        // scalar loads.
970        let mask = _mm256_movemask_epi8(mask.0) as u32;
971        let xs: [u16; 16] = std::array::from_fn(|i| {
972            let mask_bit = mask & (1 << (i * 2 + 1));
973            if mask_bit != 0 {
974                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
975                unsafe { *ptr.add(i) }
976            } else {
977                0
978            }
979        });
980        self.load_ptr(xs.as_ptr())
981    }
982
983    #[inline]
984    unsafe fn store_ptr_mask(self, x: U16x16, ptr: *mut u16, mask: M16) {
985        // There is no native masked-store instruction for i16, so fall back to
986        // scalar store.
987        let xs = Simd::to_array(x);
988        let mask = _mm256_movemask_epi8(mask.0) as u32;
989        for i in 0..16 {
990            let mask_bit = mask & (1 << (i * 2 + 1));
991            if mask_bit != 0 {
992                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
993                unsafe { *ptr.add(i) = xs[i] }
994            }
995        }
996    }
997}
998
999unsafe impl NumOps<u16> for Avx2Isa {
1000    #[inline]
1001    fn add(self, x: U16x16, y: U16x16) -> U16x16 {
1002        unsafe { _mm256_add_epi16(x.0, y.0) }.into()
1003    }
1004
1005    #[inline]
1006    fn sub(self, x: U16x16, y: U16x16) -> U16x16 {
1007        unsafe { _mm256_sub_epi16(x.0, y.0) }.into()
1008    }
1009
1010    #[inline]
1011    fn mul(self, x: U16x16, y: U16x16) -> U16x16 {
1012        unsafe { _mm256_mullo_epi16(x.0, y.0) }.into()
1013    }
1014
1015    #[inline]
1016    fn eq(self, x: U16x16, y: U16x16) -> M16 {
1017        M16(unsafe { _mm256_cmpeq_epi16(x.0, y.0) })
1018    }
1019
1020    #[inline]
1021    fn ge(self, x: U16x16, y: U16x16) -> M16 {
1022        let xy_eq = <Self as NumOps<u16>>::eq(self, x, y);
1023        let xy_gt = <Self as NumOps<u16>>::gt(self, x, y);
1024        M16(unsafe { _mm256_or_si256(xy_eq.0, xy_gt.0) })
1025    }
1026
1027    #[inline]
1028    fn gt(self, x: U16x16, y: U16x16) -> M16 {
1029        // AVX2 lacks u16 comparison. Shift both values to i16 and use signed compare.
1030        M16(unsafe {
1031            let mask = _mm256_set1_epi16(0x8000u16 as i16);
1032            let x_i16 = _mm256_xor_si256(x.0, mask);
1033            let y_i16 = _mm256_xor_si256(y.0, mask);
1034            _mm256_cmpgt_epi16(x_i16, y_i16)
1035        })
1036    }
1037}
1038
1039impl IntOps<u16> for Avx2Isa {
1040    #[inline]
1041    fn shift_left<const SHIFT: i32>(self, x: U16x16) -> U16x16 {
1042        unsafe { _mm256_slli_epi16(x.0, SHIFT) }.into()
1043    }
1044
1045    #[inline]
1046    fn shift_right<const SHIFT: i32>(self, x: U16x16) -> U16x16 {
1047        unsafe { _mm256_srli_epi16(x.0, SHIFT) }.into()
1048    }
1049}
1050
1051unsafe impl BitOps<f16> for Avx2Isa {
1052    simd_ops_common!(F16x16, M16);
1053    simd_int_ops_common!(F16x16);
1054
1055    #[inline]
1056    fn first_n_mask(self, n: usize) -> M16 {
1057        let mask: [i16; 16] = std::array::from_fn(|i| if i < n { -1 } else { 0 });
1058        M16(unsafe { _mm256_loadu_si256(mask.as_ptr() as *const __m256i) })
1059    }
1060
1061    #[inline]
1062    fn splat(self, x: f16) -> F16x16 {
1063        unsafe { _mm256_set1_epi16(x.to_bits() as i16) }.into()
1064    }
1065
1066    #[inline]
1067    unsafe fn load_ptr(self, ptr: *const f16) -> F16x16 {
1068        unsafe { _mm256_loadu_si256(ptr as *const __m256i) }.into()
1069    }
1070
1071    #[inline]
1072    fn select(self, x: F16x16, y: F16x16, mask: M16) -> F16x16 {
1073        unsafe { _mm256_blendv_epi8(y.0, x.0, mask.0) }.into()
1074    }
1075
1076    #[inline]
1077    unsafe fn store_ptr(self, x: F16x16, ptr: *mut f16) {
1078        unsafe { _mm256_storeu_si256(ptr as *mut __m256i, x.0) }
1079    }
1080
1081    #[inline]
1082    unsafe fn load_ptr_mask(self, ptr: *const f16, mask: M16) -> F16x16 {
1083        // There is no native masked-load instruction for 16-bit lanes, so fall
1084        // back to scalar loads.
1085        let mask = _mm256_movemask_epi8(mask.0) as u32;
1086        let xs: [f16; 16] = std::array::from_fn(|i| {
1087            let mask_bit = mask & (1 << (i * 2 + 1));
1088            if mask_bit != 0 {
1089                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
1090                unsafe { *ptr.add(i) }
1091            } else {
1092                f16::default()
1093            }
1094        });
1095        self.load_ptr(xs.as_ptr())
1096    }
1097
1098    #[inline]
1099    unsafe fn store_ptr_mask(self, x: F16x16, ptr: *mut f16, mask: M16) {
1100        // There is no native masked-store instruction for 16-bit lanes, so fall
1101        // back to scalar store.
1102        let xs = Simd::to_array(x);
1103        let mask = _mm256_movemask_epi8(mask.0) as u32;
1104        for i in 0..16 {
1105            let mask_bit = mask & (1 << (i * 2 + 1));
1106            if mask_bit != 0 {
1107                // Safety: Caller promises that `ptr.add(i)` is valid if mask[i] is set.
1108                unsafe { *ptr.add(i) = xs[i] }
1109            }
1110        }
1111    }
1112}
1113
1114impl Extend<f16> for Avx2Isa {
1115    type Output = F32x8;
1116
1117    #[inline]
1118    fn extend_low(self, x: F16x16) -> F32x8 {
1119        unsafe { _mm256_cvtph_ps(_mm256_castsi256_si128(x.0)).into() }
1120    }
1121
1122    #[inline]
1123    fn extend_high(self, x: F16x16) -> F32x8 {
1124        unsafe { _mm256_cvtph_ps(_mm256_extracti128_si256(x.0, 1)).into() }
1125    }
1126}
1127
1128impl NarrowSaturate<f32, f16> for Avx2Isa {
1129    type Output = F16x16;
1130
1131    #[inline]
1132    fn narrow_saturate(self, low: F32x8, high: F32x8) -> F16x16 {
1133        unsafe {
1134            let low_i128 = _mm256_cvtps_ph::<_MM_FROUND_TO_NEAREST_INT>(low.0);
1135            let high_i128 = _mm256_cvtps_ph::<_MM_FROUND_TO_NEAREST_INT>(high.0);
1136            _mm256_set_m128i(high_i128, low_i128).into()
1137        }
1138    }
1139}
1140
1141macro_rules! impl_mask {
1142    ($mask:ident, $elem:ty, $len:expr) => {
1143        #[derive(Copy, Clone, Debug)]
1144        #[repr(transparent)]
1145        pub struct $mask(__m256i);
1146
1147        impl $mask {
1148            #[allow(unused)] // Not used for M16/M8
1149            #[inline]
1150            fn as_float(self) -> __m256 {
1151                unsafe { transmute::<__m256i, __m256>(self.0) }
1152            }
1153
1154            #[allow(unused)] // Not used for M16/M8
1155            #[inline]
1156            fn from_float(m: __m256) -> Self {
1157                Self(unsafe { transmute::<__m256, __m256i>(m) })
1158            }
1159        }
1160
1161        impl Mask for $mask {
1162            type Array = [bool; $len];
1163
1164            #[inline]
1165            fn to_array(self) -> Self::Array {
1166                let array = unsafe { transmute::<Self, [$elem; $len]>(self) };
1167                std::array::from_fn(|i| array[i] != <$elem>::default())
1168            }
1169        }
1170    };
1171}
1172
1173impl_mask!(M32, u32, 8);
1174impl_mask!(M16, u16, 16);
1175impl_mask!(M8, u8, 32);
1176
1177macro_rules! impl_mask_ops {
1178    ($mask:ident) => {
1179        unsafe impl MaskOps<$mask> for Avx2Isa {
1180            #[inline]
1181            fn and(self, x: $mask, y: $mask) -> $mask {
1182                $mask(unsafe { _mm256_and_si256(x.0, y.0) })
1183            }
1184
1185            #[inline]
1186            fn any(self, x: $mask) -> bool {
1187                unsafe { _mm256_movemask_epi8(x.0) != 0 }
1188            }
1189
1190            #[inline]
1191            fn all(self, x: $mask) -> bool {
1192                unsafe { _mm256_movemask_epi8(x.0) == -1 }
1193            }
1194        }
1195    };
1196}
1197impl_mask_ops!(M32);
1198impl_mask_ops!(M16);
1199impl_mask_ops!(M8);
1200
1201impl Extend<i16> for Avx2Isa {
1202    type Output = I32x8;
1203
1204    #[inline]
1205    fn extend_low(self, x: I16x16) -> Self::Output {
1206        unsafe { _mm256_cvtepi16_epi32(_mm256_castsi256_si128(x.0)).into() }
1207    }
1208
1209    #[inline]
1210    fn extend_high(self, x: I16x16) -> Self::Output {
1211        unsafe { _mm256_cvtepi16_epi32(_mm256_extracti128_si256(x.0, 1)).into() }
1212    }
1213}
1214
1215impl Extend<i8> for Avx2Isa {
1216    type Output = I16x16;
1217
1218    #[inline]
1219    fn extend_low(self, x: I8x32) -> Self::Output {
1220        unsafe { _mm256_cvtepi8_epi16(_mm256_castsi256_si128(x.0)).into() }
1221    }
1222
1223    #[inline]
1224    fn extend_high(self, x: I8x32) -> Self::Output {
1225        unsafe { _mm256_cvtepi8_epi16(_mm256_extracti128_si256(x.0, 1)).into() }
1226    }
1227}
1228
1229impl Extend<u8> for Avx2Isa {
1230    type Output = U16x16;
1231
1232    #[inline]
1233    fn extend_low(self, x: U8x32) -> Self::Output {
1234        unsafe { _mm256_cvtepu8_epi16(_mm256_castsi256_si128(x.0)).into() }
1235    }
1236
1237    #[inline]
1238    fn extend_high(self, x: U8x32) -> Self::Output {
1239        unsafe { _mm256_cvtepu8_epi16(_mm256_extracti128_si256(x.0, 1)).into() }
1240    }
1241}
1242
1243impl IntOps<u8> for Avx2Isa {
1244    #[inline(always)]
1245    fn shift_left<const SHIFT: i32>(self, x: U8x32) -> U8x32 {
1246        let x_lo = Extend::<u8>::extend_low(self, x);
1247        let x_hi = Extend::<u8>::extend_high(self, x);
1248
1249        let u16_ops = self.u16();
1250        let y_lo = u16_ops.shift_left::<SHIFT>(x_lo);
1251        let y_hi = u16_ops.shift_left::<SHIFT>(x_hi);
1252
1253        self.narrow_truncate(y_lo, y_hi)
1254    }
1255
1256    #[inline(always)]
1257    fn shift_right<const SHIFT: i32>(self, x: U8x32) -> U8x32 {
1258        let x_lo = Extend::<u8>::extend_low(self, x);
1259        let x_hi = Extend::<u8>::extend_high(self, x);
1260
1261        let u16_ops = self.u16();
1262        let y_lo = u16_ops.shift_right::<SHIFT>(x_lo);
1263        let y_hi = u16_ops.shift_right::<SHIFT>(x_hi);
1264
1265        self.narrow_truncate(y_lo, y_hi)
1266    }
1267}
1268
1269impl Interleave<u8> for Avx2Isa {
1270    #[inline]
1271    fn interleave_low(self, a: U8x32, b: U8x32) -> U8x32 {
1272        interleave_low_x8(a.0, b.0).into()
1273    }
1274
1275    #[inline]
1276    fn interleave_high(self, a: U8x32, b: U8x32) -> U8x32 {
1277        interleave_high_x8(a.0, b.0).into()
1278    }
1279}
1280
1281/// Extract bytes at even indices.
1282///
1283/// Given an input with 16-bit lanes, this extracts truncated 8-bit values.
1284#[inline]
1285unsafe fn extract_even_bytes(vec: __m256i) -> __m128i {
1286    let lo = _mm256_extracti128_si256(vec, 0);
1287    let hi = _mm256_extracti128_si256(vec, 1);
1288
1289    // Shuffle mask that moves bytes at even indices into first half of output.
1290    // For the second half set the high bit to zero the bytes.
1291    let mask = _mm_setr_epi8(0, 2, 4, 6, 8, 10, 12, 14, -1, -1, -1, -1, -1, -1, -1, -1);
1292
1293    // Extract even bytes from each half, then concatenate.
1294    let lo_even = _mm_shuffle_epi8(lo, mask);
1295    let hi_even = _mm_shuffle_epi8(hi, mask);
1296    _mm_unpacklo_epi64(lo_even, hi_even)
1297}
1298
1299impl Narrow<I16x16> for Avx2Isa {
1300    type Output = I8x32;
1301
1302    #[inline]
1303    fn narrow_truncate(self, low: I16x16, high: I16x16) -> Self::Output {
1304        let low_even = unsafe { extract_even_bytes(low.0) };
1305        let high_even = unsafe { extract_even_bytes(high.0) };
1306        let combined = unsafe { _mm256_setr_m128i(low_even, high_even) };
1307        I8x32(combined)
1308    }
1309}
1310
1311impl Narrow<U16x16> for Avx2Isa {
1312    type Output = U8x32;
1313
1314    #[inline]
1315    fn narrow_truncate(self, low: U16x16, high: U16x16) -> Self::Output {
1316        let low_even = unsafe { extract_even_bytes(low.0) };
1317        let high_even = unsafe { extract_even_bytes(high.0) };
1318        let combined = unsafe { _mm256_setr_m128i(low_even, high_even) };
1319        U8x32(combined)
1320    }
1321}