Skip to main content

rustfft/avx/
avx_raders.rs

1use std::convert::TryInto;
2use std::sync::Arc;
3use std::{any::TypeId, arch::x86_64::*};
4
5use num_complex::Complex;
6use num_integer::{div_ceil, Integer};
7use num_traits::Zero;
8use primal_check::miller_rabin;
9use strength_reduce::StrengthReducedUsize;
10
11use crate::{array_utils, FftDirection};
12use crate::{math_utils, twiddles};
13use crate::{Direction, Fft, FftNum, Length};
14
15use super::avx_vector;
16use super::{
17    avx_vector::{AvxArray, AvxArrayMut, AvxVector, AvxVector128, AvxVector256},
18    AvxNum,
19};
20
21// This struct wraps the necessary data to compute (a * b) % divisor, where b and divisor are determined at runtime but rarely change, and a changes on every call.
22// It's written using AVX2 instructions and assumes the input a are 64-bit integers, and has a restriction that each a, b, and divisor must be 31-bit numbers or smaller.
23#[derive(Clone)]
24struct VectorizedMultiplyMod {
25    b: __m256i,
26    divisor: __m256i,
27    intermediate: __m256i,
28}
29
30impl VectorizedMultiplyMod {
31    #[target_feature(enable = "avx")]
32    unsafe fn new(b: u32, divisor: u32) -> Self {
33        assert!(
34            divisor.leading_zeros() > 0,
35            "divisor must be less than {}, got {}",
36            1 << 31,
37            divisor
38        );
39
40        let b = b % divisor;
41        let intermediate = ((b as i64) << 32) / divisor as i64;
42
43        Self {
44            b: _mm256_set1_epi64x(b as i64),
45            divisor: _mm256_set1_epi64x(divisor as i64),
46            intermediate: _mm256_set1_epi64x(intermediate),
47        }
48    }
49
50    // Input: 4 unsigned 64-bit numbers, each less than 2^30
51    // Output: (x * multiplier) % divisor for each x in input
52    #[allow(unused)]
53    #[inline(always)]
54    unsafe fn mul_rem(&self, a: __m256i) -> __m256i {
55        // Pretty hacky, but we need to prove to the compiler that each entry of the divisor is a 32-bit number, by blending the divisor vector with zeroes in the upper bits of each number.
56        // If we don't do this manually, the compiler will do it anyways, but only for _mm256_mul_epu32, not for the _mm256_sub_epi64 correction step at the end
57        // That inconstistency results in sub-optimal codegen where the compiler inserts extra code to handle the case where divisor is 64-bit. It also results in using one more register than necessary.
58        // Since we know that can't happen, we can placate the compiler by explicitly zeroing the upper 32 bit of each divisor and relying on the compiler to lift it out of the loop.
59        let masked_divisor = _mm256_blend_epi32(self.divisor, _mm256_setzero_si256(), 0xAA);
60
61        // compute the integer quotient of (a * b) / divisor. Our precomputed intermediate value lets us skip the expensive division via arithmetic strength reduction
62        let quotient = _mm256_srli_epi64(_mm256_mul_epu32(a, self.intermediate), 32);
63
64        // Now we can compute numerator - quotient * divisor to get the remanider
65        let numerator = _mm256_mul_epu32(a, self.b);
66        let quotient_product = _mm256_mul_epu32(quotient, masked_divisor);
67
68        // Standard remainder formula: remainder = numerator - quotient * divisor
69        let remainder = _mm256_sub_epi64(numerator, quotient_product);
70
71        // it's possible for the "remainder" to end up between divisor and 2 * divisor. so we'll subtract divisor from remainder, which will make some of the result negative
72        // We can then use the subtracted result as the input to a blendv. Sadly avx doesn't have a blendv_epi32 or blendv_epi64, so we're gonna do blendv_pd instead
73        // this works because blendv looks at the uppermost bit to decide which variable to use, and for a two's complement i64, the upper most bit is 1 when the number is negative!
74        // So when the subtraction result is negative, the uppermost bit is 1, which means the blend will choose the second param, which is the unsubtracted remainder
75        let casted_remainder = _mm256_castsi256_pd(remainder);
76        let subtracted_remainder = _mm256_castsi256_pd(_mm256_sub_epi64(remainder, masked_divisor));
77        let wrapped_remainder = _mm256_castpd_si256(_mm256_blendv_pd(
78            subtracted_remainder,
79            casted_remainder,
80            subtracted_remainder,
81        ));
82        wrapped_remainder
83    }
84}
85
86/// Implementation of Rader's Algorithm, using AVX2 instructions.
87///
88/// This algorithm computes a prime-sized FFT in O(nlogn) time. It does this by converting this size n FFT into a
89/// size (n - 1) which is guaranteed to be composite.
90///
91/// The worst case for this algorithm is when (n - 1) is 2 * prime, resulting in a
92/// [Cunningham Chain](https://en.wikipedia.org/wiki/Cunningham_chain)
93///
94/// Rader's Algorithm is relatively expensive compared to other FFT algorithms. Benchmarking shows that it is up to
95/// an order of magnitude slower than similar composite sizes.
96pub struct RadersAvx2<A: AvxNum, T> {
97    input_index_multiplier: VectorizedMultiplyMod,
98    input_index_init: __m256i,
99
100    output_index_mapping: Box<[__m128i]>,
101    twiddles: Box<[A::VectorType]>,
102
103    inner_fft: Arc<dyn Fft<T>>,
104
105    len: usize,
106
107    inplace_scratch_len: usize,
108    outofplace_scratch_len: usize,
109    immut_scratch_len: usize,
110    direction: FftDirection,
111
112    _phantom: std::marker::PhantomData<T>,
113}
114
115impl<A: AvxNum, T: FftNum> RadersAvx2<A, T> {
116    /// Preallocates necessary arrays and precomputes necessary data to efficiently compute the FFT
117    /// Returns Ok(instance) if this machine has the required instruction sets ("avx", "fma", and "avx2"), Err() if some instruction sets are missing
118    ///
119    /// # Panics
120    /// Panics if `inner_fft_len() + 1` is not a prime number.
121    #[inline]
122    pub fn new(inner_fft: Arc<dyn Fft<T>>) -> Result<Self, ()> {
123        // Internal sanity check: Make sure that A == T.
124        // This struct has two generic parameters A and T, but they must always be the same, and are only kept separate to help work around the lack of specialization.
125        // It would be cool if we could do this as a static_assert instead
126        let id_a = TypeId::of::<A>();
127        let id_t = TypeId::of::<T>();
128        assert_eq!(id_a, id_t);
129
130        let has_avx = is_x86_feature_detected!("avx");
131        let has_avx2 = is_x86_feature_detected!("avx2");
132        let has_fma = is_x86_feature_detected!("fma");
133        if has_avx && has_avx2 && has_fma {
134            // Safety: new_with_avx2 requires the "avx" feature set. Since we know it's present, we're safe
135            Ok(unsafe { Self::new_with_avx(inner_fft) })
136        } else {
137            Err(())
138        }
139    }
140
141    #[target_feature(enable = "avx")]
142    unsafe fn new_with_avx(inner_fft: Arc<dyn Fft<T>>) -> Self {
143        let inner_fft_len = inner_fft.len();
144        let len = inner_fft_len + 1;
145        assert!(miller_rabin(len as u64), "For raders algorithm, inner_fft.len() + 1 must be prime. Expected prime number, got {} + 1 = {}", inner_fft_len, len);
146
147        let direction = inner_fft.fft_direction();
148        let reduced_len = StrengthReducedUsize::new(len);
149
150        // compute the primitive root and its inverse for this size
151        let primitive_root = math_utils::primitive_root(len as u64).unwrap() as usize;
152
153        // compute the multiplicative inverse of primative_root mod len and vice versa.
154        // i64::extended_gcd will compute both the inverse of left mod right, and the inverse of right mod left, but we're only goingto use one of them
155        // the primtive root inverse might be negative, if so make it positive by wrapping
156        let gcd_data = i64::extended_gcd(&(primitive_root as i64), &(len as i64));
157        let primitive_root_inverse = if gcd_data.x >= 0 {
158            gcd_data.x
159        } else {
160            gcd_data.x + len as i64
161        } as usize;
162
163        // precompute the coefficients to use inside the process method
164        let inner_fft_scale = T::one() / T::from_usize(inner_fft_len).unwrap();
165        let mut inner_fft_input = vec![Complex::zero(); inner_fft_len];
166        let mut twiddle_input = 1;
167        for input_cell in &mut inner_fft_input {
168            let twiddle = twiddles::compute_twiddle(twiddle_input, len, direction);
169            *input_cell = twiddle * inner_fft_scale;
170
171            twiddle_input = (twiddle_input * primitive_root_inverse) % reduced_len;
172        }
173
174        let required_inner_scratch = inner_fft.get_inplace_scratch_len();
175        let extra_inner_scratch = if required_inner_scratch <= inner_fft_len {
176            0
177        } else {
178            required_inner_scratch
179        };
180
181        //precompute a FFT of our reordered twiddle factors
182        let mut inner_fft_scratch = vec![Zero::zero(); required_inner_scratch];
183        inner_fft.process_with_scratch(&mut inner_fft_input, &mut inner_fft_scratch);
184
185        // When computing the FFT, we'll want this array to be pre-conjugated, so conjugate it. at the same time, convert it to vectors for convenient use later.
186        let conjugation_mask =
187            AvxVector256::broadcast_complex_elements(Complex::new(A::zero(), -A::zero()));
188
189        let inner_fft_multiplier: Box<[_]> = {
190            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
191            let transmuted_inner_input: &mut [Complex<A>] =
192                array_utils::workaround_transmute_mut(&mut inner_fft_input);
193
194            transmuted_inner_input
195                .chunks(A::VectorType::COMPLEX_PER_VECTOR)
196                .map(|chunk| {
197                    let chunk_vector = match chunk.len() {
198                        1 => chunk.load_partial1_complex(0).zero_extend(),
199                        2 => {
200                            if chunk.len() == A::VectorType::COMPLEX_PER_VECTOR {
201                                chunk.load_complex(0)
202                            } else {
203                                chunk.load_partial2_complex(0).zero_extend()
204                            }
205                        }
206                        3 => chunk.load_partial3_complex(0),
207                        4 => chunk.load_complex(0),
208                        _ => unreachable!(),
209                    };
210                    AvxVector::xor(chunk_vector, conjugation_mask) // compute our conjugation by xoring our data with a precomputed mask
211                })
212                .collect()
213        };
214
215        // Set up the data for our input index remapping computation
216        const NUM_POWERS: usize = 5;
217        let mut root_powers = [0; NUM_POWERS];
218        let mut current_power = 1;
219        for i in 0..NUM_POWERS {
220            root_powers[i] = current_power;
221            current_power = (current_power * primitive_root) % reduced_len;
222        }
223
224        let (input_index_multiplier, input_index_init) = if A::VectorType::COMPLEX_PER_VECTOR == 4 {
225            (
226                VectorizedMultiplyMod::new(root_powers[4] as u32, len as u32),
227                _mm256_loadu_si256(root_powers.as_ptr().add(1) as *const __m256i),
228            )
229        } else {
230            let duplicated_powers = [
231                root_powers[1],
232                root_powers[1],
233                root_powers[2],
234                root_powers[2],
235            ];
236            (
237                VectorizedMultiplyMod::new(root_powers[2] as u32, len as u32),
238                _mm256_loadu_si256(duplicated_powers.as_ptr() as *const __m256i),
239            )
240        };
241
242        // Set up our output index remapping. Ideally we could compute the output indexes on the fly, but the output reindexing requires scatter, which doesn't exist until avx-512
243        // Instead, we can invert the scatter indexes to be gather indexes. But if there's an algorithmic way to compute this, I don't know what it is --
244        // so we won't be able to compute it on the fly with some sort of VectorizedMultiplyMod thing. Instead, we're going to precompute the inverted mapping and gather from that mapping.
245        // We want enough elements in our array to fill out an entire set of vectors so that we don't have to deal with any partial indexes etc.
246        let mapping_size = 1 + div_ceil(len, A::VectorType::COMPLEX_PER_VECTOR)
247            * A::VectorType::COMPLEX_PER_VECTOR;
248        let mut output_mapping_inverse: Vec<i32> = vec![0; mapping_size];
249        let mut output_index = 1;
250        for i in 1..len {
251            output_index = (output_index * primitive_root_inverse) % reduced_len;
252            output_mapping_inverse[output_index] = i.try_into().unwrap();
253        }
254
255        // the actual vector of indexes depends on whether we're f32 or f64
256        let output_index_mapping = if A::VectorType::COMPLEX_PER_VECTOR == 4 {
257            (&output_mapping_inverse[1..])
258                .chunks_exact(A::VectorType::COMPLEX_PER_VECTOR)
259                .map(|chunk| _mm_loadu_si128(chunk.as_ptr() as *const __m128i))
260                .collect::<Box<[__m128i]>>()
261        } else {
262            (&output_mapping_inverse[1..])
263                .chunks_exact(A::VectorType::COMPLEX_PER_VECTOR)
264                .map(|chunk| {
265                    let duplicated_indexes = [chunk[0], chunk[0], chunk[1], chunk[1]];
266                    _mm_loadu_si128(duplicated_indexes.as_ptr() as *const __m128i)
267                })
268                .collect::<Box<[__m128i]>>()
269        };
270        let inplace_scratch_len = len + extra_inner_scratch;
271        Self {
272            input_index_multiplier,
273            input_index_init,
274
275            output_index_mapping,
276
277            inner_fft: inner_fft,
278            twiddles: inner_fft_multiplier,
279
280            len,
281
282            inplace_scratch_len,
283            outofplace_scratch_len: extra_inner_scratch,
284            immut_scratch_len: inner_fft_len + required_inner_scratch + 1,
285            direction,
286
287            _phantom: std::marker::PhantomData,
288        }
289    }
290
291    // Do the necessary setup for rader's algorithm: Reorder the inputs into the output buffer, gather a sum of all inputs. Return the first input, and the aum of all inputs
292    #[target_feature(enable = "avx2", enable = "avx", enable = "fma")]
293    unsafe fn prepare_raders(&self, input: &[Complex<A>], output: &mut [Complex<A>]) {
294        let mut indexes = self.input_index_init;
295
296        let index_multiplier = self.input_index_multiplier.clone();
297
298        // loop over the output array and use AVX gathers to reorder data from the input
299        let mut chunks_iter =
300            (&mut output[1..]).chunks_exact_mut(A::VectorType::COMPLEX_PER_VECTOR);
301        for mut chunk in chunks_iter.by_ref() {
302            let gathered_elements =
303                A::VectorType::gather_complex_avx2_index64(input.as_ptr(), indexes);
304
305            // advance our indexes
306            indexes = index_multiplier.mul_rem(indexes);
307
308            // Store this chunk
309            chunk.store_complex(gathered_elements, 0);
310        }
311
312        // at this point, we either have 0 or 2 remaining elements to gather. because we know our length ends in 1 or 3. so when we subtract 1 for the inner FFT, that gives us 0 or 2
313        let mut output_remainder = chunks_iter.into_remainder();
314        if output_remainder.len() == 2 {
315            let half_data = AvxVector128::gather64_complex_avx2(
316                input.as_ptr(),
317                _mm256_castsi256_si128(indexes),
318            );
319
320            // store the remainder in the last chunk
321            output_remainder.store_partial2_complex(half_data, 0);
322        }
323    }
324
325    // Do the necessary finalization for rader's algorithm: Reorder the inputs into the output buffer, conjugating the input as we go, and add the first input value to every output value
326    #[target_feature(enable = "avx2", enable = "avx", enable = "fma")]
327    unsafe fn finalize_raders(&self, input: &[Complex<A>], output: &mut [Complex<A>]) {
328        // We need to conjugate elements as a part of the finalization step, and sadly we can't roll it into any other instructions. So we'll do it via an xor.
329        let conjugation_mask =
330            AvxVector256::broadcast_complex_elements(Complex::new(A::zero(), -A::zero()));
331
332        let mut chunks_iter =
333            (&mut output[1..]).chunks_exact_mut(A::VectorType::COMPLEX_PER_VECTOR);
334        for (i, mut chunk) in chunks_iter.by_ref().enumerate() {
335            let index_chunk = *self.output_index_mapping.get_unchecked(i);
336            let gathered_elements =
337                A::VectorType::gather_complex_avx2_index32(input.as_ptr(), index_chunk);
338
339            let conjugated_elements = AvxVector::xor(gathered_elements, conjugation_mask);
340            chunk.store_complex(conjugated_elements, 0);
341        }
342
343        // at this point, we either have 0 or 2 remaining elements to gather. because we know our length ends in 1 or 3. so when we subtract 1 for the inner FFT, that gives us 0 or 2
344        let mut output_remainder = chunks_iter.into_remainder();
345        if output_remainder.len() == 2 {
346            let index_chunk = *self
347                .output_index_mapping
348                .get_unchecked(self.output_index_mapping.len() - 1);
349            let half_data = AvxVector128::gather32_complex_avx2(input.as_ptr(), index_chunk);
350
351            let conjugated_elements = AvxVector::xor(half_data, conjugation_mask.lo());
352            output_remainder.store_partial2_complex(conjugated_elements, 0);
353        }
354    }
355
356    fn perform_fft_immut(
357        &self,
358        input: &[Complex<T>],
359        output: &mut [Complex<T>],
360        scratch: &mut [Complex<T>],
361    ) {
362        unsafe {
363            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
364            let transmuted_input: &[Complex<A>] = array_utils::workaround_transmute(input);
365            let transmuted_output: &mut [Complex<A>] =
366                array_utils::workaround_transmute_mut(output);
367            self.prepare_raders(transmuted_input, transmuted_output)
368        }
369
370        let (first_input, _) = input.split_first().unwrap();
371        let (first_output, inner_output) = output.split_first_mut().unwrap();
372        let (scratch2, extra_scratch) = scratch.split_at_mut(self.len());
373        let (_, scratch) = scratch2.split_first_mut().unwrap();
374
375        self.inner_fft.process_with_scratch(inner_output, scratch);
376
377        // inner_output[0] now contains the sum of elements 1..n. we want the sum of all inputs, so all we need to do is add the first input
378        *first_output = inner_output[0] + *first_input;
379
380        // multiply the inner result with our cached setup data
381        // also conjugate every entry. this sets us up to do an inverse FFT
382        // (because an inverse FFT is equivalent to a normal FFT where you conjugate both the inputs and outputs)
383        unsafe {
384            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
385            let transmuted_inner_input: &mut [Complex<A>] =
386                array_utils::workaround_transmute_mut(scratch);
387            let transmuted_inner_output: &mut [Complex<A>] =
388                array_utils::workaround_transmute_mut(inner_output);
389            avx_vector::pairwise_complex_mul_conjugated(
390                transmuted_inner_output,
391                transmuted_inner_input,
392                &self.twiddles,
393            )
394        };
395
396        // We need to add the first input value to all output values. We can accomplish this by adding it to the DC input of our inner ifft.
397        // Of course, we have to conjugate it, just like we conjugated the complex multiplied above
398        scratch[0] = scratch[0] + first_input.conj();
399
400        self.inner_fft.process_with_scratch(scratch, extra_scratch);
401        scratch2[0] = *first_input;
402
403        // copy the final values into the output, reordering as we go
404        unsafe {
405            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
406            let transmuted_input: &mut [Complex<A>] =
407                array_utils::workaround_transmute_mut(scratch2);
408            let transmuted_output: &mut [Complex<A>] =
409                array_utils::workaround_transmute_mut(output);
410            self.finalize_raders(transmuted_input, transmuted_output);
411        }
412    }
413
414    fn perform_fft_out_of_place(
415        &self,
416        input: &mut [Complex<T>],
417        output: &mut [Complex<T>],
418        scratch: &mut [Complex<T>],
419    ) {
420        unsafe {
421            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
422            let transmuted_input: &mut [Complex<A>] = array_utils::workaround_transmute_mut(input);
423            let transmuted_output: &mut [Complex<A>] =
424                array_utils::workaround_transmute_mut(output);
425            self.prepare_raders(transmuted_input, transmuted_output)
426        }
427
428        let (first_input, inner_input) = input.split_first_mut().unwrap();
429        let (first_output, inner_output) = output.split_first_mut().unwrap();
430
431        // perform the first of two inner FFTs
432        let inner_scratch = if scratch.len() > 0 {
433            &mut scratch[..]
434        } else {
435            &mut inner_input[..]
436        };
437        self.inner_fft
438            .process_with_scratch(inner_output, inner_scratch);
439
440        // inner_output[0] now contains the sum of elements 1..n. we want the sum of all inputs, so all we need to do is add the first input
441        *first_output = inner_output[0] + *first_input;
442
443        // multiply the inner result with our cached setup data
444        // also conjugate every entry. this sets us up to do an inverse FFT
445        // (because an inverse FFT is equivalent to a normal FFT where you conjugate both the inputs and outputs)
446        unsafe {
447            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
448            let transmuted_inner_output: &mut [Complex<A>] =
449                array_utils::workaround_transmute_mut(inner_input);
450            let transmuted_inner_input: &mut [Complex<A>] =
451                array_utils::workaround_transmute_mut(inner_output);
452            avx_vector::pairwise_complex_mul_conjugated(
453                transmuted_inner_input,
454                transmuted_inner_output,
455                &self.twiddles,
456            )
457        };
458
459        // We need to add the first input value to all output values. We can accomplish this by adding it to the DC input of our inner ifft.
460        // Of course, we have to conjugate it, just like we conjugated the complex multiplied above
461        inner_input[0] = inner_input[0] + first_input.conj();
462
463        // execute the second FFT
464        let inner_scratch = if scratch.len() > 0 {
465            scratch
466        } else {
467            &mut inner_output[..]
468        };
469        self.inner_fft
470            .process_with_scratch(inner_input, inner_scratch);
471
472        // copy the final values into the output, reordering as we go
473        unsafe {
474            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
475            let transmuted_input: &mut [Complex<A>] = array_utils::workaround_transmute_mut(input);
476            let transmuted_output: &mut [Complex<A>] =
477                array_utils::workaround_transmute_mut(output);
478            self.finalize_raders(transmuted_input, transmuted_output);
479        }
480    }
481    fn perform_fft_inplace(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
482        let (scratch, extra_scratch) = scratch.split_at_mut(self.len());
483        unsafe {
484            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
485            let transmuted_scratch: &mut [Complex<A>] =
486                array_utils::workaround_transmute_mut(scratch);
487            let transmuted_buffer: &mut [Complex<A>] =
488                array_utils::workaround_transmute_mut(buffer);
489            self.prepare_raders(transmuted_buffer, transmuted_scratch)
490        }
491
492        let first_input = buffer[0];
493
494        let truncated_scratch = &mut scratch[1..];
495
496        // perform the first of two inner FFTs
497        let inner_scratch = if extra_scratch.len() > 0 {
498            extra_scratch
499        } else {
500            &mut buffer[..]
501        };
502        self.inner_fft
503            .process_with_scratch(truncated_scratch, inner_scratch);
504
505        // truncated_scratch[0] now contains the sum of elements 1..n. we want the sum of all inputs, so all we need to do is add the first input
506        let first_output = first_input + truncated_scratch[0];
507
508        // multiply the inner result with our cached setup data
509        // also conjugate every entry. this sets us up to do an inverse FFT
510        // (because an inverse FFT is equivalent to a normal FFT where you conjugate both the inputs and outputs)
511        unsafe {
512            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
513            let transmuted_scratch: &mut [Complex<A>] =
514                array_utils::workaround_transmute_mut(truncated_scratch);
515            avx_vector::pairwise_complex_mul_assign_conjugated(transmuted_scratch, &self.twiddles)
516        };
517
518        // We need to add the first input value to all output values. We can accomplish this by adding it to the DC input of our inner ifft.
519        // Of course, we have to conjugate it, just like we conjugated the complex multiplied above
520        truncated_scratch[0] = truncated_scratch[0] + first_input.conj();
521
522        // execute the second FFT
523        self.inner_fft
524            .process_with_scratch(truncated_scratch, inner_scratch);
525
526        // copy the final values into the output, reordering as we go
527        buffer[0] = first_output;
528        unsafe {
529            // Specialization workaround: See the comments in FftPlannerAvx::new() for why these calls to array_utils::workaround_transmute are necessary
530            let transmuted_scratch: &mut [Complex<A>] =
531                array_utils::workaround_transmute_mut(scratch);
532            let transmuted_buffer: &mut [Complex<A>] =
533                array_utils::workaround_transmute_mut(buffer);
534            self.finalize_raders(transmuted_scratch, transmuted_buffer);
535        }
536    }
537}
538boilerplate_avx_fft!(
539    RadersAvx2,
540    |this: &RadersAvx2<_, _>| this.len,
541    |this: &RadersAvx2<_, _>| this.inplace_scratch_len,
542    |this: &RadersAvx2<_, _>| this.outofplace_scratch_len,
543    |this: &RadersAvx2<_, _>| this.immut_scratch_len
544);
545
546#[cfg(test)]
547mod unit_tests {
548    use num_traits::Float;
549    use rand::distributions::uniform::SampleUniform;
550
551    use super::*;
552    use crate::algorithm::Dft;
553    use crate::test_utils::check_fft_algorithm;
554    use std::sync::Arc;
555
556    #[test]
557    fn test_raders_avx_f32() {
558        for len in 3..100 {
559            if miller_rabin(len as u64) {
560                test_raders_with_length::<f32>(len, FftDirection::Forward);
561                test_raders_with_length::<f32>(len, FftDirection::Inverse);
562            }
563        }
564    }
565
566    #[test]
567    fn test_raders_avx_f64() {
568        for len in 3..100 {
569            if miller_rabin(len as u64) {
570                test_raders_with_length::<f64>(len, FftDirection::Forward);
571                test_raders_with_length::<f64>(len, FftDirection::Inverse);
572            }
573        }
574    }
575
576    fn test_raders_with_length<T: AvxNum + Float + SampleUniform>(
577        len: usize,
578        direction: FftDirection,
579    ) {
580        let inner_fft = Arc::new(Dft::new(len - 1, direction));
581        let fft = RadersAvx2::<T, T>::new(inner_fft).unwrap();
582
583        check_fft_algorithm::<T>(&fft, len, direction);
584    }
585}