yum/vfft

git clone https://git.yummers.dev/yum/vfft

yumMove large prime radices first7441799

master
branchesmasterall refs →
32.7 KiB919 linesraw
1use crate::util::*;
2use num::complex::Complex;
3use num::traits::{Float, FloatConst};
4
5fn fft_butterfly_radix_2_unsafe<T: Float + FloatConst>(
6    input: &mut [Complex<T>],
7    output: &mut [Complex<T>],
8    stride: usize,
9    big_n: usize,
10    twiddles: &[Complex<T>],
11    _: &ExecutionStage<T>,
12) {
13    let input_ptr = input.as_ptr();
14    let output_ptr = output.as_mut_ptr();
15    for start_idx in 0..stride {
16        for k in 0..big_n / 2 {
17            unsafe {
18                // Collect inputs.
19                let p = *input_ptr.add(start_idx + 2 * k * stride);
20                let q = *input_ptr.add(start_idx + (2 * k + 1) * stride);
21                // Collect relevant twiddles.
22                let ot = twiddles.get_unchecked(k * stride);
23                // Combine.
24                let r = q * ot;
25                *output_ptr.add(start_idx + k * stride) = p + r;
26                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = p - r;
27            }
28        }
29    }
30}
31
32fn fft_butterfly_radix_2_s0_unsafe<T: Float + FloatConst>(
33    input: &mut [Complex<T>],
34    output: &mut [Complex<T>],
35    _: &Matrix<Complex<T>>,
36) {
37    let input_ptr = input.as_ptr();
38    let output_ptr = output.as_mut_ptr();
39    let stride = input.len() / 2;
40    let big_n = 2;
41    for start_idx in 0..stride {
42        for k in 0..big_n / 2 {
43            unsafe {
44                // Collect inputs.
45                let p = *input_ptr.add(start_idx + 2 * k * stride);
46                let q = *input_ptr.add(start_idx + (2 * k + 1) * stride);
47                *output_ptr.add(start_idx + k * stride) = p + q;
48                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = p - q;
49            }
50        }
51    }
52}
53
54fn radix_3_dft<T: Float + FloatConst>(
55    input: &[Complex<T>; 3],
56    _: &Matrix<Complex<T>>,
57) -> [Complex<T>; 3] {
58    const N: usize = 3;
59    let zero = Complex::new(T::zero(), T::zero());
60    let mut result = [zero; N];
61    let half = T::from(0.5).unwrap();
62    let sqrt_3_over_2 =
63        T::from(0.8660254037844386467637231707529361834714026269051903140279034897).unwrap();
64    let sum = input[1] + input[2];
65    let diff = input[1] - input[2];
66    result[0] = input[0] + input[1] + input[2];
67    result[1] = input[0] - sum.scale(half) - mul_i(diff.scale(sqrt_3_over_2));
68    result[2] = input[0] - sum.scale(half) + mul_i(diff.scale(sqrt_3_over_2));
69    result
70}
71
72fn fft_butterfly_radix_3_unsafe<T: Float + FloatConst>(
73    input: &mut [Complex<T>],
74    output: &mut [Complex<T>],
75    stride: usize,
76    big_n: usize,
77    twiddles: &[Complex<T>],
78    stage: &ExecutionStage<T>,
79) {
80    let input_ptr = input.as_ptr();
81    let output_ptr = output.as_mut_ptr();
82    const N: usize = 3;
83    for start_idx in 0..stride {
84        for k in 0..big_n / N {
85            unsafe {
86                // Collect inputs.
87                let inputs: [Complex<T>; N] =
88                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
89                // Collect relevant twiddles.
90                let twiddles: [&Complex<T>; N] =
91                    std::array::from_fn(|idx| twiddles.get_unchecked(idx * k * stride));
92                // Precompute a term used in the inner loop.
93                let it: [Complex<T>; N] = std::array::from_fn(|idx| {
94                    if idx == 0 {
95                        inputs[idx]
96                    } else {
97                        inputs[idx] * twiddles[idx]
98                    }
99                });
100                let result = radix_3_dft(&it, &stage.outer_terms);
101                for row in 0..N {
102                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = result[row];
103                }
104            }
105        }
106    }
107}
108
109fn fft_butterfly_radix_3_s0_unsafe<T: Float + FloatConst>(
110    input: &mut [Complex<T>],
111    output: &mut [Complex<T>],
112    outer_terms: &Matrix<Complex<T>>,
113) {
114    const N: usize = 3;
115    let input_ptr = input.as_ptr();
116    let output_ptr = output.as_mut_ptr();
117    let stride = input.len() / N;
118    let big_n = N;
119    for start_idx in 0..stride {
120        for k in 0..big_n / N {
121            unsafe {
122                // Collect inputs.
123                let inputs: [Complex<T>; N] =
124                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
125                let result = radix_3_dft(&inputs, outer_terms);
126                for row in 0..N {
127                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = result[row];
128                }
129            }
130        }
131    }
132}
133
134fn fft_butterfly_radix_9_unsafe<T: Float + FloatConst>(
135    input: &mut [Complex<T>],
136    output: &mut [Complex<T>],
137    stride: usize,
138    big_n: usize,
139    twiddles: &[Complex<T>],
140    stage: &ExecutionStage<T>,
141) {
142    let input_ptr = input.as_ptr();
143    let output_ptr = output.as_mut_ptr();
144    const N: usize = 9;
145    for start_idx in 0..stride {
146        for k in 0..big_n / N {
147            unsafe {
148                let inputs: [Complex<T>; N] =
149                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
150                let twiddles: [&Complex<T>; N] =
151                    std::array::from_fn(|idx| twiddles.get_unchecked(idx * k * stride));
152                let it: [Complex<T>; N] = std::array::from_fn(|idx| {
153                    if idx == 0 {
154                        inputs[idx]
155                    } else {
156                        inputs[idx] * twiddles[idx]
157                    }
158                });
159
160                let b00 = [it[0], it[3], it[6]];
161                let b01 = [it[1], it[4], it[7]];
162                let b02 = [it[2], it[5], it[8]];
163
164                let b00_out = radix_3_dft(&b00, &stage.outer_terms);
165                let b01_out = radix_3_dft(&b01, &stage.outer_terms);
166                let b02_out = radix_3_dft(&b02, &stage.outer_terms);
167
168                // Recall the general way that outer terms work:
169                // First row is all 1s.
170                // Second row rotates 1/r clockwise.
171                // Third row rotates 2/r clockwise.
172                // outer_terms already takes into account the way rotations
173                // accelerate each row, so we just have to index into the right
174                // row/col.
175                let b10 = [b00_out[0], b01_out[0], b02_out[0]];
176                let b11 = [
177                    b00_out[1],
178                    b01_out[1] * stage.outer_terms.data.get_unchecked(9 + 1),
179                    b02_out[1] * stage.outer_terms.data.get_unchecked(9 + 2),
180                ];
181                let b12 = [
182                    b00_out[2],
183                    b01_out[2] * stage.outer_terms.data.get_unchecked(18 + 1),
184                    b02_out[2] * stage.outer_terms.data.get_unchecked(18 + 2),
185                ];
186
187                let r00 = radix_3_dft(&b10, &stage.outer_terms);
188                let r01 = radix_3_dft(&b11, &stage.outer_terms);
189                let r02 = radix_3_dft(&b12, &stage.outer_terms);
190
191                let results = [
192                    r00[0], r01[0], r02[0], r00[1], r01[1], r02[1], r00[2], r01[2], r02[2],
193                ];
194                for row in 0..N {
195                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = results[row];
196                }
197            }
198        }
199    }
200}
201
202fn fft_butterfly_radix_9_s0_unsafe<T: Float + FloatConst>(
203    input: &mut [Complex<T>],
204    output: &mut [Complex<T>],
205    outer_terms: &Matrix<Complex<T>>,
206) {
207    let input_ptr = input.as_ptr();
208    let output_ptr = output.as_mut_ptr();
209    const N: usize = 9;
210    let stride = input.len() / N;
211    let big_n = N;
212    for start_idx in 0..stride {
213        for k in 0..big_n / N {
214            unsafe {
215                let inputs: [Complex<T>; N] =
216                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
217
218                let b00 = [inputs[0], inputs[3], inputs[6]];
219                let b01 = [inputs[1], inputs[4], inputs[7]];
220                let b02 = [inputs[2], inputs[5], inputs[8]];
221
222                let b00_out = radix_3_dft(&b00, &outer_terms);
223                let b01_out = radix_3_dft(&b01, &outer_terms);
224                let b02_out = radix_3_dft(&b02, &outer_terms);
225
226                // Recall the general way that outer terms work:
227                // First row is all 1s.
228                // Second row rotates 1/r clockwise.
229                // Third row rotates 2/r clockwise.
230                let b10 = [b00_out[0], b01_out[0], b02_out[0]];
231                let b11 = [
232                    b00_out[1],
233                    b01_out[1] * outer_terms.data.get_unchecked(9 + 1),
234                    b02_out[1] * outer_terms.data.get_unchecked(9 + 2),
235                ];
236                let b12 = [
237                    b00_out[2],
238                    b01_out[2] * outer_terms.data.get_unchecked(18 + 1),
239                    b02_out[2] * outer_terms.data.get_unchecked(18 + 2),
240                ];
241
242                let r00 = radix_3_dft(&b10, &outer_terms);
243                let r01 = radix_3_dft(&b11, &outer_terms);
244                let r02 = radix_3_dft(&b12, &outer_terms);
245
246                let results = [
247                    r00[0], r01[0], r02[0], r00[1], r01[1], r02[1], r00[2], r01[2], r02[2],
248                ];
249                for row in 0..N {
250                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = results[row];
251                }
252            }
253        }
254    }
255}
256
257fn odd_radix_dft<const N: usize, T: Float + FloatConst>(
258    input: &[Complex<T>; N],
259    outer_terms: &Matrix<Complex<T>>,
260) -> [Complex<T>; N] {
261    let zero = Complex::new(T::zero(), T::zero());
262    let mut result = [zero; N];
263    // Precompute sums and differences.
264    let mut sums = [zero; N];
265    let mut diffs = [zero; N];
266    for col in 1..N / 2 + 1 {
267        sums[col] = input[col] + input[N - col];
268        diffs[col] = input[col] - input[N - col];
269    }
270
271    // 0,0 term is never indexed, so compute it out of the loop.
272    result[0] = input[0];
273
274    // 0th row is never indexed, so compute it out of the loop.
275    for col in 1..N / 2 + 1 {
276        // outer_terms is always 1 along the first row and column, so we can skip the multiplication.
277        result[0] = result[0] + sums[col];
278    }
279
280    for row in 1..N / 2 + 1 {
281        // Column zero is not indexed below, so precompute here. Recall the
282        // outer_terms is 1 along the first column.
283        //
284        // Per our identity, we split this sum into a "real" coefficient and an
285        // "imaginary" coefficient. Note that they are in fact both complex-valued.
286        let mut re_part = input[0];
287        let mut im_part = zero;
288
289        for col in 1..N / 2 + 1 {
290            unsafe {
291                let w = outer_terms.data.get_unchecked(row * N + col);
292                re_part = re_part + sums[col].scale(w.re);
293                im_part = im_part + diffs[col].scale(w.im);
294            }
295        }
296
297        // Then apply the identity to get both rows.
298        let i_im_part = mul_i(im_part);
299        result[row] = re_part + i_im_part;
300        result[N - row] = re_part - i_im_part;
301    }
302    result
303}
304
305fn fft_butterfly_radix_p_unsafe<const N: usize, T: Float + FloatConst>(
306    input: &mut [Complex<T>],
307    output: &mut [Complex<T>],
308    stride: usize,
309    big_n: usize,
310    twiddles: &[Complex<T>],
311    stage: &ExecutionStage<T>,
312) {
313    let input_ptr = input.as_ptr();
314    let output_ptr = output.as_mut_ptr();
315    for start_idx in 0..stride {
316        for k in 0..big_n / N {
317            unsafe {
318                // Collect inputs.
319                let inputs: [Complex<T>; N] =
320                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
321                // Collect relevant twiddles.
322                let twiddles: [&Complex<T>; N] =
323                    std::array::from_fn(|idx| twiddles.get_unchecked(idx * k * stride));
324                // Precompute a term used in the inner loop.
325                let it: [Complex<T>; N] = std::array::from_fn(|idx| {
326                    if idx == 0 {
327                        inputs[idx]
328                    } else {
329                        inputs[idx] * twiddles[idx]
330                    }
331                });
332                let result = odd_radix_dft(&it, &stage.outer_terms);
333                for row in 0..N {
334                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = result[row];
335                }
336            }
337        }
338    }
339}
340
341fn fft_butterfly_radix_p_s0_unsafe<const N: usize, T: Float + FloatConst>(
342    input: &mut [Complex<T>],
343    output: &mut [Complex<T>],
344    outer_terms: &Matrix<Complex<T>>,
345) {
346    let input_ptr = input.as_ptr();
347    let output_ptr = output.as_mut_ptr();
348    let stride = input.len() / N;
349    let big_n = N;
350    for start_idx in 0..stride {
351        for k in 0..big_n / N {
352            unsafe {
353                // Collect inputs.
354                let inputs: [Complex<T>; N] =
355                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
356                // Precompute a term used in the inner loop.
357                let result = odd_radix_dft(&inputs, outer_terms);
358                for row in 0..N {
359                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = result[row];
360                }
361            }
362        }
363    }
364}
365
366#[inline(always)]
367fn mul_i<T: Float + FloatConst>(x: Complex<T>) -> Complex<T> {
368    Complex::new(-x.im, x.re)
369}
370
371#[inline(always)]
372fn mul_ni<T: Float + FloatConst>(x: Complex<T>) -> Complex<T> {
373    Complex::new(x.im, -x.re)
374}
375
376fn fft_butterfly_radix_4_unsafe<T: Float + FloatConst>(
377    input: &mut [Complex<T>],
378    output: &mut [Complex<T>],
379    stride: usize,
380    big_n: usize,
381    twiddles: &[Complex<T>],
382    _: &ExecutionStage<T>,
383) {
384    let input_ptr = input.as_ptr();
385    let output_ptr = output.as_mut_ptr();
386    for start_idx in 0..stride {
387        for k in 0..big_n / 4 {
388            unsafe {
389                // Collect inputs.
390                let i0 = *input_ptr.add(start_idx + 4 * k * stride);
391                let i1 = *input_ptr.add(start_idx + (4 * k + 1) * stride);
392                let i2 = *input_ptr.add(start_idx + (4 * k + 2) * stride);
393                let i3 = *input_ptr.add(start_idx + (4 * k + 3) * stride);
394                // Collect relevant twiddles.
395                let ot1 = twiddles.get_unchecked(1 * k * stride);
396                let ot2 = twiddles.get_unchecked(2 * k * stride);
397                let ot3 = twiddles.get_unchecked(3 * k * stride);
398
399                let a = i0;
400                let b = ot1 * i1;
401                let c = ot2 * i2;
402                let d = ot3 * i3;
403
404                // To derive this, write the output assignments in terms of
405                // a/b/c/d, then factor out!
406                let ac_sum = a + c;
407                let ac_diff = a - c;
408                let bd_sum = b + d;
409                let bd_diff_ni = mul_ni(b - d);
410
411                *output_ptr.add(start_idx + k * stride) = ac_sum + bd_sum;
412                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = ac_diff + bd_diff_ni;
413                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = ac_sum - bd_sum;
414                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = ac_diff - bd_diff_ni;
415            }
416        }
417    }
418}
419
420fn fft_butterfly_radix_4_s0_unsafe<T: Float + FloatConst>(
421    input: &mut [Complex<T>],
422    output: &mut [Complex<T>],
423    _: &Matrix<Complex<T>>,
424) {
425    let stride = input.len() / 4;
426    let big_n = 4;
427    let input_ptr = input.as_ptr();
428    let output_ptr = output.as_mut_ptr();
429    for start_idx in 0..stride {
430        for k in 0..big_n / 4 {
431            unsafe {
432                // Collect inputs.
433                let i0 = *input_ptr.add(start_idx + 4 * k * stride);
434                let i1 = *input_ptr.add(start_idx + (4 * k + 1) * stride);
435                let i2 = *input_ptr.add(start_idx + (4 * k + 2) * stride);
436                let i3 = *input_ptr.add(start_idx + (4 * k + 3) * stride);
437
438                let a = i0;
439                let b = i1;
440                let c = i2;
441                let d = i3;
442
443                // To derive this, write the output assignments in terms of
444                // a/b/c/d, then factor out!
445                let ac_sum = a + c;
446                let ac_diff = a - c;
447                let bd_sum = b + d;
448                let bd_diff_ni = mul_ni(b - d);
449
450                *output_ptr.add(start_idx + k * stride) = ac_sum + bd_sum;
451                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = ac_diff + bd_diff_ni;
452                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = ac_sum - bd_sum;
453                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = ac_diff - bd_diff_ni;
454            }
455        }
456    }
457}
458
459#[inline(always)]
460fn rot_45<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
461    let s = T::FRAC_1_SQRT_2();
462    // The standard 2D rotation matrix gives:
463    //  [ cos(pi/4) -sin(pi/4)]   [ s -s ]
464    //  [ sin(pi/4)  cos(pi/4)] = [ s  s ]
465    Complex::<T>::new(c.re - c.im, c.re + c.im) * s
466}
467
468#[inline(always)]
469fn rot_90<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
470    // The standard 2D rotation matrix gives:
471    //  [ cos(pi/2) -sin(pi/2)]   [ 0 -1 ]
472    //  [ sin(pi/2)  cos(pi/2)] = [ 1  0 ]
473    Complex::<T>::new(-c.im, c.re)
474}
475
476#[inline(always)]
477fn rot_180<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
478    // The standard 2D rotation matrix gives:
479    //  [ cos(pi) -sin(pi)]   [ -1  0 ]
480    //  [ sin(pi)  cos(pi)] = [  0 -1 ]
481    Complex::<T>::new(-c.re, -c.im)
482}
483
484#[inline(always)]
485fn rot_270<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
486    // The standard 2D rotation matrix gives:
487    //  [ cos(3pi/2) -sin(3pi/2)]   [  0 1 ]
488    //  [ sin(3pi/2)  cos(3pi/2)] = [ -1 0 ]
489    Complex::<T>::new(c.im, -c.re)
490}
491
492fn fft_butterfly_radix_8_unsafe<T: Float + FloatConst>(
493    input: &mut [Complex<T>],
494    output: &mut [Complex<T>],
495    stride: usize,
496    big_n: usize,
497    twiddles: &[Complex<T>],
498    _: &ExecutionStage<T>,
499) {
500    let input_ptr = input.as_ptr();
501    let output_ptr = output.as_mut_ptr();
502    for start_idx in 0..stride {
503        for k in 0..big_n / 8 {
504            unsafe {
505                // Collect inputs.
506                let i0 = *input_ptr.add(start_idx + 8 * k * stride);
507                let i1 = *input_ptr.add(start_idx + (8 * k + 1) * stride);
508                let i2 = *input_ptr.add(start_idx + (8 * k + 2) * stride);
509                let i3 = *input_ptr.add(start_idx + (8 * k + 3) * stride);
510                let i4 = *input_ptr.add(start_idx + (8 * k + 4) * stride);
511                let i5 = *input_ptr.add(start_idx + (8 * k + 5) * stride);
512                let i6 = *input_ptr.add(start_idx + (8 * k + 6) * stride);
513                let i7 = *input_ptr.add(start_idx + (8 * k + 7) * stride);
514
515                // Collect relevant twiddles.
516                let ot1 = twiddles.get_unchecked(1 * k * stride);
517                let ot2 = twiddles.get_unchecked(2 * k * stride);
518                let ot3 = twiddles.get_unchecked(3 * k * stride);
519                let ot4 = twiddles.get_unchecked(4 * k * stride);
520                let ot5 = twiddles.get_unchecked(5 * k * stride);
521                let ot6 = twiddles.get_unchecked(6 * k * stride);
522                let ot7 = twiddles.get_unchecked(7 * k * stride);
523
524                let a = i0;
525                let b = ot1 * i1;
526                let c = ot2 * i2;
527                let d = ot3 * i3;
528                let e = ot4 * i4;
529                let f = ot5 * i5;
530                let g = ot6 * i6;
531                let h = ot7 * i7;
532
533                let ae_sum = a + e;
534                let ae_diff = a - e;
535                let bf_sum = b + f;
536                let bf_diff = b - f;
537                let cg_sum = c + g;
538                let cg_diff = c - g;
539                let dh_sum = d + h;
540                let dh_diff = d - h;
541
542                let w00 = ae_sum + cg_sum;
543                let w01 = ae_sum - cg_sum;
544                let w10 = ae_diff + rot_270(cg_diff);
545                let w11 = ae_diff - rot_270(cg_diff);
546                let x00 = bf_sum + dh_sum;
547                let x01 = rot_270(bf_sum) + rot_90(dh_sum);
548                let x10 = rot_45(rot_270(bf_diff) + rot_180(dh_diff));
549                let x11 = rot_45(rot_180(bf_diff) + rot_270(dh_diff));
550
551                *output_ptr.add(start_idx + k * stride) = w00 + x00;
552                *output_ptr.add(start_idx + (k + big_n / 8) * stride) = w10 + x10;
553                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = w01 + x01;
554                *output_ptr.add(start_idx + (k + 3 * big_n / 8) * stride) = w11 + x11;
555                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = w00 - x00;
556                *output_ptr.add(start_idx + (k + 5 * big_n / 8) * stride) = w10 - x10;
557                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = w01 - x01;
558                *output_ptr.add(start_idx + (k + 7 * big_n / 8) * stride) = w11 - x11;
559            }
560        }
561    }
562}
563
564fn fft_butterfly_radix_8_s0_unsafe<T: Float + FloatConst>(
565    input: &mut [Complex<T>],
566    output: &mut [Complex<T>],
567    _: &Matrix<Complex<T>>,
568) {
569    let stride = input.len() / 8;
570    let big_n = 8;
571    let input_ptr = input.as_ptr();
572    let output_ptr = output.as_mut_ptr();
573    for start_idx in 0..stride {
574        for k in 0..big_n / 8 {
575            unsafe {
576                // Collect inputs.
577                let i0 = *input_ptr.add(start_idx + 8 * k * stride);
578                let i1 = *input_ptr.add(start_idx + (8 * k + 1) * stride);
579                let i2 = *input_ptr.add(start_idx + (8 * k + 2) * stride);
580                let i3 = *input_ptr.add(start_idx + (8 * k + 3) * stride);
581                let i4 = *input_ptr.add(start_idx + (8 * k + 4) * stride);
582                let i5 = *input_ptr.add(start_idx + (8 * k + 5) * stride);
583                let i6 = *input_ptr.add(start_idx + (8 * k + 6) * stride);
584                let i7 = *input_ptr.add(start_idx + (8 * k + 7) * stride);
585
586                let a = i0;
587                let b = i1;
588                let c = i2;
589                let d = i3;
590                let e = i4;
591                let f = i5;
592                let g = i6;
593                let h = i7;
594
595                let ae_sum = a + e;
596                let ae_diff = a - e;
597                let bf_sum = b + f;
598                let bf_diff = b - f;
599                let cg_sum = c + g;
600                let cg_diff = c - g;
601                let dh_sum = d + h;
602                let dh_diff = d - h;
603
604                let w00 = ae_sum + cg_sum;
605                let w01 = ae_sum - cg_sum;
606                let w10 = ae_diff + rot_270(cg_diff);
607                let w11 = ae_diff - rot_270(cg_diff);
608                let x00 = bf_sum + dh_sum;
609                let x01 = rot_270(bf_sum) + rot_90(dh_sum);
610                let x10 = rot_45(rot_270(bf_diff) + rot_180(dh_diff));
611                let x11 = rot_45(rot_180(bf_diff) + rot_270(dh_diff));
612
613                *output_ptr.add(start_idx + k * stride) = w00 + x00;
614                *output_ptr.add(start_idx + (k + big_n / 8) * stride) = w10 + x10;
615                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = w01 + x01;
616                *output_ptr.add(start_idx + (k + 3 * big_n / 8) * stride) = w11 + x11;
617                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = w00 - x00;
618                *output_ptr.add(start_idx + (k + 5 * big_n / 8) * stride) = w10 - x10;
619                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = w01 - x01;
620                *output_ptr.add(start_idx + (k + 7 * big_n / 8) * stride) = w11 - x11;
621            }
622        }
623    }
624}
625
626struct Matrix<T> {
627    data: Vec<T>,
628    rows: usize,
629    cols: usize,
630}
631
632impl<T: Clone + Default> Matrix<T> {
633    pub fn new(rows: usize, cols: usize) -> Self {
634        let data = vec![T::default(); rows * cols];
635        Self { data, rows, cols }
636    }
637}
638
639// Reference: https://doc.rust-lang.org/std/ops/trait.Index.html
640impl<T> std::ops::Index<(usize, usize)> for Matrix<T> {
641    type Output = T;
642    fn index(&self, (r, c): (usize, usize)) -> &T {
643        &self.data[r * self.cols + c]
644    }
645}
646
647impl<T> std::ops::IndexMut<(usize, usize)> for Matrix<T> {
648    fn index_mut(&mut self, (r, c): (usize, usize)) -> &mut T {
649        &mut self.data[r * self.cols + c]
650    }
651}
652
653pub struct ExecutionPlan<T: Float + FloatConst> {
654    pub twiddles: Vec<Complex<T>>,
655    pub stages: Vec<ExecutionStage<T>>,
656    pub scratch: Vec<Complex<T>>,
657}
658
659pub struct ExecutionStage<T: Float + FloatConst> {
660    radix: usize,
661    butterfly:
662        fn(&mut [Complex<T>], &mut [Complex<T>], usize, usize, &[Complex<T>], &ExecutionStage<T>),
663    butterfly_s0: fn(&mut [Complex<T>], &mut [Complex<T>], &Matrix<Complex<T>>),
664    // TODO explain wtf this is. Link to article?
665    outer_terms: Matrix<Complex<T>>,
666}
667
668impl<T: Clone + Default + Float + FloatConst> ExecutionPlan<T> {
669    pub fn new(data: &[Complex<T>]) -> anyhow::Result<Self> {
670        let data_len = data.len();
671
672        let available_radices = [9, 7, 8, 5, 4, 3, 2];
673        let mut stages = Vec::new();
674        let mut len = data_len;
675        while len > 1 {
676            let mut got_match = false;
677            for radix in &available_radices {
678                if len % radix == 0 {
679                    let butterflies = Self::radix_to_butterfly(*radix)?;
680                    stages.push(ExecutionStage {
681                        radix: *radix,
682                        butterfly: butterflies.0,
683                        butterfly_s0: butterflies.1,
684                        outer_terms: Self::generate_radix_outer_terms(*radix),
685                    });
686                    len /= radix;
687                    got_match = true;
688                    break;
689                }
690            }
691            if !got_match {
692                anyhow::bail!("Could not factor input length {}", len);
693            }
694        }
695
696        let twiddles = precompute_twiddles::<T>(data_len);
697        let scratch: Vec<Complex<T>> = vec![Complex::new(T::zero(), T::zero()); data.len()];
698        Ok(Self {
699            twiddles,
700            stages,
701            scratch,
702        })
703    }
704
705    fn generate_radix_outer_terms(radix: usize) -> Matrix<Complex<T>> {
706        assert!(radix > 1, "radix must be greater than 1");
707        let mut data = Matrix::<Complex<T>>::new(radix, radix);
708        for row in 0..radix {
709            for col in 0..radix {
710                let k = (row * col) % radix;
711                let entry = Complex::<f64>::cis(-f64::TAU() * (k as f64) / (radix as f64));
712                data.data[row * radix + col] =
713                    Complex::<T>::new(T::from(entry.re).unwrap(), T::from(entry.im).unwrap());
714            }
715        }
716        data
717    }
718
719    fn radix_to_butterfly(
720        radix: usize,
721    ) -> anyhow::Result<(
722        fn(&mut [Complex<T>], &mut [Complex<T>], usize, usize, &[Complex<T>], &ExecutionStage<T>),
723        fn(&mut [Complex<T>], &mut [Complex<T>], &Matrix<Complex<T>>),
724    )> {
725        match radix {
726            8 => Ok((
727                fft_butterfly_radix_8_unsafe,
728                fft_butterfly_radix_8_s0_unsafe,
729            )),
730            4 => Ok((
731                fft_butterfly_radix_4_unsafe,
732                fft_butterfly_radix_4_s0_unsafe,
733            )),
734            2 => Ok((
735                fft_butterfly_radix_2_unsafe,
736                fft_butterfly_radix_2_s0_unsafe,
737            )),
738            9 => Ok((
739                fft_butterfly_radix_9_unsafe,
740                fft_butterfly_radix_9_s0_unsafe,
741            )),
742            3 => Ok((
743                fft_butterfly_radix_3_unsafe,
744                fft_butterfly_radix_3_s0_unsafe,
745            )),
746            5 => Ok((
747                fft_butterfly_radix_p_unsafe::<5, T>,
748                fft_butterfly_radix_p_s0_unsafe::<5, T>,
749            )),
750            7 => Ok((
751                fft_butterfly_radix_p_unsafe::<7, T>,
752                fft_butterfly_radix_p_s0_unsafe::<7, T>,
753            )),
754            _ => anyhow::bail!("Unsupported radix {}", radix),
755        }
756    }
757}
758
759pub fn fft<T: Float + FloatConst>(src: &mut [Complex<T>], plan: &mut ExecutionPlan<T>) {
760    let scratch = plan.scratch.as_mut_slice();
761    assert_eq!(src.len(), scratch.len());
762    assert_eq!(plan.twiddles.len(), src.len());
763
764    let n_iter = plan.stages.len();
765    if n_iter % 2 == 1 {
766        scratch.copy_from_slice(src);
767    }
768
769    let (mut input, mut output) = if n_iter % 2 == 0 {
770        (scratch, src)
771    } else {
772        (src, scratch)
773    };
774
775    let big_n = input.len();
776    let mut stride = big_n;
777    let mut big_n = 1;
778    for (stage_idx, stage) in plan.stages.iter().enumerate() {
779        stride /= stage.radix;
780        big_n *= stage.radix;
781        std::mem::swap(&mut input, &mut output);
782
783        if stage_idx == 0 {
784            (stage.butterfly_s0)(input, output, &stage.outer_terms);
785        } else {
786            (stage.butterfly)(input, output, stride, big_n, &plan.twiddles, &stage);
787        }
788    }
789}
790
791#[cfg(test)]
792mod tests {
793    use rand::RngExt;
794    use rand::SeedableRng;
795    use rand::rngs::StdRng;
796
797    use super::*;
798    use num::complex::Complex;
799    use std::time::Instant;
800
801    // Cast one float type to another, truncating if needed.
802    fn t0_to_t1<T0: Float, T1: Float>(val: T0) -> T1 {
803        return num::cast(val).unwrap();
804    }
805
806    // Casts one Complex<T> to another.
807    fn complex_to_t<T0: Float, T1: Float>(data: &[Complex<T0>]) -> Vec<Complex<T1>> {
808        data.iter()
809            .map(|c| Complex::new(t0_to_t1(c.re), t0_to_t1(c.im)))
810            .collect()
811    }
812
813    fn evaluate_results(
814        result_ref_64: &Vec<Complex<f64>>,
815        result_cur_64: &Vec<Complex<f64>>,
816        duration_dft: std::time::Duration,
817        duration_cur: std::time::Duration,
818        algo_name: &str,
819    ) {
820        assert_eq!(result_ref_64.len(), result_cur_64.len());
821
822        println!("  Algorithm:  {}", algo_name);
823        println!("    Duration:  {:?}", duration_cur);
824        println!(
825            "    Speedup:   {:?}",
826            duration_dft.as_nanos() as f32 / duration_cur.as_nanos() as f32
827        );
828
829        let mut max_err: f64 = 0.0;
830        let mut sum_err: f64 = 0.0;
831        // TODO: median, 90p, 99p
832
833        let result_64: Vec<Complex<f64>> = complex_to_t(&result_cur_64);
834
835        for i in 0..result_ref_64.len() {
836            max_err = max_err.max((result_ref_64[i] - result_64[i]).norm());
837            sum_err += (result_ref_64[i] - result_64[i]).norm();
838        }
839        println!("    Max err:   {}", max_err);
840        println!("    Avg err:   {}", sum_err / (result_ref_64.len() as f64));
841    }
842
843    fn evaluate_size(size: u32) {
844        println!("size={}", size);
845
846        let mut data: Vec<Complex<f64>> = Vec::new();
847        data.resize(size as usize, Complex::ZERO);
848        let mut rng = StdRng::seed_from_u64(42);
849        for i in 0..data.len() {
850            data[i] = Complex::new(rng.random(), rng.random());
851        }
852
853        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
854        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
855        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
856        let begin_instant = Instant::now();
857        ref_fft.process(&mut result_ref);
858        let duration_ref = begin_instant.elapsed();
859        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);
860        evaluate_results(
861            &result_ref_64,
862            &result_ref_64,
863            std::time::Duration::from_secs(0),
864            duration_ref,
865            "REF",
866        );
867
868        let mut result_rfft: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);
869        let mut rfftner = rustfft::FftPlanner::<f32>::new();
870        let rfft_fft = rfftner.plan_fft_forward(result_rfft.len());
871        let mut begin_instant = Instant::now();
872        rfft_fft.process(&mut result_rfft);
873        let duration_rfft = begin_instant.elapsed();
874        let result_rfft_64: Vec<Complex<f64>> = complex_to_t(&result_rfft);
875        evaluate_results(
876            &result_ref_64,
877            &result_rfft_64,
878            std::time::Duration::from_secs(0),
879            duration_rfft,
880            "RFFT",
881        );
882
883        let mut result_fft = complex_to_t(&data);
884        let mut plan = ExecutionPlan::<f32>::new(&result_fft).unwrap();
885        begin_instant = Instant::now();
886        fft(&mut result_fft, &mut plan);
887        let duration_fft = begin_instant.elapsed();
888        let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
889        evaluate_results(
890            &result_ref_64,
891            &result_fft_64,
892            std::time::Duration::from_secs(0),
893            duration_fft,
894            "FFT v8 (plan)",
895        );
896    }
897
898    #[test]
899    fn naive_fft_matches_ref() {
900        for exponent in 4..17 {
901            evaluate_size(2u32.pow(exponent));
902        }
903        evaluate_size(3u32.pow(3)); // 27
904        evaluate_size(3u32.pow(4)); // 81
905        evaluate_size(3u32.pow(5)); // 243
906        evaluate_size(3u32.pow(6)); // 729
907        evaluate_size(3u32.pow(7)); // 2187
908        evaluate_size(3u32.pow(8)); // 6561
909
910        evaluate_size(5u32.pow(2)); // 25
911        evaluate_size(5u32.pow(3)); // 125
912        evaluate_size(5u32.pow(4)); // 625
913        evaluate_size(5u32.pow(5)); // 3125
914
915        evaluate_size(7u32.pow(2)); // 49
916        evaluate_size(7u32.pow(3)); // 343
917        evaluate_size(7u32.pow(4)); // 2401
918    }
919}