poulpy/poulpy-cpu-avx/src/reim4/arithmetic_avx.rs
poulpy-cpu-avx/src/reim4/arithmetic_avx.rsRaw
1/// # Safety
2/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
3#[target_feature(enable = "avx")]
4pub fn reim4_extract_1blk_from_reim_contiguous_avx(m: usize, rows: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
5 use core::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd};
6
7 unsafe {
8 let mut src_ptr: *const __m256d = src.as_ptr().add(blk << 2) as *const __m256d; // src + 4*blk
9 let mut dst_ptr: *mut __m256d = dst.as_mut_ptr() as *mut __m256d;
10
11 let step: usize = m >> 2;
12
13 // Each iteration copies 4 doubles; advance src by m doubles each row
14 for _ in 0..2 * rows {
15 let v: __m256d = _mm256_loadu_pd(src_ptr as *const f64);
16 _mm256_storeu_pd(dst_ptr as *mut f64, v);
17 dst_ptr = dst_ptr.add(1);
18 src_ptr = src_ptr.add(step);
19 }
20 }
21}
22
23/// # Safety
24/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
25#[target_feature(enable = "avx")]
26pub fn reim4_save_1blk_to_reim_contiguous_avx(m: usize, rows: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
27 use core::arch::x86_64::{__m256d, _mm256_loadu_pd, _mm256_storeu_pd};
28
29 unsafe {
30 let mut src_ptr: *const __m256d = src.as_ptr() as *const __m256d;
31 let mut dst_ptr: *mut __m256d = dst.as_mut_ptr().add(blk << 2) as *mut __m256d; // dst + 4*blk
32
33 let step: usize = m >> 2;
34
35 // Each iteration copies 4 doubles; advance dst by m doubles each row
36 for _ in 0..2 * rows {
37 let v: __m256d = _mm256_loadu_pd(src_ptr as *const f64);
38 _mm256_storeu_pd(dst_ptr as *mut f64, v);
39 dst_ptr = dst_ptr.add(step);
40 src_ptr = src_ptr.add(1);
41 }
42 }
43}
44
45/// # Safety
46/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
47#[target_feature(enable = "avx2,fma")]
48pub fn reim4_save_1blk_to_reim_avx<const OVERWRITE: bool>(m: usize, blk: usize, dst: &mut [f64], src: &[f64]) {
49 use core::arch::x86_64::{__m256d, _mm256_add_pd, _mm256_loadu_pd, _mm256_storeu_pd};
50 unsafe {
51 let off: usize = blk * 4;
52 let src_ptr: *const f64 = src.as_ptr();
53
54 let s0: __m256d = _mm256_loadu_pd(src_ptr);
55 let s1: __m256d = _mm256_loadu_pd(src_ptr.add(4));
56
57 let d0_ptr: *mut f64 = dst.as_mut_ptr().add(off);
58 let d1_ptr: *mut f64 = d0_ptr.add(m);
59
60 if OVERWRITE {
61 _mm256_storeu_pd(d0_ptr, s0);
62 _mm256_storeu_pd(d1_ptr, s1);
63 } else {
64 let d0: __m256d = _mm256_loadu_pd(d0_ptr);
65 let d1: __m256d = _mm256_loadu_pd(d1_ptr);
66 _mm256_storeu_pd(d0_ptr, _mm256_add_pd(d0, s0));
67 _mm256_storeu_pd(d1_ptr, _mm256_add_pd(d1, s1));
68 }
69 }
70}
71
72/// # Safety
73/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
74#[target_feature(enable = "avx2,fma")]
75pub fn reim4_save_2blk_to_reim_avx<const OVERWRITE: bool>(
76 m: usize, //
77 blk: usize, // block index
78 dst: &mut [f64], //
79 src: &[f64], // 16 doubles [re1(4), im1(4), re2(4), im2(4)]
80) {
81 use core::arch::x86_64::{__m256d, _mm256_add_pd, _mm256_loadu_pd, _mm256_storeu_pd};
82 unsafe {
83 let off: usize = blk * 4;
84 let src_ptr: *const f64 = src.as_ptr();
85
86 let d0_ptr: *mut f64 = dst.as_mut_ptr().add(off);
87 let d1_ptr: *mut f64 = d0_ptr.add(m);
88 let d2_ptr: *mut f64 = d1_ptr.add(m);
89 let d3_ptr: *mut f64 = d2_ptr.add(m);
90
91 let s0: __m256d = _mm256_loadu_pd(src_ptr);
92 let s1: __m256d = _mm256_loadu_pd(src_ptr.add(4));
93 let s2: __m256d = _mm256_loadu_pd(src_ptr.add(8));
94 let s3: __m256d = _mm256_loadu_pd(src_ptr.add(12));
95
96 if OVERWRITE {
97 _mm256_storeu_pd(d0_ptr, s0);
98 _mm256_storeu_pd(d1_ptr, s1);
99 _mm256_storeu_pd(d2_ptr, s2);
100 _mm256_storeu_pd(d3_ptr, s3);
101 } else {
102 let d0: __m256d = _mm256_loadu_pd(d0_ptr);
103 let d1: __m256d = _mm256_loadu_pd(d1_ptr);
104 let d2: __m256d = _mm256_loadu_pd(d2_ptr);
105 let d3: __m256d = _mm256_loadu_pd(d3_ptr);
106 _mm256_storeu_pd(d0_ptr, _mm256_add_pd(d0, s0));
107 _mm256_storeu_pd(d1_ptr, _mm256_add_pd(d1, s1));
108 _mm256_storeu_pd(d2_ptr, _mm256_add_pd(d2, s2));
109 _mm256_storeu_pd(d3_ptr, _mm256_add_pd(d3, s3));
110 }
111 }
112}
113
114/// # Safety
115/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
116#[target_feature(enable = "avx2", enable = "fma")]
117pub fn reim4_vec_mat1col_product_avx(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
118 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_loadu_pd, _mm256_setzero_pd, _mm256_storeu_pd};
119
120 #[cfg(debug_assertions)]
121 {
122 assert!(dst.len() >= 8, "dst must have at least 8 doubles");
123 assert!(u.len() >= nrows * 8, "u must be at least nrows * 8 doubles");
124 assert!(v.len() >= nrows * 8, "v must be at least nrows * 8 doubles");
125 }
126
127 unsafe {
128 use std::arch::x86_64::{_mm256_add_pd, _mm256_sub_pd};
129
130 let mut re1: __m256d = _mm256_setzero_pd();
131 let mut im1: __m256d = _mm256_setzero_pd();
132 let mut re2: __m256d = _mm256_setzero_pd();
133 let mut im2: __m256d = _mm256_setzero_pd();
134
135 let mut u_ptr: *const f64 = u.as_ptr();
136 let mut v_ptr: *const f64 = v.as_ptr();
137
138 for _ in 0..nrows {
139 let ur: __m256d = _mm256_loadu_pd(u_ptr);
140 let ui: __m256d = _mm256_loadu_pd(u_ptr.add(4));
141 let vr: __m256d = _mm256_loadu_pd(v_ptr);
142 let vi: __m256d = _mm256_loadu_pd(v_ptr.add(4));
143
144 // re1 = re1 + ur*vr;
145 re1 = _mm256_fmadd_pd(ur, vr, re1);
146 // im1 = im1 + ur*d;
147 im1 = _mm256_fmadd_pd(ur, vi, im1);
148 // re2 = re2 + ui*d;
149 re2 = _mm256_fmadd_pd(ui, vi, re2);
150 // im2 = im2 + ui*vr;
151 im2 = _mm256_fmadd_pd(ui, vr, im2);
152
153 u_ptr = u_ptr.add(8);
154 v_ptr = v_ptr.add(8);
155 }
156
157 // re1 - re2
158 _mm256_storeu_pd(dst.as_mut_ptr(), _mm256_sub_pd(re1, re2));
159
160 // im1 + im2
161 _mm256_storeu_pd(dst.as_mut_ptr().add(4), _mm256_add_pd(im1, im2));
162 }
163}
164
165/// # Safety
166/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
167#[target_feature(enable = "avx2", enable = "fma")]
168pub fn reim4_vec_mat2cols_product_avx(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
169 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_setzero_pd, _mm256_storeu_pd};
170
171 #[cfg(debug_assertions)]
172 {
173 assert!(dst.len() >= 8, "dst must be at least 8 doubles but is {}", dst.len());
174 assert!(
175 u.len() >= nrows * 8,
176 "u must be at least nrows={} * 8 doubles but is {}",
177 nrows,
178 u.len()
179 );
180 assert!(
181 v.len() >= nrows * 16,
182 "v must be at least nrows={} * 16 doubles but is {}",
183 nrows,
184 v.len()
185 );
186 }
187
188 unsafe {
189 let mut re1: __m256d = _mm256_setzero_pd();
190 let mut im1: __m256d = _mm256_setzero_pd();
191 let mut re2: __m256d = _mm256_setzero_pd();
192 let mut im2: __m256d = _mm256_setzero_pd();
193
194 let mut u_ptr: *const f64 = u.as_ptr();
195 let mut v_ptr: *const f64 = v.as_ptr();
196
197 for _ in 0..nrows {
198 let ur: __m256d = _mm256_loadu_pd(u_ptr);
199 let ui: __m256d = _mm256_loadu_pd(u_ptr.add(4));
200
201 let ar: __m256d = _mm256_loadu_pd(v_ptr);
202 let ai: __m256d = _mm256_loadu_pd(v_ptr.add(4));
203 let br: __m256d = _mm256_loadu_pd(v_ptr.add(8));
204 let bi: __m256d = _mm256_loadu_pd(v_ptr.add(12));
205
206 // re1 = ui*ai - re1; re2 = ui*bi - re2;
207 re1 = _mm256_fmsub_pd(ui, ai, re1);
208 re2 = _mm256_fmsub_pd(ui, bi, re2);
209 // im1 = ur*ai + im1; im2 = ur*bi + im2;
210 im1 = _mm256_fmadd_pd(ur, ai, im1);
211 im2 = _mm256_fmadd_pd(ur, bi, im2);
212 // re1 = ur*ar - re1; re2 = ur*br - re2;
213 re1 = _mm256_fmsub_pd(ur, ar, re1);
214 re2 = _mm256_fmsub_pd(ur, br, re2);
215 // im1 = ui*ar + im1; im2 = ui*br + im2;
216 im1 = _mm256_fmadd_pd(ui, ar, im1);
217 im2 = _mm256_fmadd_pd(ui, br, im2);
218
219 u_ptr = u_ptr.add(8);
220 v_ptr = v_ptr.add(16);
221 }
222
223 _mm256_storeu_pd(dst.as_mut_ptr(), re1);
224 _mm256_storeu_pd(dst.as_mut_ptr().add(4), im1);
225 _mm256_storeu_pd(dst.as_mut_ptr().add(8), re2);
226 _mm256_storeu_pd(dst.as_mut_ptr().add(12), im2);
227 }
228}
229
230/// # Safety
231/// Caller must ensure the CPU supports AVX2 (e.g., via `is_x86_feature_detected!("avx2")`);
232#[target_feature(enable = "avx2", enable = "fma")]
233pub fn reim4_vec_mat2cols_2ndcol_product_avx(nrows: usize, dst: &mut [f64], u: &[f64], v: &[f64]) {
234 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_setzero_pd, _mm256_storeu_pd};
235
236 #[cfg(debug_assertions)]
237 {
238 assert_eq!(dst.len(), 16, "dst must have 16 doubles");
239 assert!(u.len() >= nrows * 8, "u must be at least nrows * 8 doubles");
240 assert!(v.len() >= nrows * 16, "v must be at least nrows * 16 doubles");
241 }
242
243 unsafe {
244 let mut re1: __m256d = _mm256_setzero_pd();
245 let mut im1: __m256d = _mm256_setzero_pd();
246
247 let mut u_ptr: *const f64 = u.as_ptr();
248 let mut v_ptr: *const f64 = v.as_ptr().add(8); // Offset to 2nd column
249
250 for _ in 0..nrows {
251 let ur: __m256d = _mm256_loadu_pd(u_ptr);
252 let ui: __m256d = _mm256_loadu_pd(u_ptr.add(4));
253
254 let ar: __m256d = _mm256_loadu_pd(v_ptr);
255 let ai: __m256d = _mm256_loadu_pd(v_ptr.add(4));
256
257 // re1 = ui*ai - re1;
258 re1 = _mm256_fmsub_pd(ui, ai, re1);
259 // im1 = im1 + ur*ai;
260 im1 = _mm256_fmadd_pd(ur, ai, im1);
261 // re1 = ur*ar - re1;
262 re1 = _mm256_fmsub_pd(ur, ar, re1);
263 // im1 = im1 + ui*ar;
264 im1 = _mm256_fmadd_pd(ui, ar, im1);
265
266 u_ptr = u_ptr.add(8);
267 v_ptr = v_ptr.add(16);
268 }
269
270 _mm256_storeu_pd(dst.as_mut_ptr(), re1);
271 _mm256_storeu_pd(dst.as_mut_ptr().add(4), im1);
272 }
273}
274
275/// # Safety
276/// Caller must ensure the CPU supports AVX2 and FMA (e.g. `is_x86_feature_detected!("avx2")`).
277#[target_feature(enable = "avx2", enable = "fma")]
278pub unsafe fn reim4_convolution_1coeff_avx(k: usize, dst: &mut [f64; 8], a: &[f64], a_size: usize, b: &[f64], b_size: usize) {
279 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fmsub_pd, _mm256_loadu_pd, _mm256_setzero_pd, _mm256_storeu_pd};
280
281 unsafe {
282 // Scalar guard — same semantics as reference implementation
283 if k >= a_size + b_size {
284 let zero: __m256d = _mm256_setzero_pd();
285 let dst_ptr: *mut f64 = dst.as_mut_ptr();
286 _mm256_storeu_pd(dst_ptr, zero);
287 _mm256_storeu_pd(dst_ptr.add(4), zero);
288 return;
289 }
290
291 let j_min: usize = k.saturating_sub(a_size - 1);
292 let j_max: usize = (k + 1).min(b_size);
293
294 // acc_re = dst[0..4], acc_im = dst[4..8]
295 let mut acc_re: __m256d = _mm256_setzero_pd();
296 let mut acc_im: __m256d = _mm256_setzero_pd();
297
298 let mut a_ptr: *const f64 = a.as_ptr().add(8 * (k - j_min));
299 let mut b_ptr: *const f64 = b.as_ptr().add(8 * j_min);
300
301 for _ in 0..j_max - j_min {
302 // Load a[(k - j)]
303 let ar: __m256d = _mm256_loadu_pd(a_ptr);
304 let ai: __m256d = _mm256_loadu_pd(a_ptr.add(4));
305
306 // Load b[j]
307 let br: __m256d = _mm256_loadu_pd(b_ptr);
308 let bi: __m256d = _mm256_loadu_pd(b_ptr.add(4));
309
310 // acc_re = ai*bi - acc_re
311 acc_re = _mm256_fmsub_pd(ai, bi, acc_re);
312 // acc_im = ar*bi - acc_im
313 acc_im = _mm256_fmadd_pd(ar, bi, acc_im);
314 // acc_re = ar*br - acc_re
315 acc_re = _mm256_fmsub_pd(ar, br, acc_re);
316 // acc_im = acc_im + ai*br
317 acc_im = _mm256_fmadd_pd(ai, br, acc_im);
318
319 a_ptr = a_ptr.sub(8);
320 b_ptr = b_ptr.add(8);
321 }
322
323 // Store accumulators into dst
324 _mm256_storeu_pd(dst.as_mut_ptr(), acc_re);
325 _mm256_storeu_pd(dst.as_mut_ptr().add(4), acc_im);
326 }
327}
328
329/// # Safety
330/// Caller must ensure the CPU supports AVX2 and FMA (e.g. `is_x86_feature_detected!("avx2")`).
331#[target_feature(enable = "avx2", enable = "fma")]
332pub unsafe fn reim4_convolution_2coeffs_avx(k: usize, dst: &mut [f64; 16], a: &[f64], a_size: usize, b: &[f64], b_size: usize) {
333 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_fnmadd_pd, _mm256_loadu_pd, _mm256_setzero_pd, _mm256_storeu_pd};
334
335 debug_assert!(a.len() >= 8 * a_size);
336 debug_assert!(b.len() >= 8 * b_size);
337
338 let k0: usize = k;
339 let k1: usize = k + 1;
340 let bound: usize = a_size + b_size;
341
342 // Since k is a multiple of two, if either k0 or k1 are out of range,
343 // both are.
344 if k0 >= bound {
345 unsafe {
346 let zero: __m256d = _mm256_setzero_pd();
347 let dst_ptr: *mut f64 = dst.as_mut_ptr();
348 _mm256_storeu_pd(dst_ptr, zero);
349 _mm256_storeu_pd(dst_ptr.add(4), zero);
350 _mm256_storeu_pd(dst_ptr.add(8), zero);
351 _mm256_storeu_pd(dst_ptr.add(12), zero);
352 }
353 return;
354 }
355
356 unsafe {
357 let mut acc_re_k0: __m256d = _mm256_setzero_pd();
358 let mut acc_im_k0: __m256d = _mm256_setzero_pd();
359 let mut acc_re_k1: __m256d = _mm256_setzero_pd();
360 let mut acc_im_k1: __m256d = _mm256_setzero_pd();
361
362 let j0_min: usize = (k0 + 1).saturating_sub(a_size);
363 let j0_max: usize = (k0 + 1).min(b_size);
364
365 if k1 >= bound {
366 let mut a_k0_ptr: *const f64 = a.as_ptr().add(8 * (k0 - j0_min));
367 let mut b_ptr: *const f64 = b.as_ptr().add(8 * j0_min);
368
369 // Region 1: contributions to k0 only, j ∈ [j0_min, j1_min)
370 for _ in 0..j0_max - j0_min {
371 let ar: __m256d = _mm256_loadu_pd(a_k0_ptr);
372 let ai: __m256d = _mm256_loadu_pd(a_k0_ptr.add(4));
373 let br: __m256d = _mm256_loadu_pd(b_ptr);
374 let bi: __m256d = _mm256_loadu_pd(b_ptr.add(4));
375
376 acc_re_k0 = _mm256_fmadd_pd(ar, br, acc_re_k0);
377 acc_re_k0 = _mm256_fnmadd_pd(ai, bi, acc_re_k0);
378 acc_im_k0 = _mm256_fmadd_pd(ar, bi, acc_im_k0);
379 acc_im_k0 = _mm256_fmadd_pd(ai, br, acc_im_k0);
380
381 a_k0_ptr = a_k0_ptr.sub(8);
382 b_ptr = b_ptr.add(8);
383 }
384 } else {
385 let j1_min: usize = (k1 + 1).saturating_sub(a_size);
386 let j1_max: usize = (k1 + 1).min(b_size);
387
388 let mut a_k0_ptr: *const f64 = a.as_ptr().add(8 * (k0 - j0_min));
389 let mut a_k1_ptr: *const f64 = a.as_ptr().add(8 * (k1 - j1_min));
390 let mut b_ptr: *const f64 = b.as_ptr().add(8 * j0_min);
391
392 // Region 1: contributions to k0 only, j ∈ [j0_min, j1_min)
393 for _ in 0..j1_min - j0_min {
394 let ar: __m256d = _mm256_loadu_pd(a_k0_ptr);
395 let ai: __m256d = _mm256_loadu_pd(a_k0_ptr.add(4));
396 let br: __m256d = _mm256_loadu_pd(b_ptr);
397 let bi: __m256d = _mm256_loadu_pd(b_ptr.add(4));
398
399 acc_re_k0 = _mm256_fmadd_pd(ar, br, acc_re_k0);
400 acc_re_k0 = _mm256_fnmadd_pd(ai, bi, acc_re_k0);
401 acc_im_k0 = _mm256_fmadd_pd(ar, bi, acc_im_k0);
402 acc_im_k0 = _mm256_fmadd_pd(ai, br, acc_im_k0);
403
404 a_k0_ptr = a_k0_ptr.sub(8);
405 b_ptr = b_ptr.add(8);
406 }
407
408 // Region 2: overlap, contributions to both k0 and k1, j ∈ [j1_min, j0_max)
409 // We can save one load on b.
410 for _ in 0..j0_max - j1_min {
411 let ar0: __m256d = _mm256_loadu_pd(a_k0_ptr);
412 let ai0: __m256d = _mm256_loadu_pd(a_k0_ptr.add(4));
413 let ar1: __m256d = _mm256_loadu_pd(a_k1_ptr);
414 let ai1: __m256d = _mm256_loadu_pd(a_k1_ptr.add(4));
415 let br: __m256d = _mm256_loadu_pd(b_ptr);
416 let bi: __m256d = _mm256_loadu_pd(b_ptr.add(4));
417
418 // k0
419 acc_re_k0 = _mm256_fmadd_pd(ar0, br, acc_re_k0);
420 acc_re_k0 = _mm256_fnmadd_pd(ai0, bi, acc_re_k0);
421 acc_im_k0 = _mm256_fmadd_pd(ar0, bi, acc_im_k0);
422 acc_im_k0 = _mm256_fmadd_pd(ai0, br, acc_im_k0);
423
424 // k1
425 acc_re_k1 = _mm256_fmadd_pd(ar1, br, acc_re_k1);
426 acc_re_k1 = _mm256_fnmadd_pd(ai1, bi, acc_re_k1);
427 acc_im_k1 = _mm256_fmadd_pd(ar1, bi, acc_im_k1);
428 acc_im_k1 = _mm256_fmadd_pd(ai1, br, acc_im_k1);
429
430 a_k0_ptr = a_k0_ptr.sub(8);
431 a_k1_ptr = a_k1_ptr.sub(8);
432 b_ptr = b_ptr.add(8);
433 }
434
435 // Region 3: contributions to k1 only, j ∈ [j0_max, j1_max)
436 for _ in 0..j1_max - j0_max {
437 let ar1: __m256d = _mm256_loadu_pd(a_k1_ptr);
438 let ai1: __m256d = _mm256_loadu_pd(a_k1_ptr.add(4));
439 let br: __m256d = _mm256_loadu_pd(b_ptr);
440 let bi: __m256d = _mm256_loadu_pd(b_ptr.add(4));
441
442 acc_re_k1 = _mm256_fmadd_pd(ar1, br, acc_re_k1);
443 acc_re_k1 = _mm256_fnmadd_pd(ai1, bi, acc_re_k1);
444 acc_im_k1 = _mm256_fmadd_pd(ar1, bi, acc_im_k1);
445 acc_im_k1 = _mm256_fmadd_pd(ai1, br, acc_im_k1);
446
447 a_k1_ptr = a_k1_ptr.sub(8);
448 b_ptr = b_ptr.add(8);
449 }
450 }
451
452 // Store both coefficients
453 let dst_ptr = dst.as_mut_ptr();
454 _mm256_storeu_pd(dst_ptr, acc_re_k0);
455 _mm256_storeu_pd(dst_ptr.add(4), acc_im_k0);
456 _mm256_storeu_pd(dst_ptr.add(8), acc_re_k1);
457 _mm256_storeu_pd(dst_ptr.add(12), acc_im_k1);
458 }
459}
460
461/// # Safety
462/// Caller must ensure the CPU supports AVX2 and FMA (e.g. `is_x86_feature_detected!("avx2")`).
463#[target_feature(enable = "avx2", enable = "fma")]
464pub unsafe fn reim4_convolution_by_real_const_1coeff_avx(k: usize, dst: &mut [f64; 8], a: &[f64], a_size: usize, b: &[f64]) {
465 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_loadu_pd, _mm256_set1_pd, _mm256_setzero_pd, _mm256_storeu_pd};
466
467 unsafe {
468 let b_size: usize = b.len();
469
470 if k >= a_size + b_size {
471 let zero: __m256d = _mm256_setzero_pd();
472 let dst_ptr: *mut f64 = dst.as_mut_ptr();
473 _mm256_storeu_pd(dst_ptr, zero);
474 _mm256_storeu_pd(dst_ptr.add(4), zero);
475 return;
476 }
477
478 let j_min: usize = k.saturating_sub(a_size - 1);
479 let j_max: usize = (k + 1).min(b_size);
480
481 // acc_re = dst[0..4], acc_im = dst[4..8]
482 let mut acc_re: __m256d = _mm256_setzero_pd();
483 let mut acc_im: __m256d = _mm256_setzero_pd();
484
485 let mut a_ptr: *const f64 = a.as_ptr().add(8 * (k - j_min));
486 let mut b_ptr: *const f64 = b.as_ptr().add(j_min);
487
488 for _ in 0..j_max - j_min {
489 // Load a[(k - j)]
490 let ar: __m256d = _mm256_loadu_pd(a_ptr);
491 let ai: __m256d = _mm256_loadu_pd(a_ptr.add(4));
492
493 // Load scalar b[j] and broadcast
494 let br: __m256d = _mm256_set1_pd(*b_ptr);
495
496 // Complex * real:
497 // re += ar * br
498 // im += ai * br
499 acc_re = _mm256_fmadd_pd(ar, br, acc_re);
500 acc_im = _mm256_fmadd_pd(ai, br, acc_im);
501
502 a_ptr = a_ptr.sub(8);
503 b_ptr = b_ptr.add(1);
504 }
505
506 // Store accumulators into dst
507 _mm256_storeu_pd(dst.as_mut_ptr(), acc_re);
508 _mm256_storeu_pd(dst.as_mut_ptr().add(4), acc_im);
509 }
510}
511
512/// # Safety
513/// Caller must ensure the CPU supports AVX2 and FMA (e.g. `is_x86_feature_detected!("avx2")`).
514#[target_feature(enable = "avx2", enable = "fma")]
515pub unsafe fn reim4_convolution_by_real_const_2coeffs_avx(k: usize, dst: &mut [f64; 16], a: &[f64], a_size: usize, b: &[f64]) {
516 use core::arch::x86_64::{__m256d, _mm256_fmadd_pd, _mm256_loadu_pd, _mm256_set1_pd, _mm256_setzero_pd, _mm256_storeu_pd};
517
518 let b_size: usize = b.len();
519
520 debug_assert!(a.len() >= 8 * a_size);
521
522 let k0: usize = k;
523 let k1: usize = k + 1;
524 let bound: usize = a_size + b_size;
525
526 // Since k is a multiple of two, if either k0 or k1 are out of range,
527 // both are.
528 if k0 >= bound {
529 unsafe {
530 let zero: __m256d = _mm256_setzero_pd();
531 let dst_ptr: *mut f64 = dst.as_mut_ptr();
532 _mm256_storeu_pd(dst_ptr, zero);
533 _mm256_storeu_pd(dst_ptr.add(4), zero);
534 _mm256_storeu_pd(dst_ptr.add(8), zero);
535 _mm256_storeu_pd(dst_ptr.add(12), zero);
536 }
537 return;
538 }
539
540 unsafe {
541 let mut acc_re_k0: __m256d = _mm256_setzero_pd();
542 let mut acc_im_k0: __m256d = _mm256_setzero_pd();
543 let mut acc_re_k1: __m256d = _mm256_setzero_pd();
544 let mut acc_im_k1: __m256d = _mm256_setzero_pd();
545
546 let j0_min: usize = (k0 + 1).saturating_sub(a_size);
547 let j0_max: usize = (k0 + 1).min(b_size);
548
549 if k1 >= bound {
550 let mut a_k0_ptr: *const f64 = a.as_ptr().add(8 * (k0 - j0_min));
551 let mut b_ptr: *const f64 = b.as_ptr().add(j0_min);
552
553 // Contributions to k0 only
554 for _ in 0..j0_max - j0_min {
555 let ar: __m256d = _mm256_loadu_pd(a_k0_ptr);
556 let ai: __m256d = _mm256_loadu_pd(a_k0_ptr.add(4));
557 let br: __m256d = _mm256_set1_pd(*b_ptr);
558
559 // complex * real
560 acc_re_k0 = _mm256_fmadd_pd(ar, br, acc_re_k0);
561 acc_im_k0 = _mm256_fmadd_pd(ai, br, acc_im_k0);
562
563 a_k0_ptr = a_k0_ptr.sub(8);
564 b_ptr = b_ptr.add(1);
565 }
566 } else {
567 let j1_min: usize = (k1 + 1).saturating_sub(a_size);
568 let j1_max: usize = (k1 + 1).min(b_size);
569
570 let mut a_k0_ptr: *const f64 = a.as_ptr().add(8 * (k0 - j0_min));
571 let mut a_k1_ptr: *const f64 = a.as_ptr().add(8 * (k1 - j1_min));
572 let mut b_ptr: *const f64 = b.as_ptr().add(j0_min);
573
574 // Region 1: k0 only, j ∈ [j0_min, j1_min)
575 for _ in 0..j1_min - j0_min {
576 let ar0: __m256d = _mm256_loadu_pd(a_k0_ptr);
577 let ai0: __m256d = _mm256_loadu_pd(a_k0_ptr.add(4));
578 let br: __m256d = _mm256_set1_pd(*b_ptr);
579
580 acc_re_k0 = _mm256_fmadd_pd(ar0, br, acc_re_k0);
581 acc_im_k0 = _mm256_fmadd_pd(ai0, br, acc_im_k0);
582
583 a_k0_ptr = a_k0_ptr.sub(8);
584 b_ptr = b_ptr.add(1);
585 }
586
587 // Region 2: overlap, contributions to both k0 and k1, j ∈ [j1_min, j0_max)
588 // Still “save one load on b”: we broadcast once and reuse.
589 for _ in 0..j0_max - j1_min {
590 let ar0: __m256d = _mm256_loadu_pd(a_k0_ptr);
591 let ai0: __m256d = _mm256_loadu_pd(a_k0_ptr.add(4));
592 let ar1: __m256d = _mm256_loadu_pd(a_k1_ptr);
593 let ai1: __m256d = _mm256_loadu_pd(a_k1_ptr.add(4));
594 let br: __m256d = _mm256_set1_pd(*b_ptr);
595
596 // k0
597 acc_re_k0 = _mm256_fmadd_pd(ar0, br, acc_re_k0);
598 acc_im_k0 = _mm256_fmadd_pd(ai0, br, acc_im_k0);
599
600 // k1
601 acc_re_k1 = _mm256_fmadd_pd(ar1, br, acc_re_k1);
602 acc_im_k1 = _mm256_fmadd_pd(ai1, br, acc_im_k1);
603
604 a_k0_ptr = a_k0_ptr.sub(8);
605 a_k1_ptr = a_k1_ptr.sub(8);
606 b_ptr = b_ptr.add(1);
607 }
608
609 // Region 3: k1 only, j ∈ [j0_max, j1_max)
610 for _ in 0..j1_max - j0_max {
611 let ar1: __m256d = _mm256_loadu_pd(a_k1_ptr);
612 let ai1: __m256d = _mm256_loadu_pd(a_k1_ptr.add(4));
613 let br: __m256d = _mm256_set1_pd(*b_ptr);
614
615 acc_re_k1 = _mm256_fmadd_pd(ar1, br, acc_re_k1);
616 acc_im_k1 = _mm256_fmadd_pd(ai1, br, acc_im_k1);
617
618 a_k1_ptr = a_k1_ptr.sub(8);
619 b_ptr = b_ptr.add(1);
620 }
621 }
622
623 // Store both coefficients
624 let dst_ptr = dst.as_mut_ptr();
625 _mm256_storeu_pd(dst_ptr, acc_re_k0);
626 _mm256_storeu_pd(dst_ptr.add(4), acc_im_k0);
627 _mm256_storeu_pd(dst_ptr.add(8), acc_re_k1);
628 _mm256_storeu_pd(dst_ptr.add(12), acc_im_k1);
629 }
630}