cubek-fft 0.3.0

//! Real-valued FFT with an intra-cube-parallel radix-2 Cooley-Tukey kernel.

use cubecl::prelude::*;
use cubecl::std::tensor::{
    AsView as _, AsViewExpand, AsViewMut as _, AsViewMutExpand, TensorHandle,
};

use crate::{
    fft::{
        FftMode,
        fft_parallel::{bit_reverse, fft_butterfly_parallel},
        limits::{max_shared_fft_n, max_units_per_cube},
        rfft_large::rfft_large_launch,
    },
    layout::BatchSignalLayout,
};

/// Real-valued Fast Fourier Transform.
pub fn rfft(
    client: &Client,
    signal: TensorHandle,
    dim: usize,
    dtype: ElemType,
) -> (TensorHandle, TensorHandle) {
    assert!(
        dim < signal.shape().len(),
        "dim must be between 0 and {}",
        signal.shape().len()
    );
    assert!(
        signal.shape()[dim].is_power_of_two(),
        "RFFT requires power-of-2 length"
    );
    let mut spectrum_shape = signal.shape().clone();
    spectrum_shape[dim] = signal.shape()[dim] / 2 + 1;

    let spectrum_re = TensorHandle::new_contiguous(
        spectrum_shape.clone(),
        client.empty(spectrum_shape.iter().product::<usize>() * dtype.size()),
        dtype,
    );

    let spectrum_im = TensorHandle::new_contiguous(
        spectrum_shape.clone(),
        client.empty(spectrum_shape.iter().product::<usize>() * dtype.size()),
        dtype,
    );

    rfft_launch(
        client,
        signal.binding(),
        spectrum_re.clone().binding(),
        spectrum_im.clone().binding(),
        dim,
        dtype,
    )
    .unwrap();

    (spectrum_re, spectrum_im)
}

/// Launches the RFFT kernel.
pub fn rfft_launch(
    client: &Client,
    signal: TensorBinding,
    spectrum_re: TensorBinding,
    spectrum_im: TensorBinding,
    dim: usize,
    dtype: ElemType,
) -> Result<(), LaunchError> {
    let signal_len = signal.shape[dim];
    rfft_launch_padded(
        client,
        signal,
        spectrum_re,
        spectrum_im,
        dim,
        signal_len,
        dtype,
    )
}

/// Launches the RFFT kernel while treating samples at `signal_len..n_fft` as zero.
pub fn rfft_launch_padded(
    client: &Client,
    signal: TensorBinding,
    spectrum_re: TensorBinding,
    spectrum_im: TensorBinding,
    dim: usize,
    signal_len: usize,
    dtype: ElemType,
) -> Result<(), LaunchError> {
    assert!(
        spectrum_re.shape == spectrum_im.shape,
        "spectrum real and imaginary shapes must match"
    );
    assert!(dim < signal.shape.len(), "dim must be in bounds");

    assert!(
        spectrum_re.shape[dim] >= 2,
        "RFFT spectrum dimension must contain at least DC and Nyquist bins"
    );
    let n_fft = (spectrum_re.shape[dim] - 1) * 2;
    assert!(n_fft.is_power_of_two(), "RFFT requires power-of-2 length");
    assert!(n_fft >= 2, "RFFT requires n_fft >= 2");
    assert!(
        signal_len <= signal.shape[dim],
        "signal_len ({signal_len}) must be <= signal dimension ({})",
        signal.shape[dim]
    );
    assert!(
        signal_len <= n_fft,
        "signal_len ({signal_len}) must be <= n_fft ({n_fft})"
    );

    let count: usize = signal
        .shape
        .iter()
        .enumerate()
        .filter(|(i, _)| *i != dim)
        .map(|(_, e)| *e)
        .product();
    if count == 0 {
        return Ok(());
    }

    if n_fft > max_shared_fft_n(client, dtype) {
        return rfft_large_launch(
            client,
            signal,
            spectrum_re,
            spectrum_im,
            dim,
            signal_len,
            dtype,
        );
    }

    let log2_n = n_fft.trailing_zeros() as usize;
    let threads_per_cube = (n_fft / 2).clamp(1, max_units_per_cube(client));

    let cube_dim = CubeDim::new_1d(threads_per_cube as u32);
    let cube_count = cubecl::calculate_cube_count_elemwise(client, count, CubeDim::new_single());

    rfft_kernel::launch(
        client,
        cube_count,
        cube_dim,
        signal.into_tensor_arg(),
        spectrum_re.into_tensor_arg(),
        spectrum_im.into_tensor_arg(),
        count as u32,
        signal_len as u32,
        n_fft,
        log2_n,
        threads_per_cube,
        dim,
        dtype,
    );
    Ok(())
}

#[cube(launch)]
fn rfft_kernel<F: Float>(
    signal: &Tensor<F>,
    spectrum_re: &mut Tensor<F>,
    spectrum_im: &mut Tensor<F>,
    num_windows: u32,
    signal_len: u32,
    #[comptime] n_fft: usize,
    #[comptime] log2_n: usize,
    #[comptime] threads_per_cube: usize,
    #[comptime] dim: usize,
    #[define(F)] _dtype: ElemType,
) {
    let window_index = CUBE_POS;
    if (window_index as u32) >= num_windows {
        terminate!();
    }

    let signal_view = signal.view(BatchSignalLayout::new(signal, window_index, dim));
    let mut spectrum_re_view =
        spectrum_re.view_mut(BatchSignalLayout::new(&*spectrum_re, window_index, dim));
    let mut spectrum_im_view =
        spectrum_im.view_mut(BatchSignalLayout::new(&*spectrum_im, window_index, dim));

    let mut shared_re = Shared::new_slice(n_fft);
    let mut shared_im = Shared::new_slice(n_fft);

    let mut i = UNIT_POS as usize;
    while i < n_fft {
        let j = bit_reverse(i, log2_n);
        let active = i < signal_len as usize;
        let src = select(active, i, 0);
        shared_re[j] = select(active, signal_view.read_checked(src), F::new(0.0_f32));
        shared_im[j] = F::new(0.0_f32);
        i += threads_per_cube;
    }
    sync_cube();

    fft_butterfly_parallel::<F>(
        &mut shared_re,
        &mut shared_im,
        n_fft,
        log2_n,
        threads_per_cube,
        FftMode::Forward,
    );

    let n_freq = comptime![n_fft / 2 + 1];
    let mut k = UNIT_POS as usize;
    while k < n_freq {
        spectrum_re_view.write_checked(k, shared_re[k]);
        spectrum_im_view.write_checked(k, shared_im[k]);
        k += threads_per_cube;
    }
    sync_cube();
}