diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 7d0a06507..46713d432 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -120,6 +120,10 @@ name = "contiguous" harness = false name = "fft" +[[bench]] +harness = false +name = "cfft" + [[bench]] harness = false name = "quantized_matmul" diff --git a/benchmarks/benches/cfft.rs b/benchmarks/benches/cfft.rs new file mode 100644 index 000000000..423c552fe --- /dev/null +++ b/benchmarks/benches/cfft.rs @@ -0,0 +1 @@ +benchmarks::run_bench!(cfft); diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index 21add803e..5721843f6 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -4,6 +4,9 @@ pub use cubek_attention::eval::backward::benchmarks as attention_backward; pub use cubek_attention::eval::forward::benchmarks as attention; pub use cubek_convolution::eval::benchmarks as conv2d; pub use cubek_fft::eval::benchmarks as fft; +pub mod cfft { + pub use cubek_fft::eval::benchmarks::CfftCategory as Category; +} pub use cubek_interpolate::eval::benchmarks as interpolate; pub use cubek_matmul::eval::benchmarks::gemm; pub use cubek_matmul::eval::benchmarks::gemv; @@ -26,6 +29,7 @@ pub fn all() -> &'static [&'static dyn BenchmarkCategory] { &crate::attention_backward::Category, &crate::contiguous::Category, &crate::conv2d::Category, + &crate::cfft::Category, &crate::fft::Category, &crate::gemm::Category, &crate::gemv::Category, diff --git a/crates/cubek-attention/src/backward/launch/dkdv.rs b/crates/cubek-attention/src/backward/launch/dkdv.rs index 871e3ea8e..9ec587410 100644 --- a/crates/cubek-attention/src/backward/launch/dkdv.rs +++ b/crates/cubek-attention/src/backward/launch/dkdv.rs @@ -50,7 +50,7 @@ fn flash_attention_backward_dv_kernel( let mut dv_acc = Array::new(val_dim); for dd in 0..val_dim { - dv_acc[dd] = E::new(0.0); + dv_acc[dd] = E::new(0.0_f32); } for i in 0..seq_q { @@ -60,7 +60,7 @@ fn flash_attention_backward_dv_kernel( let do_base = do_row_base + i * val_dim; let lse_i = lse[row_idx_base + i]; - let mut dot = E::new(0.0); + let mut dot = E::new(0.0_f32); for dd in 0..head_dim { dot += q[q_base + dd] * k[k_base + dd]; } @@ -115,7 +115,7 @@ fn flash_attention_backward_dk_kernel( let mut dk_acc = Array::new(head_dim); for dd in 0..head_dim { - dk_acc[dd] = E::new(0.0); + dk_acc[dd] = E::new(0.0_f32); } for i in 0..seq_q { @@ -126,13 +126,13 @@ fn flash_attention_backward_dk_kernel( let lse_i = lse[row_idx_base + i]; let d_i = d[row_idx_base + i]; - let mut dot = E::new(0.0); + let mut dot = E::new(0.0_f32); for dd in 0..head_dim { dot += q[q_base + dd] * k[k_base + dd]; } let p_ij = (dot * scale_e - lse_i).exp(); - let mut dp = E::new(0.0); + let mut dp = E::new(0.0_f32); for dd in 0..val_dim { dp += do_[do_base + dd] * v[v_base + dd]; } diff --git a/crates/cubek-attention/src/backward/launch/dq.rs b/crates/cubek-attention/src/backward/launch/dq.rs index a670f5404..7c7269a7c 100644 --- a/crates/cubek-attention/src/backward/launch/dq.rs +++ b/crates/cubek-attention/src/backward/launch/dq.rs @@ -45,7 +45,7 @@ fn flash_attention_backward_dq_kernel( let mut dq_acc = Array::new(head_dim); for dd in 0..head_dim { - dq_acc[dd] = E::new(0.0); + dq_acc[dd] = E::new(0.0_f32); } for j in 0..seq_kv { @@ -54,14 +54,14 @@ fn flash_attention_backward_dq_kernel( let k_base = k_row_base + j * head_dim; let v_base = v_row_base + j * val_dim; - let mut dot = E::new(0.0); + let mut dot = E::new(0.0_f32); for dd in 0..head_dim { dot += q[q_base + dd] * k[k_base + dd]; } let s_ij = dot * scale_e; let p_ij = (s_ij - lse_i).exp(); - let mut dp = E::new(0.0); + let mut dp = E::new(0.0_f32); for dd in 0..val_dim { dp += do_[do_base + dd] * v[v_base + dd]; } diff --git a/crates/cubek-attention/src/backward/launch/prepass.rs b/crates/cubek-attention/src/backward/launch/prepass.rs index 13b7db654..0d625e278 100644 --- a/crates/cubek-attention/src/backward/launch/prepass.rs +++ b/crates/cubek-attention/src/backward/launch/prepass.rs @@ -23,7 +23,7 @@ fn flash_attention_backward_prepass_kernel( let head_dim = o.shape(o.rank() - 1); let base = row_idx * head_dim; - let mut acc = E::new(0.0); + let mut acc = E::new(0.0_f32); for k in 0..head_dim { acc += o[base + k] * do_[base + k]; } diff --git a/crates/cubek-fft/Cargo.toml b/crates/cubek-fft/Cargo.toml index dff25b558..2f56a7fcd 100644 --- a/crates/cubek-fft/Cargo.toml +++ b/crates/cubek-fft/Cargo.toml @@ -26,6 +26,7 @@ benchmarks = ["cpu-reference", "cubecl/test-runtime"] [dependencies] cubecl = { workspace = true } +thiserror = { workspace = true } cubek-test-utils = { path = "./../cubek-test-utils/", version = "=0.3.0-pre.1", default-features = false, optional = true } num-complex = { version = "0.4.6", optional = true } diff --git a/crates/cubek-fft/src/complex.rs b/crates/cubek-fft/src/complex.rs new file mode 100644 index 000000000..3225eecc4 --- /dev/null +++ b/crates/cubek-fft/src/complex.rs @@ -0,0 +1,273 @@ +use cubecl::{ + frontend::CubePrimitive, + prelude::{ComputeClient, Runtime, StorageType, TensorBinding}, + std::tensor::TensorHandle, +}; + +use crate::FftError; + +/// A C32 tensor represented as adjacent real and imaginary F32 scalars. +#[derive(Clone)] +pub struct ComplexTensorHandle { + tensor: TensorHandle, + logical_strides: Vec, + physical_scalar_len: usize, +} + +impl ComplexTensorHandle { + /// Allocates a contiguous C32 tensor with the requested logical shape. + pub fn empty( + client: &ComputeClient, + shape: Vec, + dtype: StorageType, + ) -> Result { + ensure_c32_dtype(dtype)?; + let logical_strides = contiguous_strides(&shape)?; + let (_, physical_scalar_len) = scalar_layout(&shape, &logical_strides)?; + let byte_len = physical_scalar_len + .checked_mul(dtype.size()) + .ok_or(FftError::SizeOverflow)?; + + Self::new_strided(shape, logical_strides, client.empty(byte_len), dtype) + } + + /// Wraps a contiguous C32 buffer whose logical elements occupy adjacent scalar pairs. + pub fn new_contiguous( + shape: Vec, + handle: cubecl::server::Handle, + dtype: StorageType, + ) -> Result { + ensure_c32_dtype(dtype)?; + let logical_strides = contiguous_strides(&shape)?; + Self::new_strided(shape, logical_strides, handle, dtype) + } + + /// Wraps a C32 buffer using logical-complex-element strides. + pub fn new_strided( + shape: Vec, + logical_strides: Vec, + handle: cubecl::server::Handle, + dtype: StorageType, + ) -> Result { + ensure_c32_dtype(dtype)?; + if shape.len() != logical_strides.len() { + return Err(FftError::RankMismatch { + shape_rank: shape.len(), + stride_rank: logical_strides.len(), + }); + } + + let offset = handle.offset_start.unwrap_or_default(); + let offset_end = handle.offset_end.unwrap_or_default(); + let used_bytes = handle + .size() + .checked_sub(offset) + .and_then(|remaining| remaining.checked_sub(offset_end)) + .ok_or(FftError::InvalidBufferRange { + size: handle.size(), + offset_start: offset, + offset_end, + })?; + if !offset.is_multiple_of(dtype.size() as u64) { + return Err(FftError::MisalignedBuffer { + offset, + scalar_size: dtype.size(), + }); + } + + let (scalar_strides, physical_scalar_len) = scalar_layout(&shape, &logical_strides)?; + let available = usize::try_from(used_bytes / dtype.size() as u64) + .map_err(|_| FftError::SizeOverflow)?; + if available < physical_scalar_len { + return Err(FftError::InsufficientBuffer { + required: physical_scalar_len, + available, + }); + } + + Ok(Self { + tensor: TensorHandle::new(handle, shape, scalar_strides, dtype), + logical_strides, + physical_scalar_len, + }) + } + + /// Returns the logical complex shape. + pub fn shape(&self) -> &[usize] { + self.tensor.shape() + } + + /// Returns strides measured in logical complex elements. + pub fn strides(&self) -> &[usize] { + &self.logical_strides + } + + /// Returns physical strides measured in F32 scalars. + pub fn scalar_strides(&self) -> &[usize] { + self.tensor.strides() + } + + /// Returns the number of scalar F32 elements reachable through this layout. + pub fn physical_scalar_len(&self) -> usize { + self.physical_scalar_len + } + + /// Returns the physical scalar storage type. + pub fn dtype(&self) -> StorageType { + self.tensor.dtype + } + + /// Borrows the handle for a later CubeCL launch binding. + pub fn binding(&self) -> ComplexTensorBinding<'_, R> { + ComplexTensorBinding { handle: self } + } + + /// Returns the underlying scalar tensor metadata and allocation. + pub fn into_raw_parts(self) -> TensorHandle { + self.tensor + } +} + +/// A borrowed C32 tensor handle that can produce a CubeCL tensor binding at launch time. +pub struct ComplexTensorBinding<'a, R: Runtime> { + handle: &'a ComplexTensorHandle, +} + +#[allow(dead_code)] +impl ComplexTensorBinding<'_, R> { + pub(crate) fn shape(&self) -> &[usize] { + self.handle.shape() + } + + pub(crate) fn strides(&self) -> &[usize] { + self.handle.strides() + } + + pub(crate) fn dtype(&self) -> StorageType { + self.handle.dtype() + } + + /// Whether two bindings reference the exact same C32 handle and range. + pub(crate) fn is_same_tensor(&self, other: &Self) -> bool { + core::ptr::eq(self.handle, other.handle) + } + + pub(crate) fn tensor(&self) -> TensorBinding { + self.handle.tensor.clone().binding() + } + + pub(crate) fn ensure_unique_output(&self) -> Result<(), FftError> { + ensure_unique_output(&self.handle.tensor) + } + + pub(crate) fn output_tensor(&self) -> Result, FftError> { + self.ensure_unique_output()?; + Ok(self.tensor()) + } +} + +#[allow(dead_code)] +pub(crate) fn ensure_unique_output(tensor: &TensorHandle) -> Result<(), FftError> { + if tensor.can_mut() { + Ok(()) + } else { + Err(FftError::OverlappingBindings) + } +} + +fn ensure_c32_dtype(dtype: StorageType) -> Result<(), FftError> { + let expected = f32::as_type_native_unchecked().storage_type(); + if dtype == expected { + Ok(()) + } else { + Err(FftError::UnsupportedDtype { actual: dtype }) + } +} + +fn contiguous_strides(shape: &[usize]) -> Result, FftError> { + if shape.contains(&0) { + return Ok(vec![0; shape.len()]); + } + + let mut strides = vec![0; shape.len()]; + let mut stride = 1usize; + for (axis, extent) in shape.iter().enumerate().rev() { + strides[axis] = stride; + stride = stride.checked_mul(*extent).ok_or(FftError::SizeOverflow)?; + } + Ok(strides) +} + +fn scalar_layout( + shape: &[usize], + logical_strides: &[usize], +) -> Result<(Vec, usize), FftError> { + let scalar_strides = logical_strides + .iter() + .enumerate() + .map(|(axis, stride)| { + stride + .checked_mul(2) + .ok_or(FftError::StrideOverflow { axis }) + }) + .collect::, _>>()?; + + if shape.contains(&0) { + return Ok((scalar_strides, 0)); + } + + let last_imaginary_scalar = + shape + .iter() + .zip(&scalar_strides) + .try_fold(1usize, |offset, (extent, stride)| { + let axis_offset = (extent - 1) + .checked_mul(*stride) + .ok_or(FftError::SizeOverflow)?; + offset + .checked_add(axis_offset) + .ok_or(FftError::SizeOverflow) + })?; + let physical_scalar_len = last_imaginary_scalar + .checked_add(1) + .ok_or(FftError::SizeOverflow)?; + + Ok((scalar_strides, physical_scalar_len)) +} + +#[cfg(test)] +mod tests { + use cubecl::{Runtime, TestRuntime, frontend::CubePrimitive}; + + use super::*; + + #[test] + fn output_binding_rejects_an_aliased_allocation_before_binding() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + + let handle = client.empty(4 * dtype.size()); + let input = handle.clone(); + let aliased = + ComplexTensorHandle::::new_contiguous(vec![2], handle, dtype).unwrap(); + assert!(matches!( + aliased.binding().output_tensor(), + Err(FftError::OverlappingBindings) + )); + assert_eq!(input.size_in_used(), 4 * dtype.size() as u64); + } + + #[test] + fn invalid_handle_offset_range_returns_an_error_without_panicking() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let handle = client + .empty(4 * dtype.size()) + .offset_start(12) + .offset_end(8); + + let result = ComplexTensorHandle::::new_contiguous(vec![1], handle, dtype); + + assert!(matches!(result, Err(FftError::InvalidBufferRange { .. }))); + } +} diff --git a/crates/cubek-fft/src/error.rs b/crates/cubek-fft/src/error.rs new file mode 100644 index 000000000..3922917b4 --- /dev/null +++ b/crates/cubek-fft/src/error.rs @@ -0,0 +1,51 @@ +use cubecl::prelude::{LaunchError, StorageType}; + +#[derive(Debug, thiserror::Error)] +pub enum FftError { + #[error("unsupported FFT storage dtype {actual:?}; expected F32")] + UnsupportedDtype { actual: StorageType }, + #[error("shape rank {shape_rank} differs from stride rank {stride_rank}")] + RankMismatch { + shape_rank: usize, + stride_rank: usize, + }, + #[error("FFT axis {dim} is out of bounds for rank {rank}")] + AxisOutOfBounds { dim: usize, rank: usize }, + #[error("FFT length must be a power of two and at least 2, got {n_fft}")] + InvalidFftLength { n_fft: usize }, + #[error("FFT length {n_fft} exceeds this device's supported maximum {max_n_fft}")] + FftLengthExceedsDeviceLimit { n_fft: usize, max_n_fft: usize }, + #[error("{name}={value} is outside {min}..={max}")] + InvalidLength { + name: &'static str, + value: usize, + min: usize, + max: usize, + }, + #[error("complex buffer needs {required} scalar elements but only {available} are available")] + InsufficientBuffer { required: usize, available: usize }, + #[error("complex buffer byte offset {offset} is not aligned to scalar size {scalar_size}")] + MisalignedBuffer { offset: u64, scalar_size: usize }, + #[error( + "complex buffer offsets ({offset_start} from start, {offset_end} from end) exceed allocation size {size}" + )] + InvalidBufferRange { + size: u64, + offset_start: u64, + offset_end: u64, + }, + #[error("complex scalar stride at axis {axis} overflowed")] + StrideOverflow { axis: usize }, + #[error("complex buffer extent overflowed")] + SizeOverflow, + #[error("{name} shape {actual:?} does not match expected shape {expected:?}")] + ShapeMismatch { + name: &'static str, + actual: Vec, + expected: Vec, + }, + #[error("input and output allocations overlap")] + OverlappingBindings, + #[error(transparent)] + Launch(#[from] LaunchError), +} diff --git a/crates/cubek-fft/src/eval/benchmarks/benchmark.rs b/crates/cubek-fft/src/eval/benchmarks/benchmark.rs index c0c4112d2..47f7007e0 100644 --- a/crates/cubek-fft/src/eval/benchmarks/benchmark.rs +++ b/crates/cubek-fft/src/eval/benchmarks/benchmark.rs @@ -11,12 +11,16 @@ use cubecl::{ }; use cubek_test_utils::{RunSamples, StridedLayout, TestInput}; -use crate::eval::benchmarks::problem::FftProblem; +use crate::eval::benchmarks::problem::{CfftProblem, FftProblem}; use crate::eval::benchmarks::strategy::FftStrategy; -use crate::{FftMode, irfft_launch, rfft_launch}; +use crate::{ + ComplexTensorHandle, FftMode, FftNormalization, cfft_interleaved_launch, + fft::cfft::{CfftBindings, cfft_launch_any_size}, + irfft_interleaved_launch, irfft_launch, rfft_interleaved_launch, rfft_launch, +}; pub fn bench( - _strategy: &FftStrategy, + strategy: &FftStrategy, problem: &FftProblem, num_samples: usize, ) -> Result { @@ -26,6 +30,8 @@ pub fn bench( let bench = FftBench:: { shape: problem.shape.clone(), mode: problem.mode, + transform: TransformKind::Real, + strategy: *strategy, device, client, samples: num_samples, @@ -40,20 +46,169 @@ pub fn bench( Ok(RunSamples::new(durations)) } +pub fn bench_cfft( + strategy: &FftStrategy, + problem: &CfftProblem, + num_samples: usize, +) -> Result { + let device = ::Device::default(); + let client = ::client(&device); + let bench = FftBench:: { + shape: problem.shape.clone(), + mode: problem.mode, + transform: TransformKind::Complex, + strategy: *strategy, + device, + client, + samples: num_samples, + _e: PhantomData, + }; + let durations = bench + .run(TimingMethod::System) + .map_err(|e| format!("benchmark failed: {e}"))? + .durations; + Ok(RunSamples::new(durations)) +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum TransformKind { + Real, + Complex, +} + struct FftBench { shape: Vec, mode: FftMode, + transform: TransformKind, + strategy: FftStrategy, device: ::Device, client: ComputeClient, samples: usize, _e: PhantomData, } -#[derive(Clone)] -struct FftInput { - signal: TensorHandle, - spectrum_re: TensorHandle, - spectrum_im: TensorHandle, +enum FftInput { + SplitForward { + client: ComputeClient, + signal: TensorHandle, + spectrum_re: TensorHandle, + spectrum_im: TensorHandle, + elem: cubecl::ir::Type, + }, + SplitInverse { + client: ComputeClient, + signal: TensorHandle, + spectrum_re: TensorHandle, + spectrum_im: TensorHandle, + elem: cubecl::ir::Type, + }, + InterleavedForward { + client: ComputeClient, + signal: TensorHandle, + spectrum: ComplexTensorHandle, + }, + InterleavedInverse { + client: ComputeClient, + signal: TensorHandle, + spectrum: ComplexTensorHandle, + elem: cubecl::ir::Type, + }, + SplitComplex { + client: ComputeClient, + input_re: TensorHandle, + input_im: TensorHandle, + output_re: TensorHandle, + output_im: TensorHandle, + elem: cubecl::ir::Type, + }, + InterleavedComplex { + client: ComputeClient, + input: ComplexTensorHandle, + output: ComplexTensorHandle, + }, +} + +impl Clone for FftInput { + fn clone(&self) -> Self { + match self { + Self::SplitForward { + client, + signal, + spectrum_re, + spectrum_im, + elem, + } => Self::SplitForward { + client: client.clone(), + signal: signal.clone(), + spectrum_re: empty_handle(client, spectrum_re.shape().to_vec(), *elem), + spectrum_im: empty_handle(client, spectrum_im.shape().to_vec(), *elem), + elem: *elem, + }, + Self::SplitInverse { + client, + signal, + spectrum_re, + spectrum_im, + elem, + } => Self::SplitInverse { + client: client.clone(), + signal: empty_handle(client, signal.shape().to_vec(), *elem), + spectrum_re: spectrum_re.clone(), + spectrum_im: spectrum_im.clone(), + elem: *elem, + }, + Self::InterleavedForward { + client, + signal, + spectrum, + } => Self::InterleavedForward { + client: client.clone(), + signal: signal.clone(), + spectrum: ComplexTensorHandle::empty( + client, + spectrum.shape().to_vec(), + spectrum.dtype(), + ) + .expect("benchmark output must use a supported C32 layout"), + }, + Self::InterleavedInverse { + client, + signal, + spectrum, + elem, + } => Self::InterleavedInverse { + client: client.clone(), + signal: empty_handle(client, signal.shape().to_vec(), *elem), + spectrum: spectrum.clone(), + elem: *elem, + }, + Self::SplitComplex { + client, + input_re, + input_im, + output_re, + output_im, + elem, + } => Self::SplitComplex { + client: client.clone(), + input_re: input_re.clone(), + input_im: input_im.clone(), + output_re: empty_handle(client, output_re.shape().to_vec(), *elem), + output_im: empty_handle(client, output_im.shape().to_vec(), *elem), + elem: *elem, + }, + Self::InterleavedComplex { + client, + input, + output, + } => Self::InterleavedComplex { + client: client.clone(), + input: input.clone(), + output: ComplexTensorHandle::empty(client, output.shape().to_vec(), output.dtype()) + .expect("benchmark output must use a supported C32 layout"), + }, + } + } } fn make_uniform( @@ -77,6 +232,41 @@ fn empty_handle( TensorHandle::empty(client, shape, elem) } +fn make_interleaved_uniform( + client: &ComputeClient, + shape: Vec, + dtype: StorageType, + re_seed: u64, + im_seed: u64, +) -> ComplexTensorHandle { + let real = TestInput::builder(client.clone(), Shape::from(shape.clone())) + .layout(StridedLayout::RowMajor) + .dtype(dtype) + .uniform(re_seed, 0., 1.) + .f32_host_data(); + let imaginary = TestInput::builder(client.clone(), Shape::from(shape.clone())) + .layout(StridedLayout::RowMajor) + .dtype(dtype) + .uniform(im_seed, 0., 1.) + .f32_host_data(); + let values = real + .iter_indexed_f32() + .zip(imaginary.iter_indexed_f32()) + .flat_map(|((_, re), (_, im))| [re, im]) + .collect(); + + let mut physical_shape = shape.clone(); + let last = physical_shape.len() - 1; + physical_shape[last] *= 2; + let physical = TestInput::builder(client.clone(), Shape::from(physical_shape)) + .layout(StridedLayout::RowMajor) + .dtype(dtype) + .custom(values) + .generate_without_host_data(); + ComplexTensorHandle::new_contiguous(shape, physical.handle, dtype) + .expect("benchmark input must use a supported C32 layout") +} + impl Benchmark for FftBench { type Input = FftInput; type Output = (); @@ -86,55 +276,143 @@ impl Benchmark for FftBench { let elem = E::as_type_native_unchecked(); let storage = elem.storage_type(); - let mut shape_out = self.shape.clone(); let dim = self.shape.len() - 1; + + if self.transform == TransformKind::Complex { + return match self.strategy { + FftStrategy::Split => FftInput::SplitComplex { + client: client.clone(), + input_re: make_uniform(&client, self.shape.clone(), storage, 0), + input_im: make_uniform(&client, self.shape.clone(), storage, 1), + output_re: empty_handle(&client, self.shape.clone(), elem), + output_im: empty_handle(&client, self.shape.clone(), elem), + elem, + }, + FftStrategy::Interleaved => FftInput::InterleavedComplex { + client: client.clone(), + input: make_interleaved_uniform(&client, self.shape.clone(), storage, 0, 1), + output: ComplexTensorHandle::empty(&client, self.shape.clone(), storage) + .expect("benchmark output must use a supported C32 layout"), + }, + }; + } + + let mut shape_out = self.shape.clone(); shape_out[dim] = self.shape[dim] / 2 + 1; - match self.mode { - FftMode::Forward => { - let signal = make_uniform(&client, self.shape.clone(), storage, 0); - let spectrum_re = empty_handle(&client, shape_out.clone(), elem); - let spectrum_im = empty_handle(&client, shape_out, elem); - FftInput { - signal, - spectrum_re, - spectrum_im, - } - } - FftMode::Inverse => { - let signal = empty_handle(&client, self.shape.clone(), elem); - let spectrum_re = make_uniform(&client, shape_out.clone(), storage, 0); - let spectrum_im = make_uniform(&client, shape_out, storage, 1); - FftInput { - signal, - spectrum_re, - spectrum_im, - } - } + match self.strategy { + FftStrategy::Split => match self.mode { + FftMode::Forward => FftInput::SplitForward { + client: client.clone(), + signal: make_uniform(&client, self.shape.clone(), storage, 0), + spectrum_re: empty_handle(&client, shape_out.clone(), elem), + spectrum_im: empty_handle(&client, shape_out, elem), + elem, + }, + FftMode::Inverse => FftInput::SplitInverse { + client: client.clone(), + signal: empty_handle(&client, self.shape.clone(), elem), + spectrum_re: make_uniform(&client, shape_out.clone(), storage, 0), + spectrum_im: make_uniform(&client, shape_out, storage, 1), + elem, + }, + }, + FftStrategy::Interleaved => match self.mode { + FftMode::Forward => FftInput::InterleavedForward { + client: client.clone(), + signal: make_uniform(&client, self.shape.clone(), storage, 0), + spectrum: ComplexTensorHandle::empty(&client, shape_out, storage) + .expect("benchmark output must use a supported C32 layout"), + }, + FftMode::Inverse => FftInput::InterleavedInverse { + client: client.clone(), + signal: empty_handle(&client, self.shape.clone(), elem), + spectrum: make_interleaved_uniform(&client, shape_out, storage, 0, 1), + elem, + }, + }, } } fn execute(&self, input: Self::Input) -> Result<(), String> { let dim = self.shape.len() - 1; - match self.mode { - FftMode::Forward => rfft_launch( + match input { + FftInput::SplitForward { + signal, + spectrum_re, + spectrum_im, + .. + } => rfft_launch( &self.client, - input.signal.binding(), - input.spectrum_re.binding(), - input.spectrum_im.binding(), + signal.binding(), + spectrum_re.binding(), + spectrum_im.binding(), dim, E::as_type_native_unchecked().storage_type(), ) .map_err(|err| format!("{err}"))?, - FftMode::Inverse => irfft_launch( + FftInput::SplitInverse { + signal, + spectrum_re, + spectrum_im, + .. + } => irfft_launch( &self.client, - input.spectrum_re.binding(), - input.spectrum_im.binding(), - input.signal.binding(), + spectrum_re.binding(), + spectrum_im.binding(), + signal.binding(), dim, E::as_type_native_unchecked().storage_type(), ) .map_err(|err| format!("{err}"))?, + FftInput::InterleavedForward { + signal, spectrum, .. + } => rfft_interleaved_launch( + &self.client, + &signal, + spectrum.binding(), + dim, + FftNormalization::None, + ) + .map_err(|err| format!("{err}"))?, + FftInput::InterleavedInverse { + signal, spectrum, .. + } => irfft_interleaved_launch( + &self.client, + spectrum.binding(), + &signal, + dim, + FftNormalization::ByN, + ) + .map_err(|err| format!("{err}"))?, + FftInput::SplitComplex { + input_re, + input_im, + output_re, + output_im, + .. + } => cfft_launch_any_size( + &self.client, + CfftBindings { + input_re: input_re.binding(), + input_im: input_im.binding(), + output_re: output_re.binding(), + output_im: output_im.binding(), + }, + dim, + E::as_type_native_unchecked().storage_type(), + self.mode, + ) + .map_err(|err| format!("{err}"))?, + FftInput::InterleavedComplex { input, output, .. } => cfft_interleaved_launch( + &self.client, + input.binding(), + output.binding(), + dim, + self.mode, + FftNormalization::None, + ) + .map_err(|err| format!("{err}"))?, } Ok(()) } @@ -145,8 +423,9 @@ impl Benchmark for FftBench { fn name(&self) -> String { format!( - "fft-{}-{:?}-{:?}", + "fft-{}-{}-{:?}-{:?}", E::as_type_native_unchecked(), + self.strategy.id(), self.shape, self.mode, ) @@ -157,3 +436,45 @@ impl Benchmark for FftBench { future::block_on(self.client.sync()).unwrap() } } + +#[cfg(test)] +mod tests { + use super::*; + + fn interleaved_bench(mode: FftMode) -> FftBench { + let device = ::Device::default(); + let client = ::client(&device); + FftBench { + shape: vec![1, 8], + mode, + transform: TransformKind::Real, + strategy: FftStrategy::Interleaved, + device, + client, + samples: 1, + _e: PhantomData, + } + } + + #[test] + fn cloned_interleaved_forward_sample_has_fresh_writable_output() { + let bench = interleaved_bench(FftMode::Forward); + let prepared = bench.prepare(); + bench.execute(prepared.clone()).unwrap(); + } + + #[test] + fn cloned_interleaved_inverse_sample_has_fresh_writable_output() { + let bench = interleaved_bench(FftMode::Inverse); + let prepared = bench.prepare(); + bench.execute(prepared.clone()).unwrap(); + } + + #[test] + fn cloned_interleaved_cfft_sample_has_fresh_writable_output() { + let mut bench = interleaved_bench(FftMode::Forward); + bench.transform = TransformKind::Complex; + let prepared = bench.prepare(); + bench.execute(prepared.clone()).unwrap(); + } +} diff --git a/crates/cubek-fft/src/eval/benchmarks/correctness.rs b/crates/cubek-fft/src/eval/benchmarks/correctness.rs index 684bbb81a..7db3c6e20 100644 --- a/crates/cubek-fft/src/eval/benchmarks/correctness.rs +++ b/crates/cubek-fft/src/eval/benchmarks/correctness.rs @@ -1,11 +1,28 @@ use cubecl::{Runtime, TestRuntime}; use cubek_test_utils::{HostData, Progress}; -use crate::eval::benchmarks::problem::FftProblem; +use crate::eval::benchmarks::problem::{CfftProblem, FftProblem}; use crate::eval::benchmarks::strategy::FftStrategy; -use crate::eval::cpu_reference::{cpu_reference_result, kernel_result as fft_kernel_result}; +use crate::eval::cpu_reference::{ + complex_kernel_result, cpu_reference_complex_result, cpu_reference_result, + interleaved_kernel_result, kernel_result as fft_kernel_result, +}; pub struct FftCorrectness; +pub struct CfftCorrectness; + +#[derive(Debug, PartialEq, Eq)] +enum CorrectnessKernel { + Split, + Interleaved, +} + +fn kernel_backend(strategy: FftStrategy) -> CorrectnessKernel { + match strategy { + FftStrategy::Split => CorrectnessKernel::Split, + FftStrategy::Interleaved => CorrectnessKernel::Interleaved, + } +} impl cubek_test_utils::Correctness for FftCorrectness { type Problem = FftProblem; @@ -13,33 +30,88 @@ impl cubek_test_utils::Correctness for FftCorrectness { fn kernel_result( &self, - _strategy: &FftStrategy, + strategy: &FftStrategy, + problem: &FftProblem, + seeds: &[u64], + ) -> Result { + let device = ::Device::default(); + let client = ::client(&device); + let dim = problem.shape.len() - 1; + match kernel_backend(*strategy) { + CorrectnessKernel::Split => fft_kernel_result( + client, + problem.shape.clone(), + dim, + problem.mode, + seeds[0], + seeds[1], + ), + CorrectnessKernel::Interleaved => interleaved_kernel_result( + client, + problem.shape.clone(), + dim, + problem.mode, + seeds[0], + seeds[1], + ), + } + } + + fn reference_result( + &self, problem: &FftProblem, seeds: &[u64], + progress: Option<&Progress>, + ) -> Result { + let device = ::Device::default(); + let client = ::client(&device); + let dim = problem.shape.len() - 1; + cpu_reference_result( + client, + problem.shape.clone(), + dim, + problem.mode, + seeds[0], + seeds[1], + progress, + ) + } +} + +impl cubek_test_utils::Correctness for CfftCorrectness { + type Problem = CfftProblem; + type Strategy = FftStrategy; + + fn kernel_result( + &self, + strategy: &FftStrategy, + problem: &CfftProblem, + seeds: &[u64], ) -> Result { let device = ::Device::default(); let client = ::client(&device); let dim = problem.shape.len() - 1; - fft_kernel_result( + complex_kernel_result( client, problem.shape.clone(), dim, problem.mode, seeds[0], seeds[1], + kernel_backend(*strategy) == CorrectnessKernel::Interleaved, ) } fn reference_result( &self, - problem: &FftProblem, + problem: &CfftProblem, seeds: &[u64], progress: Option<&Progress>, ) -> Result { let device = ::Device::default(); let client = ::client(&device); let dim = problem.shape.len() - 1; - cpu_reference_result( + Ok(cpu_reference_complex_result( client, problem.shape.clone(), dim, @@ -47,6 +119,19 @@ impl cubek_test_utils::Correctness for FftCorrectness { seeds[0], seeds[1], progress, - ) + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn interleaved_strategy_selects_interleaved_correctness_kernel() { + assert_eq!( + kernel_backend(FftStrategy::Interleaved), + CorrectnessKernel::Interleaved + ); } } diff --git a/crates/cubek-fft/src/eval/benchmarks/mod.rs b/crates/cubek-fft/src/eval/benchmarks/mod.rs index 30500fb9f..d4d85a6ed 100644 --- a/crates/cubek-fft/src/eval/benchmarks/mod.rs +++ b/crates/cubek-fft/src/eval/benchmarks/mod.rs @@ -5,14 +5,15 @@ mod correctness; mod problem; mod strategy; -pub use benchmark::bench; -pub use correctness::FftCorrectness; -pub use problem::{FftProblem, problems}; +pub use benchmark::{bench, bench_cfft}; +pub use correctness::{CfftCorrectness, FftCorrectness}; +pub use problem::{CfftProblem, FftProblem, cfft_problems, problems}; pub use strategy::{FftStrategy, strategies}; use cubek_test_utils::{CatalogEntry, RunSamples}; pub struct Category; +pub struct CfftCategory; impl cubek_test_utils::Category for Category { type Problem = FftProblem; @@ -49,3 +50,40 @@ impl cubek_test_utils::Category for Category { Some(&FftCorrectness) } } + +impl cubek_test_utils::Category for CfftCategory { + type Problem = CfftProblem; + type Strategy = FftStrategy; + + fn id(&self) -> &'static str { + "cfft" + } + + fn label(&self) -> &'static str { + "CFFT" + } + + fn problems(&self) -> Vec> { + cfft_problems() + } + + fn strategies(&self) -> Vec> { + strategies() + } + + fn bench( + &self, + strategy: &FftStrategy, + problem: &CfftProblem, + num_samples: usize, + ) -> Result { + bench_cfft(strategy, problem, num_samples) + } + + fn correctness( + &self, + ) -> Option<&dyn cubek_test_utils::Correctness> + { + Some(&CfftCorrectness) + } +} diff --git a/crates/cubek-fft/src/eval/benchmarks/problem.rs b/crates/cubek-fft/src/eval/benchmarks/problem.rs index dc43c66ab..d0c051192 100644 --- a/crates/cubek-fft/src/eval/benchmarks/problem.rs +++ b/crates/cubek-fft/src/eval/benchmarks/problem.rs @@ -2,92 +2,157 @@ use cubek_test_utils::CatalogEntry; use crate::FftMode; +/// Real FFT benchmark problem. Kept source-compatible with the original +/// benchmark API so downstream users can construct it with two public fields. pub struct FftProblem { pub shape: Vec, pub mode: FftMode, } +/// Complex-to-complex FFT benchmark problem, exposed through [`CfftCategory`](super::CfftCategory). +pub struct CfftProblem { + pub shape: Vec, + pub mode: FftMode, +} + pub fn problems() -> Vec> { vec![ - CatalogEntry::new( + fft_problem( "forward_5x2x2048", "Forward (5x2x2048)", - FftProblem { - shape: vec![5, 2, 2048], - mode: FftMode::Forward, - }, + vec![5, 2, 2048], + FftMode::Forward, ), - CatalogEntry::new( + fft_problem( "inverse_5x2x2048", "Inverse (5x2x2048)", - FftProblem { - shape: vec![5, 2, 2048], - mode: FftMode::Inverse, - }, + vec![5, 2, 2048], + FftMode::Inverse, ), - CatalogEntry::new( + fft_problem( "forward_128x2048", "Forward (128x2048)", - FftProblem { - shape: vec![128, 2048], - mode: FftMode::Forward, - }, + vec![128, 2048], + FftMode::Forward, ), - CatalogEntry::new( + fft_problem( "inverse_128x2048", "Inverse (128x2048)", - FftProblem { - shape: vec![128, 2048], - mode: FftMode::Inverse, - }, + vec![128, 2048], + FftMode::Inverse, ), - CatalogEntry::new( + fft_problem( "forward_1x4096", "Forward (1x4096)", - FftProblem { - shape: vec![1, 4096], - mode: FftMode::Forward, - }, + vec![1, 4096], + FftMode::Forward, ), - CatalogEntry::new( + fft_problem( "inverse_1x4096", "Inverse (1x4096)", - FftProblem { - shape: vec![1, 4096], - mode: FftMode::Inverse, - }, + vec![1, 4096], + FftMode::Inverse, ), - CatalogEntry::new( + fft_problem( "forward_1x8192", "Forward (1x8192)", - FftProblem { - shape: vec![1, 8192], - mode: FftMode::Forward, - }, + vec![1, 8192], + FftMode::Forward, ), - CatalogEntry::new( + fft_problem( "inverse_1x8192", "Inverse (1x8192)", - FftProblem { - shape: vec![1, 8192], - mode: FftMode::Inverse, - }, + vec![1, 8192], + FftMode::Inverse, ), - CatalogEntry::new( + fft_problem( "forward_1x16384", "Forward (1x16384)", - FftProblem { - shape: vec![1, 16384], - mode: FftMode::Forward, - }, + vec![1, 16384], + FftMode::Forward, ), - CatalogEntry::new( + fft_problem( "inverse_1x16384", "Inverse (1x16384)", - FftProblem { - shape: vec![1, 16384], - mode: FftMode::Inverse, - }, + vec![1, 16384], + FftMode::Inverse, + ), + ] +} + +pub fn cfft_problems() -> Vec> { + vec![ + cfft_problem( + "forward_1x4096", + "Forward (1x4096)", + vec![1, 4096], + FftMode::Forward, + ), + cfft_problem( + "inverse_1x4096", + "Inverse (1x4096)", + vec![1, 4096], + FftMode::Inverse, + ), + cfft_problem( + "forward_1x8192", + "Forward (1x8192)", + vec![1, 8192], + FftMode::Forward, + ), + cfft_problem( + "inverse_1x8192", + "Inverse (1x8192)", + vec![1, 8192], + FftMode::Inverse, ), ] } + +fn fft_problem( + id: &str, + label: &str, + shape: Vec, + mode: FftMode, +) -> CatalogEntry { + CatalogEntry::new(id, label, FftProblem { shape, mode }) +} + +fn cfft_problem( + id: &str, + label: &str, + shape: Vec, + mode: FftMode, +) -> CatalogEntry { + CatalogEntry::new(id, label, CfftProblem { shape, mode }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fft_problem_remains_constructible_with_shape_and_mode_only() { + let problem = FftProblem { + shape: vec![1, 8], + mode: FftMode::Forward, + }; + assert_eq!(problem.shape, [1, 8]); + assert_eq!(problem.mode, FftMode::Forward); + assert!( + problems() + .iter() + .all(|entry| !entry.id.starts_with("cfft_")) + ); + } + + #[test] + fn cfft_catalog_keeps_small_and_four_step_cases_separate() { + let ids = cfft_problems() + .into_iter() + .map(|entry| entry.id) + .collect::>(); + assert!(ids.contains(&"forward_1x4096".to_string())); + assert!(ids.contains(&"inverse_1x8192".to_string())); + } +} diff --git a/crates/cubek-fft/src/eval/benchmarks/strategy.rs b/crates/cubek-fft/src/eval/benchmarks/strategy.rs index 207b654eb..ced2fe574 100644 --- a/crates/cubek-fft/src/eval/benchmarks/strategy.rs +++ b/crates/cubek-fft/src/eval/benchmarks/strategy.rs @@ -1,7 +1,23 @@ use cubek_test_utils::CatalogEntry; -pub struct FftStrategy; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FftStrategy { + Split, + Interleaved, +} + +impl FftStrategy { + pub fn id(self) -> &'static str { + match self { + Self::Split => "default", + Self::Interleaved => "interleaved", + } + } +} pub fn strategies() -> Vec> { - vec![CatalogEntry::new("default", "Default", FftStrategy)] + vec![ + CatalogEntry::new("default", "Default (split)", FftStrategy::Split), + CatalogEntry::new("interleaved", "Interleaved C32", FftStrategy::Interleaved), + ] } diff --git a/crates/cubek-fft/src/eval/cpu_reference.rs b/crates/cubek-fft/src/eval/cpu_reference.rs index af900cc8e..a7304be79 100644 --- a/crates/cubek-fft/src/eval/cpu_reference.rs +++ b/crates/cubek-fft/src/eval/cpu_reference.rs @@ -5,7 +5,7 @@ use std::f32::consts::PI; use cubecl::{ - TestRuntime, + CubeElement, TestRuntime, client::ComputeClient, frontend::CubePrimitive, zspace::{Shape, Strides}, @@ -16,7 +16,15 @@ use cubek_test_utils::{ }; use num_complex::Complex; -use crate::fft::{FftMode, irfft_launch, rfft_launch}; +use crate::{ + ComplexTensorHandle, FftNormalization, + fft::{ + FftMode, + cfft::{CfftBindings, cfft_launch_any_size}, + cfft_interleaved_launch, irfft_interleaved_launch, irfft_launch, rfft_interleaved_launch, + rfft_launch, + }, +}; /// Run the FFT kernel for `mode` against the given problem with seeded inputs /// and return its output as a [`HostData`]. @@ -112,6 +120,253 @@ pub fn kernel_result( } } +/// Interleaved C32 counterpart to [`kernel_result`] for benchmark correctness. +pub fn interleaved_kernel_result( + client: ComputeClient, + shape: Vec, + dim: usize, + mode: FftMode, + seed_lhs: u64, + seed_rhs: u64, +) -> Result { + let dtype = f32::as_type_native_unchecked().storage_type(); + + match mode { + FftMode::Forward => { + let (signal, _) = TestInput::builder(client.clone(), shape.clone()) + .dtype(dtype) + .uniform(seed_lhs, -1., 1.) + .generate_with_f32_host_data(); + let mut spectrum_shape = shape; + spectrum_shape[dim] = spectrum_shape[dim] / 2 + 1; + let spectrum = ComplexTensorHandle::empty(&client, spectrum_shape.clone(), dtype) + .map_err(|err| err.to_string())?; + + let outcome = launch_and_capture_outcome(&client, |c| { + rfft_interleaved_launch(c, &signal, spectrum.binding(), dim, FftNormalization::None) + .into() + }); + match outcome { + ExecutionOutcome::CompileError(e) => Err(format!("compile error: {e}")), + ExecutionOutcome::Executed => { + Ok(stack_interleaved(&client, spectrum, spectrum_shape)) + } + } + } + FftMode::Inverse => { + let mut spectrum_shape = shape.clone(); + spectrum_shape[dim] = shape[dim] / 2 + 1; + let (_, re) = TestInput::builder(client.clone(), Shape::from(spectrum_shape.clone())) + .dtype(dtype) + .uniform(seed_lhs, -1., 1.) + .generate_with_f32_host_data(); + let (_, im) = TestInput::builder(client.clone(), Shape::from(spectrum_shape.clone())) + .dtype(dtype) + .uniform(seed_rhs, -1., 1.) + .generate_with_f32_host_data(); + let mut interleaved = Vec::with_capacity(re.shape.num_elements() * 2); + let mut re_values = Vec::with_capacity(re.shape.num_elements()); + let mut im_values = Vec::with_capacity(im.shape.num_elements()); + pack_contiguous( + &mut re_values, + as_f32_slice(&re), + &re.strides, + &spectrum_shape, + ); + pack_contiguous( + &mut im_values, + as_f32_slice(&im), + &im.strides, + &spectrum_shape, + ); + for (re, im) in re_values.into_iter().zip(im_values) { + interleaved.extend([re, im]); + } + let spectrum = ComplexTensorHandle::new_contiguous( + spectrum_shape, + client.create_from_slice(f32::as_bytes(&interleaved)), + dtype, + ) + .map_err(|err| err.to_string())?; + let signal = TestInput::builder(client.clone(), Shape::from(shape)) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + + let outcome = launch_and_capture_outcome(&client, |c| { + irfft_interleaved_launch(c, spectrum.binding(), &signal, dim, FftNormalization::ByN) + .into() + }); + match outcome { + ExecutionOutcome::CompileError(e) => Err(format!("compile error: {e}")), + ExecutionOutcome::Executed => Ok(HostData::from_tensor_handle( + &client, + signal, + HostDataType::F32, + )), + } + } + } +} + +/// Run a split or interleaved CFFT benchmark correctness kernel. +pub fn complex_kernel_result( + client: ComputeClient, + shape: Vec, + dim: usize, + mode: FftMode, + seed_lhs: u64, + seed_rhs: u64, + interleaved: bool, +) -> Result { + let dtype = f32::as_type_native_unchecked().storage_type(); + let (input_re, _) = TestInput::builder(client.clone(), Shape::from(shape.clone())) + .dtype(dtype) + .uniform(seed_lhs, -1., 1.) + .generate_with_f32_host_data(); + let (input_im, _) = TestInput::builder(client.clone(), Shape::from(shape.clone())) + .dtype(dtype) + .uniform(seed_rhs, -1., 1.) + .generate_with_f32_host_data(); + + if interleaved { + let re = HostData::from_tensor_handle(&client, input_re, HostDataType::F32); + let im = HostData::from_tensor_handle(&client, input_im, HostDataType::F32); + let input = interleave_host_data(&client, &re, &im, shape.clone(), dtype)?; + let output = ComplexTensorHandle::empty(&client, shape.clone(), dtype) + .map_err(|err| err.to_string())?; + let outcome = launch_and_capture_outcome(&client, |c| { + cfft_interleaved_launch( + c, + input.binding(), + output.binding(), + dim, + mode, + FftNormalization::None, + ) + .into() + }); + return match outcome { + ExecutionOutcome::CompileError(e) => Err(format!("compile error: {e}")), + ExecutionOutcome::Executed => Ok(stack_interleaved(&client, output, shape)), + }; + } + + let output_re = TestInput::builder(client.clone(), Shape::from(shape.clone())) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + let output_im = TestInput::builder(client.clone(), Shape::from(shape)) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + let outcome = launch_and_capture_outcome(&client, |c| { + cfft_launch_any_size( + c, + CfftBindings { + input_re: input_re.clone().binding(), + input_im: input_im.clone().binding(), + output_re: output_re.clone().binding(), + output_im: output_im.clone().binding(), + }, + dim, + dtype, + mode, + ) + .into() + }); + match outcome { + ExecutionOutcome::CompileError(e) => Err(format!("compile error: {e}")), + ExecutionOutcome::Executed => Ok(stack_re_im( + HostData::from_tensor_handle(&client, output_re, HostDataType::F32), + HostData::from_tensor_handle(&client, output_im, HostDataType::F32), + )), + } +} + +/// CPU reference for a complex-to-complex FFT with split seeded inputs. +pub fn cpu_reference_complex_result( + client: ComputeClient, + shape: Vec, + dim: usize, + mode: FftMode, + seed_lhs: u64, + seed_rhs: u64, + progress: Option<&Progress>, +) -> HostData { + let dtype = f32::as_type_native_unchecked().storage_type(); + let (_, re) = TestInput::builder(client.clone(), Shape::from(shape.clone())) + .dtype(dtype) + .uniform(seed_lhs, -1., 1.) + .generate_with_f32_host_data(); + let (_, im) = TestInput::builder(client, Shape::from(shape)) + .dtype(dtype) + .uniform(seed_rhs, -1., 1.) + .generate_with_f32_host_data(); + let (re, im) = cfft_ref(&re, &im, dim, mode, progress); + stack_re_im(re, im) +} + +fn interleave_host_data( + client: &ComputeClient, + re: &HostData, + im: &HostData, + shape: Vec, + dtype: cubecl::prelude::StorageType, +) -> Result, String> { + let mut re_values = Vec::with_capacity(re.shape.num_elements()); + let mut im_values = Vec::with_capacity(im.shape.num_elements()); + pack_contiguous(&mut re_values, as_f32_slice(re), &re.strides, &shape); + pack_contiguous(&mut im_values, as_f32_slice(im), &im.strides, &shape); + let mut values = Vec::with_capacity(re_values.len() * 2); + for (re, im) in re_values.into_iter().zip(im_values) { + values.extend([re, im]); + } + ComplexTensorHandle::new_contiguous( + shape, + client.create_from_slice(f32::as_bytes(&values)), + dtype, + ) + .map_err(|err| err.to_string()) +} + +fn as_f32_slice(host: &HostData) -> &[f32] { + match &host.data { + HostDataVec::F32(values) => values, + _ => unreachable!("FFT correctness data is always F32"), + } +} + +fn stack_interleaved( + client: &ComputeClient, + spectrum: ComplexTensorHandle, + logical_shape: Vec, +) -> HostData { + let raw = spectrum.into_raw_parts(); + let bytes = client.read_one(raw.handle).unwrap(); + let scalars = f32::from_bytes(&bytes); + let mut re = Vec::with_capacity(scalars.len() / 2); + let mut im = Vec::with_capacity(scalars.len() / 2); + for pair in scalars.chunks_exact(2) { + re.push(pair[0]); + im.push(pair[1]); + } + let shape = Shape::from(logical_shape); + let strides = StridedLayout::RowMajor.compute_strides(&shape); + stack_re_im( + HostData { + data: HostDataVec::F32(re), + shape: shape.clone(), + strides: strides.clone(), + }, + HostData { + data: HostDataVec::F32(im), + shape, + strides, + }, + ) +} + /// CPU-only counterpart to [`kernel_result`]: generate the same seeded inputs /// and run the recursive Cooley-Tukey reference. Returns the stacked re/im /// pair for [`FftMode::Forward`] and the reconstructed signal for @@ -246,6 +501,56 @@ fn fft_recursive(x: &mut [Complex], fft_mode: FftMode) { } } +fn cfft_ref( + re: &HostData, + im: &HostData, + dim: usize, + mode: FftMode, + progress: Option<&Progress>, +) -> (HostData, HostData) { + let shape = re.shape.as_slice(); + let n_fft = shape[dim]; + let num_windows = re.shape.num_elements() / n_fft; + let strides = StridedLayout::RowMajor.compute_strides(&re.shape); + if let Some(progress) = progress { + progress.set_total(num_windows as u64); + } + + let mut out_re = vec![0.0; re.shape.num_elements()]; + let mut out_im = vec![0.0; re.shape.num_elements()]; + for window in 0..num_windows { + let mut coords = get_coords(window, shape, dim); + let mut values = Vec::with_capacity(n_fft); + for i in 0..n_fft { + coords[dim] = i; + values.push(Complex::new(re.get_f32(&coords), im.get_f32(&coords))); + } + fft_recursive(&mut values, mode); + for (i, value) in values.into_iter().enumerate() { + coords[dim] = i; + let index = compute_index(&strides, &coords); + out_re[index] = value.re; + out_im[index] = value.im; + } + if let Some(progress) = progress { + progress.bump(); + } + } + + ( + HostData { + data: HostDataVec::F32(out_re), + shape: re.shape.clone(), + strides: strides.clone(), + }, + HostData { + data: HostDataVec::F32(out_im), + shape: re.shape.clone(), + strides, + }, + ) +} + /// Reference IRFFT: reconstruct real signal from first n/2 + 1 complex bins. pub fn irfft_ref( re: &HostData, diff --git a/crates/cubek-fft/src/fft/cfft.rs b/crates/cubek-fft/src/fft/cfft.rs index 410f38e75..1c31aebc6 100644 --- a/crates/cubek-fft/src/fft/cfft.rs +++ b/crates/cubek-fft/src/fft/cfft.rs @@ -391,7 +391,7 @@ fn cfft_four_step_radix1_kernel( /// /// Grid: `count * N1` cubes. `CUBE_POS = window * N1 + k1`. #[cube(launch)] -fn cfft_four_step_radix2_kernel( +pub(crate) fn cfft_four_step_radix2_kernel( scratch_re: &mut Tensor, scratch_im: &mut Tensor, num_cubes: u32, diff --git a/crates/cubek-fft/src/fft/cfft_interleaved.rs b/crates/cubek-fft/src/fft/cfft_interleaved.rs new file mode 100644 index 000000000..f279601f5 --- /dev/null +++ b/crates/cubek-fft/src/fft/cfft_interleaved.rs @@ -0,0 +1,504 @@ +use core::f32::consts::PI; + +use cubecl::prelude::*; +use cubecl::std::tensor::{ + AsView as _, AsViewExpand, AsViewMut as _, AsViewMutExpand, TensorHandle, +}; + +use crate::{ + ComplexTensorBinding, ComplexTensorHandle, FftError, FftNormalization, + fft::{ + FftMode, + cfft::{cfft_four_step_radix2_kernel, factor_four_step}, + fft_parallel::{bit_reverse, fft_butterfly_parallel}, + limits::{max_shared_fft_n, max_units_per_cube}, + }, + interleaved_layout::InterleavedBatchSignalLayout, +}; + +/// Runs a C32 FFT over an interleaved complex tensor along `dim`. +pub fn cfft_interleaved( + input: ComplexTensorHandle, + dim: usize, + mode: FftMode, + normalization: FftNormalization, +) -> Result, FftError> { + let shape = input.shape().to_vec(); + let client = R::client(&Default::default()); + let plan = cfft_plan(&client, &shape, dim)?; + normalization.scale_f32(plan.n_fft)?; + + let strides = input.strides().to_vec(); + ensure_non_overlapping_output_layout(&shape, &strides)?; + let dtype = input.dtype(); + let byte_len = input + .physical_scalar_len() + .checked_mul(dtype.size()) + .ok_or(FftError::SizeOverflow)?; + let output = ComplexTensorHandle::new_strided(shape, strides, client.empty(byte_len), dtype)?; + + cfft_interleaved_launch( + &client, + input.binding(), + output.binding(), + dim, + mode, + normalization, + )?; + Ok(output) +} + +/// Launches a C32 FFT into an interleaved output tensor. +pub fn cfft_interleaved_launch( + client: &ComputeClient, + input: ComplexTensorBinding<'_, R>, + output: ComplexTensorBinding<'_, R>, + dim: usize, + mode: FftMode, + normalization: FftNormalization, +) -> Result<(), FftError> { + let input_shape = input.shape(); + let plan = cfft_plan(client, input_shape, dim)?; + if output.shape() != input_shape { + return Err(FftError::ShapeMismatch { + name: "output", + actual: output.shape().to_vec(), + expected: input_shape.to_vec(), + }); + } + if output.dtype() != input.dtype() { + return Err(FftError::UnsupportedDtype { + actual: output.dtype(), + }); + } + if input.is_same_tensor(&output) { + return Err(FftError::OverlappingBindings); + } + ensure_non_overlapping_output_layout(output.shape(), output.strides())?; + + normalization.scale_f32(plan.n_fft)?; + + output.ensure_unique_output()?; + if plan.count == 0 { + return Ok(()); + } + + if plan.n_fft <= max_shared_fft_n(client) { + cfft_interleaved_shared_launch(client, input, output, dim, mode, normalization, plan) + } else { + cfft_interleaved_four_step_launch(client, input, output, dim, mode, normalization, plan) + } +} + +struct CfftPlan { + n_fft: usize, + count: usize, + count_u32: u32, +} + +fn cfft_plan( + client: &ComputeClient, + shape: &[usize], + dim: usize, +) -> Result { + validate_fft_shape(shape, dim)?; + let n_fft = shape[dim]; + let max_n = max_shared_fft_n(client); + let max_four_step_n = max_n.saturating_mul(max_n); + if n_fft > max_four_step_n { + return Err(FftError::InvalidFftLength { n_fft }); + } + let count = shape + .iter() + .enumerate() + .filter(|(axis, _)| *axis != dim) + .try_fold(1usize, |count, (_, extent)| { + count.checked_mul(*extent).ok_or(FftError::SizeOverflow) + })?; + let count_u32 = u32::try_from(count).map_err(|_| FftError::SizeOverflow)?; + Ok(CfftPlan { + n_fft, + count, + count_u32, + }) +} + +fn cfft_interleaved_shared_launch( + client: &ComputeClient, + input: ComplexTensorBinding<'_, R>, + output: ComplexTensorBinding<'_, R>, + dim: usize, + mode: FftMode, + normalization: FftNormalization, + plan: CfftPlan, +) -> Result<(), FftError> { + let log2_n = plan.n_fft.trailing_zeros() as usize; + let threads_per_cube = (plan.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, plan.count, CubeDim::new_single()); + + cfft_interleaved_shared_kernel::launch::( + client, + cube_count, + cube_dim, + input.tensor().into_tensor_arg(), + output.tensor().into_tensor_arg(), + plan.count_u32, + plan.n_fft, + log2_n, + threads_per_cube, + dim, + mode, + normalization, + ); + Ok(()) +} + +fn cfft_interleaved_four_step_launch( + client: &ComputeClient, + input: ComplexTensorBinding<'_, R>, + output: ComplexTensorBinding<'_, R>, + dim: usize, + mode: FftMode, + normalization: FftNormalization, + plan: CfftPlan, +) -> Result<(), FftError> { + let max_n = max_shared_fft_n(client); + let max_four_step_n = max_n.saturating_mul(max_n); + if plan.n_fft > max_four_step_n { + return Err(FftError::InvalidFftLength { n_fft: plan.n_fft }); + } + let (n1, n2) = factor_four_step(plan.n_fft, max_n); + let total = plan + .count + .checked_mul(plan.n_fft) + .ok_or(FftError::SizeOverflow)?; + let total_u32 = u32::try_from(total).map_err(|_| FftError::SizeOverflow)?; + let num_radix1_cubes = plan.count.checked_mul(n2).ok_or(FftError::SizeOverflow)?; + let num_radix1_cubes_u32 = + u32::try_from(num_radix1_cubes).map_err(|_| FftError::SizeOverflow)?; + let num_radix2_cubes = plan.count.checked_mul(n1).ok_or(FftError::SizeOverflow)?; + let num_radix2_cubes_u32 = + u32::try_from(num_radix2_cubes).map_err(|_| FftError::SizeOverflow)?; + let byte_len = total + .checked_mul(core::mem::size_of::()) + .ok_or(FftError::SizeOverflow)?; + let shape = input.shape().to_vec(); + let dtype = input.dtype(); + let scratch_re = + TensorHandle::::new_contiguous(shape.clone(), client.empty(byte_len), dtype); + let scratch_im = TensorHandle::::new_contiguous(shape, client.empty(byte_len), dtype); + let max_units = max_units_per_cube(client); + + { + let threads_per_cube = (n1 / 2).clamp(1, max_units); + let cube_dim = CubeDim::new_1d(threads_per_cube as u32); + let cube_count = + cubecl::calculate_cube_count_elemwise(client, num_radix1_cubes, CubeDim::new_single()); + cfft_interleaved_four_step_radix1_kernel::launch::( + client, + cube_count, + cube_dim, + input.tensor().into_tensor_arg(), + scratch_re.clone().binding().into_tensor_arg(), + scratch_im.clone().binding().into_tensor_arg(), + num_radix1_cubes_u32, + n1, + n2, + n1.trailing_zeros() as usize, + threads_per_cube, + dim, + mode, + ); + } + + { + let threads_per_cube = (n2 / 2).clamp(1, max_units); + let cube_dim = CubeDim::new_1d(threads_per_cube as u32); + let cube_count = + cubecl::calculate_cube_count_elemwise(client, num_radix2_cubes, CubeDim::new_single()); + cfft_four_step_radix2_kernel::launch::( + client, + cube_count, + cube_dim, + scratch_re.clone().binding().into_tensor_arg(), + scratch_im.clone().binding().into_tensor_arg(), + num_radix2_cubes_u32, + n1, + n2, + n2.trailing_zeros() as usize, + threads_per_cube, + dim, + mode, + ); + } + + let cube_dim = CubeDim::new_1d(256); + let cube_count = cubecl::calculate_cube_count_elemwise(client, total, cube_dim); + cfft_interleaved_four_step_transpose_kernel::launch::( + client, + cube_count, + cube_dim, + scratch_re.binding().into_tensor_arg(), + scratch_im.binding().into_tensor_arg(), + output.tensor().into_tensor_arg(), + total_u32, + n1, + n2, + dim, + normalization, + ); + Ok(()) +} + +fn ensure_non_overlapping_output_layout( + shape: &[usize], + strides: &[usize], +) -> Result<(), FftError> { + if shape.contains(&0) { + return Ok(()); + } + + let mut axes = shape + .iter() + .zip(strides) + .filter_map(|(extent, stride)| (*extent > 1).then_some((*stride, *extent))) + .collect::>(); + axes.sort_unstable_by_key(|(stride, _)| *stride); + + let mut span = 0usize; + for (stride, extent) in axes { + if stride <= span { + return Err(FftError::OverlappingBindings); + } + span = span + .checked_add( + (extent - 1) + .checked_mul(stride) + .ok_or(FftError::SizeOverflow)?, + ) + .ok_or(FftError::SizeOverflow)?; + } + Ok(()) +} + +fn validate_fft_shape(shape: &[usize], dim: usize) -> Result<(), FftError> { + if dim >= shape.len() { + return Err(FftError::AxisOutOfBounds { + dim, + rank: shape.len(), + }); + } + let n_fft = shape[dim]; + if n_fft < 2 || !n_fft.is_power_of_two() { + return Err(FftError::InvalidFftLength { n_fft }); + } + Ok(()) +} + +#[cube(launch)] +fn cfft_interleaved_shared_kernel( + input: &Tensor, + output: &mut Tensor, + num_windows: u32, + #[comptime] n_fft: usize, + #[comptime] log2_n: usize, + #[comptime] threads_per_cube: usize, + #[comptime] dim: usize, + #[comptime] mode: FftMode, + #[comptime] normalization: FftNormalization, +) { + let window_index = CUBE_POS; + if (window_index as u32) >= num_windows { + terminate!(); + } + + let input_re = input.view(InterleavedBatchSignalLayout::new( + input, + window_index, + dim, + 0usize, + )); + let input_im = input.view(InterleavedBatchSignalLayout::new( + input, + window_index, + dim, + 1usize, + )); + 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); + shared_re[j] = input_re.read_checked(i); + shared_im[j] = input_im.read_checked(i); + i += threads_per_cube; + } + sync_cube(); + + fft_butterfly_parallel::( + &mut shared_re, + &mut shared_im, + n_fft, + log2_n, + threads_per_cube, + mode, + ); + + let scale = match normalization { + FftNormalization::None => F::new(1.0_f32), + FftNormalization::ByN => F::new(1.0_f32) / F::cast_from(n_fft), + FftNormalization::Ortho => F::new(1.0_f32) / F::cast_from(n_fft).sqrt(), + }; + { + let mut output_re = output.view_mut(InterleavedBatchSignalLayout::new( + &*output, + window_index, + dim, + 0usize, + )); + let mut k = UNIT_POS as usize; + while k < n_fft { + output_re.write_checked(k, shared_re[k] * scale); + k += threads_per_cube; + } + } + + let mut output_im = output.view_mut(InterleavedBatchSignalLayout::new( + &*output, + window_index, + dim, + 1usize, + )); + let mut k = UNIT_POS as usize; + while k < n_fft { + output_im.write_checked(k, shared_im[k] * scale); + k += threads_per_cube; + } +} + +/// First four-step pass over the strided N1 dimension of each C32 window. +/// Reads interleaved scalar pairs and writes split scratch with the fused +/// Cooley-Tukey twiddle. +#[cube(launch)] +fn cfft_interleaved_four_step_radix1_kernel( + input: &Tensor, + scratch_re: &mut Tensor, + scratch_im: &mut Tensor, + num_cubes: u32, + #[comptime] n1: usize, + #[comptime] n2: usize, + #[comptime] log2_n1: usize, + #[comptime] threads_per_cube: usize, + #[comptime] dim: usize, + #[comptime] mode: FftMode, +) { + let cube_pos = CUBE_POS; + if cube_pos >= num_cubes as usize { + terminate!(); + } + + let window = cube_pos / n2; + let n2_idx = cube_pos - window * n2; + let input_re = input.view(InterleavedBatchSignalLayout::new( + input, window, dim, 0usize, + )); + let input_im = input.view(InterleavedBatchSignalLayout::new( + input, window, dim, 1usize, + )); + let mut scratch_re_view = scratch_re.view_mut(crate::layout::BatchSignalLayout::new( + &*scratch_re, + window, + dim, + )); + let mut scratch_im_view = scratch_im.view_mut(crate::layout::BatchSignalLayout::new( + &*scratch_im, + window, + dim, + )); + let mut shared_re = Shared::new_slice(n1); + let mut shared_im = Shared::new_slice(n1); + + let mut i = UNIT_POS as usize; + while i < n1 { + let j = bit_reverse(i, log2_n1); + let flat = i * n2 + n2_idx; + shared_re[j] = input_re.read_checked(flat); + shared_im[j] = input_im.read_checked(flat); + i += threads_per_cube; + } + sync_cube(); + + fft_butterfly_parallel::( + &mut shared_re, + &mut shared_im, + n1, + log2_n1, + threads_per_cube, + mode, + ); + + let sign = F::new(mode.sign()); + let n_total = comptime![n1 * n2]; + let two_pi = F::new(2.0 * PI); + let mut k1 = UNIT_POS as usize; + while k1 < n1 { + let theta = sign * two_pi * F::cast_from(k1 * n2_idx) / F::cast_from(n_total); + let w_re = theta.cos(); + let w_im = theta.sin(); + let ar = shared_re[k1]; + let ai = shared_im[k1]; + let flat = k1 * n2 + n2_idx; + scratch_re_view.write_checked(flat, w_re * ar - w_im * ai); + scratch_im_view.write_checked(flat, w_re * ai + w_im * ar); + k1 += threads_per_cube; + } +} + +/// Final four-step transpose writes adjacent interleaved output scalars and +/// applies the requested normalization as part of the global store. +#[cube(launch)] +fn cfft_interleaved_four_step_transpose_kernel( + scratch_re: &Tensor, + scratch_im: &Tensor, + output: &mut Tensor, + total: u32, + #[comptime] n1: usize, + #[comptime] n2: usize, + #[comptime] dim: usize, + #[comptime] normalization: FftNormalization, +) { + let pos = ABSOLUTE_POS; + if pos >= total as usize { + terminate!(); + } + + let n_fft = comptime![n1 * n2]; + let inner = pos % n_fft; + let window = pos / n_fft; + let scratch_re_view = scratch_re.view(crate::layout::BatchSignalLayout::new( + scratch_re, window, dim, + )); + let scratch_im_view = scratch_im.view(crate::layout::BatchSignalLayout::new( + scratch_im, window, dim, + )); + let k2 = inner / n1; + let k1 = inner - k2 * n1; + let src = k1 * n2 + k2; + let scale = match normalization { + FftNormalization::None => F::new(1.0_f32), + FftNormalization::ByN => F::new(1.0_f32) / F::cast_from(n_fft), + FftNormalization::Ortho => F::new(1.0_f32) / F::cast_from(n_fft).sqrt(), + }; + { + let mut output_re = output.view_mut(InterleavedBatchSignalLayout::new( + &*output, window, dim, 0usize, + )); + output_re.write_checked(inner, scratch_re_view.read_checked(src) * scale); + } + let mut output_im = output.view_mut(InterleavedBatchSignalLayout::new( + &*output, window, dim, 1usize, + )); + output_im.write_checked(inner, scratch_im_view.read_checked(src) * scale); +} diff --git a/crates/cubek-fft/src/fft/fft_inner.rs b/crates/cubek-fft/src/fft/fft_inner.rs index 9cd8f4be1..4be3f5eac 100644 --- a/crates/cubek-fft/src/fft/fft_inner.rs +++ b/crates/cubek-fft/src/fft/fft_inner.rs @@ -91,8 +91,8 @@ fn fft_butterfly_stages( let mut k = 0; while k < n { - let mut w_re = F::new(1.0); - let mut w_im = F::new(0.0); + let mut w_re = F::new(1.0_f32); + let mut w_im = F::new(0.0_f32); let mut j = 0; while j < half_m { diff --git a/crates/cubek-fft/src/fft/irfft.rs b/crates/cubek-fft/src/fft/irfft.rs index fdc9d765c..060178150 100644 --- a/crates/cubek-fft/src/fft/irfft.rs +++ b/crates/cubek-fft/src/fft/irfft.rs @@ -183,12 +183,16 @@ fn irfft_kernel( let src_bin = select(k < n_freq, k, n_fft - k); let active = src_bin < spec_bins as usize; let src_bin = select(active, src_bin, 0); - let im_sign = select(k < n_freq, F::new(1.0), F::new(-1.0)); - shared_re[dst] = select(active, spectrum_re_view.read_checked(src_bin), F::new(0.0)); + let im_sign = select(k < n_freq, F::new(1.0_f32), F::new(-1.0_f32)); + shared_re[dst] = select( + active, + spectrum_re_view.read_checked(src_bin), + F::new(0.0_f32), + ); shared_im[dst] = select( active, spectrum_im_view.read_checked(src_bin) * im_sign, - F::new(0.0), + F::new(0.0_f32), ); k += threads_per_cube; } @@ -203,7 +207,7 @@ fn irfft_kernel( FftMode::Inverse, ); - let scale = F::new(1.0) / F::cast_from(n_fft); + let scale = F::new(1.0_f32) / F::cast_from(n_fft); let mut i = UNIT_POS as usize; while i < n_fft { signal_view.write_checked(i, shared_re[i] * scale); diff --git a/crates/cubek-fft/src/fft/irfft_interleaved.rs b/crates/cubek-fft/src/fft/irfft_interleaved.rs new file mode 100644 index 000000000..6245b2893 --- /dev/null +++ b/crates/cubek-fft/src/fft/irfft_interleaved.rs @@ -0,0 +1,302 @@ +//! Inverse real-valued FFT from interleaved C32 input. + +use cubecl::prelude::*; +use cubecl::std::tensor::{ + AsView as _, AsViewExpand, AsViewMut as _, AsViewMutExpand, TensorHandle, +}; + +use crate::{ + ComplexTensorBinding, ComplexTensorHandle, FftError, FftNormalization, + complex::ensure_unique_output, + fft::{ + FftMode, + fft_parallel::{bit_reverse, fft_butterfly_parallel}, + limits::{ensure_packed_cfft_supported, max_shared_fft_n, max_units_per_cube}, + rfft_large::irfft_interleaved_large_launch, + }, + interleaved_layout::InterleavedBatchSignalLayout, + layout::BatchSignalLayout, +}; + +/// Runs an inverse real F32 FFT from an interleaved C32 half-spectrum. +pub fn irfft_interleaved( + spectrum: ComplexTensorHandle, + dim: usize, + normalization: FftNormalization, +) -> Result, FftError> { + let client = R::client(&Default::default()); + let spectrum_shape = spectrum.shape(); + let n_freq = *spectrum_shape.get(dim).ok_or(FftError::AxisOutOfBounds { + dim, + rank: spectrum_shape.len(), + })?; + let n_fft = n_freq + .checked_sub(1) + .and_then(|n| n.checked_mul(2)) + .ok_or(FftError::SizeOverflow)?; + ensure_packed_cfft_supported(&client, n_fft)?; + let mut signal_shape = spectrum_shape.to_vec(); + signal_shape[dim] = n_fft; + irfft_plan( + &spectrum.binding(), + &signal_shape, + spectrum.dtype(), + dim, + n_freq, + )?; + let elements = signal_shape.iter().try_fold(1usize, |total, extent| { + total.checked_mul(*extent).ok_or(FftError::SizeOverflow) + })?; + let byte_len = elements + .checked_mul(spectrum.dtype().size()) + .ok_or(FftError::SizeOverflow)?; + let signal = + TensorHandle::new_contiguous(signal_shape, client.empty(byte_len), spectrum.dtype()); + irfft_interleaved_launch_padded( + &client, + spectrum.binding(), + &signal, + dim, + n_freq, + normalization, + )?; + Ok(signal) +} + +/// Launches an inverse real F32 FFT into caller-provided real output. +pub fn irfft_interleaved_launch( + client: &ComputeClient, + spectrum: ComplexTensorBinding<'_, R>, + signal: &TensorHandle, + dim: usize, + normalization: FftNormalization, +) -> Result<(), FftError> { + let spec_bins = *spectrum.shape().get(dim).ok_or(FftError::AxisOutOfBounds { + dim, + rank: spectrum.shape().len(), + })?; + irfft_interleaved_launch_padded(client, spectrum, signal, dim, spec_bins, normalization) +} + +/// Launches an interleaved IRFFT while treating bins at `spec_bins..n_freq` as zero. +pub fn irfft_interleaved_launch_padded( + client: &ComputeClient, + spectrum: ComplexTensorBinding<'_, R>, + signal: &TensorHandle, + dim: usize, + spec_bins: usize, + normalization: FftNormalization, +) -> Result<(), FftError> { + let plan = irfft_plan(&spectrum, signal.shape(), signal.dtype, dim, spec_bins)?; + ensure_packed_cfft_supported(client, plan.n_fft)?; + ensure_non_overlapping_output_layout(signal.shape(), signal.strides())?; + ensure_unique_output(signal)?; + if plan.count == 0 { + return Ok(()); + } + if plan.n_fft > max_shared_fft_n(client) { + return irfft_interleaved_large_launch( + client, + spectrum, + signal, + dim, + spec_bins, + normalization, + plan.n_fft, + plan.count, + ); + } + + let log2_n = plan.n_fft.trailing_zeros() as usize; + let threads_per_cube = (plan.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, plan.count, CubeDim::new_single()); + irfft_interleaved_kernel::launch::( + client, + cube_count, + cube_dim, + spectrum.tensor().into_tensor_arg(), + signal.clone().binding().into_tensor_arg(), + plan.count_u32, + spec_bins as u32, + plan.n_fft, + log2_n, + threads_per_cube, + dim, + normalization, + ); + Ok(()) +} + +struct IrfftPlan { + n_fft: usize, + count: usize, + count_u32: u32, +} + +fn irfft_plan( + spectrum: &ComplexTensorBinding<'_, R>, + signal_shape: &[usize], + signal_dtype: StorageType, + dim: usize, + spec_bins: usize, +) -> Result { + if spectrum.dtype() != f32::as_type_native_unchecked().storage_type() { + return Err(FftError::UnsupportedDtype { + actual: spectrum.dtype(), + }); + } + if signal_dtype != f32::as_type_native_unchecked().storage_type() { + return Err(FftError::UnsupportedDtype { + actual: signal_dtype, + }); + } + if signal_shape.len() != spectrum.shape().len() { + return Err(FftError::ShapeMismatch { + name: "signal", + actual: signal_shape.to_vec(), + expected: spectrum.shape().to_vec(), + }); + } + let n_fft = *signal_shape.get(dim).ok_or(FftError::AxisOutOfBounds { + dim, + rank: signal_shape.len(), + })?; + if n_fft < 2 || !n_fft.is_power_of_two() { + return Err(FftError::InvalidFftLength { n_fft }); + } + let n_freq = n_fft / 2 + 1; + let mut expected_spectrum_shape = signal_shape.to_vec(); + expected_spectrum_shape[dim] = spectrum.shape()[dim]; + if spectrum.shape() != expected_spectrum_shape { + return Err(FftError::ShapeMismatch { + name: "spectrum", + actual: spectrum.shape().to_vec(), + expected: expected_spectrum_shape, + }); + } + if spec_bins == 0 || spec_bins > spectrum.shape()[dim] || spec_bins > n_freq { + return Err(FftError::InvalidLength { + name: "spec_bins", + value: spec_bins, + min: 1, + max: spectrum.shape()[dim].min(n_freq), + }); + } + let count = signal_shape + .iter() + .enumerate() + .filter(|(axis, _)| *axis != dim) + .try_fold(1usize, |count, (_, extent)| { + count.checked_mul(*extent).ok_or(FftError::SizeOverflow) + })?; + let count_u32 = u32::try_from(count).map_err(|_| FftError::SizeOverflow)?; + Ok(IrfftPlan { + n_fft, + count, + count_u32, + }) +} + +fn ensure_non_overlapping_output_layout( + shape: &[usize], + strides: &[usize], +) -> Result<(), FftError> { + if shape.contains(&0) { + return Ok(()); + } + let mut axes = shape + .iter() + .zip(strides) + .filter_map(|(extent, stride)| (*extent > 1).then_some((*stride, *extent))) + .collect::>(); + axes.sort_unstable_by_key(|(stride, _)| *stride); + + let mut span = 0usize; + for (stride, extent) in axes { + if stride <= span { + return Err(FftError::OverlappingBindings); + } + span = span + .checked_add( + (extent - 1) + .checked_mul(stride) + .ok_or(FftError::SizeOverflow)?, + ) + .ok_or(FftError::SizeOverflow)?; + } + Ok(()) +} + +#[cube(launch)] +fn irfft_interleaved_kernel( + spectrum: &Tensor, + signal: &mut Tensor, + num_windows: u32, + spec_bins: u32, + #[comptime] n_fft: usize, + #[comptime] log2_n: usize, + #[comptime] threads_per_cube: usize, + #[comptime] dim: usize, + #[comptime] normalization: FftNormalization, +) { + let window_index = CUBE_POS; + if (window_index as u32) >= num_windows { + terminate!(); + } + + let spectrum_re = spectrum.view(InterleavedBatchSignalLayout::new( + spectrum, + window_index, + dim, + 0usize, + )); + let spectrum_im = spectrum.view(InterleavedBatchSignalLayout::new( + spectrum, + window_index, + dim, + 1usize, + )); + let mut signal_view = signal.view_mut(BatchSignalLayout::new(&*signal, window_index, dim)); + let mut shared_re = Shared::new_slice(n_fft); + let mut shared_im = Shared::new_slice(n_fft); + let n_freq = comptime![n_fft / 2 + 1]; + + let mut k = UNIT_POS as usize; + while k < n_fft { + let dst = bit_reverse(k, log2_n); + let src_bin = select(k < n_freq, k, n_fft - k); + let active = src_bin < spec_bins as usize; + let src_bin = select(active, src_bin, 0); + let im_sign = select(k < n_freq, F::new(1.0_f32), F::new(-1.0_f32)); + shared_re[dst] = select(active, spectrum_re.read_checked(src_bin), F::new(0.0_f32)); + shared_im[dst] = select( + active, + spectrum_im.read_checked(src_bin) * im_sign, + F::new(0.0_f32), + ); + k += threads_per_cube; + } + sync_cube(); + + fft_butterfly_parallel::( + &mut shared_re, + &mut shared_im, + n_fft, + log2_n, + threads_per_cube, + FftMode::Inverse, + ); + + let scale = match normalization { + FftNormalization::None => F::new(1.0_f32), + FftNormalization::ByN => F::new(1.0_f32) / F::cast_from(n_fft), + FftNormalization::Ortho => F::new(1.0_f32) / F::cast_from(n_fft).sqrt(), + }; + let mut i = UNIT_POS as usize; + while i < n_fft { + signal_view.write_checked(i, shared_re[i] * scale); + i += threads_per_cube; + } +} diff --git a/crates/cubek-fft/src/fft/limits.rs b/crates/cubek-fft/src/fft/limits.rs index 43a06ee5a..b6496ded5 100644 --- a/crates/cubek-fft/src/fft/limits.rs +++ b/crates/cubek-fft/src/fft/limits.rs @@ -2,6 +2,8 @@ use cubecl::prelude::*; +use crate::FftError; + /// Largest power-of-two `n_fft` such that a shared-memory radix-2 butterfly /// over `f32` fits in this device's per-cube shared memory. /// @@ -21,6 +23,24 @@ pub(crate) fn max_units_per_cube(client: &ComputeClient) -> usize client.properties().hardware.max_units_per_cube as usize } +/// Reject real transforms whose packed CFFT cannot be factored into two +/// device-supported shared-memory FFTs. +pub(crate) fn ensure_packed_cfft_supported( + client: &ComputeClient, + n_fft: usize, +) -> Result<(), FftError> { + let max_shared = max_shared_fft_n(client); + let max_packed = max_shared.saturating_mul(max_shared); + if n_fft / 2 > max_packed { + Err(FftError::FftLengthExceedsDeviceLimit { + n_fft, + max_n_fft: max_packed.saturating_mul(2), + }) + } else { + Ok(()) + } +} + fn floor_power_of_two(n: usize) -> usize { assert!(n > 0, "device reports zero shared memory / units"); if n.is_power_of_two() { diff --git a/crates/cubek-fft/src/fft/mod.rs b/crates/cubek-fft/src/fft/mod.rs index 4746d8399..cad73d384 100644 --- a/crates/cubek-fft/src/fft/mod.rs +++ b/crates/cubek-fft/src/fft/mod.rs @@ -1,11 +1,17 @@ -mod cfft; +pub(crate) mod cfft; +mod cfft_interleaved; mod fft_inner; mod fft_parallel; mod irfft; +mod irfft_interleaved; mod limits; mod rfft; +mod rfft_interleaved; mod rfft_large; +pub use cfft_interleaved::*; pub use fft_inner::*; pub use irfft::*; +pub use irfft_interleaved::*; pub use rfft::*; +pub use rfft_interleaved::*; diff --git a/crates/cubek-fft/src/fft/rfft.rs b/crates/cubek-fft/src/fft/rfft.rs index 87448a015..e805ef6aa 100644 --- a/crates/cubek-fft/src/fft/rfft.rs +++ b/crates/cubek-fft/src/fft/rfft.rs @@ -191,8 +191,8 @@ fn rfft_kernel( 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)); - shared_im[j] = F::new(0.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(); diff --git a/crates/cubek-fft/src/fft/rfft_interleaved.rs b/crates/cubek-fft/src/fft/rfft_interleaved.rs new file mode 100644 index 000000000..7f62ff26c --- /dev/null +++ b/crates/cubek-fft/src/fft/rfft_interleaved.rs @@ -0,0 +1,322 @@ +//! Real-valued FFT into interleaved C32 output. + +use cubecl::prelude::*; +use cubecl::std::tensor::{ + AsView as _, AsViewExpand, AsViewMut as _, AsViewMutExpand, TensorHandle, +}; + +use crate::{ + ComplexTensorBinding, ComplexTensorHandle, FftError, FftNormalization, + fft::{ + FftMode, + fft_parallel::{bit_reverse, fft_butterfly_parallel}, + limits::{ensure_packed_cfft_supported, max_shared_fft_n, max_units_per_cube}, + rfft_large::rfft_interleaved_large_launch, + }, + interleaved_layout::InterleavedBatchSignalLayout, + layout::BatchSignalLayout, +}; + +/// Runs a real F32 FFT over `signal` and returns an interleaved C32 spectrum. +pub fn rfft_interleaved( + signal: TensorHandle, + dim: usize, + normalization: FftNormalization, +) -> Result, FftError> { + let shape = signal.shape().to_vec(); + let n_fft = validate_real_signal(signal.dtype, &shape, dim)?; + normalization.scale_f32(n_fft)?; + + let client = R::client(&Default::default()); + ensure_packed_cfft_supported(&client, n_fft)?; + let mut spectrum_shape = shape; + spectrum_shape[dim] = n_fft / 2 + 1; + let spectrum = ComplexTensorHandle::empty(&client, spectrum_shape, signal.dtype)?; + + rfft_interleaved_launch(&client, &signal, spectrum.binding(), dim, normalization)?; + Ok(spectrum) +} + +/// Launches a real F32 FFT into caller-provided interleaved C32 output. +pub fn rfft_interleaved_launch( + client: &ComputeClient, + signal: &TensorHandle, + spectrum: ComplexTensorBinding<'_, R>, + dim: usize, + normalization: FftNormalization, +) -> Result<(), FftError> { + let signal_len = signal + .shape() + .get(dim) + .copied() + .ok_or(FftError::AxisOutOfBounds { + dim, + rank: signal.shape().len(), + })?; + rfft_interleaved_launch_padded(client, signal, spectrum, dim, signal_len, normalization) +} + +/// Launches an interleaved RFFT while treating samples at `signal_len..n_fft` as zero. +pub fn rfft_interleaved_launch_padded( + client: &ComputeClient, + signal: &TensorHandle, + spectrum: ComplexTensorBinding<'_, R>, + dim: usize, + signal_len: usize, + normalization: FftNormalization, +) -> Result<(), FftError> { + let plan = rfft_plan(signal, &spectrum, dim, signal_len)?; + normalization.scale_f32(plan.n_fft)?; + ensure_packed_cfft_supported(client, plan.n_fft)?; + + spectrum.ensure_unique_output()?; + if plan.count == 0 { + return Ok(()); + } + if plan.n_fft > max_shared_fft_n(client) { + return rfft_interleaved_large_launch( + client, + signal, + spectrum, + dim, + signal_len, + normalization, + plan.n_fft, + plan.count, + ); + } + + let log2_n = plan.n_fft.trailing_zeros() as usize; + let threads_per_cube = (plan.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, plan.count, CubeDim::new_single()); + + rfft_interleaved_kernel::launch::( + client, + cube_count, + cube_dim, + signal.clone().binding().into_tensor_arg(), + spectrum.tensor().into_tensor_arg(), + plan.count_u32, + signal_len as u32, + plan.n_fft, + log2_n, + threads_per_cube, + dim, + normalization, + ); + Ok(()) +} + +struct RfftPlan { + n_fft: usize, + count: usize, + count_u32: u32, +} + +fn rfft_plan( + signal: &TensorHandle, + spectrum: &ComplexTensorBinding<'_, R>, + dim: usize, + signal_len: usize, +) -> Result { + let signal_shape = signal.shape(); + validate_signal_dtype_and_axis(signal.dtype, signal_shape, dim)?; + if spectrum.dtype() != f32::as_type_native_unchecked().storage_type() { + return Err(FftError::UnsupportedDtype { + actual: spectrum.dtype(), + }); + } + if spectrum.shape().len() != signal_shape.len() { + return Err(FftError::ShapeMismatch { + name: "spectrum", + actual: spectrum.shape().to_vec(), + expected: signal_shape.to_vec(), + }); + } + + let n_freq = spectrum.shape()[dim]; + if n_freq < 2 { + return Err(FftError::InvalidFftLength { n_fft: 0 }); + } + let n_fft = n_freq + .checked_sub(1) + .and_then(|n| n.checked_mul(2)) + .ok_or(FftError::SizeOverflow)?; + if n_fft < 2 || !n_fft.is_power_of_two() { + return Err(FftError::InvalidFftLength { n_fft }); + } + + let mut expected_shape = signal_shape.to_vec(); + expected_shape[dim] = n_freq; + if spectrum.shape() != expected_shape { + return Err(FftError::ShapeMismatch { + name: "spectrum", + actual: spectrum.shape().to_vec(), + expected: expected_shape, + }); + } + ensure_non_overlapping_output_layout(spectrum.shape(), spectrum.strides())?; + if signal_len > signal_shape[dim] { + return Err(FftError::InvalidLength { + name: "signal_len", + value: signal_len, + min: 0, + max: signal_shape[dim], + }); + } + if signal_len > n_fft { + return Err(FftError::InvalidLength { + name: "signal_len", + value: signal_len, + min: 0, + max: n_fft, + }); + } + + let count = signal_shape + .iter() + .enumerate() + .filter(|(axis, _)| *axis != dim) + .try_fold(1usize, |count, (_, extent)| { + count.checked_mul(*extent).ok_or(FftError::SizeOverflow) + })?; + let count_u32 = u32::try_from(count).map_err(|_| FftError::SizeOverflow)?; + Ok(RfftPlan { + n_fft, + count, + count_u32, + }) +} + +fn ensure_non_overlapping_output_layout( + shape: &[usize], + strides: &[usize], +) -> Result<(), FftError> { + if shape.contains(&0) { + return Ok(()); + } + + let mut axes = shape + .iter() + .zip(strides) + .filter_map(|(extent, stride)| (*extent > 1).then_some((*stride, *extent))) + .collect::>(); + axes.sort_unstable_by_key(|(stride, _)| *stride); + + let mut span = 0usize; + for (stride, extent) in axes { + if stride <= span { + return Err(FftError::OverlappingBindings); + } + span = span + .checked_add( + (extent - 1) + .checked_mul(stride) + .ok_or(FftError::SizeOverflow)?, + ) + .ok_or(FftError::SizeOverflow)?; + } + Ok(()) +} + +fn validate_real_signal( + dtype: StorageType, + shape: &[usize], + dim: usize, +) -> Result { + validate_signal_dtype_and_axis(dtype, shape, dim)?; + let n_fft = shape[dim]; + if n_fft < 2 || !n_fft.is_power_of_two() { + return Err(FftError::InvalidFftLength { n_fft }); + } + Ok(n_fft) +} + +fn validate_signal_dtype_and_axis( + dtype: StorageType, + shape: &[usize], + dim: usize, +) -> Result<(), FftError> { + if dtype != f32::as_type_native_unchecked().storage_type() { + return Err(FftError::UnsupportedDtype { actual: dtype }); + } + shape.get(dim).ok_or(FftError::AxisOutOfBounds { + dim, + rank: shape.len(), + })?; + Ok(()) +} + +#[cube(launch)] +fn rfft_interleaved_kernel( + signal: &Tensor, + spectrum: &mut Tensor, + num_windows: u32, + signal_len: u32, + #[comptime] n_fft: usize, + #[comptime] log2_n: usize, + #[comptime] threads_per_cube: usize, + #[comptime] dim: usize, + #[comptime] normalization: FftNormalization, +) { + 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 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::( + &mut shared_re, + &mut shared_im, + n_fft, + log2_n, + threads_per_cube, + FftMode::Forward, + ); + + let scale = match normalization { + FftNormalization::None => F::new(1.0_f32), + FftNormalization::ByN => F::new(1.0_f32) / F::cast_from(n_fft), + FftNormalization::Ortho => F::new(1.0_f32) / F::cast_from(n_fft).sqrt(), + }; + let n_freq = comptime![n_fft / 2 + 1]; + { + let mut spectrum_re = spectrum.view_mut(InterleavedBatchSignalLayout::new( + &*spectrum, + window_index, + dim, + 0usize, + )); + let mut k = UNIT_POS as usize; + while k < n_freq { + spectrum_re.write_checked(k, shared_re[k] * scale); + k += threads_per_cube; + } + } + let mut spectrum_im = spectrum.view_mut(InterleavedBatchSignalLayout::new( + &*spectrum, + window_index, + dim, + 1usize, + )); + let mut k = UNIT_POS as usize; + while k < n_freq { + spectrum_im.write_checked(k, shared_im[k] * scale); + k += threads_per_cube; + } +} diff --git a/crates/cubek-fft/src/fft/rfft_large.rs b/crates/cubek-fft/src/fft/rfft_large.rs index a2dfbc6f2..961988e1e 100644 --- a/crates/cubek-fft/src/fft/rfft_large.rs +++ b/crates/cubek-fft/src/fft/rfft_large.rs @@ -40,10 +40,12 @@ use cubecl::std::tensor::{ }; use crate::{ + ComplexTensorBinding, FftError, FftNormalization, fft::{ FftMode, cfft::{CfftBindings, cfft_launch_any_size}, }, + interleaved_layout::InterleavedBatchSignalLayout, layout::BatchSignalLayout, }; @@ -253,6 +255,181 @@ pub(crate) fn irfft_large_launch( Ok(()) } +/// Forward large-`n_fft` RFFT into an interleaved C32 half-spectrum. +/// +/// Packed CFFT buffers remain split; only the final post-processing pass writes +/// the interleaved real and imaginary component views. +#[allow(clippy::too_many_arguments)] +pub(crate) fn rfft_interleaved_large_launch( + client: &ComputeClient, + signal: &TensorHandle, + spectrum: ComplexTensorBinding<'_, R>, + dim: usize, + signal_len: usize, + normalization: FftNormalization, + n_fft: usize, + count: usize, +) -> Result<(), FftError> { + let m = n_fft / 2; + let packed_elems = count.checked_mul(m).ok_or(FftError::SizeOverflow)?; + let total_u32 = u32::try_from(packed_elems).map_err(|_| FftError::SizeOverflow)?; + let signal_len_u32 = u32::try_from(signal_len).map_err(|_| FftError::SizeOverflow)?; + let n_freq = m.checked_add(1).ok_or(FftError::SizeOverflow)?; + let post_total = count.checked_mul(n_freq).ok_or(FftError::SizeOverflow)?; + let post_total_u32 = u32::try_from(post_total).map_err(|_| FftError::SizeOverflow)?; + let byte_len = packed_elems + .checked_mul(signal.dtype.size()) + .ok_or(FftError::SizeOverflow)?; + let packed_shape = signal + .shape() + .iter() + .enumerate() + .map(|(axis, &extent)| if axis == dim { m } else { extent }) + .collect::>(); + let packed_re = TensorHandle::::new_contiguous( + packed_shape.clone(), + client.empty(byte_len), + signal.dtype, + ); + let packed_im = + TensorHandle::::new_contiguous(packed_shape, client.empty(byte_len), signal.dtype); + + let cube_dim = CubeDim::new_1d(256); + let cube_count = cubecl::calculate_cube_count_elemwise(client, packed_elems, cube_dim); + rfft_pack_kernel::launch::( + client, + cube_count, + cube_dim, + signal.clone().binding().into_tensor_arg(), + packed_re.clone().binding().into_tensor_arg(), + packed_im.clone().binding().into_tensor_arg(), + total_u32, + signal_len_u32, + m, + dim, + ); + + cfft_launch_any_size::( + client, + CfftBindings { + input_re: packed_re.clone().binding(), + input_im: packed_im.clone().binding(), + output_re: packed_re.clone().binding(), + output_im: packed_im.clone().binding(), + }, + dim, + signal.dtype, + FftMode::Forward, + )?; + + let cube_count = cubecl::calculate_cube_count_elemwise(client, post_total, cube_dim); + rfft_post_interleaved_kernel::launch::( + client, + cube_count, + cube_dim, + packed_re.binding().into_tensor_arg(), + packed_im.binding().into_tensor_arg(), + spectrum.tensor().into_tensor_arg(), + post_total_u32, + n_fft, + m, + dim, + normalization, + ); + Ok(()) +} + +/// Inverse large-`n_fft` RFFT from an interleaved C32 half-spectrum. +/// +/// The pre-process reads interleaved component views into split packed CFFT +/// buffers; the unpack fuses the public inverse-normalization adjustment. +#[allow(clippy::too_many_arguments)] +pub(crate) fn irfft_interleaved_large_launch( + client: &ComputeClient, + spectrum: ComplexTensorBinding<'_, R>, + signal: &TensorHandle, + dim: usize, + spec_bins: usize, + normalization: FftNormalization, + n_fft: usize, + count: usize, +) -> Result<(), FftError> { + let m = n_fft / 2; + let packed_elems = count.checked_mul(m).ok_or(FftError::SizeOverflow)?; + let total_u32 = u32::try_from(packed_elems).map_err(|_| FftError::SizeOverflow)?; + let spec_bins_u32 = u32::try_from(spec_bins).map_err(|_| FftError::SizeOverflow)?; + let byte_len = packed_elems + .checked_mul(signal.dtype.size()) + .ok_or(FftError::SizeOverflow)?; + let packed_shape = signal + .shape() + .iter() + .enumerate() + .map(|(axis, &extent)| if axis == dim { m } else { extent }) + .collect::>(); + let packed_in_re = TensorHandle::::new_contiguous( + packed_shape.clone(), + client.empty(byte_len), + signal.dtype, + ); + let packed_in_im = TensorHandle::::new_contiguous( + packed_shape.clone(), + client.empty(byte_len), + signal.dtype, + ); + let packed_out_re = TensorHandle::::new_contiguous( + packed_shape.clone(), + client.empty(byte_len), + signal.dtype, + ); + let packed_out_im = + TensorHandle::::new_contiguous(packed_shape, client.empty(byte_len), signal.dtype); + + let cube_dim = CubeDim::new_1d(256); + let cube_count = cubecl::calculate_cube_count_elemwise(client, packed_elems, cube_dim); + irfft_pre_interleaved_kernel::launch::( + client, + cube_count, + cube_dim, + spectrum.tensor().into_tensor_arg(), + packed_in_re.clone().binding().into_tensor_arg(), + packed_in_im.clone().binding().into_tensor_arg(), + total_u32, + spec_bins_u32, + n_fft, + m, + dim, + ); + + cfft_launch_any_size::( + client, + CfftBindings { + input_re: packed_in_re.binding(), + input_im: packed_in_im.binding(), + output_re: packed_out_re.clone().binding(), + output_im: packed_out_im.clone().binding(), + }, + dim, + signal.dtype, + FftMode::Inverse, + )?; + + let cube_count = cubecl::calculate_cube_count_elemwise(client, packed_elems, cube_dim); + irfft_unpack_interleaved_kernel::launch::( + client, + cube_count, + cube_dim, + packed_out_re.binding().into_tensor_arg(), + packed_out_im.binding().into_tensor_arg(), + signal.clone().binding().into_tensor_arg(), + total_u32, + m, + dim, + normalization, + ); + Ok(()) +} + // --- pack / post / pre / unpack kernels -------------------------------- /// `y[k] = x[2k] + i * x[2k+1]`, one thread per `k`. @@ -283,11 +460,11 @@ fn rfft_pack_kernel( let odd = select(odd_active, odd, 0); packed_re_view.write_checked( k, - select(even_active, signal_view.read_checked(even), F::new(0.0)), + select(even_active, signal_view.read_checked(even), F::new(0.0_f32)), ); packed_im_view.write_checked( k, - select(odd_active, signal_view.read_checked(odd), F::new(0.0)), + select(odd_active, signal_view.read_checked(odd), F::new(0.0_f32)), ); } @@ -330,12 +507,12 @@ fn rfft_post_kernel( let y0_re = packed_re_view.read_checked(0); let y0_im = packed_im_view.read_checked(0); spectrum_re_view.write_checked(k, y0_re + y0_im); - spectrum_im_view.write_checked(k, F::new(0.0)); + spectrum_im_view.write_checked(k, F::new(0.0_f32)); } else if k == m { let y0_re = packed_re_view.read_checked(0); let y0_im = packed_im_view.read_checked(0); spectrum_re_view.write_checked(k, y0_re - y0_im); - spectrum_im_view.write_checked(k, F::new(0.0)); + spectrum_im_view.write_checked(k, F::new(0.0_f32)); } else { let a_re = packed_re_view.read_checked(k); let a_im = packed_im_view.read_checked(k); @@ -353,15 +530,82 @@ fn rfft_post_kernel( // 1 - i*W = (1 + s) - i*c // 1 + i*W = (1 - s) + i*c // 2 X[k] = A*(1 - i*W) + B*(1 + i*W) - let one_plus_s = F::new(1.0) + s; - let one_minus_s = F::new(1.0) - s; - let x_re = F::new(0.5) * (a_re * one_plus_s + a_im * c + b_re * one_minus_s - b_im * c); - let x_im = F::new(0.5) * (a_im * one_plus_s - a_re * c + b_re * c + b_im * one_minus_s); + let one_plus_s = F::new(1.0_f32) + s; + let one_minus_s = F::new(1.0_f32) - s; + let x_re = F::new(0.5_f32) * (a_re * one_plus_s + a_im * c + b_re * one_minus_s - b_im * c); + let x_im = F::new(0.5_f32) * (a_im * one_plus_s - a_re * c + b_re * c + b_im * one_minus_s); spectrum_re_view.write_checked(k, x_re); spectrum_im_view.write_checked(k, x_im); } } +/// Interleaved-output variant of `rfft_post_kernel`. +#[cube(launch)] +fn rfft_post_interleaved_kernel( + packed_re: &Tensor, + packed_im: &Tensor, + spectrum: &mut Tensor, + total: u32, + #[comptime] n_fft: usize, + #[comptime] m: usize, + #[comptime] dim: usize, + #[comptime] normalization: FftNormalization, +) { + let pos = ABSOLUTE_POS; + if pos >= total as usize { + terminate!(); + } + let n_freq = comptime![m + 1]; + let k = pos % n_freq; + let window = pos / n_freq; + let packed_re_view = packed_re.view(BatchSignalLayout::new(packed_re, window, dim)); + let packed_im_view = packed_im.view(BatchSignalLayout::new(packed_im, window, dim)); + let scale = match normalization { + FftNormalization::None => F::new(1.0_f32), + FftNormalization::ByN => F::new(1.0_f32) / F::cast_from(n_fft), + FftNormalization::Ortho => F::new(1.0_f32) / F::cast_from(n_fft).sqrt(), + }; + + let (x_re, x_im) = if k == 0 { + let y0_re = packed_re_view.read_checked(0); + let y0_im = packed_im_view.read_checked(0); + (y0_re + y0_im, F::new(0.0_f32)) + } else if k == m { + let y0_re = packed_re_view.read_checked(0); + let y0_im = packed_im_view.read_checked(0); + (y0_re - y0_im, F::new(0.0_f32)) + } else { + let a_re = packed_re_view.read_checked(k); + let a_im = packed_im_view.read_checked(k); + let b_re = packed_re_view.read_checked(m - k); + let b_im = -packed_im_view.read_checked(m - k); + let theta = -F::new(2.0 * PI) * F::cast_from(k) / F::cast_from(n_fft); + let c = theta.cos(); + let s = theta.sin(); + ( + F::new(0.5_f32) + * (a_re * (F::new(1.0_f32) + s) + a_im * c + b_re * (F::new(1.0_f32) - s) + - b_im * c), + F::new(0.5_f32) + * (a_im * (F::new(1.0_f32) + s) - a_re * c + + b_re * c + + b_im * (F::new(1.0_f32) - s)), + ) + }; + { + let mut spectrum_re = spectrum.view_mut(InterleavedBatchSignalLayout::new( + &*spectrum, window, dim, 0usize, + )); + spectrum_re.write_checked(k, x_re * scale); + } + { + let mut spectrum_im = spectrum.view_mut(InterleavedBatchSignalLayout::new( + &*spectrum, window, dim, 1usize, + )); + spectrum_im.write_checked(k, x_im * scale); + } +} + /// Build packed `Y[0..M]` from half-spectrum `X[0..N/2+1]` for the /// packed-real inverse path. Inverse of `rfft_post_kernel`. /// @@ -400,26 +644,30 @@ fn irfft_pre_kernel( let has_nyquist = m < spec_bins as usize; let x0_re = spectrum_re_view.read_checked(0); let xm = select(has_nyquist, m, 0); - let xm_re = select(has_nyquist, spectrum_re_view.read_checked(xm), F::new(0.0)); - packed_re_view.write_checked(k, F::new(0.5) * (x0_re + xm_re)); - packed_im_view.write_checked(k, F::new(0.5) * (x0_re - xm_re)); + let xm_re = select( + has_nyquist, + spectrum_re_view.read_checked(xm), + F::new(0.0_f32), + ); + packed_re_view.write_checked(k, F::new(0.5_f32) * (x0_re + xm_re)); + packed_im_view.write_checked(k, F::new(0.5_f32) * (x0_re - xm_re)); } else { let active = k < spec_bins as usize; let src = select(active, k, 0); - let x_re = select(active, spectrum_re_view.read_checked(src), F::new(0.0)); - let x_im = select(active, spectrum_im_view.read_checked(src), F::new(0.0)); + let x_re = select(active, spectrum_re_view.read_checked(src), F::new(0.0_f32)); + let x_im = select(active, spectrum_im_view.read_checked(src), F::new(0.0_f32)); let mirror = m - k; let mirror_active = mirror < spec_bins as usize; let mirror = select(mirror_active, mirror, 0); let xm_re = select( mirror_active, spectrum_re_view.read_checked(mirror), - F::new(0.0), + F::new(0.0_f32), ); let xm_im_raw = select( mirror_active, spectrum_im_view.read_checked(mirror), - F::new(0.0), + F::new(0.0_f32), ); let xm_im = -xm_im_raw; // conj(X[M-k]) @@ -437,10 +685,76 @@ fn irfft_pre_kernel( // Re = xm_re*(1+s) + xm_im*c (xm_im is already negated here) // Im = -xm_re*c + xm_im*(1+s) // 2 Y[k] = A*(1 + i*W) + B*(1 - i*W). - let one_plus_s = F::new(1.0) + s; - let one_minus_s = F::new(1.0) - s; - let y_re = F::new(0.5) * (x_re * one_minus_s - x_im * c + xm_re * one_plus_s + xm_im * c); - let y_im = F::new(0.5) * (x_im * one_minus_s + x_re * c - xm_re * c + xm_im * one_plus_s); + let one_plus_s = F::new(1.0_f32) + s; + let one_minus_s = F::new(1.0_f32) - s; + let y_re = + F::new(0.5_f32) * (x_re * one_minus_s - x_im * c + xm_re * one_plus_s + xm_im * c); + let y_im = + F::new(0.5_f32) * (x_im * one_minus_s + x_re * c - xm_re * c + xm_im * one_plus_s); + packed_re_view.write_checked(k, y_re); + packed_im_view.write_checked(k, y_im); + } +} + +/// Interleaved-input variant of `irfft_pre_kernel`. +#[cube(launch)] +fn irfft_pre_interleaved_kernel( + spectrum: &Tensor, + packed_re: &mut Tensor, + packed_im: &mut Tensor, + total: u32, + spec_bins: u32, + #[comptime] n_fft: usize, + #[comptime] m: usize, + #[comptime] dim: usize, +) { + let pos = ABSOLUTE_POS; + if pos >= total as usize { + terminate!(); + } + let k = pos % m; + let window = pos / m; + let spectrum_re = spectrum.view(InterleavedBatchSignalLayout::new( + spectrum, window, dim, 0usize, + )); + let spectrum_im = spectrum.view(InterleavedBatchSignalLayout::new( + spectrum, window, dim, 1usize, + )); + let mut packed_re_view = packed_re.view_mut(BatchSignalLayout::new(&*packed_re, window, dim)); + let mut packed_im_view = packed_im.view_mut(BatchSignalLayout::new(&*packed_im, window, dim)); + + if k == 0 { + let has_nyquist = m < spec_bins as usize; + let x0_re = spectrum_re.read_checked(0); + let xm = select(has_nyquist, m, 0); + let xm_re = select(has_nyquist, spectrum_re.read_checked(xm), F::new(0.0_f32)); + packed_re_view.write_checked(k, F::new(0.5_f32) * (x0_re + xm_re)); + packed_im_view.write_checked(k, F::new(0.5_f32) * (x0_re - xm_re)); + } else { + let active = k < spec_bins as usize; + let src = select(active, k, 0); + let x_re = select(active, spectrum_re.read_checked(src), F::new(0.0_f32)); + let x_im = select(active, spectrum_im.read_checked(src), F::new(0.0_f32)); + let mirror = m - k; + let mirror_active = mirror < spec_bins as usize; + let mirror = select(mirror_active, mirror, 0); + let xm_re = select( + mirror_active, + spectrum_re.read_checked(mirror), + F::new(0.0_f32), + ); + let xm_im = -select( + mirror_active, + spectrum_im.read_checked(mirror), + F::new(0.0_f32), + ); + let theta = F::new(2.0 * PI) * F::cast_from(k) / F::cast_from(n_fft); + let c = theta.cos(); + let s = theta.sin(); + let y_re = F::new(0.5_f32) + * (x_re * (F::new(1.0_f32) - s) - x_im * c + xm_re * (F::new(1.0_f32) + s) + xm_im * c); + let y_im = F::new(0.5_f32) + * (x_im * (F::new(1.0_f32) - s) + x_re * c - xm_re * c + xm_im * (F::new(1.0_f32) + s)); packed_re_view.write_checked(k, y_re); packed_im_view.write_checked(k, y_im); } @@ -466,7 +780,38 @@ fn irfft_unpack_kernel( let packed_re_view = packed_re.view(BatchSignalLayout::new(packed_re, window, dim)); let packed_im_view = packed_im.view(BatchSignalLayout::new(packed_im, window, dim)); let mut signal_view = signal.view_mut(BatchSignalLayout::new(&*signal, window, dim)); - let scale = F::new(1.0) / F::cast_from(m); + let scale = F::new(1.0_f32) / F::cast_from(m); + signal_view.write_checked(2 * k, packed_re_view.read_checked(k) * scale); + signal_view.write_checked(2 * k + 1, packed_im_view.read_checked(k) * scale); +} + +/// Interleaved IRFFT unpack with the public normalization fused into stores. +#[cube(launch)] +fn irfft_unpack_interleaved_kernel( + packed_re: &Tensor, + packed_im: &Tensor, + signal: &mut Tensor, + total: u32, + #[comptime] m: usize, + #[comptime] dim: usize, + #[comptime] normalization: FftNormalization, +) { + let pos = ABSOLUTE_POS; + if pos >= total as usize { + terminate!(); + } + let k = pos % m; + let window = pos / m; + let packed_re_view = packed_re.view(BatchSignalLayout::new(packed_re, window, dim)); + let packed_im_view = packed_im.view(BatchSignalLayout::new(packed_im, window, dim)); + let mut signal_view = signal.view_mut(BatchSignalLayout::new(&*signal, window, dim)); + let n_fft = comptime![2 * m]; + let adjustment = match normalization { + FftNormalization::None => F::cast_from(n_fft), + FftNormalization::ByN => F::new(1.0_f32), + FftNormalization::Ortho => F::cast_from(n_fft).sqrt(), + }; + let scale = adjustment / F::cast_from(m); signal_view.write_checked(2 * k, packed_re_view.read_checked(k) * scale); signal_view.write_checked(2 * k + 1, packed_im_view.read_checked(k) * scale); } diff --git a/crates/cubek-fft/src/interleaved_layout.rs b/crates/cubek-fft/src/interleaved_layout.rs new file mode 100644 index 000000000..f78ea523d --- /dev/null +++ b/crates/cubek-fft/src/interleaved_layout.rs @@ -0,0 +1,66 @@ +use cubecl::{ + prelude::*, + std::tensor::layout::{Coords1d, Layout, LayoutExpand}, +}; + +/// A one-dimensional component view over an interleaved C32 tensor window. +#[derive(CubeType, Clone, Copy)] +pub(crate) struct InterleavedBatchSignalLayout { + num_samples: usize, + stride_samples: usize, + batch_offset: usize, + component: usize, +} + +#[cube] +impl InterleavedBatchSignalLayout { + pub fn new( + tensor: &Tensor, + batch_index: usize, + #[comptime] dim: usize, + #[comptime] component: usize, + ) -> Self { + let rank = tensor.rank(); + let mut batch_offset = 0; + let mut temp_idx = batch_index; + + for axis in 0..rank { + if axis != dim { + let size = tensor.shape(axis); + let stride = tensor.stride(axis); + let coord = temp_idx % size; + batch_offset += coord * stride; + temp_idx /= size; + } + } + + InterleavedBatchSignalLayout { + num_samples: tensor.shape(dim), + stride_samples: tensor.stride(dim), + batch_offset, + component, + } + } +} + +#[cube] +impl Layout for InterleavedBatchSignalLayout { + type Coordinates = Coords1d; + type SourceCoordinates = Coords1d; + + fn to_source_pos(&self, coords: Self::Coordinates) -> usize { + self.batch_offset + coords * self.stride_samples + self.component + } + + fn to_source_pos_checked(&self, coords: Self::Coordinates) -> (usize, bool) { + (self.to_source_pos(coords), self.is_in_bounds(coords)) + } + + fn shape(&self) -> Self::Coordinates { + self.num_samples + } + + fn is_in_bounds(&self, pos: Self::Coordinates) -> bool { + pos < self.num_samples + } +} diff --git a/crates/cubek-fft/src/lib.rs b/crates/cubek-fft/src/lib.rs index d6a74912a..ffb015563 100644 --- a/crates/cubek-fft/src/lib.rs +++ b/crates/cubek-fft/src/lib.rs @@ -1,7 +1,55 @@ +//! FFT primitives for CubeCL runtimes. +//! +//! # Interleaved C32 ABI +//! +//! The interleaved entry points accept [`ComplexTensorHandle`] values with one +//! logical C32 element per complex value. Physically, each logical element is +//! stored as two adjacent F32 scalars in `[re, im]` order. Shapes and the +//! strides accepted by [`ComplexTensorHandle::new_strided`] are therefore in +//! logical complex-element units; [`ComplexTensorHandle::scalar_strides`] is +//! the corresponding physical F32-scalar stride (`2 * logical_stride`). +//! `num_complex::Complex32` is `#[repr(C)]` with adjacent `re: f32` and +//! `im: f32` fields, so a contiguous slice has the same bytes as this ABI +//! when its starting address satisfies `align_of::()` (currently +//! the same four-byte alignment enforced for F32 offsets). This describes +//! layout compatibility only; callers must still use a sound ownership-aware +//! cast or copy when moving between Rust slices and runtime buffers. +//! +//! RFFT inputs and IRFFT outputs retain their ordinary real logical shape. +//! The interleaved RFFT output and IRFFT input use the same logical shape with +//! the transformed axis shortened to `n_fft / 2 + 1`. Caller-provided outputs +//! must be unique, non-overlapping allocations; aliased bindings and +//! overlapping output layouts are rejected. +//! +//! | [`FftNormalization`] | Scale applied to every transform direction | +//! | --- | --- | +//! | `None` | `1` | +//! | `ByN` | `1 / n_fft` | +//! | `Ortho` | `1 / sqrt(n_fft)` | +//! +//! Interleaved CFFT chooses the shared-memory small FFT path up to the active +//! device's shared-size limit and the four-step path above it. The allocating +//! helpers use the runtime's default device; use the `*_launch` functions with +//! an explicit `ComputeClient` to select a device. This release supports only +//! F32 real tensors and C32 interleaved tensors. F64/C64 support is deferred. +//! +//! Profiling an interleaved launch should show no standalone interleaved +//! conversion or scaling kernel. Large real transforms still use their +//! algorithmic real-to-packed-CFFT and packed-CFFT-to-real stages internally; +//! their external complex boundary reads or writes `[re, im]` directly and +//! final scaling is fused into the existing store. + +mod complex; +mod error; mod fft; +mod interleaved_layout; mod layout; +mod normalization; +pub use complex::{ComplexTensorBinding, ComplexTensorHandle}; +pub use error::FftError; pub use fft::*; +pub use normalization::FftNormalization; #[cfg(any(feature = "cpu-reference", feature = "benchmarks"))] pub mod eval; diff --git a/crates/cubek-fft/src/normalization.rs b/crates/cubek-fft/src/normalization.rs new file mode 100644 index 000000000..200a93f1e --- /dev/null +++ b/crates/cubek-fft/src/normalization.rs @@ -0,0 +1,21 @@ +use crate::FftError; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum FftNormalization { + None, + ByN, + Ortho, +} + +impl FftNormalization { + pub fn scale_f32(self, n_fft: usize) -> Result { + if n_fft < 2 || !n_fft.is_power_of_two() { + return Err(FftError::InvalidFftLength { n_fft }); + } + Ok(match self { + Self::None => 1.0, + Self::ByN => 1.0 / n_fft as f32, + Self::Ortho => 1.0 / (n_fft as f32).sqrt(), + }) + } +} diff --git a/crates/cubek-fft/tests/fft/bench_catalog.rs b/crates/cubek-fft/tests/fft/bench_catalog.rs index c74a9a88b..5018da7a9 100644 --- a/crates/cubek-fft/tests/fft/bench_catalog.rs +++ b/crates/cubek-fft/tests/fft/bench_catalog.rs @@ -1,7 +1,9 @@ //! Correctness over the FFT benchmark catalogue. #![cfg(feature = "benchmarks")] -use cubek_fft::eval::benchmarks::{FftCorrectness, FftProblem, FftStrategy}; +use cubek_fft::eval::benchmarks::{ + CfftCorrectness, CfftProblem, FftCorrectness, FftProblem, FftStrategy, +}; use cubek_test_utils::{CatalogEntry, Correctness, TestOutcome, assert_equals_approx}; const SEEDS: [u64; 2] = [12, 34]; @@ -9,6 +11,32 @@ const SEEDS: [u64; 2] = [12, 34]; /// FFT is f32 round-trip, so we can be tighter than the matmul side. const FFT_EPS: f32 = 1e-3; +#[test] +fn bench_catalog_exposes_split_and_interleaved_strategies_for_large_problems() { + use cubek_fft::eval::benchmarks::{cfft_problems, problems, strategies}; + + let strategy_ids = strategies() + .into_iter() + .map(|entry| entry.id) + .collect::>(); + assert_eq!(strategy_ids, ["default", "interleaved"]); + + let problem_ids = problems() + .into_iter() + .map(|entry| entry.id) + .collect::>(); + assert!(problem_ids.contains(&"forward_1x4096".to_string())); + assert!(problem_ids.contains(&"inverse_1x8192".to_string())); + assert!(problem_ids.iter().all(|id| !id.starts_with("cfft_"))); + + let cfft_problem_ids = cfft_problems() + .into_iter() + .map(|entry| entry.id) + .collect::>(); + assert!(cfft_problem_ids.contains(&"forward_1x4096".to_string())); + assert!(cfft_problem_ids.contains(&"inverse_1x8192".to_string())); +} + fn lookup(entries: Vec>, id: &str) -> T { entries .into_iter() @@ -36,6 +64,23 @@ fn run(strategy_id: &str, problem_id: &str) { .enforce(); } +fn run_cfft(strategy_id: &str, problem_id: &str) { + use cubek_fft::eval::benchmarks::{cfft_problems, strategies}; + + let strategy: FftStrategy = lookup(strategies(), strategy_id); + let problem: CfftProblem = lookup(cfft_problems(), problem_id); + let actual = match CfftCorrectness.kernel_result(&strategy, &problem, &SEEDS) { + Ok(host) => host, + Err(e) => return TestOutcome::CompileError(e).enforce(), + }; + let expected = CfftCorrectness + .reference_result(&problem, &SEEDS, None) + .unwrap_or_else(|e| panic!("CFFT reference failed for {problem_id}: {e}")); + assert_equals_approx(&actual, &expected, FFT_EPS) + .as_test_outcome() + .enforce(); +} + #[test] fn forward_5x2x2048_default() { run("default", "forward_5x2x2048"); @@ -55,3 +100,33 @@ fn forward_1x4096_default() { fn forward_1x16384_default() { run("default", "forward_1x16384"); } + +#[test] +fn forward_1x4096_interleaved() { + run("interleaved", "forward_1x4096"); +} + +#[test] +fn inverse_5x2x2048_interleaved() { + run("interleaved", "inverse_5x2x2048"); +} + +#[test] +fn cfft_forward_1x4096_default() { + run_cfft("default", "forward_1x4096"); +} + +#[test] +fn cfft_forward_1x4096_interleaved() { + run_cfft("interleaved", "forward_1x4096"); +} + +#[test] +fn cfft_forward_1x8192_default() { + run_cfft("default", "forward_1x8192"); +} + +#[test] +fn cfft_forward_1x8192_interleaved() { + run_cfft("interleaved", "forward_1x8192"); +} diff --git a/crates/cubek-fft/tests/fft/interleaved_cfft.rs b/crates/cubek-fft/tests/fft/interleaved_cfft.rs new file mode 100644 index 000000000..5afcd473a --- /dev/null +++ b/crates/cubek-fft/tests/fft/interleaved_cfft.rs @@ -0,0 +1,305 @@ +use cubecl::{CubeElement, Runtime, TestRuntime, client::ComputeClient, frontend::CubePrimitive}; +use cubek_fft::{ + ComplexTensorHandle, FftError, FftMode, FftNormalization, cfft_interleaved, + cfft_interleaved_launch, +}; + +fn contiguous_c32( + client: &ComputeClient, + shape: Vec, + values: &[f32], +) -> ComplexTensorHandle { + let dtype = f32::as_type_native_unchecked().storage_type(); + ComplexTensorHandle::new_contiguous( + shape, + client.create_from_slice(f32::as_bytes(values)), + dtype, + ) + .unwrap() +} + +fn scalar_buffer( + client: &ComputeClient, + tensor: ComplexTensorHandle, +) -> Vec { + let raw = tensor.into_raw_parts(); + f32::from_bytes(&client.read_one(raw.handle).unwrap()).to_vec() +} + +fn assert_complex_scalars_approx( + client: &ComputeClient, + tensor: ComplexTensorHandle, + expected: &[f32], + epsilon: f32, +) { + let actual = scalar_buffer(client, tensor); + assert_eq!(actual.len(), expected.len()); + for (index, (actual, expected)) in actual.iter().zip(expected).enumerate() { + assert!( + (actual - expected).abs() <= epsilon, + "scalar {index}: got {actual}, expected {expected}" + ); + } +} + +fn assert_logical_c32_approx( + client: &ComputeClient, + tensor: ComplexTensorHandle, + expected: &[f32], + epsilon: f32, +) { + let shape = tensor.shape().to_vec(); + let scalar_strides = tensor.scalar_strides().to_vec(); + let scalars = scalar_buffer(client, tensor); + for logical in 0..shape.iter().product::() { + let mut remaining = logical; + let mut scalar_index = 0; + for axis in (0..shape.len()).rev() { + let coord = remaining % shape[axis]; + remaining /= shape[axis]; + scalar_index += coord * scalar_strides[axis]; + } + for component in 0..2 { + let actual = scalars[scalar_index + component]; + let expected = expected[logical * 2 + component]; + assert!( + (actual - expected).abs() <= epsilon, + "logical {logical}, component {component}: got {actual}, expected {expected}" + ); + } + } +} + +fn values_for(shape: &[usize]) -> Vec { + (0..shape.iter().product::()) + .flat_map(|i| [i as f32 + 0.25, -(i as f32) + 0.5]) + .collect() +} + +fn run_round_trip( + client: &ComputeClient, + shape: Vec, + dim: usize, + normalization: FftNormalization, + epsilon: f32, +) { + let values = values_for(&shape); + let input = contiguous_c32(client, shape, &values); + let forward_normalization = match normalization { + FftNormalization::ByN => FftNormalization::None, + normalization => normalization, + }; + let spectrum = cfft_interleaved(input, dim, FftMode::Forward, forward_normalization).unwrap(); + let inverse_normalization = match normalization { + FftNormalization::None => FftNormalization::ByN, + normalization => normalization, + }; + let result = cfft_interleaved(spectrum, dim, FftMode::Inverse, inverse_normalization).unwrap(); + assert_complex_scalars_approx(client, result, &values, epsilon); +} + +fn round_trip(shape: Vec, dim: usize, normalization: FftNormalization) { + let client = ::client(&Default::default()); + run_round_trip(&client, shape, dim, normalization, 1e-4); +} + +fn test_max_shared_fft_n(client: &ComputeClient) -> usize { + let max_elems = + client.properties().hardware.max_shared_memory_size / (2 * core::mem::size_of::()); + if max_elems.is_power_of_two() { + max_elems + } else { + max_elems.next_power_of_two() >> 1 + } +} + +#[cfg(feature = "heavy")] +fn first_four_step_n(client: &ComputeClient) -> usize { + 2 * test_max_shared_fft_n(client) +} + +#[test] +fn cfft_interleaved_small_round_trip_preserves_c32_order() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let values: Vec = (0..8) + .flat_map(|i| [i as f32 + 0.25, -(i as f32)]) + .collect(); + let input = ComplexTensorHandle::new_contiguous( + vec![1, 8], + client.create_from_slice(f32::as_bytes(&values)), + dtype, + ) + .unwrap(); + let spectrum = cfft_interleaved(input, 1, FftMode::Forward, FftNormalization::None).unwrap(); + let result = cfft_interleaved(spectrum, 1, FftMode::Inverse, FftNormalization::ByN).unwrap(); + assert_complex_scalars_approx(&client, result, &values, 1e-4); +} + +#[test] +fn cfft_interleaved_supports_axis_zero() { + round_trip(vec![8, 2], 0, FftNormalization::None); +} + +#[test] +fn cfft_interleaved_supports_middle_axis() { + round_trip(vec![2, 8, 3], 1, FftNormalization::None); +} + +#[test] +fn cfft_interleaved_supports_batched_windows() { + round_trip(vec![3, 8], 1, FftNormalization::None); +} + +#[test] +fn cfft_interleaved_preserves_scalar_strided_logical_layout() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let shape = vec![2, 8, 3]; + let strides = vec![30, 3, 1]; + let values = values_for(&shape); + let mut physical = vec![0.0; 108]; + for logical in 0..shape.iter().product::() { + let c0 = logical / (shape[1] * shape[2]); + let c1 = (logical / shape[2]) % shape[1]; + let c2 = logical % shape[2]; + let scalar = 2 * (c0 * strides[0] + c1 * strides[1] + c2 * strides[2]); + physical[scalar] = values[2 * logical]; + physical[scalar + 1] = values[2 * logical + 1]; + } + let input = ComplexTensorHandle::new_strided( + shape, + strides, + client.create_from_slice(f32::as_bytes(&physical)), + dtype, + ) + .unwrap(); + let spectrum = cfft_interleaved(input, 1, FftMode::Forward, FftNormalization::None).unwrap(); + let result = cfft_interleaved(spectrum, 1, FftMode::Inverse, FftNormalization::ByN).unwrap(); + assert_logical_c32_approx(&client, result, &values, 1e-4); +} + +#[test] +fn cfft_interleaved_supports_minimum_n_fft() { + round_trip(vec![1, 2], 1, FftNormalization::None); +} + +#[test] +fn cfft_interleaved_ortho_round_trip() { + round_trip(vec![1, 8], 1, FftNormalization::Ortho); +} + +#[test] +fn cfft_interleaved_shared_memory_boundary_round_trip() { + let client = ::client(&Default::default()); + let n_fft = test_max_shared_fft_n(&client); + run_round_trip(&client, vec![1, n_fft, 1], 1, FftNormalization::ByN, 0.03); +} + +#[test] +#[cfg(feature = "heavy")] +fn cfft_interleaved_first_four_step_round_trip() { + let client = ::client(&Default::default()); + let n_fft = first_four_step_n(&client); + run_round_trip(&client, vec![1, n_fft, 1], 1, FftNormalization::ByN, 0.03); +} + +#[test] +fn cfft_interleaved_rejects_invalid_axis() { + let client = ::client(&Default::default()); + let input = contiguous_c32(&client, vec![8], &values_for(&[8])); + assert!(matches!( + cfft_interleaved(input, 1, FftMode::Forward, FftNormalization::None), + Err(FftError::AxisOutOfBounds { dim: 1, rank: 1 }) + )); +} + +#[test] +fn cfft_interleaved_rejects_invalid_length() { + let client = ::client(&Default::default()); + let input = contiguous_c32(&client, vec![3], &values_for(&[3])); + assert!(matches!( + cfft_interleaved(input, 0, FftMode::Forward, FftNormalization::None), + Err(FftError::InvalidFftLength { n_fft: 3 }) + )); +} + +#[test] +fn cfft_interleaved_launch_rejects_shape_mismatch() { + let client = ::client(&Default::default()); + let input = contiguous_c32(&client, vec![1, 8], &values_for(&[1, 8])); + let output = contiguous_c32(&client, vec![2, 4], &values_for(&[2, 4])); + assert!(matches!( + cfft_interleaved_launch( + &client, + input.binding(), + output.binding(), + 1, + FftMode::Forward, + FftNormalization::None, + ), + Err(FftError::ShapeMismatch { name: "output", .. }) + )); +} + +#[test] +fn cfft_interleaved_launch_rejects_aliased_output() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let handle = client.create_from_slice(f32::as_bytes(&values_for(&[8]))); + let input = ComplexTensorHandle::new_contiguous(vec![8], handle.clone(), dtype).unwrap(); + let output = ComplexTensorHandle::new_contiguous(vec![8], handle, dtype).unwrap(); + assert!(matches!( + cfft_interleaved_launch( + &client, + input.binding(), + output.binding(), + 0, + FftMode::Forward, + FftNormalization::None, + ), + Err(FftError::OverlappingBindings) + )); +} + +#[test] +fn cfft_interleaved_launch_rejects_same_tensor_as_input_and_output() { + let client = ::client(&Default::default()); + let tensor = contiguous_c32(&client, vec![8], &values_for(&[8])); + assert!(matches!( + cfft_interleaved_launch( + &client, + tensor.binding(), + tensor.binding(), + 0, + FftMode::Forward, + FftNormalization::None, + ), + Err(FftError::OverlappingBindings) + )); +} + +#[test] +fn cfft_interleaved_launch_rejects_overlapping_output_layout() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let input = contiguous_c32(&client, vec![2, 2], &values_for(&[2, 2])); + let output = ComplexTensorHandle::new_strided( + vec![2, 2], + vec![0, 1], + client.empty(4 * dtype.size()), + dtype, + ) + .unwrap(); + assert!(matches!( + cfft_interleaved_launch( + &client, + input.binding(), + output.binding(), + 1, + FftMode::Forward, + FftNormalization::None, + ), + Err(FftError::OverlappingBindings) + )); +} diff --git a/crates/cubek-fft/tests/fft/interleaved_irfft.rs b/crates/cubek-fft/tests/fft/interleaved_irfft.rs new file mode 100644 index 000000000..8a0dcc582 --- /dev/null +++ b/crates/cubek-fft/tests/fft/interleaved_irfft.rs @@ -0,0 +1,325 @@ +use cubecl::std::tensor::TensorHandle; +use cubecl::{CubeElement, Runtime, TestRuntime, client::ComputeClient, frontend::CubePrimitive}; +use cubek_fft::eval::cpu_reference::irfft_ref; +use cubek_fft::{ + ComplexTensorHandle, FftError, FftNormalization, irfft_interleaved, irfft_interleaved_launch, + irfft_interleaved_launch_padded, +}; +use cubek_test_utils::{HostData, HostDataType, HostDataVec}; + +fn real_tensor( + client: &ComputeClient, + shape: Vec, + values: &[f32], +) -> TensorHandle { + TensorHandle::new_contiguous( + shape, + client.create_from_slice(f32::as_bytes(values)), + f32::as_type_native_unchecked().storage_type(), + ) +} + +fn scalar_buffer( + client: &ComputeClient, + tensor: ComplexTensorHandle, +) -> Vec { + let raw = tensor.into_raw_parts(); + f32::from_bytes(&client.read_one(raw.handle).unwrap()).to_vec() +} + +fn to_f32(data: HostData) -> Vec { + match data.data { + HostDataVec::F32(values) => values, + _ => panic!("expected F32 host data"), + } +} + +fn spectrum_values(shape: &[usize]) -> Vec { + (0..shape.iter().product::()) + .flat_map(|index| { + let value = ((index as f32 + 0.25) * 0.37).sin(); + [value, value * 0.5 - 0.25] + }) + .collect() +} + +fn expected_signal( + client: &ComputeClient, + spectrum: ComplexTensorHandle, + dim: usize, + normalization: FftNormalization, +) -> Vec { + let shape = spectrum.shape().to_vec(); + let n_fft = (shape[dim] - 1) * 2; + let scalars = scalar_buffer(client, spectrum); + let re = scalars.iter().step_by(2).copied().collect::>(); + let im = scalars + .iter() + .skip(1) + .step_by(2) + .copied() + .collect::>(); + let re = real_tensor(client, shape.clone(), &re); + let im = real_tensor(client, shape, &im); + let by_n = to_f32(irfft_ref( + &HostData::from_tensor_handle(client, re, HostDataType::F32), + &HostData::from_tensor_handle(client, im, HostDataType::F32), + dim, + None, + )); + let multiplier = match normalization { + FftNormalization::ByN => 1.0, + FftNormalization::None => n_fft as f32, + FftNormalization::Ortho => (n_fft as f32).sqrt(), + }; + by_n.into_iter().map(|value| value * multiplier).collect() +} + +fn assert_scalars_approx(actual: &[f32], expected: &[f32]) { + assert_eq!(actual.len(), expected.len()); + for (index, (actual, expected)) in actual.iter().zip(expected).enumerate() { + assert!( + (actual - expected).abs() <= 1e-4, + "scalar {index}: got {actual}, expected {expected}" + ); + } +} + +fn run_allocating_case(shape: Vec, dim: usize, normalization: FftNormalization) { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let expected = expected_signal( + &client, + ComplexTensorHandle::::new_contiguous( + shape.clone(), + client.create_from_slice(f32::as_bytes(&spectrum_values(&shape))), + dtype, + ) + .unwrap(), + dim, + normalization, + ); + let spectrum = ComplexTensorHandle::::new_contiguous( + shape.clone(), + client.create_from_slice(f32::as_bytes(&spectrum_values(&shape))), + dtype, + ) + .unwrap(); + + let signal = irfft_interleaved(spectrum, dim, normalization).unwrap(); + let mut expected_shape = shape; + expected_shape[dim] = (expected_shape[dim] - 1) * 2; + assert_eq!(signal.shape().as_slice(), expected_shape); + assert_scalars_approx( + f32::from_bytes(&client.read_one(signal.handle).unwrap()), + &expected, + ); +} + +#[test] +fn irfft_interleaved_axis_last_matches_reference() { + run_allocating_case(vec![2, 5], 1, FftNormalization::ByN); +} + +#[test] +fn irfft_interleaved_axis_zero_and_middle_match_reference_with_trailing_batches() { + run_allocating_case(vec![5, 2, 3], 0, FftNormalization::ByN); + run_allocating_case(vec![2, 5, 3], 1, FftNormalization::ByN); +} + +#[test] +fn irfft_interleaved_applies_all_normalizations_at_final_real_stores() { + for normalization in [ + FftNormalization::None, + FftNormalization::ByN, + FftNormalization::Ortho, + ] { + run_allocating_case(vec![2, 5, 3], 1, normalization); + } +} + +#[test] +fn irfft_interleaved_padded_dc_only_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let virtual_shape = vec![2, 1, 3]; + let mut materialized_shape = virtual_shape.clone(); + materialized_shape[1] = 5; + let signal_shape = vec![2, 8, 3]; + let virtual_spectrum = ComplexTensorHandle::new_contiguous( + virtual_shape.clone(), + client.create_from_slice(f32::as_bytes(&spectrum_values(&virtual_shape))), + dtype, + ) + .unwrap(); + let mut materialized_values = vec![0.0; materialized_shape.iter().product::() * 2]; + for before in 0..virtual_shape[0] { + for after in 0..virtual_shape[2] { + let virtual_offset = (before * virtual_shape[1] * virtual_shape[2] + after) * 2; + let materialized_offset = + (before * materialized_shape[1] * materialized_shape[2] + after) * 2; + materialized_values[materialized_offset..materialized_offset + 2].copy_from_slice( + &spectrum_values(&virtual_shape)[virtual_offset..virtual_offset + 2], + ); + } + } + let materialized_spectrum = ComplexTensorHandle::new_contiguous( + materialized_shape, + client.create_from_slice(f32::as_bytes(&materialized_values)), + dtype, + ) + .unwrap(); + let virtual_signal = real_tensor(&client, signal_shape.clone(), &[0.0; 48]); + let materialized_signal = real_tensor(&client, signal_shape, &[0.0; 48]); + + irfft_interleaved_launch_padded( + &client, + virtual_spectrum.binding(), + &virtual_signal, + 1, + 1, + FftNormalization::Ortho, + ) + .unwrap(); + irfft_interleaved_launch( + &client, + materialized_spectrum.binding(), + &materialized_signal, + 1, + FftNormalization::Ortho, + ) + .unwrap(); + + assert_scalars_approx( + f32::from_bytes(&client.read_one(virtual_signal.handle).unwrap()), + f32::from_bytes(&client.read_one(materialized_signal.handle).unwrap()), + ); +} + +#[test] +fn irfft_interleaved_launch_rejects_non_f32_output() { + let client = ::client(&Default::default()); + let spectrum_shape = vec![2, 5]; + let spectrum = ComplexTensorHandle::new_contiguous( + spectrum_shape.clone(), + client.create_from_slice(f32::as_bytes(&spectrum_values(&spectrum_shape))), + f32::as_type_native_unchecked().storage_type(), + ) + .unwrap(); + let output = TensorHandle::new_contiguous( + vec![2, 8], + client.empty(16 * i32::as_type_native_unchecked().storage_type().size()), + i32::as_type_native_unchecked().storage_type(), + ); + + assert!(matches!( + irfft_interleaved_launch( + &client, + spectrum.binding(), + &output, + 1, + FftNormalization::ByN, + ), + Err(FftError::UnsupportedDtype { .. }) + )); +} + +#[test] +fn irfft_interleaved_launch_rejects_shared_output_allocation() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let allocation = client.empty(20 * dtype.size()); + let spectrum = + ComplexTensorHandle::new_contiguous(vec![2, 5], allocation.clone(), dtype).unwrap(); + let output = TensorHandle::new_contiguous(vec![2, 8], allocation, dtype); + + assert!(matches!( + irfft_interleaved_launch( + &client, + spectrum.binding(), + &output, + 1, + FftNormalization::ByN, + ), + Err(FftError::OverlappingBindings) + )); +} + +#[test] +#[cfg(feature = "heavy")] +fn interleaved_irfft_large_multi_bin_virtual_padding_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let max_elems = + client.properties().hardware.max_shared_memory_size / (2 * core::mem::size_of::()); + let max_shared_fft_n = if max_elems.is_power_of_two() { + max_elems + } else { + max_elems.next_power_of_two() >> 1 + }; + let n_fft = 2 * max_shared_fft_n; + let n_freq = n_fft / 2 + 1; + let spec_bins = n_freq / 2; + let virtual_shape = vec![2, spec_bins, 3]; + let materialized_shape = vec![2, n_freq, 3]; + let signal_shape = vec![2, n_fft, 3]; + let virtual_values = spectrum_values(&virtual_shape); + let virtual_spectrum = ComplexTensorHandle::new_contiguous( + virtual_shape.clone(), + client.create_from_slice(f32::as_bytes(&virtual_values)), + dtype, + ) + .unwrap(); + let mut materialized_values = vec![0.0; materialized_shape.iter().product::() * 2]; + for before in 0..virtual_shape[0] { + for bin in 0..virtual_shape[1] { + for after in 0..virtual_shape[2] { + let virtual_offset = + ((before * virtual_shape[1] + bin) * virtual_shape[2] + after) * 2; + let materialized_offset = + ((before * materialized_shape[1] + bin) * materialized_shape[2] + after) * 2; + materialized_values[materialized_offset..materialized_offset + 2] + .copy_from_slice(&virtual_values[virtual_offset..virtual_offset + 2]); + } + } + } + let materialized_spectrum = ComplexTensorHandle::new_contiguous( + materialized_shape, + client.create_from_slice(f32::as_bytes(&materialized_values)), + dtype, + ) + .unwrap(); + let virtual_signal = real_tensor( + &client, + signal_shape.clone(), + &vec![0.0; signal_shape.iter().product()], + ); + let materialized_signal = real_tensor( + &client, + signal_shape.clone(), + &vec![0.0; signal_shape.iter().product()], + ); + + irfft_interleaved_launch_padded( + &client, + virtual_spectrum.binding(), + &virtual_signal, + 1, + spec_bins, + FftNormalization::Ortho, + ) + .unwrap(); + irfft_interleaved_launch( + &client, + materialized_spectrum.binding(), + &materialized_signal, + 1, + FftNormalization::Ortho, + ) + .unwrap(); + + assert_scalars_approx( + f32::from_bytes(&client.read_one(virtual_signal.handle).unwrap()), + f32::from_bytes(&client.read_one(materialized_signal.handle).unwrap()), + ); +} diff --git a/crates/cubek-fft/tests/fft/interleaved_rfft.rs b/crates/cubek-fft/tests/fft/interleaved_rfft.rs new file mode 100644 index 000000000..bdeb599e2 --- /dev/null +++ b/crates/cubek-fft/tests/fft/interleaved_rfft.rs @@ -0,0 +1,285 @@ +use cubecl::std::tensor::TensorHandle; +use cubecl::{CubeElement, Runtime, TestRuntime, client::ComputeClient, frontend::CubePrimitive}; +use cubek_fft::eval::cpu_reference::rfft_ref; +use cubek_fft::{ + ComplexTensorHandle, FftError, FftNormalization, rfft_interleaved, rfft_interleaved_launch, + rfft_interleaved_launch_padded, +}; +use cubek_test_utils::{HostData, HostDataType, HostDataVec}; + +#[cfg(feature = "heavy")] +use cubek_fft::irfft_interleaved; + +fn real_tensor( + client: &ComputeClient, + shape: Vec, + values: &[f32], +) -> TensorHandle { + TensorHandle::new_contiguous( + shape, + client.create_from_slice(f32::as_bytes(values)), + f32::as_type_native_unchecked().storage_type(), + ) +} + +fn scalar_buffer( + client: &ComputeClient, + tensor: ComplexTensorHandle, +) -> Vec { + let raw = tensor.into_raw_parts(); + f32::from_bytes(&client.read_one(raw.handle).unwrap()).to_vec() +} + +fn to_f32(data: HostData) -> Vec { + match data.data { + HostDataVec::F32(values) => values, + _ => panic!("expected F32 host data"), + } +} + +fn expected_interleaved(signal: HostData, dim: usize, normalization: FftNormalization) -> Vec { + let n_fft = signal.shape.as_slice()[dim]; + let scale = normalization.scale_f32(n_fft).unwrap(); + let (re, im) = rfft_ref(&signal, dim, None); + to_f32(re) + .into_iter() + .zip(to_f32(im)) + .flat_map(|(re, im)| [re * scale, im * scale]) + .collect() +} + +fn values_for(shape: &[usize]) -> Vec { + (0..shape.iter().product::()) + .map(|index| ((index as f32 + 0.25) * 0.37).sin()) + .collect() +} + +fn assert_scalars_approx(actual: &[f32], expected: &[f32]) { + assert_eq!(actual.len(), expected.len()); + for (index, (actual, expected)) in actual.iter().zip(expected).enumerate() { + assert!( + (actual - expected).abs() <= 1e-4, + "scalar {index}: got {actual}, expected {expected}" + ); + } +} + +fn run_allocating_case(shape: Vec, dim: usize, normalization: FftNormalization) { + let client = ::client(&Default::default()); + let values = values_for(&shape); + let signal = real_tensor(&client, shape.clone(), &values); + let expected = expected_interleaved( + HostData::from_tensor_handle(&client, signal.clone(), HostDataType::F32), + dim, + normalization, + ); + + let spectrum = rfft_interleaved(signal, dim, normalization).unwrap(); + let mut expected_shape = shape; + expected_shape[dim] = expected_shape[dim] / 2 + 1; + assert_eq!(spectrum.shape(), expected_shape); + assert_scalars_approx(&scalar_buffer(&client, spectrum), &expected); +} + +#[cfg(feature = "heavy")] +fn first_large_n(client: &ComputeClient) -> usize { + let max_elems = + client.properties().hardware.max_shared_memory_size / (2 * core::mem::size_of::()); + let max_shared_fft_n = if max_elems.is_power_of_two() { + max_elems + } else { + max_elems.next_power_of_two() >> 1 + }; + 2 * max_shared_fft_n +} + +#[cfg(feature = "heavy")] +fn run_large_round_trip(shape: Vec, dim: usize) { + let client = ::client(&Default::default()); + let values = values_for(&shape); + let signal = real_tensor(&client, shape, &values); + let spectrum = rfft_interleaved(signal, dim, FftNormalization::None).unwrap(); + let reconstructed = irfft_interleaved(spectrum, dim, FftNormalization::ByN).unwrap(); + let reconstructed_bytes = client.read_one(reconstructed.handle).unwrap(); + let reconstructed = f32::from_bytes(&reconstructed_bytes); + for (index, (actual, expected)) in reconstructed.iter().zip(&values).enumerate() { + assert!( + (actual - expected).abs() <= 0.04, + "sample {index}: got {actual}, expected {expected}" + ); + } +} + +#[test] +fn rfft_interleaved_axis_last_matches_reference() { + run_allocating_case(vec![2, 8], 1, FftNormalization::None); +} + +#[test] +fn rfft_interleaved_axis_zero_and_middle_match_reference_with_trailing_batches() { + run_allocating_case(vec![8, 2, 3], 0, FftNormalization::None); + run_allocating_case(vec![2, 8, 3], 1, FftNormalization::None); +} + +#[test] +fn rfft_interleaved_applies_all_normalizations_at_direct_c32_stores() { + for normalization in [ + FftNormalization::None, + FftNormalization::ByN, + FftNormalization::Ortho, + ] { + run_allocating_case(vec![2, 8, 3], 1, normalization); + } +} + +#[test] +fn rfft_interleaved_padded_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let virtual_shape = vec![2, 5, 3]; + let mut padded_shape = virtual_shape.clone(); + padded_shape[1] = 8; + let mut spectrum_shape = virtual_shape.clone(); + spectrum_shape[1] = 5; + + let virtual_signal = real_tensor(&client, virtual_shape.clone(), &values_for(&virtual_shape)); + let mut padded_values = vec![0.0; padded_shape.iter().product()]; + for batch_before in 0..virtual_shape[0] { + for sample in 0..virtual_shape[1] { + for batch_after in 0..virtual_shape[2] { + padded_values + [(batch_before * padded_shape[1] + sample) * padded_shape[2] + batch_after] = + values_for(&virtual_shape)[(batch_before * virtual_shape[1] + sample) + * virtual_shape[2] + + batch_after]; + } + } + } + let padded_signal = real_tensor(&client, padded_shape, &padded_values); + let virtual_spectrum = + ComplexTensorHandle::empty(&client, spectrum_shape.clone(), dtype).unwrap(); + let materialized_spectrum = ComplexTensorHandle::empty(&client, spectrum_shape, dtype).unwrap(); + + rfft_interleaved_launch_padded( + &client, + &virtual_signal, + virtual_spectrum.binding(), + 1, + virtual_shape[1], + FftNormalization::Ortho, + ) + .unwrap(); + rfft_interleaved_launch( + &client, + &padded_signal, + materialized_spectrum.binding(), + 1, + FftNormalization::Ortho, + ) + .unwrap(); + + assert_scalars_approx( + &scalar_buffer(&client, virtual_spectrum), + &scalar_buffer(&client, materialized_spectrum), + ); +} + +#[test] +fn rfft_interleaved_launch_rejects_overlapping_output_layout() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let signal = real_tensor(&client, vec![2, 8], &values_for(&[2, 8])); + let spectrum = ComplexTensorHandle::new_strided( + vec![2, 5], + vec![0, 1], + client.empty(10 * dtype.size()), + dtype, + ) + .unwrap(); + + assert!(matches!( + rfft_interleaved_launch( + &client, + &signal, + spectrum.binding(), + 1, + FftNormalization::None, + ), + Err(FftError::OverlappingBindings) + )); +} + +#[test] +#[cfg(feature = "heavy")] +fn interleaved_rfft_and_irfft_first_large_round_trip() { + let client = ::client(&Default::default()); + let n_fft = first_large_n(&client); + run_large_round_trip(vec![2, n_fft], 1); +} + +#[test] +#[cfg(feature = "heavy")] +fn interleaved_rfft_and_irfft_batched_large_round_trip() { + let client = ::client(&Default::default()); + let n_fft = first_large_n(&client); + run_large_round_trip(vec![3, n_fft], 1); +} + +#[test] +#[cfg(feature = "heavy")] +fn interleaved_rfft_and_irfft_strided_axis_large_round_trip() { + let client = ::client(&Default::default()); + let n_fft = first_large_n(&client); + run_large_round_trip(vec![2, n_fft, 3], 1); +} + +#[test] +#[cfg(feature = "heavy")] +fn interleaved_rfft_large_virtual_padding_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let n_fft = first_large_n(&client); + let virtual_shape = vec![2, n_fft / 2, 3]; + let mut padded_shape = virtual_shape.clone(); + padded_shape[1] = n_fft; + let mut spectrum_shape = virtual_shape.clone(); + spectrum_shape[1] = n_fft / 2 + 1; + let virtual_values = values_for(&virtual_shape); + let virtual_signal = real_tensor(&client, virtual_shape.clone(), &virtual_values); + let mut padded_values = vec![0.0; padded_shape.iter().product()]; + for before in 0..virtual_shape[0] { + for sample in 0..virtual_shape[1] { + for after in 0..virtual_shape[2] { + padded_values[(before * n_fft + sample) * virtual_shape[2] + after] = + virtual_values[(before * virtual_shape[1] + sample) * virtual_shape[2] + after]; + } + } + } + let padded_signal = real_tensor(&client, padded_shape, &padded_values); + let virtual_spectrum = + ComplexTensorHandle::empty(&client, spectrum_shape.clone(), dtype).unwrap(); + let materialized_spectrum = ComplexTensorHandle::empty(&client, spectrum_shape, dtype).unwrap(); + + rfft_interleaved_launch_padded( + &client, + &virtual_signal, + virtual_spectrum.binding(), + 1, + virtual_shape[1], + FftNormalization::Ortho, + ) + .unwrap(); + rfft_interleaved_launch( + &client, + &padded_signal, + materialized_spectrum.binding(), + 1, + FftNormalization::Ortho, + ) + .unwrap(); + + assert_scalars_approx( + &scalar_buffer(&client, virtual_spectrum), + &scalar_buffer(&client, materialized_spectrum), + ); +} diff --git a/crates/cubek-fft/tests/fft/interleaved_validation.rs b/crates/cubek-fft/tests/fft/interleaved_validation.rs new file mode 100644 index 000000000..43ace6c20 --- /dev/null +++ b/crates/cubek-fft/tests/fft/interleaved_validation.rs @@ -0,0 +1,172 @@ +use cubecl::{Runtime, TestRuntime, frontend::CubePrimitive, std::tensor::TensorHandle}; +use cubek_fft::{ + ComplexTensorHandle, FftError, FftNormalization, irfft_interleaved, irfft_interleaved_launch, + rfft_interleaved, rfft_interleaved_launch, +}; + +#[test] +fn contiguous_c32_uses_two_adjacent_scalars_per_logical_element() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let complex = ComplexTensorHandle::::empty(&client, vec![2, 3], dtype).unwrap(); + assert_eq!(complex.shape(), &[2, 3]); + assert_eq!(complex.strides(), &[3, 1]); + assert_eq!(complex.scalar_strides(), &[6, 2]); + assert_eq!(complex.physical_scalar_len(), 12); +} + +#[test] +fn c32_rejects_wrong_dtype_and_short_buffer() { + let client = ::client(&Default::default()); + let f32_dtype = f32::as_type_native_unchecked().storage_type(); + let f64_dtype = f64::as_type_native_unchecked().storage_type(); + let wrong = ComplexTensorHandle::::new_contiguous( + vec![4], + client.empty(8 * f64_dtype.size()), + f64_dtype, + ); + assert!(matches!(wrong, Err(FftError::UnsupportedDtype { .. }))); + let short = ComplexTensorHandle::::new_contiguous( + vec![4], + client.empty(7 * f32_dtype.size()), + f32_dtype, + ); + assert!(matches!(short, Err(FftError::InsufficientBuffer { .. }))); +} + +#[test] +fn normalization_scales_are_direction_independent() { + assert_eq!(FftNormalization::None.scale_f32(16).unwrap(), 1.0); + assert_eq!(FftNormalization::ByN.scale_f32(16).unwrap(), 1.0 / 16.0); + assert_eq!(FftNormalization::Ortho.scale_f32(16).unwrap(), 0.25); +} + +#[test] +fn c32_metadata_errors_are_typed() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let rank = ComplexTensorHandle::::new_strided( + vec![2], + vec![], + client.empty(8 * dtype.size()), + dtype, + ); + assert!(matches!(rank, Err(FftError::RankMismatch { .. }))); + + let misaligned = ComplexTensorHandle::::new_contiguous( + vec![2], + client.empty(5 * dtype.size()).offset_start(1), + dtype, + ); + assert!(matches!(misaligned, Err(FftError::MisalignedBuffer { .. }))); + + let stride_overflow = ComplexTensorHandle::::new_strided( + vec![2], + vec![usize::MAX], + client.empty(4 * dtype.size()), + dtype, + ); + assert!(matches!( + stride_overflow, + Err(FftError::StrideOverflow { axis: 0 }) + )); + + let extent_overflow = ComplexTensorHandle::::new_strided( + vec![usize::MAX, 2], + vec![1, 1], + client.empty(4 * dtype.size()), + dtype, + ); + assert!(matches!(extent_overflow, Err(FftError::SizeOverflow))); +} + +#[test] +fn zero_sized_c32_shape_has_no_physical_scalars() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let complex = + ComplexTensorHandle::::new_contiguous(vec![0, 3], client.empty(0), dtype) + .unwrap(); + assert_eq!(complex.physical_scalar_len(), 0); +} + +#[test] +fn non_contiguous_c32_extent_includes_the_last_imaginary_scalar() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let complex = ComplexTensorHandle::::new_strided( + vec![2, 3], + vec![5, 1], + client.empty(16 * dtype.size()), + dtype, + ) + .unwrap(); + assert_eq!(complex.scalar_strides(), &[10, 2]); + assert_eq!(complex.physical_scalar_len(), 16); +} + +fn first_unsupported_real_fft_n() -> usize { + let client = ::client(&Default::default()); + let max_elems = + client.properties().hardware.max_shared_memory_size / (2 * core::mem::size_of::()); + let max_shared = if max_elems.is_power_of_two() { + max_elems + } else { + max_elems.next_power_of_two() >> 1 + }; + max_shared.saturating_mul(max_shared).saturating_mul(4) +} + +#[test] +fn oversized_rfft_is_rejected_for_allocating_and_caller_owned_apis_without_allocating_data() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let n_fft = first_unsupported_real_fft_n(); + let signal = TensorHandle::new_contiguous(vec![0, n_fft], client.empty(0), dtype); + + assert!(matches!( + rfft_interleaved(signal.clone(), 1, FftNormalization::None), + Err(FftError::FftLengthExceedsDeviceLimit { .. }) + )); + + let spectrum = + ComplexTensorHandle::new_contiguous(vec![0, n_fft / 2 + 1], client.empty(0), dtype) + .unwrap(); + assert!(matches!( + rfft_interleaved_launch( + &client, + &signal, + spectrum.binding(), + 1, + FftNormalization::None, + ), + Err(FftError::FftLengthExceedsDeviceLimit { .. }) + )); +} + +#[test] +fn oversized_irfft_is_rejected_for_allocating_and_caller_owned_apis_without_allocating_data() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let n_fft = first_unsupported_real_fft_n(); + let spectrum = + ComplexTensorHandle::new_contiguous(vec![0, n_fft / 2 + 1], client.empty(0), dtype) + .unwrap(); + + assert!(matches!( + irfft_interleaved(spectrum.clone(), 1, FftNormalization::ByN), + Err(FftError::FftLengthExceedsDeviceLimit { .. }) + )); + + let signal = TensorHandle::new_contiguous(vec![0, n_fft], client.empty(0), dtype); + assert!(matches!( + irfft_interleaved_launch( + &client, + spectrum.binding(), + &signal, + 1, + FftNormalization::ByN, + ), + Err(FftError::FftLengthExceedsDeviceLimit { .. }) + )); +} diff --git a/crates/cubek-fft/tests/fft/mod.rs b/crates/cubek-fft/tests/fft/mod.rs index fc95afc5d..929fffa0b 100644 --- a/crates/cubek-fft/tests/fft/mod.rs +++ b/crates/cubek-fft/tests/fft/mod.rs @@ -1,5 +1,9 @@ #[cfg(feature = "benchmarks")] mod bench_catalog; +mod interleaved_cfft; +mod interleaved_irfft; +mod interleaved_rfft; +mod interleaved_validation; mod irfft; mod rfft; mod round_trip; diff --git a/crates/cubek-interpolate/src/components/global/nearest_backward.rs b/crates/cubek-interpolate/src/components/global/nearest_backward.rs index 4afd9c4cd..522e4ef4f 100644 --- a/crates/cubek-interpolate/src/components/global/nearest_backward.rs +++ b/crates/cubek-interpolate/src/components/global/nearest_backward.rs @@ -68,7 +68,7 @@ fn start_index( NearestMode::Exact => { let num = F::cast_from(input_index * output_size); let den = F::cast_from(input_size); - let div = (num / den).ceil() - F::new(0.5); + let div = (num / den).ceil() - F::new(0.5_f32); let mask = F::cast_from((div >= F::zero()) as usize); usize::cast_from(div.ceil() * mask) diff --git a/crates/cubek-interpolate/src/definition/base.rs b/crates/cubek-interpolate/src/definition/base.rs index b0ca295eb..8e73a34c3 100644 --- a/crates/cubek-interpolate/src/definition/base.rs +++ b/crates/cubek-interpolate/src/definition/base.rs @@ -255,7 +255,7 @@ pub fn compute_value_default( } if I::REQUIRES_BOUND_CHECK { - let epsilon = Vector::cast_from(P::EA::new(1e-7)); + let epsilon = Vector::cast_from(P::EA::new(1e-7_f32)); Vector::cast_from(final_value / total_weight.max(epsilon)) } else { Vector::cast_from(final_value) diff --git a/crates/cubek-interpolate/src/definition/modes/bicubic.rs b/crates/cubek-interpolate/src/definition/modes/bicubic.rs index 4ec5884c6..7e6cb01e9 100644 --- a/crates/cubek-interpolate/src/definition/modes/bicubic.rs +++ b/crates/cubek-interpolate/src/definition/modes/bicubic.rs @@ -11,22 +11,23 @@ impl Interpolate for Bicubic { const REQUIRES_BOUND_CHECK: bool = false; fn compute_weight(x: EA) -> EA { - let a = EA::new(-0.75); + let a = EA::new(-0.75_f32); let abs_x = x.abs(); let x2 = abs_x * abs_x; let x3 = x2 * abs_x; // Convolution 1 (|x| <= 1.0) - let w1 = (a + EA::new(2.0)) * x3 - (a + EA::new(3.0)) * x2 + EA::new(1.0); + let w1 = (a + EA::new(2.0_f32)) * x3 - (a + EA::new(3.0_f32)) * x2 + EA::new(1.0_f32); // Convolution 2 (1.0 < |x| <= 2.0) - let w2 = a * x3 - EA::new(5.0) * a * x2 + EA::new(8.0) * a * abs_x - EA::new(4.0) * a; + let w2 = a * x3 - EA::new(5.0_f32) * a * x2 + EA::new(8.0_f32) * a * abs_x + - EA::new(4.0_f32) * a; select( - abs_x <= EA::new(1.0), + abs_x <= EA::new(1.0_f32), w1, - select(abs_x <= EA::new(2.0), w2, EA::new(0.0)), + select(abs_x <= EA::new(2.0_f32), w2, EA::new(0.0_f32)), ) } } diff --git a/crates/cubek-interpolate/src/definition/modes/bilinear.rs b/crates/cubek-interpolate/src/definition/modes/bilinear.rs index 0a05fa4bb..e00542ff6 100644 --- a/crates/cubek-interpolate/src/definition/modes/bilinear.rs +++ b/crates/cubek-interpolate/src/definition/modes/bilinear.rs @@ -12,6 +12,10 @@ impl Interpolate for Bilinear { fn compute_weight(x: EA) -> EA { let abs_x = x.abs(); - select(abs_x < EA::new(1.0), EA::new(1.0) - abs_x, EA::new(0.0)) + select( + abs_x < EA::new(1.0_f32), + EA::new(1.0_f32) - abs_x, + EA::new(0.0_f32), + ) } } diff --git a/crates/cubek-interpolate/src/definition/modes/lanczos3.rs b/crates/cubek-interpolate/src/definition/modes/lanczos3.rs index 640089f28..b3bf116dd 100644 --- a/crates/cubek-interpolate/src/definition/modes/lanczos3.rs +++ b/crates/cubek-interpolate/src/definition/modes/lanczos3.rs @@ -13,16 +13,16 @@ impl Interpolate for Lanczos3 { fn compute_weight(x: EA) -> EA { let abs_x = x.abs(); let pi_x = EA::cast_from(core::f32::consts::PI) * x; - let denom = (pi_x * pi_x) / EA::new(3.0); - let safe_denom = select(abs_x < EA::new(1e-7), EA::new(1.0), denom); + let denom = (pi_x * pi_x) / EA::new(3.0_f32); + let safe_denom = select(abs_x < EA::new(1e-7_f32), EA::new(1.0_f32), denom); select( - abs_x < EA::new(1e-7), - EA::new(1.0), + abs_x < EA::new(1e-7_f32), + EA::new(1.0_f32), select( - abs_x < EA::new(3.0), - (pi_x.sin() * (pi_x / EA::new(3.0)).sin()) / safe_denom, - EA::new(0.0), + abs_x < EA::new(3.0_f32), + (pi_x.sin() * (pi_x / EA::new(3.0_f32)).sin()) / safe_denom, + EA::new(0.0_f32), ), ) } diff --git a/crates/cubek-interpolate/src/definition/modes/nearest.rs b/crates/cubek-interpolate/src/definition/modes/nearest.rs index 8d2571546..81118fae9 100644 --- a/crates/cubek-interpolate/src/definition/modes/nearest.rs +++ b/crates/cubek-interpolate/src/definition/modes/nearest.rs @@ -11,6 +11,6 @@ impl Interpolate for Nearest { const REQUIRES_BOUND_CHECK: bool = false; fn compute_weight(_x: EA) -> EA { - EA::new(1.0) + EA::new(1.0_f32) } } diff --git a/crates/cubek-resample/src/definition/kernel.rs b/crates/cubek-resample/src/definition/kernel.rs index 9eadd66d6..ecd8c72da 100644 --- a/crates/cubek-resample/src/definition/kernel.rs +++ b/crates/cubek-resample/src/definition/kernel.rs @@ -62,7 +62,7 @@ impl Kernel { #[comptime] vectorized_axis: usize, #[comptime] lane: usize, ) -> F { - let mut weight = F::new(1.0); + let mut weight = F::new(1.0_f32); #[unroll] for axis_idx in 0..comptime!(config.resample_axes.len()) { @@ -90,7 +90,7 @@ impl Kernel { #[cube] fn weight_1d(frac: F, #[comptime] kernel: &Kernel) -> F { match kernel { - Kernel::Uniform { scale } => F::new(1.0) / F::cast_from(*scale), + Kernel::Uniform { scale } => F::new(1.0_f32) / F::cast_from(*scale), Kernel::Linear => linear_weight::(frac), Kernel::Cubic { a_numerator, @@ -104,7 +104,11 @@ fn weight_1d(frac: F, #[comptime] kernel: &Kernel) -> F { #[cube] fn linear_weight(frac: F) -> F { let abs_frac = frac.abs(); - select(abs_frac < F::new(1.0), F::new(1.0) - abs_frac, F::new(0.0)) + select( + abs_frac < F::new(1.0_f32), + F::new(1.0_f32) - abs_frac, + F::new(0.0_f32), + ) } /// Computes the cubic weight for a given fractional position. @@ -121,15 +125,16 @@ fn cubic_weight( let frac3 = frac2 * abs_frac; // Convolution 1 (|x| <= 1.0) - let w1 = (a + F::new(2.0)) * frac3 - (a + F::new(3.0)) * frac2 + F::new(1.0); + let w1 = (a + F::new(2.0_f32)) * frac3 - (a + F::new(3.0_f32)) * frac2 + F::new(1.0_f32); // Convolution 2 (1.0 < |x| <= 2.0) - let w2 = a * frac3 - F::new(5.0) * a * frac2 + F::new(8.0) * a * abs_frac - F::new(4.0) * a; + let w2 = a * frac3 - F::new(5.0_f32) * a * frac2 + F::new(8.0_f32) * a * abs_frac + - F::new(4.0_f32) * a; select( - abs_frac <= F::new(1.0), + abs_frac <= F::new(1.0_f32), w1, - select(abs_frac <= F::new(2.0), w2, F::new(0.0)), + select(abs_frac <= F::new(2.0_f32), w2, F::new(0.0_f32)), ) } @@ -139,15 +144,15 @@ fn lanczos_weight(frac: F, #[comptime] lobes: u8) -> F { let abs_frac = frac.abs(); let pi_frac = F::cast_from(core::f32::consts::PI) * frac; let denom = (pi_frac * pi_frac) / F::cast_from(lobes); - let safe_denom = select(abs_frac < F::new(1e-7), F::new(1.0), denom); + let safe_denom = select(abs_frac < F::new(1e-7_f32), F::new(1.0_f32), denom); select( - abs_frac < F::new(1e-7), - F::new(1.0), + abs_frac < F::new(1e-7_f32), + F::new(1.0_f32), select( abs_frac < F::cast_from(lobes), (pi_frac.sin() * (pi_frac / F::cast_from(lobes)).sin()) / safe_denom, - F::new(0.0), + F::new(0.0_f32), ), ) } diff --git a/crates/cubek-resample/src/definition/semiring.rs b/crates/cubek-resample/src/definition/semiring.rs index 87a25ba07..54e5587cf 100644 --- a/crates/cubek-resample/src/definition/semiring.rs +++ b/crates/cubek-resample/src/definition/semiring.rs @@ -16,7 +16,7 @@ impl Semiring { /// Get the identity element for the semiring. pub fn identity(#[comptime] this: &Self) -> Vector { match this { - Semiring::Linear => Vector::new(F::new(0.0)), + Semiring::Linear => Vector::new(F::new(0.0_f32)), Semiring::Tropical | Semiring::Log => Vector::min_value(), } } @@ -45,7 +45,7 @@ impl Semiring { Semiring::Log => { let m = accumulator.max(value); let diff = (accumulator - value).abs(); - let zero = Vector::new(F::new(0.0)); + let zero = Vector::new(F::new(0.0_f32)); m + (zero - diff).exp().log1p() } } diff --git a/docs/superpowers/plans/2026-07-18-interleaved-c32-fft.md b/docs/superpowers/plans/2026-07-18-interleaved-c32-fft.md new file mode 100644 index 000000000..1dfc6754e --- /dev/null +++ b/docs/superpowers/plans/2026-07-18-interleaved-c32-fft.md @@ -0,0 +1,700 @@ +# Interleaved C32 FFT Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add public, panic-free interleaved F32/C32 CFFT, RFFT, and IRFFT APIs to `cubek-fft`, covering small and four-step paths without standalone pack/unpack kernels. + +**Architecture:** `ComplexTensorHandle` stores one scalar allocation with logical complex shape and scalar-space strides; `ComplexTensorBinding` borrows it for launches. Component-aware layouts map each logical complex coordinate to adjacent even/odd F32 positions. Global I/O is interleaved, while shared memory and internal four-step/packed scratch remain split. + +**Tech Stack:** Rust 2024, CubeCL/CubeK, `thiserror`, `num-complex` CPU references, Cargo integration tests, WGPU/Metal test runtime. + +## Global Constraints + +- Target `tensor4all/cubek:main`; reference issue #6 without closing it. +- Existing split APIs, signatures, numerical behavior, and normalization behavior remain unchanged. +- The first pull request supports F32/C32 only; F64/C64 execution is a follow-up pull request. +- C32 physical order is `[re0, im0, re1, im1, ...]`; no visible trailing dimension of length two. +- Reuse `max_shared_fft_n(client)` for small/four-step selection; do not add a fixed 4096 threshold. +- No standalone pack/unpack or scale kernel, no host staging, and no CPU fallback. Caller-owned launch APIs perform no hidden output allocation; allocating convenience APIs allocate only their documented output. +- Internal split shared memory and split four-step/packed scratch remain unchanged. +- New public APIs return `Result<_, FftError>` and do not validate user input with `assert!` or `unwrap`. +- Out-of-place launches reject aliased writable output before the first launch or scratch allocation. +- Run `codegraph sync` after structural changes so the repository index remains current locally; never commit `.codegraph/`. + +## File map + +- Create `crates/cubek-fft/src/complex.rs`: complex handle/binding, physical extent validation, output uniqueness check. +- Create `crates/cubek-fft/src/error.rs`: public typed FFT validation/launch errors. +- Create `crates/cubek-fft/src/normalization.rs`: public normalization enum and scale helpers. +- Create `crates/cubek-fft/src/interleaved_layout.rs`: CubeCL component-aware real/imaginary layouts. +- Create `crates/cubek-fft/src/fft/cfft_interleaved.rs`: public CFFT API plus small/four-step interleaved kernels. +- Create `crates/cubek-fft/src/fft/rfft_interleaved.rs`: public RFFT API plus small interleaved output path. +- Create `crates/cubek-fft/src/fft/irfft_interleaved.rs`: public IRFFT API plus small interleaved input path. +- Modify `crates/cubek-fft/src/fft/rfft_large.rs`: add interleaved large RFFT/IRFFT boundary kernels while retaining split scratch. +- Modify `crates/cubek-fft/src/lib.rs` and `crates/cubek-fft/src/fft/mod.rs`: export the additive API. +- Modify `crates/cubek-fft/Cargo.toml`: add the workspace `thiserror` dependency. +- Create `crates/cubek-fft/tests/fft/interleaved_cfft.rs`, `interleaved_rfft.rs`, `interleaved_irfft.rs`, and `interleaved_validation.rs`. +- Modify `crates/cubek-fft/tests/fft/mod.rs`: register the new tests. +- Modify `crates/cubek-fft/src/eval/benchmarks/{strategy.rs,benchmark.rs,problem.rs}` and FFT documentation for interleaved benchmark coverage. + +--- + +### Task 1: Complex tensor ABI, errors, and normalization + +**Files:** +- Create: `crates/cubek-fft/src/complex.rs` +- Create: `crates/cubek-fft/src/error.rs` +- Create: `crates/cubek-fft/src/normalization.rs` +- Create: `crates/cubek-fft/tests/fft/interleaved_validation.rs` +- Modify: `crates/cubek-fft/src/lib.rs` +- Modify: `crates/cubek-fft/Cargo.toml` +- Modify: `crates/cubek-fft/tests/fft/mod.rs` + +**Interfaces:** +- Produces `ComplexTensorHandle`, `ComplexTensorBinding<'a, R>`, `FftError`, and `FftNormalization` for all later tasks. +- `ComplexTensorHandle` stores a `TensorHandle` whose shape is logical and whose strides are physical scalar strides. +- Public constructors accept logical complex strides and multiply them by two with checked arithmetic. + +- [ ] **Step 1: Write failing ABI and normalization tests** + +Add these tests to `tests/fft/interleaved_validation.rs`: + +```rust +use cubecl::{CubeElement, Runtime, TestRuntime, frontend::CubePrimitive}; +use cubek_fft::{ComplexTensorHandle, FftError, FftNormalization}; + +#[test] +fn contiguous_c32_uses_two_adjacent_scalars_per_logical_element() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let complex = ComplexTensorHandle::::empty(&client, vec![2, 3], dtype).unwrap(); + assert_eq!(complex.shape(), &[2, 3]); + assert_eq!(complex.strides(), &[3, 1]); + assert_eq!(complex.scalar_strides(), &[6, 2]); + assert_eq!(complex.physical_scalar_len(), 12); +} + +#[test] +fn c32_rejects_wrong_dtype_and_short_buffer() { + let client = ::client(&Default::default()); + let f32_dtype = f32::as_type_native_unchecked().storage_type(); + let f64_dtype = f64::as_type_native_unchecked().storage_type(); + let wrong = ComplexTensorHandle::::new_contiguous( + vec![4], client.empty(8 * f64_dtype.size()), f64_dtype, + ); + assert!(matches!(wrong, Err(FftError::UnsupportedDtype { .. }))); + let short = ComplexTensorHandle::::new_contiguous( + vec![4], client.empty(7 * f32_dtype.size()), f32_dtype, + ); + assert!(matches!(short, Err(FftError::InsufficientBuffer { .. }))); +} + +#[test] +fn normalization_scales_are_direction_independent() { + assert_eq!(FftNormalization::None.scale_f32(16).unwrap(), 1.0); + assert_eq!(FftNormalization::ByN.scale_f32(16).unwrap(), 1.0 / 16.0); + assert_eq!(FftNormalization::Ortho.scale_f32(16).unwrap(), 0.25); +} + +#[test] +fn c32_metadata_errors_are_typed() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let rank = ComplexTensorHandle::::new_strided( + vec![2], vec![], client.empty(8 * dtype.size()), dtype, + ); + assert!(matches!(rank, Err(FftError::RankMismatch { .. }))); + + let misaligned = ComplexTensorHandle::::new_contiguous( + vec![2], client.empty(5 * dtype.size()).offset_start(1), dtype, + ); + assert!(matches!(misaligned, Err(FftError::MisalignedBuffer { .. }))); + + let stride_overflow = ComplexTensorHandle::::new_strided( + vec![2], vec![usize::MAX], client.empty(4 * dtype.size()), dtype, + ); + assert!(matches!(stride_overflow, Err(FftError::StrideOverflow { axis: 0 }))); + + let extent_overflow = ComplexTensorHandle::::new_strided( + vec![usize::MAX, 2], vec![1, 1], client.empty(4 * dtype.size()), dtype, + ); + assert!(matches!(extent_overflow, Err(FftError::SizeOverflow))); +} +``` + +- [ ] **Step 2: Run the tests and verify RED** + +Run: + +```bash +cargo test -p cubek-fft --test lib interleaved_validation -- --nocapture +``` + +Expected: compilation fails because `ComplexTensorHandle`, `FftError`, and `FftNormalization` do not exist. + +- [ ] **Step 3: Implement the public host-side types** + +Add `thiserror = { workspace = true }` to `cubek-fft` dependencies. Implement these exact public surfaces: + +```rust +#[derive(Debug, thiserror::Error)] +pub enum FftError { + #[error("unsupported FFT storage dtype {actual:?}; expected F32")] + UnsupportedDtype { actual: StorageType }, + #[error("shape rank {shape_rank} differs from stride rank {stride_rank}")] + RankMismatch { shape_rank: usize, stride_rank: usize }, + #[error("FFT axis {dim} is out of bounds for rank {rank}")] + AxisOutOfBounds { dim: usize, rank: usize }, + #[error("FFT length must be a power of two and at least 2, got {n_fft}")] + InvalidFftLength { n_fft: usize }, + #[error("{name}={value} is outside {min}..={max}")] + InvalidLength { name: &'static str, value: usize, min: usize, max: usize }, + #[error("complex buffer needs {required} scalar elements but only {available} are available")] + InsufficientBuffer { required: usize, available: usize }, + #[error("complex buffer byte offset {offset} is not aligned to scalar size {scalar_size}")] + MisalignedBuffer { offset: u64, scalar_size: usize }, + #[error("complex scalar stride at axis {axis} overflowed")] + StrideOverflow { axis: usize }, + #[error("complex buffer extent overflowed")] + SizeOverflow, + #[error("{name} shape {actual:?} does not match expected shape {expected:?}")] + ShapeMismatch { name: &'static str, actual: Vec, expected: Vec }, + #[error("input and output allocations overlap")] + OverlappingBindings, + #[error(transparent)] + Launch(#[from] cubecl::prelude::LaunchError), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum FftNormalization { None, ByN, Ortho } + +impl FftNormalization { + pub fn scale_f32(self, n_fft: usize) -> Result { + if n_fft < 2 || !n_fft.is_power_of_two() { + return Err(FftError::InvalidFftLength { n_fft }); + } + Ok(match self { + Self::None => 1.0, + Self::ByN => 1.0 / n_fft as f32, + Self::Ortho => 1.0 / (n_fft as f32).sqrt(), + }) + } +} +``` + +Implement `ComplexTensorHandle::empty`, `new_contiguous`, `new_strided`, `shape`, `strides`, `scalar_strides`, `physical_scalar_len`, `dtype`, `binding`, and `into_raw_parts`. `ComplexTensorBinding<'a, R>` borrows the handle and exposes crate-private `tensor()` for launch code. Use checked multiplication/addition for scalar strides and extent. + +Before production code, extend the failing test file with zero-sized shape and a successful non-contiguous extent case. Assert `physical_scalar_len() == 0` for the former and the checked maximum reachable imaginary scalar plus one for the latter. + +Implement crate-private output validation with the CubeCL handle-count contract: + +```rust +pub(crate) fn ensure_unique_output(tensor: &TensorHandle) -> Result<(), FftError> { + let probe = tensor.handle.clone(); + if probe.can_mut() { Ok(()) } else { Err(FftError::OverlappingBindings) } +} +``` + +The probe makes a unique output have two handles (accepted) and an aliased input/output have at least three (rejected). Perform this check before cloning a handle into a CubeCL `TensorBinding`. + +- [ ] **Step 4: Run focused and crate tests and verify GREEN** + +```bash +cargo test -p cubek-fft --test lib interleaved_validation -- --nocapture +cargo test -p cubek-fft +cargo fmt --all -- --check +``` + +Expected: the new focused tests pass; all existing 14 FFT integration tests still pass; formatting is clean. + +- [ ] **Step 5: Commit the host-side ABI** + +```bash +git add crates/cubek-fft/Cargo.toml crates/cubek-fft/src/lib.rs \ + crates/cubek-fft/src/complex.rs crates/cubek-fft/src/error.rs \ + crates/cubek-fft/src/normalization.rs crates/cubek-fft/tests/fft/mod.rs \ + crates/cubek-fft/tests/fft/interleaved_validation.rs +git commit -m "feat(fft): add interleaved complex tensor ABI" +``` + +--- + +### Task 2: Component layout and small interleaved CFFT + +**Files:** +- Create: `crates/cubek-fft/src/interleaved_layout.rs` +- Create: `crates/cubek-fft/src/fft/cfft_interleaved.rs` +- Create: `crates/cubek-fft/tests/fft/interleaved_cfft.rs` +- Modify: `crates/cubek-fft/src/lib.rs` +- Modify: `crates/cubek-fft/src/fft/mod.rs` +- Modify: `crates/cubek-fft/tests/fft/mod.rs` + +**Interfaces:** +- Consumes Task 1 types. +- Produces `cfft_interleaved`, `cfft_interleaved_launch`, and the component layout used by RFFT/IRFFT. + +- [ ] **Step 1: Write a failing small CFFT round-trip test** + +Create an 8-element C32 input with known scalar order, launch forward `None`, then inverse `ByN`, and compare the returned raw scalar buffer to the input: + +```rust +#[test] +fn cfft_interleaved_small_round_trip_preserves_c32_order() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + let values: Vec = (0..8).flat_map(|i| [i as f32 + 0.25, -(i as f32)]).collect(); + let input = ComplexTensorHandle::new_contiguous( + vec![1, 8], client.create_from_slice(f32::as_bytes(&values)), dtype, + ).unwrap(); + let spectrum = cfft_interleaved(input, 1, FftMode::Forward, FftNormalization::None).unwrap(); + let result = cfft_interleaved(spectrum, 1, FftMode::Inverse, FftNormalization::ByN).unwrap(); + assert_complex_scalars_approx(&client, result, &values, 1e-4); +} +``` + +Also add axis-0, middle-axis, batched, scalar-strided logical layout, minimum `n_fft = 2`, `Ortho` round-trip, invalid-axis, invalid-length, shape mismatch, and aliased-output tests using a shared helper in the same file. + +- [ ] **Step 2: Run the small CFFT test and verify RED** + +```bash +cargo test -p cubek-fft --test lib cfft_interleaved_small_round_trip_preserves_c32_order -- --nocapture +``` + +Expected: compilation fails because `cfft_interleaved` and its launch function do not exist. + +- [ ] **Step 3: Implement component views and small CFFT** + +Implement `InterleavedBatchSignalLayout` with scalar-space metadata: + +```rust +#[derive(CubeType, Clone, Copy)] +pub(crate) struct InterleavedBatchSignalLayout { + num_samples: usize, + stride_samples: usize, + batch_offset: usize, + component: usize, +} +``` + +`to_source_pos(coords)` returns `batch_offset + coords * stride_samples + component`; `shape()` returns the logical axis length. Build real and imaginary views with components 0 and 1. + +Expose these signatures: + +```rust +pub fn cfft_interleaved( + input: ComplexTensorHandle, dim: usize, mode: FftMode, + normalization: FftNormalization, +) -> Result, FftError>; + +pub fn cfft_interleaved_launch( + client: &ComputeClient, input: ComplexTensorBinding<'_, R>, + output: ComplexTensorBinding<'_, R>, dim: usize, mode: FftMode, + normalization: FftNormalization, +) -> Result<(), FftError>; +``` + +Validate all metadata and output uniqueness before cloning either tensor handle. For `n_fft <= max_shared_fft_n(client)`, launch a new small kernel that reads component views into existing split shared arrays, calls `fft_butterfly_parallel`, and writes `shared_re[k] * scale` and `shared_im[k] * scale` to component views. Keep the existing split kernel unchanged. + +- [ ] **Step 4: Verify small paths and compatibility** + +```bash +cargo test -p cubek-fft --test lib interleaved_cfft -- --nocapture +cargo test -p cubek-fft +cargo fmt --all -- --check +``` + +Expected: all light interleaved CFFT cases and the existing split suite pass. + +- [ ] **Step 5: Commit small CFFT** + +```bash +git add crates/cubek-fft/src/lib.rs crates/cubek-fft/src/interleaved_layout.rs \ + crates/cubek-fft/src/fft/mod.rs crates/cubek-fft/src/fft/cfft_interleaved.rs \ + crates/cubek-fft/tests/fft/mod.rs crates/cubek-fft/tests/fft/interleaved_cfft.rs +git commit -m "feat(fft): add small interleaved CFFT" +``` + +--- + +### Task 3: Four-step interleaved CFFT + +**Files:** +- Modify: `crates/cubek-fft/src/fft/cfft.rs` +- Modify: `crates/cubek-fft/src/fft/cfft_interleaved.rs` +- Modify: `crates/cubek-fft/tests/fft/interleaved_cfft.rs` + +**Interfaces:** +- Extends Task 2 launch functions without changing signatures. +- Reuses `factor_four_step`, `max_shared_fft_n(client)`, and existing split scratch allocation. + +- [ ] **Step 1: Add failing first-four-step and boundary tests** + +Add this test-only calculation, which mirrors the launch decision from reported hardware properties without exposing an internal selector: + +```rust +fn test_max_shared_fft_n(client: &ComputeClient) -> usize { + let max_elems = client.properties().hardware.max_shared_memory_size + / (2 * core::mem::size_of::()); + if max_elems.is_power_of_two() { + max_elems + } else { + max_elems.next_power_of_two() >> 1 + } +} + +fn first_four_step_n(client: &ComputeClient) -> usize { + 2 * test_max_shared_fft_n(client) +} +``` + +Test `n_fft = test_max_shared_fft_n(&client)` and `n_fft = first_four_step_n(&client)`. Mark the four-step numerical case `#[cfg(feature = "heavy")]` when the device-reported limit makes it expensive. + +```rust +#[test] +#[cfg(feature = "heavy")] +fn cfft_interleaved_first_four_step_round_trip() { + let client = ::client(&Default::default()); + let n_fft = first_four_step_n(&client); + run_round_trip(&client, vec![1, n_fft, 1], 1, FftNormalization::ByN, 0.03); +} +``` + +- [ ] **Step 2: Verify the large test fails for the missing dispatch** + +```bash +cargo test -p cubek-fft --features heavy --test lib cfft_interleaved_first_four_step_round_trip -- --nocapture +``` + +Expected: the launch returns a typed unsupported-size/internal-path error or the test fails because only the small path exists. + +- [ ] **Step 3: Implement the hybrid four-step path** + +For `n_fft > max_shared_fft_n(client)`: + +1. Allocate the existing split `scratch_re` and `scratch_im` with logical shape. +2. First radix kernel reads interleaved real/imag component views and writes split scratch with the existing post-twiddle. +3. Change `cfft_four_step_radix2_kernel` in `cfft.rs` from private to `pub(crate)` and call the same generated launch module from the interleaved path. Its split, in-place signature remains unchanged. +4. Final transpose reads split scratch and writes adjacent interleaved output scalars multiplied by the requested scale. + +Factor both dimensions with `factor_four_step(n_fft, max_shared_fft_n(client))`; return `FftError::InvalidFftLength`/`SizeOverflow` instead of asserting on user-controlled lengths. + +- [ ] **Step 4: Verify small/four-step boundary and numerical tests** + +```bash +cargo test -p cubek-fft --features heavy --test lib interleaved_cfft -- --nocapture +cargo test -p cubek-fft --features heavy +cargo fmt --all -- --check +``` + +Expected: boundary selection, both numerical paths, and existing heavy split tests pass. + +- [ ] **Step 5: Commit four-step CFFT** + +```bash +git add crates/cubek-fft/src/fft/cfft.rs crates/cubek-fft/src/fft/cfft_interleaved.rs \ + crates/cubek-fft/tests/fft/interleaved_cfft.rs +git commit -m "feat(fft): add four-step interleaved CFFT" +``` + +--- + +### Task 4: Small and padded interleaved RFFT + +**Files:** +- Create: `crates/cubek-fft/src/fft/rfft_interleaved.rs` +- Create: `crates/cubek-fft/tests/fft/interleaved_rfft.rs` +- Modify: `crates/cubek-fft/src/fft/mod.rs` +- Modify: `crates/cubek-fft/tests/fft/mod.rs` + +**Interfaces:** +- Produces allocating, caller-owned, and padded interleaved RFFT functions. + +- [ ] **Step 1: Add failing numerical, layout, normalization, and padding tests** + +Expose and test: + +```rust +pub fn rfft_interleaved( + signal: TensorHandle, dim: usize, normalization: FftNormalization, +) -> Result, FftError>; + +pub fn rfft_interleaved_launch( + client: &ComputeClient, signal: &TensorHandle, + spectrum: ComplexTensorBinding<'_, R>, dim: usize, + normalization: FftNormalization, +) -> Result<(), FftError>; + +pub fn rfft_interleaved_launch_padded( + client: &ComputeClient, signal: &TensorHandle, + spectrum: ComplexTensorBinding<'_, R>, dim: usize, signal_len: usize, + normalization: FftNormalization, +) -> Result<(), FftError>; +``` + +Test axis 0/middle/last, trailing batches, direct `[re, im]` output order against `rfft_ref`, all normalization variants, and virtual padding matching materialized zeros. + +- [ ] **Step 2: Run a focused test and verify RED** + +```bash +cargo test -p cubek-fft --test lib rfft_interleaved_axis_last_matches_reference -- --nocapture +``` + +Expected: compilation fails because the RFFT interleaved API is missing. + +- [ ] **Step 3: Implement small RFFT direct stores** + +Copy only the small-path control flow needed from `rfft.rs`. Keep real input/shared arrays/butterfly unchanged. Replace the two split output tensors with one complex binding and two `InterleavedBatchSignalLayout` component views. Multiply both components by `normalization.scale_f32(n_fft)?` in the final store. Validate dtype, shape, axis, `signal_len`, output extent, and output uniqueness before launch. + +- [ ] **Step 4: Run RFFT and regression suites** + +```bash +cargo test -p cubek-fft --test lib interleaved_rfft -- --nocapture +cargo test -p cubek-fft +cargo fmt --all -- --check +``` + +Expected: all new light/padding tests and existing split tests pass. + +- [ ] **Step 5: Commit interleaved RFFT small path** + +```bash +git add crates/cubek-fft/src/fft/mod.rs crates/cubek-fft/src/fft/rfft_interleaved.rs \ + crates/cubek-fft/tests/fft/mod.rs crates/cubek-fft/tests/fft/interleaved_rfft.rs +git commit -m "feat(fft): add small interleaved RFFT" +``` + +--- + +### Task 5: Small and padded interleaved IRFFT + +**Files:** +- Create: `crates/cubek-fft/src/fft/irfft_interleaved.rs` +- Create: `crates/cubek-fft/tests/fft/interleaved_irfft.rs` +- Modify: `crates/cubek-fft/src/fft/mod.rs` +- Modify: `crates/cubek-fft/tests/fft/mod.rs` + +**Interfaces:** +- Produces allocating, caller-owned, and padded interleaved IRFFT functions. + +- [ ] **Step 1: Add failing numerical, normalization, and padding tests** + +Use these public signatures: + +```rust +pub fn irfft_interleaved( + spectrum: ComplexTensorHandle, dim: usize, + normalization: FftNormalization, +) -> Result, FftError>; + +pub fn irfft_interleaved_launch( + client: &ComputeClient, spectrum: ComplexTensorBinding<'_, R>, + signal: &TensorHandle, dim: usize, normalization: FftNormalization, +) -> Result<(), FftError>; + +pub fn irfft_interleaved_launch_padded( + client: &ComputeClient, spectrum: ComplexTensorBinding<'_, R>, + signal: &TensorHandle, dim: usize, spec_bins: usize, + normalization: FftNormalization, +) -> Result<(), FftError>; +``` + +Test axis 0/middle/last, trailing batches, all normalization variants, DC-only input, and virtual spectrum padding. For normalization expectations, derive `None` and `Ortho` from the existing `ByN` CPU reference by multiplying by `n_fft` and `sqrt(n_fft)` respectively. + +- [ ] **Step 2: Run a focused test and verify RED** + +```bash +cargo test -p cubek-fft --test lib irfft_interleaved_axis_last_matches_reference -- --nocapture +``` + +Expected: compilation fails because the IRFFT interleaved API is missing. + +- [ ] **Step 3: Implement small IRFFT direct loads** + +Copy only the small-path control flow needed from `irfft.rs`. Read the half-spectrum from real and imaginary component views, reconstruct conjugate bins in the existing split shared arrays, run the inverse butterfly, and write the real signal with the selected scale. Do not change the existing split IRFFT's hard-coded `1 / n_fft` behavior. Validate output uniqueness before converting its handle to a CubeCL binding. + +- [ ] **Step 4: Run IRFFT and full light suites** + +```bash +cargo test -p cubek-fft --test lib interleaved_irfft -- --nocapture +cargo test -p cubek-fft +cargo fmt --all -- --check +``` + +Expected: new IRFFT tests and all existing split tests pass. + +- [ ] **Step 5: Commit interleaved IRFFT small path** + +```bash +git add crates/cubek-fft/src/fft/mod.rs crates/cubek-fft/src/fft/irfft_interleaved.rs \ + crates/cubek-fft/tests/fft/mod.rs crates/cubek-fft/tests/fft/interleaved_irfft.rs +git commit -m "feat(fft): add small interleaved IRFFT" +``` + +--- + +### Task 6: Large interleaved RFFT and IRFFT + +**Files:** +- Modify: `crates/cubek-fft/src/fft/rfft_large.rs` +- Modify: `crates/cubek-fft/src/fft/rfft_interleaved.rs` +- Modify: `crates/cubek-fft/src/fft/irfft_interleaved.rs` +- Modify: `crates/cubek-fft/tests/fft/interleaved_rfft.rs` +- Modify: `crates/cubek-fft/tests/fft/interleaved_irfft.rs` + +**Interfaces:** +- Extends Tasks 4 and 5 without public signature changes. +- Keeps packed CFFT buffers split and changes only the external post/pre boundary kernels. + +- [ ] **Step 1: Add failing heavy large/padded tests** + +Add the first-large-size, batched large, strided-axis large, and virtual padding cases for both directions. Derive the first large size as `2 * max_shared_fft_n(client)` rather than assuming 8192. + +```rust +#[test] +#[cfg(feature = "heavy")] +fn interleaved_rfft_and_irfft_first_large_round_trip() { + let client = ::client(&Default::default()); + let n_fft = first_four_step_n(&client); + run_real_round_trip(&client, vec![2, n_fft], 1, FftNormalization::None, + FftNormalization::ByN, 0.04); +} +``` + +- [ ] **Step 2: Run heavy focused tests and verify RED** + +```bash +cargo test -p cubek-fft --features heavy --test lib interleaved_rfft_and_irfft_first_large_round_trip -- --nocapture +``` + +Expected: launch fails because Tasks 4/5 have no large interleaved dispatch. + +- [ ] **Step 3: Add direct interleaved large boundary kernels** + +For large RFFT, retain real-to-split packing and split CFFT. Change only the final postprocess kernel to write real/imaginary component views in one complex output, applying the selected scale. + +For large IRFFT, change only the pre-process kernel to read component views into split packed input. Retain split inverse CFFT and real unpack. The existing unpack implements `ByN`; multiply its store by this adjustment: + +```rust +let adjustment = match normalization { + FftNormalization::None => n_fft as f32, + FftNormalization::ByN => 1.0, + FftNormalization::Ortho => (n_fft as f32).sqrt(), +}; +``` + +This preserves the existing algebra while producing the public normalization contract. Fuse the adjustment into the unpack store. Do not add conversion or scaling launches. + +- [ ] **Step 4: Run all heavy FFT tests** + +```bash +cargo test -p cubek-fft --features heavy --test lib interleaved -- --nocapture +cargo test -p cubek-fft --features heavy +cargo fmt --all -- --check +``` + +Expected: small, four-step, padding, normalization, and existing heavy split tests all pass. + +- [ ] **Step 5: Commit large RFFT/IRFFT** + +```bash +git add crates/cubek-fft/src/fft/rfft_large.rs \ + crates/cubek-fft/src/fft/rfft_interleaved.rs \ + crates/cubek-fft/src/fft/irfft_interleaved.rs \ + crates/cubek-fft/tests/fft/interleaved_rfft.rs \ + crates/cubek-fft/tests/fft/interleaved_irfft.rs +git commit -m "feat(fft): add large interleaved real FFTs" +``` + +--- + +### Task 7: Benchmarks, documentation, review, PR, and merge + +**Files:** +- Modify: `crates/cubek-fft/src/eval/benchmarks/strategy.rs` +- Modify: `crates/cubek-fft/src/eval/benchmarks/benchmark.rs` +- Modify: `crates/cubek-fft/src/eval/benchmarks/problem.rs` +- Modify: `crates/cubek-fft/src/lib.rs`: add crate-level ABI documentation. +- Modify: `crates/cubek-fft/tests/fft/interleaved_validation.rs` + +**Interfaces:** +- Consumes the complete interleaved API. +- Produces benchmark catalogue entries and final user-facing ABI documentation. + +- [ ] **Step 1: Add failing benchmark catalogue assertions** + +Extend the benchmark strategy to distinguish split and interleaved: + +```rust +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FftStrategy { Split, Interleaved } + +pub fn strategies() -> Vec> { + vec![ + CatalogEntry::new("default", "Default (split)", FftStrategy::Split), + CatalogEntry::new("interleaved", "Interleaved C32", FftStrategy::Interleaved), + ] +} +``` + +Update the existing benchmark catalogue test to assert the backward-compatible `default` ID, the new `interleaved` ID, and representative 4096/8192 problem IDs. + +- [ ] **Step 2: Run benchmark catalogue tests and verify RED** + +```bash +cargo test -p cubek-fft --features benchmarks --test lib bench_catalog -- --nocapture +``` + +Expected: the catalogue assertion fails because only the default strategy exists. + +- [ ] **Step 3: Implement benchmark dispatch and ABI documentation** + +Make `FftBench::prepare` allocate either split tensors or one `ComplexTensorHandle` according to `FftStrategy`. Dispatch `execute` to the corresponding split or interleaved launch API. Preserve the same seeded data and shapes so results are comparable. Include the strategy ID in the benchmark name. + +Document the exact physical order, logical shape/stride units, normalization table, output uniqueness requirement, small/four-step device selector, F32/C32 backend scope, and F64/C64 deferral. State explicitly that profiling should show the algorithm kernels only and no standalone pack/unpack pass. + +- [ ] **Step 4: Run fresh final verification** + +Run all commands from the worktree and inspect every exit code: + +```bash +cargo fmt --all -- --check +cargo clippy -p cubek-fft --all-targets --features heavy,benchmarks -- -D warnings +cargo test -p cubek-fft +cargo test -p cubek-fft --features heavy +cargo test -p cubek-fft --features benchmarks --test lib bench_catalog -- --nocapture +cargo test -p cubek-fft --features cubecl/wgpu +git diff --check origin/main...HEAD +``` + +On the Apple host, record the WGPU adapter as Metal in the test output or a short PR note. Expected: formatting/clippy are clean, every test command has zero failures, and the diff check emits no errors. + +- [ ] **Step 5: Commit docs/benchmarks and request code review** + +```bash +git add crates/cubek-fft/src/eval/benchmarks crates/cubek-fft/src/lib.rs \ + crates/cubek-fft/tests/fft/interleaved_validation.rs +git commit -m "docs(fft): document and benchmark interleaved C32" +``` + +Run the requesting-code-review workflow against `origin/main..HEAD`; fix every critical and important finding with a new failing regression test and a separate commit. + +- [ ] **Step 6: Push and create the PR** + +```bash +git push -u origin codex/fft-interleaved-c32 +``` + +Create a PR targeting `main` titled `FFT: add interleaved F32/C32 APIs`. In the body include: + +- `Refs #6` rather than `Closes #6`; +- the hybrid global-interleaved/internal-split architecture; +- normalization and typed-error behavior; +- exact local verification commands/results; +- Metal/WGPU test evidence; +- the remaining F64/C64 follow-up. + +- [ ] **Step 7: Drive CI and review to merge** + +Use the GitHub CI workflow to inspect every required check. Diagnose failures from logs, reproduce locally where possible, add a failing regression test, fix, rerun the full relevant command, and push normally. Address actionable review threads and keep the worktree alive. + +When all required checks are green, the PR is mergeable, and no critical/important review finding remains, merge the PR using the repository's allowed merge method. Re-fetch `main`, verify the PR state is `MERGED`, and update issue #6 with the merged PR link plus an unchecked F64/C64 follow-up item. Do not close issue #6. diff --git a/docs/superpowers/specs/2026-07-18-interleaved-c32-fft-design.md b/docs/superpowers/specs/2026-07-18-interleaved-c32-fft-design.md new file mode 100644 index 000000000..07cc73ec5 --- /dev/null +++ b/docs/superpowers/specs/2026-07-18-interleaved-c32-fft-design.md @@ -0,0 +1,179 @@ +# Interleaved C32 FFT Design + +## Scope + +This design covers the first implementation pull request for tensor4all/cubek#6. It adds an interleaved F32/C32 global-memory ABI to CFFT, RFFT, and IRFFT on `main`. F64/C64 execution is deliberately deferred to a second pull request because Metal does not provide native `f64` shader arithmetic. + +The first pull request must be complete for F32/C32: public APIs, caller-owned launches, small and four-step kernels, normalization, validation, tests, documentation, and benchmarks. The existing split real/imaginary APIs and their behavior remain unchanged. + +## Chosen approach + +Use hybrid direct kernels: + +- Global complex input and output buffers are interleaved as `[re0, im0, re1, im1, ...]`. +- Kernel loads expose separate real and imaginary scalar views of the interleaved buffer. +- Existing split shared-memory arrays and split four-step scratch buffers remain unchanged. +- Final kernel stores write interleaved output directly. +- No standalone interleaved conversion or scaling kernel is introduced. Large real transforms retain the packed-real algorithm stages required to lower through CFFT. + +This avoids extra global-memory passes while limiting changes to the global load/store boundaries. Fully interleaving shared memory and scratch is outside this pull request. + +## Public data model + +Add `ComplexTensorHandle` and `ComplexTensorBinding` as dedicated wrappers for complex tensors. A wrapper owns or borrows one scalar buffer while keeping these concepts distinct: + +- logical complex shape; +- logical complex strides; +- scalar storage dtype; +- physical scalar buffer length. + +For C32, logical complex element `i` occupies scalar offsets `2 * i` and `2 * i + 1`. Logical strides are measured in complex elements and are converted to scalar offsets at the kernel boundary. The public logical shape never contains a synthetic trailing dimension of length two. + +Constructors validate rank, shape/stride rank agreement, offset and extent arithmetic, scalar alignment, and that the physical buffer covers the final reachable real/imaginary scalar. Contiguous allocation uses exactly `2 * logical_element_count * sizeof(f32)` bytes. Binding construction is fallible and never silently changes dtype or shape. + +`num_complex::Complex32` is `#[repr(C)]` with adjacent `re: f32` and `im: f32` fields. A contiguous slice therefore has the same byte layout as a contiguous C32 buffer when the start address satisfies `align_of::()` (four bytes for the pinned `num-complex` representation). This is a layout statement, not permission to create aliased or otherwise unsound Rust views; callers remain responsible for ownership and valid casting. The crate does not add a general CubeCL complex scalar type. + +## Public FFT APIs + +Keep all existing split APIs source-compatible. Add clearly separate interleaved APIs: + +- allocating convenience functions: `cfft_interleaved`, `rfft_interleaved`, and `irfft_interleaved`; +- caller-owned functions: `cfft_interleaved_launch`, `rfft_interleaved_launch`, and `irfft_interleaved_launch`; +- padded caller-owned functions: `rfft_interleaved_launch_padded` and `irfft_interleaved_launch_padded`. + +The new functions return `Result` and do not call `assert!` or `unwrap` for user input validation. The launch APIs accept caller-owned input and output bindings. They perform no host transfer and no hidden output allocation. The large path may allocate the same internal split scratch currently allocated by the split implementation; exposing reusable scratch is a later optimization, not part of this ABI change. + +Out-of-place operation is guaranteed. Unsupported overlapping input/output bindings return a typed error before launch. In-place CFFT is deferred until its aliasing contract is designed independently. + +The first pull request accepts only the F32/C32 dtype pair. The wrappers and enum-based API must not preclude adding F64/C64 dispatch later, but no dummy F64 path is added. + +## Normalization + +Add the public enum: + +```rust +pub enum FftNormalization { + None, + ByN, + Ortho, +} +``` + +The scale for an FFT of length `n_fft` is: + +- `None`: `1.0`; +- `ByN`: `1.0 / n_fft`; +- `Ortho`: `1.0 / sqrt(n_fft)`. + +The enum is accepted by all new forward and inverse interleaved APIs and has the same meaning in either direction. Scaling is fused into the final output store. No separate scaling kernel is launched. + +Existing split APIs retain their current behavior: CFFT and RFFT are unscaled, while IRFFT scales by `1 / n_fft`. + +## Kernel data flow + +### Small CFFT + +The interleaved input is viewed through component-aware layouts. The real layout maps a logical complex coordinate to the even scalar position and the imaginary layout maps it to the adjacent odd scalar position. Loads place values in the existing bit-reversed split shared-memory arrays. The existing butterfly implementation is unchanged. Final real and imaginary values are scaled and stored through the two component views of the interleaved output. + +### Four-step CFFT + +The first radix stage reads interleaved global input through component-aware views and writes the existing split real/imaginary scratch buffers. The second radix stage continues to operate in place on split scratch. The final transpose/reorder reads split scratch and writes scaled interleaved global output. There is no extra conversion pass. + +### RFFT + +The real F32 input and internal split shared memory are unchanged. The final half-spectrum store writes scaled values directly to the real and imaginary component views of one C32 output binding. + +### IRFFT + +The half-spectrum load reads real and imaginary component views from one C32 binding, reconstructs conjugate bins in split shared memory, and runs the existing inverse butterfly. The final real F32 store applies the selected normalization. + +### Large RFFT and IRFFT + +The current large paths lower through CFFT. Their boundary stages are changed to consume or produce the interleaved binding while retaining split internal scratch. They must not route through a standalone interleaved-to-split conversion kernel. + +## Small versus large selection + +Interleaved and split transforms share the existing `max_shared_fft_n(client)` selector. Do not introduce a second threshold or a hard-coded 4096 constant. + +The selector derives the largest power-of-two transform whose two F32 shared arrays fit in `client.properties().hardware.max_shared_memory_size`. On a device with 32 KiB available, this yields 4096. A transform at or below the result uses the small shared-memory path; the next power of two uses the four-step path. + +Four-step factorization continues to keep both factors at or below the same device-derived shared-memory limit and chooses a balanced power-of-two split. + +## Validation and errors + +Add a public `FftError` type with variants that distinguish at least: + +- unsupported scalar dtype or dtype pair; +- rank, shape, or stride mismatch; +- transform axis out of bounds; +- non-power-of-two length or `n_fft < 2`; +- invalid `signal_len` or `spec_bins`; +- insufficient physical buffer bytes or invalid scalar alignment; +- unsupported input/output overlap; +- size arithmetic overflow; +- underlying CubeCL launch/setup failure where the lower layer returns an error. + +All validation happens before the first kernel launch or scratch allocation. Empty batches remain successful no-ops after metadata validation. + +## Compatibility + +- Existing split symbols remain available with unchanged signatures. +- Existing split numerical and normalization behavior remains unchanged. +- Existing arbitrary transform axis, batches, strided tensor layout, and virtual padding behavior is preserved. +- The interleaved API is additive and can be adopted independently by tenferro. +- Issue #6 remains open after this pull request because F64/C64 and the backend capability matrix are handled by the second pull request. + +## Testing + +Follow test-driven development. Each public behavior is introduced by a failing test before production code. + +Host-side tests cover: + +- contiguous Complex32 byte ordering and logical shape; +- physical byte extent calculation for strided logical layouts; +- all typed validation failures; +- normalization scale calculation; +- small/four-step selector boundary using mocked hardware properties or a pure selector helper. + +Backend numerical tests cover CFFT, RFFT, and IRFFT with: + +- transform axis 0, a middle axis, and the last axis; +- single and batched inputs; +- contiguous and currently supported strided layouts; +- virtual padding and truncation parameters; +- `None`, `ByN`, and `Ortho` normalization; +- minimum FFT length; +- the largest small transform and first four-step transform supported by the test device; +- forward/inverse round trips and comparison with the CPU reference; +- direct inspection of interleaved output scalar order. + +Run the F32/C32 suite on CubeCL WGPU/Metal locally and on the repository's CI-supported runtime. Tests that require more device memory may be categorized consistently with the existing heavy/extended test features. + +Regression tests also prove that existing split APIs still compile and retain their current results. + +## Benchmarks and documentation + +Add interleaved CFFT/RFFT/IRFFT cases to the existing FFT benchmark catalogue. Compare them with the existing split implementation at representative small and four-step sizes. Record kernel count or profiling evidence showing that the interleaved path contains no standalone interleaved conversion or scaling kernel. Large real FFTs still contain the packed-real algorithm stages required to lower through CFFT. + +Document: + +- the logical-versus-physical complex tensor contract; +- Rust `Complex32` byte compatibility; +- normalization semantics; +- supported F32/C32 backend behavior; +- the explicit deferral of F64/C64 to the follow-up pull request. + +## Pull request and merge criteria + +The pull request targets `tensor4all/cubek:main` and references issue #6 without closing it. Before merge: + +- formatting and lint checks pass; +- the full `cubek-fft` unit/integration suite passes; +- WGPU/Metal F32/C32 tests pass on the Apple development host; +- repository-required GitHub Actions checks pass; +- no unresolved critical or important review findings remain; +- the issue checklist is updated to identify the merged F32/C32 phase and the remaining F64/C64 phase. + +## Deferred work + +The follow-up pull request adds F64/C64 to the same API and layout model on runtimes with native F64 support. Unsupported runtimes, including Metal, must reject F64/C64 before launch with a typed capability error. It must not narrow to F32, stage through the host, or fall back to CPU execution.