1use std::arch::x86_64::*;
2use std::marker::PhantomData;
3use std::mem::MaybeUninit;
4
5use num_complex::Complex;
6
7use crate::array_utils::DoubleBuf;
8use crate::{common::FftNum, twiddles};
9use crate::{Direction, Fft, FftDirection, Length};
10
11use super::avx32_utils;
12use super::avx_vector::{self, AvxArray};
13use super::avx_vector::{AvxArrayMut, AvxVector, AvxVector128, AvxVector256, Rotation90};
14
15macro_rules! boilerplate_fft_simd_butterfly {
19 ($struct_name:ident, $len:expr) => {
20 impl $struct_name<f32> {
21 #[inline]
22 pub fn is_supported_by_cpu() -> bool {
23 is_x86_feature_detected!("avx") && is_x86_feature_detected!("fma")
24 }
25 #[inline]
26 pub fn new(direction: FftDirection) -> Result<Self, ()> {
27 if Self::is_supported_by_cpu() {
28 Ok(unsafe { Self::new_with_avx(direction) })
30 } else {
31 Err(())
32 }
33 }
34 }
35
36 impl<T: FftNum> Fft<T> for $struct_name<f32> {
37 fn process_immutable_with_scratch(
38 &self,
39 input: &[Complex<T>],
40 output: &mut [Complex<T>],
41 _scratch: &mut [Complex<T>],
42 ) {
43 unsafe {
44 let simd_input = crate::array_utils::workaround_transmute(input);
45 let simd_output = crate::array_utils::workaround_transmute_mut(output);
46 super::avx_fft_helper_immut(
47 simd_input,
48 simd_output,
49 &mut [],
50 self.len(),
51 0,
52 |input, output, _| self.perform_fft_f32(DoubleBuf { input, output }),
53 );
54 }
55 }
56 fn process_outofplace_with_scratch(
57 &self,
58 input: &mut [Complex<T>],
59 output: &mut [Complex<T>],
60 _scratch: &mut [Complex<T>],
61 ) {
62 unsafe {
63 let simd_input = crate::array_utils::workaround_transmute_mut(input);
64 let simd_output = crate::array_utils::workaround_transmute_mut(output);
65 super::avx_fft_helper_outofplace(
66 simd_input,
67 simd_output,
68 &mut [],
69 self.len(),
70 0,
71 |input, output, _| self.perform_fft_f32(DoubleBuf { input, output }),
72 );
73 }
74 }
75 fn process_with_scratch(&self, buffer: &mut [Complex<T>], _scratch: &mut [Complex<T>]) {
76 unsafe {
77 let simd_buffer = crate::array_utils::workaround_transmute_mut(buffer);
78 super::avx_fft_helper_inplace(
79 simd_buffer,
80 &mut [],
81 self.len(),
82 0,
83 |chunk, _| self.perform_fft_f32(chunk),
84 )
85 }
86 }
87 #[inline(always)]
88 fn get_inplace_scratch_len(&self) -> usize {
89 0
90 }
91 #[inline(always)]
92 fn get_outofplace_scratch_len(&self) -> usize {
93 0
94 }
95 #[inline(always)]
96 fn get_immutable_scratch_len(&self) -> usize {
97 0
98 }
99 }
100 impl<T> Length for $struct_name<T> {
101 #[inline(always)]
102 fn len(&self) -> usize {
103 $len
104 }
105 }
106 impl<T> Direction for $struct_name<T> {
107 #[inline(always)]
108 fn fft_direction(&self) -> FftDirection {
109 self.direction
110 }
111 }
112 };
113}
114
115macro_rules! boilerplate_fft_simd_butterfly_with_scratch {
119 ($struct_name:ident, $len:expr) => {
120 impl $struct_name<f32> {
121 #[inline]
122 pub fn new(direction: FftDirection) -> Result<Self, ()> {
123 let has_avx = is_x86_feature_detected!("avx");
124 let has_fma = is_x86_feature_detected!("fma");
125 if has_avx && has_fma {
126 Ok(unsafe { Self::new_with_avx(direction) })
128 } else {
129 Err(())
130 }
131 }
132 }
133 impl<T: FftNum> $struct_name<T> {
134 #[inline]
135 fn perform_fft_inplace(
136 &self,
137 buffer: &mut [Complex<f32>],
138 scratch: &mut [Complex<f32>],
139 ) {
140 unsafe { self.column_butterflies_and_transpose(buffer, scratch) };
143
144 unsafe {
147 self.row_butterflies(DoubleBuf {
148 input: scratch,
149 output: buffer,
150 })
151 };
152 }
153
154 #[inline]
155 fn perform_fft_immut(&self, input: &[Complex<f32>], output: &mut [Complex<f32>]) {
156 unsafe { self.column_butterflies_and_transpose(input, output) };
159
160 unsafe { self.row_butterflies(output) };
163 }
164 }
165 impl<T: FftNum> Fft<T> for $struct_name<f32> {
166 fn process_immutable_with_scratch(
167 &self,
168 input: &[Complex<T>],
169 output: &mut [Complex<T>],
170 _scratch: &mut [Complex<T>],
171 ) {
172 unsafe {
173 let simd_input = crate::array_utils::workaround_transmute(input);
174 let simd_output = crate::array_utils::workaround_transmute_mut(output);
175 super::avx_fft_helper_immut(
176 simd_input,
177 simd_output,
178 &mut [],
179 self.len(),
180 0,
181 |input, output, _| self.perform_fft_immut(input, output),
182 );
183 }
184 }
185 fn process_outofplace_with_scratch(
186 &self,
187 input: &mut [Complex<T>],
188 output: &mut [Complex<T>],
189 _scratch: &mut [Complex<T>],
190 ) {
191 unsafe {
192 let simd_input = crate::array_utils::workaround_transmute_mut(input);
193 let simd_output = crate::array_utils::workaround_transmute_mut(output);
194 super::avx_fft_helper_outofplace(
195 simd_input,
196 simd_output,
197 &mut [],
198 self.len(),
199 0,
200 |input, output, _| self.perform_fft_immut(input, output),
201 );
202 }
203 }
204 fn process_with_scratch(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
205 unsafe {
206 let simd_buffer = crate::array_utils::workaround_transmute_mut(buffer);
207 let simd_scratch = crate::array_utils::workaround_transmute_mut(scratch);
208 super::avx_fft_helper_inplace(
209 simd_buffer,
210 simd_scratch,
211 self.len(),
212 self.len(),
213 |chunk, scratch| self.perform_fft_inplace(chunk, scratch),
214 )
215 }
216 }
217 #[inline(always)]
218 fn get_inplace_scratch_len(&self) -> usize {
219 $len
220 }
221 #[inline(always)]
222 fn get_outofplace_scratch_len(&self) -> usize {
223 0
224 }
225 #[inline(always)]
226 fn get_immutable_scratch_len(&self) -> usize {
227 0
228 }
229 }
230 impl<T> Length for $struct_name<T> {
231 #[inline(always)]
232 fn len(&self) -> usize {
233 $len
234 }
235 }
236 impl<T> Direction for $struct_name<T> {
237 #[inline(always)]
238 fn fft_direction(&self) -> FftDirection {
239 self.direction
240 }
241 }
242 };
243}
244
245macro_rules! gen_butterfly_twiddles_interleaved_columns {
246 ($num_rows:expr, $num_cols:expr, $skip_cols:expr, $direction: expr) => {{
247 const FFT_LEN: usize = $num_rows * $num_cols;
248 const TWIDDLE_ROWS: usize = $num_rows - 1;
249 const TWIDDLE_COLS: usize = $num_cols - $skip_cols;
250 const TWIDDLE_VECTOR_COLS: usize = TWIDDLE_COLS / 4;
251 const TWIDDLE_VECTOR_COUNT: usize = TWIDDLE_VECTOR_COLS * TWIDDLE_ROWS;
252 let mut twiddles = [AvxVector::zero(); TWIDDLE_VECTOR_COUNT];
253 for index in 0..TWIDDLE_VECTOR_COUNT {
254 let y = (index / TWIDDLE_VECTOR_COLS) + 1;
255 let x = (index % TWIDDLE_VECTOR_COLS) * 4 + $skip_cols;
256
257 twiddles[index] = AvxVector::make_mixedradix_twiddle_chunk(x, y, FFT_LEN, $direction);
258 }
259 twiddles
260 }};
261}
262
263macro_rules! gen_butterfly_twiddles_separated_columns {
264 ($num_rows:expr, $num_cols:expr, $skip_cols:expr, $direction: expr) => {{
265 const FFT_LEN: usize = $num_rows * $num_cols;
266 const TWIDDLE_ROWS: usize = $num_rows - 1;
267 const TWIDDLE_COLS: usize = $num_cols - $skip_cols;
268 const TWIDDLE_VECTOR_COLS: usize = TWIDDLE_COLS / 4;
269 const TWIDDLE_VECTOR_COUNT: usize = TWIDDLE_VECTOR_COLS * TWIDDLE_ROWS;
270 let mut twiddles = [AvxVector::zero(); TWIDDLE_VECTOR_COUNT];
271 for index in 0..TWIDDLE_VECTOR_COUNT {
272 let y = (index % TWIDDLE_ROWS) + 1;
273 let x = (index / TWIDDLE_ROWS) * 4 + $skip_cols;
274
275 twiddles[index] = AvxVector::make_mixedradix_twiddle_chunk(x, y, FFT_LEN, $direction);
276 }
277 twiddles
278 }};
279}
280
281pub struct Butterfly5Avx<T> {
282 twiddles: [__m128; 3],
283 direction: FftDirection,
284 _phantom_t: std::marker::PhantomData<T>,
285}
286boilerplate_fft_simd_butterfly!(Butterfly5Avx, 5);
287impl Butterfly5Avx<f32> {
288 #[target_feature(enable = "avx")]
289 unsafe fn new_with_avx(direction: FftDirection) -> Self {
290 let twiddle1 = twiddles::compute_twiddle(1, 5, direction);
291 let twiddle2 = twiddles::compute_twiddle(2, 5, direction);
292 Self {
293 twiddles: [
294 _mm_set_ps(twiddle1.im, twiddle1.im, twiddle1.re, twiddle1.re),
295 _mm_set_ps(twiddle2.im, twiddle2.im, twiddle2.re, twiddle2.re),
296 _mm_set_ps(-twiddle1.im, -twiddle1.im, twiddle1.re, twiddle1.re),
297 ],
298 direction,
299 _phantom_t: PhantomData,
300 }
301 }
302}
303impl<T> Butterfly5Avx<T> {
304 #[target_feature(enable = "avx", enable = "fma")]
305 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
306 let input0 = _mm_castpd_ps(_mm_load1_pd(buffer.input_ptr() as *const f64)); let input12 = buffer.load_partial2_complex(1);
308 let input34 = buffer.load_partial2_complex(3);
309
310 let input43 = AvxVector::reverse_complex_elements(input34);
312
313 let [sum12, diff43] = AvxVector::column_butterfly2([input12, input43]);
315
316 let rotation = AvxVector::make_rotation90(FftDirection::Inverse);
317 let rotated43 = AvxVector::rotate90(diff43, rotation);
318
319 let [mid14, mid23] = AvxVector::unpack_complex([sum12, rotated43]);
320
321 let sum1234 = AvxVector::add(mid14, mid23);
323 let output0 = AvxVector::add(input0, sum1234);
324
325 let twiddled14_mid = AvxVector::mul(mid14, self.twiddles[0]);
327 let twiddled23_mid = AvxVector::mul(mid14, self.twiddles[1]);
328 let twiddled14 = AvxVector::fmadd(mid23, self.twiddles[1], twiddled14_mid);
329 let twiddled23 = AvxVector::fmadd(mid23, self.twiddles[2], twiddled23_mid);
330
331 let [twiddled12, twiddled43] = AvxVector::unpack_complex([twiddled14, twiddled23]);
333 let [output12, output43] = AvxVector::column_butterfly2([twiddled12, twiddled43]);
334
335 let final12 = AvxVector::add(input0, output12);
337 let output34 = AvxVector::reverse_complex_elements(output43);
338 let final34 = AvxVector::add(input0, output34);
339
340 buffer.store_partial1_complex(output0, 0);
341 buffer.store_partial2_complex(final12, 1);
342 buffer.store_partial2_complex(final34, 3);
343 }
344}
345
346pub struct Butterfly7Avx<T> {
347 twiddles: [__m128; 5],
348 direction: FftDirection,
349 _phantom_t: std::marker::PhantomData<T>,
350}
351boilerplate_fft_simd_butterfly!(Butterfly7Avx, 7);
352impl Butterfly7Avx<f32> {
353 #[target_feature(enable = "avx")]
354 unsafe fn new_with_avx(direction: FftDirection) -> Self {
355 let twiddle1 = twiddles::compute_twiddle(1, 7, direction);
356 let twiddle2 = twiddles::compute_twiddle(2, 7, direction);
357 let twiddle3 = twiddles::compute_twiddle(3, 7, direction);
358 Self {
359 twiddles: [
360 _mm_set_ps(twiddle1.im, twiddle1.im, twiddle1.re, twiddle1.re),
361 _mm_set_ps(twiddle2.im, twiddle2.im, twiddle2.re, twiddle2.re),
362 _mm_set_ps(twiddle3.im, twiddle3.im, twiddle3.re, twiddle3.re),
363 _mm_set_ps(-twiddle3.im, -twiddle3.im, twiddle3.re, twiddle3.re),
364 _mm_set_ps(-twiddle1.im, -twiddle1.im, twiddle1.re, twiddle1.re),
365 ],
366 direction,
367 _phantom_t: PhantomData,
368 }
369 }
370}
371impl<T> Butterfly7Avx<T> {
372 #[target_feature(enable = "avx", enable = "fma")]
373 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
374 let input0 = _mm_castpd_ps(_mm_load1_pd(buffer.input_ptr() as *const f64));
376
377 let input123 = buffer.load_complex(1);
381 let input456 = buffer.load_complex(3);
382
383 let input654 = AvxVector::reverse_complex_elements(input456);
385
386 let [sum123, diff654] = AvxVector::column_butterfly2([input123, input654]);
388 let rotation = AvxVector::make_rotation90(FftDirection::Inverse);
389 let rotated654 = AvxVector::rotate90(diff654, rotation);
390
391 let [mid1634, mid25] = AvxVector::unpack_complex([sum123, rotated654]);
392
393 let mid16 = mid1634.lo();
394 let mid25 = mid25.lo();
395 let mid34 = mid1634.hi();
396
397 let output0_left = AvxVector::add(mid16, mid25);
399 let output0_right = AvxVector::add(input0, mid34);
400 let output0 = AvxVector::add(output0_left, output0_right);
401 buffer.store_partial1_complex(output0, 0);
402
403 _mm256_zeroupper();
404
405 let twiddled16_intermediate1 = AvxVector::mul(mid16, self.twiddles[0]);
407 let twiddled25_intermediate1 = AvxVector::mul(mid16, self.twiddles[1]);
408 let twiddled34_intermediate1 = AvxVector::mul(mid16, self.twiddles[2]);
409
410 let twiddled16_intermediate2 =
411 AvxVector::fmadd(mid25, self.twiddles[1], twiddled16_intermediate1);
412 let twiddled25_intermediate2 =
413 AvxVector::fmadd(mid25, self.twiddles[3], twiddled25_intermediate1);
414 let twiddled34_intermediate2 =
415 AvxVector::fmadd(mid25, self.twiddles[4], twiddled34_intermediate1);
416
417 let twiddled16 = AvxVector::fmadd(mid34, self.twiddles[2], twiddled16_intermediate2);
418 let twiddled25 = AvxVector::fmadd(mid34, self.twiddles[4], twiddled25_intermediate2);
419 let twiddled34 = AvxVector::fmadd(mid34, self.twiddles[1], twiddled34_intermediate2);
420
421 let [twiddled12, twiddled65] = AvxVector::unpack_complex([twiddled16, twiddled25]);
423 let [twiddled33, twiddled44] = AvxVector::unpack_complex([twiddled34, twiddled34]);
424
425 let twiddled033 = AvxVector::add(twiddled33, input0);
427
428 let [output12, output65] = AvxVector::column_butterfly2([twiddled12, twiddled65]);
429 let [output033, output044] = AvxVector::column_butterfly2([twiddled033, twiddled44]);
430 let output56 = AvxVector::reverse_complex_elements(output65);
431
432 buffer.store_partial2_complex(AvxVector::add(output12, input0), 1);
433 buffer.store_partial1_complex(output033, 3);
434 buffer.store_partial1_complex(output044, 4);
435 buffer.store_partial2_complex(AvxVector::add(output56, input0), 5);
436 }
437}
438
439pub struct Butterfly11Avx<T> {
440 twiddles: [__m256; 10],
441 twiddle_lo_4: __m128,
442 twiddle_lo_9: __m128,
443 twiddle_lo_3: __m128,
444 twiddle_lo_8: __m128,
445 twiddle_lo_2: __m128,
446 direction: FftDirection,
447 _phantom_t: std::marker::PhantomData<T>,
448}
449boilerplate_fft_simd_butterfly!(Butterfly11Avx, 11);
450impl Butterfly11Avx<f32> {
451 #[target_feature(enable = "avx")]
452 unsafe fn new_with_avx(direction: FftDirection) -> Self {
453 let twiddle1 = twiddles::compute_twiddle(1, 11, direction);
454 let twiddle2 = twiddles::compute_twiddle(2, 11, direction);
455 let twiddle3 = twiddles::compute_twiddle(3, 11, direction);
456 let twiddle4 = twiddles::compute_twiddle(4, 11, direction);
457 let twiddle5 = twiddles::compute_twiddle(5, 11, direction);
458
459 let twiddles_lo = [
460 _mm_set_ps(twiddle1.im, twiddle1.im, twiddle1.re, twiddle1.re),
461 _mm_set_ps(twiddle2.im, twiddle2.im, twiddle2.re, twiddle2.re),
462 _mm_set_ps(twiddle3.im, twiddle3.im, twiddle3.re, twiddle3.re),
463 _mm_set_ps(twiddle4.im, twiddle4.im, twiddle4.re, twiddle4.re),
464 _mm_set_ps(twiddle5.im, twiddle5.im, twiddle5.re, twiddle5.re),
465 _mm_set_ps(-twiddle5.im, -twiddle5.im, twiddle5.re, twiddle5.re),
466 _mm_set_ps(-twiddle4.im, -twiddle4.im, twiddle4.re, twiddle4.re),
467 _mm_set_ps(-twiddle3.im, -twiddle3.im, twiddle3.re, twiddle3.re),
468 _mm_set_ps(-twiddle2.im, -twiddle2.im, twiddle2.re, twiddle2.re),
469 _mm_set_ps(-twiddle1.im, -twiddle1.im, twiddle1.re, twiddle1.re),
470 ];
471
472 Self {
473 twiddles: [
474 AvxVector256::merge(twiddles_lo[0], twiddles_lo[2]),
475 AvxVector256::merge(twiddles_lo[1], twiddles_lo[3]),
476 AvxVector256::merge(twiddles_lo[1], twiddles_lo[5]),
477 AvxVector256::merge(twiddles_lo[3], twiddles_lo[7]),
478 AvxVector256::merge(twiddles_lo[2], twiddles_lo[8]),
479 AvxVector256::merge(twiddles_lo[5], twiddles_lo[0]),
480 AvxVector256::merge(twiddles_lo[3], twiddles_lo[0]),
481 AvxVector256::merge(twiddles_lo[7], twiddles_lo[4]),
482 AvxVector256::merge(twiddles_lo[4], twiddles_lo[3]),
483 AvxVector256::merge(twiddles_lo[9], twiddles_lo[8]),
484 ],
485 twiddle_lo_4: twiddles_lo[4],
486 twiddle_lo_9: twiddles_lo[9],
487 twiddle_lo_3: twiddles_lo[3],
488 twiddle_lo_8: twiddles_lo[8],
489 twiddle_lo_2: twiddles_lo[2],
490 direction,
491 _phantom_t: PhantomData,
492 }
493 }
494}
495impl<T> Butterfly11Avx<T> {
496 #[target_feature(enable = "avx", enable = "fma")]
497 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
498 let input0 = _mm_castpd_ps(_mm_load1_pd(buffer.input_ptr() as *const f64)); let input1234 = buffer.load_complex(1);
500 let input56 = buffer.load_partial2_complex(5);
501 let input78910 = buffer.load_complex(7);
502
503 let [input55, input66] = AvxVector::unpack_complex([input56, input56]);
505 let input10987 = AvxVector::reverse_complex_elements(input78910);
506
507 let [sum1234, diff10987] = AvxVector::column_butterfly2([input1234, input10987]);
509 let [sum55, diff66] = AvxVector::column_butterfly2([input55, input66]);
510
511 let rotation = AvxVector::make_rotation90(FftDirection::Inverse);
512 let rotated10987 = AvxVector::rotate90(diff10987, rotation);
513 let rotated66 = AvxVector::rotate90(diff66, rotation.lo());
514
515 let [mid11038, mid2947] = AvxVector::unpack_complex([sum1234, rotated10987]);
517
518 let mid110: __m256 = AvxVector256::merge(mid11038.lo(), mid11038.lo());
519 let mid29: __m256 = AvxVector256::merge(mid2947.lo(), mid2947.lo());
520 let mid38: __m256 = AvxVector256::merge(mid11038.hi(), mid11038.hi());
521 let mid47: __m256 = AvxVector256::merge(mid2947.hi(), mid2947.hi());
522 let mid56 = AvxVector::unpacklo_complex([sum55, rotated66]);
523 let mid56: __m256 = AvxVector256::merge(mid56, mid56);
524
525 let mid12910 = AvxVector::add(mid110.lo(), mid29.lo());
527 let mid3478 = AvxVector::add(mid38.lo(), mid47.lo());
528 let output0_left = AvxVector::add(input0, mid56.lo());
529 let output0_right = AvxVector::add(mid12910, mid3478);
530 let output0 = AvxVector::add(output0_left, output0_right);
531 buffer.store_partial1_complex(output0, 0);
532
533 let zero = _mm_setzero_pd();
536 let input0 = _mm_castpd_ps(_mm_move_sd(zero, _mm_castps_pd(input0)));
537 let input0 = AvxVector256::merge(input0, input0);
538
539 let twiddled11038 = AvxVector::fmadd(mid110, self.twiddles[0], input0);
541 let twiddled2947 = AvxVector::fmadd(mid110, self.twiddles[1], input0);
542 let twiddled56 = AvxVector::fmadd(mid110.lo(), self.twiddle_lo_4, input0.lo());
543
544 let twiddled11038 = AvxVector::fmadd(mid29, self.twiddles[2], twiddled11038);
545 let twiddled2947 = AvxVector::fmadd(mid29, self.twiddles[3], twiddled2947);
546 let twiddled56 = AvxVector::fmadd(mid29.lo(), self.twiddle_lo_9, twiddled56);
547
548 let twiddled11038 = AvxVector::fmadd(mid38, self.twiddles[4], twiddled11038);
549 let twiddled2947 = AvxVector::fmadd(mid38, self.twiddles[5], twiddled2947);
550 let twiddled56 = AvxVector::fmadd(mid38.lo(), self.twiddle_lo_3, twiddled56);
551
552 let twiddled11038 = AvxVector::fmadd(mid47, self.twiddles[6], twiddled11038);
553 let twiddled2947 = AvxVector::fmadd(mid47, self.twiddles[7], twiddled2947);
554 let twiddled56 = AvxVector::fmadd(mid47.lo(), self.twiddle_lo_8, twiddled56);
555
556 let twiddled11038 = AvxVector::fmadd(mid56, self.twiddles[8], twiddled11038);
557 let twiddled2947 = AvxVector::fmadd(mid56, self.twiddles[9], twiddled2947);
558 let twiddled56 = AvxVector::fmadd(mid56.lo(), self.twiddle_lo_2, twiddled56);
559
560 let [twiddled1234, twiddled10987] =
562 AvxVector::unpack_complex([twiddled11038, twiddled2947]);
563 let [twiddled55, twiddled66] = AvxVector::unpack_complex([twiddled56, twiddled56]);
564
565 let [output1234, output10987] = AvxVector::column_butterfly2([twiddled1234, twiddled10987]);
566 let [output55, output66] = AvxVector::column_butterfly2([twiddled55, twiddled66]);
567 let output78910 = AvxVector::reverse_complex_elements(output10987);
568
569 buffer.store_complex(output1234, 1);
570 buffer.store_partial1_complex(output55, 5);
571 buffer.store_partial1_complex(output66, 6);
572 buffer.store_complex(output78910, 7);
573 }
574}
575
576pub struct Butterfly8Avx<T> {
577 twiddles: __m256,
578 twiddles_butterfly4: __m256,
579 direction: FftDirection,
580 _phantom_t: std::marker::PhantomData<T>,
581}
582boilerplate_fft_simd_butterfly!(Butterfly8Avx, 8);
583impl Butterfly8Avx<f32> {
584 #[target_feature(enable = "avx")]
585 unsafe fn new_with_avx(direction: FftDirection) -> Self {
586 Self {
587 twiddles: AvxVector::make_mixedradix_twiddle_chunk(0, 1, 8, direction),
588 twiddles_butterfly4: match direction {
589 FftDirection::Forward => [
590 Complex::new(0.0f32, 0.0),
591 Complex::new(0.0, -0.0),
592 Complex::new(0.0, 0.0),
593 Complex::new(0.0, -0.0),
594 ]
595 .as_slice()
596 .load_complex(0),
597 FftDirection::Inverse => [
598 Complex::new(0.0f32, 0.0),
599 Complex::new(-0.0, 0.0),
600 Complex::new(0.0, 0.0),
601 Complex::new(-0.0, 0.0),
602 ]
603 .as_slice()
604 .load_complex(0),
605 },
606 direction,
607 _phantom_t: PhantomData,
608 }
609 }
610}
611impl<T> Butterfly8Avx<T> {
612 #[target_feature(enable = "avx", enable = "fma")]
613 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
614 let row0 = buffer.load_complex(0);
615 let row1 = buffer.load_complex(4);
616
617 let [intermediate0, intermediate1_pretwiddle] = AvxVector::column_butterfly2([row0, row1]);
619
620 let intermediate1 = AvxVector::mul_complex(intermediate1_pretwiddle, self.twiddles);
622
623 let permuted0 = _mm256_permute2f128_ps(intermediate0, intermediate1, 0x20);
626 let permuted1 = _mm256_permute2f128_ps(intermediate0, intermediate1, 0x31);
627
628 let [postbutterfly0, postbutterfly1_pretwiddle] =
630 AvxVector::column_butterfly2([permuted0, permuted1]);
631
632 let elements_swapped = _mm256_permute_ps(postbutterfly1_pretwiddle, 0xB4);
636
637 let postbutterfly1 = AvxVector::xor(elements_swapped, self.twiddles_butterfly4);
639
640 let unpermuted0 = _mm256_permute2f128_ps(postbutterfly0, postbutterfly1, 0x20);
642 let unpermuted1 = _mm256_permute2f128_ps(postbutterfly0, postbutterfly1, 0x31);
643 let unpacked = AvxVector::unpack_complex([unpermuted0, unpermuted1]);
644
645 let [output0, output1] = AvxVector::column_butterfly2(unpacked);
646
647 buffer.store_complex(output0, 0);
648 buffer.store_complex(output1, 4);
649 }
650}
651
652pub struct Butterfly9Avx<T> {
653 twiddles: __m256,
654 twiddles_butterfly3: __m256,
655 direction: FftDirection,
656 _phantom_t: std::marker::PhantomData<T>,
657}
658boilerplate_fft_simd_butterfly!(Butterfly9Avx, 9);
659impl Butterfly9Avx<f32> {
660 #[target_feature(enable = "avx")]
661 unsafe fn new_with_avx(direction: FftDirection) -> Self {
662 let twiddles: [Complex<f32>; 4] = [
663 twiddles::compute_twiddle(1, 9, direction),
664 twiddles::compute_twiddle(2, 9, direction),
665 twiddles::compute_twiddle(2, 9, direction),
666 twiddles::compute_twiddle(4, 9, direction),
667 ];
668 Self {
669 twiddles: twiddles.as_slice().load_complex(0),
670 twiddles_butterfly3: AvxVector::broadcast_twiddle(1, 3, direction),
671 direction,
672 _phantom_t: PhantomData,
673 }
674 }
675}
676impl<T> Butterfly9Avx<T> {
677 #[target_feature(enable = "avx", enable = "fma")]
678 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
679 let input0_lo = _mm_castpd_ps(_mm_load1_pd(buffer.input_ptr() as *const f64));
683 let input0_hi = buffer.load_partial2_complex(1);
684 let input0 = AvxVector256::merge(input0_lo, input0_hi);
685 let input1 = buffer.load_complex(2);
686 let input2 = buffer.load_complex(5);
687
688 let [mid0, mid1, mid2] =
690 AvxVector::column_butterfly3([input0, input1, input2], self.twiddles_butterfly3);
691
692 let twiddle_data = _mm256_permute2f128_ps(mid1, mid2, 0x31);
694 let twiddled = AvxVector::mul_complex(twiddle_data, self.twiddles);
695
696 let permute0 = _mm256_permute2f128_ps(mid0, mid2, 0x20);
700 let permute1 = _mm256_permute2f128_ps(mid1, mid1, 0x20);
701 let permute2 = _mm256_permute2f128_ps(mid0, twiddled, 0x31);
702 let permute3 = _mm256_permute2f128_ps(twiddled, twiddled, 0x20);
703
704 let transposed0 = AvxVector::unpackhi_complex([permute0, permute1]);
705 let [transposed1, transposed2] = AvxVector::unpack_complex([permute2, permute3]);
706
707 let output_rows = AvxVector::column_butterfly3(
709 [transposed0, transposed1, transposed2],
710 self.twiddles_butterfly3,
711 );
712
713 let swapped1 = _mm256_permute_ps(output_rows[1], 0x4E); let packed1 = _mm256_permute2f128_ps(swapped1, output_rows[2], 0x21);
716 buffer.store_complex(packed1, 4);
717
718 let zero_swapped1_lo = AvxVector256::merge(AvxVector::zero(), swapped1.lo());
720 let packed0 = _mm256_blend_ps(output_rows[0], zero_swapped1_lo, 0xC0);
721 buffer.store_complex(packed0, 0);
722
723 buffer.store_partial1_complex(output_rows[2].hi(), 8);
725 }
726}
727
728pub struct Butterfly12Avx<T> {
729 twiddles: [__m256; 2],
730 twiddles_butterfly3: __m256,
731 twiddles_butterfly4: Rotation90<__m256>,
732 direction: FftDirection,
733 _phantom_t: std::marker::PhantomData<T>,
734}
735boilerplate_fft_simd_butterfly!(Butterfly12Avx, 12);
736impl Butterfly12Avx<f32> {
737 #[target_feature(enable = "avx")]
738 unsafe fn new_with_avx(direction: FftDirection) -> Self {
739 let twiddles = [
740 Complex {
741 re: 1.0f32,
742 im: 0.0,
743 },
744 Complex { re: 1.0, im: 0.0 },
745 twiddles::compute_twiddle(2, 12, direction),
746 twiddles::compute_twiddle(4, 12, direction),
747 twiddles::compute_twiddle(1, 12, direction),
749 twiddles::compute_twiddle(2, 12, direction),
750 twiddles::compute_twiddle(3, 12, direction),
751 twiddles::compute_twiddle(6, 12, direction),
752 ];
753 Self {
754 twiddles: [
755 twiddles.as_slice().load_complex(0),
756 twiddles.as_slice().load_complex(4),
757 ],
758 twiddles_butterfly3: AvxVector::broadcast_twiddle(1, 3, direction),
759 twiddles_butterfly4: AvxVector::make_rotation90(direction),
760 direction,
761 _phantom_t: PhantomData,
762 }
763 }
764}
765impl<T> Butterfly12Avx<T> {
766 #[target_feature(enable = "avx", enable = "fma")]
767 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
768 let input0_lo = _mm_castpd_ps(_mm_load1_pd(buffer.input_ptr() as *const f64));
773 let input0_hi = buffer.load_partial2_complex(1);
774 let input_rows = [
775 AvxVector256::merge(input0_lo, input0_hi),
776 buffer.load_complex(2),
777 buffer.load_complex(5),
778 buffer.load_complex(8),
779 ];
780
781 let mut mid = AvxVector::column_butterfly4(input_rows, self.twiddles_butterfly4);
783
784 mid[2] = AvxVector::mul_complex(mid[2], self.twiddles[0]);
788 let merged_mid13 = _mm256_permute2f128_ps(mid[1], mid[3], 0x31);
789 let twiddled13 = AvxVector::mul_complex(self.twiddles[1], merged_mid13);
790
791 let transposed = {
793 let permute0 = _mm256_permute2f128_ps(mid[0], mid[2], 0x20);
794 let permute1 = _mm256_permute2f128_ps(mid[1], mid[3], 0x20);
795 let permute2 = _mm256_permute2f128_ps(mid[0], mid[2], 0x31);
796 let permute3 = twiddled13; let unpacked1 = AvxVector::unpackhi_complex([permute0, permute1]);
799 let [unpacked2, unpacked3] = AvxVector::unpack_complex([permute2, permute3]);
800
801 [unpacked1, unpacked2, unpacked3]
802 };
803
804 let output_rows = AvxVector::column_butterfly3(transposed, self.twiddles_butterfly3);
806
807 buffer.store_complex(output_rows[0], 0);
808 buffer.store_complex(output_rows[1], 4);
809 buffer.store_complex(output_rows[2], 8);
810 }
811}
812
813pub struct Butterfly16Avx<T> {
814 twiddles: [__m256; 3],
815 twiddles_butterfly4: Rotation90<__m256>,
816 direction: FftDirection,
817 _phantom_t: std::marker::PhantomData<T>,
818}
819boilerplate_fft_simd_butterfly!(Butterfly16Avx, 16);
820impl Butterfly16Avx<f32> {
821 #[target_feature(enable = "avx")]
822 unsafe fn new_with_avx(direction: FftDirection) -> Self {
823 Self {
824 twiddles: gen_butterfly_twiddles_interleaved_columns!(4, 4, 0, direction),
825 twiddles_butterfly4: AvxVector::make_rotation90(direction),
826 direction,
827 _phantom_t: PhantomData,
828 }
829 }
830}
831impl<T> Butterfly16Avx<T> {
832 #[target_feature(enable = "avx", enable = "fma")]
833 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
834 let rows = [
836 buffer.load_complex(0),
837 buffer.load_complex(4),
838 buffer.load_complex(8),
839 buffer.load_complex(12),
840 ];
841
842 let mut mid = AvxVector::column_butterfly4(rows, self.twiddles_butterfly4);
844
845 for r in 1..4 {
847 mid[r] = AvxVector::mul_complex(mid[r], self.twiddles[r - 1]);
848 }
849
850 let transposed = avx32_utils::transpose_4x4_f32(mid);
852
853 let output_rows = AvxVector::column_butterfly4(transposed, self.twiddles_butterfly4);
855
856 buffer.store_complex(output_rows[0], 0);
858 buffer.store_complex(output_rows[1], 4);
859 buffer.store_complex(output_rows[2], 8);
860 buffer.store_complex(output_rows[3], 12);
861 }
862}
863
864pub struct Butterfly24Avx<T> {
865 twiddles: [__m256; 5],
866 twiddles_butterfly3: __m256,
867 twiddles_butterfly4: Rotation90<__m256>,
868 direction: FftDirection,
869 _phantom_t: std::marker::PhantomData<T>,
870}
871boilerplate_fft_simd_butterfly!(Butterfly24Avx, 24);
872impl Butterfly24Avx<f32> {
873 #[target_feature(enable = "avx")]
874 unsafe fn new_with_avx(direction: FftDirection) -> Self {
875 Self {
876 twiddles: gen_butterfly_twiddles_interleaved_columns!(6, 4, 0, direction),
877 twiddles_butterfly3: AvxVector::broadcast_twiddle(1, 3, direction),
878 twiddles_butterfly4: AvxVector::make_rotation90(direction),
879 direction,
880 _phantom_t: PhantomData,
881 }
882 }
883}
884impl<T> Butterfly24Avx<T> {
885 #[target_feature(enable = "avx", enable = "fma")]
886 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
887 let rows = [
889 buffer.load_complex(0),
890 buffer.load_complex(4),
891 buffer.load_complex(8),
892 buffer.load_complex(12),
893 buffer.load_complex(16),
894 buffer.load_complex(20),
895 ];
896
897 let mut mid = AvxVector256::column_butterfly6(rows, self.twiddles_butterfly3);
899
900 for r in 1..6 {
902 mid[r] = AvxVector::mul_complex(mid[r], self.twiddles[r - 1]);
903 }
904
905 let (transposed0, transposed1) = avx32_utils::transpose_4x6_to_6x4_f32(mid);
907
908 let output0 = AvxVector::column_butterfly4(transposed0, self.twiddles_butterfly4);
910 let output1 = AvxVector::column_butterfly4(transposed1, self.twiddles_butterfly4);
911
912 for r in 0..4 {
914 buffer.store_complex(output0[r], 6 * r);
915 buffer.store_partial2_complex(output1[r].lo(), r * 6 + 4);
916 }
917 }
918}
919
920pub struct Butterfly27Avx<T> {
921 twiddles: [__m256; 4],
922 twiddles_butterfly9: [__m256; 3],
923 twiddles_butterfly3: __m256,
924 direction: FftDirection,
925 _phantom_t: std::marker::PhantomData<T>,
926}
927boilerplate_fft_simd_butterfly!(Butterfly27Avx, 27);
928impl Butterfly27Avx<f32> {
929 #[target_feature(enable = "avx")]
930 unsafe fn new_with_avx(direction: FftDirection) -> Self {
931 Self {
932 twiddles: gen_butterfly_twiddles_interleaved_columns!(3, 9, 1, direction),
933 twiddles_butterfly9: [
934 AvxVector::broadcast_twiddle(1, 9, direction),
935 AvxVector::broadcast_twiddle(2, 9, direction),
936 AvxVector::broadcast_twiddle(4, 9, direction),
937 ],
938 twiddles_butterfly3: AvxVector::broadcast_twiddle(1, 3, direction),
939 direction,
940 _phantom_t: PhantomData,
941 }
942 }
943}
944impl<T> Butterfly27Avx<T> {
945 #[target_feature(enable = "avx", enable = "fma")]
946 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
947 let mut rows0 = [AvxVector::zero(); 3];
950 let mut rows1 = [AvxVector::zero(); 3];
951 let mut rows2 = [AvxVector::zero(); 3];
952 for r in 0..3 {
953 rows0[r] = buffer.load_partial2_complex(r * 9);
954 rows1[r] = buffer.load_complex(r * 9 + 1);
955 rows2[r] = buffer.load_complex(r * 9 + 5);
956 }
957
958 let mid0 = AvxVector::column_butterfly3(rows0, self.twiddles_butterfly3.lo());
960 let mut mid1 = AvxVector::column_butterfly3(rows1, self.twiddles_butterfly3);
961 let mut mid2 = AvxVector::column_butterfly3(rows2, self.twiddles_butterfly3);
962
963 mid1[1] = AvxVector::mul_complex(mid1[1], self.twiddles[0]);
965 mid2[1] = AvxVector::mul_complex(mid2[1], self.twiddles[1]);
966 mid1[2] = AvxVector::mul_complex(mid1[2], self.twiddles[2]);
967 mid2[2] = AvxVector::mul_complex(mid2[2], self.twiddles[3]);
968
969 let transposed = avx32_utils::transpose_9x3_to_3x9_emptycolumn1_f32(mid0, mid1, mid2);
971
972 let output_rows = AvxVector256::column_butterfly9(
974 transposed,
975 self.twiddles_butterfly9,
976 self.twiddles_butterfly3,
977 );
978
979 buffer.store_partial3_complex(output_rows[0], 0);
981 buffer.store_partial3_complex(output_rows[1], 3);
982 buffer.store_partial3_complex(output_rows[2], 6);
983 buffer.store_partial3_complex(output_rows[3], 9);
984 buffer.store_partial3_complex(output_rows[4], 12);
985 buffer.store_partial3_complex(output_rows[5], 15);
986 buffer.store_partial3_complex(output_rows[6], 18);
987 buffer.store_partial3_complex(output_rows[7], 21);
988 buffer.store_partial3_complex(output_rows[8], 24);
989 }
990}
991
992pub struct Butterfly32Avx<T> {
993 twiddles: [__m256; 6],
994 twiddles_butterfly4: Rotation90<__m256>,
995 direction: FftDirection,
996 _phantom_t: std::marker::PhantomData<T>,
997}
998boilerplate_fft_simd_butterfly!(Butterfly32Avx, 32);
999impl Butterfly32Avx<f32> {
1000 #[target_feature(enable = "avx")]
1001 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1002 Self {
1003 twiddles: gen_butterfly_twiddles_interleaved_columns!(4, 8, 0, direction),
1004 twiddles_butterfly4: AvxVector::make_rotation90(direction),
1005 direction,
1006 _phantom_t: PhantomData,
1007 }
1008 }
1009}
1010impl<T> Butterfly32Avx<T> {
1011 #[target_feature(enable = "avx", enable = "fma")]
1012 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
1013 let mut rows0 = [AvxVector::zero(); 4];
1014 let mut rows1 = [AvxVector::zero(); 4];
1015 for r in 0..4 {
1016 rows0[r] = buffer.load_complex(8 * r);
1017 rows1[r] = buffer.load_complex(8 * r + 4);
1018 }
1019
1020 let mut mid0 = AvxVector::column_butterfly4(rows0, self.twiddles_butterfly4);
1022 let mut mid1 = AvxVector::column_butterfly4(rows1, self.twiddles_butterfly4);
1023
1024 for r in 1..4 {
1026 mid0[r] = AvxVector::mul_complex(mid0[r], self.twiddles[2 * r - 2]);
1027 mid1[r] = AvxVector::mul_complex(mid1[r], self.twiddles[2 * r - 1]);
1028 }
1029
1030 let transposed = avx32_utils::transpose_8x4_to_4x8_f32(mid0, mid1);
1032
1033 let output_rows = AvxVector::column_butterfly8(transposed, self.twiddles_butterfly4);
1035
1036 buffer.store_complex(output_rows[0], 0);
1038 buffer.store_complex(output_rows[1], 1 * 4);
1039 buffer.store_complex(output_rows[2], 2 * 4);
1040 buffer.store_complex(output_rows[3], 3 * 4);
1041 buffer.store_complex(output_rows[4], 4 * 4);
1042 buffer.store_complex(output_rows[5], 5 * 4);
1043 buffer.store_complex(output_rows[6], 6 * 4);
1044 buffer.store_complex(output_rows[7], 7 * 4);
1045 }
1046}
1047
1048pub struct Butterfly36Avx<T> {
1049 twiddles: [__m256; 6],
1050 twiddles_butterfly9: [__m256; 3],
1051 twiddles_butterfly3: __m256,
1052 twiddles_butterfly4: Rotation90<__m256>,
1053 direction: FftDirection,
1054 _phantom_t: std::marker::PhantomData<T>,
1055}
1056boilerplate_fft_simd_butterfly!(Butterfly36Avx, 36);
1057impl Butterfly36Avx<f32> {
1058 #[target_feature(enable = "avx")]
1059 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1060 Self {
1061 twiddles: gen_butterfly_twiddles_interleaved_columns!(4, 9, 1, direction),
1062 twiddles_butterfly9: [
1063 AvxVector::broadcast_twiddle(1, 9, direction),
1064 AvxVector::broadcast_twiddle(2, 9, direction),
1065 AvxVector::broadcast_twiddle(4, 9, direction),
1066 ],
1067 twiddles_butterfly3: AvxVector::broadcast_twiddle(1, 3, direction),
1068 twiddles_butterfly4: AvxVector::make_rotation90(direction),
1069 direction,
1070 _phantom_t: PhantomData,
1071 }
1072 }
1073}
1074impl<T> Butterfly36Avx<T> {
1075 #[target_feature(enable = "avx", enable = "fma")]
1076 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
1077 let mut rows0 = [AvxVector::zero(); 4];
1080 let mut rows1 = [AvxVector::zero(); 4];
1081 let mut rows2 = [AvxVector::zero(); 4];
1082 for r in 0..4 {
1083 rows0[r] = buffer.load_partial2_complex(r * 9);
1084 rows1[r] = buffer.load_complex(r * 9 + 1);
1085 rows2[r] = buffer.load_complex(r * 9 + 5);
1086 }
1087
1088 let mid0 = AvxVector::column_butterfly4(rows0, self.twiddles_butterfly4.lo());
1090 let mut mid1 = AvxVector::column_butterfly4(rows1, self.twiddles_butterfly4);
1091 let mut mid2 = AvxVector::column_butterfly4(rows2, self.twiddles_butterfly4);
1092
1093 for r in 1..4 {
1095 mid1[r] = AvxVector::mul_complex(mid1[r], self.twiddles[2 * r - 2]);
1096 mid2[r] = AvxVector::mul_complex(mid2[r], self.twiddles[2 * r - 1]);
1097 }
1098
1099 let transposed = avx32_utils::transpose_9x4_to_4x9_emptycolumn1_f32(mid0, mid1, mid2);
1101
1102 let output_rows = AvxVector256::column_butterfly9(
1104 transposed,
1105 self.twiddles_butterfly9,
1106 self.twiddles_butterfly3,
1107 );
1108
1109 for r in 0..3 {
1110 buffer.store_complex(output_rows[r * 3], r * 12);
1111 buffer.store_complex(output_rows[r * 3 + 1], r * 12 + 4);
1112 buffer.store_complex(output_rows[r * 3 + 2], r * 12 + 8);
1113 }
1114 }
1115}
1116
1117pub struct Butterfly48Avx<T> {
1118 twiddles: [__m256; 9],
1119 twiddles_butterfly3: __m256,
1120 twiddles_butterfly4: Rotation90<__m256>,
1121 direction: FftDirection,
1122 _phantom_t: std::marker::PhantomData<T>,
1123}
1124boilerplate_fft_simd_butterfly!(Butterfly48Avx, 48);
1125impl Butterfly48Avx<f32> {
1126 #[target_feature(enable = "avx")]
1127 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1128 Self {
1129 twiddles: gen_butterfly_twiddles_interleaved_columns!(4, 12, 0, direction),
1130 twiddles_butterfly3: AvxVector::broadcast_twiddle(1, 3, direction),
1131 twiddles_butterfly4: AvxVector::make_rotation90(direction),
1132 direction,
1133 _phantom_t: PhantomData,
1134 }
1135 }
1136}
1137impl<T> Butterfly48Avx<T> {
1138 #[target_feature(enable = "avx", enable = "fma")]
1139 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
1140 let mut rows0 = [AvxVector::zero(); 4];
1141 let mut rows1 = [AvxVector::zero(); 4];
1142 let mut rows2 = [AvxVector::zero(); 4];
1143 for r in 0..4 {
1144 rows0[r] = buffer.load_complex(12 * r);
1145 rows1[r] = buffer.load_complex(12 * r + 4);
1146 rows2[r] = buffer.load_complex(12 * r + 8);
1147 }
1148
1149 let mut mid0 = AvxVector::column_butterfly4(rows0, self.twiddles_butterfly4);
1151 let mut mid1 = AvxVector::column_butterfly4(rows1, self.twiddles_butterfly4);
1152 let mut mid2 = AvxVector::column_butterfly4(rows2, self.twiddles_butterfly4);
1153
1154 for r in 1..4 {
1156 mid0[r] = AvxVector::mul_complex(mid0[r], self.twiddles[3 * r - 3]);
1157 mid1[r] = AvxVector::mul_complex(mid1[r], self.twiddles[3 * r - 2]);
1158 mid2[r] = AvxVector::mul_complex(mid2[r], self.twiddles[3 * r - 1]);
1159 }
1160
1161 let transposed = avx32_utils::transpose_12x4_to_4x12_f32(mid0, mid1, mid2);
1163
1164 let output_rows = AvxVector256::column_butterfly12(
1166 transposed,
1167 self.twiddles_butterfly3,
1168 self.twiddles_butterfly4,
1169 );
1170
1171 buffer.store_complex(output_rows[0], 0);
1173 buffer.store_complex(output_rows[1], 4);
1174 buffer.store_complex(output_rows[2], 8);
1175 buffer.store_complex(output_rows[3], 12);
1176 buffer.store_complex(output_rows[4], 16);
1177 buffer.store_complex(output_rows[5], 20);
1178 buffer.store_complex(output_rows[6], 24);
1179 buffer.store_complex(output_rows[7], 28);
1180 buffer.store_complex(output_rows[8], 32);
1181 buffer.store_complex(output_rows[9], 36);
1182 buffer.store_complex(output_rows[10], 40);
1183 buffer.store_complex(output_rows[11], 44);
1184 }
1185}
1186
1187pub struct Butterfly54Avx<T> {
1188 twiddles: [__m256; 10],
1189 twiddles_butterfly9: [__m256; 3],
1190 twiddles_butterfly9_lo: [__m256; 2],
1191 twiddles_butterfly3: __m256,
1192 direction: FftDirection,
1193 _phantom_t: std::marker::PhantomData<T>,
1194}
1195boilerplate_fft_simd_butterfly!(Butterfly54Avx, 54);
1196impl Butterfly54Avx<f32> {
1197 #[target_feature(enable = "avx")]
1198 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1199 let twiddle1 = __m128::broadcast_twiddle(1, 9, direction);
1200 let twiddle2 = __m128::broadcast_twiddle(2, 9, direction);
1201 let twiddle4 = __m128::broadcast_twiddle(4, 9, direction);
1202
1203 Self {
1204 twiddles: gen_butterfly_twiddles_interleaved_columns!(6, 9, 1, direction),
1205 twiddles_butterfly9: [
1206 AvxVector::broadcast_twiddle(1, 9, direction),
1207 AvxVector::broadcast_twiddle(2, 9, direction),
1208 AvxVector::broadcast_twiddle(4, 9, direction),
1209 ],
1210 twiddles_butterfly9_lo: [
1211 AvxVector256::merge(twiddle1, twiddle2),
1212 AvxVector256::merge(twiddle2, twiddle4),
1213 ],
1214 twiddles_butterfly3: AvxVector::broadcast_twiddle(1, 3, direction),
1215 direction,
1216 _phantom_t: PhantomData,
1217 }
1218 }
1219}
1220impl<T> Butterfly54Avx<T> {
1221 #[target_feature(enable = "avx", enable = "fma")]
1222 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
1223 let mut rows0 = [AvxVector::zero(); 6];
1228 for r in 0..3 {
1229 rows0[r * 2] = buffer.load_partial2_complex(r * 18);
1230 rows0[r * 2 + 1] = buffer.load_partial2_complex(r * 18 + 9);
1231 }
1232 let mid0 = AvxVector128::column_butterfly6(rows0, self.twiddles_butterfly3);
1233
1234 let mut rows1 = [AvxVector::zero(); 6];
1236 for r in 0..3 {
1237 rows1[r * 2] = buffer.load_complex(r * 18 + 1);
1238 rows1[r * 2 + 1] = buffer.load_complex(r * 18 + 10);
1239 }
1240 let mut mid1 = AvxVector256::column_butterfly6(rows1, self.twiddles_butterfly3);
1241 for r in 1..6 {
1242 mid1[r] = AvxVector::mul_complex(mid1[r], self.twiddles[2 * r - 2]);
1243 }
1244
1245 let mut rows2 = [AvxVector::zero(); 6];
1247 for r in 0..3 {
1248 rows2[r * 2] = buffer.load_complex(r * 18 + 5);
1249 rows2[r * 2 + 1] = buffer.load_complex(r * 18 + 14);
1250 }
1251 let mut mid2 = AvxVector256::column_butterfly6(rows2, self.twiddles_butterfly3);
1252 for r in 1..6 {
1253 mid2[r] = AvxVector::mul_complex(mid2[r], self.twiddles[2 * r - 1]);
1254 }
1255
1256 let (transposed0, transposed1) =
1258 avx32_utils::transpose_9x6_to_6x9_emptycolumn1_f32(mid0, mid1, mid2);
1259
1260 let output_rows1 = AvxVector128::column_butterfly9(
1263 transposed1,
1264 self.twiddles_butterfly9_lo,
1265 self.twiddles_butterfly3,
1266 );
1267 for r in 0..9 {
1268 buffer.store_partial2_complex(output_rows1[r], r * 6 + 4);
1269 }
1270
1271 let output_rows0 = AvxVector256::column_butterfly9(
1273 transposed0,
1274 self.twiddles_butterfly9,
1275 self.twiddles_butterfly3,
1276 );
1277 for r in 0..9 {
1278 buffer.store_complex(output_rows0[r], r * 6);
1279 }
1280 }
1281}
1282
1283pub struct Butterfly64Avx<T> {
1284 twiddles: [__m256; 14],
1285 twiddles_butterfly4: Rotation90<__m256>,
1286 direction: FftDirection,
1287 _phantom_t: std::marker::PhantomData<T>,
1288}
1289boilerplate_fft_simd_butterfly!(Butterfly64Avx, 64);
1290impl Butterfly64Avx<f32> {
1291 #[target_feature(enable = "avx")]
1292 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1293 Self {
1294 twiddles: gen_butterfly_twiddles_separated_columns!(8, 8, 0, direction),
1295 twiddles_butterfly4: AvxVector::make_rotation90(direction),
1296 direction,
1297 _phantom_t: PhantomData,
1298 }
1299 }
1300}
1301impl<T> Butterfly64Avx<T> {
1302 #[target_feature(enable = "avx", enable = "fma")]
1303 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
1304 let mut rows0 = [AvxVector::zero(); 8];
1308 for r in 0..8 {
1309 rows0[r] = buffer.load_complex(8 * r);
1310 }
1311 let mut mid0 = AvxVector::column_butterfly8(rows0, self.twiddles_butterfly4);
1312 for r in 1..8 {
1313 mid0[r] = AvxVector::mul_complex(mid0[r], self.twiddles[r - 1]);
1314 }
1315
1316 let mut rows1 = [AvxVector::zero(); 8];
1318 for r in 0..8 {
1319 rows1[r] = buffer.load_complex(8 * r + 4);
1320 }
1321 let mut mid1 = AvxVector::column_butterfly8(rows1, self.twiddles_butterfly4);
1322 for r in 1..8 {
1323 mid1[r] = AvxVector::mul_complex(mid1[r], self.twiddles[r - 1 + 7]);
1324 }
1325
1326 let (transposed0, transposed1) = avx32_utils::transpose_8x8_f32(mid0, mid1);
1328
1329 let output0 = AvxVector::column_butterfly8(transposed0, self.twiddles_butterfly4);
1332 for r in 0..8 {
1333 buffer.store_complex(output0[r], 8 * r);
1334 }
1335
1336 let output1 = AvxVector::column_butterfly8(transposed1, self.twiddles_butterfly4);
1337 for r in 0..8 {
1338 buffer.store_complex(output1[r], 8 * r + 4);
1339 }
1340 }
1341}
1342
1343pub struct Butterfly72Avx<T> {
1344 twiddles: [__m256; 15],
1345 twiddles_butterfly4: Rotation90<__m256>,
1346 twiddles_butterfly3: __m256,
1347 direction: FftDirection,
1348 _phantom_t: std::marker::PhantomData<T>,
1349}
1350boilerplate_fft_simd_butterfly!(Butterfly72Avx, 72);
1351impl Butterfly72Avx<f32> {
1352 #[target_feature(enable = "avx")]
1353 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1354 Self {
1355 twiddles: gen_butterfly_twiddles_separated_columns!(6, 12, 0, direction),
1356 twiddles_butterfly4: AvxVector::make_rotation90(direction),
1357 twiddles_butterfly3: AvxVector::broadcast_twiddle(1, 3, direction),
1358 direction,
1359 _phantom_t: PhantomData,
1360 }
1361 }
1362}
1363impl<T> Butterfly72Avx<T> {
1364 #[target_feature(enable = "avx", enable = "fma")]
1365 unsafe fn perform_fft_f32(&self, mut buffer: impl AvxArrayMut<f32>) {
1366 let mut rows0 = [AvxVector::zero(); 6];
1370 for r in 0..6 {
1371 rows0[r] = buffer.load_complex(12 * r);
1372 }
1373 let mut mid0 = AvxVector256::column_butterfly6(rows0, self.twiddles_butterfly3);
1374 for r in 1..6 {
1375 mid0[r] = AvxVector::mul_complex(mid0[r], self.twiddles[r - 1]);
1376 }
1377
1378 let mut rows1 = [AvxVector::zero(); 6];
1380 for r in 0..6 {
1381 rows1[r] = buffer.load_complex(12 * r + 4);
1382 }
1383 let mut mid1 = AvxVector256::column_butterfly6(rows1, self.twiddles_butterfly3);
1384 for r in 1..6 {
1385 mid1[r] = AvxVector::mul_complex(mid1[r], self.twiddles[r - 1 + 5]);
1386 }
1387
1388 let mut rows2 = [AvxVector::zero(); 6];
1390 for r in 0..6 {
1391 rows2[r] = buffer.load_complex(12 * r + 8);
1392 }
1393 let mut mid2 = AvxVector256::column_butterfly6(rows2, self.twiddles_butterfly3);
1394 for r in 1..6 {
1395 mid2[r] = AvxVector::mul_complex(mid2[r], self.twiddles[r - 1 + 10]);
1396 }
1397
1398 let (transposed0, transposed1) = avx32_utils::transpose_12x6_to_6x12_f32(mid0, mid1, mid2);
1400
1401 let output0 = AvxVector128::column_butterfly12(
1404 transposed0,
1405 self.twiddles_butterfly3,
1406 self.twiddles_butterfly4,
1407 );
1408 for r in 0..12 {
1409 buffer.store_partial2_complex(output0[r], 6 * r);
1410 }
1411
1412 let output1 = AvxVector256::column_butterfly12(
1413 transposed1,
1414 self.twiddles_butterfly3,
1415 self.twiddles_butterfly4,
1416 );
1417 for r in 0..12 {
1418 buffer.store_complex(output1[r], 6 * r + 2);
1419 }
1420 }
1421}
1422
1423pub struct Butterfly128Avx<T> {
1424 twiddles: [__m256; 28],
1425 twiddles_butterfly16: [__m256; 2],
1426 twiddles_butterfly4: Rotation90<__m256>,
1427 direction: FftDirection,
1428 _phantom_t: std::marker::PhantomData<T>,
1429}
1430boilerplate_fft_simd_butterfly_with_scratch!(Butterfly128Avx, 128);
1431impl Butterfly128Avx<f32> {
1432 #[target_feature(enable = "avx")]
1433 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1434 Self {
1435 twiddles: gen_butterfly_twiddles_separated_columns!(8, 16, 0, direction),
1436 twiddles_butterfly16: [
1437 AvxVector::broadcast_twiddle(1, 16, direction),
1438 AvxVector::broadcast_twiddle(3, 16, direction),
1439 ],
1440 twiddles_butterfly4: AvxVector::make_rotation90(direction),
1441 direction,
1442 _phantom_t: PhantomData,
1443 }
1444 }
1445}
1446impl<T> Butterfly128Avx<T> {
1447 #[target_feature(enable = "avx", enable = "fma")]
1448 unsafe fn column_butterflies_and_transpose(
1449 &self,
1450 input: &[Complex<f32>],
1451 mut output: &mut [Complex<f32>],
1452 ) {
1453 for columnset in 0..4 {
1460 let mut rows = [AvxVector::zero(); 8];
1461 for r in 0..8 {
1462 rows[r] = input.load_complex(columnset * 4 + 16 * r);
1463 }
1464 let mut mid = AvxVector::column_butterfly8(rows, self.twiddles_butterfly4);
1466
1467 for r in 1..8 {
1469 mid[r] = AvxVector::mul_complex(mid[r], self.twiddles[r - 1 + 7 * columnset]);
1470 }
1471
1472 let transposed = AvxVector::transpose8_packed(mid);
1474
1475 for i in 0..4 {
1477 output.store_complex(transposed[i * 2], columnset * 32 + i * 8);
1478 output.store_complex(transposed[i * 2 + 1], columnset * 32 + i * 8 + 4);
1479 }
1480 }
1481 }
1482
1483 #[target_feature(enable = "avx", enable = "fma")]
1484 unsafe fn row_butterflies(&self, mut buffer: impl AvxArrayMut<f32>) {
1485 for columnset in 0usize..2 {
1489 column_butterfly16_loadfn!(
1490 |index: usize| buffer.load_complex(columnset * 4 + index * 8),
1491 |data, index| buffer.store_complex(data, columnset * 4 + index * 8),
1492 self.twiddles_butterfly16,
1493 self.twiddles_butterfly4
1494 );
1495 }
1496 }
1497}
1498
1499#[allow(non_camel_case_types)]
1500pub struct Butterfly256Avx<T> {
1501 twiddles: [__m256; 56],
1502 twiddles_butterfly32: [__m256; 6],
1503 twiddles_butterfly4: Rotation90<__m256>,
1504 direction: FftDirection,
1505 _phantom_t: std::marker::PhantomData<T>,
1506}
1507boilerplate_fft_simd_butterfly_with_scratch!(Butterfly256Avx, 256);
1508impl Butterfly256Avx<f32> {
1509 #[target_feature(enable = "avx")]
1510 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1511 Self {
1512 twiddles: gen_butterfly_twiddles_separated_columns!(8, 32, 0, direction),
1513 twiddles_butterfly32: [
1514 AvxVector::broadcast_twiddle(1, 32, direction),
1515 AvxVector::broadcast_twiddle(2, 32, direction),
1516 AvxVector::broadcast_twiddle(3, 32, direction),
1517 AvxVector::broadcast_twiddle(5, 32, direction),
1518 AvxVector::broadcast_twiddle(6, 32, direction),
1519 AvxVector::broadcast_twiddle(7, 32, direction),
1520 ],
1521 twiddles_butterfly4: AvxVector::make_rotation90(direction),
1522 direction,
1523 _phantom_t: PhantomData,
1524 }
1525 }
1526}
1527impl<T> Butterfly256Avx<T> {
1528 #[target_feature(enable = "avx", enable = "fma")]
1529 unsafe fn column_butterflies_and_transpose(
1530 &self,
1531 input: &[Complex<f32>],
1532 mut output: &mut [Complex<f32>],
1533 ) {
1534 for columnset in 0..8 {
1541 let mut rows = [AvxVector::zero(); 8];
1542 for r in 0..8 {
1543 rows[r] = input.load_complex(columnset * 4 + 32 * r);
1544 }
1545 let mut mid = AvxVector::column_butterfly8(rows, self.twiddles_butterfly4);
1546 for r in 1..8 {
1547 mid[r] = AvxVector::mul_complex(mid[r], self.twiddles[r - 1 + 7 * columnset]);
1548 }
1549
1550 let transposed = AvxVector::transpose8_packed(mid);
1552
1553 for i in 0..4 {
1554 output.store_complex(transposed[i * 2], columnset * 32 + i * 8);
1555 output.store_complex(transposed[i * 2 + 1], columnset * 32 + i * 8 + 4);
1556 }
1557 }
1558 }
1559
1560 #[target_feature(enable = "avx", enable = "fma")]
1561 unsafe fn row_butterflies(&self, mut buffer: impl AvxArrayMut<f32>) {
1562 for columnset in 0..2 {
1566 column_butterfly32_loadfn!(
1567 |index: usize| buffer.load_complex(columnset * 4 + index * 8),
1568 |data, index| buffer.store_complex(data, columnset * 4 + index * 8),
1569 self.twiddles_butterfly32,
1570 self.twiddles_butterfly4
1571 );
1572 }
1573 }
1574}
1575
1576pub struct Butterfly512Avx<T> {
1577 twiddles: [__m256; 120],
1578 twiddles_butterfly32: [__m256; 6],
1579 twiddles_butterfly16: [__m256; 2],
1580 twiddles_butterfly4: Rotation90<__m256>,
1581 direction: FftDirection,
1582 _phantom_t: std::marker::PhantomData<T>,
1583}
1584boilerplate_fft_simd_butterfly_with_scratch!(Butterfly512Avx, 512);
1585impl Butterfly512Avx<f32> {
1586 #[target_feature(enable = "avx")]
1587 unsafe fn new_with_avx(direction: FftDirection) -> Self {
1588 Self {
1589 twiddles: gen_butterfly_twiddles_separated_columns!(16, 32, 0, direction),
1590 twiddles_butterfly32: [
1591 AvxVector::broadcast_twiddle(1, 32, direction),
1592 AvxVector::broadcast_twiddle(2, 32, direction),
1593 AvxVector::broadcast_twiddle(3, 32, direction),
1594 AvxVector::broadcast_twiddle(5, 32, direction),
1595 AvxVector::broadcast_twiddle(6, 32, direction),
1596 AvxVector::broadcast_twiddle(7, 32, direction),
1597 ],
1598 twiddles_butterfly16: [
1599 AvxVector::broadcast_twiddle(1, 16, direction),
1600 AvxVector::broadcast_twiddle(3, 16, direction),
1601 ],
1602 twiddles_butterfly4: AvxVector::make_rotation90(direction),
1603 direction,
1604 _phantom_t: PhantomData,
1605 }
1606 }
1607}
1608impl<T> Butterfly512Avx<T> {
1609 #[target_feature(enable = "avx", enable = "fma")]
1610 unsafe fn column_butterflies_and_transpose(
1611 &self,
1612 input: &[Complex<f32>],
1613 mut output: &mut [Complex<f32>],
1614 ) {
1615 const TWIDDLES_PER_COLUMN: usize = 15;
1623 for (columnset, twiddle_chunk) in
1624 self.twiddles.chunks_exact(TWIDDLES_PER_COLUMN).enumerate()
1625 {
1626 let mut mid_uninit: [MaybeUninit<__m256>; 16] = [MaybeUninit::<__m256>::uninit(); 16];
1629
1630 column_butterfly16_loadfn!(
1631 |index: usize| input.load_complex(columnset * 4 + 32 * index),
1632 |data, index: usize| {
1633 mid_uninit[index].as_mut_ptr().write(data);
1634 },
1635 self.twiddles_butterfly16,
1636 self.twiddles_butterfly4
1637 );
1638
1639 for chunk in 0..4 {
1642 let twiddled = [
1643 if chunk > 0 {
1644 AvxVector::mul_complex(
1645 mid_uninit[4 * chunk].assume_init(),
1646 twiddle_chunk[4 * chunk - 1],
1647 )
1648 } else {
1649 mid_uninit[4 * chunk].assume_init()
1650 },
1651 AvxVector::mul_complex(
1652 mid_uninit[4 * chunk + 1].assume_init(),
1653 twiddle_chunk[4 * chunk],
1654 ),
1655 AvxVector::mul_complex(
1656 mid_uninit[4 * chunk + 2].assume_init(),
1657 twiddle_chunk[4 * chunk + 1],
1658 ),
1659 AvxVector::mul_complex(
1660 mid_uninit[4 * chunk + 3].assume_init(),
1661 twiddle_chunk[4 * chunk + 2],
1662 ),
1663 ];
1664
1665 let transposed = AvxVector::transpose4_packed(twiddled);
1666
1667 output.store_complex(transposed[0], columnset * 64 + 0 * 16 + 4 * chunk);
1668 output.store_complex(transposed[1], columnset * 64 + 1 * 16 + 4 * chunk);
1669 output.store_complex(transposed[2], columnset * 64 + 2 * 16 + 4 * chunk);
1670 output.store_complex(transposed[3], columnset * 64 + 3 * 16 + 4 * chunk);
1671 }
1672 }
1673 }
1674
1675 #[target_feature(enable = "avx", enable = "fma")]
1676 unsafe fn row_butterflies(&self, mut buffer: impl AvxArrayMut<f32>) {
1677 for columnset in 0..4 {
1681 column_butterfly32_loadfn!(
1682 |index: usize| buffer.load_complex(columnset * 4 + index * 16),
1683 |data, index| buffer.store_complex(data, columnset * 4 + index * 16),
1684 self.twiddles_butterfly32,
1685 self.twiddles_butterfly4
1686 );
1687 }
1688 }
1689}
1690
1691#[cfg(test)]
1692mod unit_tests {
1693 use super::*;
1694 use crate::test_utils::check_fft_algorithm;
1695
1696 macro_rules! test_avx_butterfly {
1697 ($test_name:ident, $struct_name:ident, $size:expr) => (
1698 #[test]
1699 fn $test_name() {
1700 let butterfly = $struct_name::<f32>::new(FftDirection::Forward).expect("Can't run test because this machine doesn't have the required instruction sets");
1701 check_fft_algorithm(&butterfly as &dyn Fft<f32>, $size, FftDirection::Forward);
1702
1703 let butterfly_inverse = $struct_name::<f32>::new(FftDirection::Inverse).expect("Can't run test because this machine doesn't have the required instruction sets");
1704 check_fft_algorithm(&butterfly_inverse as &dyn Fft<f32>, $size, FftDirection::Inverse);
1705 }
1706 )
1707 }
1708 test_avx_butterfly!(test_avx_butterfly5, Butterfly5Avx, 5);
1709 test_avx_butterfly!(test_avx_butterfly7, Butterfly7Avx, 7);
1710 test_avx_butterfly!(test_avx_butterfly8, Butterfly8Avx, 8);
1711 test_avx_butterfly!(test_avx_butterfly9, Butterfly9Avx, 9);
1712 test_avx_butterfly!(test_avx_butterfly11, Butterfly11Avx, 11);
1713 test_avx_butterfly!(test_avx_butterfly12, Butterfly12Avx, 12);
1714 test_avx_butterfly!(test_avx_butterfly16, Butterfly16Avx, 16);
1715 test_avx_butterfly!(test_avx_butterfly24, Butterfly24Avx, 24);
1716 test_avx_butterfly!(test_avx_butterfly27, Butterfly27Avx, 27);
1717 test_avx_butterfly!(test_avx_butterfly32, Butterfly32Avx, 32);
1718 test_avx_butterfly!(test_avx_butterfly36, Butterfly36Avx, 36);
1719 test_avx_butterfly!(test_avx_butterfly48, Butterfly48Avx, 48);
1720 test_avx_butterfly!(test_avx_butterfly54, Butterfly54Avx, 54);
1721 test_avx_butterfly!(test_avx_butterfly64, Butterfly64Avx, 64);
1722 test_avx_butterfly!(test_avx_butterfly72, Butterfly72Avx, 72);
1723 test_avx_butterfly!(test_avx_butterfly128, Butterfly128Avx, 128);
1724 test_avx_butterfly!(test_avx_butterfly256, Butterfly256Avx, 256);
1725 test_avx_butterfly!(test_avx_butterfly512, Butterfly512Avx, 512);
1726}