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};
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}
75impl<R: Runtime> RealFftPlan<R> {
76    pub fn new(client: ComputeClient<R>, n: usize, mode: FftMode) -> Result<Self, LaunchError> {
77        let m = exact_convolution_len(n)?;
78        // Pin an implicit client to its current queue without introducing a
79        // new stream or changing any allocation's ownership.
80        let client = client.fixed_execution_queue();
81        let tables = if let Some(m) = m {
82            let chirp_re = alloc(&client, 1, n)?;
83            let chirp_im = alloc(&client, 1, n)?;
84            let spectrum_re = alloc(&client, 1, m)?;
85            let spectrum_im = alloc(&client, 1, m)?;
86            let block = RudaDim::new_1d(256);
87            let grid = ruda_kernel::dsl::calculate_ruda_count_elemwise(&client, m, block);
88            make_chirp::launch::<R>(&client, grid, block,
89                chirp_re.clone().into_arg(), chirp_im.clone().into_arg(),
90                spectrum_re.clone().into_arg(), spectrum_im.clone().into_arg(), n, m, mode);
91            cfft_launch_any_size(&client, CfftBindings {
92                input_re: spectrum_re.clone().binding(), input_im: spectrum_im.clone().binding(),
93                output_re: spectrum_re.clone().binding(), output_im: spectrum_im.clone().binding(),
94            }, 1, f32::as_type_native_unchecked().storage_type(), FftMode::Forward)?;
95            Some(Tables { m, chirp_re, chirp_im, spectrum_re, spectrum_im })
96        } else { None };
97        Ok(Self { client, n, mode, tables, workspace: None })
98    }
99
100    pub fn len(&self) -> usize { self.n }
101    pub fn is_empty(&self) -> bool { false }
102    pub fn convolution_len(&self) -> Option<usize> { self.tables.as_ref().map(|t| t.m) }
103    /// Bytes retained by this plan, excluding output tensors, driver metadata,
104    /// compiled kernels and temporary initialization scratch.
105    pub fn retained_bytes(&self) -> usize {
106        let tables = self.tables.as_ref().map_or(0, |t| 8 * (self.n + t.m));
107        tables + self.workspace.as_ref().map_or(0, |w| {
108            let m = self.tables.as_ref().unwrap().m;
109            w.rows * m * if w.four_step.is_some() { 16 } else { 8 }
110        })
111    }
112    /// Release scratch; chirp tables remain reusable. Runtime stream ordering
113    /// handles deferred release. This is not a device-wide synchronization.
114    pub fn clear_workspace(&mut self) { self.workspace = None; }
115
116    fn prepare_workspace(&mut self, rows: usize) -> Result<(), LaunchError> {
117        if self.workspace.as_ref().is_some_and(|w| w.rows == rows) { return Ok(()); }
118        let m = self.tables.as_ref().unwrap().m;
119        // Drop the old workspace before reserving a differently sized one.
120        // The allocator still honors pending users on its execution queue.
121        self.workspace = None;
122        let re = alloc(&self.client, rows, m)?;
123        let im = alloc(&self.client, rows, m)?;
124        let four_step = if m > MAX_SHARED_N_FFT {
125            Some((alloc(&self.client, rows, m)?, alloc(&self.client, rows, m)?))
126        } else { None };
127        self.workspace = Some(Workspace { rows, re, im, four_step });
128        Ok(())
129    }
130
131    fn convolution(&self) -> Result<(), LaunchError> {
132        let t = self.tables.as_ref().unwrap();
133        let w = self.workspace.as_ref().unwrap();
134        let dtype = f32::as_type_native_unchecked().storage_type();
135        let bindings = || CfftBindings {
136            input_re: w.re.clone().binding(), input_im: w.im.clone().binding(),
137            output_re: w.re.clone().binding(), output_im: w.im.clone().binding(),
138        };
139        cfft_launch_with_scratch(&self.client, bindings(), 1, dtype, FftMode::Forward,
140            w.four_step.clone())?;
141        let total = w.rows * t.m;
142        let block = RudaDim::new_1d(256);
143        let grid = ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, total, block);
144        multiply_spectrum::launch::<R>(&self.client, grid, block,
145            w.re.clone().into_arg(), w.im.clone().into_arg(),
146            t.spectrum_re.clone().into_arg(), t.spectrum_im.clone().into_arg(),
147            total as u32, t.m);
148        cfft_launch_with_scratch(&self.client, bindings(), 1, dtype, FftMode::Inverse,
149            w.four_step.clone())
150    }
151
152    /// Unnormalized N-point RFFT. Samples at `used..N` are zero, without
153    /// constructing a padded input tensor. `used` may be zero.
154    pub fn forward(&mut self, signal: TensorBinding<R>, real: TensorBinding<R>,
155        imag: TensorBinding<R>, dim: usize, used: usize) -> Result<(), LaunchError>
156    {
157        if self.mode != FftMode::Forward { return Err(invalid("forward called on inverse FFT plan")); }
158        let rows = validate(&signal, &real, &imag, dim, self.n, used, false)?;
159        if rows == 0 { return Ok(()); }
160        let block = RudaDim::new_1d(256);
161        if self.n == 1 {
162            scalar_forward::launch::<R>(&self.client,
163                ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, rows, block), block,
164                signal.into_tensor_arg(), real.into_tensor_arg(), imag.into_tensor_arg(),
165                rows as u32, used as u32, dim);
166            return Ok(());
167        }
168        if self.tables.is_none() {
169            return rfft_launch_padded(&self.client, signal, real, imag, dim, used,
170                f32::as_type_native_unchecked().storage_type());
171        }
172        self.prepare_workspace(rows)?;
173        let t = self.tables.as_ref().unwrap();
174        let w = self.workspace.as_ref().unwrap();
175        let total = rows * t.m;
176        prepare_forward::launch::<R>(&self.client,
177            ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, total, block), block,
178            signal.into_tensor_arg(), t.chirp_re.clone().into_arg(), t.chirp_im.clone().into_arg(),
179            w.re.clone().into_arg(), w.im.clone().into_arg(), total as u32, used as u32,
180            self.n, t.m, dim);
181        self.convolution()?;
182        let bins = self.n / 2 + 1;
183        finish_forward::launch::<R>(&self.client,
184            ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, rows * bins, block), block,
185            w.re.clone().into_arg(), w.im.clone().into_arg(),
186            t.chirp_re.clone().into_arg(), t.chirp_im.clone().into_arg(),
187            real.into_tensor_arg(), imag.into_tensor_arg(), (rows * bins) as u32,
188            self.n, t.m, dim);
189        Ok(())
190    }
191
192    /// N-point IRFFT normalized by 1/N. DC, and Nyquist for even N, are
193    /// real-only. Bins at `used..N/2+1` are treated as zero.
194    pub fn inverse(&mut self, real: TensorBinding<R>, imag: TensorBinding<R>,
195        signal: TensorBinding<R>, dim: usize, used: usize) -> Result<(), LaunchError>
196    {
197        if self.mode != FftMode::Inverse { return Err(invalid("inverse called on forward FFT plan")); }
198        let rows = validate(&signal, &real, &imag, dim, self.n, used, true)?;
199        if rows == 0 { return Ok(()); }
200        let block = RudaDim::new_1d(256);
201        if self.n == 1 {
202            scalar_inverse::launch::<R>(&self.client,
203                ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, rows, block), block,
204                real.into_tensor_arg(), signal.into_tensor_arg(), rows as u32, used as u32, dim);
205            return Ok(());
206        }
207        if self.tables.is_none() {
208            return irfft_launch_padded(&self.client, real, imag, signal, dim, used,
209                f32::as_type_native_unchecked().storage_type());
210        }
211        self.prepare_workspace(rows)?;
212        let t = self.tables.as_ref().unwrap();
213        let w = self.workspace.as_ref().unwrap();
214        let total = rows * t.m;
215        prepare_inverse::launch::<R>(&self.client,
216            ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, total, block), block,
217            real.into_tensor_arg(), imag.into_tensor_arg(),
218            t.chirp_re.clone().into_arg(), t.chirp_im.clone().into_arg(),
219            w.re.clone().into_arg(), w.im.clone().into_arg(), total as u32, used as u32,
220            self.n, t.m, dim);
221        self.convolution()?;
222        finish_inverse::launch::<R>(&self.client,
223            ruda_kernel::dsl::calculate_ruda_count_elemwise(&self.client, rows * self.n, block), block,
224            w.re.clone().into_arg(), w.im.clone().into_arg(),
225            t.chirp_re.clone().into_arg(), t.chirp_im.clone().into_arg(),
226            signal.into_tensor_arg(), (rows * self.n) as u32, self.n, t.m, dim);
227        Ok(())
228    }
229}
230
231fn check_binding<R: Runtime>(binding: &TensorBinding<R>) -> Result<(), LaunchError> {
232    if binding.shape.len() != binding.strides.len() || binding.shape.is_empty() {
233        return Err(invalid("FFT needs non-scalar, well-formed tensor metadata"));
234    }
235    if binding.shape.contains(&0) { return Ok(()); }
236    let last = binding.shape.iter().zip(binding.strides.iter()).try_fold(0usize,
237        |sum, (&size, &stride)| (size - 1).checked_mul(stride).and_then(|v| sum.checked_add(v)))
238        .filter(|&v| v < u32::MAX as usize)
239        .and_then(|v| v.checked_add(1)).and_then(|v| v.checked_mul(4))
240        .ok_or_else(|| invalid("FFT tensor span overflows F32/U32 addressing"))?;
241    let available = binding.handle.size.checked_sub(binding.handle.offset_start.unwrap_or(0))
242        .and_then(|v| v.checked_sub(binding.handle.offset_end.unwrap_or(0)))
243        .ok_or_else(|| invalid("invalid FFT buffer offsets"))?;
244    if last as u64 > available { return Err(invalid("FFT tensor strides exceed buffer size")); }
245    Ok(())
246}
247
248fn validate<R: Runtime>(signal: &TensorBinding<R>, real: &TensorBinding<R>,
249    imag: &TensorBinding<R>, dim: usize, n: usize, used: usize, inverse: bool)
250    -> Result<usize, LaunchError>
251{
252    for binding in [signal, real, imag] { check_binding(binding)?; }
253    if dim >= signal.shape.len() || signal.shape.len() != real.shape.len()
254        || real.shape != imag.shape {
255        return Err(invalid("invalid FFT dimension, rank or real/imag shapes"));
256    }
257    let bins = n / 2 + 1;
258    if (!inverse && real.shape[dim] != bins) || (inverse && signal.shape[dim] != n)
259        || (!inverse && (used > signal.shape[dim] || used > n))
260        || (inverse && (used == 0 || used > real.shape[dim] || used > bins)) {
261        return Err(invalid("FFT input/output length does not match plan"));
262    }
263    let mut rows = 1usize;
264    for (axis, (&a, &b)) in signal.shape.iter().zip(real.shape.iter()).enumerate() {
265        if axis != dim {
266            if a != b { return Err(invalid("FFT batch shapes must match")); }
267            rows = rows.checked_mul(a).ok_or_else(|| invalid("FFT batch size overflow"))?;
268        }
269    }
270    rows.checked_mul(n.max(bins)).filter(|&v| v <= u32::MAX as usize)
271        .ok_or_else(|| invalid("FFT output exceeds 32-bit indexing"))?;
272    // Outputs may be pitched or permuted, but must not have overlapping
273    // elements. This sufficient test accepts dense permutations and padding.
274    for out in if inverse { vec![signal] } else { vec![real, imag] } {
275        let mut axes: Vec<_> = out.shape.iter().zip(out.strides.iter())
276            .filter(|(size, _)| **size > 1).map(|(&size, &stride)| (stride, size)).collect();
277        axes.sort_unstable();
278        let mut span = 1usize;
279        for (stride, size) in axes {
280            if stride < span { return Err(invalid("FFT output layout has overlapping elements")); }
281            span = (size - 1).checked_mul(stride).and_then(|v| v.checked_add(span))
282                .ok_or_else(|| invalid("FFT output layout overflow"))?;
283        }
284    }
285    Ok(rows)
286}
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291    #[test]
292    fn exact_plan_length_limits() {
293        for n in [1, 2, 4, 1024, 8192] { assert_eq!(exact_convolution_len(n).unwrap(), None); }
294        for (n, m) in [(3, 8), (6, 16), (7, 16), (1009, 2048), (4097, 16384)] {
295            assert_eq!(exact_convolution_len(n).unwrap(), Some(m));
296        }
297        assert!(exact_convolution_len(0).is_err());
298        assert!(exact_convolution_len(usize::MAX).is_err());
299        assert!(exact_convolution_len(MAX_CONVOLUTION).is_ok());
300        assert!(exact_convolution_len(MAX_CONVOLUTION - 1).is_err());
301    }
302}
rufft/fft/exact.rs

rufft/fft/
exact.rs
