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}