poulpy/poulpy-backend/src/cpu_fft64_avx/reim/fft_vec_avx2_fma.rs
poulpy-backend/src/cpu_fft64_avx/reim/fft_vec_avx2_fma.rsRaw
1/// # Safety
2/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
3#[target_feature(enable = "avx2,fma")]
4pub fn reim_add_avx2_fma(res: &mut [f64], a: &[f64], b: &[f64]) {
5 #[cfg(debug_assertions)]
6 {
7 assert_eq!(a.len(), res.len());
8 assert_eq!(b.len(), res.len());
9 }
10
11 use std::arch::x86_64::{__m256d, _mm256_add_pd, _mm256_loadu_pd, _mm256_storeu_pd};
12
13 let span: usize = res.len() >> 2;
14
15 unsafe {
16 let mut rr: *mut f64 = res.as_mut_ptr();
17 let mut aa: *const f64 = a.as_ptr();
18 let mut bb: *const f64 = b.as_ptr();
19
20 for _ in 0..span {
21 let a_256: __m256d = _mm256_loadu_pd(aa);
22 let b_256: __m256d = _mm256_loadu_pd(bb);
23 _mm256_storeu_pd(rr, _mm256_add_pd(a_256, b_256));
24 rr = rr.add(4);
25 aa = aa.add(4);
26 bb = bb.add(4);
27 }
28 }
29}
30
31/// # Safety
32/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
33#[target_feature(enable = "avx2,fma")]
34pub fn reim_add_inplace_avx2_fma(res: &mut [f64], a: &[f64]) {
35 #[cfg(debug_assertions)]
36 {
37 assert_eq!(a.len(), res.len());
38 }
39
40 use std::arch::x86_64::{__m256d, _mm256_add_pd, _mm256_loadu_pd, _mm256_storeu_pd};
41
42 let span: usize = res.len() >> 2;
43
44 unsafe {
45 let mut rr: *mut f64 = res.as_mut_ptr();
46 let mut aa: *const f64 = a.as_ptr();
47
48 for _ in 0..span {
49 let a_256: __m256d = _mm256_loadu_pd(aa);
50 let r_256: __m256d = _mm256_loadu_pd(rr);
51 _mm256_storeu_pd(rr, _mm256_add_pd(r_256, a_256));
52 rr = rr.add(4);
53 aa = aa.add(4);
54 }
55 }
56}
57
58/// # Safety
59/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
60#[target_feature(enable = "avx2,fma")]
61pub fn reim_sub_avx2_fma(res: &mut [f64], a: &[f64], b: &[f64]) {
62 #[cfg(debug_assertions)]
63 {
64 assert_eq!(a.len(), res.len());
65 assert_eq!(b.len(), res.len());
66 }
67
68 use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_sub_pd};
69
70 let span: usize = res.len() >> 2;
71
72 unsafe {
73 let mut rr: *mut f64 = res.as_mut_ptr();
74 let mut aa: *const f64 = a.as_ptr();
75 let mut bb: *const f64 = b.as_ptr();
76
77 for _ in 0..span {
78 let a_256: __m256d = _mm256_loadu_pd(aa);
79 let b_256: __m256d = _mm256_loadu_pd(bb);
80 _mm256_storeu_pd(rr, _mm256_sub_pd(a_256, b_256));
81 rr = rr.add(4);
82 aa = aa.add(4);
83 bb = bb.add(4);
84 }
85 }
86}
87
88/// # Safety
89/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
90#[target_feature(enable = "avx2,fma")]
91pub fn reim_sub_ab_inplace_avx2_fma(res: &mut [f64], a: &[f64]) {
92 #[cfg(debug_assertions)]
93 {
94 assert_eq!(a.len(), res.len());
95 }
96
97 use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_sub_pd};
98
99 let span: usize = res.len() >> 2;
100
101 unsafe {
102 let mut rr: *mut f64 = res.as_mut_ptr();
103 let mut aa: *const f64 = a.as_ptr();
104
105 for _ in 0..span {
106 let a_256: __m256d = _mm256_loadu_pd(aa);
107 let r_256: __m256d = _mm256_loadu_pd(rr);
108 _mm256_storeu_pd(rr, _mm256_sub_pd(r_256, a_256));
109 rr = rr.add(4);
110 aa = aa.add(4);
111 }
112 }
113}
114
115/// # Safety
116/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
117#[target_feature(enable = "avx2,fma")]
118pub fn reim_sub_ba_inplace_avx2_fma(res: &mut [f64], a: &[f64]) {
119 #[cfg(debug_assertions)]
120 {
121 assert_eq!(a.len(), res.len());
122 }
123
124 use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_sub_pd};
125
126 let span: usize = res.len() >> 2;
127
128 unsafe {
129 let mut rr: *mut f64 = res.as_mut_ptr();
130 let mut aa: *const f64 = a.as_ptr();
131
132 for _ in 0..span {
133 let a_256: __m256d = _mm256_loadu_pd(aa);
134 let r_256: __m256d = _mm256_loadu_pd(rr);
135 _mm256_storeu_pd(rr, _mm256_sub_pd(a_256, r_256));
136 rr = rr.add(4);
137 aa = aa.add(4);
138 }
139 }
140}
141
142/// # Safety
143/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
144#[target_feature(enable = "avx2,fma")]
145pub fn reim_negate_avx2_fma(res: &mut [f64], a: &[f64]) {
146 #[cfg(debug_assertions)]
147 {
148 assert_eq!(a.len(), res.len());
149 }
150
151 use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_xor_pd};
152
153 let span: usize = res.len() >> 2;
154
155 unsafe {
156 use std::arch::x86_64::_mm256_set1_pd;
157
158 let mut rr: *mut f64 = res.as_mut_ptr();
159 let mut aa: *const f64 = a.as_ptr();
160
161 let neg0: __m256d = _mm256_set1_pd(-0.0);
162
163 for _ in 0..span {
164 let a_256: __m256d = _mm256_loadu_pd(aa);
165 _mm256_storeu_pd(rr, _mm256_xor_pd(a_256, neg0));
166 rr = rr.add(4);
167 aa = aa.add(4);
168 }
169 }
170}
171
172/// # Safety
173/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
174#[target_feature(enable = "avx2,fma")]
175pub fn reim_negate_inplace_avx2_fma(res: &mut [f64]) {
176 use std::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd, _mm256_xor_pd};
177
178 let span: usize = res.len() >> 2;
179
180 unsafe {
181 use std::arch::x86_64::_mm256_set1_pd;
182
183 let mut rr: *mut f64 = res.as_mut_ptr();
184 let neg0: __m256d = _mm256_set1_pd(-0.0);
185
186 for _ in 0..span {
187 let r_256: __m256d = _mm256_loadu_pd(rr);
188 _mm256_storeu_pd(rr, _mm256_xor_pd(r_256, neg0));
189 rr = rr.add(4);
190 }
191 }
192}
193
194/// # Safety
195/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
196#[target_feature(enable = "avx2,fma")]
197pub fn reim_addmul_avx2_fma(res: &mut [f64], a: &[f64], b: &[f64]) {
198 #[cfg(debug_assertions)]
199 {
200 assert_eq!(a.len(), res.len());
201 assert_eq!(b.len(), res.len());
202 }
203
204 let m: usize = res.len() >> 1;
205
206 let (rr, ri) = res.split_at_mut(m);
207 let (ar, ai) = a.split_at(m);
208 let (br, bi) = b.split_at(m);
209
210 unsafe {
211 let mut rr_ptr: *mut f64 = rr.as_mut_ptr();
212 let mut ri_ptr: *mut f64 = ri.as_mut_ptr();
213 let mut ar_ptr: *const f64 = ar.as_ptr();
214 let mut ai_ptr: *const f64 = ai.as_ptr();
215 let mut br_ptr: *const f64 = br.as_ptr();
216 let mut bi_ptr: *const f64 = bi.as_ptr();
217
218 use std::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_storeu_pd};
219
220 for _ in 0..(m >> 2) {
221 let mut rr: __m256d = _mm256_loadu_pd(rr_ptr);
222 let mut ri: __m256d = _mm256_loadu_pd(ri_ptr);
223 let ar: __m256d = _mm256_loadu_pd(ar_ptr);
224 let ai: __m256d = _mm256_loadu_pd(ai_ptr);
225 let br: __m256d = _mm256_loadu_pd(br_ptr);
226 let bi: __m256d = _mm256_loadu_pd(bi_ptr);
227
228 rr = _mm256_fmsub_pd(ai, bi, rr);
229 rr = _mm256_fmsub_pd(ar, br, rr);
230 ri = _mm256_fmadd_pd(ar, bi, ri);
231 ri = _mm256_fmadd_pd(ai, br, ri);
232
233 _mm256_storeu_pd(rr_ptr, rr);
234 _mm256_storeu_pd(ri_ptr, ri);
235
236 rr_ptr = rr_ptr.add(4);
237 ri_ptr = ri_ptr.add(4);
238 ar_ptr = ar_ptr.add(4);
239 ai_ptr = ai_ptr.add(4);
240 br_ptr = br_ptr.add(4);
241 bi_ptr = bi_ptr.add(4);
242 }
243 }
244}
245
246/// # Safety
247/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
248#[target_feature(enable = "avx2,fma")]
249pub fn reim_mul_avx2_fma(res: &mut [f64], a: &[f64], b: &[f64]) {
250 #[cfg(debug_assertions)]
251 {
252 assert_eq!(a.len(), res.len());
253 assert_eq!(b.len(), res.len());
254 }
255
256 let m: usize = res.len() >> 1;
257
258 let (rr, ri) = res.split_at_mut(m);
259 let (ar, ai) = a.split_at(m);
260 let (br, bi) = b.split_at(m);
261
262 unsafe {
263 let mut rr_ptr: *mut f64 = rr.as_mut_ptr();
264 let mut ri_ptr: *mut f64 = ri.as_mut_ptr();
265 let mut ar_ptr: *const f64 = ar.as_ptr();
266 let mut ai_ptr: *const f64 = ai.as_ptr();
267 let mut br_ptr: *const f64 = br.as_ptr();
268 let mut bi_ptr: *const f64 = bi.as_ptr();
269
270 use std::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_mul_pd, _mm256_storeu_pd};
271
272 for _ in 0..(m >> 2) {
273 let ar: __m256d = _mm256_loadu_pd(ar_ptr);
274 let ai: __m256d = _mm256_loadu_pd(ai_ptr);
275 let br: __m256d = _mm256_loadu_pd(br_ptr);
276 let bi: __m256d = _mm256_loadu_pd(bi_ptr);
277
278 let t1: __m256d = _mm256_mul_pd(ai, bi);
279 let t2: __m256d = _mm256_mul_pd(ar, bi);
280
281 let rr: __m256d = _mm256_fmsub_pd(ar, br, t1);
282 let ri: __m256d = _mm256_fmadd_pd(ai, br, t2);
283
284 _mm256_storeu_pd(rr_ptr, rr);
285 _mm256_storeu_pd(ri_ptr, ri);
286
287 rr_ptr = rr_ptr.add(4);
288 ri_ptr = ri_ptr.add(4);
289 ar_ptr = ar_ptr.add(4);
290 ai_ptr = ai_ptr.add(4);
291 br_ptr = br_ptr.add(4);
292 bi_ptr = bi_ptr.add(4);
293 }
294 }
295}
296
297/// # Safety
298/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
299#[target_feature(enable = "avx2,fma")]
300pub fn reim_mul_inplace_avx2_fma(res: &mut [f64], a: &[f64]) {
301 #[cfg(debug_assertions)]
302 {
303 assert_eq!(a.len(), res.len());
304 }
305
306 let m: usize = res.len() >> 1;
307
308 let (rr, ri) = res.split_at_mut(m);
309 let (ar, ai) = a.split_at(m);
310
311 unsafe {
312 let mut rr_ptr: *mut f64 = rr.as_mut_ptr();
313 let mut ri_ptr: *mut f64 = ri.as_mut_ptr();
314 let mut ar_ptr: *const f64 = ar.as_ptr();
315 let mut ai_ptr: *const f64 = ai.as_ptr();
316
317 use std::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_mul_pd, _mm256_storeu_pd};
318
319 for _ in 0..(m >> 2) {
320 let ar: __m256d = _mm256_loadu_pd(ar_ptr);
321 let ai: __m256d = _mm256_loadu_pd(ai_ptr);
322 let br: __m256d = _mm256_loadu_pd(rr_ptr);
323 let bi: __m256d = _mm256_loadu_pd(ri_ptr);
324
325 let t1: __m256d = _mm256_mul_pd(ai, bi);
326 let t2: __m256d = _mm256_mul_pd(ar, bi);
327
328 let rr = _mm256_fmsub_pd(ar, br, t1);
329 let ri = _mm256_fmadd_pd(ai, br, t2);
330
331 _mm256_storeu_pd(rr_ptr, rr);
332 _mm256_storeu_pd(ri_ptr, ri);
333
334 rr_ptr = rr_ptr.add(4);
335 ri_ptr = ri_ptr.add(4);
336 ar_ptr = ar_ptr.add(4);
337 ai_ptr = ai_ptr.add(4);
338 }
339 }
340}