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