Skip to main content

rustfft/avx/
avx32_butterflies.rs

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
15// Safety: This macro will call `self::perform_fft_f32()` which probably has a #[target_feature(enable = "...")] annotation on it.
16// Calling functions with that annotation is unsafe, because it doesn't actually check if the CPU has the required features.
17// Callers of this macro must guarantee that users can't even obtain an instance of $struct_name if their CPU doesn't have the required CPU features.
18macro_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                    // Safety: new_internal requires the "avx" feature set. Since we know it's present, we're safe
29                    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
115// Safety: This macro will call `self::column_butterflies_and_transpose and self::row_butterflies()` which probably has a #[target_feature(enable = "...")] annotation on it.
116// Calling functions with that annotation is unsafe, because it doesn't actually check if the CPU has the required features.
117// Callers of this macro must guarantee that users can't even obtain an instance of $struct_name if their CPU doesn't have the required CPU features.
118macro_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                    // Safety: new_internal requires the "avx" feature set. Since we know it's present, we're safe
127                    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                // Perform the column FFTs
141                // Safety: self.perform_column_butterflies() requres the "avx" and "fma" instruction sets, and we return Err() in our constructor if the instructions aren't available
142                unsafe { self.column_butterflies_and_transpose(buffer, scratch) };
143
144                // process the row FFTs, and copy from the scratch back to the buffer as we go
145                // Safety: self.transpose() requres the "avx" instruction set, and we return Err() in our constructor if the instructions aren't available
146                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                // Perform the column FFTs
157                // Safety: self.perform_column_butterflies() requres the "avx" and "fma" instruction sets, and we return Err() in our constructor if the instructions aren't available
158                unsafe { self.column_butterflies_and_transpose(input, output) };
159
160                // process the row FFTs in-place in the output buffer
161                // Safety: self.transpose() requres the "avx" instruction set, and we return Err() in our constructor if the instructions aren't available
162                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)); // load the first element of the input, and duplicate it into both complex number slots of input0
307        let input12 = buffer.load_partial2_complex(1);
308        let input34 = buffer.load_partial2_complex(3);
309
310        // swap elements for inputs 3 and 4
311        let input43 = AvxVector::reverse_complex_elements(input34);
312
313        // do some prep work before we can start applying twiddle factors
314        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        // to compute the first output, compute the sum of all elements. mid14[0] and mid23[0] already have the sum of 1+4 and 2+3 respectively, so if we add them, we'll get the sum of all 4
322        let sum1234 = AvxVector::add(mid14, mid23);
323        let output0 = AvxVector::add(input0, sum1234);
324
325        // apply twiddle factors
326        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        // unpack the data for the last butterfly 2
332        let [twiddled12, twiddled43] = AvxVector::unpack_complex([twiddled14, twiddled23]);
333        let [output12, output43] = AvxVector::column_butterfly2([twiddled12, twiddled43]);
334
335        // swap the elements in output43 before writing them out, and add the first input to everything
336        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        // load the first element of the input, and duplicate it into both complex number slots of input0
375        let input0 = _mm_castpd_ps(_mm_load1_pd(buffer.input_ptr() as *const f64));
376
377        // we want to load 3 elements into 123 and 3 elements into 456, but we can only load 4, so we're going to do slightly overlapping reads here
378        // we have to reverse 456 immediately after loading, and that'll be easiest if we load the 456 into the latter 3 slots of the register, rather than the front 3 slots
379        // as a bonus, that also means we don't need masked reads or anything
380        let input123 = buffer.load_complex(1);
381        let input456 = buffer.load_complex(3);
382
383        // reverse the order of input456
384        let input654 = AvxVector::reverse_complex_elements(input456);
385
386        // do some prep work before we can start applying twiddle factors
387        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        // to compute the first output, compute the sum of all elements. mid16[0], mid25[0], and mid34[0] already have the sum of 1+6, 2+5 and 3+4 respectively, so if we add them, we'll get 1+2+3+4+5+6
398        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        // apply twiddle factors
406        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        // unpack the data for the last butterfly 2
422        let [twiddled12, twiddled65] = AvxVector::unpack_complex([twiddled16, twiddled25]);
423        let [twiddled33, twiddled44] = AvxVector::unpack_complex([twiddled34, twiddled34]);
424
425        // we can save one add if we add input0 to twiddled33 now. normally we'd add input0 to the final output, but the arrangement of data makes that a little awkward
426        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)); // load the first element of the input, and duplicate it into both complex number slots of input0
499        let input1234 = buffer.load_complex(1);
500        let input56 = buffer.load_partial2_complex(5);
501        let input78910 = buffer.load_complex(7);
502
503        // reverse the order of input78910, and separate
504        let [input55, input66] = AvxVector::unpack_complex([input56, input56]);
505        let input10987 = AvxVector::reverse_complex_elements(input78910);
506
507        // do some initial butterflies and rotations
508        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        // arrange the data into the format to apply twiddles
516        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        // to compute the first output, compute the sum of all elements. mid16[0], mid25[0], and mid34[0] already have the sum of 1+6, 2+5 and 3+4 respectively, so if we add them, we'll get 1+2+3+4+5+6
526        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        // we need to add the first input to each of our 5 twiddles values -- but right now, input0 is duplicated into both slots
534        // but we only want to add it once, so zero the second element
535        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        // apply twiddle factors
540        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        // unpack the data for the last butterfly 2
561        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        // Do our butterfly 2's down the columns
618        let [intermediate0, intermediate1_pretwiddle] = AvxVector::column_butterfly2([row0, row1]);
619
620        // Apply the size-8 twiddle factors
621        let intermediate1 = AvxVector::mul_complex(intermediate1_pretwiddle, self.twiddles);
622
623        // Rearrange the data before we do our butterfly 4s. This swaps the last 2 elements of butterfly0 with the first two elements of butterfly1
624        // The result is that we can then do a 4x butterfly 2, apply twiddles, use unpack instructions to transpose to the final output, then do another 4x butterfly 2
625        let permuted0 = _mm256_permute2f128_ps(intermediate0, intermediate1, 0x20);
626        let permuted1 = _mm256_permute2f128_ps(intermediate0, intermediate1, 0x31);
627
628        // Do the first set of butterfly 2's
629        let [postbutterfly0, postbutterfly1_pretwiddle] =
630            AvxVector::column_butterfly2([permuted0, permuted1]);
631
632        // Which negative we blend in depends on whether we're forward or direction
633        // Our goal is to swap the reals with the imaginaries, then negate either the reals or the imaginaries, based on whether we're an direction or not
634        // but we can't use the AvxVector swap_complex_components function, because we only want to swap the odd reals with the odd imaginaries
635        let elements_swapped = _mm256_permute_ps(postbutterfly1_pretwiddle, 0xB4);
636
637        // We can negate the elements we want by xoring the row with a pre-set vector
638        let postbutterfly1 = AvxVector::xor(elements_swapped, self.twiddles_butterfly4);
639
640        // use unpack instructions to transpose, and to prepare for the final butterfly 2's
641        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        // we're going to load these elements in a peculiar way. instead of loading a row into the first 3 element of each register and leaving the last element empty
680        // we're leaving the first element empty and putting the data in the last 3 elements. this will let us do 3 total complex multiplies instead of 4.
681
682        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        // We're going to treat our input as a 3x3 2d array. First, do 3 butterfly 3's down the columns of that array.
689        let [mid0, mid1, mid2] =
690            AvxVector::column_butterfly3([input0, input1, input2], self.twiddles_butterfly3);
691
692        // merge the twiddle-able data into a single avx vector
693        let twiddle_data = _mm256_permute2f128_ps(mid1, mid2, 0x31);
694        let twiddled = AvxVector::mul_complex(twiddle_data, self.twiddles);
695
696        // Transpose our 3x3 array. We could use the 4x4 transpose with an empty bottom row, which would result in an empty last column
697        // but it turns out that it'll make our packing process later simpler if we duplicate the second row into the last row
698        // which will result in duplicating the second column into the last column after the transpose
699        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        // more size 3 buterflies
708        let output_rows = AvxVector::column_butterfly3(
709            [transposed0, transposed1, transposed2],
710            self.twiddles_butterfly3,
711        );
712
713        // the elements of row 1 are in pretty much the worst possible order, thankfully we can fix that with just a couple instructions
714        let swapped1 = _mm256_permute_ps(output_rows[1], 0x4E); // swap even and odd complex numbers
715        let packed1 = _mm256_permute2f128_ps(swapped1, output_rows[2], 0x21);
716        buffer.store_complex(packed1, 4);
717
718        // merge just the high element of swapped_lo into the high element of row 0
719        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        // The last element can just be written on its own
724        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            // note that these twiddles are deliberately in a weird order, see perform_fft_f32 for why
748            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        // we're going to load these elements in a peculiar way. instead of loading a row into the first 3 element of each register and leaving the last element empty
769        // we're leaving the first element empty and putting the data in the last 3 elements. this will save us a complex multiply.
770
771        // for everything but the first element, we can do overlapping reads. for the first element, an "overlapping read" would have us reading from index -1, so instead we have to shuffle some data around
772        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        // 3 butterfly 4's down the columns
782        let mut mid = AvxVector::column_butterfly4(input_rows, self.twiddles_butterfly4);
783
784        // Multiply in our twiddle factors. mid2 will be normal, but for mid1 and mid3, we're going to merge the twiddle-able parts into a single vector,
785        // and do a single complex multiply on it. this transformation saves a complex multiply and costs nothing,
786        // because we needthe second halves of mid1 and mid3 in a single vector for the transpose afterward anyways, so we would have done this permute2f128 operation either way
787        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        // Transpose our 3x4 array into a 4x3. we're doing a custom transpose here because we have to re-distribute the merged twiddled23 back out, and we can roll that into the transpose to make it free
792        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; // normally we'd need to do a permute here, but we can skip it because we already did it for twiddle factors
797
798            let unpacked1 = AvxVector::unpackhi_complex([permute0, permute1]);
799            let [unpacked2, unpacked3] = AvxVector::unpack_complex([permute2, permute3]);
800
801            [unpacked1, unpacked2, unpacked3]
802        };
803
804        // Do 4 butterfly 3's down the columns of our transposed array
805        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        // Manually unrolling this loop because writing a "for r in 0..4" loop results in slow codegen that makes the whole thing take 1.5x longer :(
835        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        // We're going to treat our input as a 4x4 2d array. First, do 4 butterfly 4's down the columns of that array.
843        let mut mid = AvxVector::column_butterfly4(rows, self.twiddles_butterfly4);
844
845        // apply twiddle factors
846        for r in 1..4 {
847            mid[r] = AvxVector::mul_complex(mid[r], self.twiddles[r - 1]);
848        }
849
850        // Transpose our 4x4 array
851        let transposed = avx32_utils::transpose_4x4_f32(mid);
852
853        // Do 4 butterfly 4's down the columns of our transposed array
854        let output_rows = AvxVector::column_butterfly4(transposed, self.twiddles_butterfly4);
855
856        // Manually unrolling this loop because writing a "for r in 0..4" loop results in slow codegen that makes the whole thing take 1.5x longer :(
857        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        // Manually unrolling this loop because writing a "for r in 0..6" loop results in slow codegen that makes the whole thing take 1.5x longer :(
888        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        // We're going to treat our input as a 4x6 2d array. First, do 4 butterfly 6's down the columns of that array.
898        let mut mid = AvxVector256::column_butterfly6(rows, self.twiddles_butterfly3);
899
900        // apply twiddle factors
901        for r in 1..6 {
902            mid[r] = AvxVector::mul_complex(mid[r], self.twiddles[r - 1]);
903        }
904
905        // Transpose our 6x4 array into a 4x6.
906        let (transposed0, transposed1) = avx32_utils::transpose_4x6_to_6x4_f32(mid);
907
908        // Do 6 butterfly 4's down the columns of our transposed array
909        let output0 = AvxVector::column_butterfly4(transposed0, self.twiddles_butterfly4);
910        let output1 = AvxVector::column_butterfly4(transposed1, self.twiddles_butterfly4);
911
912        // the upper two elements of output1 are empty, so only store half the data for it
913        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        // we're going to load our data in a peculiar way. we're going to load the first column on its own as a column of __m128.
948        // it's faster to just load the first 2 columns into these m128s than trying to worry about masks, etc, so the second column will piggyback along and we just won't use it
949        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        // butterfly 3s down the columns
959        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        // apply twiddle factors
964        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        // transpose 9x3 to 3x9. this will be a little awkward because of rows0 containing garbage data, so use a transpose function that knows to ignore it
970        let transposed = avx32_utils::transpose_9x3_to_3x9_emptycolumn1_f32(mid0, mid1, mid2);
971
972        // butterfly 9s down the rows
973        let output_rows = AvxVector256::column_butterfly9(
974            transposed,
975            self.twiddles_butterfly9,
976            self.twiddles_butterfly3,
977        );
978
979        // Our last column is empty, so it's a bit awkward to write out to memory. We could pack it in fewer vectors, but benchmarking shows it's simpler and just as fast to just brute-force it with partial writes
980        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        // We're going to treat our input as a 8x4 2d array. First, do 8 butterfly 4's down the columns of that array.
1021        let mut mid0 = AvxVector::column_butterfly4(rows0, self.twiddles_butterfly4);
1022        let mut mid1 = AvxVector::column_butterfly4(rows1, self.twiddles_butterfly4);
1023
1024        // apply twiddle factors
1025        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        // Transpose our 8x4 array to an 4x8 array
1031        let transposed = avx32_utils::transpose_8x4_to_4x8_f32(mid0, mid1);
1032
1033        // Do 4 butterfly 8's down the columns of our transpsed array
1034        let output_rows = AvxVector::column_butterfly8(transposed, self.twiddles_butterfly4);
1035
1036        // Manually unrolling this loop because writing a "for r in 0..8" loop results in slow codegen that makes the whole thing take 1.5x longer :(
1037        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        // we're going to load our data in a peculiar way. we're going to load the first column on its own as a column of __m128.
1078        // it's faster to just load the first 2 columns into these m128s than trying to worry about masks, etc, so the second column will piggyback along and we just won't use it
1079        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        // butterfly 4s down the columns
1089        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        // apply twiddle factors
1094        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        // transpose 9x4 to 4x9. this will be a little awkward because of rows0 containing garbage data, so use a transpose function that knows to ignore it
1100        let transposed = avx32_utils::transpose_9x4_to_4x9_emptycolumn1_f32(mid0, mid1, mid2);
1101
1102        // butterfly 9s down the rows
1103        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        // We're going to treat our input as a 12x4 2d array. First, do 12 butterfly 4's down the columns of that array.
1150        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        // apply twiddle factors
1155        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        // Transpose our 12x4 array into a 4x12.
1162        let transposed = avx32_utils::transpose_12x4_to_4x12_f32(mid0, mid1, mid2);
1163
1164        // Do 4 butterfly 12's down the columns of our transposed array
1165        let output_rows = AvxVector256::column_butterfly12(
1166            transposed,
1167            self.twiddles_butterfly3,
1168            self.twiddles_butterfly4,
1169        );
1170
1171        // Manually unrolling this loop because writing a "for r in 0..12" loop results in slow codegen that makes the whole thing take 1.5x longer :(
1172        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        // we're going to load our data in a peculiar way. we're going to load the first column on its own as a column of __m128.
1224        // it's faster to just load the first 2 columns into these m128s than trying to worry about masks, etc, so the second column will piggyback along and we just won't use it
1225        //
1226        // we have too much data to fit into registers all at once, so split up our data processing so that we entirely finish with one "rows_" array before moving to the next
1227        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        // next set of butterfly 6's
1235        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        // final set of butterfly 6's
1246        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        // transpose 9x6 to 6x9. this will be a little awkward because of rows0 containing garbage data, so use a transpose function that knows to ignore it
1257        let (transposed0, transposed1) =
1258            avx32_utils::transpose_9x6_to_6x9_emptycolumn1_f32(mid0, mid1, mid2);
1259
1260        // butterfly 9s down the rows
1261        // process the other half
1262        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        // we have too much data to fit into registers all at once, do one set of butterfly 9's and output them before even starting on the others, to make it easier for the compiler to figure out what to spill
1272        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        // We're going to treat our input as a 8x8 2d array. First, do 8 butterfly 8's down the columns of that array.
1305        // We can't fit the whole problem into AVX registers at once, so we'll have to spill some things.
1306        // By computing a sizeable chunk and not referencing any of it for a while, we're making it easy for the compiler to decide what to spill
1307        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        // One half is done, so the compiler can spill everything above this. Now do the other set of columns
1317        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        // Transpose our 8x8 array
1327        let (transposed0, transposed1) = avx32_utils::transpose_8x8_f32(mid0, mid1);
1328
1329        // Do 8 butterfly 8's down the columns of our transposed array, and store the results
1330        // Same thing as above - Do the half of the butterfly 8's separately to give the compiler a better hint about what to spill
1331        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        // We're going to treat our input as a 12x6 2d array. First, do butterfly 6's down the columns of that array.
1367        // We can't fit the whole problem into AVX registers at once, so we'll have to spill some things.
1368        // By computing a sizeable chunk and not referencing any of it for a while, we're making it easy for the compiler to decide what to spill
1369        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        // One third is done, so the compiler can spill everything above this. Now do the middle set of columns
1379        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        // two thirds are done, so the compiler can spill everything above this. Now do the final set of columns
1389        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        // Transpose our 12x6 array to 6x12 array
1399        let (transposed0, transposed1) = avx32_utils::transpose_12x6_to_6x12_f32(mid0, mid1, mid2);
1400
1401        // Do butterfly 12's down the columns of our transposed array, and store the results
1402        // Same thing as above - Do the half of the butterfly 12's separately to give the compiler a better hint about what to spill
1403        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        // A size-128 FFT is way too big to fit in registers, so instead we're going to compute it in two phases, storing in scratch in between.
1454
1455        // First phase is to treat this size-128 array like a 16x8 2D array, and do butterfly 8's down the columns
1456        // Then, apply twiddle factors, and finally transpose into the scratch space
1457
1458        // But again, we don't have enough registers to load it all at once, so only load one column of AVX vectors at a time
1459        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            // apply butterfly 8
1465            let mut mid = AvxVector::column_butterfly8(rows, self.twiddles_butterfly4);
1466
1467            // apply twiddle factors
1468            for r in 1..8 {
1469                mid[r] = AvxVector::mul_complex(mid[r], self.twiddles[r - 1 + 7 * columnset]);
1470            }
1471
1472            // transpose
1473            let transposed = AvxVector::transpose8_packed(mid);
1474
1475            // write out
1476            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        // Second phase: Butterfly 16's down the columns of our transposed array.
1486        // Thankfully, during the first phase, we set everything up so that all we have to do here is compute the size-16 FFT columns and write them back out where we got them
1487        // We're also using a customized butterfly16 function that is smarter about when it loads/stores data, to reduce register spilling
1488        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        // A size-256 FFT is way too big to fit in registers, so instead we're going to compute it in two phases, storing in scratch in between.
1535
1536        // First phase is to treeat this size-256 array like a 32x8 2D array, and do butterfly 8's down the columns
1537        // Then, apply twiddle factors, and finally transpose into the scratch space
1538
1539        // But again, we don't have enough registers to load it all at once, so only load one column of AVX vectors at a time
1540        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            // Before writing to the scratch, transpose this chunk of the array
1551            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        // Second phase: Butterfly 32's down the columns of our transposed array.
1563        // Thankfully, during the first phase, we set everything up so that all we have to do here is compute the size-32 FFT columns and write them back out where we got them
1564        // We're also using a customized butterfly32 function that is smarter about when it loads/stores data, to reduce register spilling
1565        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        // A size-512 FFT is way too big to fit in registers, so instead we're going to compute it in two phases, storing in scratch in between.
1616
1617        // First phase is to treat this size-512 array like a 32x16 2D array, and do butterfly 16's down the columns
1618        // Then, apply twiddle factors, and finally transpose into the scratch space
1619
1620        // But again, we don't have enough registers to load it all at once, so only load one column of AVX vectors at a time
1621        // We're also using a customized butterfly16 function that is smarter about when it loads/stores data, to reduce register spilling
1622        const TWIDDLES_PER_COLUMN: usize = 15;
1623        for (columnset, twiddle_chunk) in
1624            self.twiddles.chunks_exact(TWIDDLES_PER_COLUMN).enumerate()
1625        {
1626            // Sadly we have to use MaybeUninit here. If we init an array like normal with AvxVector::Zero(), the compiler can't seem to figure out that it can
1627            // eliminate the dead stores of zeroes to the stack. By using uninit here, we avoid those unnecessary writes
1628            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            // Apply twiddle factors, transpose, and store. Traditionally we apply all the twiddle factors at once and then do all the transposes at once,
1640            // But our data is pushing the limit of what we can store in registers, so the idea here is to get the data out the door with as few spills to the stack as possible
1641            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        // Second phase: Butterfly 32's down the columns of our transposed array.
1678        // Thankfully, during the first phase, we set everything up so that all we have to do here is compute the size-32 FFT columns and write them back out where we got them
1679        // We're also using a customized butterfly32 function that is smarter about when it loads/stores data, to reduce register spilling
1680        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}