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