1use core::arch::x86_64::*;
  2
  3#[inline]
  4fn inv_mod_pow2(p: usize, bits: u32) -> usize {
  5    // Compute p^{-1} mod 2^bits (p must be odd) through Hensel lifting.
  6    debug_assert!(p % 2 == 1);
  7    let mut x: usize = 1usize; // inverse mod 2
  8    let mut i: u32 = 1;
  9    while i < bits {
 10        // x <- x * (2 - p*x)  mod 2^(2^i)  (wrapping arithmetic)
 11        x = x.wrapping_mul(2usize.wrapping_sub(p.wrapping_mul(x)));
 12        i <<= 1;
 13    }
 14    x & ((1usize << bits) - 1)
 15}
 16
 17/// # Safety
 18/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
 19/// all inputs must have the same length and must not alias.
 20#[cfg(target_arch = "x86_64")]
 21#[target_feature(enable = "avx2", enable = "fma")]
 22pub fn znx_automorphism_avx(p: i64, res: &mut [i64], a: &[i64]) {
 23    debug_assert_eq!(res.len(), a.len());
 24    let n: usize = res.len();
 25    if n == 0 {
 26        return;
 27    }
 28    debug_assert!(n.is_power_of_two(), "n must be power of two");
 29    debug_assert!(p & 1 == 1, "p must be odd (invertible mod 2n)");
 30
 31    if n < 4 {
 32        use poulpy_hal::reference::znx::znx_automorphism_ref;
 33
 34        znx_automorphism_ref(p, res, a);
 35        return;
 36    }
 37
 38    unsafe {
 39        let two_n: usize = n << 1;
 40        let span: usize = n >> 2;
 41        let bits: u32 = (two_n as u64).trailing_zeros();
 42        let mask_2n: usize = two_n - 1;
 43        let mask_1n: usize = n - 1;
 44
 45        // p mod 2n (positive)
 46        let p_2n: usize = (((p & mask_2n as i64) + two_n as i64) as usize) & mask_2n;
 47
 48        // p^-1 mod 2n
 49        let inv: usize = inv_mod_pow2(p_2n, bits);
 50
 51        // Broadcast constants
 52        let n_minus1_vec: __m256i = _mm256_set1_epi64x((n as i64) - 1);
 53        let mask_2n_vec: __m256i = _mm256_set1_epi64x(mask_2n as i64);
 54        let mask_1n_vec: __m256i = _mm256_set1_epi64x(mask_1n as i64);
 55
 56        // Lane offsets [0, inv, 2*inv, 3*inv] (mod 2n)
 57        let lane_offsets: __m256i = _mm256_set_epi64x(
 58            ((inv * 3) & mask_2n) as i64,
 59            ((inv * 2) & mask_2n) as i64,
 60            inv as i64,
 61            0i64,
 62        );
 63
 64        // t_base = (j * inv) mod 2n.
 65        let mut t_base: usize = 0;
 66        let step: usize = (inv << 2) & mask_2n;
 67
 68        let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
 69        let aa: *const i64 = a.as_ptr();
 70
 71        for _ in 0..span {
 72            // t_vec = (t_base + [0, inv, 2*inv, 3*inv]) & (2n-1)
 73            let t_base_vec: __m256i = _mm256_set1_epi64x(t_base as i64);
 74            let t_vec: __m256i = _mm256_and_si256(_mm256_add_epi64(t_base_vec, lane_offsets), mask_2n_vec);
 75
 76            // idx = t_vec & (n-1)
 77            let idx_vec: __m256i = _mm256_and_si256(t_vec, mask_1n_vec);
 78
 79            // sign = t >= n ? -1 : 0  (mask of all-ones where negate)
 80            let sign_mask: __m256i = _mm256_cmpgt_epi64(t_vec, n_minus1_vec);
 81
 82            // gather a[idx] (scale = 8 bytes per i64)
 83            let vals: __m256i = _mm256_i64gather_epi64(aa, idx_vec, 8);
 84
 85            // Conditional negate: (vals ^ sign_mask) - sign_mask
 86            let vals_x: __m256i = _mm256_xor_si256(vals, sign_mask);
 87            let out: __m256i = _mm256_sub_epi64(vals_x, sign_mask);
 88
 89            // store to res[j..j+4]
 90            _mm256_storeu_si256(rr, out);
 91
 92            // advance
 93            rr = rr.add(1);
 94            t_base = (t_base + step) & mask_2n;
 95        }
 96    }
 97}
 98
 99/// # Safety
100/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
101/// all inputs must have the same length and must not alias.
102#[cfg(all(test, any(target_arch = "x86_64", target_arch = "x86")))]
103mod tests {
104    use poulpy_hal::reference::znx::znx_automorphism_ref;
105
106    use super::*;
107
108    #[target_feature(enable = "avx2", enable = "fma")]
109    fn test_znx_automorphism_internal() {
110        let a: [i64; 16] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
111
112        let p: i64 = -5;
113
114        let mut r0: Vec<i64> = vec![0i64; a.len()];
115        let mut r1: Vec<i64> = vec![0i64; a.len()];
116
117        znx_automorphism_ref(p, &mut r0, &a);
118        znx_automorphism_avx(p, &mut r1, &a);
119
120        assert_eq!(r0, r1);
121    }
122
123    #[test]
124    fn test_znx_automorphism_avx() {
125        if !std::is_x86_feature_detected!("avx2") {
126            eprintln!("skipping: CPU lacks avx2");
127            return;
128        };
129        unsafe {
130            test_znx_automorphism_internal();
131        }
132    }
133}