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#[target_feature(enable = "avx2")]
  5pub fn znx_sub_avx(res: &mut [i64], a: &[i64], b: &[i64]) {
  6    #[cfg(debug_assertions)]
  7    {
  8        assert_eq!(res.len(), a.len());
  9        assert_eq!(res.len(), b.len());
 10    }
 11
 12    use core::arch::x86_64::{__m256i, _mm256_loadu_si256, _mm256_storeu_si256, _mm256_sub_epi64};
 13
 14    let n: usize = res.len();
 15
 16    let span: usize = n >> 2;
 17
 18    let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
 19    let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
 20    let mut bb: *const __m256i = b.as_ptr() as *const __m256i;
 21
 22    unsafe {
 23        for _ in 0..span {
 24            let sum: __m256i = _mm256_sub_epi64(_mm256_loadu_si256(aa), _mm256_loadu_si256(bb));
 25            _mm256_storeu_si256(rr, sum);
 26            rr = rr.add(1);
 27            aa = aa.add(1);
 28            bb = bb.add(1);
 29        }
 30    }
 31
 32    // tail
 33    if !res.len().is_multiple_of(4) {
 34        use poulpy_hal::reference::znx::znx_sub_ref;
 35
 36        znx_sub_ref(&mut res[span << 2..], &a[span << 2..], &b[span << 2..]);
 37    }
 38}
 39
 40/// # Safety
 41/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
 42/// all inputs must have the same length and must not alias.
 43#[target_feature(enable = "avx2")]
 44pub fn znx_sub_inplace_avx(res: &mut [i64], a: &[i64]) {
 45    #[cfg(debug_assertions)]
 46    {
 47        assert_eq!(res.len(), a.len());
 48    }
 49
 50    use core::arch::x86_64::{__m256i, _mm256_loadu_si256, _mm256_storeu_si256, _mm256_sub_epi64};
 51
 52    let n: usize = res.len();
 53
 54    let span: usize = n >> 2;
 55
 56    let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
 57    let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
 58
 59    unsafe {
 60        for _ in 0..span {
 61            let sum: __m256i = _mm256_sub_epi64(_mm256_loadu_si256(rr), _mm256_loadu_si256(aa));
 62            _mm256_storeu_si256(rr, sum);
 63            rr = rr.add(1);
 64            aa = aa.add(1);
 65        }
 66    }
 67
 68    // tail
 69    if !res.len().is_multiple_of(4) {
 70        use poulpy_hal::reference::znx::znx_sub_inplace_ref;
 71
 72        znx_sub_inplace_ref(&mut res[span << 2..], &a[span << 2..]);
 73    }
 74}
 75
 76/// # Safety
 77/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
 78/// all inputs must have the same length and must not alias.
 79#[target_feature(enable = "avx2")]
 80pub fn znx_sub_negate_inplace_avx(res: &mut [i64], a: &[i64]) {
 81    #[cfg(debug_assertions)]
 82    {
 83        assert_eq!(res.len(), a.len());
 84    }
 85
 86    use core::arch::x86_64::{__m256i, _mm256_loadu_si256, _mm256_storeu_si256, _mm256_sub_epi64};
 87
 88    let n: usize = res.len();
 89
 90    let span: usize = n >> 2;
 91
 92    let mut rr: *mut __m256i = res.as_mut_ptr() as *mut __m256i;
 93    let mut aa: *const __m256i = a.as_ptr() as *const __m256i;
 94
 95    unsafe {
 96        for _ in 0..span {
 97            let sum: __m256i = _mm256_sub_epi64(_mm256_loadu_si256(aa), _mm256_loadu_si256(rr));
 98            _mm256_storeu_si256(rr, sum);
 99            rr = rr.add(1);
100            aa = aa.add(1);
101        }
102    }
103
104    // tail
105    if !res.len().is_multiple_of(4) {
106        use poulpy_hal::reference::znx::znx_sub_negate_inplace_ref;
107
108        znx_sub_negate_inplace_ref(&mut res[span << 2..], &a[span << 2..]);
109    }
110}