Skip to main content

rustfft/avx/
avx_vector.rs

1use std::arch::x86_64::*;
2use std::fmt::Debug;
3use std::ops::{Deref, DerefMut};
4
5use num_complex::Complex;
6use num_traits::Zero;
7
8use crate::{array_utils::DoubleBuf, twiddles, FftDirection};
9
10use super::AvxNum;
11
12/// A SIMD vector of complex numbers, stored with the real values and imaginary values interleaved.
13/// Implemented for __m128, __m128d, __m256, __m256d, but these all require the AVX instruction set.
14///
15/// The goal of this trait is to reduce code duplication by letting code be generic over the vector type
16pub trait AvxVector: Copy + Debug + Send + Sync {
17    const COMPLEX_PER_VECTOR: usize;
18
19    // useful constants
20    unsafe fn zero() -> Self;
21    unsafe fn half_root2() -> Self; // an entire vector filled with 0.5.sqrt()
22
23    // Basic operations that map directly to 1-2 AVX intrinsics
24    unsafe fn add(left: Self, right: Self) -> Self;
25    unsafe fn sub(left: Self, right: Self) -> Self;
26    unsafe fn xor(left: Self, right: Self) -> Self;
27    unsafe fn neg(self) -> Self;
28    unsafe fn mul(left: Self, right: Self) -> Self;
29    unsafe fn fmadd(left: Self, right: Self, add: Self) -> Self;
30    unsafe fn fnmadd(left: Self, right: Self, add: Self) -> Self;
31    unsafe fn fmaddsub(left: Self, right: Self, add: Self) -> Self;
32    unsafe fn fmsubadd(left: Self, right: Self, add: Self) -> Self;
33
34    // More basic operations that end up being implemented in 1-2 intrinsics, but unlike the ones above, these have higher-level meaning than just arithmetic
35    /// Swap each real number with its corresponding imaginary number
36    unsafe fn swap_complex_components(self) -> Self;
37
38    /// first return is the reals duplicated into the imaginaries, second return is the imaginaries duplicated into the reals
39    unsafe fn duplicate_complex_components(self) -> (Self, Self);
40
41    /// Reverse the order of complex numbers in the vector, so that the last is the first and the first is the last
42    unsafe fn reverse_complex_elements(self) -> Self;
43
44    /// Copies the even elements of rows[1] into the corresponding odd elements of rows[0] and returns the result.
45    unsafe fn unpacklo_complex(rows: [Self; 2]) -> Self;
46    /// Copies the odd elements of rows[0] into the corresponding even elements of rows[1] and returns the result.
47    unsafe fn unpackhi_complex(rows: [Self; 2]) -> Self;
48
49    #[inline(always)]
50    unsafe fn unpack_complex(rows: [Self; 2]) -> [Self; 2] {
51        [Self::unpacklo_complex(rows), Self::unpackhi_complex(rows)]
52    }
53
54    /// Fill a vector by computing a twiddle factor and repeating it across the whole vector
55    unsafe fn broadcast_twiddle(index: usize, len: usize, direction: FftDirection) -> Self;
56
57    /// create a Rotator90 instance to rotate complex numbers either 90 or 270 degrees, based on the value of `inverse`
58    unsafe fn make_rotation90(direction: FftDirection) -> Rotation90<Self>;
59
60    /// Generates a chunk of twiddle factors starting at (X,Y) and incrementing X `COMPLEX_PER_VECTOR` times.
61    /// The result will be [twiddle(x*y, len), twiddle((x+1)*y, len), twiddle((x+2)*y, len), ...] for as many complex numbers fit in a vector
62    unsafe fn make_mixedradix_twiddle_chunk(
63        x: usize,
64        y: usize,
65        len: usize,
66        direction: FftDirection,
67    ) -> Self;
68
69    /// Packed transposes. Used by mixed radix. These all take a NxC array, where C is COMPLEX_PER_VECTOR, and transpose it to a CxN array.
70    /// But they also pack the result into as few vectors as possible, with the goal of writing the transposed data out contiguously.
71    unsafe fn transpose2_packed(rows: [Self; 2]) -> [Self; 2];
72    unsafe fn transpose3_packed(rows: [Self; 3]) -> [Self; 3];
73    unsafe fn transpose4_packed(rows: [Self; 4]) -> [Self; 4];
74    unsafe fn transpose5_packed(rows: [Self; 5]) -> [Self; 5];
75    unsafe fn transpose6_packed(rows: [Self; 6]) -> [Self; 6];
76    unsafe fn transpose7_packed(rows: [Self; 7]) -> [Self; 7];
77    unsafe fn transpose8_packed(rows: [Self; 8]) -> [Self; 8];
78    unsafe fn transpose9_packed(rows: [Self; 9]) -> [Self; 9];
79    unsafe fn transpose11_packed(rows: [Self; 11]) -> [Self; 11];
80    unsafe fn transpose12_packed(rows: [Self; 12]) -> [Self; 12];
81    unsafe fn transpose16_packed(rows: [Self; 16]) -> [Self; 16];
82
83    /// Pairwise multiply the complex numbers in `left` with the complex numbers in `right`.
84    #[inline(always)]
85    unsafe fn mul_complex(left: Self, right: Self) -> Self {
86        // Extract the real and imaginary components from left into 2 separate registers
87        let (left_real, left_imag) = Self::duplicate_complex_components(left);
88
89        // create a shuffled version of right where the imaginary values are swapped with the reals
90        let right_shuffled = Self::swap_complex_components(right);
91
92        // multiply our duplicated imaginary left vector by our shuffled right vector. that will give us the right side of the traditional complex multiplication formula
93        let output_right = Self::mul(left_imag, right_shuffled);
94
95        // use a FMA instruction to multiply together left side of the complex multiplication formula, then alternatingly add and subtract the left side from the right
96        Self::fmaddsub(left_real, right, output_right)
97    }
98
99    #[inline(always)]
100    unsafe fn rotate90(self, rotation: Rotation90<Self>) -> Self {
101        // Use the pre-computed vector stored in the Rotation90 instance to negate either the reals or imaginaries
102        let negated = Self::xor(self, rotation.0);
103
104        // Our goal is to swap the reals with the imaginaries, then negate either the reals or the imaginaries, based on whether we're an inverse or not
105        Self::swap_complex_components(negated)
106    }
107
108    #[inline(always)]
109    unsafe fn column_butterfly2(rows: [Self; 2]) -> [Self; 2] {
110        [Self::add(rows[0], rows[1]), Self::sub(rows[0], rows[1])]
111    }
112
113    #[inline(always)]
114    unsafe fn column_butterfly3(rows: [Self; 3], twiddles: Self) -> [Self; 3] {
115        // This algorithm is derived directly from the definition of the Dft of size 3
116        // We'd theoretically have to do 4 complex multiplications, but all of the twiddles we'd be multiplying by are conjugates of each other
117        // By doing some algebra to expand the complex multiplications and factor out the multiplications, we get this
118
119        let [mut mid1, mid2] = Self::column_butterfly2([rows[1], rows[2]]);
120        let output0 = Self::add(rows[0], mid1);
121
122        let (twiddle_real, twiddle_imag) = Self::duplicate_complex_components(twiddles);
123
124        mid1 = Self::fmadd(mid1, twiddle_real, rows[0]);
125
126        let rotation = Self::make_rotation90(FftDirection::Inverse);
127        let mid2_rotated = Self::rotate90(mid2, rotation);
128
129        let output1 = Self::fmadd(mid2_rotated, twiddle_imag, mid1);
130        let output2 = Self::fnmadd(mid2_rotated, twiddle_imag, mid1);
131
132        [output0, output1, output2]
133    }
134
135    #[inline(always)]
136    unsafe fn column_butterfly4(rows: [Self; 4], rotation: Rotation90<Self>) -> [Self; 4] {
137        // Algorithm: 2x2 mixed radix
138
139        // Perform the first set of size-2 FFTs.
140        let [mid0, mid2] = Self::column_butterfly2([rows[0], rows[2]]);
141        let [mid1, mid3] = Self::column_butterfly2([rows[1], rows[3]]);
142
143        // Apply twiddle factors (in this case just a rotation)
144        let mid3_rotated = mid3.rotate90(rotation);
145
146        // Transpose the data and do size-2 FFTs down the columns
147        let [output0, output1] = Self::column_butterfly2([mid0, mid1]);
148        let [output2, output3] = Self::column_butterfly2([mid2, mid3_rotated]);
149
150        // Swap outputs 1 and 2 in the output to do a square transpose
151        [output0, output2, output1, output3]
152    }
153
154    #[inline(always)]
155    unsafe fn column_butterfly5(rows: [Self; 5], twiddles: [Self; 2]) -> [Self; 5] {
156        // This algorithm is derived directly from the definition of the Dft of size 5
157        // We'd theoretically have to do 16 complex multiplications for the Dft, but many of the twiddles we'd be multiplying by are conjugates of each other
158        // By doing some algebra to expand the complex multiplications and factor out the real multiplications, we get this faster formula where we only do the equivalent of 4 multiplications
159
160        // do some prep work before we can start applying twiddle factors
161        let [sum1, diff4] = Self::column_butterfly2([rows[1], rows[4]]);
162        let [sum2, diff3] = Self::column_butterfly2([rows[2], rows[3]]);
163
164        let rotation = Self::make_rotation90(FftDirection::Inverse);
165        let rotated4 = Self::rotate90(diff4, rotation);
166        let rotated3 = Self::rotate90(diff3, rotation);
167
168        // to compute the first output, compute the sum of all elements. sum1 and sum2 already have the sum of 1+4 and 2+3 respectively, so if we add them, we'll get the sum of all 4
169        let sum1234 = Self::add(sum1, sum2);
170        let output0 = Self::add(rows[0], sum1234);
171
172        // apply twiddle factors
173        let (twiddles0_re, twiddles0_im) = Self::duplicate_complex_components(twiddles[0]);
174        let (twiddles1_re, twiddles1_im) = Self::duplicate_complex_components(twiddles[1]);
175
176        let twiddled1_mid = Self::fmadd(twiddles0_re, sum1, rows[0]);
177        let twiddled2_mid = Self::fmadd(twiddles1_re, sum1, rows[0]);
178        let twiddled3_mid = Self::mul(twiddles1_im, rotated4);
179        let twiddled4_mid = Self::mul(twiddles0_im, rotated4);
180
181        let twiddled1 = Self::fmadd(twiddles1_re, sum2, twiddled1_mid);
182        let twiddled2 = Self::fmadd(twiddles0_re, sum2, twiddled2_mid);
183        let twiddled3 = Self::fnmadd(twiddles0_im, rotated3, twiddled3_mid); // fnmadd instead of fmadd because we're actually re-using twiddle0 here. remember that this algorithm is all about factoring out conjugated multiplications -- this negation of the twiddle0 imaginaries is a reflection of one of those conugations
184        let twiddled4 = Self::fmadd(twiddles1_im, rotated3, twiddled4_mid);
185
186        // Post-processing to mix the twiddle factors between the rest of the output
187        let [output1, output4] = Self::column_butterfly2([twiddled1, twiddled4]);
188        let [output2, output3] = Self::column_butterfly2([twiddled2, twiddled3]);
189
190        [output0, output1, output2, output3, output4]
191    }
192
193    #[inline(always)]
194    unsafe fn column_butterfly7(rows: [Self; 7], twiddles: [Self; 3]) -> [Self; 7] {
195        // This algorithm is derived directly from the definition of the Dft of size 7
196        // We'd theoretically have to do 36 complex multiplications for the Dft, but many of the twiddles we'd be multiplying by are conjugates of each other
197        // By doing some algebra to expand the complex multiplications and factor out the real multiplications, we get this faster formula where we only do the equivalent of 9 multiplications
198
199        // do some prep work before we can start applying twiddle factors
200        let [sum1, diff6] = Self::column_butterfly2([rows[1], rows[6]]);
201        let [sum2, diff5] = Self::column_butterfly2([rows[2], rows[5]]);
202        let [sum3, diff4] = Self::column_butterfly2([rows[3], rows[4]]);
203
204        let rotation = Self::make_rotation90(FftDirection::Inverse);
205        let rotated4 = Self::rotate90(diff4, rotation);
206        let rotated5 = Self::rotate90(diff5, rotation);
207        let rotated6 = Self::rotate90(diff6, rotation);
208
209        // to compute the first output, compute the sum of all elements. sum1, sum2, and sum3 already have the sum of 1+6 and 2+5 and 3+4 respectively, so if we add them, we'll get the sum of all 6
210        let output0_left = Self::add(sum1, sum2);
211        let output0_right = Self::add(sum3, rows[0]);
212        let output0 = Self::add(output0_left, output0_right);
213
214        // apply twiddle factors. This is probably pushing the limit of how much we should do with this technique.
215        // We probably shouldn't do a size-11 FFT with this technique, for example, because this block of multiplies would grow quadratically
216        let (twiddles0_re, twiddles0_im) = Self::duplicate_complex_components(twiddles[0]);
217        let (twiddles1_re, twiddles1_im) = Self::duplicate_complex_components(twiddles[1]);
218        let (twiddles2_re, twiddles2_im) = Self::duplicate_complex_components(twiddles[2]);
219
220        // Let's do a plain 7-point Dft
221        // | X0 |   | W0  W0  W0  W0  W0  W0  W0  |   | x0 |
222        // | X1 |   | W0  W1  W2  W3  W4  W5  W6  |   | x1 |
223        // | X2 |   | W0  W2  W4  W6  W8  W10 W12 |   | x2 |
224        // | X3 |   | W0  W3  W6  W9  W12 W15 W18 |   | x3 |
225        // | X4 |   | W0  W4  W8  W12 W16 W20 W24 |   | x4 |
226        // | X5 |   | W0  W5  W10 W15 W20 W25 W30 |   | x5 |
227        // | X6 |   | W0  W6  W12 W18 W24 W30 W36 |   | x6 |
228        // where Wn = exp(-2*pi*n/7) for a forward transform, and exp(+2*pi*n/7) for an inverse.
229
230        // Next, take advantage of the fact that twiddle factor indexes for a size-7 Dft are cyclical mod 7
231        // | X0 |   | W0  W0  W0  W0  W0  W0  W0  |   | x0 |
232        // | X1 |   | W0  W1  W2  W3  W4  W5  W6  |   | x1 |
233        // | X2 |   | W0  W2  W4  W6  W1  W3  W5  |   | x2 |
234        // | X3 |   | W0  W3  W6  W2  W5  W1  W4  |   | x3 |
235        // | X4 |   | W0  W4  W1  W5  W2  W6  W3  |   | x4 |
236        // | X5 |   | W0  W5  W3  W1  W6  W4  W2  |   | x5 |
237        // | X6 |   | W0  W6  W5  W4  W3  W2  W1  |   | x6 |
238
239        // Finally, take advantage of the fact that for a size-7 Dft,
240        // twiddles 4 through 6 are conjugates of twiddes 3 through 0 (Asterisk marks conjugates)
241        // | X0 |   | W0  W0  W0  W0  W0  W0  W0  |   | x0 |
242        // | X1 |   | W0  W1  W2  W3  W3* W2* W1* |   | x1 |
243        // | X2 |   | W0  W2  W3* W1* W1  W3  W2* |   | x2 |
244        // | X3 |   | W0  W3  W1* W2  W2* W1  W3* |   | x3 |
245        // | X4 |   | W0  W3* W1  W2* W2  W1* W3  |   | x4 |
246        // | X5 |   | W0  W2* W3  W1  W1* W3* W2  |   | x5 |
247        // | X6 |   | W0  W1* W2* W3* W3  W2  W1  |   | x6 |
248
249        let twiddled1_mid = Self::fmadd(twiddles0_re, sum1, rows[0]);
250        let twiddled2_mid = Self::fmadd(twiddles1_re, sum1, rows[0]);
251        let twiddled3_mid = Self::fmadd(twiddles2_re, sum1, rows[0]);
252        let twiddled4_mid = Self::mul(twiddles2_im, rotated6);
253        let twiddled5_mid = Self::mul(twiddles1_im, rotated6);
254        let twiddled6_mid = Self::mul(twiddles0_im, rotated6);
255
256        let twiddled1_mid2 = Self::fmadd(twiddles1_re, sum2, twiddled1_mid);
257        let twiddled2_mid2 = Self::fmadd(twiddles2_re, sum2, twiddled2_mid);
258        let twiddled3_mid2 = Self::fmadd(twiddles0_re, sum2, twiddled3_mid);
259        let twiddled4_mid2 = Self::fnmadd(twiddles0_im, rotated5, twiddled4_mid); // fnmadd instead of fmadd because we're actually re-using twiddle0 here. remember that this algorithm is all about factoring out conjugated multiplications -- this negation of the twiddle0 imaginaries is a reflection of one of those conugations
260        let twiddled5_mid2 = Self::fnmadd(twiddles2_im, rotated5, twiddled5_mid);
261        let twiddled6_mid2 = Self::fmadd(twiddles1_im, rotated5, twiddled6_mid);
262
263        let twiddled1 = Self::fmadd(twiddles2_re, sum3, twiddled1_mid2);
264        let twiddled2 = Self::fmadd(twiddles0_re, sum3, twiddled2_mid2);
265        let twiddled3 = Self::fmadd(twiddles1_re, sum3, twiddled3_mid2);
266        let twiddled4 = Self::fmadd(twiddles1_im, rotated4, twiddled4_mid2);
267        let twiddled5 = Self::fnmadd(twiddles0_im, rotated4, twiddled5_mid2);
268        let twiddled6 = Self::fmadd(twiddles2_im, rotated4, twiddled6_mid2);
269
270        // Post-processing to mix the twiddle factors between the rest of the output
271        let [output1, output6] = Self::column_butterfly2([twiddled1, twiddled6]);
272        let [output2, output5] = Self::column_butterfly2([twiddled2, twiddled5]);
273        let [output3, output4] = Self::column_butterfly2([twiddled3, twiddled4]);
274
275        [
276            output0, output1, output2, output3, output4, output5, output6,
277        ]
278    }
279
280    #[inline(always)]
281    unsafe fn column_butterfly8(rows: [Self; 8], rotation: Rotation90<Self>) -> [Self; 8] {
282        // Algorithm: 4x2 mixed radix
283
284        // Size-4 FFTs down the columns
285        let mid0 = Self::column_butterfly4([rows[0], rows[2], rows[4], rows[6]], rotation);
286        let mut mid1 = Self::column_butterfly4([rows[1], rows[3], rows[5], rows[7]], rotation);
287
288        // Apply twiddle factors
289        mid1[1] = apply_butterfly8_twiddle1(mid1[1], rotation);
290        mid1[2] = mid1[2].rotate90(rotation);
291        mid1[3] = apply_butterfly8_twiddle3(mid1[3], rotation);
292
293        // Transpose the data and do size-2 FFTs down the columns
294        let [output0, output1] = Self::column_butterfly2([mid0[0], mid1[0]]);
295        let [output2, output3] = Self::column_butterfly2([mid0[1], mid1[1]]);
296        let [output4, output5] = Self::column_butterfly2([mid0[2], mid1[2]]);
297        let [output6, output7] = Self::column_butterfly2([mid0[3], mid1[3]]);
298
299        [
300            output0, output2, output4, output6, output1, output3, output5, output7,
301        ]
302    }
303
304    #[inline(always)]
305    unsafe fn column_butterfly11(rows: [Self; 11], twiddles: [Self; 5]) -> [Self; 11] {
306        // This algorithm is derived directly from the definition of the Dft of size 11
307        // We'd theoretically have to do 100 complex multiplications for the Dft, but many of the twiddles we'd be multiplying by are conjugates of each other
308        // By doing some algebra to expand the complex multiplications and factor out the real multiplications, we get this faster formula where we only do the equivalent of 9 multiplications
309
310        // do some prep work before we can start applying twiddle factors
311        let [sum1, diff10] = Self::column_butterfly2([rows[1], rows[10]]);
312        let [sum2, diff9] = Self::column_butterfly2([rows[2], rows[9]]);
313        let [sum3, diff8] = Self::column_butterfly2([rows[3], rows[8]]);
314        let [sum4, diff7] = Self::column_butterfly2([rows[4], rows[7]]);
315        let [sum5, diff6] = Self::column_butterfly2([rows[5], rows[6]]);
316
317        let rotation = Self::make_rotation90(FftDirection::Inverse);
318        let rotated10 = Self::rotate90(diff10, rotation);
319        let rotated9 = Self::rotate90(diff9, rotation);
320        let rotated8 = Self::rotate90(diff8, rotation);
321        let rotated7 = Self::rotate90(diff7, rotation);
322        let rotated6 = Self::rotate90(diff6, rotation);
323
324        // to compute the first output, compute the sum of all elements. sum1, sum2, and sum3 already have the sum of 1+6 and 2+5 and 3+4 respectively, so if we add them, we'll get the sum of all 6
325        let sum01 = Self::add(rows[0], sum1);
326        let sum23 = Self::add(sum2, sum3);
327        let sum45 = Self::add(sum4, sum5);
328        let sum0123 = Self::add(sum01, sum23);
329        let output0 = Self::add(sum0123, sum45);
330
331        // apply twiddle factors. This is probably pushing the limit of how much we should do with this technique.
332        // We probably shouldn't do a size-11 FFT with this technique, for example, because this block of multiplies would grow quadratically
333        let (twiddles0_re, twiddles0_im) = Self::duplicate_complex_components(twiddles[0]);
334        let (twiddles1_re, twiddles1_im) = Self::duplicate_complex_components(twiddles[1]);
335        let (twiddles2_re, twiddles2_im) = Self::duplicate_complex_components(twiddles[2]);
336        let (twiddles3_re, twiddles3_im) = Self::duplicate_complex_components(twiddles[3]);
337        let (twiddles4_re, twiddles4_im) = Self::duplicate_complex_components(twiddles[4]);
338
339        let twiddled1 = Self::fmadd(twiddles0_re, sum1, rows[0]);
340        let twiddled2 = Self::fmadd(twiddles1_re, sum1, rows[0]);
341        let twiddled3 = Self::fmadd(twiddles2_re, sum1, rows[0]);
342        let twiddled4 = Self::fmadd(twiddles3_re, sum1, rows[0]);
343        let twiddled5 = Self::fmadd(twiddles4_re, sum1, rows[0]);
344        let twiddled6 = Self::mul(twiddles4_im, rotated10);
345        let twiddled7 = Self::mul(twiddles3_im, rotated10);
346        let twiddled8 = Self::mul(twiddles2_im, rotated10);
347        let twiddled9 = Self::mul(twiddles1_im, rotated10);
348        let twiddled10 = Self::mul(twiddles0_im, rotated10);
349
350        let twiddled1 = Self::fmadd(twiddles1_re, sum2, twiddled1);
351        let twiddled2 = Self::fmadd(twiddles3_re, sum2, twiddled2);
352        let twiddled3 = Self::fmadd(twiddles4_re, sum2, twiddled3);
353        let twiddled4 = Self::fmadd(twiddles2_re, sum2, twiddled4);
354        let twiddled5 = Self::fmadd(twiddles0_re, sum2, twiddled5);
355        let twiddled6 = Self::fnmadd(twiddles0_im, rotated9, twiddled6);
356        let twiddled7 = Self::fnmadd(twiddles2_im, rotated9, twiddled7);
357        let twiddled8 = Self::fnmadd(twiddles4_im, rotated9, twiddled8);
358        let twiddled9 = Self::fmadd(twiddles3_im, rotated9, twiddled9);
359        let twiddled10 = Self::fmadd(twiddles1_im, rotated9, twiddled10);
360
361        let twiddled1 = Self::fmadd(twiddles2_re, sum3, twiddled1);
362        let twiddled2 = Self::fmadd(twiddles4_re, sum3, twiddled2);
363        let twiddled3 = Self::fmadd(twiddles1_re, sum3, twiddled3);
364        let twiddled4 = Self::fmadd(twiddles0_re, sum3, twiddled4);
365        let twiddled5 = Self::fmadd(twiddles3_re, sum3, twiddled5);
366        let twiddled6 = Self::fmadd(twiddles3_im, rotated8, twiddled6);
367        let twiddled7 = Self::fmadd(twiddles0_im, rotated8, twiddled7);
368        let twiddled8 = Self::fnmadd(twiddles1_im, rotated8, twiddled8);
369        let twiddled9 = Self::fnmadd(twiddles4_im, rotated8, twiddled9);
370        let twiddled10 = Self::fmadd(twiddles2_im, rotated8, twiddled10);
371
372        let twiddled1 = Self::fmadd(twiddles3_re, sum4, twiddled1);
373        let twiddled2 = Self::fmadd(twiddles2_re, sum4, twiddled2);
374        let twiddled3 = Self::fmadd(twiddles0_re, sum4, twiddled3);
375        let twiddled4 = Self::fmadd(twiddles4_re, sum4, twiddled4);
376        let twiddled5 = Self::fmadd(twiddles1_re, sum4, twiddled5);
377        let twiddled6 = Self::fnmadd(twiddles1_im, rotated7, twiddled6);
378        let twiddled7 = Self::fmadd(twiddles4_im, rotated7, twiddled7);
379        let twiddled8 = Self::fmadd(twiddles0_im, rotated7, twiddled8);
380        let twiddled9 = Self::fnmadd(twiddles2_im, rotated7, twiddled9);
381        let twiddled10 = Self::fmadd(twiddles3_im, rotated7, twiddled10);
382
383        let twiddled1 = Self::fmadd(twiddles4_re, sum5, twiddled1);
384        let twiddled2 = Self::fmadd(twiddles0_re, sum5, twiddled2);
385        let twiddled3 = Self::fmadd(twiddles3_re, sum5, twiddled3);
386        let twiddled4 = Self::fmadd(twiddles1_re, sum5, twiddled4);
387        let twiddled5 = Self::fmadd(twiddles2_re, sum5, twiddled5);
388        let twiddled6 = Self::fmadd(twiddles2_im, rotated6, twiddled6);
389        let twiddled7 = Self::fnmadd(twiddles1_im, rotated6, twiddled7);
390        let twiddled8 = Self::fmadd(twiddles3_im, rotated6, twiddled8);
391        let twiddled9 = Self::fnmadd(twiddles0_im, rotated6, twiddled9);
392        let twiddled10 = Self::fmadd(twiddles4_im, rotated6, twiddled10);
393
394        // Post-processing to mix the twiddle factors between the rest of the output
395        let [output1, output10] = Self::column_butterfly2([twiddled1, twiddled10]);
396        let [output2, output9] = Self::column_butterfly2([twiddled2, twiddled9]);
397        let [output3, output8] = Self::column_butterfly2([twiddled3, twiddled8]);
398        let [output4, output7] = Self::column_butterfly2([twiddled4, twiddled7]);
399        let [output5, output6] = Self::column_butterfly2([twiddled5, twiddled6]);
400
401        [
402            output0, output1, output2, output3, output4, output5, output6, output7, output8,
403            output9, output10,
404        ]
405    }
406}
407
408/// A 256-bit SIMD vector of complex numbers, stored with the real values and imaginary values interleaved.
409/// Implemented for __m256, __m256d
410///
411/// This trait implements things specific to 256-types, like splitting a 256 vector into 128 vectors
412/// For compiler-placation reasons, all interactions/awareness the scalar type go here
413pub trait AvxVector256: AvxVector {
414    type HalfVector: AvxVector128<FullVector = Self>;
415    type ScalarType: AvxNum<VectorType = Self>;
416
417    unsafe fn lo(self) -> Self::HalfVector;
418    unsafe fn hi(self) -> Self::HalfVector;
419    unsafe fn merge(lo: Self::HalfVector, hi: Self::HalfVector) -> Self;
420
421    /// Fill a vector by repeating the provided complex number as many times as possible
422    unsafe fn broadcast_complex_elements(value: Complex<Self::ScalarType>) -> Self;
423
424    // loads/stores of complex numbers
425    unsafe fn load_complex(ptr: *const Complex<Self::ScalarType>) -> Self;
426    unsafe fn store_complex(ptr: *mut Complex<Self::ScalarType>, data: Self);
427
428    // Gather 4 complex numbers (for f32) or 2 complex numbers (for f64) using 4 i32 indexes (for index32) or 4 i64 indexes (for index64).
429    // For f32, there should be 1 index per complex. For f64, there should be 2 indexes, each duplicated
430    // (So to load the complex<f64> at index 5 and 7, the index vector should contain 5,5,7,7. this api sucks but it's internal so whatever.)
431    unsafe fn gather_complex_avx2_index32(
432        ptr: *const Complex<Self::ScalarType>,
433        indexes: __m128i,
434    ) -> Self;
435    unsafe fn gather_complex_avx2_index64(
436        ptr: *const Complex<Self::ScalarType>,
437        indexes: __m256i,
438    ) -> Self;
439
440    // loads/stores of partial vectors of complex numbers. When loading, empty elements are zeroed
441    // unimplemented!() if Self::COMPLEX_PER_VECTOR is not greater than the partial count
442    unsafe fn load_partial1_complex(ptr: *const Complex<Self::ScalarType>) -> Self::HalfVector;
443    unsafe fn load_partial2_complex(ptr: *const Complex<Self::ScalarType>) -> Self::HalfVector;
444    unsafe fn load_partial3_complex(ptr: *const Complex<Self::ScalarType>) -> Self;
445    unsafe fn store_partial1_complex(ptr: *mut Complex<Self::ScalarType>, data: Self::HalfVector);
446    unsafe fn store_partial2_complex(ptr: *mut Complex<Self::ScalarType>, data: Self::HalfVector);
447    unsafe fn store_partial3_complex(ptr: *mut Complex<Self::ScalarType>, data: Self);
448
449    #[inline(always)]
450    unsafe fn column_butterfly6(rows: [Self; 6], twiddles: Self) -> [Self; 6] {
451        // Algorithm: 3x2 good-thomas
452
453        // Size-3 FFTs down the columns of our reordered array
454        let mid0 = Self::column_butterfly3([rows[0], rows[2], rows[4]], twiddles);
455        let mid1 = Self::column_butterfly3([rows[3], rows[5], rows[1]], twiddles);
456
457        // We normally would put twiddle factors right here, but since this is good-thomas algorithm, we don't need twiddle factors
458
459        // Transpose the data and do size-2 FFTs down the columns
460        let [output0, output1] = Self::column_butterfly2([mid0[0], mid1[0]]);
461        let [output2, output3] = Self::column_butterfly2([mid0[1], mid1[1]]);
462        let [output4, output5] = Self::column_butterfly2([mid0[2], mid1[2]]);
463
464        // Reorder into output
465        [output0, output3, output4, output1, output2, output5]
466    }
467
468    #[inline(always)]
469    unsafe fn column_butterfly9(
470        rows: [Self; 9],
471        twiddles: [Self; 3],
472        butterfly3_twiddles: Self,
473    ) -> [Self; 9] {
474        // Algorithm: 3x3 mixed radix
475
476        // Size-3 FFTs down the columns
477        let mid0 = Self::column_butterfly3([rows[0], rows[3], rows[6]], butterfly3_twiddles);
478        let mut mid1 = Self::column_butterfly3([rows[1], rows[4], rows[7]], butterfly3_twiddles);
479        let mut mid2 = Self::column_butterfly3([rows[2], rows[5], rows[8]], butterfly3_twiddles);
480
481        // Apply twiddle factors. Note that we're re-using twiddles[1]
482        mid1[1] = Self::mul_complex(twiddles[0], mid1[1]);
483        mid1[2] = Self::mul_complex(twiddles[1], mid1[2]);
484        mid2[1] = Self::mul_complex(twiddles[1], mid2[1]);
485        mid2[2] = Self::mul_complex(twiddles[2], mid2[2]);
486
487        let [output0, output1, output2] =
488            Self::column_butterfly3([mid0[0], mid1[0], mid2[0]], butterfly3_twiddles);
489        let [output3, output4, output5] =
490            Self::column_butterfly3([mid0[1], mid1[1], mid2[1]], butterfly3_twiddles);
491        let [output6, output7, output8] =
492            Self::column_butterfly3([mid0[2], mid1[2], mid2[2]], butterfly3_twiddles);
493
494        [
495            output0, output3, output6, output1, output4, output7, output2, output5, output8,
496        ]
497    }
498
499    #[inline(always)]
500    unsafe fn column_butterfly12(
501        rows: [Self; 12],
502        butterfly3_twiddles: Self,
503        rotation: Rotation90<Self>,
504    ) -> [Self; 12] {
505        // Algorithm: 4x3 good-thomas
506
507        // Size-4 FFTs down the columns of our reordered array
508        let mid0 = Self::column_butterfly4([rows[0], rows[3], rows[6], rows[9]], rotation);
509        let mid1 = Self::column_butterfly4([rows[4], rows[7], rows[10], rows[1]], rotation);
510        let mid2 = Self::column_butterfly4([rows[8], rows[11], rows[2], rows[5]], rotation);
511
512        // Since this is good-thomas algorithm, we don't need twiddle factors
513
514        // Transpose the data and do size-2 FFTs down the columns
515        let [output0, output1, output2] =
516            Self::column_butterfly3([mid0[0], mid1[0], mid2[0]], butterfly3_twiddles);
517        let [output3, output4, output5] =
518            Self::column_butterfly3([mid0[1], mid1[1], mid2[1]], butterfly3_twiddles);
519        let [output6, output7, output8] =
520            Self::column_butterfly3([mid0[2], mid1[2], mid2[2]], butterfly3_twiddles);
521        let [output9, output10, output11] =
522            Self::column_butterfly3([mid0[3], mid1[3], mid2[3]], butterfly3_twiddles);
523
524        [
525            output0, output4, output8, output9, output1, output5, output6, output10, output2,
526            output3, output7, output11,
527        ]
528    }
529}
530
531/// A 128-bit SIMD vector of complex numbers, stored with the real values and imaginary values interleaved.
532/// Implemented for __m128, __m128d, but these are all oriented around AVX, so don't call methods on these from a SSE-only context
533///
534/// This trait implements things specific to 128-types, like merging 2 128 vectors into a 256 vector
535pub trait AvxVector128: AvxVector {
536    type FullVector: AvxVector256<HalfVector = Self>;
537
538    unsafe fn merge(lo: Self, hi: Self) -> Self::FullVector;
539    unsafe fn zero_extend(self) -> Self::FullVector;
540
541    unsafe fn lo(input: Self::FullVector) -> Self;
542    unsafe fn hi(input: Self::FullVector) -> Self;
543    unsafe fn split(input: Self::FullVector) -> (Self, Self) {
544        (Self::lo(input), Self::hi(input))
545    }
546    unsafe fn lo_rotation(input: Rotation90<Self::FullVector>) -> Rotation90<Self>;
547
548    /// Fill a vector by repeating the provided complex number as many times as possible
549    unsafe fn broadcast_complex_elements(
550        value: Complex<<<Self as AvxVector128>::FullVector as AvxVector256>::ScalarType>,
551    ) -> Self;
552
553    // Gather 2 complex numbers (for f32) or 1 complex number (for f64) using 2 i32 indexes (for gather32) or 2 i64 indexes (for gather64).
554    // For f32, there should be 1 index per complex. For f64, there should be 2 indexes, each duplicated
555    // (So to load the complex<f64> at index 5, the index vector should contain 5,5. this api sucks but it's internal so whatever.)
556    unsafe fn gather32_complex_avx2(
557        ptr: *const Complex<<Self::FullVector as AvxVector256>::ScalarType>,
558        indexes: __m128i,
559    ) -> Self;
560    unsafe fn gather64_complex_avx2(
561        ptr: *const Complex<<Self::FullVector as AvxVector256>::ScalarType>,
562        indexes: __m128i,
563    ) -> Self;
564
565    #[inline(always)]
566    unsafe fn column_butterfly6(rows: [Self; 6], twiddles: Self::FullVector) -> [Self; 6] {
567        // Algorithm: 3x2 good-thomas
568
569        // if we merge some of our 128 registers into 256 registers, we can do 1 inner butterfly3 instead of 2
570        let rows03 = Self::merge(rows[0], rows[3]);
571        let rows25 = Self::merge(rows[2], rows[5]);
572        let rows41 = Self::merge(rows[4], rows[1]);
573
574        // Size-3 FFTs down the columns of our reordered array
575        let mid = Self::FullVector::column_butterfly3([rows03, rows25, rows41], twiddles);
576
577        // We normally would put twiddle factors right here, but since this is good-thomas algorithm, we don't need twiddle factors
578
579        // we can't use our merged columns anymore. so split them back into half vectors
580        let (mid0_0, mid1_0) = Self::split(mid[0]);
581        let (mid0_1, mid1_1) = Self::split(mid[1]);
582        let (mid0_2, mid1_2) = Self::split(mid[2]);
583
584        // Transpose the data and do size-2 FFTs down the columns
585        let [output0, output1] = Self::column_butterfly2([mid0_0, mid1_0]);
586        let [output2, output3] = Self::column_butterfly2([mid0_1, mid1_1]);
587        let [output4, output5] = Self::column_butterfly2([mid0_2, mid1_2]);
588
589        // Reorder into output
590        [output0, output3, output4, output1, output2, output5]
591    }
592
593    #[inline(always)]
594    unsafe fn column_butterfly9(
595        rows: [Self; 9],
596        twiddles_merged: [Self::FullVector; 2],
597        butterfly3_twiddles: Self::FullVector,
598    ) -> [Self; 9] {
599        // Algorithm: 3x3 mixed radix
600
601        // if we merge some of our 128 registers into 256 registers, we can do 2 inner butterfly3's instead of 3
602        let rows12 = Self::merge(rows[1], rows[2]);
603        let rows45 = Self::merge(rows[4], rows[5]);
604        let rows78 = Self::merge(rows[7], rows[8]);
605
606        let mid0 =
607            Self::column_butterfly3([rows[0], rows[3], rows[6]], Self::lo(butterfly3_twiddles));
608        let mut mid12 =
609            Self::FullVector::column_butterfly3([rows12, rows45, rows78], butterfly3_twiddles);
610
611        // Apply twiddle factors. we're applying them on the merged set of vectors, so we need slightly different twiddle factors
612        mid12[1] = Self::FullVector::mul_complex(twiddles_merged[0], mid12[1]);
613        mid12[2] = Self::FullVector::mul_complex(twiddles_merged[1], mid12[2]);
614
615        // we can't use our merged columns anymore. so split them back into half vectors
616        let (mid1_0, mid2_0) = Self::split(mid12[0]);
617        let (mid1_1, mid2_1) = Self::split(mid12[1]);
618        let (mid1_2, mid2_2) = Self::split(mid12[2]);
619
620        // Re-merge our half vectors into different, transposed full vectors. Thankfully the compiler is smart enough to combine these inserts and extracts into permutes
621        let transposed12 = Self::merge(mid0[1], mid0[2]);
622        let transposed45 = Self::merge(mid1_1, mid1_2);
623        let transposed78 = Self::merge(mid2_1, mid2_2);
624
625        let [output0, output1, output2] =
626            Self::column_butterfly3([mid0[0], mid1_0, mid2_0], Self::lo(butterfly3_twiddles));
627        let [output36, output47, output58] = Self::FullVector::column_butterfly3(
628            [transposed12, transposed45, transposed78],
629            butterfly3_twiddles,
630        );
631
632        // Finally, extract our second set of merged columns
633        let (output3, output6) = Self::split(output36);
634        let (output4, output7) = Self::split(output47);
635        let (output5, output8) = Self::split(output58);
636
637        [
638            output0, output3, output6, output1, output4, output7, output2, output5, output8,
639        ]
640    }
641
642    #[inline(always)]
643    unsafe fn column_butterfly12(
644        rows: [Self; 12],
645        butterfly3_twiddles: Self::FullVector,
646        rotation: Rotation90<Self::FullVector>,
647    ) -> [Self; 12] {
648        // Algorithm: 4x3 good-thomas
649
650        // if we merge some of our 128 registers into 256 registers, we can do 2 inner butterfly4's instead of 3
651        let rows48 = Self::merge(rows[4], rows[8]);
652        let rows711 = Self::merge(rows[7], rows[11]);
653        let rows102 = Self::merge(rows[10], rows[2]);
654        let rows15 = Self::merge(rows[1], rows[5]);
655
656        // Size-4 FFTs down the columns of our reordered array
657        let mid0 = Self::column_butterfly4(
658            [rows[0], rows[3], rows[6], rows[9]],
659            Self::lo_rotation(rotation),
660        );
661        let mid12 =
662            Self::FullVector::column_butterfly4([rows48, rows711, rows102, rows15], rotation);
663
664        // We normally would put twiddle factors right here, but since this is good-thomas algorithm, we don't need twiddle factors
665
666        // we can't use our merged columns anymore. so split them back into half vectors
667        let (mid1_0, mid2_0) = Self::split(mid12[0]);
668        let (mid1_1, mid2_1) = Self::split(mid12[1]);
669        let (mid1_2, mid2_2) = Self::split(mid12[2]);
670        let (mid1_3, mid2_3) = Self::split(mid12[3]);
671
672        // Re-merge our half vectors into different, transposed full vectors. This will let us do 2 inner butterfly 3's instead of 4!
673        // Thankfully the compiler is smart enough to combine these inserts and extracts into permutes
674        let transposed03 = Self::merge(mid0[0], mid0[1]);
675        let transposed14 = Self::merge(mid1_0, mid1_1);
676        let transposed25 = Self::merge(mid2_0, mid2_1);
677
678        let transposed69 = Self::merge(mid0[2], mid0[3]);
679        let transposed710 = Self::merge(mid1_2, mid1_3);
680        let transposed811 = Self::merge(mid2_2, mid2_3);
681
682        // Transpose the data and do size-2 FFTs down the columns
683        let [output03, output14, output25] = Self::FullVector::column_butterfly3(
684            [transposed03, transposed14, transposed25],
685            butterfly3_twiddles,
686        );
687        let [output69, output710, output811] = Self::FullVector::column_butterfly3(
688            [transposed69, transposed710, transposed811],
689            butterfly3_twiddles,
690        );
691
692        // Finally, extract our second set of merged columns
693        let (output0, output3) = Self::split(output03);
694        let (output1, output4) = Self::split(output14);
695        let (output2, output5) = Self::split(output25);
696        let (output6, output9) = Self::split(output69);
697        let (output7, output10) = Self::split(output710);
698        let (output8, output11) = Self::split(output811);
699
700        [
701            output0, output4, output8, output9, output1, output5, output6, output10, output2,
702            output3, output7, output11,
703        ]
704    }
705}
706
707#[inline(always)]
708pub unsafe fn apply_butterfly8_twiddle1<V: AvxVector>(input: V, rotation: Rotation90<V>) -> V {
709    let rotated = input.rotate90(rotation);
710    let combined = V::add(rotated, input);
711    V::mul(V::half_root2(), combined)
712}
713#[inline(always)]
714pub unsafe fn apply_butterfly8_twiddle3<V: AvxVector>(input: V, rotation: Rotation90<V>) -> V {
715    let rotated = input.rotate90(rotation);
716    let combined = V::sub(rotated, input);
717    V::mul(V::half_root2(), combined)
718}
719
720#[repr(transparent)]
721#[derive(Clone, Copy, Debug)]
722pub struct Rotation90<V>(V);
723impl<V: AvxVector256> Rotation90<V> {
724    #[inline(always)]
725    pub unsafe fn lo(self) -> Rotation90<V::HalfVector> {
726        Rotation90(self.0.lo())
727    }
728}
729
730impl AvxVector for __m256 {
731    const COMPLEX_PER_VECTOR: usize = 4;
732
733    #[inline(always)]
734    unsafe fn zero() -> Self {
735        _mm256_setzero_ps()
736    }
737    #[inline(always)]
738    unsafe fn half_root2() -> Self {
739        // note: we're computing a square root here, but checking the assembly says the compiler is smart enough to turn this into a constant
740        _mm256_broadcast_ss(&0.5f32.sqrt())
741    }
742
743    #[inline(always)]
744    unsafe fn xor(left: Self, right: Self) -> Self {
745        _mm256_xor_ps(left, right)
746    }
747    #[inline(always)]
748    unsafe fn neg(self) -> Self {
749        _mm256_xor_ps(self, _mm256_broadcast_ss(&-0.0))
750    }
751    #[inline(always)]
752    unsafe fn add(left: Self, right: Self) -> Self {
753        _mm256_add_ps(left, right)
754    }
755    #[inline(always)]
756    unsafe fn sub(left: Self, right: Self) -> Self {
757        _mm256_sub_ps(left, right)
758    }
759    #[inline(always)]
760    unsafe fn mul(left: Self, right: Self) -> Self {
761        _mm256_mul_ps(left, right)
762    }
763    #[inline(always)]
764    unsafe fn fmadd(left: Self, right: Self, add: Self) -> Self {
765        _mm256_fmadd_ps(left, right, add)
766    }
767    #[inline(always)]
768    unsafe fn fnmadd(left: Self, right: Self, add: Self) -> Self {
769        _mm256_fnmadd_ps(left, right, add)
770    }
771    #[inline(always)]
772    unsafe fn fmaddsub(left: Self, right: Self, add: Self) -> Self {
773        _mm256_fmaddsub_ps(left, right, add)
774    }
775    #[inline(always)]
776    unsafe fn fmsubadd(left: Self, right: Self, add: Self) -> Self {
777        _mm256_fmsubadd_ps(left, right, add)
778    }
779    #[inline(always)]
780    unsafe fn reverse_complex_elements(self) -> Self {
781        // swap the elements in-lane
782        let permuted = _mm256_permute_ps(self, 0x4E);
783        // swap the lanes
784        _mm256_permute2f128_ps(permuted, permuted, 0x01)
785    }
786    #[inline(always)]
787    unsafe fn unpacklo_complex(rows: [Self; 2]) -> Self {
788        let row0_double = _mm256_castps_pd(rows[0]);
789        let row1_double = _mm256_castps_pd(rows[1]);
790        let unpacked = _mm256_unpacklo_pd(row0_double, row1_double);
791        _mm256_castpd_ps(unpacked)
792    }
793    #[inline(always)]
794    unsafe fn unpackhi_complex(rows: [Self; 2]) -> Self {
795        let row0_double = _mm256_castps_pd(rows[0]);
796        let row1_double = _mm256_castps_pd(rows[1]);
797        let unpacked = _mm256_unpackhi_pd(row0_double, row1_double);
798        _mm256_castpd_ps(unpacked)
799    }
800
801    #[inline(always)]
802    unsafe fn swap_complex_components(self) -> Self {
803        _mm256_permute_ps(self, 0xB1)
804    }
805
806    #[inline(always)]
807    unsafe fn duplicate_complex_components(self) -> (Self, Self) {
808        (_mm256_moveldup_ps(self), _mm256_movehdup_ps(self))
809    }
810
811    #[inline(always)]
812    unsafe fn make_rotation90(direction: FftDirection) -> Rotation90<Self> {
813        let broadcast = match direction {
814            FftDirection::Forward => Complex::new(-0.0, 0.0),
815            FftDirection::Inverse => Complex::new(0.0, -0.0),
816        };
817        Rotation90(Self::broadcast_complex_elements(broadcast))
818    }
819
820    #[inline(always)]
821    unsafe fn make_mixedradix_twiddle_chunk(
822        x: usize,
823        y: usize,
824        len: usize,
825        direction: FftDirection,
826    ) -> Self {
827        let mut twiddle_chunk = [Complex::<f32>::zero(); Self::COMPLEX_PER_VECTOR];
828        for i in 0..Self::COMPLEX_PER_VECTOR {
829            twiddle_chunk[i] = twiddles::compute_twiddle(y * (x + i), len, direction);
830        }
831
832        twiddle_chunk.as_slice().load_complex(0)
833    }
834
835    #[inline(always)]
836    unsafe fn broadcast_twiddle(index: usize, len: usize, direction: FftDirection) -> Self {
837        Self::broadcast_complex_elements(twiddles::compute_twiddle(index, len, direction))
838    }
839
840    #[inline(always)]
841    unsafe fn transpose2_packed(rows: [Self; 2]) -> [Self; 2] {
842        let unpacked = Self::unpack_complex(rows);
843        let output0 = _mm256_permute2f128_ps(unpacked[0], unpacked[1], 0x20);
844        let output1 = _mm256_permute2f128_ps(unpacked[0], unpacked[1], 0x31);
845
846        [output0, output1]
847    }
848    #[inline(always)]
849    unsafe fn transpose3_packed(rows: [Self; 3]) -> [Self; 3] {
850        let unpacked0 = Self::unpacklo_complex([rows[0], rows[1]]);
851        let unpacked2 = Self::unpackhi_complex([rows[1], rows[2]]);
852
853        // output0 and output2 each need to swap some elements. thankfully we can blend those elements into the same intermediate value, and then do a permute 128 from there
854        let blended = _mm256_blend_ps(rows[0], rows[2], 0x33);
855
856        let output1 = _mm256_permute2f128_ps(unpacked0, unpacked2, 0x12);
857
858        let output0 = _mm256_permute2f128_ps(unpacked0, blended, 0x20);
859        let output2 = _mm256_permute2f128_ps(unpacked2, blended, 0x13);
860
861        [output0, output1, output2]
862    }
863    #[inline(always)]
864    unsafe fn transpose4_packed(rows: [Self; 4]) -> [Self; 4] {
865        let permute0 = _mm256_permute2f128_ps(rows[0], rows[2], 0x20);
866        let permute1 = _mm256_permute2f128_ps(rows[1], rows[3], 0x20);
867        let permute2 = _mm256_permute2f128_ps(rows[0], rows[2], 0x31);
868        let permute3 = _mm256_permute2f128_ps(rows[1], rows[3], 0x31);
869
870        let [unpacked0, unpacked1] = Self::unpack_complex([permute0, permute1]);
871        let [unpacked2, unpacked3] = Self::unpack_complex([permute2, permute3]);
872
873        [unpacked0, unpacked1, unpacked2, unpacked3]
874    }
875    #[inline(always)]
876    unsafe fn transpose5_packed(rows: [Self; 5]) -> [Self; 5] {
877        let unpacked0 = Self::unpacklo_complex([rows[0], rows[1]]);
878        let unpacked1 = Self::unpackhi_complex([rows[1], rows[2]]);
879        let unpacked2 = Self::unpacklo_complex([rows[2], rows[3]]);
880        let unpacked3 = Self::unpackhi_complex([rows[3], rows[4]]);
881        let blended04 = _mm256_blend_ps(rows[0], rows[4], 0x33);
882
883        [
884            _mm256_permute2f128_ps(unpacked0, unpacked2, 0x20),
885            _mm256_permute2f128_ps(blended04, unpacked1, 0x20),
886            _mm256_blend_ps(unpacked0, unpacked3, 0x0f),
887            _mm256_permute2f128_ps(unpacked2, blended04, 0x31),
888            _mm256_permute2f128_ps(unpacked1, unpacked3, 0x31),
889        ]
890    }
891    #[inline(always)]
892    unsafe fn transpose6_packed(rows: [Self; 6]) -> [Self; 6] {
893        let [unpacked0, unpacked1] = Self::unpack_complex([rows[0], rows[1]]);
894        let [unpacked2, unpacked3] = Self::unpack_complex([rows[2], rows[3]]);
895        let [unpacked4, unpacked5] = Self::unpack_complex([rows[4], rows[5]]);
896
897        [
898            _mm256_permute2f128_ps(unpacked0, unpacked2, 0x20),
899            _mm256_permute2f128_ps(unpacked1, unpacked4, 0x02),
900            _mm256_permute2f128_ps(unpacked3, unpacked5, 0x20),
901            _mm256_permute2f128_ps(unpacked0, unpacked2, 0x31),
902            _mm256_permute2f128_ps(unpacked1, unpacked4, 0x13),
903            _mm256_permute2f128_ps(unpacked3, unpacked5, 0x31),
904        ]
905    }
906    #[inline(always)]
907    unsafe fn transpose7_packed(rows: [Self; 7]) -> [Self; 7] {
908        let unpacked0 = Self::unpacklo_complex([rows[0], rows[1]]);
909        let unpacked1 = Self::unpackhi_complex([rows[1], rows[2]]);
910        let unpacked2 = Self::unpacklo_complex([rows[2], rows[3]]);
911        let unpacked3 = Self::unpackhi_complex([rows[3], rows[4]]);
912        let unpacked4 = Self::unpacklo_complex([rows[4], rows[5]]);
913        let unpacked5 = Self::unpackhi_complex([rows[5], rows[6]]);
914        let blended06 = _mm256_blend_ps(rows[0], rows[6], 0x33);
915
916        [
917            _mm256_permute2f128_ps(unpacked0, unpacked2, 0x20),
918            _mm256_permute2f128_ps(unpacked4, blended06, 0x20),
919            _mm256_permute2f128_ps(unpacked1, unpacked3, 0x20),
920            _mm256_blend_ps(unpacked0, unpacked5, 0x0f),
921            _mm256_permute2f128_ps(unpacked2, unpacked4, 0x31),
922            _mm256_permute2f128_ps(blended06, unpacked1, 0x31),
923            _mm256_permute2f128_ps(unpacked3, unpacked5, 0x31),
924        ]
925    }
926    #[inline(always)]
927    unsafe fn transpose8_packed(rows: [Self; 8]) -> [Self; 8] {
928        let chunk0 = [rows[0], rows[1], rows[2], rows[3]];
929        let chunk1 = [rows[4], rows[5], rows[6], rows[7]];
930
931        let output0 = Self::transpose4_packed(chunk0);
932        let output1 = Self::transpose4_packed(chunk1);
933
934        [
935            output0[0], output1[0], output0[1], output1[1], output0[2], output1[2], output0[3],
936            output1[3],
937        ]
938    }
939    #[inline(always)]
940    unsafe fn transpose9_packed(rows: [Self; 9]) -> [Self; 9] {
941        let unpacked0 = Self::unpacklo_complex([rows[0], rows[1]]);
942        let unpacked1 = Self::unpackhi_complex([rows[1], rows[2]]);
943        let unpacked2 = Self::unpacklo_complex([rows[2], rows[3]]);
944        let unpacked3 = Self::unpackhi_complex([rows[3], rows[4]]);
945        let unpacked5 = Self::unpacklo_complex([rows[4], rows[5]]);
946        let unpacked6 = Self::unpackhi_complex([rows[5], rows[6]]);
947        let unpacked7 = Self::unpacklo_complex([rows[6], rows[7]]);
948        let unpacked8 = Self::unpackhi_complex([rows[7], rows[8]]);
949        let blended9 = _mm256_blend_ps(rows[0], rows[8], 0x33);
950
951        [
952            _mm256_permute2f128_ps(unpacked0, unpacked2, 0x20),
953            _mm256_permute2f128_ps(unpacked5, unpacked7, 0x20),
954            _mm256_permute2f128_ps(blended9, unpacked1, 0x20),
955            _mm256_permute2f128_ps(unpacked3, unpacked6, 0x20),
956            _mm256_blend_ps(unpacked0, unpacked8, 0x0f),
957            _mm256_permute2f128_ps(unpacked2, unpacked5, 0x31),
958            _mm256_permute2f128_ps(unpacked7, blended9, 0x31),
959            _mm256_permute2f128_ps(unpacked1, unpacked3, 0x31),
960            _mm256_permute2f128_ps(unpacked6, unpacked8, 0x31),
961        ]
962    }
963    #[inline(always)]
964    unsafe fn transpose11_packed(rows: [Self; 11]) -> [Self; 11] {
965        let unpacked0 = Self::unpacklo_complex([rows[0], rows[1]]);
966        let unpacked1 = Self::unpackhi_complex([rows[1], rows[2]]);
967        let unpacked2 = Self::unpacklo_complex([rows[2], rows[3]]);
968        let unpacked3 = Self::unpackhi_complex([rows[3], rows[4]]);
969        let unpacked4 = Self::unpacklo_complex([rows[4], rows[5]]);
970        let unpacked5 = Self::unpackhi_complex([rows[5], rows[6]]);
971        let unpacked6 = Self::unpacklo_complex([rows[6], rows[7]]);
972        let unpacked7 = Self::unpackhi_complex([rows[7], rows[8]]);
973        let unpacked8 = Self::unpacklo_complex([rows[8], rows[9]]);
974        let unpacked9 = Self::unpackhi_complex([rows[9], rows[10]]);
975        let blended10 = _mm256_blend_ps(rows[0], rows[10], 0x33);
976
977        [
978            _mm256_permute2f128_ps(unpacked0, unpacked2, 0x20),
979            _mm256_permute2f128_ps(unpacked4, unpacked6, 0x20),
980            _mm256_permute2f128_ps(unpacked8, blended10, 0x20),
981            _mm256_permute2f128_ps(unpacked1, unpacked3, 0x20),
982            _mm256_permute2f128_ps(unpacked5, unpacked7, 0x20),
983            _mm256_blend_ps(unpacked0, unpacked9, 0x0f),
984            _mm256_permute2f128_ps(unpacked2, unpacked4, 0x31),
985            _mm256_permute2f128_ps(unpacked6, unpacked8, 0x31),
986            _mm256_permute2f128_ps(blended10, unpacked1, 0x31),
987            _mm256_permute2f128_ps(unpacked3, unpacked5, 0x31),
988            _mm256_permute2f128_ps(unpacked7, unpacked9, 0x31),
989        ]
990    }
991    #[inline(always)]
992    unsafe fn transpose12_packed(rows: [Self; 12]) -> [Self; 12] {
993        let chunk0 = [rows[0], rows[1], rows[2], rows[3]];
994        let chunk1 = [rows[4], rows[5], rows[6], rows[7]];
995        let chunk2 = [rows[8], rows[9], rows[10], rows[11]];
996
997        let output0 = Self::transpose4_packed(chunk0);
998        let output1 = Self::transpose4_packed(chunk1);
999        let output2 = Self::transpose4_packed(chunk2);
1000
1001        [
1002            output0[0], output1[0], output2[0], output0[1], output1[1], output2[1], output0[2],
1003            output1[2], output2[2], output0[3], output1[3], output2[3],
1004        ]
1005    }
1006    #[inline(always)]
1007    unsafe fn transpose16_packed(rows: [Self; 16]) -> [Self; 16] {
1008        let chunk0 = [
1009            rows[0], rows[1], rows[2], rows[3], rows[4], rows[5], rows[6], rows[7],
1010        ];
1011        let chunk1 = [
1012            rows[8], rows[9], rows[10], rows[11], rows[12], rows[13], rows[14], rows[15],
1013        ];
1014
1015        let output0 = Self::transpose8_packed(chunk0);
1016        let output1 = Self::transpose8_packed(chunk1);
1017
1018        [
1019            output0[0], output0[1], output1[0], output1[1], output0[2], output0[3], output1[2],
1020            output1[3], output0[4], output0[5], output1[4], output1[5], output0[6], output0[7],
1021            output1[6], output1[7],
1022        ]
1023    }
1024}
1025impl AvxVector256 for __m256 {
1026    type ScalarType = f32;
1027    type HalfVector = __m128;
1028
1029    #[inline(always)]
1030    unsafe fn lo(self) -> Self::HalfVector {
1031        _mm256_castps256_ps128(self)
1032    }
1033    #[inline(always)]
1034    unsafe fn hi(self) -> Self::HalfVector {
1035        _mm256_extractf128_ps(self, 1)
1036    }
1037    #[inline(always)]
1038    unsafe fn merge(lo: Self::HalfVector, hi: Self::HalfVector) -> Self {
1039        _mm256_insertf128_ps(_mm256_castps128_ps256(lo), hi, 1)
1040    }
1041
1042    #[inline(always)]
1043    unsafe fn broadcast_complex_elements(value: Complex<Self::ScalarType>) -> Self {
1044        _mm256_set_ps(
1045            value.im, value.re, value.im, value.re, value.im, value.re, value.im, value.re,
1046        )
1047    }
1048
1049    #[inline(always)]
1050    unsafe fn load_complex(ptr: *const Complex<Self::ScalarType>) -> Self {
1051        _mm256_loadu_ps(ptr as *const Self::ScalarType)
1052    }
1053    #[inline(always)]
1054    unsafe fn store_complex(ptr: *mut Complex<Self::ScalarType>, data: Self) {
1055        _mm256_storeu_ps(ptr as *mut Self::ScalarType, data)
1056    }
1057    #[inline(always)]
1058    unsafe fn gather_complex_avx2_index32(
1059        ptr: *const Complex<Self::ScalarType>,
1060        indexes: __m128i,
1061    ) -> Self {
1062        _mm256_castpd_ps(_mm256_i32gather_pd(ptr as *const f64, indexes, 8))
1063    }
1064    #[inline(always)]
1065    unsafe fn gather_complex_avx2_index64(
1066        ptr: *const Complex<Self::ScalarType>,
1067        indexes: __m256i,
1068    ) -> Self {
1069        _mm256_castpd_ps(_mm256_i64gather_pd(ptr as *const f64, indexes, 8))
1070    }
1071
1072    #[inline(always)]
1073    unsafe fn load_partial1_complex(ptr: *const Complex<Self::ScalarType>) -> Self::HalfVector {
1074        let data = _mm_load_sd(ptr as *const f64);
1075        _mm_castpd_ps(data)
1076    }
1077    #[inline(always)]
1078    unsafe fn load_partial2_complex(ptr: *const Complex<Self::ScalarType>) -> Self::HalfVector {
1079        _mm_loadu_ps(ptr as *const f32)
1080    }
1081    #[inline(always)]
1082    unsafe fn load_partial3_complex(ptr: *const Complex<Self::ScalarType>) -> Self {
1083        let lo = Self::load_partial2_complex(ptr);
1084        let hi = Self::load_partial1_complex(ptr.add(2));
1085        Self::merge(lo, hi)
1086    }
1087    #[inline(always)]
1088    unsafe fn store_partial1_complex(ptr: *mut Complex<Self::ScalarType>, data: Self::HalfVector) {
1089        _mm_store_sd(ptr as *mut f64, _mm_castps_pd(data));
1090    }
1091    #[inline(always)]
1092    unsafe fn store_partial2_complex(ptr: *mut Complex<Self::ScalarType>, data: Self::HalfVector) {
1093        _mm_storeu_ps(ptr as *mut f32, data);
1094    }
1095    #[inline(always)]
1096    unsafe fn store_partial3_complex(ptr: *mut Complex<Self::ScalarType>, data: Self) {
1097        Self::store_partial2_complex(ptr, data.lo());
1098        Self::store_partial1_complex(ptr.add(2), data.hi());
1099    }
1100}
1101
1102impl AvxVector for __m128 {
1103    const COMPLEX_PER_VECTOR: usize = 2;
1104
1105    #[inline(always)]
1106    unsafe fn zero() -> Self {
1107        _mm_setzero_ps()
1108    }
1109    #[inline(always)]
1110    unsafe fn half_root2() -> Self {
1111        // note: we're computing a square root here, but checking the assembly says the compiler is smart enough to turn this into a constant
1112        _mm_broadcast_ss(&0.5f32.sqrt())
1113    }
1114
1115    #[inline(always)]
1116    unsafe fn xor(left: Self, right: Self) -> Self {
1117        _mm_xor_ps(left, right)
1118    }
1119    #[inline(always)]
1120    unsafe fn neg(self) -> Self {
1121        _mm_xor_ps(self, _mm_broadcast_ss(&-0.0))
1122    }
1123    #[inline(always)]
1124    unsafe fn add(left: Self, right: Self) -> Self {
1125        _mm_add_ps(left, right)
1126    }
1127    #[inline(always)]
1128    unsafe fn sub(left: Self, right: Self) -> Self {
1129        _mm_sub_ps(left, right)
1130    }
1131    #[inline(always)]
1132    unsafe fn mul(left: Self, right: Self) -> Self {
1133        _mm_mul_ps(left, right)
1134    }
1135    #[inline(always)]
1136    unsafe fn fmadd(left: Self, right: Self, add: Self) -> Self {
1137        _mm_fmadd_ps(left, right, add)
1138    }
1139    #[inline(always)]
1140    unsafe fn fnmadd(left: Self, right: Self, add: Self) -> Self {
1141        _mm_fnmadd_ps(left, right, add)
1142    }
1143    #[inline(always)]
1144    unsafe fn fmaddsub(left: Self, right: Self, add: Self) -> Self {
1145        _mm_fmaddsub_ps(left, right, add)
1146    }
1147    #[inline(always)]
1148    unsafe fn fmsubadd(left: Self, right: Self, add: Self) -> Self {
1149        _mm_fmsubadd_ps(left, right, add)
1150    }
1151
1152    #[inline(always)]
1153    unsafe fn reverse_complex_elements(self) -> Self {
1154        // swap the elements in-lane
1155        _mm_permute_ps(self, 0x4E)
1156    }
1157
1158    #[inline(always)]
1159    unsafe fn unpacklo_complex(rows: [Self; 2]) -> Self {
1160        let row0_double = _mm_castps_pd(rows[0]);
1161        let row1_double = _mm_castps_pd(rows[1]);
1162        let unpacked = _mm_unpacklo_pd(row0_double, row1_double);
1163        _mm_castpd_ps(unpacked)
1164    }
1165    #[inline(always)]
1166    unsafe fn unpackhi_complex(rows: [Self; 2]) -> Self {
1167        let row0_double = _mm_castps_pd(rows[0]);
1168        let row1_double = _mm_castps_pd(rows[1]);
1169        let unpacked = _mm_unpackhi_pd(row0_double, row1_double);
1170        _mm_castpd_ps(unpacked)
1171    }
1172
1173    #[inline(always)]
1174    unsafe fn swap_complex_components(self) -> Self {
1175        _mm_permute_ps(self, 0xB1)
1176    }
1177    #[inline(always)]
1178    unsafe fn duplicate_complex_components(self) -> (Self, Self) {
1179        (_mm_moveldup_ps(self), _mm_movehdup_ps(self))
1180    }
1181
1182    #[inline(always)]
1183    unsafe fn make_rotation90(direction: FftDirection) -> Rotation90<Self> {
1184        let broadcast = match direction {
1185            FftDirection::Forward => Complex::new(-0.0, 0.0),
1186            FftDirection::Inverse => Complex::new(0.0, -0.0),
1187        };
1188        Rotation90(Self::broadcast_complex_elements(broadcast))
1189    }
1190    #[inline(always)]
1191    unsafe fn make_mixedradix_twiddle_chunk(
1192        x: usize,
1193        y: usize,
1194        len: usize,
1195        direction: FftDirection,
1196    ) -> Self {
1197        let mut twiddle_chunk = [Complex::<f32>::zero(); Self::COMPLEX_PER_VECTOR];
1198        for i in 0..Self::COMPLEX_PER_VECTOR {
1199            twiddle_chunk[i] = twiddles::compute_twiddle(y * (x + i), len, direction);
1200        }
1201
1202        _mm_loadu_ps(twiddle_chunk.as_ptr() as *const f32)
1203    }
1204    #[inline(always)]
1205    unsafe fn broadcast_twiddle(index: usize, len: usize, direction: FftDirection) -> Self {
1206        Self::broadcast_complex_elements(twiddles::compute_twiddle(index, len, direction))
1207    }
1208
1209    #[inline(always)]
1210    unsafe fn transpose2_packed(rows: [Self; 2]) -> [Self; 2] {
1211        Self::unpack_complex(rows)
1212    }
1213    #[inline(always)]
1214    unsafe fn transpose3_packed(rows: [Self; 3]) -> [Self; 3] {
1215        let unpacked0 = Self::unpacklo_complex([rows[0], rows[1]]);
1216        let blended = _mm_blend_ps(rows[0], rows[2], 0x03);
1217        let unpacked2 = Self::unpackhi_complex([rows[1], rows[2]]);
1218
1219        [unpacked0, blended, unpacked2]
1220    }
1221    #[inline(always)]
1222    unsafe fn transpose4_packed(rows: [Self; 4]) -> [Self; 4] {
1223        let [unpacked0, unpacked1] = Self::unpack_complex([rows[0], rows[1]]);
1224        let [unpacked2, unpacked3] = Self::unpack_complex([rows[2], rows[3]]);
1225
1226        [unpacked0, unpacked2, unpacked1, unpacked3]
1227    }
1228    #[inline(always)]
1229    unsafe fn transpose5_packed(rows: [Self; 5]) -> [Self; 5] {
1230        [
1231            Self::unpacklo_complex([rows[0], rows[1]]),
1232            Self::unpacklo_complex([rows[2], rows[3]]),
1233            _mm_blend_ps(rows[0], rows[4], 0x03),
1234            Self::unpackhi_complex([rows[1], rows[2]]),
1235            Self::unpackhi_complex([rows[3], rows[4]]),
1236        ]
1237    }
1238    #[inline(always)]
1239    unsafe fn transpose6_packed(rows: [Self; 6]) -> [Self; 6] {
1240        let [unpacked0, unpacked1] = Self::unpack_complex([rows[0], rows[1]]);
1241        let [unpacked2, unpacked3] = Self::unpack_complex([rows[2], rows[3]]);
1242        let [unpacked4, unpacked5] = Self::unpack_complex([rows[4], rows[5]]);
1243
1244        [
1245            unpacked0, unpacked2, unpacked4, unpacked1, unpacked3, unpacked5,
1246        ]
1247    }
1248    #[inline(always)]
1249    unsafe fn transpose7_packed(rows: [Self; 7]) -> [Self; 7] {
1250        [
1251            Self::unpacklo_complex([rows[0], rows[1]]),
1252            Self::unpacklo_complex([rows[2], rows[3]]),
1253            Self::unpacklo_complex([rows[4], rows[5]]),
1254            _mm_shuffle_ps(rows[6], rows[0], 0xE4),
1255            Self::unpackhi_complex([rows[1], rows[2]]),
1256            Self::unpackhi_complex([rows[3], rows[4]]),
1257            Self::unpackhi_complex([rows[5], rows[6]]),
1258        ]
1259    }
1260    #[inline(always)]
1261    unsafe fn transpose8_packed(rows: [Self; 8]) -> [Self; 8] {
1262        let chunk0 = [rows[0], rows[1], rows[2], rows[3]];
1263        let chunk1 = [rows[4], rows[5], rows[6], rows[7]];
1264
1265        let output0 = Self::transpose4_packed(chunk0);
1266        let output1 = Self::transpose4_packed(chunk1);
1267
1268        [
1269            output0[0], output0[1], output1[0], output1[1], output0[2], output0[3], output1[2],
1270            output1[3],
1271        ]
1272    }
1273    #[inline(always)]
1274    unsafe fn transpose9_packed(rows: [Self; 9]) -> [Self; 9] {
1275        [
1276            Self::unpacklo_complex([rows[0], rows[1]]),
1277            Self::unpacklo_complex([rows[2], rows[3]]),
1278            Self::unpacklo_complex([rows[4], rows[5]]),
1279            Self::unpacklo_complex([rows[6], rows[7]]),
1280            _mm_shuffle_ps(rows[8], rows[0], 0xE4),
1281            Self::unpackhi_complex([rows[1], rows[2]]),
1282            Self::unpackhi_complex([rows[3], rows[4]]),
1283            Self::unpackhi_complex([rows[5], rows[6]]),
1284            Self::unpackhi_complex([rows[7], rows[8]]),
1285        ]
1286    }
1287    #[inline(always)]
1288    unsafe fn transpose11_packed(rows: [Self; 11]) -> [Self; 11] {
1289        [
1290            Self::unpacklo_complex([rows[0], rows[1]]),
1291            Self::unpacklo_complex([rows[2], rows[3]]),
1292            Self::unpacklo_complex([rows[4], rows[5]]),
1293            Self::unpacklo_complex([rows[6], rows[7]]),
1294            Self::unpacklo_complex([rows[8], rows[9]]),
1295            _mm_shuffle_ps(rows[10], rows[0], 0xE4),
1296            Self::unpackhi_complex([rows[1], rows[2]]),
1297            Self::unpackhi_complex([rows[3], rows[4]]),
1298            Self::unpackhi_complex([rows[5], rows[6]]),
1299            Self::unpackhi_complex([rows[7], rows[8]]),
1300            Self::unpackhi_complex([rows[9], rows[10]]),
1301        ]
1302    }
1303    #[inline(always)]
1304    unsafe fn transpose12_packed(rows: [Self; 12]) -> [Self; 12] {
1305        let chunk0 = [rows[0], rows[1], rows[2], rows[3]];
1306        let chunk1 = [rows[4], rows[5], rows[6], rows[7]];
1307        let chunk2 = [rows[8], rows[9], rows[10], rows[11]];
1308
1309        let output0 = Self::transpose4_packed(chunk0);
1310        let output1 = Self::transpose4_packed(chunk1);
1311        let output2 = Self::transpose4_packed(chunk2);
1312
1313        [
1314            output0[0], output0[1], output1[0], output1[1], output2[0], output2[1], output0[2],
1315            output0[3], output1[2], output1[3], output2[2], output2[3],
1316        ]
1317    }
1318    #[inline(always)]
1319    unsafe fn transpose16_packed(rows: [Self; 16]) -> [Self; 16] {
1320        let chunk0 = [
1321            rows[0], rows[1], rows[2], rows[3], rows[4], rows[5], rows[6], rows[7],
1322        ];
1323        let chunk1 = [
1324            rows[8], rows[9], rows[10], rows[11], rows[12], rows[13], rows[14], rows[15],
1325        ];
1326
1327        let output0 = Self::transpose8_packed(chunk0);
1328        let output1 = Self::transpose8_packed(chunk1);
1329
1330        [
1331            output0[0], output0[1], output0[2], output0[3], output1[0], output1[1], output1[2],
1332            output1[3], output0[4], output0[5], output0[6], output0[7], output1[4], output1[5],
1333            output1[6], output1[7],
1334        ]
1335    }
1336}
1337impl AvxVector128 for __m128 {
1338    type FullVector = __m256;
1339
1340    #[inline(always)]
1341    unsafe fn lo(input: Self::FullVector) -> Self {
1342        _mm256_castps256_ps128(input)
1343    }
1344    #[inline(always)]
1345    unsafe fn hi(input: Self::FullVector) -> Self {
1346        _mm256_extractf128_ps(input, 1)
1347    }
1348    #[inline(always)]
1349    unsafe fn merge(lo: Self, hi: Self) -> Self::FullVector {
1350        _mm256_insertf128_ps(_mm256_castps128_ps256(lo), hi, 1)
1351    }
1352    #[inline(always)]
1353    unsafe fn zero_extend(self) -> Self::FullVector {
1354        _mm256_zextps128_ps256(self)
1355    }
1356    #[inline(always)]
1357    unsafe fn lo_rotation(input: Rotation90<Self::FullVector>) -> Rotation90<Self> {
1358        input.lo()
1359    }
1360    #[inline(always)]
1361    unsafe fn broadcast_complex_elements(value: Complex<f32>) -> Self {
1362        _mm_set_ps(value.im, value.re, value.im, value.re)
1363    }
1364    #[inline(always)]
1365    unsafe fn gather32_complex_avx2(ptr: *const Complex<f32>, indexes: __m128i) -> Self {
1366        _mm_castpd_ps(_mm_i32gather_pd(ptr as *const f64, indexes, 8))
1367    }
1368    #[inline(always)]
1369    unsafe fn gather64_complex_avx2(ptr: *const Complex<f32>, indexes: __m128i) -> Self {
1370        _mm_castpd_ps(_mm_i64gather_pd(ptr as *const f64, indexes, 8))
1371    }
1372}
1373
1374impl AvxVector for __m256d {
1375    const COMPLEX_PER_VECTOR: usize = 2;
1376
1377    #[inline(always)]
1378    unsafe fn zero() -> Self {
1379        _mm256_setzero_pd()
1380    }
1381    #[inline(always)]
1382    unsafe fn half_root2() -> Self {
1383        // note: we're computing a square root here, but checking the assembly says the compiler is smart enough to turn this into a constant
1384        _mm256_broadcast_sd(&0.5f64.sqrt())
1385    }
1386
1387    #[inline(always)]
1388    unsafe fn xor(left: Self, right: Self) -> Self {
1389        _mm256_xor_pd(left, right)
1390    }
1391    #[inline(always)]
1392    unsafe fn neg(self) -> Self {
1393        _mm256_xor_pd(self, _mm256_broadcast_sd(&-0.0))
1394    }
1395    #[inline(always)]
1396    unsafe fn add(left: Self, right: Self) -> Self {
1397        _mm256_add_pd(left, right)
1398    }
1399    #[inline(always)]
1400    unsafe fn sub(left: Self, right: Self) -> Self {
1401        _mm256_sub_pd(left, right)
1402    }
1403    #[inline(always)]
1404    unsafe fn mul(left: Self, right: Self) -> Self {
1405        _mm256_mul_pd(left, right)
1406    }
1407    #[inline(always)]
1408    unsafe fn fmadd(left: Self, right: Self, add: Self) -> Self {
1409        _mm256_fmadd_pd(left, right, add)
1410    }
1411    #[inline(always)]
1412    unsafe fn fnmadd(left: Self, right: Self, add: Self) -> Self {
1413        _mm256_fnmadd_pd(left, right, add)
1414    }
1415    #[inline(always)]
1416    unsafe fn fmaddsub(left: Self, right: Self, add: Self) -> Self {
1417        _mm256_fmaddsub_pd(left, right, add)
1418    }
1419    #[inline(always)]
1420    unsafe fn fmsubadd(left: Self, right: Self, add: Self) -> Self {
1421        _mm256_fmsubadd_pd(left, right, add)
1422    }
1423
1424    #[inline(always)]
1425    unsafe fn reverse_complex_elements(self) -> Self {
1426        _mm256_permute2f128_pd(self, self, 0x01)
1427    }
1428    #[inline(always)]
1429    unsafe fn unpacklo_complex(rows: [Self; 2]) -> Self {
1430        _mm256_permute2f128_pd(rows[0], rows[1], 0x20)
1431    }
1432    #[inline(always)]
1433    unsafe fn unpackhi_complex(rows: [Self; 2]) -> Self {
1434        _mm256_permute2f128_pd(rows[0], rows[1], 0x31)
1435    }
1436
1437    #[inline(always)]
1438    unsafe fn swap_complex_components(self) -> Self {
1439        _mm256_permute_pd(self, 0x05)
1440    }
1441    #[inline(always)]
1442    unsafe fn duplicate_complex_components(self) -> (Self, Self) {
1443        (_mm256_movedup_pd(self), _mm256_permute_pd(self, 0x0F))
1444    }
1445
1446    #[inline(always)]
1447    unsafe fn make_rotation90(direction: FftDirection) -> Rotation90<Self> {
1448        let broadcast = match direction {
1449            FftDirection::Forward => Complex::new(-0.0, 0.0),
1450            FftDirection::Inverse => Complex::new(0.0, -0.0),
1451        };
1452        Rotation90(Self::broadcast_complex_elements(broadcast))
1453    }
1454    #[inline(always)]
1455    unsafe fn make_mixedradix_twiddle_chunk(
1456        x: usize,
1457        y: usize,
1458        len: usize,
1459        direction: FftDirection,
1460    ) -> Self {
1461        let mut twiddle_chunk = [Complex::<f64>::zero(); Self::COMPLEX_PER_VECTOR];
1462        for i in 0..Self::COMPLEX_PER_VECTOR {
1463            twiddle_chunk[i] = twiddles::compute_twiddle(y * (x + i), len, direction);
1464        }
1465
1466        twiddle_chunk.as_slice().load_complex(0)
1467    }
1468    #[inline(always)]
1469    unsafe fn broadcast_twiddle(index: usize, len: usize, direction: FftDirection) -> Self {
1470        Self::broadcast_complex_elements(twiddles::compute_twiddle(index, len, direction))
1471    }
1472
1473    #[inline(always)]
1474    unsafe fn transpose2_packed(rows: [Self; 2]) -> [Self; 2] {
1475        Self::unpack_complex(rows)
1476    }
1477    #[inline(always)]
1478    unsafe fn transpose3_packed(rows: [Self; 3]) -> [Self; 3] {
1479        let unpacked0 = Self::unpacklo_complex([rows[0], rows[1]]);
1480        let blended = _mm256_blend_pd(rows[0], rows[2], 0x03);
1481        let unpacked2 = Self::unpackhi_complex([rows[1], rows[2]]);
1482
1483        [unpacked0, blended, unpacked2]
1484    }
1485    #[inline(always)]
1486    unsafe fn transpose4_packed(rows: [Self; 4]) -> [Self; 4] {
1487        let [unpacked0, unpacked1] = Self::unpack_complex([rows[0], rows[1]]);
1488        let [unpacked2, unpacked3] = Self::unpack_complex([rows[2], rows[3]]);
1489
1490        [unpacked0, unpacked2, unpacked1, unpacked3]
1491    }
1492    #[inline(always)]
1493    unsafe fn transpose5_packed(rows: [Self; 5]) -> [Self; 5] {
1494        [
1495            Self::unpacklo_complex([rows[0], rows[1]]),
1496            Self::unpacklo_complex([rows[2], rows[3]]),
1497            _mm256_blend_pd(rows[0], rows[4], 0x03),
1498            Self::unpackhi_complex([rows[1], rows[2]]),
1499            Self::unpackhi_complex([rows[3], rows[4]]),
1500        ]
1501    }
1502    #[inline(always)]
1503    unsafe fn transpose6_packed(rows: [Self; 6]) -> [Self; 6] {
1504        let [unpacked0, unpacked1] = Self::unpack_complex([rows[0], rows[1]]);
1505        let [unpacked2, unpacked3] = Self::unpack_complex([rows[2], rows[3]]);
1506        let [unpacked4, unpacked5] = Self::unpack_complex([rows[4], rows[5]]);
1507
1508        [
1509            unpacked0, unpacked2, unpacked4, unpacked1, unpacked3, unpacked5,
1510        ]
1511    }
1512    #[inline(always)]
1513    unsafe fn transpose7_packed(rows: [Self; 7]) -> [Self; 7] {
1514        [
1515            Self::unpacklo_complex([rows[0], rows[1]]),
1516            Self::unpacklo_complex([rows[2], rows[3]]),
1517            Self::unpacklo_complex([rows[4], rows[5]]),
1518            _mm256_blend_pd(rows[0], rows[6], 0x03),
1519            Self::unpackhi_complex([rows[1], rows[2]]),
1520            Self::unpackhi_complex([rows[3], rows[4]]),
1521            Self::unpackhi_complex([rows[5], rows[6]]),
1522        ]
1523    }
1524    #[inline(always)]
1525    unsafe fn transpose8_packed(rows: [Self; 8]) -> [Self; 8] {
1526        let chunk0 = [rows[0], rows[1], rows[2], rows[3]];
1527        let chunk1 = [rows[4], rows[5], rows[6], rows[7]];
1528
1529        let output0 = Self::transpose4_packed(chunk0);
1530        let output1 = Self::transpose4_packed(chunk1);
1531
1532        [
1533            output0[0], output0[1], output1[0], output1[1], output0[2], output0[3], output1[2],
1534            output1[3],
1535        ]
1536    }
1537    #[inline(always)]
1538    unsafe fn transpose9_packed(rows: [Self; 9]) -> [Self; 9] {
1539        [
1540            _mm256_permute2f128_pd(rows[0], rows[1], 0x20),
1541            _mm256_permute2f128_pd(rows[2], rows[3], 0x20),
1542            _mm256_permute2f128_pd(rows[4], rows[5], 0x20),
1543            _mm256_permute2f128_pd(rows[6], rows[7], 0x20),
1544            _mm256_permute2f128_pd(rows[8], rows[0], 0x30),
1545            _mm256_permute2f128_pd(rows[1], rows[2], 0x31),
1546            _mm256_permute2f128_pd(rows[3], rows[4], 0x31),
1547            _mm256_permute2f128_pd(rows[5], rows[6], 0x31),
1548            _mm256_permute2f128_pd(rows[7], rows[8], 0x31),
1549        ]
1550    }
1551    #[inline(always)]
1552    unsafe fn transpose11_packed(rows: [Self; 11]) -> [Self; 11] {
1553        [
1554            _mm256_permute2f128_pd(rows[0], rows[1], 0x20),
1555            _mm256_permute2f128_pd(rows[2], rows[3], 0x20),
1556            _mm256_permute2f128_pd(rows[4], rows[5], 0x20),
1557            _mm256_permute2f128_pd(rows[6], rows[7], 0x20),
1558            _mm256_permute2f128_pd(rows[8], rows[9], 0x20),
1559            _mm256_permute2f128_pd(rows[10], rows[0], 0x30),
1560            _mm256_permute2f128_pd(rows[1], rows[2], 0x31),
1561            _mm256_permute2f128_pd(rows[3], rows[4], 0x31),
1562            _mm256_permute2f128_pd(rows[5], rows[6], 0x31),
1563            _mm256_permute2f128_pd(rows[7], rows[8], 0x31),
1564            _mm256_permute2f128_pd(rows[9], rows[10], 0x31),
1565        ]
1566    }
1567    #[inline(always)]
1568    unsafe fn transpose12_packed(rows: [Self; 12]) -> [Self; 12] {
1569        let chunk0 = [rows[0], rows[1], rows[2], rows[3]];
1570        let chunk1 = [rows[4], rows[5], rows[6], rows[7]];
1571        let chunk2 = [rows[8], rows[9], rows[10], rows[11]];
1572
1573        let output0 = Self::transpose4_packed(chunk0);
1574        let output1 = Self::transpose4_packed(chunk1);
1575        let output2 = Self::transpose4_packed(chunk2);
1576
1577        [
1578            output0[0], output0[1], output1[0], output1[1], output2[0], output2[1], output0[2],
1579            output0[3], output1[2], output1[3], output2[2], output2[3],
1580        ]
1581    }
1582    #[inline(always)]
1583    unsafe fn transpose16_packed(rows: [Self; 16]) -> [Self; 16] {
1584        let chunk0 = [
1585            rows[0], rows[1], rows[2], rows[3], rows[4], rows[5], rows[6], rows[7],
1586        ];
1587        let chunk1 = [
1588            rows[8], rows[9], rows[10], rows[11], rows[12], rows[13], rows[14], rows[15],
1589        ];
1590
1591        let output0 = Self::transpose8_packed(chunk0);
1592        let output1 = Self::transpose8_packed(chunk1);
1593
1594        [
1595            output0[0], output0[1], output0[2], output0[3], output1[0], output1[1], output1[2],
1596            output1[3], output0[4], output0[5], output0[6], output0[7], output1[4], output1[5],
1597            output1[6], output1[7],
1598        ]
1599    }
1600}
1601impl AvxVector256 for __m256d {
1602    type ScalarType = f64;
1603    type HalfVector = __m128d;
1604
1605    #[inline(always)]
1606    unsafe fn lo(self) -> Self::HalfVector {
1607        _mm256_castpd256_pd128(self)
1608    }
1609    #[inline(always)]
1610    unsafe fn hi(self) -> Self::HalfVector {
1611        _mm256_extractf128_pd(self, 1)
1612    }
1613    #[inline(always)]
1614    unsafe fn merge(lo: Self::HalfVector, hi: Self::HalfVector) -> Self {
1615        _mm256_insertf128_pd(_mm256_castpd128_pd256(lo), hi, 1)
1616    }
1617
1618    #[inline(always)]
1619    unsafe fn broadcast_complex_elements(value: Complex<Self::ScalarType>) -> Self {
1620        _mm256_set_pd(value.im, value.re, value.im, value.re)
1621    }
1622
1623    #[inline(always)]
1624    unsafe fn load_complex(ptr: *const Complex<Self::ScalarType>) -> Self {
1625        _mm256_loadu_pd(ptr as *const Self::ScalarType)
1626    }
1627    #[inline(always)]
1628    unsafe fn store_complex(ptr: *mut Complex<Self::ScalarType>, data: Self) {
1629        _mm256_storeu_pd(ptr as *mut Self::ScalarType, data)
1630    }
1631    #[inline(always)]
1632    unsafe fn gather_complex_avx2_index32(
1633        ptr: *const Complex<Self::ScalarType>,
1634        indexes: __m128i,
1635    ) -> Self {
1636        let offsets = _mm_set_epi32(1, 0, 1, 0);
1637        let shifted = _mm_slli_epi32(indexes, 1);
1638        let modified_indexes = _mm_add_epi32(offsets, shifted);
1639
1640        _mm256_i32gather_pd(ptr as *const f64, modified_indexes, 8)
1641    }
1642    #[inline(always)]
1643    unsafe fn gather_complex_avx2_index64(
1644        ptr: *const Complex<Self::ScalarType>,
1645        indexes: __m256i,
1646    ) -> Self {
1647        let offsets = _mm256_set_epi64x(1, 0, 1, 0);
1648        let shifted = _mm256_slli_epi64(indexes, 1);
1649        let modified_indexes = _mm256_add_epi64(offsets, shifted);
1650
1651        _mm256_i64gather_pd(ptr as *const f64, modified_indexes, 8)
1652    }
1653    #[inline(always)]
1654    unsafe fn load_partial1_complex(ptr: *const Complex<Self::ScalarType>) -> Self::HalfVector {
1655        _mm_loadu_pd(ptr as *const f64)
1656    }
1657    #[inline(always)]
1658    unsafe fn load_partial2_complex(_ptr: *const Complex<Self::ScalarType>) -> Self::HalfVector {
1659        unimplemented!("Impossible to do a partial load of 2 complex f64's")
1660    }
1661    #[inline(always)]
1662    unsafe fn load_partial3_complex(_ptr: *const Complex<Self::ScalarType>) -> Self {
1663        unimplemented!("Impossible to do a partial load of 3 complex f64's")
1664    }
1665    #[inline(always)]
1666    unsafe fn store_partial1_complex(ptr: *mut Complex<Self::ScalarType>, data: Self::HalfVector) {
1667        _mm_storeu_pd(ptr as *mut f64, data);
1668    }
1669    #[inline(always)]
1670    unsafe fn store_partial2_complex(
1671        _ptr: *mut Complex<Self::ScalarType>,
1672        _data: Self::HalfVector,
1673    ) {
1674        unimplemented!("Impossible to do a partial store of 2 complex f64's")
1675    }
1676    #[inline(always)]
1677    unsafe fn store_partial3_complex(_ptr: *mut Complex<Self::ScalarType>, _data: Self) {
1678        unimplemented!("Impossible to do a partial store of 3 complex f64's")
1679    }
1680}
1681
1682impl AvxVector for __m128d {
1683    const COMPLEX_PER_VECTOR: usize = 1;
1684
1685    #[inline(always)]
1686    unsafe fn zero() -> Self {
1687        _mm_setzero_pd()
1688    }
1689    #[inline(always)]
1690    unsafe fn half_root2() -> Self {
1691        // note: we're computing a square root here, but checking the assembly says the compiler is smart enough to turn this into a constant
1692        _mm_load1_pd(&0.5f64.sqrt())
1693    }
1694
1695    #[inline(always)]
1696    unsafe fn xor(left: Self, right: Self) -> Self {
1697        _mm_xor_pd(left, right)
1698    }
1699    #[inline(always)]
1700    unsafe fn neg(self) -> Self {
1701        _mm_xor_pd(self, _mm_load1_pd(&-0.0))
1702    }
1703    #[inline(always)]
1704    unsafe fn add(left: Self, right: Self) -> Self {
1705        _mm_add_pd(left, right)
1706    }
1707    #[inline(always)]
1708    unsafe fn sub(left: Self, right: Self) -> Self {
1709        _mm_sub_pd(left, right)
1710    }
1711    #[inline(always)]
1712    unsafe fn mul(left: Self, right: Self) -> Self {
1713        _mm_mul_pd(left, right)
1714    }
1715    #[inline(always)]
1716    unsafe fn fmadd(left: Self, right: Self, add: Self) -> Self {
1717        _mm_fmadd_pd(left, right, add)
1718    }
1719    #[inline(always)]
1720    unsafe fn fnmadd(left: Self, right: Self, add: Self) -> Self {
1721        _mm_fnmadd_pd(left, right, add)
1722    }
1723    #[inline(always)]
1724    unsafe fn fmaddsub(left: Self, right: Self, add: Self) -> Self {
1725        _mm_fmaddsub_pd(left, right, add)
1726    }
1727    #[inline(always)]
1728    unsafe fn fmsubadd(left: Self, right: Self, add: Self) -> Self {
1729        _mm_fmsubadd_pd(left, right, add)
1730    }
1731
1732    #[inline(always)]
1733    unsafe fn reverse_complex_elements(self) -> Self {
1734        // nothing to reverse
1735        self
1736    }
1737    #[inline(always)]
1738    unsafe fn unpacklo_complex(_rows: [Self; 2]) -> Self {
1739        unimplemented!(); // this operation doesn't make sense with one element. TODO: I don't know if it would be more useful to error here or to just return the inputs unchanged. If returning the inputs is useful, do that.
1740    }
1741    #[inline(always)]
1742    unsafe fn unpackhi_complex(_rows: [Self; 2]) -> Self {
1743        unimplemented!(); // this operation doesn't make sense with one element. TODO: I don't know if it would be more useful to error here or to just return the inputs unchanged. If returning the inputs is useful, do that.
1744    }
1745
1746    #[inline(always)]
1747    unsafe fn swap_complex_components(self) -> Self {
1748        _mm_permute_pd(self, 0x01)
1749    }
1750    #[inline(always)]
1751    unsafe fn duplicate_complex_components(self) -> (Self, Self) {
1752        (_mm_movedup_pd(self), _mm_permute_pd(self, 0x03))
1753    }
1754
1755    #[inline(always)]
1756    unsafe fn make_rotation90(direction: FftDirection) -> Rotation90<Self> {
1757        let broadcast = match direction {
1758            FftDirection::Forward => Complex::new(-0.0, 0.0),
1759            FftDirection::Inverse => Complex::new(0.0, -0.0),
1760        };
1761        Rotation90(Self::broadcast_complex_elements(broadcast))
1762    }
1763    #[inline(always)]
1764    unsafe fn make_mixedradix_twiddle_chunk(
1765        x: usize,
1766        y: usize,
1767        len: usize,
1768        direction: FftDirection,
1769    ) -> Self {
1770        let mut twiddle_chunk = [Complex::<f64>::zero(); Self::COMPLEX_PER_VECTOR];
1771        for i in 0..Self::COMPLEX_PER_VECTOR {
1772            twiddle_chunk[i] = twiddles::compute_twiddle(y * (x + i), len, direction);
1773        }
1774
1775        _mm_loadu_pd(twiddle_chunk.as_ptr() as *const f64)
1776    }
1777    #[inline(always)]
1778    unsafe fn broadcast_twiddle(index: usize, len: usize, direction: FftDirection) -> Self {
1779        Self::broadcast_complex_elements(twiddles::compute_twiddle(index, len, direction))
1780    }
1781
1782    #[inline(always)]
1783    unsafe fn transpose2_packed(rows: [Self; 2]) -> [Self; 2] {
1784        rows
1785    }
1786    #[inline(always)]
1787    unsafe fn transpose3_packed(rows: [Self; 3]) -> [Self; 3] {
1788        rows
1789    }
1790    #[inline(always)]
1791    unsafe fn transpose4_packed(rows: [Self; 4]) -> [Self; 4] {
1792        rows
1793    }
1794    #[inline(always)]
1795    unsafe fn transpose5_packed(rows: [Self; 5]) -> [Self; 5] {
1796        rows
1797    }
1798    #[inline(always)]
1799    unsafe fn transpose6_packed(rows: [Self; 6]) -> [Self; 6] {
1800        rows
1801    }
1802    #[inline(always)]
1803    unsafe fn transpose7_packed(rows: [Self; 7]) -> [Self; 7] {
1804        rows
1805    }
1806    #[inline(always)]
1807    unsafe fn transpose8_packed(rows: [Self; 8]) -> [Self; 8] {
1808        rows
1809    }
1810    #[inline(always)]
1811    unsafe fn transpose9_packed(rows: [Self; 9]) -> [Self; 9] {
1812        rows
1813    }
1814    #[inline(always)]
1815    unsafe fn transpose11_packed(rows: [Self; 11]) -> [Self; 11] {
1816        rows
1817    }
1818    #[inline(always)]
1819    unsafe fn transpose12_packed(rows: [Self; 12]) -> [Self; 12] {
1820        rows
1821    }
1822    #[inline(always)]
1823    unsafe fn transpose16_packed(rows: [Self; 16]) -> [Self; 16] {
1824        rows
1825    }
1826}
1827impl AvxVector128 for __m128d {
1828    type FullVector = __m256d;
1829
1830    #[inline(always)]
1831    unsafe fn lo(input: Self::FullVector) -> Self {
1832        _mm256_castpd256_pd128(input)
1833    }
1834    #[inline(always)]
1835    unsafe fn hi(input: Self::FullVector) -> Self {
1836        _mm256_extractf128_pd(input, 1)
1837    }
1838    #[inline(always)]
1839    unsafe fn merge(lo: Self, hi: Self) -> Self::FullVector {
1840        _mm256_insertf128_pd(_mm256_castpd128_pd256(lo), hi, 1)
1841    }
1842    #[inline(always)]
1843    unsafe fn zero_extend(self) -> Self::FullVector {
1844        _mm256_zextpd128_pd256(self)
1845    }
1846    #[inline(always)]
1847    unsafe fn lo_rotation(input: Rotation90<Self::FullVector>) -> Rotation90<Self> {
1848        input.lo()
1849    }
1850    #[inline(always)]
1851    unsafe fn broadcast_complex_elements(value: Complex<f64>) -> Self {
1852        _mm_set_pd(value.im, value.re)
1853    }
1854    #[inline(always)]
1855    unsafe fn gather32_complex_avx2(ptr: *const Complex<f64>, indexes: __m128i) -> Self {
1856        let mut index_storage: [i32; 4] = [0; 4];
1857        _mm_storeu_si128(index_storage.as_mut_ptr() as *mut __m128i, indexes);
1858
1859        _mm_loadu_pd(ptr.offset(index_storage[0] as isize) as *const f64)
1860    }
1861    #[inline(always)]
1862    unsafe fn gather64_complex_avx2(ptr: *const Complex<f64>, indexes: __m128i) -> Self {
1863        let mut index_storage: [i64; 4] = [0; 4];
1864        _mm_storeu_si128(index_storage.as_mut_ptr() as *mut __m128i, indexes);
1865
1866        _mm_loadu_pd(ptr.offset(index_storage[0] as isize) as *const f64)
1867    }
1868}
1869
1870pub trait AvxArray<T: AvxNum>: Deref {
1871    unsafe fn load_complex(&self, index: usize) -> T::VectorType;
1872    unsafe fn load_partial1_complex(
1873        &self,
1874        index: usize,
1875    ) -> <T::VectorType as AvxVector256>::HalfVector;
1876    unsafe fn load_partial2_complex(
1877        &self,
1878        index: usize,
1879    ) -> <T::VectorType as AvxVector256>::HalfVector;
1880    unsafe fn load_partial3_complex(&self, index: usize) -> T::VectorType;
1881
1882    // some avx operations need bespoke one-off things that don't fit into the methods above, so we should provide an escape hatch for them
1883    fn input_ptr(&self) -> *const Complex<T>;
1884}
1885pub trait AvxArrayMut<T: AvxNum>: AvxArray<T> + DerefMut {
1886    unsafe fn store_complex(&mut self, data: T::VectorType, index: usize);
1887    unsafe fn store_partial1_complex(
1888        &mut self,
1889        data: <T::VectorType as AvxVector256>::HalfVector,
1890        index: usize,
1891    );
1892    unsafe fn store_partial2_complex(
1893        &mut self,
1894        data: <T::VectorType as AvxVector256>::HalfVector,
1895        index: usize,
1896    );
1897    unsafe fn store_partial3_complex(&mut self, data: T::VectorType, index: usize);
1898}
1899
1900impl<T: AvxNum> AvxArray<T> for &[Complex<T>] {
1901    #[inline(always)]
1902    unsafe fn load_complex(&self, index: usize) -> T::VectorType {
1903        debug_assert!(self.len() >= index + T::VectorType::COMPLEX_PER_VECTOR);
1904        T::VectorType::load_complex(self.as_ptr().add(index))
1905    }
1906    #[inline(always)]
1907    unsafe fn load_partial1_complex(
1908        &self,
1909        index: usize,
1910    ) -> <T::VectorType as AvxVector256>::HalfVector {
1911        debug_assert!(self.len() >= index + 1);
1912        T::VectorType::load_partial1_complex(self.as_ptr().add(index))
1913    }
1914    #[inline(always)]
1915    unsafe fn load_partial2_complex(
1916        &self,
1917        index: usize,
1918    ) -> <T::VectorType as AvxVector256>::HalfVector {
1919        debug_assert!(self.len() >= index + 2);
1920        T::VectorType::load_partial2_complex(self.as_ptr().add(index))
1921    }
1922    #[inline(always)]
1923    unsafe fn load_partial3_complex(&self, index: usize) -> T::VectorType {
1924        debug_assert!(self.len() >= index + 3);
1925        T::VectorType::load_partial3_complex(self.as_ptr().add(index))
1926    }
1927    #[inline(always)]
1928    fn input_ptr(&self) -> *const Complex<T> {
1929        self.as_ptr()
1930    }
1931}
1932impl<T: AvxNum> AvxArray<T> for &mut [Complex<T>] {
1933    #[inline(always)]
1934    unsafe fn load_complex(&self, index: usize) -> T::VectorType {
1935        debug_assert!(self.len() >= index + T::VectorType::COMPLEX_PER_VECTOR);
1936        T::VectorType::load_complex(self.as_ptr().add(index))
1937    }
1938    #[inline(always)]
1939    unsafe fn load_partial1_complex(
1940        &self,
1941        index: usize,
1942    ) -> <T::VectorType as AvxVector256>::HalfVector {
1943        debug_assert!(self.len() >= index + 1);
1944        T::VectorType::load_partial1_complex(self.as_ptr().add(index))
1945    }
1946    #[inline(always)]
1947    unsafe fn load_partial2_complex(
1948        &self,
1949        index: usize,
1950    ) -> <T::VectorType as AvxVector256>::HalfVector {
1951        debug_assert!(self.len() >= index + 2);
1952        T::VectorType::load_partial2_complex(self.as_ptr().add(index))
1953    }
1954    #[inline(always)]
1955    unsafe fn load_partial3_complex(&self, index: usize) -> T::VectorType {
1956        debug_assert!(self.len() >= index + 3);
1957        T::VectorType::load_partial3_complex(self.as_ptr().add(index))
1958    }
1959    #[inline(always)]
1960    fn input_ptr(&self) -> *const Complex<T> {
1961        self.as_ptr()
1962    }
1963}
1964impl<'a, T: AvxNum> AvxArray<T> for DoubleBuf<'a, T>
1965where
1966    &'a [Complex<T>]: AvxArray<T>,
1967{
1968    #[inline(always)]
1969    unsafe fn load_complex(&self, index: usize) -> T::VectorType {
1970        self.input.load_complex(index)
1971    }
1972    #[inline(always)]
1973    unsafe fn load_partial1_complex(
1974        &self,
1975        index: usize,
1976    ) -> <T::VectorType as AvxVector256>::HalfVector {
1977        self.input.load_partial1_complex(index)
1978    }
1979    #[inline(always)]
1980    unsafe fn load_partial2_complex(
1981        &self,
1982        index: usize,
1983    ) -> <T::VectorType as AvxVector256>::HalfVector {
1984        self.input.load_partial2_complex(index)
1985    }
1986    #[inline(always)]
1987    unsafe fn load_partial3_complex(&self, index: usize) -> T::VectorType {
1988        self.input.load_partial3_complex(index)
1989    }
1990    #[inline(always)]
1991    fn input_ptr(&self) -> *const Complex<T> {
1992        self.input.input_ptr()
1993    }
1994}
1995
1996impl<T: AvxNum> AvxArrayMut<T> for &mut [Complex<T>] {
1997    #[inline(always)]
1998    unsafe fn store_complex(&mut self, data: T::VectorType, index: usize) {
1999        debug_assert!(self.len() >= index + T::VectorType::COMPLEX_PER_VECTOR);
2000        T::VectorType::store_complex(self.as_mut_ptr().add(index), data);
2001    }
2002    #[inline(always)]
2003    unsafe fn store_partial1_complex(
2004        &mut self,
2005        data: <T::VectorType as AvxVector256>::HalfVector,
2006        index: usize,
2007    ) {
2008        debug_assert!(self.len() >= index + 1);
2009        T::VectorType::store_partial1_complex(self.as_mut_ptr().add(index), data);
2010    }
2011    #[inline(always)]
2012    unsafe fn store_partial2_complex(
2013        &mut self,
2014        data: <T::VectorType as AvxVector256>::HalfVector,
2015        index: usize,
2016    ) {
2017        debug_assert!(self.len() >= index + 2);
2018        T::VectorType::store_partial2_complex(self.as_mut_ptr().add(index), data);
2019    }
2020    #[inline(always)]
2021    unsafe fn store_partial3_complex(&mut self, data: T::VectorType, index: usize) {
2022        debug_assert!(self.len() >= index + 3);
2023        T::VectorType::store_partial3_complex(self.as_mut_ptr().add(index), data);
2024    }
2025}
2026impl<'a, T: AvxNum> AvxArrayMut<T> for DoubleBuf<'a, T>
2027where
2028    Self: AvxArray<T>,
2029    &'a mut [Complex<T>]: AvxArrayMut<T>,
2030{
2031    #[inline(always)]
2032    unsafe fn store_complex(&mut self, data: T::VectorType, index: usize) {
2033        self.output.store_complex(data, index);
2034    }
2035    #[inline(always)]
2036    unsafe fn store_partial1_complex(
2037        &mut self,
2038        data: <T::VectorType as AvxVector256>::HalfVector,
2039        index: usize,
2040    ) {
2041        self.output.store_partial1_complex(data, index);
2042    }
2043    #[inline(always)]
2044    unsafe fn store_partial2_complex(
2045        &mut self,
2046        data: <T::VectorType as AvxVector256>::HalfVector,
2047        index: usize,
2048    ) {
2049        self.output.store_partial2_complex(data, index);
2050    }
2051    #[inline(always)]
2052    unsafe fn store_partial3_complex(&mut self, data: T::VectorType, index: usize) {
2053        self.output.store_partial3_complex(data, index);
2054    }
2055}
2056
2057// A custom butterfly-16 function that calls a lambda to load/store data instead of taking an array
2058// This is particularly useful for butterfly 16, because the whole problem doesn't fit into registers, and the compiler isn't smart enough to only load data when it's needed
2059// So the version that takes an array ends up loading data and immediately re-storing it on the stack. By lazily loading and storing exactly when we need to, we can avoid some data reshuffling
2060macro_rules! column_butterfly16_loadfn{
2061    ($load_expr: expr, $store_expr: expr, $twiddles: expr, $rotation: expr) => (
2062        // Size-4 FFTs down the columns
2063        let input1 = [$load_expr(1), $load_expr(5), $load_expr(9), $load_expr(13)];
2064        let mut mid1 = AvxVector::column_butterfly4(input1, $rotation);
2065
2066        mid1[1] = AvxVector::mul_complex(mid1[1], $twiddles[0]);
2067        mid1[2] = avx_vector::apply_butterfly8_twiddle1(mid1[2], $rotation);
2068        mid1[3] = AvxVector::mul_complex(mid1[3], $twiddles[1]);
2069
2070        let input2 = [$load_expr(2), $load_expr(6), $load_expr(10), $load_expr(14)];
2071        let mut mid2 = AvxVector::column_butterfly4(input2, $rotation);
2072
2073        mid2[1] = avx_vector::apply_butterfly8_twiddle1(mid2[1], $rotation);
2074        mid2[2] = mid2[2].rotate90($rotation);
2075        mid2[3] = avx_vector::apply_butterfly8_twiddle3(mid2[3], $rotation);
2076
2077        let input3 = [$load_expr(3), $load_expr(7), $load_expr(11), $load_expr(15)];
2078        let mut mid3 = AvxVector::column_butterfly4(input3, $rotation);
2079
2080        mid3[1] = AvxVector::mul_complex(mid3[1], $twiddles[1]);
2081        mid3[2] = avx_vector::apply_butterfly8_twiddle3(mid3[2], $rotation);
2082        mid3[3] = AvxVector::mul_complex(mid3[3], $twiddles[0].neg());
2083
2084        // do the first row last, because it doesn't need twiddles and therefore requires fewer intermediates
2085        let input0 = [$load_expr(0), $load_expr(4), $load_expr(8), $load_expr(12)];
2086        let mid0     = AvxVector::column_butterfly4(input0, $rotation);
2087
2088        // All of the data is now in the right format to just do a bunch of butterfly 8's.
2089        // Write the data out to the final output as we go so that the compiler can stop worrying about finding stack space for it
2090        for i in 0..4 {
2091            let output = AvxVector::column_butterfly4([mid0[i], mid1[i], mid2[i], mid3[i]], $rotation);
2092            $store_expr(output[0], i);
2093            $store_expr(output[1], i + 4);
2094            $store_expr(output[2], i + 8);
2095            $store_expr(output[3], i + 12);
2096        }
2097    )
2098}
2099
2100// A custom butterfly-32 function that calls a lambda to load/store data instead of taking an array
2101// This is particularly useful for butterfly 32, because the whole problem doesn't fit into registers, and the compiler isn't smart enough to only load data when it's needed
2102// So the version that takes an array ends up loading data and immediately re-storing it on the stack. By lazily loading and storing exactly when we need to, we can avoid some data reshuffling
2103macro_rules! column_butterfly32_loadfn{
2104    ($load_expr: expr, $store_expr: expr, $twiddles: expr, $rotation: expr) => (
2105        // Size-4 FFTs down the columns
2106        let input1 = [$load_expr(1), $load_expr(9), $load_expr(17), $load_expr(25)];
2107        let mut mid1     = AvxVector::column_butterfly4(input1, $rotation);
2108
2109        mid1[1] = AvxVector::mul_complex(mid1[1], $twiddles[0]);
2110        mid1[2] = AvxVector::mul_complex(mid1[2], $twiddles[1]);
2111        mid1[3] = AvxVector::mul_complex(mid1[3], $twiddles[2]);
2112
2113        let input2 = [$load_expr(2), $load_expr(10), $load_expr(18), $load_expr(26)];
2114        let mut mid2     = AvxVector::column_butterfly4(input2, $rotation);
2115
2116        mid2[1] = AvxVector::mul_complex(mid2[1], $twiddles[1]);
2117        mid2[2] = avx_vector::apply_butterfly8_twiddle1(mid2[2], $rotation);
2118        mid2[3] = AvxVector::mul_complex(mid2[3], $twiddles[4]);
2119
2120        let input3 = [$load_expr(3), $load_expr(11), $load_expr(19), $load_expr(27)];
2121        let mut mid3     = AvxVector::column_butterfly4(input3, $rotation);
2122
2123        mid3[1] = AvxVector::mul_complex(mid3[1], $twiddles[2]);
2124        mid3[2] = AvxVector::mul_complex(mid3[2], $twiddles[4]);
2125        mid3[3] = AvxVector::mul_complex(mid3[3], $twiddles[0].rotate90($rotation));
2126
2127        let input4 = [$load_expr(4), $load_expr(12), $load_expr(20), $load_expr(28)];
2128        let mut mid4     = AvxVector::column_butterfly4(input4, $rotation);
2129
2130        mid4[1] = avx_vector::apply_butterfly8_twiddle1(mid4[1], $rotation);
2131        mid4[2] = mid4[2].rotate90($rotation);
2132        mid4[3] = avx_vector::apply_butterfly8_twiddle3(mid4[3], $rotation);
2133
2134        let input5 = [$load_expr(5), $load_expr(13), $load_expr(21), $load_expr(29)];
2135        let mut mid5     = AvxVector::column_butterfly4(input5, $rotation);
2136
2137        mid5[1] = AvxVector::mul_complex(mid5[1], $twiddles[3]);
2138        mid5[2] = AvxVector::mul_complex(mid5[2], $twiddles[1].rotate90($rotation));
2139        mid5[3] = AvxVector::mul_complex(mid5[3], $twiddles[5].rotate90($rotation));
2140
2141        let input6 = [$load_expr(6), $load_expr(14), $load_expr(22), $load_expr(30)];
2142        let mut mid6     = AvxVector::column_butterfly4(input6, $rotation);
2143
2144        mid6[1] = AvxVector::mul_complex(mid6[1], $twiddles[4]);
2145        mid6[2] = avx_vector::apply_butterfly8_twiddle3(mid6[2], $rotation);
2146        mid6[3] = AvxVector::mul_complex(mid6[3], $twiddles[1].neg());
2147
2148        let input7 = [$load_expr(7), $load_expr(15), $load_expr(23), $load_expr(31)];
2149        let mut mid7     = AvxVector::column_butterfly4(input7, $rotation);
2150
2151        mid7[1] = AvxVector::mul_complex(mid7[1], $twiddles[5]);
2152        mid7[2] = AvxVector::mul_complex(mid7[2], $twiddles[4].rotate90($rotation));
2153        mid7[3] = AvxVector::mul_complex(mid7[3], $twiddles[3].neg());
2154
2155        let input0 = [$load_expr(0), $load_expr(8), $load_expr(16), $load_expr(24)];
2156        let mid0     = AvxVector::column_butterfly4(input0, $rotation);
2157
2158        // All of the data is now in the right format to just do a bunch of butterfly 8's in a loop.
2159        // Write the data out to the final output as we go so that the compiler can stop worrying about finding stack space for it
2160        for i in 0..4 {
2161            let output = AvxVector::column_butterfly8([mid0[i], mid1[i], mid2[i], mid3[i], mid4[i], mid5[i], mid6[i], mid7[i]], $rotation);
2162            $store_expr(output[0], i);
2163            $store_expr(output[1], i + 4);
2164            $store_expr(output[2], i + 8);
2165            $store_expr(output[3], i + 12);
2166            $store_expr(output[4], i + 16);
2167            $store_expr(output[5], i + 20);
2168            $store_expr(output[6], i + 24);
2169            $store_expr(output[7], i + 28);
2170        }
2171    )
2172}
2173
2174/// Multiply the complex numbers in `left` by the complex numbers in `right`.
2175/// This is exactly the same as `mul_complex` in `AvxVector`, but this implementation also conjugates the `left` input before multiplying
2176#[inline(always)]
2177unsafe fn mul_complex_conjugated<V: AvxVector>(left: V, right: V) -> V {
2178    // Extract the real and imaginary components from left into 2 separate registers
2179    let (left_real, left_imag) = V::duplicate_complex_components(left);
2180
2181    // create a shuffled version of right where the imaginary values are swapped with the reals
2182    let right_shuffled = V::swap_complex_components(right);
2183
2184    // multiply our duplicated imaginary left vector by our shuffled right vector. that will give us the right side of the traditional complex multiplication formula
2185    let output_right = V::mul(left_imag, right_shuffled);
2186
2187    // use a FMA instruction to multiply together left side of the complex multiplication formula, then alternatingly add and subtract the left side from the right
2188    // By using subadd instead of addsub, we can conjugate the left side for free.
2189    V::fmsubadd(left_real, right, output_right)
2190}
2191
2192// compute buffer[i] = buffer[i].conj() * multiplier[i] pairwise complex multiplication for each element.
2193#[target_feature(enable = "avx", enable = "fma")]
2194pub unsafe fn pairwise_complex_mul_assign_conjugated<T: AvxNum>(
2195    mut buffer: &mut [Complex<T>],
2196    multiplier: &[T::VectorType],
2197) {
2198    assert!(multiplier.len() * T::VectorType::COMPLEX_PER_VECTOR >= buffer.len()); // Assert to convince the compiler to omit bounds checks inside the loop
2199
2200    for (i, mut buffer_chunk) in buffer
2201        .chunks_exact_mut(T::VectorType::COMPLEX_PER_VECTOR)
2202        .enumerate()
2203    {
2204        let left = buffer_chunk.load_complex(0);
2205
2206        // Do a complex multiplication between `left` and `right`
2207        let product = mul_complex_conjugated(left, multiplier[i]);
2208
2209        // Store the result
2210        buffer_chunk.store_complex(product, 0);
2211    }
2212
2213    // Process the remainder, if there is one
2214    let remainder_count = buffer.len() % T::VectorType::COMPLEX_PER_VECTOR;
2215    if remainder_count > 0 {
2216        let remainder_index = buffer.len() - remainder_count;
2217        let remainder_multiplier = multiplier.last().unwrap();
2218        match remainder_count {
2219            1 => {
2220                let left = buffer.load_partial1_complex(remainder_index);
2221                let product = mul_complex_conjugated(left, remainder_multiplier.lo());
2222                buffer.store_partial1_complex(product, remainder_index);
2223            }
2224            2 => {
2225                let left = buffer.load_partial2_complex(remainder_index);
2226                let product = mul_complex_conjugated(left, remainder_multiplier.lo());
2227                buffer.store_partial2_complex(product, remainder_index);
2228            }
2229            3 => {
2230                let left = buffer.load_partial3_complex(remainder_index);
2231                let product = mul_complex_conjugated(left, *remainder_multiplier);
2232                buffer.store_partial3_complex(product, remainder_index);
2233            }
2234            _ => unreachable!(),
2235        }
2236    }
2237}
2238
2239// compute output[i] = input[i].conj() * multiplier[i] pairwise complex multiplication for each element.
2240#[target_feature(enable = "avx", enable = "fma")]
2241pub unsafe fn pairwise_complex_mul_conjugated<T: AvxNum>(
2242    input: &[Complex<T>],
2243    mut output: &mut [Complex<T>],
2244    multiplier: &[T::VectorType],
2245) {
2246    assert!(
2247        multiplier.len() * T::VectorType::COMPLEX_PER_VECTOR >= input.len(),
2248        "multiplier len = {}, input len = {}",
2249        multiplier.len(),
2250        input.len()
2251    ); // Assert to convince the compiler to omit bounds checks inside the loop
2252    assert_eq!(input.len(), output.len()); // Assert to convince the compiler to omit bounds checks inside the loop
2253    let main_loop_count = input.len() / T::VectorType::COMPLEX_PER_VECTOR;
2254    let remainder_count = input.len() % T::VectorType::COMPLEX_PER_VECTOR;
2255
2256    for (i, m) in (&multiplier[..main_loop_count]).iter().enumerate() {
2257        let left = input.load_complex(i * T::VectorType::COMPLEX_PER_VECTOR);
2258
2259        // Do a complex multiplication between `left` and `right`
2260        let product = mul_complex_conjugated(left, *m);
2261
2262        // Store the result
2263        output.store_complex(product, i * T::VectorType::COMPLEX_PER_VECTOR);
2264    }
2265
2266    // Process the remainder, if there is one
2267    if remainder_count > 0 {
2268        let remainder_index = input.len() - remainder_count;
2269        let remainder_multiplier = multiplier.last().unwrap();
2270        match remainder_count {
2271            1 => {
2272                let left = input.load_partial1_complex(remainder_index);
2273                let product = mul_complex_conjugated(left, remainder_multiplier.lo());
2274                output.store_partial1_complex(product, remainder_index);
2275            }
2276            2 => {
2277                let left = input.load_partial2_complex(remainder_index);
2278                let product = mul_complex_conjugated(left, remainder_multiplier.lo());
2279                output.store_partial2_complex(product, remainder_index);
2280            }
2281            3 => {
2282                let left = input.load_partial3_complex(remainder_index);
2283                let product = mul_complex_conjugated(left, *remainder_multiplier);
2284                output.store_partial3_complex(product, remainder_index);
2285            }
2286            _ => unreachable!(),
2287        }
2288    }
2289}