arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/poulpy-backend/src/cpu_fft64_avx/znx_avx/sub.rs
poulpy-backend/src/cpu_fft64_avx/znx_avx/sub.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_sub_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_loadu_si256, _mm256_storeu_si256, _mm256_sub_epi64};
 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_sub_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_sub_ref;
 36
 37        znx_sub_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_sub_ab_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_loadu_si256, _mm256_storeu_si256, _mm256_sub_epi64};
 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_sub_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_sub_ab_inplace_ref;
 73
 74        znx_sub_ab_inplace_ref(&mut res[span << 2..], &a[span << 2..]);
 75    }
 76}
 77
 78/// # Safety
 79/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
 80/// all inputs must have the same length and must not alias.
 81#[cfg(target_arch = "x86_64")]
 82#[target_feature(enable = "avx2")]
 83pub fn znx_sub_ba_inplace_avx(res: &mut [i64], a: &[i64]) {
 84    #[cfg(debug_assertions)]
 85    {
 86        assert_eq!(res.len(), a.len());
 87    }
 88
 89    use core::arch::x86_64::{__m256i, _mm256_loadu_si256, _mm256_storeu_si256, _mm256_sub_epi64};
 90
 91    let n: usize = res.len();
 92
 93    let span: usize = n >> 2;
 94
 95    let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
 96    let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
 97
 98    unsafe {
 99        for _ in 0..span {
100            let sum: __m256i = _mm256_sub_epi64(_mm256_loadu_si256(aa), _mm256_loadu_si256(rr));
101            _mm256_storeu_si256(rr, sum);
102            rr = rr.add(1);
103            aa = aa.add(1);
104        }
105    }
106
107    // tail
108    if !res.len().is_multiple_of(4) {
109        use poulpy_hal::reference::znx::znx_sub_ba_inplace_ref;
110
111        znx_sub_ba_inplace_ref(&mut res[span << 2..], &a[span << 2..]);
112    }
113}