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}