arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-cpu-avx/src/znx_avx/neg.rs
poulpy-cpu-avx/src/znx_avx/neg.rsRaw
 1/// # Safety
 2/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
 3/// all inputs must have the same length and must not alias.
 4#[target_feature(enable = "avx2")]
 5pub fn znx_negate_avx(res: &mut [i64], src: &[i64]) {
 6    #[cfg(debug_assertions)]
 7    {
 8        assert_eq!(res.len(), src.len())
 9    }
10
11    let n: usize = res.len();
12
13    use std::arch::x86_64::{__m256i, _mm256_loadu_si256, _mm256_setzero_si256, _mm256_storeu_si256, _mm256_sub_epi64};
14    let span: usize = n >> 2;
15
16    unsafe {
17        let mut aa: *const __m256i = src.as_ptr() as *const __m256i;
18        let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
19        let zero: __m256i = _mm256_setzero_si256();
20        for _ in 0..span {
21            let v: __m256i = _mm256_loadu_si256(aa);
22            let neg: __m256i = _mm256_sub_epi64(zero, v);
23            _mm256_storeu_si256(rr, neg);
24            aa = aa.add(1);
25            rr = rr.add(1);
26        }
27    }
28
29    if !res.len().is_multiple_of(4) {
30        use poulpy_hal::reference::znx::znx_negate_ref;
31
32        znx_negate_ref(&mut res[span << 2..], &src[span << 2..])
33    }
34}
35
36/// # Safety
37/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
38/// all inputs must have the same length and must not alias.
39#[target_feature(enable = "avx2")]
40pub fn znx_negate_inplace_avx(res: &mut [i64]) {
41    let n: usize = res.len();
42
43    use std::arch::x86_64::{__m256i, _mm256_loadu_si256, _mm256_setzero_si256, _mm256_storeu_si256, _mm256_sub_epi64};
44    let span: usize = n >> 2;
45
46    unsafe {
47        let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
48        let zero: __m256i = _mm256_setzero_si256();
49        for _ in 0..span {
50            let v: __m256i = _mm256_loadu_si256(rr);
51            let neg: __m256i = _mm256_sub_epi64(zero, v);
52            _mm256_storeu_si256(rr, neg);
53            rr = rr.add(1);
54        }
55    }
56
57    if !res.len().is_multiple_of(4) {
58        use poulpy_hal::reference::znx::znx_negate_inplace_ref;
59
60        znx_negate_inplace_ref(&mut res[span << 2..])
61    }
62}