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