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