From 7aecbe7aa0f9629297ad66f8dd7185a3d4658e93 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Sun, 19 Jul 2026 09:02:58 +0900 Subject: [PATCH 1/5] feat(fft): backport interleaved APIs to t4a CubeCL --- Cargo.toml | 2 +- crates/cubek-fft/Cargo.toml | 9 +- crates/cubek-fft/src/complex.rs | 273 +++++++ crates/cubek-fft/src/error.rs | 51 ++ crates/cubek-fft/src/eval/cpu_reference.rs | 712 ++++++++++++++++++ crates/cubek-fft/src/eval/mod.rs | 2 + crates/cubek-fft/src/fft/cfft.rs | 103 ++- crates/cubek-fft/src/fft/cfft_interleaved.rs | 504 +++++++++++++ crates/cubek-fft/src/fft/fft_inner.rs | 22 +- crates/cubek-fft/src/fft/irfft.rs | 28 +- crates/cubek-fft/src/fft/irfft_interleaved.rs | 302 ++++++++ crates/cubek-fft/src/fft/limits.rs | 51 ++ crates/cubek-fft/src/fft/mod.rs | 9 +- crates/cubek-fft/src/fft/rfft.rs | 24 +- crates/cubek-fft/src/fft/rfft_interleaved.rs | 322 ++++++++ crates/cubek-fft/src/fft/rfft_large.rs | 449 +++++++++-- crates/cubek-fft/src/interleaved_layout.rs | 66 ++ crates/cubek-fft/src/lib.rs | 50 +- crates/cubek-fft/src/normalization.rs | 21 + .../cubek-fft/tests/fft/interleaved_cfft.rs | 305 ++++++++ .../cubek-fft/tests/fft/interleaved_irfft.rs | 325 ++++++++ .../cubek-fft/tests/fft/interleaved_rfft.rs | 285 +++++++ .../tests/fft/interleaved_validation.rs | 172 +++++ crates/cubek-fft/tests/fft/irfft.rs | 368 +++++++++ crates/cubek-fft/tests/fft/mod.rs | 7 + crates/cubek-fft/tests/fft/rfft.rs | 473 ++++++++++++ crates/cubek-fft/tests/fft/round_trip.rs | 32 + crates/cubek-fft/tests/lib.rs | 2 +- .../2026-07-19-t4a-cubek-fft-git-compat.md | 65 ++ 29 files changed, 4889 insertions(+), 145 deletions(-) create mode 100644 crates/cubek-fft/src/complex.rs create mode 100644 crates/cubek-fft/src/error.rs create mode 100644 crates/cubek-fft/src/eval/cpu_reference.rs create mode 100644 crates/cubek-fft/src/eval/mod.rs create mode 100644 crates/cubek-fft/src/fft/cfft_interleaved.rs create mode 100644 crates/cubek-fft/src/fft/irfft_interleaved.rs create mode 100644 crates/cubek-fft/src/fft/limits.rs create mode 100644 crates/cubek-fft/src/fft/rfft_interleaved.rs create mode 100644 crates/cubek-fft/src/interleaved_layout.rs create mode 100644 crates/cubek-fft/src/normalization.rs create mode 100644 crates/cubek-fft/tests/fft/interleaved_cfft.rs create mode 100644 crates/cubek-fft/tests/fft/interleaved_irfft.rs create mode 100644 crates/cubek-fft/tests/fft/interleaved_rfft.rs create mode 100644 crates/cubek-fft/tests/fft/interleaved_validation.rs create mode 100644 crates/cubek-fft/tests/fft/irfft.rs create mode 100644 crates/cubek-fft/tests/fft/mod.rs create mode 100644 crates/cubek-fft/tests/fft/rfft.rs create mode 100644 crates/cubek-fft/tests/fft/round_trip.rs create mode 100644 docs/superpowers/plans/2026-07-19-t4a-cubek-fft-git-compat.md diff --git a/Cargo.toml b/Cargo.toml index 5023454ce..e72e2bf63 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ resolver = "2" members = [ + "crates/cubek-fft", "crates/cubek-matmul", "crates/cubek-quant", "crates/cubek-random", @@ -18,7 +19,6 @@ exclude = [ "crates/cubek", "crates/cubek-attention", "crates/cubek-convolution", - "crates/cubek-fft", "crates/cubek-interpolate", "crates/cubek-reduce", ] diff --git a/crates/cubek-fft/Cargo.toml b/crates/cubek-fft/Cargo.toml index b21211892..bebe1c56c 100644 --- a/crates/cubek-fft/Cargo.toml +++ b/crates/cubek-fft/Cargo.toml @@ -23,12 +23,13 @@ cpu-reference = ["dep:num-complex", "dep:cubek-test-utils"] [dependencies] cubecl = { workspace = true } -cubek-test-utils = { package = "t4a-cubek-test-utils", path = "./../cubek-test-utils/", version = "0.2.0", default-features = false, optional = true } -num-complex = { workspace = true, optional = true } +thiserror = { workspace = true } +cubek-test-utils = { package = "t4a-cubek-test-utils", path = "./../cubek-test-utils/", version = "=0.2.0", default-features = false, optional = true } +num-complex = { version = "0.4.6", optional = true } [dev-dependencies] -num-complex = { workspace = true } +num-complex = "0.4.6" cubecl = { workspace = true, features = ["test-runtime"] } cubecl-common = { workspace = true } cubek-fft = { path = ".", features = ["cpu-reference"] } -cubek-test-utils = { package = "t4a-cubek-test-utils", path = "./../cubek-test-utils/", version = "0.2.0", default-features = false } +cubek-test-utils = { package = "t4a-cubek-test-utils", path = "./../cubek-test-utils/", version = "=0.2.0", default-features = false } 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/cpu_reference.rs b/crates/cubek-fft/src/eval/cpu_reference.rs new file mode 100644 index 000000000..47194764a --- /dev/null +++ b/crates/cubek-fft/src/eval/cpu_reference.rs @@ -0,0 +1,712 @@ +//! CPU reference and seeded "produce a HostData" primitives for FFT. + +#![allow(clippy::needless_range_loop)] + +use std::f32::consts::PI; + +use cubecl::{ + CubeElement, TestRuntime, + client::ComputeClient, + frontend::CubePrimitive, + zspace::{Shape, Strides}, +}; +use cubek_test_utils::{ + ExecutionOutcome, HostData, HostDataType, HostDataVec, Progress, TestInput, + launch_and_capture_outcome, +}; +use num_complex::Complex; + +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, + }, +}; + +fn row_major_strides(shape: &[usize]) -> Strides { + let mut strides = vec![0; shape.len()]; + let mut stride = 1usize; + for (axis, extent) in shape.iter().enumerate().rev() { + strides[axis] = stride; + stride = stride.saturating_mul(*extent); + } + strides.into() +} + +/// Run the FFT kernel for `mode` against the given problem with seeded inputs +/// and return its output as a [`HostData`]. +pub fn 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.clone(); + spectrum_shape[dim] = shape[dim] / 2 + 1; + + let re = TestInput::builder(client.clone(), spectrum_shape.clone()) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + let im = TestInput::builder(client.clone(), spectrum_shape.clone()) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + + let outcome = launch_and_capture_outcome(&client, |c| { + rfft_launch::( + c, + signal.clone().binding(), + re.clone().binding(), + im.clone().binding(), + dim, + dtype, + ) + .into() + }); + + match outcome { + ExecutionOutcome::CompileError(e) => Err(format!("compile error: {e}")), + ExecutionOutcome::Executed => { + let re_host = HostData::from_tensor_handle(&client, re, HostDataType::F32); + let im_host = HostData::from_tensor_handle(&client, im, HostDataType::F32); + Ok(stack_re_im(re_host, im_host)) + } + } + } + FftMode::Inverse => { + let mut spectrum_shape = shape.clone(); + spectrum_shape[dim] = shape[dim] / 2 + 1; + + let (re, _) = TestInput::builder(client.clone(), spectrum_shape.clone()) + .dtype(dtype) + .uniform(seed_lhs, -1., 1.) + .generate_with_f32_host_data(); + let (im, _) = TestInput::builder(client.clone(), spectrum_shape.clone()) + .dtype(dtype) + .uniform(seed_rhs, -1., 1.) + .generate_with_f32_host_data(); + + let signal = TestInput::builder(client.clone(), shape.clone()) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + + let outcome = launch_and_capture_outcome(&client, |c| { + irfft_launch::( + c, + re.binding(), + im.binding(), + signal.clone().binding(), + dim, + dtype, + ) + .into() + }); + + match outcome { + ExecutionOutcome::CompileError(e) => Err(format!("compile error: {e}")), + ExecutionOutcome::Executed => Ok(HostData::from_tensor_handle( + &client, + signal, + HostDataType::F32, + )), + } + } + } +} + +/// 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 = row_major_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 +/// [`FftMode::Inverse`]. +pub fn cpu_reference_result( + client: ComputeClient, + shape: Vec, + dim: usize, + mode: FftMode, + seed_lhs: u64, + seed_rhs: u64, + progress: Option<&Progress>, +) -> 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 (re, im) = rfft_ref(&signal, dim, progress); + Ok(stack_re_im(re, im)) + } + FftMode::Inverse => { + let mut spectrum_shape = shape.clone(); + spectrum_shape[dim] = shape[dim] / 2 + 1; + + let (_, re) = TestInput::builder(client.clone(), spectrum_shape.clone()) + .dtype(dtype) + .uniform(seed_lhs, -1., 1.) + .generate_with_f32_host_data(); + let (_, im) = TestInput::builder(client.clone(), spectrum_shape.clone()) + .dtype(dtype) + .uniform(seed_rhs, -1., 1.) + .generate_with_f32_host_data(); + + Ok(irfft_ref(&re, &im, dim, progress)) + } + } +} + +/// Number of progress bumps the FFT reference will produce. Granularity is one +/// bump per FFT window — the inner `fft_recursive` dominates runtime, so +/// per-bin bumps would be noisy. +pub fn cpu_reference_total(shape: &[usize], dim: usize, mode: FftMode) -> u64 { + let sample_window = shape[dim]; + let num_freq_bins = sample_window / 2 + 1; + let total: usize = shape.iter().product(); + match mode { + FftMode::Forward => (total / sample_window) as u64, + FftMode::Inverse => { + // For inverse the input shape passed in is the *signal* shape; + // count windows over the spectrum shape (`num_freq_bins` along `dim`). + let mut spec_total = total; + spec_total = spec_total / sample_window * num_freq_bins; + (spec_total / num_freq_bins) as u64 + } + } +} + +/// Stack two equal-shape `HostData` blobs along a fresh leading dim of size 2. +/// Index `0` along that dim is `re`, index `1` is `im`. Used so the forward +/// mode can produce a single comparable [`HostData`] from a (re, im) pair. +fn stack_re_im(re: HostData, im: HostData) -> HostData { + assert_eq!(re.shape, im.shape, "re/im shape mismatch"); + let inner_shape = re.shape.as_slice().to_vec(); + let inner_numel: usize = inner_shape.iter().product(); + + let HostDataVec::F32(re_vec) = re.data else { + panic!("re must be F32"); + }; + let HostDataVec::F32(im_vec) = im.data else { + panic!("im must be F32"); + }; + + let re_strides_slice: &[usize] = &re.strides; + let im_strides_slice: &[usize] = &im.strides; + let mut packed = Vec::with_capacity(inner_numel * 2); + pack_contiguous(&mut packed, &re_vec, re_strides_slice, &inner_shape); + pack_contiguous(&mut packed, &im_vec, im_strides_slice, &inner_shape); + + let mut out_shape_vec = vec![2]; + out_shape_vec.extend(inner_shape); + let out_shape = Shape::from(out_shape_vec); + let strides = row_major_strides(&out_shape); + + HostData { + data: HostDataVec::F32(packed), + shape: out_shape, + strides, + } +} + +fn pack_contiguous(out: &mut Vec, data: &[f32], strides: &[usize], shape: &[usize]) { + let mut idx = vec![0usize; shape.len()]; + let total: usize = shape.iter().product(); + for _ in 0..total { + let mut linear = 0; + for (s, c) in strides.iter().zip(idx.iter()) { + linear += s * c; + } + out.push(data[linear]); + + for d in (0..shape.len()).rev() { + idx[d] += 1; + if idx[d] < shape[d] { + break; + } + idx[d] = 0; + } + } +} + +/// Recursive Cooley-Tukey FFT for complex inputs (length must be power of 2). +fn fft_recursive(x: &mut [Complex], fft_mode: FftMode) { + let n = x.len(); + if n <= 1 { + return; + } + + let mut even: Vec<_> = x.iter().step_by(2).cloned().collect(); + let mut odd: Vec<_> = x.iter().skip(1).step_by(2).cloned().collect(); + + fft_recursive(&mut even, fft_mode); + fft_recursive(&mut odd, fft_mode); + + for k in 0..n / 2 { + let t = Complex::from_polar(1.0, fft_mode.sign() * 2.0 * PI * k as f32 / n as f32) * odd[k]; + x[k] = even[k] + t; + x[k + n / 2] = even[k] - t; + } +} + +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 = row_major_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, + im: &HostData, + dim: usize, + progress: Option<&Progress>, +) -> HostData { + let in_shape = re.shape.as_slice(); + let num_freq_bins = in_shape[dim]; + let sample_window = (num_freq_bins - 1) * 2; + assert!( + sample_window.is_power_of_two(), + "Requires power-of-2 sample_window length" + ); + + let mut out_shape_vec = in_shape.to_vec(); + out_shape_vec[dim] = sample_window; + let out_shape = Shape::from(out_shape_vec); + let num_windows = re.shape.num_elements() / num_freq_bins; + let out_strides = row_major_strides(&out_shape); + + if let Some(p) = progress { + p.set_total(num_windows as u64); + } + + let mut flattened = vec![0.0; out_shape.num_elements()]; + + for l in 0..num_windows { + let mut coords = get_coords(l, in_shape, dim); + let mut spectrum = vec![Complex::new(0.0, 0.0); sample_window]; + + for k in 0..num_freq_bins { + coords[dim] = k; + let r = re.get_f32(&coords); + let i = im.get_f32(&coords); + spectrum[k] = Complex::new(r, i); + } + + for k in 1..num_freq_bins - 1 { + spectrum[sample_window - k] = spectrum[k].conj(); + } + + fft_recursive(&mut spectrum, FftMode::Inverse); + + for i in 0..sample_window { + coords[dim] = i; + let flat_idx = compute_index(&out_strides, coords.as_slice()); + + flattened[flat_idx] = spectrum[i].re / sample_window as f32; + } + if let Some(p) = progress { + p.bump(); + } + } + + HostData { + data: HostDataVec::F32(flattened), + shape: out_shape, + strides: out_strides, + } +} + +/// Reference RFFT: input real slice, output first n/2 + 1 complex numbers. +pub fn rfft_ref( + signal: &HostData, + dim: usize, + progress: Option<&Progress>, +) -> (HostData, HostData) { + let in_shape = signal.shape.as_slice(); + let sample_window = in_shape[dim]; + let num_freq_bins = sample_window / 2 + 1; + assert!( + sample_window.is_power_of_two(), + "Requires power-of-2 sample_window length" + ); + + let mut out_shape_vec = in_shape.to_vec(); + out_shape_vec[dim] = num_freq_bins; + let out_shape = Shape::from(out_shape_vec); + let num_windows = signal.shape.num_elements() / sample_window; + let out_strides = row_major_strides(&out_shape); + + if let Some(p) = progress { + p.set_total(num_windows as u64); + } + + let mut re_data = vec![0.0; out_shape.num_elements()]; + let mut im_data = vec![0.0; out_shape.num_elements()]; + for l in 0..num_windows { + let mut coords = get_coords(l, in_shape, dim); + let mut spectrum = Vec::with_capacity(sample_window); + for i in 0..sample_window { + coords[dim] = i; + let v = signal.get_f32(&coords); + spectrum.push(Complex::new(v, 0.)); + } + + fft_recursive(&mut spectrum, FftMode::Forward); + for k in 0..num_freq_bins { + coords[dim] = k; + let flat_idx = compute_index(&out_strides, coords.as_slice()); + re_data[flat_idx] = spectrum[k].re; + im_data[flat_idx] = spectrum[k].im; + } + if let Some(p) = progress { + p.bump(); + } + } + + ( + HostData { + data: HostDataVec::F32(re_data), + shape: out_shape.clone(), + strides: out_strides.clone(), + }, + HostData { + data: HostDataVec::F32(im_data), + shape: out_shape, + strides: out_strides, + }, + ) +} + +fn get_coords(lane_idx: usize, shape: &[usize], dim: usize) -> Vec { + let mut coords = vec![0; shape.len()]; + let mut temp = lane_idx; + for i in (0..shape.len()).rev() { + if i == dim { + continue; + } + coords[i] = temp % shape[i]; + temp /= shape[i]; + } + coords +} + +fn compute_index(strides: &Strides, coords: &[usize]) -> usize { + assert_eq!( + coords.len(), + strides.rank(), + "Coordinate rank must match stride rank", + ); + + coords + .iter() + .zip(strides.iter()) + .map(|(&c, &s)| c * s) + .sum() +} diff --git a/crates/cubek-fft/src/eval/mod.rs b/crates/cubek-fft/src/eval/mod.rs new file mode 100644 index 000000000..c078f692f --- /dev/null +++ b/crates/cubek-fft/src/eval/mod.rs @@ -0,0 +1,2 @@ +#[cfg(feature = "cpu-reference")] +pub mod cpu_reference; diff --git a/crates/cubek-fft/src/fft/cfft.rs b/crates/cubek-fft/src/fft/cfft.rs index 914655e24..8f806eaed 100644 --- a/crates/cubek-fft/src/fft/cfft.rs +++ b/crates/cubek-fft/src/fft/cfft.rs @@ -2,16 +2,17 @@ //! //! Two flavours live here: //! -//! * A shared-memory kernel (`cfft_kernel`) for any size up to -//! [`MAX_SHARED_N_FFT`]. Structurally identical to `rfft_kernel` / -//! `irfft_kernel` except the imaginary input is read rather than zeroed -//! and all `N` bins are written out (no Hermitian truncation). -//! * A four-step Cooley-Tukey orchestrator for `N > MAX_SHARED_N_FFT` that -//! factors `N = N1 * N2` with both factors `<= MAX_SHARED_N_FFT`. Each -//! sub-FFT reuses the shared-memory butterfly via a dedicated "strided / -//! twiddled" radix kernel (`cfft_four_step_radix_kernel`). A single -//! transpose kernel converts the (N1, N2) internal layout to natural -//! linear bin order on the way out. +//! * A shared-memory kernel (`cfft_kernel`) for any size up to the device's +//! per-cube shared-memory budget (see [`max_shared_fft_n`]). Structurally +//! identical to `rfft_kernel` / `irfft_kernel` except the imaginary input +//! is read rather than zeroed and all `N` bins are written out (no +//! Hermitian truncation). +//! * A four-step Cooley-Tukey orchestrator for larger `N` that factors +//! `N = N1 * N2` with both factors inside that budget. Each sub-FFT +//! reuses the shared-memory butterfly via a dedicated "strided / twiddled" +//! radix kernel (`cfft_four_step_radix_kernel`). A single transpose kernel +//! converts the (N1, N2) internal layout to natural linear bin order on +//! the way out. //! //! The public API of this module is the single //! [`cfft_launch_any_size`] function which picks the right path. @@ -30,21 +31,11 @@ use crate::{ fft::{ FftMode, fft_parallel::{bit_reverse, fft_butterfly_parallel}, + limits::{max_shared_fft_n, max_units_per_cube}, }, layout::BatchSignalLayout, }; -/// Portable size limit for the single-pass shared-memory path. The kernel -/// allocates two `f32` shared buffers of length `n_fft`, so `n_fft = 4096` -/// uses 2 * 4096 * 4 bytes = 32 KiB. Larger sizes use the packed-real / -/// four-step path instead of relying on backend-specific larger workgroup -/// memory limits. -pub(crate) const MAX_SHARED_N_FFT: usize = 4096; - -/// Portable cap on the number of units in one cube. Larger FFTs still cover -/// all bins by having each unit process multiple indices. -const MAX_UNITS_PER_CUBE: usize = 256; - pub(crate) struct CfftBindings { pub(crate) input_re: TensorBinding, pub(crate) input_im: TensorBinding, @@ -61,15 +52,15 @@ struct CfftPlan { } /// Factor `n_fft = N1 * N2` for the four-step FFT. Both factors are powers -/// of two, both `<= MAX_SHARED_N_FFT`, and the split is as balanced as +/// of two, both `<= max_shared_n_fft`, and the split is as balanced as /// possible. -pub(crate) fn factor_four_step(n_fft: usize) -> (usize, usize) { +pub(crate) fn factor_four_step(n_fft: usize, max_shared_n_fft: usize) -> (usize, usize) { assert!( n_fft.is_power_of_two(), "four-step needs power-of-two n_fft" ); let log2_n = n_fft.trailing_zeros() as usize; - let max_log2 = MAX_SHARED_N_FFT.trailing_zeros() as usize; + let max_log2 = max_shared_n_fft.trailing_zeros() as usize; // Balanced split, then push each factor up to the shared-mem cap if the // other factor would otherwise exceed it. let log2_n1 = log2_n / 2; @@ -81,7 +72,7 @@ pub(crate) fn factor_four_step(n_fft: usize) -> (usize, usize) { }; assert!( log2_n1 <= max_log2 && log2_n2 <= max_log2, - "four-step cannot handle n_fft = {n_fft} with MAX_SHARED_N_FFT = {MAX_SHARED_N_FFT}", + "four-step cannot handle n_fft = {n_fft} with max shared-mem n = {max_shared_n_fft}", ); (1 << log2_n1, 1 << log2_n2) } @@ -119,7 +110,7 @@ pub(crate) fn cfft_launch_any_size( fft_mode, }; - if n_fft <= MAX_SHARED_N_FFT { + if n_fft <= max_shared_fft_n(client) { cfft_shared_launch::(client, bindings, plan) } else { cfft_four_step_launch::(client, bindings, dtype, plan) @@ -132,7 +123,7 @@ fn cfft_shared_launch( plan: CfftPlan, ) -> Result<(), LaunchError> { let log2_n = plan.n_fft.trailing_zeros() as usize; - let threads_per_cube = (plan.n_fft / 2).clamp(1, MAX_UNITS_PER_CUBE); + 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()); @@ -175,10 +166,8 @@ fn cfft_shared_kernel( let input_re_view = input_re.view(BatchSignalLayout::new(input_re, window_index, dim)); let input_im_view = input_im.view(BatchSignalLayout::new(input_im, window_index, dim)); - let mut output_re_view = - output_re.view_mut(BatchSignalLayout::new(output_re, window_index, dim)); - let mut output_im_view = - output_im.view_mut(BatchSignalLayout::new(output_im, window_index, dim)); + let output_re_view = output_re.view_mut(BatchSignalLayout::new(&*output_re, window_index, dim)); + let output_im_view = output_im.view_mut(BatchSignalLayout::new(&*output_im, window_index, dim)); let mut shared_re = SharedMemory::::new(n_fft); let mut shared_im = SharedMemory::::new(n_fft); @@ -186,8 +175,8 @@ fn cfft_shared_kernel( let mut i = UNIT_POS as usize; while i < n_fft { let j = bit_reverse(i, log2_n); - shared_re[j] = input_re_view[i]; - shared_im[j] = input_im_view[i]; + shared_re[j] = input_re_view.read_checked(i); + shared_im[j] = input_im_view.read_checked(i); i += threads_per_cube; } sync_cube(); @@ -203,15 +192,15 @@ fn cfft_shared_kernel( let mut k = UNIT_POS as usize; while k < n_fft { - output_re_view[k] = shared_re[k]; - output_im_view[k] = shared_im[k]; + output_re_view.write_checked(k, shared_re[k]); + output_im_view.write_checked(k, shared_im[k]); k += threads_per_cube; } } // --- Four-step path ---------------------------------------------------- -/// Four-step complex FFT for `n_fft > MAX_SHARED_N_FFT`. +/// Four-step complex FFT for `n_fft > max_shared_fft_n(client)`. /// /// Layout convention: each window's `n_fft` axis is viewed as /// `(N1, N2)` row-major with the flat index `n = n1 * N2 + n2`. After the @@ -224,7 +213,9 @@ fn cfft_four_step_launch( dtype: StorageType, plan: CfftPlan, ) -> Result<(), LaunchError> { - let (n1, n2) = factor_four_step(plan.n_fft); + let max_n = max_shared_fft_n(client); + let max_units = max_units_per_cube(client); + let (n1, n2) = factor_four_step(plan.n_fft, max_n); // Scratch buffer, same shape as input. Two passes ping-pong through // scratch and output; the transpose at the end lands in `output`. @@ -245,7 +236,7 @@ fn cfft_four_step_launch( // (window, n2). Reads from `input_*`, writes to `scratch_*` with fused // twiddle multiplication by W_N^{k1 * n2} for the inter-stage factor. { - let threads_per_cube = (n1 / 2).clamp(1, MAX_UNITS_PER_CUBE); + let threads_per_cube = (n1 / 2).clamp(1, max_units); let log2_n1 = n1.trailing_zeros() as usize; let cube_dim = CubeDim::new_1d(threads_per_cube as u32); let cube_count = @@ -272,7 +263,7 @@ fn cfft_four_step_launch( // Step 2: contiguous FFT_{N2} along the n2 axis of (N1, N2). One cube // per (window, k1). Reads/writes scratch in place. { - let threads_per_cube = (n2 / 2).clamp(1, MAX_UNITS_PER_CUBE); + let threads_per_cube = (n2 / 2).clamp(1, max_units); let log2_n2 = n2.trailing_zeros() as usize; let cube_dim = CubeDim::new_1d(threads_per_cube as u32); let cube_count = @@ -345,8 +336,8 @@ fn cfft_four_step_radix1_kernel( let n2_idx = cube_pos - window * n2; let input_re_view = input_re.view(BatchSignalLayout::new(input_re, window, dim)); let input_im_view = input_im.view(BatchSignalLayout::new(input_im, window, dim)); - let mut scratch_re_view = scratch_re.view_mut(BatchSignalLayout::new(scratch_re, window, dim)); - let mut scratch_im_view = scratch_im.view_mut(BatchSignalLayout::new(scratch_im, window, dim)); + let scratch_re_view = scratch_re.view_mut(BatchSignalLayout::new(&*scratch_re, window, dim)); + let scratch_im_view = scratch_im.view_mut(BatchSignalLayout::new(&*scratch_im, window, dim)); let mut shared_re = SharedMemory::::new(n1); let mut shared_im = SharedMemory::::new(n1); @@ -357,8 +348,8 @@ fn cfft_four_step_radix1_kernel( while i < n1 { let j = bit_reverse(i, log2_n1); let flat = i * n2 + n2_idx; - shared_re[j] = input_re_view[flat]; - shared_im[j] = input_im_view[flat]; + shared_re[j] = input_re_view.read_checked(flat); + shared_im[j] = input_im_view.read_checked(flat); i += threads_per_cube; } sync_cube(); @@ -386,8 +377,8 @@ fn cfft_four_step_radix1_kernel( let ar = shared_re[k1]; let ai = shared_im[k1]; let flat = k1 * n2 + n2_idx; - scratch_re_view[flat] = w_re * ar - w_im * ai; - scratch_im_view[flat] = w_re * ai + w_im * ar; + 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; } } @@ -396,7 +387,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, @@ -415,8 +406,8 @@ fn cfft_four_step_radix2_kernel( let window = cube_pos / n1; let k1 = cube_pos - window * n1; let row_base = k1 * n2; - let mut scratch_re_view = scratch_re.view_mut(BatchSignalLayout::new(scratch_re, window, dim)); - let mut scratch_im_view = scratch_im.view_mut(BatchSignalLayout::new(scratch_im, window, dim)); + let scratch_re_view = scratch_re.view_mut(BatchSignalLayout::new(&*scratch_re, window, dim)); + let scratch_im_view = scratch_im.view_mut(BatchSignalLayout::new(&*scratch_im, window, dim)); let mut shared_re = SharedMemory::::new(n2); let mut shared_im = SharedMemory::::new(n2); @@ -424,8 +415,8 @@ fn cfft_four_step_radix2_kernel( let mut i = UNIT_POS as usize; while i < n2 { let j = bit_reverse(i, log2_n2); - shared_re[j] = scratch_re_view[row_base + i]; - shared_im[j] = scratch_im_view[row_base + i]; + shared_re[j] = scratch_re_view.read_checked(row_base + i); + shared_im[j] = scratch_im_view.read_checked(row_base + i); i += threads_per_cube; } sync_cube(); @@ -441,8 +432,8 @@ fn cfft_four_step_radix2_kernel( let mut k2 = UNIT_POS as usize; while k2 < n2 { - scratch_re_view[row_base + k2] = shared_re[k2]; - scratch_im_view[row_base + k2] = shared_im[k2]; + scratch_re_view.write_checked(row_base + k2, shared_re[k2]); + scratch_im_view.write_checked(row_base + k2, shared_im[k2]); k2 += threads_per_cube; } } @@ -474,13 +465,13 @@ fn cfft_four_step_transpose_kernel( let window = pos_u / m; let scratch_re_view = scratch_re.view(BatchSignalLayout::new(scratch_re, window, dim)); let scratch_im_view = scratch_im.view(BatchSignalLayout::new(scratch_im, window, dim)); - let mut output_re_view = output_re.view_mut(BatchSignalLayout::new(output_re, window, dim)); - let mut output_im_view = output_im.view_mut(BatchSignalLayout::new(output_im, window, dim)); + let output_re_view = output_re.view_mut(BatchSignalLayout::new(&*output_re, window, dim)); + let output_im_view = output_im.view_mut(BatchSignalLayout::new(&*output_im, window, dim)); // pos's inner index is the destination linear index k = k1 + k2 * N1. let k2 = inner / n1; let k1 = inner - k2 * n1; let src = k1 * n2 + k2; - output_re_view[inner] = scratch_re_view[src]; - output_im_view[inner] = scratch_im_view[src]; + output_re_view.write_checked(inner, scratch_re_view.read_checked(src)); + output_im_view.write_checked(inner, scratch_im_view.read_checked(src)); } 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..4a072dbd3 --- /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 = SharedMemory::::new(n_fft); + let mut shared_im = SharedMemory::::new(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 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 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 scratch_re_view = scratch_re.view_mut(crate::layout::BatchSignalLayout::new( + &*scratch_re, + window, + dim, + )); + let scratch_im_view = scratch_im.view_mut(crate::layout::BatchSignalLayout::new( + &*scratch_im, + window, + dim, + )); + let mut shared_re = SharedMemory::::new(n1); + let mut shared_im = SharedMemory::::new(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 output_re = output.view_mut(InterleavedBatchSignalLayout::new( + &*output, window, dim, 0usize, + )); + output_re.write_checked(inner, scratch_re_view.read_checked(src) * scale); + } + let 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 d6c8dec7c..3236791ec 100644 --- a/crates/cubek-fft/src/fft/fft_inner.rs +++ b/crates/cubek-fft/src/fft/fft_inner.rs @@ -64,9 +64,9 @@ fn bit_reverse_permutation( #[cube] /// Swap two elements of a 1D array. fn swap(view_1d: &mut View, i: usize, j: usize) { - let tmp = view_1d[i]; - view_1d[i] = view_1d[j]; - view_1d[j] = tmp; + let tmp = view_1d.read(i); + view_1d.write(i, view_1d.read(j)); + view_1d.write(j, tmp); } #[cube] @@ -91,25 +91,25 @@ 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 { let i0 = k + j; let i1 = i0 + half_m; - let a = (spectrum_re[i0], spectrum_im[i0]); - let b = (spectrum_re[i1], spectrum_im[i1]); + let a = (spectrum_re.read(i0), spectrum_im.read(i0)); + let b = (spectrum_re.read(i1), spectrum_im.read(i1)); let t = complex_mul::((w_re, w_im), b); let out0 = complex_add::(a, t); let out1 = complex_sub::(a, t); - spectrum_re[i0] = out0.0; - spectrum_im[i0] = out0.1; - spectrum_re[i1] = out1.0; - spectrum_im[i1] = out1.1; + spectrum_re.write(i0, out0.0); + spectrum_im.write(i0, out0.1); + spectrum_re.write(i1, out1.0); + spectrum_im.write(i1, out1.1); let new_w = complex_mul::((w_re, w_im), (wm_re, wm_in)); w_re = new_w.0; diff --git a/crates/cubek-fft/src/fft/irfft.rs b/crates/cubek-fft/src/fft/irfft.rs index 2dee78d11..430a2783c 100644 --- a/crates/cubek-fft/src/fft/irfft.rs +++ b/crates/cubek-fft/src/fft/irfft.rs @@ -9,14 +9,12 @@ use crate::{ fft::{ FftMode, fft_parallel::{bit_reverse, fft_butterfly_parallel}, - rfft::SHARED_MEM_CAP, + limits::{max_shared_fft_n, max_units_per_cube}, rfft_large::irfft_large_launch, }, layout::BatchSignalLayout, }; -const MAX_UNITS_PER_CUBE: usize = 256; - /// Inverse Real-valued Fast Fourier Transform. pub fn irfft( spectrum_re: TensorHandle, @@ -118,7 +116,7 @@ pub fn irfft_launch_padded( return Ok(()); } - if n_fft > SHARED_MEM_CAP { + if n_fft > max_shared_fft_n(client) { return irfft_large_launch::( client, spectrum_re, @@ -131,7 +129,7 @@ pub fn irfft_launch_padded( } let log2_n = n_fft.trailing_zeros() as usize; - let threads_per_cube = (n_fft / 2).clamp(1, MAX_UNITS_PER_CUBE); + let threads_per_cube = (n_fft / 2).clamp(1, max_units_per_cube(client)); let cube_dim = CubeDim::new_1d(threads_per_cube as u32); let cube_count = cubecl::calculate_cube_count_elemwise(client, count, CubeDim::new_single()); @@ -172,7 +170,7 @@ fn irfft_kernel( let spectrum_re_view = spectrum_re.view(BatchSignalLayout::new(spectrum_re, window_index, dim)); let spectrum_im_view = spectrum_im.view(BatchSignalLayout::new(spectrum_im, window_index, dim)); - let mut signal_view = signal.view_mut(BatchSignalLayout::new(signal, window_index, dim)); + let signal_view = signal.view_mut(BatchSignalLayout::new(&*signal, window_index, dim)); let mut shared_re = SharedMemory::::new(n_fft); let mut shared_im = SharedMemory::::new(n_fft); @@ -185,9 +183,17 @@ 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[src_bin], F::new(0.0)); - shared_im[dst] = select(active, spectrum_im_view[src_bin] * im_sign, 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_f32), + ); k += threads_per_cube; } sync_cube(); @@ -201,10 +207,10 @@ 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[i] = shared_re[i] * scale; + signal_view.write_checked(i, shared_re[i] * scale); i += threads_per_cube; } sync_cube(); 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..f48a5bc76 --- /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 signal_view = signal.view_mut(BatchSignalLayout::new(&*signal, window_index, dim)); + let mut shared_re = SharedMemory::::new(n_fft); + let mut shared_im = SharedMemory::::new(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 new file mode 100644 index 000000000..b6496ded5 --- /dev/null +++ b/crates/cubek-fft/src/fft/limits.rs @@ -0,0 +1,51 @@ +//! Per-device limits used when launching the FFT kernels. + +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. +/// +/// Every shared-memory FFT kernel in this crate allocates two +/// `SharedMemory::new(n_fft)` buffers (one for the real part, one for +/// the imaginary part), so the byte budget is +/// `2 * size_of::() * n_fft <= hardware.max_shared_memory_size`. +/// We floor to a power of two because the butterfly requires it. +pub(crate) fn max_shared_fft_n(client: &ComputeClient) -> usize { + let max_smem = client.properties().hardware.max_shared_memory_size; + let max_elems = max_smem / (2 * core::mem::size_of::()); + floor_power_of_two(max_elems) +} + +/// Hardware-reported maximum number of units (threads) per cube. +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() { + n + } else { + n.next_power_of_two() >> 1 + } +} diff --git a/crates/cubek-fft/src/fft/mod.rs b/crates/cubek-fft/src/fft/mod.rs index ec91a0a03..cad73d384 100644 --- a/crates/cubek-fft/src/fft/mod.rs +++ b/crates/cubek-fft/src/fft/mod.rs @@ -1,10 +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 b57db6a27..2e32a4855 100644 --- a/crates/cubek-fft/src/fft/rfft.rs +++ b/crates/cubek-fft/src/fft/rfft.rs @@ -9,14 +9,12 @@ use crate::{ fft::{ FftMode, fft_parallel::{bit_reverse, fft_butterfly_parallel}, + limits::{max_shared_fft_n, max_units_per_cube}, rfft_large::rfft_large_launch, }, layout::BatchSignalLayout, }; -const MAX_UNITS_PER_CUBE: usize = 256; -pub(crate) const SHARED_MEM_CAP: usize = 4096; - /// Real-valued Fast Fourier Transform. pub fn rfft( signal: TensorHandle, @@ -127,7 +125,7 @@ pub fn rfft_launch_padded( return Ok(()); } - if n_fft > SHARED_MEM_CAP { + if n_fft > max_shared_fft_n(client) { return rfft_large_launch::( client, signal, @@ -140,7 +138,7 @@ pub fn rfft_launch_padded( } let log2_n = n_fft.trailing_zeros() as usize; - let threads_per_cube = (n_fft / 2).clamp(1, MAX_UNITS_PER_CUBE); + let threads_per_cube = (n_fft / 2).clamp(1, max_units_per_cube(client)); let cube_dim = CubeDim::new_1d(threads_per_cube as u32); let cube_count = cubecl::calculate_cube_count_elemwise(client, count, CubeDim::new_single()); @@ -180,10 +178,10 @@ fn rfft_kernel( } let signal_view = signal.view(BatchSignalLayout::new(signal, window_index, dim)); - let mut spectrum_re_view = - spectrum_re.view_mut(BatchSignalLayout::new(spectrum_re, window_index, dim)); - let mut spectrum_im_view = - spectrum_im.view_mut(BatchSignalLayout::new(spectrum_im, window_index, dim)); + let spectrum_re_view = + spectrum_re.view_mut(BatchSignalLayout::new(&*spectrum_re, window_index, dim)); + let spectrum_im_view = + spectrum_im.view_mut(BatchSignalLayout::new(&*spectrum_im, window_index, dim)); let mut shared_re = SharedMemory::::new(n_fft); let mut shared_im = SharedMemory::::new(n_fft); @@ -193,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[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(); @@ -211,8 +209,8 @@ fn rfft_kernel( let n_freq = comptime![n_fft / 2 + 1]; let mut k = UNIT_POS as usize; while k < n_freq { - spectrum_re_view[k] = shared_re[k]; - spectrum_im_view[k] = shared_im[k]; + spectrum_re_view.write_checked(k, shared_re[k]); + spectrum_im_view.write_checked(k, shared_im[k]); k += 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..3ba5f6e7e --- /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 = SharedMemory::::new(n_fft); + let mut shared_im = SharedMemory::::new(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 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 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 9d3a360d8..0f6fb66d3 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`. @@ -273,16 +450,22 @@ fn rfft_pack_kernel( let k = pos % m; let window = pos / m; let signal_view = signal.view(BatchSignalLayout::new(signal, window, dim)); - 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)); + let packed_re_view = packed_re.view_mut(BatchSignalLayout::new(&*packed_re, window, dim)); + let packed_im_view = packed_im.view_mut(BatchSignalLayout::new(&*packed_im, window, dim)); let even = 2 * k; let odd = even + 1; let even_active = even < signal_len as usize; let odd_active = odd < signal_len as usize; let even = select(even_active, even, 0); let odd = select(odd_active, odd, 0); - packed_re_view[k] = select(even_active, signal_view[even], F::new(0.0)); - packed_im_view[k] = select(odd_active, signal_view[odd], F::new(0.0)); + packed_re_view.write_checked( + k, + 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_f32)), + ); } /// Recover `X[0..N/2+1]` from `Y[0..M]` for the packed-real forward path. @@ -315,26 +498,24 @@ fn rfft_post_kernel( 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 mut spectrum_re_view = - spectrum_re.view_mut(BatchSignalLayout::new(spectrum_re, window, dim)); - let mut spectrum_im_view = - spectrum_im.view_mut(BatchSignalLayout::new(spectrum_im, window, dim)); + let spectrum_re_view = spectrum_re.view_mut(BatchSignalLayout::new(&*spectrum_re, window, dim)); + let spectrum_im_view = spectrum_im.view_mut(BatchSignalLayout::new(&*spectrum_im, window, dim)); if k == 0 { - let y0_re = packed_re_view[0]; - let y0_im = packed_im_view[0]; - spectrum_re_view[k] = y0_re + y0_im; - spectrum_im_view[k] = F::new(0.0); + 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_f32)); } else if k == m { - let y0_re = packed_re_view[0]; - let y0_im = packed_im_view[0]; - spectrum_re_view[k] = y0_re - y0_im; - spectrum_im_view[k] = F::new(0.0); + 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_f32)); } else { - let a_re = packed_re_view[k]; - let a_im = packed_im_view[k]; - let b_re = packed_re_view[m - k]; - let b_im_raw = packed_im_view[m - k]; + 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_raw = packed_im_view.read_checked(m - k); let b_im = -b_im_raw; // conj(Y[M-k]) // Forward twiddle W_N^k = cos(-2π k / N) + i sin(-2π k / N). @@ -347,12 +528,79 @@ 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); - spectrum_re_view[k] = x_re; - spectrum_im_view[k] = x_im; + 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 spectrum_re = spectrum.view_mut(InterleavedBatchSignalLayout::new( + &*spectrum, window, dim, 0usize, + )); + spectrum_re.write_checked(k, x_re * scale); + } + { + let spectrum_im = spectrum.view_mut(InterleavedBatchSignalLayout::new( + &*spectrum, window, dim, 1usize, + )); + spectrum_im.write_checked(k, x_im * scale); } } @@ -387,26 +635,38 @@ fn irfft_pre_kernel( let window = pos / m; let spectrum_re_view = spectrum_re.view(BatchSignalLayout::new(spectrum_re, window, dim)); let spectrum_im_view = spectrum_im.view(BatchSignalLayout::new(spectrum_im, window, dim)); - 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)); + let packed_re_view = packed_re.view_mut(BatchSignalLayout::new(&*packed_re, window, dim)); + let 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_view[0]; + 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[xm], F::new(0.0)); - packed_re_view[k] = F::new(0.5) * (x0_re + xm_re); - packed_im_view[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[src], F::new(0.0)); - let x_im = select(active, spectrum_im_view[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[mirror], F::new(0.0)); - let xm_im_raw = select(mirror_active, spectrum_im_view[mirror], F::new(0.0)); + let xm_re = select( + mirror_active, + spectrum_re_view.read_checked(mirror), + F::new(0.0_f32), + ); + let xm_im_raw = select( + mirror_active, + spectrum_im_view.read_checked(mirror), + F::new(0.0_f32), + ); let xm_im = -xm_im_raw; // conj(X[M-k]) // Inverse twiddle W_N^{-k} = cos(2π k / N) + i sin(2π k / N). @@ -423,12 +683,78 @@ 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); - packed_re_view[k] = y_re; - packed_im_view[k] = y_im; + 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 packed_re_view = packed_re.view_mut(BatchSignalLayout::new(&*packed_re, window, dim)); + let 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); } } @@ -451,8 +777,39 @@ fn irfft_unpack_kernel( 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 scale = F::new(1.0) / F::cast_from(m); - signal_view[2 * k] = packed_re_view[k] * scale; - signal_view[2 * k + 1] = packed_im_view[k] * scale; + let signal_view = signal.view_mut(BatchSignalLayout::new(&*signal, window, dim)); + 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 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 8cc64a8f8..177f2d5c6 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(feature = "cpu-reference")] -pub mod cpu_reference; +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/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/irfft.rs b/crates/cubek-fft/tests/fft/irfft.rs new file mode 100644 index 000000000..277e1f227 --- /dev/null +++ b/crates/cubek-fft/tests/fft/irfft.rs @@ -0,0 +1,368 @@ +use cubecl::CubeElement; +use cubecl::{ + client::ComputeClient, + frontend::CubePrimitive, + prelude::StorageType, + std::tensor::TensorHandle, + {Runtime, TestRuntime}, +}; +use cubek_fft::{irfft_launch, irfft_launch_padded}; +use cubek_test_utils::{ + self, ExecutionOutcome, HostData, HostDataType, TestInput, TestOutcome, ValidationResult, + assert_equals_approx, launch_and_capture_outcome, +}; + +use cubek_fft::eval::cpu_reference::irfft_ref; + +fn test_launch(client: ComputeClient, spectrum_shape: Vec, dim: usize) { + let dtype = f32::as_type_native_unchecked().storage_type(); + let mut signal_shape = spectrum_shape.clone(); + signal_shape[dim] = (spectrum_shape[dim] - 1) * 2; + + let (random_spectrum_re_handle, random_spectrum_re_data) = + TestInput::builder(client.clone(), spectrum_shape.clone()) + .dtype(dtype) + .uniform(43, -1., 1.) + .generate_with_f32_host_data(); + + let (random_spectrum_im_handle, random_spectrum_im_data) = + TestInput::builder(client.clone(), spectrum_shape) + .dtype(dtype) + .uniform(44, -1., 1.) + .generate_with_f32_host_data(); + + let signal_handle = TestInput::builder(client.clone(), signal_shape) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + + let re_binding = random_spectrum_re_handle.binding(); + let im_binding = random_spectrum_im_handle.binding(); + let signal_binding = signal_handle.clone().binding(); + + let outcome = launch_and_capture_outcome(&client, |c| { + irfft_launch::(c, re_binding, im_binding, signal_binding, dim, dtype).into() + }); + + match outcome { + ExecutionOutcome::Executed => assert_irfft_result( + &client, + random_spectrum_re_data, + random_spectrum_im_data, + signal_handle, + dim, + ) + .as_test_outcome(), + ExecutionOutcome::CompileError(e) => TestOutcome::CompileError(e), + } + .enforce(); +} + +fn test_launch_padded( + client: ComputeClient, + spectrum_shape: Vec, + dim: usize, + n_fft: usize, +) { + let dtype = f32::as_type_native_unchecked().storage_type(); + let spec_bins = spectrum_shape[dim]; + let n_freq = n_fft / 2 + 1; + + let mut full_spectrum_shape = spectrum_shape.clone(); + full_spectrum_shape[dim] = n_freq; + let mut signal_shape = spectrum_shape.clone(); + signal_shape[dim] = n_fft; + + let virtual_re = tensor_from_data( + &client, + spectrum_shape.clone(), + &data_for_shape(&spectrum_shape), + dtype, + ); + let virtual_im = tensor_from_data( + &client, + spectrum_shape.clone(), + &data_for_shape(&spectrum_shape), + dtype, + ); + let padded_re = tensor_from_data( + &client, + full_spectrum_shape.clone(), + &padded_data(&spectrum_shape, dim, n_freq), + dtype, + ); + let padded_im = tensor_from_data( + &client, + full_spectrum_shape, + &padded_data(&spectrum_shape, dim, n_freq), + dtype, + ); + + let virtual_signal = empty_tensor(&client, signal_shape.clone(), dtype); + let padded_signal = empty_tensor(&client, signal_shape, dtype); + + let virtual_re_binding = virtual_re.binding(); + let virtual_im_binding = virtual_im.binding(); + let virtual_signal_binding = virtual_signal.clone().binding(); + let padded_re_binding = padded_re.binding(); + let padded_im_binding = padded_im.binding(); + let padded_signal_binding = padded_signal.clone().binding(); + + let outcome = launch_and_capture_outcome(&client, |c| { + if let Err(e) = irfft_launch_padded::( + c, + virtual_re_binding, + virtual_im_binding, + virtual_signal_binding, + dim, + spec_bins, + dtype, + ) { + return ExecutionOutcome::CompileError(format!("virtual launch failed: {e}")); + } + irfft_launch::( + c, + padded_re_binding, + padded_im_binding, + padded_signal_binding, + dim, + dtype, + ) + .into() + }); + + match outcome { + ExecutionOutcome::Executed => { + let actual = HostData::from_tensor_handle(&client, virtual_signal, HostDataType::F32); + let expected = HostData::from_tensor_handle(&client, padded_signal, HostDataType::F32); + assert_equals_approx(&actual, &expected, 1e-4).as_test_outcome() + } + ExecutionOutcome::CompileError(e) => TestOutcome::CompileError(e), + } + .enforce(); +} + +fn assert_irfft_result( + client: &ComputeClient, + spectrum_re: HostData, + spectrum_im: HostData, + signal: TensorHandle, + dim: usize, +) -> ValidationResult { + let epsilon = 0.01; + let expected_signal = irfft_ref(&spectrum_re, &spectrum_im, dim, None); + let actual_signal = HostData::from_tensor_handle(client, signal, HostDataType::F32); + + assert_equals_approx(&actual_signal, &expected_signal, epsilon) +} + +fn coords_from_index(mut index: usize, shape: &[usize]) -> Vec { + let mut coords = vec![0; shape.len()]; + for axis in (0..shape.len()).rev() { + coords[axis] = index % shape[axis]; + index /= shape[axis]; + } + coords +} + +fn sample_value(coords: &[usize]) -> f32 { + coords + .iter() + .enumerate() + .map(|(axis, coord)| (axis as f32 + 1.0) * (*coord as f32 + 0.25)) + .sum::() + .sin() +} + +fn data_for_shape(shape: &[usize]) -> Vec { + (0..shape.iter().product::()) + .map(|index| sample_value(&coords_from_index(index, shape))) + .collect() +} + +fn padded_data(shape: &[usize], dim: usize, target_len: usize) -> Vec { + let mut padded_shape = shape.to_vec(); + padded_shape[dim] = target_len; + + (0..padded_shape.iter().product::()) + .map(|index| { + let coords = coords_from_index(index, &padded_shape); + if coords[dim] < shape[dim] { + sample_value(&coords) + } else { + 0.0 + } + }) + .collect() +} + +fn tensor_from_data( + client: &ComputeClient, + shape: Vec, + data: &[f32], + dtype: StorageType, +) -> TensorHandle { + TensorHandle::::new_contiguous( + shape, + client.create_from_slice(f32::as_bytes(data)), + dtype, + ) +} + +fn empty_tensor( + client: &ComputeClient, + shape: Vec, + dtype: StorageType, +) -> TensorHandle { + let elems = shape.iter().product::(); + TensorHandle::::new_contiguous(shape, client.empty(elems * dtype.size()), dtype) +} + +#[test] +fn irfft_light_axis_last() { + let client = ::client(&Default::default()); + let spectrum_shape = [1, 5].to_vec(); + let dim = spectrum_shape.len() - 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +fn irfft_light_axis_1_strided() { + let client = ::client(&Default::default()); + let spectrum_shape = [2, 5, 1].to_vec(); + let dim = 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +fn irfft_light_axis_1_strided_trailing_batch() { + let client = ::client(&Default::default()); + let spectrum_shape = [3, 5, 2].to_vec(); + let dim = 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +fn irfft_light_axis_0_strided() { + let client = ::client(&Default::default()); + let spectrum_shape = [5, 2].to_vec(); + let dim = 0; + test_launch(client, spectrum_shape, dim); +} + +#[test] +fn irfft_light_axis_last_n16() { + let client = ::client(&Default::default()); + let spectrum_shape = [1, 9].to_vec(); + let dim = spectrum_shape.len() - 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +fn irfft_virtual_padding_axis_1_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + test_launch_padded(client, vec![2, 3, 3], 1, 8); +} + +#[test] +fn irfft_virtual_padding_dc_only_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + test_launch_padded(client, vec![2, 1, 3], 1, 8); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_3d_last_axis() { + let client = ::client(&Default::default()); + let spectrum_shape = [5, 2, 1025].to_vec(); + let dim = spectrum_shape.len() - 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_3d_axis_0() { + let client = ::client(&Default::default()); + let spectrum_shape = [33, 2, 1024].to_vec(); + let dim = 0; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_3d_axis_1() { + let client = ::client(&Default::default()); + let spectrum_shape = [33, 5, 1024].to_vec(); + let dim = 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_4d_axis_2() { + let client = ::client(&Default::default()); + let spectrum_shape = [12, 8, 513, 4].to_vec(); + let dim = 2; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_shared_memory_cap_axis_1_strided() { + let client = ::client(&Default::default()); + let spectrum_shape = [1, 2049, 1].to_vec(); + let dim = 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_large_axis_1_strided() { + let client = ::client(&Default::default()); + let spectrum_shape = [1, 4097, 1].to_vec(); + let dim = 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_four_step_axis_1_strided() { + let client = ::client(&Default::default()); + let spectrum_shape = [1, 8193, 1].to_vec(); + let dim = 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_batched_large_axis_last() { + let client = ::client(&Default::default()); + let spectrum_shape = [3, 4097].to_vec(); + let dim = spectrum_shape.len() - 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_large_virtual_padding_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + test_launch_padded(client, vec![1, 3000], 1, 8192); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_3d_batch_singleton_dim() { + let client = ::client(&Default::default()); + let spectrum_shape = [22, 1, 1025].to_vec(); + let dim = spectrum_shape.len() - 1; + test_launch(client, spectrum_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn irfft_dispatch_more_than_wgpu_x_axis_limit() { + let client = ::client(&Default::default()); + let spectrum_shape = [65_536, 2].to_vec(); + let dim = spectrum_shape.len() - 1; + test_launch(client, spectrum_shape, dim); +} diff --git a/crates/cubek-fft/tests/fft/mod.rs b/crates/cubek-fft/tests/fft/mod.rs new file mode 100644 index 000000000..437045e00 --- /dev/null +++ b/crates/cubek-fft/tests/fft/mod.rs @@ -0,0 +1,7 @@ +mod interleaved_cfft; +mod interleaved_irfft; +mod interleaved_rfft; +mod interleaved_validation; +mod irfft; +mod rfft; +mod round_trip; diff --git a/crates/cubek-fft/tests/fft/rfft.rs b/crates/cubek-fft/tests/fft/rfft.rs new file mode 100644 index 000000000..6c26a75c4 --- /dev/null +++ b/crates/cubek-fft/tests/fft/rfft.rs @@ -0,0 +1,473 @@ +use cubecl::CubeElement; +use cubecl::{ + client::ComputeClient, + frontend::CubePrimitive, + prelude::StorageType, + std::tensor::TensorHandle, + {Runtime, TestRuntime}, +}; +use cubek_fft::{rfft_launch, rfft_launch_padded}; +#[cfg(feature = "heavy")] +use cubek_test_utils::HostDataVec; +use cubek_test_utils::{ + self, ExecutionOutcome, HostData, HostDataType, TestInput, TestOutcome, ValidationResult, + assert_equals_approx, launch_and_capture_outcome, +}; + +use cubek_fft::eval::cpu_reference::rfft_ref; + +fn test_launch(client: ComputeClient, signal_shape: Vec, dim: usize) { + let dtype = f32::as_type_native_unchecked().storage_type(); + let mut spectrum_shape = signal_shape.clone(); + spectrum_shape[dim] = signal_shape[dim] / 2 + 1; + + let (white_noise_handle, white_noise_data) = + TestInput::builder(client.clone(), signal_shape.clone()) + .dtype(dtype) + .uniform(42, -1., 1.) + .generate_with_f32_host_data(); + + let spectrum_re_handle = TestInput::builder(client.clone(), spectrum_shape.to_vec()) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + + let spectrum_im_handle = TestInput::builder(client.clone(), spectrum_shape.to_vec()) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + + let signal_binding = white_noise_handle.binding(); + let re_binding = spectrum_re_handle.clone().binding(); + let im_binding = spectrum_im_handle.clone().binding(); + + let outcome = launch_and_capture_outcome(&client, |c| { + rfft_launch::(c, signal_binding, re_binding, im_binding, dim, dtype).into() + }); + + match outcome { + ExecutionOutcome::Executed => assert_rfft_result( + &client, + white_noise_data, + spectrum_re_handle, + spectrum_im_handle, + dim, + ) + .as_test_outcome(), + ExecutionOutcome::CompileError(e) => TestOutcome::CompileError(e), + } + .enforce(); +} + +fn test_launch_padded( + client: ComputeClient, + signal_shape: Vec, + dim: usize, + signal_len: usize, + n_fft: usize, +) { + let dtype = f32::as_type_native_unchecked().storage_type(); + let n_freq = n_fft / 2 + 1; + + let mut spectrum_shape = signal_shape.clone(); + spectrum_shape[dim] = n_freq; + let mut padded_shape = signal_shape.clone(); + padded_shape[dim] = n_fft; + + let virtual_signal = tensor_from_data( + &client, + signal_shape.clone(), + &data_for_shape_with_len(&signal_shape, dim, signal_len), + dtype, + ); + let padded_signal = tensor_from_data( + &client, + padded_shape, + &padded_data(&signal_shape, dim, signal_len, n_fft), + dtype, + ); + + let virtual_re = empty_tensor(&client, spectrum_shape.clone(), dtype); + let virtual_im = empty_tensor(&client, spectrum_shape.clone(), dtype); + let padded_re = empty_tensor(&client, spectrum_shape.clone(), dtype); + let padded_im = empty_tensor(&client, spectrum_shape, dtype); + + let virtual_signal_binding = virtual_signal.binding(); + let virtual_re_binding = virtual_re.clone().binding(); + let virtual_im_binding = virtual_im.clone().binding(); + let padded_signal_binding = padded_signal.binding(); + let padded_re_binding = padded_re.clone().binding(); + let padded_im_binding = padded_im.clone().binding(); + + let outcome = launch_and_capture_outcome(&client, |c| { + if let Err(e) = rfft_launch_padded::( + c, + virtual_signal_binding, + virtual_re_binding, + virtual_im_binding, + dim, + signal_len, + dtype, + ) { + return ExecutionOutcome::CompileError(format!("virtual launch failed: {e}")); + } + rfft_launch::( + c, + padded_signal_binding, + padded_re_binding, + padded_im_binding, + dim, + dtype, + ) + .into() + }); + + match outcome { + ExecutionOutcome::Executed => { + let actual_re = HostData::from_tensor_handle(&client, virtual_re, HostDataType::F32); + let actual_im = HostData::from_tensor_handle(&client, virtual_im, HostDataType::F32); + let expected_re = HostData::from_tensor_handle(&client, padded_re, HostDataType::F32); + let expected_im = HostData::from_tensor_handle(&client, padded_im, HostDataType::F32); + combine_re_im( + assert_equals_approx(&actual_re, &expected_re, 1e-4), + assert_equals_approx(&actual_im, &expected_im, 1e-4), + ) + .as_test_outcome() + } + ExecutionOutcome::CompileError(e) => TestOutcome::CompileError(e), + } + .enforce(); +} + +pub fn assert_rfft_result( + client: &ComputeClient, + signal: HostData, + spectrum_re: TensorHandle, + spectrum_im: TensorHandle, + dim: usize, +) -> ValidationResult { + // big epsilon because with wgpu, compute is less precise + let epsilon = 0.4; + let (expected_re, expected_im) = rfft_ref(&signal, dim, None); + + let actual_spectrum_re = HostData::from_tensor_handle(client, spectrum_re, HostDataType::F32); + let actual_spectrum_im = HostData::from_tensor_handle(client, spectrum_im, HostDataType::F32); + + combine_re_im( + assert_equals_approx(&actual_spectrum_re, &expected_re, epsilon), + assert_equals_approx(&actual_spectrum_im, &expected_im, epsilon), + ) +} + +fn combine_re_im(re: ValidationResult, im: ValidationResult) -> ValidationResult { + use ValidationResult::*; + match (re, im) { + (Fail(e), _) | (_, Fail(e)) => Fail(e), + (Error(e), _) | (_, Error(e)) => Error(e), + (Skipped(r1), Skipped(r2)) => Skipped(format!("{r1}, {r2}")), + (Skipped(r), Pass) | (Pass, Skipped(r)) => Skipped(r), + (Pass, Pass) => Pass, + } +} + +#[cfg(feature = "heavy")] +fn to_f32(host: HostData) -> Vec { + match host.data { + HostDataVec::F32(v) => v, + _ => panic!("expected f32 host data"), + } +} + +fn coords_from_index(mut index: usize, shape: &[usize]) -> Vec { + let mut coords = vec![0; shape.len()]; + for axis in (0..shape.len()).rev() { + coords[axis] = index % shape[axis]; + index /= shape[axis]; + } + coords +} + +fn sample_value(coords: &[usize]) -> f32 { + coords + .iter() + .enumerate() + .map(|(axis, coord)| (axis as f32 + 1.0) * (*coord as f32 + 0.25)) + .sum::() + .sin() +} + +fn data_for_shape_with_len(shape: &[usize], dim: usize, signal_len: usize) -> Vec { + (0..shape.iter().product::()) + .map(|index| { + let coords = coords_from_index(index, shape); + if coords[dim] < signal_len { + sample_value(&coords) + } else { + 0.0 + } + }) + .collect() +} + +fn padded_data(shape: &[usize], dim: usize, signal_len: usize, target_len: usize) -> Vec { + let mut padded_shape = shape.to_vec(); + padded_shape[dim] = target_len; + + (0..padded_shape.iter().product::()) + .map(|index| { + let coords = coords_from_index(index, &padded_shape); + if coords[dim] < signal_len { + sample_value(&coords) + } else { + 0.0 + } + }) + .collect() +} + +fn tensor_from_data( + client: &ComputeClient, + shape: Vec, + data: &[f32], + dtype: StorageType, +) -> TensorHandle { + TensorHandle::::new_contiguous( + shape, + client.create_from_slice(f32::as_bytes(data)), + dtype, + ) +} + +fn empty_tensor( + client: &ComputeClient, + shape: Vec, + dtype: StorageType, +) -> TensorHandle { + let elems = shape.iter().product::(); + TensorHandle::::new_contiguous(shape, client.empty(elems * dtype.size()), dtype) +} + +#[test] +fn rfft_light_axis_last() { + let client = ::client(&Default::default()); + let signal_shape = [1, 8].to_vec(); + let dim = signal_shape.len() - 1; + test_launch(client, signal_shape, dim); +} + +#[test] +fn rfft_light_axis_1_strided() { + let client = ::client(&Default::default()); + let signal_shape = [2, 8, 1].to_vec(); + let dim = 1; + test_launch(client, signal_shape, dim); +} + +#[test] +fn rfft_light_axis_1_strided_trailing_batch() { + let client = ::client(&Default::default()); + let signal_shape = [3, 8, 2].to_vec(); + let dim = 1; + test_launch(client, signal_shape, dim); +} + +#[test] +fn rfft_light_axis_0_strided() { + let client = ::client(&Default::default()); + let signal_shape = [8, 2].to_vec(); + let dim = 0; + test_launch(client, signal_shape, dim); +} + +#[test] +fn rfft_light_axis_last_n16() { + let client = ::client(&Default::default()); + let signal_shape = [1, 16].to_vec(); + let dim = signal_shape.len() - 1; + test_launch(client, signal_shape, dim); +} + +#[test] +fn rfft_virtual_padding_axis_1_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + test_launch_padded(client, vec![2, 5, 3], 1, 5, 8); +} + +#[test] +fn rfft_virtual_padding_ignores_tail_after_signal_len() { + let client = ::client(&Default::default()); + test_launch_padded(client, vec![2, 7, 3], 1, 5, 8); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_3d_axis_last() { + let client = ::client(&Default::default()); + let signal_shape = [5, 2, 2048].to_vec(); + let dim = signal_shape.len() - 1; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_3d_axis_1_strided() { + let client = ::client(&Default::default()); + let signal_shape = [5, 64, 1000].to_vec(); + let dim = 1; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_3d_axis_0_strided() { + let client = ::client(&Default::default()); + let signal_shape = [128, 6, 1000].to_vec(); + let dim = 0; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_4d_axis_1_strided() { + let client = ::client(&Default::default()); + let signal_shape = [5, 256, 6, 42].to_vec(); + let dim = 1; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_shared_memory_cap_axis_1_strided() { + let client = ::client(&Default::default()); + let signal_shape = [1, 4096, 1].to_vec(); + let dim = 1; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_large_axis_1_strided() { + let client = ::client(&Default::default()); + let signal_shape = [1, 8192, 1].to_vec(); + let dim = 1; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_four_step_axis_1_strided() { + let client = ::client(&Default::default()); + let signal_shape = [1, 16384, 1].to_vec(); + let dim = 1; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_batched_large_axis_last() { + let client = ::client(&Default::default()); + let signal_shape = [3, 8192].to_vec(); + let dim = signal_shape.len() - 1; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_large_virtual_padding_matches_materialized_zero_padding() { + let client = ::client(&Default::default()); + test_launch_padded(client, vec![1, 5000], 1, 5000, 8192); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_nyquist_bin_large_sizes() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + + for &n_fft in &[8192usize, 16384] { + let batch = 2; + let n_freq = n_fft / 2 + 1; + let signal_shape = [batch, n_fft].to_vec(); + let spectrum_shape = [batch, n_freq].to_vec(); + + let signal_data: Vec = (0..batch) + .flat_map(|_| (0..n_fft).map(|i| if i % 2 == 0 { 1.0 } else { -1.0 })) + .collect(); + let signal_handle = client.create_from_slice(f32::as_bytes(&signal_data)); + let signal = + TensorHandle::::new_contiguous(signal_shape, signal_handle, dtype); + let spectrum_re = TestInput::builder(client.clone(), spectrum_shape.clone()) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + let spectrum_im = TestInput::builder(client.clone(), spectrum_shape) + .dtype(dtype) + .zeros() + .generate_without_host_data(); + + let signal_binding = signal.binding(); + let re_binding = spectrum_re.clone().binding(); + let im_binding = spectrum_im.clone().binding(); + + let outcome = launch_and_capture_outcome(&client, |c| { + rfft_launch::(c, signal_binding, re_binding, im_binding, 1, dtype).into() + }); + + let outcome = match outcome { + ExecutionOutcome::Executed => { + let re = to_f32(HostData::from_tensor_handle( + &client, + spectrum_re, + HostDataType::F32, + )); + let im = to_f32(HostData::from_tensor_handle( + &client, + spectrum_im, + HostDataType::F32, + )); + let mut result = ValidationResult::Pass; + 'check: for b in 0..batch { + let base = b * n_freq; + for k in 0..n_freq { + let expected = if k == n_fft / 2 { n_fft as f32 } else { 0.0 }; + if (re[base + k] - expected).abs() >= 1.0 { + result = ValidationResult::Fail(format!( + "n_fft={n_fft}, batch={b}, bin={k}: real={}, want {expected}", + re[base + k] + )); + break 'check; + } + if im[base + k].abs() >= 1.0 { + result = ValidationResult::Fail(format!( + "n_fft={n_fft}, batch={b}, bin={k}: imag={}", + im[base + k] + )); + break 'check; + } + } + } + result.as_test_outcome() + } + ExecutionOutcome::CompileError(e) => TestOutcome::CompileError(e), + }; + outcome.enforce(); + } +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_3d_batch_singleton_dim() { + let client = ::client(&Default::default()); + let signal_shape = [22, 1, 2048].to_vec(); + let dim = signal_shape.len() - 1; + test_launch(client, signal_shape, dim); +} + +#[test] +#[cfg(feature = "heavy")] +fn rfft_dispatch_more_than_wgpu_x_axis_limit() { + let client = ::client(&Default::default()); + let signal_shape = [65_536, 2].to_vec(); + let dim = signal_shape.len() - 1; + test_launch(client, signal_shape, dim); +} diff --git a/crates/cubek-fft/tests/fft/round_trip.rs b/crates/cubek-fft/tests/fft/round_trip.rs new file mode 100644 index 000000000..4e0ff64ca --- /dev/null +++ b/crates/cubek-fft/tests/fft/round_trip.rs @@ -0,0 +1,32 @@ +#[cfg(feature = "heavy")] +use cubecl::{Runtime, TestRuntime, prelude::CubePrimitive}; +#[cfg(feature = "heavy")] +use cubek_fft::{irfft, rfft}; +//use cubefx_engine::{SignalSpec, phase_shift_effect}; +#[cfg(feature = "heavy")] +use cubek_test_utils::{HostData, TestInput, assert_equals_approx}; + +#[test] +#[cfg(feature = "heavy")] +fn large_fft_roundtrip() { + let client = ::client(&Default::default()); + let dtype = f32::as_type_native_unchecked().storage_type(); + + let shape = [431, 2, 2048]; + + let (original_signal, signal_data) = TestInput::builder(client.clone(), shape) + .dtype(dtype) + .uniform(42, -1., 1.) + .generate_with_f32_host_data(); + + let (spectrum_re, spectrum_im) = rfft(original_signal, shape.len() - 1, dtype); + let signal_back = irfft(spectrum_re, spectrum_im, shape.len() - 1, dtype); + + assert_equals_approx( + &HostData::from_tensor_handle(&client, signal_back, cubek_test_utils::HostDataType::F32), + &signal_data, + 0.03, + ) + .as_test_outcome() + .enforce(); +} diff --git a/crates/cubek-fft/tests/lib.rs b/crates/cubek-fft/tests/lib.rs index 7a4d0e3b2..4375d49f1 100644 --- a/crates/cubek-fft/tests/lib.rs +++ b/crates/cubek-fft/tests/lib.rs @@ -1 +1 @@ -mod suite; +mod fft; diff --git a/docs/superpowers/plans/2026-07-19-t4a-cubek-fft-git-compat.md b/docs/superpowers/plans/2026-07-19-t4a-cubek-fft-git-compat.md new file mode 100644 index 000000000..c6ffa76b9 --- /dev/null +++ b/docs/superpowers/plans/2026-07-19-t4a-cubek-fft-git-compat.md @@ -0,0 +1,65 @@ +# CubeK FFT Git Compatibility 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:** Make the merged interleaved F32/C32 CubeK FFT implementation available as a git dependency on the tensor4all CubeCL 0.10 fork used by tenferro. + +**Architecture:** Backport only the `cubek-fft` implementation and tests from CubeK main onto `release/t4a-0.2.0`, keep the package private/non-released, and adapt the small amount of newer CubeCL syntax to the 0.10 API (`SharedMemory` and 0.10 tensor views). Add the crate to the workspace so inherited metadata/dependencies resolve for git consumers. + +**Tech Stack:** Rust, CubeK FFT, `t4a-cubecl` 0.10, CubeCL CPU test runtime. + +## Global Constraints + +- Preserve the merged public interleaved APIs: `ComplexTensorHandle`, `FftNormalization`, `cfft_interleaved*`, `rfft_interleaved*`, and `irfft_interleaved*`. +- Support only F32 real and interleaved C32 complex tensors; do not add F64/C64. +- Preserve the same device-derived threshold for small and large CFFT/RFFT/IRFFT paths. +- Do not add standalone conversion or scaling kernels. +- Keep `cubek-fft` unpublished (`publish = false`); tenferro consumes this branch by git revision. +- Use the `t4a-cubecl` 0.10 workspace dependencies and `t4a-cubek-test-utils` 0.2 package. +- Do not backport the newer benchmark-catalog framework, which is not available in `t4a-cubek-test-utils` 0.2 and is not required by tenferro runtime execution. +- Do not merge or release from this task. + +--- + +### Task 1: Backport interleaved FFT to the t4a 0.2 dependency line + +**Files:** +- Modify: `Cargo.toml` +- Modify: `crates/cubek-fft/**` +- Create: `docs/superpowers/plans/2026-07-19-t4a-cubek-fft-git-compat.md` + +**Interfaces:** +- Produces a git-resolvable private `cubek-fft` package using the same `t4a-cubecl` types as tenferro. +- Produces the same public FFT API and F32/C32 semantics as CubeK PRs #7 and #8. + +- [ ] **Step 1: Add `crates/cubek-fft` to the t4a workspace** + + Add the crate to `workspace.members` and remove it from `workspace.exclude`. Keep its package name `cubek-fft`, set `publish = false`, point repository metadata at `tensor4all/cubek`, and use `t4a-cubek-test-utils =0.2.0`. + +- [ ] **Step 2: Backport the merged CubeK FFT source and tests** + + Bring the runtime portion of `crates/cubek-fft` from commit `e66fe2e60f9dfb4c77be5c68cc9979bbdbabc354`, including interleaved ABI validation, CFFT/RFFT/IRFFT kernels, normalization, small/large selection, CPU references, and correctness tests. Omit the newer benchmark-catalog modules that require post-0.2 test utilities. + +- [ ] **Step 3: Adapt only CubeCL API syntax** + + Replace newer `Shared<[F]>`/`Shared::new_slice` usage with CubeCL 0.10 `SharedMemory`/`SharedMemory::new`, and newer mutable-view aliases with `View`. Do not alter algorithms, thresholds, layouts, normalization, or launch signatures. + +- [ ] **Step 4: Verify compilation and correctness** + + Run: + + ```bash + cargo fmt --all -- --check + cargo check -p cubek-fft + cargo test -p cubek-fft + cargo clippy -p cubek-fft --all-targets --all-features -- -D warnings + ``` + + Expected: all commands pass, including small and large interleaved CPU-runtime tests. + +- [ ] **Step 5: Commit** + + ```bash + git add Cargo.toml Cargo.lock crates/cubek-fft docs/superpowers/plans/2026-07-19-t4a-cubek-fft-git-compat.md + git commit -m "feat(fft): backport interleaved APIs to t4a CubeCL" + ``` From 09135c96519bc28f6ffa5472a9c3fc7897775907 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Sun, 19 Jul 2026 09:10:00 +0900 Subject: [PATCH 2/5] fix(fft): restore interleaved CPU barriers --- .superpowers/sdd/task-1-fix-report.md | 23 +++++++++++++++++++ crates/cubek-fft/src/fft/cfft_interleaved.rs | 1 + crates/cubek-fft/src/fft/irfft_interleaved.rs | 1 + crates/cubek-fft/src/fft/rfft_interleaved.rs | 1 + 4 files changed, 26 insertions(+) create mode 100644 .superpowers/sdd/task-1-fix-report.md diff --git a/.superpowers/sdd/task-1-fix-report.md b/.superpowers/sdd/task-1-fix-report.md new file mode 100644 index 000000000..6b8576bd0 --- /dev/null +++ b/.superpowers/sdd/task-1-fix-report.md @@ -0,0 +1,23 @@ +# Task 1 review-fix report + +Status: DONE + +Review finding addressed: + +- Restored `sync_cube()` immediately after the final global-write loops in the + small/shared CFFT, RFFT, and IRFFT interleaved kernels. +- The three insertion points match upstream commit + `e66fe2e60f9dfb4c77be5c68cc9979bbdbabc354` (CubeK PR #8); no algorithm, + threshold, layout, or launch-interface changes were made. + +Verification: + +- Targeted small CFFT round-trip test: passed. +- Targeted RFFT reference test: passed. +- Targeted IRFFT reference test: passed. +- `cargo test -p cubek-fft --all-features`: passed (2 unit + 77 integration + doctests). +- `cargo clippy -p cubek-fft --all-targets --all-features -- -D warnings`: passed. +- `cargo fmt --all -- --check`: passed. +- `git diff --check`: passed. + +Concerns: none. diff --git a/crates/cubek-fft/src/fft/cfft_interleaved.rs b/crates/cubek-fft/src/fft/cfft_interleaved.rs index 4a072dbd3..d0c330357 100644 --- a/crates/cubek-fft/src/fft/cfft_interleaved.rs +++ b/crates/cubek-fft/src/fft/cfft_interleaved.rs @@ -376,6 +376,7 @@ fn cfft_interleaved_shared_kernel( output_im.write_checked(k, shared_im[k] * scale); k += threads_per_cube; } + sync_cube(); } /// First four-step pass over the strided N1 dimension of each C32 window. diff --git a/crates/cubek-fft/src/fft/irfft_interleaved.rs b/crates/cubek-fft/src/fft/irfft_interleaved.rs index f48a5bc76..aa9dd9ec9 100644 --- a/crates/cubek-fft/src/fft/irfft_interleaved.rs +++ b/crates/cubek-fft/src/fft/irfft_interleaved.rs @@ -299,4 +299,5 @@ fn irfft_interleaved_kernel( signal_view.write_checked(i, shared_re[i] * scale); 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 index 3ba5f6e7e..0eb9a7fd1 100644 --- a/crates/cubek-fft/src/fft/rfft_interleaved.rs +++ b/crates/cubek-fft/src/fft/rfft_interleaved.rs @@ -319,4 +319,5 @@ fn rfft_interleaved_kernel( spectrum_im.write_checked(k, shared_im[k] * scale); k += threads_per_cube; } + sync_cube(); } From fb5dbc8e994bb3023bfefe38abe0140f29cbb15e Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Sun, 19 Jul 2026 09:49:41 +0900 Subject: [PATCH 3/5] chore: pin host-visible CubeCL revision --- Cargo.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e72e2bf63..9fe411fe6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,8 +35,8 @@ version = "0.2.0" # PR/CI builds use the tensor4all CubeCL commit that adds the `t4a-cubecl*` # package manifests. Published packages use the registry packages at the same # exact versions. -cubecl = { package = "t4a-cubecl", git = "https://github.com/tensor4all/cubecl.git", rev = "6424d9da407d4bc10034aee3bae52705eb5d2da0", version = "=0.10.0", default-features = false } -cubecl-common = { package = "t4a-cubecl-common", git = "https://github.com/tensor4all/cubecl.git", rev = "6424d9da407d4bc10034aee3bae52705eb5d2da0", version = "=0.10.0", default-features = false } +cubecl = { package = "t4a-cubecl", git = "https://github.com/tensor4all/cubecl.git", rev = "11b52669f13e27bbe188f988fd696df6d989a562", version = "=0.10.0", default-features = false } +cubecl-common = { package = "t4a-cubecl-common", git = "https://github.com/tensor4all/cubecl.git", rev = "11b52669f13e27bbe188f988fd696df6d989a562", version = "=0.10.0", default-features = false } derive-new = { version = "0.7.0", default-features = false } log = { default-features = false, version = "0.4.22" } From bd6742618e26eaa6c6972415c8a4df46e98ff42d Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Sun, 19 Jul 2026 10:54:11 +0900 Subject: [PATCH 4/5] fix(matmul): preserve thread-safe error sources --- .../src/components/stage/unit/setup.rs | 2 +- .../src/components/tile/attention.rs | 14 +- .../src/routines/blackbox_accelerated.rs | 10 +- crates/cubek-attention/src/routines/unit.rs | 6 +- .../src/components/stage/reader.rs | 2 +- .../src/kernels/backward_data/launch.rs | 2 +- .../batch/gemv_plane_parallel/config.rs | 10 +- .../batch/gemv_plane_parallel/setup.rs | 96 +++++++----- .../batch/gemv_unit_perpendicular/setup.rs | 40 +++-- .../src/components/batch/naive/setup.rs | 22 +-- .../global/multi_stage/ordered/setup.rs | 8 +- .../global/read/strategy/async_full_cyclic.rs | 6 +- .../read/strategy/async_full_strided.rs | 6 +- .../read/strategy/async_partial_cyclic.rs | 6 +- .../read/strategy/async_partial_strided.rs | 6 +- .../components/global/read/strategy/base.rs | 42 ++++-- .../global/read/strategy/sync_full_cyclic.rs | 2 +- .../global/read/strategy/sync_full_strided.rs | 2 +- .../read/strategy/sync_partial_cyclic.rs | 2 +- .../global/single_stage/simple/setup.rs | 8 +- .../components/global/specialization/roles.rs | 2 +- .../stage/matmul/plane_partitioned/setup.rs | 19 ++- .../stage/matmul/unit_partitioned/setup.rs | 18 ++- .../src/components/stage/memory/layout.rs | 4 +- crates/cubek-matmul/src/components/tile.rs | 140 +++++++++++------- crates/cubek-matmul/src/definition/error.rs | 39 ++++- crates/cubek-matmul/src/launch/complex.rs | 46 ++++-- .../launch/launch_vecmat_plane_parallel.rs | 26 ++-- .../launch_vecmat_unit_perpendicular.rs | 34 +++-- crates/cubek-std/src/cube_dim_resource.rs | 2 +- crates/cubek-std/src/error.rs | 65 +++++++- crates/cubek-std/src/matrix_layout.rs | 4 +- 32 files changed, 450 insertions(+), 241 deletions(-) diff --git a/crates/cubek-attention/src/components/stage/unit/setup.rs b/crates/cubek-attention/src/components/stage/unit/setup.rs index 3215cdc51..b64928979 100644 --- a/crates/cubek-attention/src/components/stage/unit/setup.rs +++ b/crates/cubek-attention/src/components/stage/unit/setup.rs @@ -58,7 +58,7 @@ impl> StageAttentio CubeDimResource::Units(units * blueprint.tiling_scheme.stage_size.seq_q) } _ => { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Error: Expected unit tile attention, got a plane tile attention".to_string(), ))); } diff --git a/crates/cubek-attention/src/components/tile/attention.rs b/crates/cubek-attention/src/components/tile/attention.rs index 962808a25..7687556f8 100644 --- a/crates/cubek-attention/src/components/tile/attention.rs +++ b/crates/cubek-attention/src/components/tile/attention.rs @@ -184,7 +184,7 @@ fn validate_unit( let check_divisible = |dim: u32, vec_size: u32, name: &str, vec_name: &str| -> Result<(), AttentionSetupError> { if !dim.is_multiple_of(vec_size) { - return Err(AttentionSetupError::InvalidConfig(Box::new(format!( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new(format!( "Tile's {} ({:?}) must be divisible by {} vector size ({:?})", name, dim, vec_name, vec_size )))); @@ -223,7 +223,7 @@ fn validate_blackbox( dtypes: &AttentionElems, ) -> Result<(), AttentionSetupError> { if dtypes.query_global != dtypes.query_tile { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Query global and tile types must be the same because no stage to cast in between", ))); } @@ -268,7 +268,7 @@ fn validate_blackbox( } if line_sizes_mask > 1 { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Line size mask > 1 not supported yet on accelerated tile attention", ))); } @@ -278,13 +278,13 @@ fn validate_blackbox( let softmax_total = softmax_num_rows * softmax_num_cols; if !softmax_total.is_multiple_of(cfg.plane_dim) { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Softmax size should be divisible by plane dim", ))); } if cfg.inner_layout == InnerLayout::Contiguous && softmax_num_rows > cfg.plane_dim { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "More than one row per unit not supported with this inner layout", ))); } @@ -292,13 +292,13 @@ fn validate_blackbox( if cfg.inner_layout == InnerLayout::SplitRows && !softmax_total.is_multiple_of(2 * cfg.plane_dim) { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "With split rows, units must have two elements each", ))); } if cfg.tile_size.head_dim < cfg.tile_size.val_dim { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Can't have tile head_dim < tile val dim (not sure why)", ))); } diff --git a/crates/cubek-attention/src/routines/blackbox_accelerated.rs b/crates/cubek-attention/src/routines/blackbox_accelerated.rs index 6e01e1ef4..2285e9745 100644 --- a/crates/cubek-attention/src/routines/blackbox_accelerated.rs +++ b/crates/cubek-attention/src/routines/blackbox_accelerated.rs @@ -151,13 +151,13 @@ fn blueprint( .map_err(map_err)?; if tile_size_score_matmul.m != values_matmul.m { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Seq_q mismatch: `m` of score_matmul does not match `m` of values_matmul. ", ))); } if tile_size_score_matmul.n != values_matmul.k { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Seq_kv mismatch: `n` of score_matmul does not match `k` of values_matmul. ", ))); } @@ -208,13 +208,13 @@ fn validate( if !(problem.dims.seq_q as u32) .is_multiple_of(blueprint.tiling_scheme.elements_in_stage_seq_q()) { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Stage seq_q must divide problem seq_q".to_string(), ))); } if !(problem.dims.head_dim as u32).is_multiple_of(blueprint.tiling_scheme.tile_size.head_dim) { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Tile size head dim must divide problem head dim".to_string(), ))); } @@ -222,7 +222,7 @@ fn validate( if blueprint.tiling_scheme.partition_size.head_dim * blueprint.tiling_scheme.tile_size.head_dim != problem.dims.head_dim as u32 { - return Err(AttentionSetupError::InvalidConfig(Box::new(format!( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new(format!( "Tiling scheme's total head dim ({}) does not match problem's head dim ({})", blueprint.tiling_scheme.partition_size.head_dim * blueprint.tiling_scheme.tile_size.head_dim, diff --git a/crates/cubek-attention/src/routines/unit.rs b/crates/cubek-attention/src/routines/unit.rs index 79de0b127..de876fbc6 100644 --- a/crates/cubek-attention/src/routines/unit.rs +++ b/crates/cubek-attention/src/routines/unit.rs @@ -68,7 +68,7 @@ impl Routine for UnitRoutine { CubeDimResource::Units(units * blueprint.tiling_scheme.stage_size.seq_q) } _ => { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Error: Expected unit tile attention, got a plane tile attention".to_string(), ))); } @@ -136,7 +136,7 @@ fn validate( blueprint: AttentionBlueprint, ) -> Result { if !(problem.dims.head_dim as u32).is_multiple_of(blueprint.tiling_scheme.tile_size.head_dim) { - return Err(AttentionSetupError::InvalidConfig(Box::new( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Tile size head dim must divide problem head dim".to_string(), ))); } @@ -144,7 +144,7 @@ fn validate( if blueprint.tiling_scheme.partition_size.head_dim * blueprint.tiling_scheme.tile_size.head_dim != problem.dims.head_dim as u32 { - return Err(AttentionSetupError::InvalidConfig(Box::new(format!( + return Err(AttentionSetupError::InvalidConfig(cubek_std::InvalidConfigError::new(format!( "Tiling scheme's total head dim ({}) does not match problem's head dim ({})", blueprint.tiling_scheme.partition_size.head_dim * blueprint.tiling_scheme.tile_size.head_dim, diff --git a/crates/cubek-convolution/src/components/stage/reader.rs b/crates/cubek-convolution/src/components/stage/reader.rs index ae06b1d48..4f6a814fe 100644 --- a/crates/cubek-convolution/src/components/stage/reader.rs +++ b/crates/cubek-convolution/src/components/stage/reader.rs @@ -47,7 +47,7 @@ impl TilingValidation for BiasTilingLayout { fn check(config: StageMemoryConfig) -> Result<(), InvalidConfigError> { let stage_width = config.elements_per_stage_along_col(); if config.vector_size > stage_width { - return Err(Box::new(format!( + return Err(cubek_std::InvalidConfigError::new(format!( "Invalid vector size. Got {:?} which should not be >{:?}", config.vector_size, stage_width, ))); diff --git a/crates/cubek-convolution/src/kernels/backward_data/launch.rs b/crates/cubek-convolution/src/kernels/backward_data/launch.rs index 813636c53..74dd8fd06 100644 --- a/crates/cubek-convolution/src/kernels/backward_data/launch.rs +++ b/crates/cubek-convolution/src/kernels/backward_data/launch.rs @@ -195,7 +195,7 @@ where /// rejected. #[allow(dead_code)] pub(crate) fn unsupported_tma_error() -> ConvSetupError { - ConvSetupError::Matmul(MatmulSetupError::InvalidConfig(Box::new( + ConvSetupError::Matmul(MatmulSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Data backprop doesn't yet work with current TMA tiling strategy", ))) } diff --git a/crates/cubek-matmul/src/components/batch/gemv_plane_parallel/config.rs b/crates/cubek-matmul/src/components/batch/gemv_plane_parallel/config.rs index 71cdb76ba..be72d758e 100644 --- a/crates/cubek-matmul/src/components/batch/gemv_plane_parallel/config.rs +++ b/crates/cubek-matmul/src/components/batch/gemv_plane_parallel/config.rs @@ -32,10 +32,12 @@ impl GemvKind { MatrixLayout::RowMajor => GemvKind::MatVecRowMajor, }) } else { - Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Problem is not a valid GEMV, got (m,n,k)=({:?},{:?},{:?})", - problem.m, problem.n, problem.k - )))) + Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Problem is not a valid GEMV, got (m,n,k)=({:?},{:?},{:?})", + problem.m, problem.n, problem.k + )), + )) } } } diff --git a/crates/cubek-matmul/src/components/batch/gemv_plane_parallel/setup.rs b/crates/cubek-matmul/src/components/batch/gemv_plane_parallel/setup.rs index 5089e79d9..2d7e911aa 100644 --- a/crates/cubek-matmul/src/components/batch/gemv_plane_parallel/setup.rs +++ b/crates/cubek-matmul/src/components/batch/gemv_plane_parallel/setup.rs @@ -143,41 +143,51 @@ impl BatchMatmulFamily<()> for GemvPlaneParallelFamily { vector_sizes: &MatmulVectorSizes, ) -> Result<(), MatmulSetupError> { if vector_sizes.lhs != vector_sizes.rhs { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Lhs and Rhs vector sizes must be equal, got lhs:{:?}, rhs:{:?}", - vector_sizes.lhs, vector_sizes.rhs - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Lhs and Rhs vector sizes must be equal, got lhs:{:?}, rhs:{:?}", + vector_sizes.lhs, vector_sizes.rhs + )), + )); } if vector_sizes.out != 1 { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Out vector size must be 1, got {:?}", - vector_sizes.out, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Out vector size must be 1, got {:?}", + vector_sizes.out, + )), + )); } let plane_dim = client.properties().hardware.plane_size_max as usize; if blueprint.tile_dim != plane_dim * vector_sizes.lhs { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile dim must equal plane_dim * vector_size, got {:?} != {:?} * {:?}", - blueprint.tile_dim, plane_dim, vector_sizes.lhs, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile dim must equal plane_dim * vector_size, got {:?} != {:?} * {:?}", + blueprint.tile_dim, plane_dim, vector_sizes.lhs, + )), + )); } if !problem.k.is_multiple_of(blueprint.tile_dim) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Problem dimensions k={:?} must be divisible by tile dim ({:?})", - problem.k, blueprint.tile_dim, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Problem dimensions k={:?} must be divisible by tile dim ({:?})", + problem.k, blueprint.tile_dim, + )), + )); } match blueprint.kind { GemvKind::VecMatRowMajor => { if !problem.n.is_multiple_of(blueprint.tile_dim) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "For VecMatTransposeSwap, problem.n ({:?}) must be divisible by tile_dim ({:?})", - problem.n, blueprint.tile_dim, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "For VecMatTransposeSwap, problem.n ({:?}) must be divisible by tile_dim ({:?})", + problem.n, blueprint.tile_dim, + )), + )); } if blueprint.tile_dim @@ -186,22 +196,26 @@ impl BatchMatmulFamily<()> for GemvPlaneParallelFamily { * vector_sizes.rhs > client.properties().hardware.max_shared_memory_size { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Requesting too much shared memory, requested {:?}, max {:?}", - blueprint.tile_dim - * blueprint.tile_dim - * dtypes.rhs_global.size() - * vector_sizes.rhs, - client.properties().hardware.max_shared_memory_size - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Requesting too much shared memory, requested {:?}, max {:?}", + blueprint.tile_dim + * blueprint.tile_dim + * dtypes.rhs_global.size() + * vector_sizes.rhs, + client.properties().hardware.max_shared_memory_size + )), + )); } } GemvKind::MatVecColMajor => { if !problem.m.is_multiple_of(blueprint.tile_dim) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "For MatVecTransposeSwap, problem.m ({:?}) must be divisible by tile_dim ({:?})", - problem.m, blueprint.tile_dim, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "For MatVecTransposeSwap, problem.m ({:?}) must be divisible by tile_dim ({:?})", + problem.m, blueprint.tile_dim, + )), + )); } if blueprint.tile_dim @@ -210,14 +224,16 @@ impl BatchMatmulFamily<()> for GemvPlaneParallelFamily { * vector_sizes.lhs > client.properties().hardware.max_shared_memory_size { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Requesting too much shared memory, requested {:?}, max {:?}", - blueprint.tile_dim - * blueprint.tile_dim - * dtypes.lhs_global.size() - * vector_sizes.lhs, - client.properties().hardware.max_shared_memory_size - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Requesting too much shared memory, requested {:?}, max {:?}", + blueprint.tile_dim + * blueprint.tile_dim + * dtypes.lhs_global.size() + * vector_sizes.lhs, + client.properties().hardware.max_shared_memory_size + )), + )); } } _ => {} diff --git a/crates/cubek-matmul/src/components/batch/gemv_unit_perpendicular/setup.rs b/crates/cubek-matmul/src/components/batch/gemv_unit_perpendicular/setup.rs index 52d85fac3..ce60d0562 100644 --- a/crates/cubek-matmul/src/components/batch/gemv_unit_perpendicular/setup.rs +++ b/crates/cubek-matmul/src/components/batch/gemv_unit_perpendicular/setup.rs @@ -147,34 +147,42 @@ impl BatchMatmulFamily<()> for VecMatUnitPerpendicularFamily { ) -> Result<(), MatmulSetupError> { let vector_size = vector_sizes.lhs; if !(vector_size == vector_sizes.rhs && vector_size == vector_sizes.out) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "All vector sizes must be equal, got lhs:{:?}, rhs:{:?}, out:{:?}", - vector_size, vector_sizes.rhs, vector_sizes.out - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "All vector sizes must be equal, got lhs:{:?}, rhs:{:?}, out:{:?}", + vector_size, vector_sizes.rhs, vector_sizes.out + )), + )); } let plane_dim = client.properties().hardware.plane_size_max as usize; if blueprint.tile_dim != plane_dim * vector_size { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile dim must equal plane_dim * vector_size, got {:?} != {:?} * {:?}", - blueprint.tile_dim, plane_dim, vector_size, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile dim must equal plane_dim * vector_size, got {:?} != {:?} * {:?}", + blueprint.tile_dim, plane_dim, vector_size, + )), + )); } if !problem.k.is_multiple_of(vector_size) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Problem dimension k={:?} must be divisible by vector size ({:?})", - problem.k, vector_size, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Problem dimension k={:?} must be divisible by vector size ({:?})", + problem.k, vector_size, + )), + )); } let aligned_k = problem.k.is_multiple_of(blueprint.tile_dim); let aligned_n = problem.n.is_multiple_of(blueprint.tile_dim); if (!aligned_k || !aligned_n) && blueprint.check_bounds != CheckBounds::Checked { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Problem dimensions n={:?}, k={:?} not divisible by tile dim ({:?}) require CheckBounds::Checked", - problem.n, problem.k, blueprint.tile_dim, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Problem dimensions n={:?}, k={:?} not divisible by tile dim ({:?}) require CheckBounds::Checked", + problem.n, problem.k, blueprint.tile_dim, + )), + )); } Ok(()) diff --git a/crates/cubek-matmul/src/components/batch/naive/setup.rs b/crates/cubek-matmul/src/components/batch/naive/setup.rs index 7a02feb3d..012dffe63 100644 --- a/crates/cubek-matmul/src/components/batch/naive/setup.rs +++ b/crates/cubek-matmul/src/components/batch/naive/setup.rs @@ -133,9 +133,9 @@ impl BatchMatmulFamily<()> for NaiveBatchMatmulFamily { vector_sizes: &MatmulVectorSizes, ) -> Result<(), MatmulSetupError> { if blueprint.vector_size_out > 1 { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Vector size on output not supported", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new("Vector size on output not supported"), + )); } if let Some(scheme) = problem.lhs_scheme @@ -145,9 +145,11 @@ impl BatchMatmulFamily<()> for NaiveBatchMatmulFamily { let block_size = block_size.to_dim_vec(2.max(block_size.len())); let block_width = block_size[block_size.len() - 1] as usize; if !block_width.is_multiple_of(vector_size) { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Block size isn't a multiple of load size on lhs", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "Block size isn't a multiple of load size on lhs", + ), + )); } } @@ -158,9 +160,11 @@ impl BatchMatmulFamily<()> for NaiveBatchMatmulFamily { let block_size = block_size.to_dim_vec(2.max(block_size.len())); let block_width = block_size[block_size.len() - 2] as usize; if !block_width.is_multiple_of(vector_size) { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Block size isn't a multiple of load size on rhs", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "Block size isn't a multiple of load size on rhs", + ), + )); } } diff --git a/crates/cubek-matmul/src/components/global/multi_stage/ordered/setup.rs b/crates/cubek-matmul/src/components/global/multi_stage/ordered/setup.rs index efeb703ab..3fd799854 100644 --- a/crates/cubek-matmul/src/components/global/multi_stage/ordered/setup.rs +++ b/crates/cubek-matmul/src/components/global/multi_stage/ordered/setup.rs @@ -215,9 +215,11 @@ where RL::validate_with_problem(problem, dtypes, StageIdent::Rhs)?; if blueprint.tiling_scheme.partitions_per_stage_along_n() > 1 { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Ordered does not support number of stage partitions > 1 in n", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "Ordered does not support number of stage partitions > 1 in n", + ), + )); } SMM::validate_blueprint(client, blueprint, dtypes, vector_sizes) diff --git a/crates/cubek-matmul/src/components/global/read/strategy/async_full_cyclic.rs b/crates/cubek-matmul/src/components/global/read/strategy/async_full_cyclic.rs index fae153946..1de93e754 100644 --- a/crates/cubek-matmul/src/components/global/read/strategy/async_full_cyclic.rs +++ b/crates/cubek-matmul/src/components/global/read/strategy/async_full_cyclic.rs @@ -50,7 +50,7 @@ impl LoadingValidation for AsyncFullCyclicLoading { let total_units = config.loading_units_count(); if !num_stage_vectors.is_multiple_of(total_units) { - return Err(Box::new(format!( + return Err(cubek_std::InvalidConfigError::new(format!( "Too many data will be loaded, resulting in out of bounds. Try setting vector size and number of planes so that total unit count {total_units:?} divides number of vectors in stage.", ))); @@ -63,7 +63,9 @@ impl LoadingValidation for AsyncFullCyclicLoading { .elements_per_tile_along_contiguous_dim() .is_multiple_of(vector_size) { - return Err(Box::new("Tile size isn't divisible by copy vector size")); + return Err(cubek_std::InvalidConfigError::new( + "Tile size isn't divisible by copy vector size", + )); } validate_swizzle_atom_size(config.smem_config)?; diff --git a/crates/cubek-matmul/src/components/global/read/strategy/async_full_strided.rs b/crates/cubek-matmul/src/components/global/read/strategy/async_full_strided.rs index fa78809b7..31b30772d 100644 --- a/crates/cubek-matmul/src/components/global/read/strategy/async_full_strided.rs +++ b/crates/cubek-matmul/src/components/global/read/strategy/async_full_strided.rs @@ -42,14 +42,16 @@ impl LoadingValidation for AsyncFullStridedLoading { .elements_per_stage_along_contiguous_dim() .is_multiple_of(vector_size) { - return Err(Box::new("Stage size isn't divisible by copy vector size")); + return Err(cubek_std::InvalidConfigError::new( + "Stage size isn't divisible by copy vector size", + )); } let num_stage_vectors = config.smem_config.elements_per_stage() / vector_size; let total_units = config.loading_units_count(); if !num_stage_vectors.is_multiple_of(total_units) { - return Err(Box::new(format!( + return Err(cubek_std::InvalidConfigError::new(format!( "Too many data will be loaded, resulting in out of bounds. Try setting vector size and number of planes so that total unit count {total_units:?} divides number of vectors in stage.", ))); diff --git a/crates/cubek-matmul/src/components/global/read/strategy/async_partial_cyclic.rs b/crates/cubek-matmul/src/components/global/read/strategy/async_partial_cyclic.rs index 57147af22..7aa41db2e 100644 --- a/crates/cubek-matmul/src/components/global/read/strategy/async_partial_cyclic.rs +++ b/crates/cubek-matmul/src/components/global/read/strategy/async_partial_cyclic.rs @@ -73,7 +73,7 @@ impl LoadingValidation for AsyncPartialCyclicLoading { let num_stage_elements = config.smem_config.elements_per_stage(); if max_position > num_stage_elements { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Too many data will be loaded, resulting in out-of-bounds", )); } @@ -85,7 +85,9 @@ impl LoadingValidation for AsyncPartialCyclicLoading { .elements_per_tile_along_contiguous_dim() .is_multiple_of(vector_size) { - return Err(Box::new("Tile size isn't divisible by copy vector size")); + return Err(cubek_std::InvalidConfigError::new( + "Tile size isn't divisible by copy vector size", + )); } validate_swizzle_atom_size(config.smem_config)?; diff --git a/crates/cubek-matmul/src/components/global/read/strategy/async_partial_strided.rs b/crates/cubek-matmul/src/components/global/read/strategy/async_partial_strided.rs index fd9d20f74..430e25872 100644 --- a/crates/cubek-matmul/src/components/global/read/strategy/async_partial_strided.rs +++ b/crates/cubek-matmul/src/components/global/read/strategy/async_partial_strided.rs @@ -48,14 +48,16 @@ impl LoadingValidation for AsyncPartialStridedLoading { .elements_per_stage_along_contiguous_dim() .is_multiple_of(vector_size) { - return Err(Box::new("Stage size isn't divisible by copy vector size")); + return Err(cubek_std::InvalidConfigError::new( + "Stage size isn't divisible by copy vector size", + )); } let num_stage_vectors = config.smem_config.elements_per_stage() / vector_size; let total_units = config.loading_units_count(); if !num_stage_vectors.is_multiple_of(total_units) { - return Err(Box::new(format!( + return Err(cubek_std::InvalidConfigError::new(format!( "Too many data will be loaded, resulting in out of bounds. Try setting vector size and number of planes so that total unit count {total_units:?} divides number of vectors in stage.", ))); diff --git a/crates/cubek-matmul/src/components/global/read/strategy/base.rs b/crates/cubek-matmul/src/components/global/read/strategy/base.rs index 570613a40..9e49d551f 100644 --- a/crates/cubek-matmul/src/components/global/read/strategy/base.rs +++ b/crates/cubek-matmul/src/components/global/read/strategy/base.rs @@ -79,7 +79,7 @@ pub fn validate_async_barrier(device_props: &DeviceProperties) -> Result<(), Inv .features .supports_type(OpaqueType::Barrier(BarrierLevel::Cube)) { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Async barrier instructions are not available on the current device", )); } @@ -94,13 +94,13 @@ pub fn validate_async_copy( dtype_stage: &StorageType, ) -> Result<(), InvalidConfigError> { if !device_props.features.copy_async { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Async copy instructions are not available on the current device", )); } if dtype_global.size() != dtype_stage.size() { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Async copy requires stage and global types to be the same", )); } @@ -108,7 +108,7 @@ pub fn validate_async_copy( if matches!(dtype_global, StorageType::Packed(_, _)) && !matches!(dtype_stage, StorageType::Packed(_, _)) { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Async copy doesn't support dequantizing on global read", )); } @@ -119,7 +119,9 @@ pub fn validate_async_copy( /// Validates if swizzling is disabled, for loaders that can't support it. pub fn validate_noswizzle(config: StageMemoryConfig) -> Result<(), InvalidConfigError> { if config.swizzle != SwizzleMode::None { - return Err(Box::new("This loader doesn't support swizzling")); + return Err(cubek_std::InvalidConfigError::new( + "This loader doesn't support swizzling", + )); } Ok(()) @@ -134,7 +136,9 @@ pub fn validate_swizzle_atom_size(config: StageMemoryConfig) -> Result<(), Inval let vector_bytes = config.dtype.size() * config.vector_size as usize; if vector_bytes > config.swizzle.atom_size() { - return Err(Box::new("Load atom can't be larger than swizzle atom")); + return Err(cubek_std::InvalidConfigError::new( + "Load atom can't be larger than swizzle atom", + )); } Ok(()) @@ -148,7 +152,7 @@ pub fn validate_tma( global_dtype: &StorageType, ) -> Result<(), InvalidConfigError> { if !device_props.features.supports_type(SemanticType::TensorMap) { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Tensor memory accelerator features are not available on the current device", )); } @@ -156,7 +160,7 @@ pub fn validate_tma( let stage_dtype = smem_config.dtype; if global_dtype.size() != stage_dtype.size() { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "TMA requires stage and global types to be the same", )); } @@ -164,7 +168,9 @@ pub fn validate_tma( if matches!(global_dtype, StorageType::Packed(_, _)) && !matches!(stage_dtype, StorageType::Packed(_, _)) { - return Err(Box::new("TMA doesn't support dequantizing on global read")); + return Err(cubek_std::InvalidConfigError::new( + "TMA doesn't support dequantizing on global read", + )); } if matches!(smem_config.swizzle, SwizzleMode::None) { @@ -180,7 +186,9 @@ pub fn validate_tma( // Slightly tighter than the actual requirements, but simple enough and is always followed by // selection. Getting illegal memory access if this isn't followed for some reason. if row_bytes as usize != smem_config.swizzle.span_size() { - return Err(Box::new("Swizzling size must be equal to row size for TMA")); + return Err(cubek_std::InvalidConfigError::new( + "Swizzling size must be equal to row size for TMA", + )); } Ok(()) @@ -198,7 +206,7 @@ pub fn validate_async_copy_with_problem( }; if is_quantized { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Async copy doesn't support dequantizing on global read", )); } @@ -210,7 +218,7 @@ pub fn validate_async_copy_with_problem( }; if stride_align_bits(strides, layout, &dtypes.global(ident.into())) < 4 { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Async copy requires strides to be aligned to 16 bytes", )); } @@ -230,7 +238,9 @@ pub fn validate_tma_with_problem( }; if is_quantized { - return Err(Box::new("TMA doesn't support dequantizing on global read")); + return Err(cubek_std::InvalidConfigError::new( + "TMA doesn't support dequantizing on global read", + )); } let (strides, layout) = match ident { @@ -240,14 +250,16 @@ pub fn validate_tma_with_problem( }; if stride_align_bits(strides, layout, &dtypes.global(ident.into())) < 4 { - return Err(Box::new("TMA requires strides to be aligned to 16 bytes")); + return Err(cubek_std::InvalidConfigError::new( + "TMA requires strides to be aligned to 16 bytes", + )); } if problem.lhs_batches != problem.rhs_batches && problem.lhs_batches.iter().product::() != 1 && problem.rhs_batches.iter().product::() != 1 { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "TMA doesn't support mixing broadcast and non-broadcast dims", )); } diff --git a/crates/cubek-matmul/src/components/global/read/strategy/sync_full_cyclic.rs b/crates/cubek-matmul/src/components/global/read/strategy/sync_full_cyclic.rs index 82a3e0ee6..2e30b5705 100644 --- a/crates/cubek-matmul/src/components/global/read/strategy/sync_full_cyclic.rs +++ b/crates/cubek-matmul/src/components/global/read/strategy/sync_full_cyclic.rs @@ -35,7 +35,7 @@ impl LoadingValidation for SyncFullCyclicLoading { let total_units = config.loading_units_count(); if !num_stage_vectors.is_multiple_of(total_units) { - return Err(Box::new(format!( + return Err(cubek_std::InvalidConfigError::new(format!( "Too many data will be loaded, resulting in out of bounds. Try setting vector size and number of planes so that total unit count {total_units:?} divides number of vectors in stage.", ))); diff --git a/crates/cubek-matmul/src/components/global/read/strategy/sync_full_strided.rs b/crates/cubek-matmul/src/components/global/read/strategy/sync_full_strided.rs index e28c98c17..494531e01 100644 --- a/crates/cubek-matmul/src/components/global/read/strategy/sync_full_strided.rs +++ b/crates/cubek-matmul/src/components/global/read/strategy/sync_full_strided.rs @@ -29,7 +29,7 @@ impl LoadingValidation for SyncFullStridedLoading { let total_units = config.loading_units_count(); if !num_stage_vectors.is_multiple_of(total_units) { - return Err(Box::new(format!( + return Err(cubek_std::InvalidConfigError::new(format!( "Too many data will be loaded, resulting in out of bounds. Try setting vector size and number of planes so that total unit count {total_units:?} divides number of vectors in stage.", ))); diff --git a/crates/cubek-matmul/src/components/global/read/strategy/sync_partial_cyclic.rs b/crates/cubek-matmul/src/components/global/read/strategy/sync_partial_cyclic.rs index 8e0fba4f9..fe5e30aa6 100644 --- a/crates/cubek-matmul/src/components/global/read/strategy/sync_partial_cyclic.rs +++ b/crates/cubek-matmul/src/components/global/read/strategy/sync_partial_cyclic.rs @@ -46,7 +46,7 @@ impl LoadingValidation for SyncPartialCyclicLoading { let num_stage_elements = config.smem_config.elements_per_stage(); if max_position > num_stage_elements { - return Err(Box::new( + return Err(cubek_std::InvalidConfigError::new( "Too many data will be loaded, resulting in out-of-bounds", )); } diff --git a/crates/cubek-matmul/src/components/global/single_stage/simple/setup.rs b/crates/cubek-matmul/src/components/global/single_stage/simple/setup.rs index 8bf3edddd..174d703a0 100644 --- a/crates/cubek-matmul/src/components/global/single_stage/simple/setup.rs +++ b/crates/cubek-matmul/src/components/global/single_stage/simple/setup.rs @@ -186,9 +186,11 @@ where let resources = if !blueprint.load_flows.has_specialization() { SMM::cubedim_resource(blueprint) } else { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Specialization is unavailable for simple matmul.", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "Specialization is unavailable for simple matmul.", + ), + )); }?; Ok(resources) diff --git a/crates/cubek-matmul/src/components/global/specialization/roles.rs b/crates/cubek-matmul/src/components/global/specialization/roles.rs index 8ae0bca25..647898e84 100644 --- a/crates/cubek-matmul/src/components/global/specialization/roles.rs +++ b/crates/cubek-matmul/src/components/global/specialization/roles.rs @@ -21,7 +21,7 @@ pub fn make_plane_flow_config( None => { if load_flows.has_specialization() { - return Err(MatmulSetupError::InvalidConfig(Box::new( + return Err(MatmulSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( "Error: Load specialization config has specialization but no reader tasks were given." .to_string(), ))); diff --git a/crates/cubek-matmul/src/components/stage/matmul/plane_partitioned/setup.rs b/crates/cubek-matmul/src/components/stage/matmul/plane_partitioned/setup.rs index 11add18fe..82cf78370 100644 --- a/crates/cubek-matmul/src/components/stage/matmul/plane_partitioned/setup.rs +++ b/crates/cubek-matmul/src/components/stage/matmul/plane_partitioned/setup.rs @@ -148,7 +148,7 @@ impl StageM * blueprint.tiling_scheme.partitions_per_stage_along_n(), )) } else { - Err(Box::new( + Err(cubek_std::InvalidConfigError::new( "Error: Tried to use a plane stage matmul with a unit tile matmul.".to_string(), )) } @@ -166,17 +166,22 @@ impl StageM Self::cubedim_resource(blueprint)?.num_planes(blueprint.plane_dim)?; if num_compute_planes != num_planes_needed { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Error: Number of compute planes {num_compute_planes} should be {num_planes_needed}." - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Error: Number of compute planes {num_compute_planes} should be {num_planes_needed}." + )), + )); } if blueprint.partition_buffering == PartitionBuffering::Double && blueprint.tiling_scheme.tiles_per_stage_partition_along_n() < 2 { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Error: Tried doing double buffering with only one tile to compute.".to_string(), - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "Error: Tried doing double buffering with only one tile to compute." + .to_string(), + ), + )); } blueprint diff --git a/crates/cubek-matmul/src/components/stage/matmul/unit_partitioned/setup.rs b/crates/cubek-matmul/src/components/stage/matmul/unit_partitioned/setup.rs index cb32b82e4..a0fc7cde3 100644 --- a/crates/cubek-matmul/src/components/stage/matmul/unit_partitioned/setup.rs +++ b/crates/cubek-matmul/src/components/stage/matmul/unit_partitioned/setup.rs @@ -140,7 +140,7 @@ impl StageMatmulFamily * blueprint.tiling_scheme.partitions_per_stage_along_n(), )) } else { - Err(Box::new( + Err(cubek_std::InvalidConfigError::new( "Error: Tried to use a unit stage matmul with a plane tile matmul.".to_string(), )) } @@ -159,17 +159,21 @@ impl StageMatmulFamily let num_units = blueprint.plane_dim * num_compute_planes; if num_units != working_units { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Error: Number of units {num_units} should be {working_units}." - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Error: Number of units {num_units} should be {working_units}." + )), + )); } if blueprint.partition_buffering == PartitionBuffering::Double && blueprint.tiling_scheme.tiles_per_stage_partition_along_n() < 2 { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Error: Tried doing partition double buffering with only one tile to compute.", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "Error: Tried doing partition double buffering with only one tile to compute.", + ), + )); } blueprint diff --git a/crates/cubek-matmul/src/components/stage/memory/layout.rs b/crates/cubek-matmul/src/components/stage/memory/layout.rs index bbd61e1af..98a578c41 100644 --- a/crates/cubek-matmul/src/components/stage/memory/layout.rs +++ b/crates/cubek-matmul/src/components/stage/memory/layout.rs @@ -348,7 +348,7 @@ impl TilingValidation for ContiguousTilingLayout { fn check(config: StageMemoryConfig) -> Result<(), InvalidConfigError> { let tile_width = config.elements_per_tile_along_contiguous_dim(); if config.vector_size > tile_width { - return Err(Box::new(format!( + return Err(cubek_std::InvalidConfigError::new(format!( "Invalid vector size. Got {:?} which should not be >{:?}", config.vector_size, tile_width, ))); @@ -441,7 +441,7 @@ impl TilingValidation for StridedTilingLayout { fn check(config: StageMemoryConfig) -> Result<(), InvalidConfigError> { let stage_width = config.elements_per_stage_along_contiguous_dim(); if config.vector_size > stage_width { - return Err(Box::new(format!( + return Err(cubek_std::InvalidConfigError::new(format!( "Invalid vector size. Got {:?} which should not be >{:?}", config.vector_size, stage_width, ))); diff --git a/crates/cubek-matmul/src/components/tile.rs b/crates/cubek-matmul/src/components/tile.rs index f506a45e5..ea3026f44 100644 --- a/crates/cubek-matmul/src/components/tile.rs +++ b/crates/cubek-matmul/src/components/tile.rs @@ -299,9 +299,9 @@ fn validate_cmma( } if blueprint.swizzle_modes.has_swizzle() { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "This tile matmul doesn't support swizzling", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new("This tile matmul doesn't support swizzling"), + )); } Ok(()) @@ -363,40 +363,50 @@ fn validate_register( match blueprint.lhs_layout { MatrixLayout::RowMajor => { if !k.is_multiple_of(lhs) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile shape in vectorized axis k({k:?}) should be divisible by vector size lhs({lhs:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile shape in vectorized axis k({k:?}) should be divisible by vector size lhs({lhs:?})" + )), + )); } } MatrixLayout::ColMajor => { if !m.is_multiple_of(lhs) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile shape in vectorized axis m({m:?}) should be divisible by vector size lhs({lhs:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile shape in vectorized axis m({m:?}) should be divisible by vector size lhs({lhs:?})" + )), + )); } } } match blueprint.rhs_layout { MatrixLayout::RowMajor => { if !n.is_multiple_of(rhs) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile shape in vectorized axis n({n:?}) should be divisible by vector size rhs({rhs:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile shape in vectorized axis n({n:?}) should be divisible by vector size rhs({rhs:?})" + )), + )); } } MatrixLayout::ColMajor => { if !k.is_multiple_of(rhs) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile shape in vectorized axis k({k:?}) should be divisible by vector size rhs({rhs:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile shape in vectorized axis k({k:?}) should be divisible by vector size rhs({rhs:?})" + )), + )); } } } if !n.is_multiple_of(out) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile shape in vectorized axis n({n:?}) should be divisible by vector size out({out:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile shape in vectorized axis n({n:?}) should be divisible by vector size out({out:?})" + )), + )); } Ok(()) @@ -411,15 +421,15 @@ fn validate_plane_vec( check_types_available(client, dtypes, true)?; if blueprint.lhs_layout != MatrixLayout::RowMajor { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Only Row Major layout is supported for Lhs", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new("Only Row Major layout is supported for Lhs"), + )); } if blueprint.rhs_layout != MatrixLayout::ColMajor { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Only Col Major layout is supported for Rhs", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new("Only Col Major layout is supported for Rhs"), + )); } let m = blueprint.tiling_scheme.tile_size.m(); @@ -431,28 +441,34 @@ fn validate_plane_vec( let out_vector = vector_sizes.out as u32; if m != 1 { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Only m=1 is supported, got m={m:?}", - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!("Only m=1 is supported, got m={m:?}",)), + )); } if lhs_vector != rhs_vector { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Lhs and Rhs must have same vector size, got lhs={lhs_vector:?} and rhs={rhs_vector:?}", - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Lhs and Rhs must have same vector size, got lhs={lhs_vector:?} and rhs={rhs_vector:?}", + )), + )); } if k != blueprint.plane_dim * lhs_vector { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "k must be equal to plane_dim times vector size (of both lhs and rhs), got k={:?}, plane_dim={:?} vector_size={:?}", - k, blueprint.plane_dim, lhs_vector - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "k must be equal to plane_dim times vector size (of both lhs and rhs), got k={:?}, plane_dim={:?} vector_size={:?}", + k, blueprint.plane_dim, lhs_vector + )), + )); } if !n.is_multiple_of(out_vector) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "n must be divisible by out vector size, got n={n:?}, out_vector_size={out_vector:?}", - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "n must be divisible by out vector size, got n={n:?}, out_vector_size={out_vector:?}", + )), + )); } Ok(()) @@ -472,10 +488,12 @@ fn validate_interleaved( let plane_dim = blueprint.plane_dim; if !k.is_multiple_of(plane_dim) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "k must be divisible by plane_dim. Got k={:?}, plane_dim={:?}", - k, plane_dim, - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "k must be divisible by plane_dim. Got k={:?}, plane_dim={:?}", + k, plane_dim, + )), + )); } let k_local = k / plane_dim; @@ -487,40 +505,50 @@ fn validate_interleaved( match blueprint.lhs_layout { MatrixLayout::RowMajor => { if !k_local.is_multiple_of(lhs) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Local shape in vectorized axis k ({k_local:?}) should be divisible by vector size lhs ({lhs:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Local shape in vectorized axis k ({k_local:?}) should be divisible by vector size lhs ({lhs:?})" + )), + )); } } MatrixLayout::ColMajor => { if !m.is_multiple_of(lhs) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile shape in vectorized axis m ({m:?}) should be divisible by vector size lhs ({lhs:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile shape in vectorized axis m ({m:?}) should be divisible by vector size lhs ({lhs:?})" + )), + )); } } } match blueprint.rhs_layout { MatrixLayout::RowMajor => { if !n.is_multiple_of(rhs) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile shape in vectorized axis n ({n:?}) should be divisible by vector size rhs ({rhs:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile shape in vectorized axis n ({n:?}) should be divisible by vector size rhs ({rhs:?})" + )), + )); } } MatrixLayout::ColMajor => { if !k_local.is_multiple_of(rhs) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Local shape in vectorized axis k ({k_local:?}) should be divisible by vector size rhs ({rhs:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Local shape in vectorized axis k ({k_local:?}) should be divisible by vector size rhs ({rhs:?})" + )), + )); } } } if !n.is_multiple_of(out) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Tile shape in vectorized axis n ({n:?}) should be divisible by vector size out ({out:?})" - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Tile shape in vectorized axis n ({n:?}) should be divisible by vector size out ({out:?})" + )), + )); } Ok(()) diff --git a/crates/cubek-matmul/src/definition/error.rs b/crates/cubek-matmul/src/definition/error.rs index 8c746d788..0eaaecfaf 100644 --- a/crates/cubek-matmul/src/definition/error.rs +++ b/crates/cubek-matmul/src/definition/error.rs @@ -1,6 +1,9 @@ use cubecl::{CubeCount, CubeDim, VectorizationError, ir::StorageType, server::LaunchError}; use cubek_std::{InvalidConfigError, MatrixLayout, TileSize}; -use std::fmt::{Debug, Display}; +use std::{ + error::Error, + fmt::{Debug, Display}, +}; /// Errors that can occur during the setup phase of a matmul operation. pub enum MatmulSetupError { @@ -94,6 +97,15 @@ impl Display for MatmulSetupError { } } +impl Error for MatmulSetupError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + Self::InvalidConfig(error) => Some(error), + _ => None, + } + } +} + impl Debug for MatmulSetupError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { @@ -192,3 +204,28 @@ impl Debug for MatmulAvailabilityError { } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::error::Error; + + fn assert_error_send_sync() {} + + #[test] + fn setup_error_is_a_thread_safe_error() { + assert_error_send_sync::(); + } + + #[test] + fn invalid_config_is_exposed_as_the_error_source() { + let error = MatmulSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( + "invalid test config", + )); + + assert_eq!( + error.source().map(ToString::to_string).as_deref(), + Some("invalid test config") + ); + } +} diff --git a/crates/cubek-matmul/src/launch/complex.rs b/crates/cubek-matmul/src/launch/complex.rs index f991898af..4bd2e0720 100644 --- a/crates/cubek-matmul/src/launch/complex.rs +++ b/crates/cubek-matmul/src/launch/complex.rs @@ -252,12 +252,16 @@ fn validate_normal_input( ) -> Result<(), MatmulSetupError> { match input { InputBinding::Normal(_, dtype) if *dtype == c32_storage_type() => Ok(()), - InputBinding::Normal(_, dtype) => Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "complex GEMM {name} must use C32 storage, got {dtype:?}" - )))), - InputBinding::Quantized { .. } => Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "complex GEMM {name} does not support quantized input" - )))), + InputBinding::Normal(_, dtype) => Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "complex GEMM {name} must use C32 storage, got {dtype:?}" + )), + )), + InputBinding::Quantized { .. } => Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "complex GEMM {name} does not support quantized input" + )), + )), } } @@ -267,19 +271,23 @@ fn validate_c32_globals(dtypes: &MatmulElems) -> Result<(), MatmulSetupError> { || dtypes.rhs_global != c32_type || dtypes.acc_global != c32_type { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "complex GEMM dtypes must be C32, got lhs={:?}, rhs={:?}, out={:?}", - dtypes.lhs_global, dtypes.rhs_global, dtypes.acc_global - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "complex GEMM dtypes must be C32, got lhs={:?}, rhs={:?}, out={:?}", + dtypes.lhs_global, dtypes.rhs_global, dtypes.acc_global + )), + )); } Ok(()) } fn validate_rank(name: &'static str, shape: &Shape) -> Result<(), MatmulSetupError> { if shape.len() < 2 { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "complex GEMM {name} must have rank at least 2", - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "complex GEMM {name} must have rank at least 2", + )), + )); } Ok(()) } @@ -319,12 +327,16 @@ fn dense_matrix_strides( } let mut batch_stride = rows.checked_mul(cols).ok_or_else(|| { - MatmulSetupError::InvalidConfig(Box::new("complex GEMM batch stride overflow")) + MatmulSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( + "complex GEMM batch stride overflow", + )) })?; for axis in (0..rank - 2).rev() { strides[axis] = batch_stride; batch_stride = batch_stride.checked_mul(shape[axis]).ok_or_else(|| { - MatmulSetupError::InvalidConfig(Box::new("complex GEMM batch stride overflow")) + MatmulSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( + "complex GEMM batch stride overflow", + )) })?; } Ok(Strides::new(&strides)) @@ -352,7 +364,9 @@ fn metadata_with_backing( fn logical_len(shape: &[usize]) -> Result { shape.iter().try_fold(1usize, |acc, dim| { acc.checked_mul(*dim).ok_or_else(|| { - MatmulSetupError::InvalidConfig(Box::new("complex GEMM logical length overflow")) + MatmulSetupError::InvalidConfig(cubek_std::InvalidConfigError::new( + "complex GEMM logical length overflow", + )) }) }) } diff --git a/crates/cubek-matmul/src/launch/launch_vecmat_plane_parallel.rs b/crates/cubek-matmul/src/launch/launch_vecmat_plane_parallel.rs index 07d045dd8..0fb4da0d3 100644 --- a/crates/cubek-matmul/src/launch/launch_vecmat_plane_parallel.rs +++ b/crates/cubek-matmul/src/launch/launch_vecmat_plane_parallel.rs @@ -57,10 +57,12 @@ pub fn launch_ref( let plane_size = client.properties().hardware.plane_size_max as usize; if !k.is_multiple_of(plane_size) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Dimension k={} must be a multiple of plane size {}", - k, plane_size - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Dimension k={} must be a multiple of plane size {}", + k, plane_size + )), + )); } let lhs_vector_size = vector_size_for(client, &lhs, dtypes.lhs_global.size(), plane_size, k)?; @@ -120,13 +122,17 @@ pub fn launch_ref( if device_settings.plane_dim > 1 { if matches!(expand_info.blueprint.kind, GemvKind::MatVecColMajor) { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "On GPU, MatVec plane parallel only supports row major lhs for now", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "On GPU, MatVec plane parallel only supports row major lhs for now", + ), + )); } else if matches!(expand_info.blueprint.kind, GemvKind::VecMatRowMajor) { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "On GPU, Vecmat plane parallel only supports col major rhs for now", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "On GPU, Vecmat plane parallel only supports col major rhs for now", + ), + )); } } diff --git a/crates/cubek-matmul/src/launch/launch_vecmat_unit_perpendicular.rs b/crates/cubek-matmul/src/launch/launch_vecmat_unit_perpendicular.rs index b091a41ac..af07e8982 100644 --- a/crates/cubek-matmul/src/launch/launch_vecmat_unit_perpendicular.rs +++ b/crates/cubek-matmul/src/launch/launch_vecmat_unit_perpendicular.rs @@ -60,9 +60,9 @@ pub fn launch_ref( let k = lhs.shape().to_vec()[rank - 1]; if m != 1 { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "m must equal 1 to qualify as a vecmat problem", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new("m must equal 1 to qualify as a vecmat problem"), + )); } let rhs_shape = rhs.shape(); @@ -70,17 +70,21 @@ pub fn launch_ref( let plane_size = client.properties().hardware.plane_size_max as usize; if !k.is_multiple_of(plane_size) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Lhs dimension k={} must be a multiple of plane size {}", - k, plane_size - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Lhs dimension k={} must be a multiple of plane size {}", + k, plane_size + )), + )); } if !n.is_multiple_of(plane_size) { - return Err(MatmulSetupError::InvalidConfig(Box::new(format!( - "Rhs dimension n={} must be a multiple of plane size {}", - n, plane_size - )))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new(format!( + "Rhs dimension n={} must be a multiple of plane size {}", + n, plane_size + )), + )); } let lhs_vector_size = vector_size_for(client, &lhs, dtypes.lhs_global.size(), plane_size, k)?; @@ -119,9 +123,11 @@ pub fn launch_ref( ); if problem.rhs_layout != MatrixLayout::RowMajor { - return Err(MatmulSetupError::InvalidConfig(Box::new( - "Vecmat unit perpendicular only supports row major rhs for now", - ))); + return Err(MatmulSetupError::InvalidConfig( + cubek_std::InvalidConfigError::new( + "Vecmat unit perpendicular only supports row major rhs for now", + ), + )); } let device_settings = GemvUnitPerpendicularRoutine::device_settings(client, vector_sizes); diff --git a/crates/cubek-std/src/cube_dim_resource.rs b/crates/cubek-std/src/cube_dim_resource.rs index 0ec22a28d..507e4423e 100644 --- a/crates/cubek-std/src/cube_dim_resource.rs +++ b/crates/cubek-std/src/cube_dim_resource.rs @@ -77,7 +77,7 @@ impl CubeDimResource { if units % plane_dim == 0 { Ok(CubeDimResource::Planes(units / plane_dim)) } else { - Err(Box::new(format!( + Err(crate::InvalidConfigError::new(format!( "Number of units {units:?} should be divisible by plane_dim {plane_dim:?}" ))) } diff --git a/crates/cubek-std/src/error.rs b/crates/cubek-std/src/error.rs index 27ab6db63..27b4680a6 100644 --- a/crates/cubek-std/src/error.rs +++ b/crates/cubek-std/src/error.rs @@ -1,17 +1,60 @@ -use std::fmt::Display; +use std::{ + error::Error, + fmt::{Debug, Display}, +}; /// Error that arises from invalid configurations -pub type InvalidConfigError = Box; +pub struct InvalidConfigError { + diagnostic: Box, +} + +trait InvalidConfigDiagnostic: Debug + Display + Send + Sync + 'static {} + +impl InvalidConfigDiagnostic for T where T: Debug + Display + Send + Sync + 'static {} + +impl InvalidConfigError { + /// Wrap a diagnostic value in a thread-safe configuration error. + pub fn new(diagnostic: T) -> Self + where + T: Debug + Display + Send + Sync + 'static, + { + Self { + diagnostic: Box::new(diagnostic), + } + } +} + +impl Debug for InvalidConfigError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + Debug::fmt(&self.diagnostic, f) + } +} + +impl Display for InvalidConfigError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + Display::fmt(&self.diagnostic, f) + } +} + +impl Error for InvalidConfigError {} /// Error that arises from invalid configurations pub struct FormattedConfigError { - func: Box String>, + func: Box String + Send + Sync>, +} + +impl Debug for FormattedConfigError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("FormattedConfigError") + .field(&(self.func)()) + .finish() + } } impl FormattedConfigError { #[allow(clippy::new_ret_no_self)] - pub fn new String + 'static>(func: F) -> Box { - Box::new(Self { + pub fn new String + Send + Sync + 'static>(func: F) -> InvalidConfigError { + InvalidConfigError::new(Self { func: Box::new(func), }) } @@ -23,3 +66,15 @@ impl Display for FormattedConfigError { write!(f, "{string}") } } + +#[cfg(test)] +mod tests { + use super::*; + + fn assert_error_send_sync() {} + + #[test] + fn invalid_config_error_is_a_thread_safe_error() { + assert_error_send_sync::(); + } +} diff --git a/crates/cubek-std/src/matrix_layout.rs b/crates/cubek-std/src/matrix_layout.rs index 74718142f..c58e9b75e 100644 --- a/crates/cubek-std/src/matrix_layout.rs +++ b/crates/cubek-std/src/matrix_layout.rs @@ -34,7 +34,7 @@ impl MatrixLayout { return Ok(MatrixLayout::ColMajor); } - return Err(Box::new(format!( + return Err(crate::InvalidConfigError::new(format!( "Invalid or non-contiguous matrix layout: packing_dim={packing_dim:?}" ))); } @@ -64,7 +64,7 @@ impl MatrixLayout { return Ok(MatrixLayout::ColMajor); } - Err(Box::new(format!( + Err(crate::InvalidConfigError::new(format!( "Invalid or non-contiguous matrix layout: shape={shape:?}, strides={strides:?}", ))) } From 43e8521885f141cb8ccdf99a766bfde118412010 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Sun, 19 Jul 2026 10:59:44 +0900 Subject: [PATCH 5/5] fix(matmul): retain launch error source --- crates/cubek-matmul/src/definition/error.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/crates/cubek-matmul/src/definition/error.rs b/crates/cubek-matmul/src/definition/error.rs index 0eaaecfaf..76dcc5d03 100644 --- a/crates/cubek-matmul/src/definition/error.rs +++ b/crates/cubek-matmul/src/definition/error.rs @@ -101,6 +101,7 @@ impl Error for MatmulSetupError { fn source(&self) -> Option<&(dyn Error + 'static)> { match self { Self::InvalidConfig(error) => Some(error), + Self::Launch(error) => Some(error), _ => None, } } @@ -208,6 +209,7 @@ impl Debug for MatmulAvailabilityError { #[cfg(test)] mod tests { use super::*; + use cubecl::{CompilationError, backtrace::BackTrace}; use std::error::Error; fn assert_error_send_sync() {} @@ -228,4 +230,19 @@ mod tests { Some("invalid test config") ); } + + #[test] + fn launch_error_is_exposed_as_the_error_source() { + let launch_error = LaunchError::CompilationError(CompilationError::Generic { + reason: "test compilation failure".into(), + backtrace: BackTrace::default(), + }); + let expected = launch_error.to_string(); + let error = MatmulSetupError::Launch(launch_error); + + assert_eq!( + error.source().map(ToString::to_string).as_deref(), + Some(expected.as_str()) + ); + } }