ruda-fft 0.2.0

Fast Fourier transforms for Ruda.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

use ruda_core::tensor::{DType, Shape, TensorMetadata};
use ruda_kernel::dsl::prelude::*;
use ruda_kernel::tensor::{RudaTensor, allocation::empty_device_dtype, reshape::reshape};
use crate::{FftMode, RealFftPlan};

fn tune_shared_fft<R: Runtime>(inputs: Vec<RudaTensor<R>>, dim: usize, length: usize, used: usize,
    inverse: bool) -> Option<Vec<RudaTensor<R>>> {
    if !ruda_kernel::tensor::tuning::is_enabled() || !length.is_power_of_two() || !(2..=4096).contains(&length) {
        return None;
    }
    let hardware = &inputs[0].client.properties().hardware;
    let maximum = ((length / 2).clamp(1, 256) as u32).min(hardware.max_ruda_dim.0).min(hardware.max_units_per_ruda);
    let mut candidates = vec![("original_shared_schedule", 0)];
    for (name, units) in [("units_1", 1), ("units_8", 8), ("units_32", 32), ("units_64", 64), ("units_128", 128)] {
        if units <= maximum { candidates.push((name, units)); }
    }
    ruda_kernel::tensor::tuning::execute_variants(inputs, if inverse { "fft_irfft_shared" } else { "fft_rfft_shared" },
        format!("axis={dim};length={length};used={used};normalized_inverse={inverse}"), candidates, move |values, units| {
            let input = &values[0];
            let mut shape = input.shape();
            shape[dim] = if inverse { length } else { length / 2 + 1 };
            let output = empty_device_dtype(input.client.clone(), input.device.clone(), shape.clone(), DType::F32);
            let dtype = f32::as_type_native_unchecked().storage_type();
            if inverse {
                crate::fft::irfft_launch_padded_with_units(&input.client, input.clone().binding(), values[1].clone().binding(),
                    output.clone().binding(), dim, used, dtype, units).map_err(|error| format!("{error:?}"))?;
                Ok(vec![output])
            } else {
                let imag = empty_device_dtype(input.client.clone(), input.device.clone(), shape, DType::F32);
                crate::fft::rfft_launch_padded_with_units(&input.client, input.clone().binding(), output.clone().binding(),
                    imag.clone().binding(), dim, used, dtype, units).map_err(|error| format!("{error:?}"))?;
                Ok(vec![output, imag])
            }
        }).unwrap_or_else(|error| panic!("FFT autotune failed without replay: {error}"))
}

/// Legacy compatibility API: the requested signal length is rounded up to a
/// power of two. Use `rfft_exact` for an actual N-point DFT. Virtual padding
/// avoids allocating and copying a padded signal.
pub fn rfft<R: Runtime>(signal: RudaTensor<R>, dim: usize, n: Option<usize>)
    -> (RudaTensor<R>, RudaTensor<R>)
{
    forward(signal, dim, n, false)
}

/// Exact N-point real FFT, returning floor(N/2)+1 F32 frequency bins. A supplied
/// N truncates or virtually zero-pads the input, but does not change the DFT's
/// length. Non-power-of-two N uses the device Bluestein path.
pub fn rfft_exact<R: Runtime>(signal: RudaTensor<R>, dim: usize, n: Option<usize>)
    -> (RudaTensor<R>, RudaTensor<R>)
{
    forward(signal, dim, n, true)
}

fn forward<R: Runtime>(signal: RudaTensor<R>, dim: usize, n: Option<usize>, exact: bool)
    -> (RudaTensor<R>, RudaTensor<R>)
{
    assert_eq!(signal.dtype, DType::F32, "ruFFT device kernels currently require F32 storage");
    let shape = signal.shape();
    assert!(dim < shape.len(), "rfft: dimension out of bounds");
    let requested = n.unwrap_or(shape[dim]);
    assert!(requested > 0, "rfft: transform length must be positive");
    let length = if exact { requested } else {
        requested.checked_next_power_of_two().expect("rfft length overflow")
    };
    let used = requested.min(shape[dim]);
    let rank_one = shape.len() == 1;
    let (signal, dim) = if rank_one {
        (reshape(signal, Shape::new([1, shape[0]])), 1)
    } else { (signal, dim) };
    if let Some(outputs) = tune_shared_fft(vec![signal.clone()], dim, length, used, false) {
        let mut outputs = outputs.into_iter();
        let real = outputs.next().expect("real spectrum output");
        let imag = outputs.next().expect("imaginary spectrum output");
        return if rank_one { (reshape(real, Shape::new([length / 2 + 1])), reshape(imag, Shape::new([length / 2 + 1]))) }
            else { (real, imag) };
    }
    let mut plan = RealFftPlan::new(signal.client.clone(), length, FftMode::Forward)
        .unwrap_or_else(|e| panic!("rfft plan failed (requested={requested}, actual={length}): {e}"));
    let mut out_shape = signal.shape();
    out_shape[dim] = length / 2 + 1;
    let real = empty_device_dtype(signal.client.clone(), signal.device.clone(), out_shape.clone(), DType::F32);
    let imag = empty_device_dtype(signal.client.clone(), signal.device.clone(), out_shape, DType::F32);
    plan.forward(signal.binding(), real.clone().binding(), imag.clone().binding(), dim, used)
        .unwrap_or_else(|e| panic!("rfft launch failed (requested={requested}, actual={length}): {e}"));
    if rank_one {
        (reshape(real, Shape::new([length / 2 + 1])), reshape(imag, Shape::new([length / 2 + 1])))
    } else { (real, imag) }
}

/// Legacy inverse: compute the next-power-of-two inverse, then crop to N.
/// Prefer `irfft_exact` when the spectrum describes an actual N-point DFT.
pub fn irfft<R: Runtime>(real: RudaTensor<R>, imag: RudaTensor<R>, dim: usize, n: Option<usize>)
    -> RudaTensor<R>
{
    inverse(real, imag, dim, n, false)
}

/// Exact normalized N-point inverse real FFT. For odd N, pass Some(N): the
/// half-spectrum alone cannot distinguish an odd from an even original length.
pub fn irfft_exact<R: Runtime>(real: RudaTensor<R>, imag: RudaTensor<R>, dim: usize, n: Option<usize>)
    -> RudaTensor<R>
{
    inverse(real, imag, dim, n, true)
}

fn inverse<R: Runtime>(real: RudaTensor<R>, imag: RudaTensor<R>, dim: usize, n: Option<usize>, exact: bool)
    -> RudaTensor<R>
{
    assert_eq!(real.dtype, DType::F32, "ruFFT device kernels currently require F32 storage");
    assert_eq!(imag.dtype, real.dtype, "irfft: real and imaginary dtypes differ");
    assert_eq!(real.client.device_id(), imag.client.device_id(), "irfft: input devices differ");
    let shape = real.shape();
    assert_eq!(shape, imag.shape(), "irfft: real and imaginary shapes differ");
    assert!(dim < shape.len(), "irfft: dimension out of bounds");
    assert!(shape[dim] > 0, "irfft: spectrum must contain at least one bin");
    let inferred = (shape[dim] - 1).checked_mul(2).expect("irfft inferred length overflow");
    let requested = n.unwrap_or(inferred);
    assert!(requested > 0, "irfft: positive N is required (use Some(1) for one bin)");
    let length = if exact { requested } else {
        requested.checked_next_power_of_two().expect("irfft length overflow")
    };
    let used = shape[dim].min(length / 2 + 1);
    let rank_one = shape.len() == 1;
    let (real, imag, dim) = if rank_one {
        (reshape(real, Shape::new([1, shape[0]])), reshape(imag, Shape::new([1, shape[0]])), 1)
    } else { (real, imag, dim) };
    let mut out_shape = real.shape();
    out_shape[dim] = length;
    let output = if let Some(mut outputs) = tune_shared_fft(vec![real.clone(), imag.clone()], dim, length, used, true) {
        outputs.remove(0)
    } else {
        let mut plan = RealFftPlan::new(real.client.clone(), length, FftMode::Inverse)
            .unwrap_or_else(|e| panic!("irfft plan failed (requested={requested}, actual={length}): {e}"));
        let output = empty_device_dtype(real.client.clone(), real.device.clone(), out_shape.clone(), DType::F32);
        plan.inverse(real.binding(), imag.binding(), output.clone().binding(), dim, used)
            .unwrap_or_else(|e| panic!("irfft launch failed (requested={requested}, actual={length}): {e}"));
        output
    };
    // Cropping is part of the old API only. Neither path pads/copies spectra.
    let output = if length != requested {
        let ranges: Vec<_> = out_shape.iter().enumerate()
            .map(|(axis, &size)| 0..if axis == dim { requested } else { size }).collect();
        ruprim::indexing::slice(output, &ranges)
    } else { output };
    if rank_one { reshape(output, Shape::new([requested])) } else { output }
}