1#[cfg(target_arch = "x86_64")]
 2#[target_feature(enable = "avx2")]
 3pub unsafe fn znx_switch_ring_avx(res: &mut [i64], a: &[i64]) {
 4    unsafe {
 5        use core::arch::x86_64::*;
 6
 7        let (n_in, n_out) = (a.len(), res.len());
 8
 9        #[cfg(debug_assertions)]
10        {
11            assert!(n_in.is_power_of_two());
12            assert!(n_in.max(n_out).is_multiple_of(n_in.min(n_out)))
13        }
14
15        if n_in == n_out {
16            use poulpy_hal::reference::znx::znx_copy_ref;
17
18            znx_copy_ref(res, a);
19            return;
20        }
21
22        if n_in > n_out {
23            // Downsample: res[k] = a[k * gap_in], contiguous stores
24            let gap_in: usize = n_in / n_out;
25
26            // index vector: [0*gap, 1*gap, 2*gap, 3*gap] * gap_in
27            let step: __m256i = _mm256_setr_epi64x(0, gap_in as i64, 2 * gap_in as i64, 3 * gap_in as i64);
28
29            let span: usize = n_out >> 2;
30            let bump: __m256i = _mm256_set1_epi64x(4 * gap_in as i64);
31
32            let mut res_4xi64: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
33            let a_ptr: *const i64 = a.as_ptr();
34
35            let mut base: __m256i = _mm256_setzero_si256(); // starts at 0*gap
36
37            for _ in 0..span {
38                // idx = base + step
39                let idx: __m256i = _mm256_add_epi64(base, step);
40
41                // gather 4 spaced i64 (scale=8 bytes)
42                let v: __m256i = _mm256_i64gather_epi64(a_ptr, idx, 8);
43
44                // store contiguously
45                _mm256_storeu_si256(res_4xi64, v);
46
47                base = _mm256_add_epi64(base, bump);
48                res_4xi64 = res_4xi64.add(1);
49            }
50        } else {
51            // Upsample: res[k * gap_out] = a[k], i.e. res has holes;
52
53            use poulpy_hal::reference::znx::znx_zero_ref;
54            let gap_out = n_out / n_in;
55
56            // zero then scatter scalar stores
57            znx_zero_ref(res);
58
59            let mut a_4xi64: *const __m256i = a.as_ptr() as *const __m256i;
60
61            for i in (0..n_in).step_by(4) {
62                // Load contiguously 4 inputs
63                let v = _mm256_loadu_si256(a_4xi64);
64
65                // extract 4 lanes (pextrq). This is still the best we can do on AVX2.
66                let x0: i64 = _mm256_extract_epi64(v, 0);
67                let x1: i64 = _mm256_extract_epi64(v, 1);
68                let x2: i64 = _mm256_extract_epi64(v, 2);
69                let x3: i64 = _mm256_extract_epi64(v, 3);
70
71                // starting output pointer for this group
72                let mut p: *mut i64 = res.as_mut_ptr().add(i * gap_out);
73
74                // four strided stores with pointer bump (avoid mul each time)
75                *p = x0;
76                p = p.add(gap_out);
77                *p = x1;
78                p = p.add(gap_out);
79                *p = x2;
80                p = p.add(gap_out);
81                *p = x3;
82
83                a_4xi64 = a_4xi64.add(1)
84            }
85        }
86    }
87}