rufft/fft/
fft_inner.rs

1use ruda_kernel::dsl as kernel_dsl;
2use std::f32::consts::PI;
3
4use ruda_kernel::dsl::prelude::*;
5use ruda_kernel::library::tensor::View;
6use ruda_kernel::library::tensor::layout::Coords1d;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9/// Whether doing RFFT or IRFFT.
10pub enum FftMode {
11    Forward,
12    Inverse,
13}
14
15impl FftMode {
16    pub fn sign(&self) -> f32 {
17        match self {
18            FftMode::Forward => -1.,
19            FftMode::Inverse => 1.,
20        }
21    }
22}
23
24#[ruda]
25/// In-place FFT of a 1D complex signal.
26/// Reorders input with bit-reversal and applies butterfly stages
27pub(crate) fn fft_inner_compute<F: Float>(
28    spectrum_re: &mut View<F, Coords1d, ReadWrite>,
29    spectrum_im: &mut View<F, Coords1d, ReadWrite>,
30    #[comptime] fft_mode: FftMode,
31) {
32    let num_samples = spectrum_re.shape();
33
34    bit_reverse_permutation(spectrum_re, spectrum_im, num_samples);
35
36    fft_butterfly_stages(spectrum_re, spectrum_im, fft_mode);
37}
38
39#[ruda]
40/// In-place bit-reversal permutation.
41///
42/// Reorders elements so index `i` maps to the index formed by
43/// reversing the `log2(n)` bits of `i`.
44fn bit_reverse_permutation<F: Float>(
45    view_re: &mut View<F, Coords1d, ReadWrite>,
46    view_im: &mut View<F, Coords1d, ReadWrite>,
47    n: usize,
48) {
49    let mut j = 0;
50    for i in 0..n {
51        if i < j {
52            swap(view_re, i, j);
53            swap(view_im, i, j);
54        }
55        let mut m = n >> 1;
56        while m >= 1 && j >= m {
57            j -= m;
58            m >>= 1;
59        }
60        j += m;
61    }
62}
63
64#[ruda]
65/// Swap two elements of a 1D array.
66fn swap<F: Float>(view_1d: &mut View<F, Coords1d, ReadWrite>, i: usize, j: usize) {
67    let tmp = view_1d[i];
68    view_1d[i] = view_1d[j];
69    view_1d[j] = tmp;
70}
71
72#[ruda]
73/// Iterative radix-2 FFT butterfly computation.
74/// Combines pairs of elements using twiddle factors to compute higher-level FFT outputs.
75fn fft_butterfly_stages<F: Float>(
76    spectrum_re: &mut View<F, Coords1d, ReadWrite>,
77    spectrum_im: &mut View<F, Coords1d, ReadWrite>,
78    #[comptime] fft_mode: FftMode,
79) {
80    let n = spectrum_re.shape();
81    let mut m = 2;
82
83    while m <= n {
84        let half_m = m >> 1;
85
86        // twiddle base: exp(-2πi / m)
87        let theta = F::new(fft_mode.sign() * 2.0 * PI) / F::cast_from(m);
88
89        let wm_re = theta.cos();
90        let wm_in = theta.sin();
91
92        let mut k = 0;
93        while k < n {
94            let mut w_re = F::new(1.0);
95            let mut w_im = F::new(0.0);
96
97            let mut j = 0;
98            while j < half_m {
99                let i0 = k + j;
100                let i1 = i0 + half_m;
101
102                let a = (spectrum_re[i0], spectrum_im[i0]);
103                let b = (spectrum_re[i1], spectrum_im[i1]);
104
105                let t = complex_mul::<F>((w_re, w_im), b);
106                let out0 = complex_add::<F>(a, t);
107                let out1 = complex_sub::<F>(a, t);
108
109                spectrum_re[i0] = out0.0;
110                spectrum_im[i0] = out0.1;
111                spectrum_re[i1] = out1.0;
112                spectrum_im[i1] = out1.1;
113
114                let new_w = complex_mul::<F>((w_re, w_im), (wm_re, wm_in));
115                w_re = new_w.0;
116                w_im = new_w.1;
117
118                j += 1;
119            }
120
121            k += m;
122        }
123
124        m <<= 1;
125    }
126}
127
128#[ruda]
129/// Addition on a complex number encoded as a pair of floats
130fn complex_add<F: Float>(a: (F, F), b: (F, F)) -> (F, F) {
131    (a.0 + b.0, a.1 + b.1)
132}
133
134#[ruda]
135/// Subtraction on a complex number encoded as a pair of floats
136fn complex_sub<F: Float>(a: (F, F), b: (F, F)) -> (F, F) {
137    (a.0 - b.0, a.1 - b.1)
138}
139
140#[ruda]
141/// Multiplication on a complex number encoded as a pair of floats
142fn complex_mul<F: Float>(a: (F, F), b: (F, F)) -> (F, F) {
143    (a.0 * b.0 - a.1 * b.1, a.0 * b.1 + a.1 * b.0)
144}
rufft/fft/fft_inner.rs

rufft/fft/
fft_inner.rs
