rufft/fft/
exact.rs

1//! Exact-length F32 real transforms. Non-radix-2 lengths use Bluestein's
2//! convolution, never an N-point transform silently replaced by a larger DFT.
3//!
4//! This plan owns chirps and reusable device scratch. It must be used on the
5//! execution queue on which it was created. It does not perform CPU arithmetic
6//! on signal data or implicitly synchronize the device.
7use ruda_kernel::dsl as kernel_dsl;
8use ruda_kernel::dsl::prelude::*;
9use ruda_kernel::dsl::backtrace::BackTrace;
10use ruda_kernel::library::tensor::TensorHandle;
11use super::{FftMode, rfft_launch_padded, irfft_launch_padded};
12use super::cfft::{CfftBindings, MAX_SHARED_N_FFT, cfft_launch_any_size,
13    cfft_launch_with_scratch, cfft_inverse_product_with_scratch};
14
15mod kernels;
16use kernels::*;
17
18fn invalid(reason: impl Into<String>) -> LaunchError {
19    LaunchError::Unknown { reason: reason.into(), backtrace: BackTrace::capture() }
20}
21
22/// Bounded implementation limit, not a device memory availability guarantee.
23const MAX_CONVOLUTION: usize = MAX_SHARED_N_FFT * MAX_SHARED_N_FFT;
24
25/// Bluestein workspace length. Power-of-two real FFTs use the existing packed
26/// real path, which can support up to twice this complex FFT limit.
27pub fn exact_convolution_len(n: usize) -> Result<Option<usize>, LaunchError> {
28    if n == 0 { return Err(invalid("FFT length must be positive")); }
29    if n.is_power_of_two() {
30        if n > MAX_CONVOLUTION * 2 { return Err(invalid("radix-2 FFT exceeds the implemented four-step limit")); }
31        return Ok(None);
32    }
33    let m = n.checked_mul(2).and_then(|v| v.checked_sub(1))
34        .and_then(usize::checked_next_power_of_two)
35        .filter(|&v| v <= MAX_CONVOLUTION)
36        .ok_or_else(|| invalid("Bluestein convolution exceeds the implemented four-step limit"))?;
37    Ok(Some(m))
38}
39
40fn alloc<R: Runtime>(client: &ComputeClient<R>, rows: usize, cols: usize)
41    -> Result<TensorHandle<R>, LaunchError>
42{
43    let elements = rows.checked_mul(cols).filter(|&v| v <= u32::MAX as usize)
44        .ok_or_else(|| invalid("FFT workspace exceeds 32-bit indexing"))?;
45    let bytes = elements.checked_mul(4).ok_or_else(|| invalid("FFT workspace byte-size overflow"))?;
46    Ok(TensorHandle::new_contiguous(vec![rows, cols], client.empty(bytes),
47        f32::as_type_native_unchecked().storage_type()))
48}
49
50struct Tables<R: Runtime> {
51    m: usize,
52    chirp_re: TensorHandle<R>,
53    chirp_im: TensorHandle<R>,
54    spectrum_re: TensorHandle<R>,
55    spectrum_im: TensorHandle<R>,
56}
57struct Workspace<R: Runtime> {
58    rows: usize,
59    re: TensorHandle<R>,
60    im: TensorHandle<R>,
61    four_step: Option<(TensorHandle<R>, TensorHandle<R>)>,
62}
63
64/// Reusable exact-length F32 FFT plan. The direction and N are fixed; batches
65/// may change, in which case the batched workspace is replaced, not grown
66/// without bound. Caller-owned output bindings must be non-overlapping.
67/// Raw bindings carry no dtype/device tag: pass F32 storage from this client.
68pub struct RealFftPlan<R: Runtime> {
69    client: ComputeClient<R>,
70    n: usize,
71    mode: FftMode,
72    tables: Option<Tables<R>>,
73    workspace: Option<Workspace<R>>,
74    fuse_spectrum: bool,
75}
76impl<R: Runtime> RealFftPlan<R> {
77    pub fn new(client: ComputeClient<R>, n: usize, mode: FftMode) -> Result<Self, LaunchError> {
78        let m = exact_convolution_len(n)?;
79        // Pin an implicit client to its current queue without introducing a
80        // new stream or changing any allocation's ownership.
81        let client = client.fixed_execution_queue();
82        let tables = if let Some(m) = m {
83            let chirp_re = alloc(&client, 1, n)?;
84            let chirp_im = alloc(&client, 1, n)?;
85            let spectrum_re = alloc(&client, 1, m)?;
86            let spectrum_im = alloc(&client, 1, m)?;
87            let block = RudaDim::new_1d(256);
88            let grid = ruda_kernel::dsl::calculate_ruda_count_elemwise(&client, m, block);
89            make_chirp::launch::<R>(&client, grid, block,
90                chirp_re.clone().into_arg(), chirp_im.clone().into_arg(),
91                spectrum_re.clone().into_arg(), spectrum_im.clone().into_arg(), n, m, mode);
92            cfft_launch_any_size(&client, CfftBindings {
93                input_re: spectrum_re.clone().binding(), input_im: spectrum_im.clone().binding(),
94                output_re: spectrum_re.clone().binding(), output_im: spectrum_im.clone().binding(),
95            }, 1, f32::as_type_native_unchecked().storage_type(), FftMode::Forward)?;
96            Some(Tables { m, chirp_re, chirp_im, spectrum_re, spectrum_im })
97        } else { None };
98        Ok(Self { client, n, mode, tables, workspace: None, fuse_spectrum: false })
99    }
100
101    /// Opt into the candidate v22 fused spectrum/inverse-FFT input stage.
102    /// False keeps the separately launched product validated on the v21 T4 run.
103    /// No effect on N=1 or power-of-two transforms; no additional GPU allocation.
104    /// Call between submissions on this plan's fixed execution queue.
105    pub fn set_spectrum_fusion(&mut self, enabled: bool) { self.fuse_spectrum = enabled; }
106
107    /// Whether this plan will use the fused product (only Bluestein lengths).
108    pub fn spectrum_fusion_enabled(&self) -> bool {
109        self.fuse_spectrum && self.tables.is_some()
110    }
111
112    pub fn len(&self) -> usize { self.n }
113    pub fn is_empty(&self) -> bool { false }
114    pub fn convolution_len(&self) -> Option<usize> { self.tables.as_ref().map(|t| t.m) }
115    /// Bytes retained by this plan, excluding output tensors, driver metadata,
116    /// compiled kernels and temporary initialization scratch.
117    pub fn retained_bytes(&self) -> usize {
118        let tables = self.tables.as_ref().map_or(0, |t| 8 * (self.n + t.m));
119        tables + self.workspace.as_ref().map_or(0, |w| {
120            let m = self.tables.as_ref().unwrap().m;
121            w.rows * m * if w.four_step.is_some() { 16 } else { 8 }
122        })
123    }
124    /// Release scratch; chirp tables remain reusable. Runtime stream ordering
125    /// handles deferred release. This is not a device-wide synchronization.
126    pub fn clear_workspace(&mut self) { self.workspace = None; }
127
128    fn prepare_workspace(&mut self, rows: usize) -> Result<(), LaunchError> {
129        if self.workspace.as_ref().is_some_and(|w| w.rows == rows) { return Ok(()); }
130        let m = self.tables.as_ref().unwrap().m;
131        // Drop the old workspace before reserving a differently sized one.
132        // The allocator still honors pending users on its execution queue.
133        self.workspace = None;
134        let re = alloc(&self.client, rows, m)?;
135        let im = alloc(&self.client, rows, m)?;
136        let four_step = if m > MAX_SHARED_N_FFT {
137            Some((alloc(&self.client, rows, m)?, alloc(&self.client, rows, m)?))
138        } else { None };
139        self.workspace = Some(Workspace { rows, re, im, four_step });
140        Ok(())
141    }
142
143    fn convolution(&self) -> Result<(), LaunchError> {
144        let t = self.tables.as_ref().unwrap();
145        let w = self.workspace.as_ref().unwrap();
146        let dtype = f32::as_type_native_unchecked().storage_type();
147        let bindings = || CfftBindings {
148            input_re: w.re.clone().binding(), input_im: w.im.clone().binding(),
149            output_re: w.re.clone().binding(), output_im: w.im.clone().binding(),
150        };
151        cfft_launch_with_scratch(&self.client, bindings(), 1, dtype, FftMode::Forward,
152            w.four_step.clone())?;
153        if self.fuse_spectrum {
154            return cfft_inverse_product_with_scratch(&self.client, bindings(), 1, dtype,
155                (t.spectrum_re.clone().binding(), t.spectrum_im.clone().binding()),
156                w.four_step.clone());
157        }
158        let total = w.rows * t.m;
159        let block = RudaDim::new_1d(256);
160        let grid = ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, total, block);
161        multiply_spectrum::launch::<R>(&self.client, grid, block,
162            w.re.clone().into_arg(), w.im.clone().into_arg(),
163            t.spectrum_re.clone().into_arg(), t.spectrum_im.clone().into_arg(),
164            total as u32, t.m);
165        cfft_launch_with_scratch(&self.client, bindings(), 1, dtype, FftMode::Inverse,
166            w.four_step.clone())
167    }
168
169    /// Unnormalized N-point RFFT. Samples at `used..N` are zero, without
170    /// constructing a padded input tensor. `used` may be zero.
171    pub fn forward(&mut self, signal: TensorBinding<R>, real: TensorBinding<R>,
172        imag: TensorBinding<R>, dim: usize, used: usize) -> Result<(), LaunchError>
173    {
174        if self.mode != FftMode::Forward { return Err(invalid("forward called on inverse FFT plan")); }
175        let rows = validate(&signal, &real, &imag, dim, self.n, used, false)?;
176        if rows == 0 { return Ok(()); }
177        let block = RudaDim::new_1d(256);
178        if self.n == 1 {
179            scalar_forward::launch::<R>(&self.client,
180                ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, rows, block), block,
181                signal.into_tensor_arg(), real.into_tensor_arg(), imag.into_tensor_arg(),
182                rows as u32, used as u32, dim);
183            return Ok(());
184        }
185        if self.tables.is_none() {
186            return rfft_launch_padded(&self.client, signal, real, imag, dim, used,
187                f32::as_type_native_unchecked().storage_type());
188        }
189        self.prepare_workspace(rows)?;
190        let t = self.tables.as_ref().unwrap();
191        let w = self.workspace.as_ref().unwrap();
192        let total = rows * t.m;
193        prepare_forward::launch::<R>(&self.client,
194            ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, total, block), block,
195            signal.into_tensor_arg(), t.chirp_re.clone().into_arg(), t.chirp_im.clone().into_arg(),
196            w.re.clone().into_arg(), w.im.clone().into_arg(), total as u32, used as u32,
197            self.n, t.m, dim);
198        self.convolution()?;
199        let bins = self.n / 2 + 1;
200        finish_forward::launch::<R>(&self.client,
201            ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, rows * bins, block), block,
202            w.re.clone().into_arg(), w.im.clone().into_arg(),
203            t.chirp_re.clone().into_arg(), t.chirp_im.clone().into_arg(),
204            real.into_tensor_arg(), imag.into_tensor_arg(), (rows * bins) as u32,
205            self.n, t.m, dim);
206        Ok(())
207    }
208
209    /// N-point IRFFT normalized by 1/N. DC, and Nyquist for even N, are
210    /// real-only. Bins at `used..N/2+1` are treated as zero.
211    pub fn inverse(&mut self, real: TensorBinding<R>, imag: TensorBinding<R>,
212        signal: TensorBinding<R>, dim: usize, used: usize) -> Result<(), LaunchError>
213    {
214        if self.mode != FftMode::Inverse { return Err(invalid("inverse called on forward FFT plan")); }
215        let rows = validate(&signal, &real, &imag, dim, self.n, used, true)?;
216        if rows == 0 { return Ok(()); }
217        let block = RudaDim::new_1d(256);
218        if self.n == 1 {
219            scalar_inverse::launch::<R>(&self.client,
220                ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, rows, block), block,
221                real.into_tensor_arg(), signal.into_tensor_arg(), rows as u32, used as u32, dim);
222            return Ok(());
223        }
224        if self.tables.is_none() {
225            return irfft_launch_padded(&self.client, real, imag, signal, dim, used,
226                f32::as_type_native_unchecked().storage_type());
227        }
228        self.prepare_workspace(rows)?;
229        let t = self.tables.as_ref().unwrap();
230        let w = self.workspace.as_ref().unwrap();
231        let total = rows * t.m;
232        prepare_inverse::launch::<R>(&self.client,
233            ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, total, block), block,
234            real.into_tensor_arg(), imag.into_tensor_arg(),
235            t.chirp_re.clone().into_arg(), t.chirp_im.clone().into_arg(),
236            w.re.clone().into_arg(), w.im.clone().into_arg(), total as u32, used as u32,
237            self.n, t.m, dim);
238        self.convolution()?;
239        finish_inverse::launch::<R>(&self.client,
240            ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, rows * self.n, block), block,
241            w.re.clone().into_arg(), w.im.clone().into_arg(),
242            t.chirp_re.clone().into_arg(), t.chirp_im.clone().into_arg(),
243            signal.into_tensor_arg(), (rows * self.n) as u32, self.n, t.m, dim);
244        Ok(())
245    }
246}
247
248fn check_binding<R: Runtime>(binding: &TensorBinding<R>) -> Result<(), LaunchError> {
249    if binding.shape.len() != binding.strides.len() || binding.shape.is_empty() {
250        return Err(invalid("FFT needs non-scalar, well-formed tensor metadata"));
251    }
252    if binding.shape.contains(&0) { return Ok(()); }
253    let last = binding.shape.iter().zip(binding.strides.iter()).try_fold(0usize,
254        |sum, (&size, &stride)| (size - 1).checked_mul(stride).and_then(|v| sum.checked_add(v)))
255        .filter(|&v| v < u32::MAX as usize)
256        .and_then(|v| v.checked_add(1)).and_then(|v| v.checked_mul(4))
257        .ok_or_else(|| invalid("FFT tensor span overflows F32/U32 addressing"))?;
258    let available = binding.handle.size.checked_sub(binding.handle.offset_start.unwrap_or(0))
259        .and_then(|v| v.checked_sub(binding.handle.offset_end.unwrap_or(0)))
260        .ok_or_else(|| invalid("invalid FFT buffer offsets"))?;
261    if last as u64 > available { return Err(invalid("FFT tensor strides exceed buffer size")); }
262    Ok(())
263}
264
265fn validate<R: Runtime>(signal: &TensorBinding<R>, real: &TensorBinding<R>,
266    imag: &TensorBinding<R>, dim: usize, n: usize, used: usize, inverse: bool)
267    -> Result<usize, LaunchError>
268{
269    for binding in [signal, real, imag] { check_binding(binding)?; }
270    if dim >= signal.shape.len() || signal.shape.len() != real.shape.len()
271        || real.shape != imag.shape {
272        return Err(invalid("invalid FFT dimension, rank or real/imag shapes"));
273    }
274    let bins = n / 2 + 1;
275    if (!inverse && real.shape[dim] != bins) || (inverse && signal.shape[dim] != n)
276        || (!inverse && (used > signal.shape[dim] || used > n))
277        || (inverse && (used == 0 || used > real.shape[dim] || used > bins)) {
278        return Err(invalid("FFT input/output length does not match plan"));
279    }
280    let mut rows = 1usize;
281    for (axis, (&a, &b)) in signal.shape.iter().zip(real.shape.iter()).enumerate() {
282        if axis != dim {
283            if a != b { return Err(invalid("FFT batch shapes must match")); }
284            rows = rows.checked_mul(a).ok_or_else(|| invalid("FFT batch size overflow"))?;
285        }
286    }
287    rows.checked_mul(n.max(bins)).filter(|&v| v <= u32::MAX as usize)
288        .ok_or_else(|| invalid("FFT output exceeds 32-bit indexing"))?;
289    // Outputs may be pitched or permuted, but must not have overlapping
290    // elements. This sufficient test accepts dense permutations and padding.
291    for out in if inverse { vec![signal] } else { vec![real, imag] } {
292        let mut axes: Vec<_> = out.shape.iter().zip(out.strides.iter())
293            .filter(|(size, _)| **size > 1).map(|(&size, &stride)| (stride, size)).collect();
294        axes.sort_unstable();
295        let mut span = 1usize;
296        for (stride, size) in axes {
297            if stride < span { return Err(invalid("FFT output layout has overlapping elements")); }
298            span = (size - 1).checked_mul(stride).and_then(|v| v.checked_add(span))
299                .ok_or_else(|| invalid("FFT output layout overflow"))?;
300        }
301    }
302    Ok(rows)
303}
304
305#[cfg(test)]
306mod tests {
307    use super::*;
308    #[test]
309    fn exact_plan_length_limits() {
310        for n in [1, 2, 4, 1024, 8192] { assert_eq!(exact_convolution_len(n).unwrap(), None); }
311        for (n, m) in [(3, 8), (6, 16), (7, 16), (1009, 2048), (4097, 16384)] {
312            assert_eq!(exact_convolution_len(n).unwrap(), Some(m));
313        }
314        assert!(exact_convolution_len(0).is_err());
315        assert!(exact_convolution_len(usize::MAX).is_err());
316        assert!(exact_convolution_len(MAX_CONVOLUTION).is_ok());
317        assert!(exact_convolution_len(MAX_CONVOLUTION - 1).is_err());
318    }
319}
rufft/fft/exact.rs

rufft/fft/
exact.rs
