FFT: add interleaved F32/C32 APIs

benchmarks/Cargo.toml

@@ -120,6 +120,10 @@ name = "contiguous"
 harness = false
 name = "fft"
 
+[[bench]]
+harness = false
+name = "cfft"
+
 [[bench]]
 harness = false
 name = "quantized_matmul"

benchmarks/benches/cfft.rs

@@ -0,0 +1 @@
+benchmarks::run_bench!(cfft);

benchmarks/src/lib.rs

@@ -4,6 +4,9 @@ pub use cubek_attention::eval::backward::benchmarks as attention_backward;
 pub use cubek_attention::eval::forward::benchmarks as attention;
 pub use cubek_convolution::eval::benchmarks as conv2d;
 pub use cubek_fft::eval::benchmarks as fft;
+pub mod cfft {
+    pub use cubek_fft::eval::benchmarks::CfftCategory as Category;
+}
 pub use cubek_interpolate::eval::benchmarks as interpolate;
 pub use cubek_matmul::eval::benchmarks::gemm;
 pub use cubek_matmul::eval::benchmarks::gemv;
@@ -26,6 +29,7 @@ pub fn all() -> &'static [&'static dyn BenchmarkCategory] {
         &crate::attention_backward::Category,
         &crate::contiguous::Category,
         &crate::conv2d::Category,
+        &crate::cfft::Category,
         &crate::fft::Category,
         &crate::gemm::Category,
         &crate::gemv::Category,

crates/cubek-attention/src/backward/launch/dkdv.rs

@@ -50,7 +50,7 @@ fn flash_attention_backward_dv_kernel<E: Float>(
 
     let mut dv_acc = Array::new(val_dim);
     for dd in 0..val_dim {
-        dv_acc[dd] = E::new(0.0);
+        dv_acc[dd] = E::new(0.0_f32);
     }
 
     for i in 0..seq_q {
@@ -60,7 +60,7 @@ fn flash_attention_backward_dv_kernel<E: Float>(
             let do_base = do_row_base + i * val_dim;
             let lse_i = lse[row_idx_base + i];
 
-            let mut dot = E::new(0.0);
+            let mut dot = E::new(0.0_f32);
             for dd in 0..head_dim {
                 dot += q[q_base + dd] * k[k_base + dd];
             }
@@ -115,7 +115,7 @@ fn flash_attention_backward_dk_kernel<E: Float>(
 
     let mut dk_acc = Array::new(head_dim);
     for dd in 0..head_dim {
-        dk_acc[dd] = E::new(0.0);
+        dk_acc[dd] = E::new(0.0_f32);
     }
 
     for i in 0..seq_q {
@@ -126,13 +126,13 @@ fn flash_attention_backward_dk_kernel<E: Float>(
             let lse_i = lse[row_idx_base + i];
             let d_i = d[row_idx_base + i];
 
-            let mut dot = E::new(0.0);
+            let mut dot = E::new(0.0_f32);
             for dd in 0..head_dim {
                 dot += q[q_base + dd] * k[k_base + dd];
             }
             let p_ij = (dot * scale_e - lse_i).exp();
 
-            let mut dp = E::new(0.0);
+            let mut dp = E::new(0.0_f32);
             for dd in 0..val_dim {
                 dp += do_[do_base + dd] * v[v_base + dd];
             }

crates/cubek-attention/src/backward/launch/dq.rs

@@ -45,7 +45,7 @@ fn flash_attention_backward_dq_kernel<E: Float>(
 
     let mut dq_acc = Array::new(head_dim);
     for dd in 0..head_dim {
-        dq_acc[dd] = E::new(0.0);
+        dq_acc[dd] = E::new(0.0_f32);
     }
 
     for j in 0..seq_kv {
@@ -54,14 +54,14 @@ fn flash_attention_backward_dq_kernel<E: Float>(
             let k_base = k_row_base + j * head_dim;
             let v_base = v_row_base + j * val_dim;
 
-            let mut dot = E::new(0.0);
+            let mut dot = E::new(0.0_f32);
             for dd in 0..head_dim {
                 dot += q[q_base + dd] * k[k_base + dd];
             }
             let s_ij = dot * scale_e;
             let p_ij = (s_ij - lse_i).exp();
 
-            let mut dp = E::new(0.0);
+            let mut dp = E::new(0.0_f32);
             for dd in 0..val_dim {
                 dp += do_[do_base + dd] * v[v_base + dd];
             }

crates/cubek-attention/src/backward/launch/prepass.rs

@@ -23,7 +23,7 @@ fn flash_attention_backward_prepass_kernel<E: Float>(
     let head_dim = o.shape(o.rank() - 1);
     let base = row_idx * head_dim;
 
-    let mut acc = E::new(0.0);
+    let mut acc = E::new(0.0_f32);
     for k in 0..head_dim {
         acc += o[base + k] * do_[base + k];
     }

crates/cubek-fft/Cargo.toml

@@ -26,6 +26,7 @@ benchmarks = ["cpu-reference", "cubecl/test-runtime"]
 
 [dependencies]
 cubecl = { workspace = true }
+thiserror = { workspace = true }
 cubek-test-utils = { path = "./../cubek-test-utils/", version = "=0.3.0-pre.1", default-features = false, optional = true }
 num-complex = { version = "0.4.6", optional = true }
 

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<R: Runtime> {
+    tensor: TensorHandle<R>,
+    logical_strides: Vec<usize>,
+    physical_scalar_len: usize,
+}
+
+impl<R: Runtime> ComplexTensorHandle<R> {
+    /// Allocates a contiguous C32 tensor with the requested logical shape.
+    pub fn empty(
+        client: &ComputeClient<R>,
+        shape: Vec<usize>,
+        dtype: StorageType,
+    ) -> Result<Self, FftError> {
+        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<usize>,
+        handle: cubecl::server::Handle,
+        dtype: StorageType,
+    ) -> Result<Self, FftError> {
+        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<usize>,
+        logical_strides: Vec<usize>,
+        handle: cubecl::server::Handle,
+        dtype: StorageType,
+    ) -> Result<Self, FftError> {
+        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<R> {
+        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<R>,
+}
+
+#[allow(dead_code)]
+impl<R: Runtime> 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<R> {
+        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<TensorBinding<R>, FftError> {
+        self.ensure_unique_output()?;
+        Ok(self.tensor())
+    }
+}
+
+#[allow(dead_code)]
+pub(crate) fn ensure_unique_output<R: Runtime>(tensor: &TensorHandle<R>) -> 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<Vec<usize>, 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>, usize), FftError> {
+    let scalar_strides = logical_strides
+        .iter()
+        .enumerate()
+        .map(|(axis, stride)| {
+            stride
+                .checked_mul(2)
+                .ok_or(FftError::StrideOverflow { axis })
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+
+    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 = <TestRuntime as Runtime>::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::<TestRuntime>::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 = <TestRuntime as Runtime>::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::<TestRuntime>::new_contiguous(vec![1], handle, dtype);
+
+        assert!(matches!(result, Err(FftError::InvalidBufferRange { .. })));
+    }
+}