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