rustfft/sse/
sse_radix4.rs

1use num_complex::Complex;
2
3use std::any::TypeId;
4use std::sync::Arc;
5
6use crate::array_utils::{bitreversed_transpose, workaround_transmute_mut};
7use crate::{common::FftNum, FftDirection};
8use crate::{Direction, Fft, Length};
9
10use super::SseNum;
11
12use super::sse_vector::{Rotation90, SseArray, SseArrayMut, SseVector};
13
14/// FFT algorithm optimized for power-of-two sizes, SSE accelerated version.
15/// This is designed to be used via a Planner, and not created directly.
16
17pub struct SseRadix4<S: SseNum, T> {
18    twiddles: Box<[S::VectorType]>,
19    rotation: Rotation90<S::VectorType>,
20
21    base_fft: Arc<dyn Fft<T>>,
22    base_len: usize,
23
24    len: usize,
25    direction: FftDirection,
26}
27
28impl<S: SseNum, T: FftNum> SseRadix4<S, T> {
29    /// Constructs a new SseRadix4 which computes FFTs of size 4^k * base_fft.len()
30    #[inline]
31    pub fn new(k: u32, base_fft: Arc<dyn Fft<T>>) -> Result<Self, ()> {
32        // Internal sanity check: Make sure that S == T.
33        // This struct has two generic parameters S and T, but they must always be the same, and are only kept separate to help work around the lack of specialization.
34        // It would be cool if we could do this as a static_assert instead
35        let id_a = TypeId::of::<S>();
36        let id_t = TypeId::of::<T>();
37        assert_eq!(id_a, id_t);
38
39        let has_sse = is_x86_feature_detected!("sse4.1");
40        if has_sse {
41            // Safety: new_with_sse requires the "sse4.1" feature set. Since we know it's present, we're safe
42            Ok(unsafe { Self::new_with_sse(k, base_fft) })
43        } else {
44            Err(())
45        }
46    }
47
48    #[target_feature(enable = "sse4.1")]
49    unsafe fn new_with_sse(k: u32, base_fft: Arc<dyn Fft<T>>) -> Self {
50        let direction = base_fft.fft_direction();
51        let base_len = base_fft.len();
52
53        // note that we can eventually release this restriction - we just need to update the rest of the code in here to handle remainders
54        assert!(base_len % (2 * S::VectorType::COMPLEX_PER_VECTOR) == 0 && base_len > 0);
55
56        let len = base_len * (1 << (k * 2));
57
58        // precompute the twiddle factors this algorithm will use.
59        // we're doing the same precomputation of twiddle factors as the mixed radix algorithm where width=4 and height=len/4
60        // but mixed radix only does one step and then calls itself recusrively, and this algorithm does every layer all the way down
61        // so we're going to pack all the "layers" of twiddle factors into a single array, starting with the bottom layer and going up
62        const ROW_COUNT: usize = 4;
63        let mut cross_fft_len = base_len * ROW_COUNT;
64        let mut twiddle_factors = Vec::with_capacity(len * 2);
65        while cross_fft_len <= len {
66            let num_scalar_columns = cross_fft_len / ROW_COUNT;
67            let num_vector_columns = num_scalar_columns / S::VectorType::COMPLEX_PER_VECTOR;
68
69            for i in 0..num_vector_columns {
70                for k in 1..ROW_COUNT {
71                    twiddle_factors.push(SseVector::make_mixedradix_twiddle_chunk(
72                        i * S::VectorType::COMPLEX_PER_VECTOR,
73                        k,
74                        cross_fft_len,
75                        direction,
76                    ));
77                }
78            }
79            cross_fft_len *= ROW_COUNT;
80        }
81
82        Self {
83            twiddles: twiddle_factors.into_boxed_slice(),
84            rotation: SseVector::make_rotate90(direction),
85
86            base_fft,
87            base_len,
88
89            len,
90            direction,
91        }
92    }
93
94    #[target_feature(enable = "sse4.1")]
95    unsafe fn perform_fft_immut(
96        &self,
97        input: &[Complex<T>],
98        output: &mut [Complex<T>],
99        _scratch: &mut [Complex<T>],
100    ) {
101        // copy the data into the output vector
102        if self.len() == self.base_len {
103            output.copy_from_slice(input);
104        } else {
105            bitreversed_transpose::<Complex<T>, 4>(self.base_len, input, output);
106        }
107
108        // Base-level FFTs
109        self.base_fft.process_with_scratch(output, &mut []);
110
111        // cross-FFTs
112        const ROW_COUNT: usize = 4;
113        let mut cross_fft_len = self.base_len * ROW_COUNT;
114        let mut layer_twiddles: &[S::VectorType] = &self.twiddles;
115
116        while cross_fft_len <= input.len() {
117            let num_rows = input.len() / cross_fft_len;
118            let num_scalar_columns = cross_fft_len / ROW_COUNT;
119            let num_vector_columns = num_scalar_columns / S::VectorType::COMPLEX_PER_VECTOR;
120
121            for i in 0..num_rows {
122                butterfly_4::<S, T>(
123                    &mut output[i * cross_fft_len..],
124                    layer_twiddles,
125                    num_scalar_columns,
126                    &self.rotation,
127                )
128            }
129
130            // skip past all the twiddle factors used in this layer
131            let twiddle_offset = num_vector_columns * (ROW_COUNT - 1);
132            layer_twiddles = &layer_twiddles[twiddle_offset..];
133
134            cross_fft_len *= ROW_COUNT;
135        }
136    }
137
138    #[target_feature(enable = "sse4.1")]
139    unsafe fn perform_fft_out_of_place(
140        &self,
141        input: &[Complex<T>],
142        output: &mut [Complex<T>],
143        _scratch: &mut [Complex<T>],
144    ) {
145        self.perform_fft_immut(input, output, _scratch);
146    }
147}
148boilerplate_fft_sse_oop!(SseRadix4, |this: &SseRadix4<_, _>| this.len);
149
150#[target_feature(enable = "sse4.1")]
151unsafe fn butterfly_4<S: SseNum, T: FftNum>(
152    data: &mut [Complex<T>],
153    twiddles: &[S::VectorType],
154    num_ffts: usize,
155    rotation: &Rotation90<S::VectorType>,
156) {
157    let unroll_offset = S::VectorType::COMPLEX_PER_VECTOR;
158
159    let mut idx = 0usize;
160    let mut buffer: &mut [Complex<S>] = workaround_transmute_mut(data);
161    for tw in twiddles
162        .chunks_exact(6)
163        .take(num_ffts / (S::VectorType::COMPLEX_PER_VECTOR * 2))
164    {
165        let mut scratcha = [
166            buffer.load_complex(idx + 0 * num_ffts),
167            buffer.load_complex(idx + 1 * num_ffts),
168            buffer.load_complex(idx + 2 * num_ffts),
169            buffer.load_complex(idx + 3 * num_ffts),
170        ];
171        let mut scratchb = [
172            buffer.load_complex(idx + 0 * num_ffts + unroll_offset),
173            buffer.load_complex(idx + 1 * num_ffts + unroll_offset),
174            buffer.load_complex(idx + 2 * num_ffts + unroll_offset),
175            buffer.load_complex(idx + 3 * num_ffts + unroll_offset),
176        ];
177
178        scratcha[1] = SseVector::mul_complex(scratcha[1], tw[0]);
179        scratcha[2] = SseVector::mul_complex(scratcha[2], tw[1]);
180        scratcha[3] = SseVector::mul_complex(scratcha[3], tw[2]);
181        scratchb[1] = SseVector::mul_complex(scratchb[1], tw[3]);
182        scratchb[2] = SseVector::mul_complex(scratchb[2], tw[4]);
183        scratchb[3] = SseVector::mul_complex(scratchb[3], tw[5]);
184
185        let scratcha = SseVector::column_butterfly4(scratcha, *rotation);
186        let scratchb = SseVector::column_butterfly4(scratchb, *rotation);
187
188        buffer.store_complex(scratcha[0], idx + 0 * num_ffts);
189        buffer.store_complex(scratchb[0], idx + 0 * num_ffts + unroll_offset);
190        buffer.store_complex(scratcha[1], idx + 1 * num_ffts);
191        buffer.store_complex(scratchb[1], idx + 1 * num_ffts + unroll_offset);
192        buffer.store_complex(scratcha[2], idx + 2 * num_ffts);
193        buffer.store_complex(scratchb[2], idx + 2 * num_ffts + unroll_offset);
194        buffer.store_complex(scratcha[3], idx + 3 * num_ffts);
195        buffer.store_complex(scratchb[3], idx + 3 * num_ffts + unroll_offset);
196
197        idx += S::VectorType::COMPLEX_PER_VECTOR * 2;
198    }
199}
200
201#[cfg(test)]
202mod unit_tests {
203    use super::*;
204    use crate::test_utils::{check_fft_algorithm, construct_base};
205
206    #[test]
207    fn test_sse_radix4_64() {
208        for base in [2, 4, 6, 8, 12, 16] {
209            let base_forward = construct_base(base, FftDirection::Forward);
210            let base_inverse = construct_base(base, FftDirection::Inverse);
211            for k in 0..4 {
212                test_sse_radix4_64_with_base(k, Arc::clone(&base_forward));
213                test_sse_radix4_64_with_base(k, Arc::clone(&base_inverse));
214            }
215        }
216    }
217
218    fn test_sse_radix4_64_with_base(k: u32, base_fft: Arc<dyn Fft<f64>>) {
219        let len = base_fft.len() * 4usize.pow(k);
220        let direction = base_fft.fft_direction();
221        let fft = SseRadix4::<f64, f64>::new(k, base_fft).unwrap();
222        check_fft_algorithm::<f64>(&fft, len, direction);
223    }
224
225    #[test]
226    fn test_sse_radix4_32() {
227        for base in [4, 8, 12, 16] {
228            let base_forward = construct_base(base, FftDirection::Forward);
229            let base_inverse = construct_base(base, FftDirection::Inverse);
230            for k in 0..4 {
231                test_sse_radix4_32_with_base(k, Arc::clone(&base_forward));
232                test_sse_radix4_32_with_base(k, Arc::clone(&base_inverse));
233            }
234        }
235    }
236
237    fn test_sse_radix4_32_with_base(k: u32, base_fft: Arc<dyn Fft<f32>>) {
238        let len = base_fft.len() * 4usize.pow(k);
239        let direction = base_fft.fft_direction();
240        let fft = SseRadix4::<f32, f32>::new(k, base_fft).unwrap();
241        check_fft_algorithm::<f32>(&fft, len, direction);
242    }
243}
rustfft/sse/sse_radix4.rs

rustfft/sse/
sse_radix4.rs
