FFT: expose interleaved C32 APIs with F32 kernels

[shinaoka]

Context

cubek-fft currently exposes split real/imaginary buffers and its kernels are effectively hard-coded to f32. The downstream Apple/Metal FFT work in tensor4all/tenferro-rs#1394 needs a stable complex ABI that can be shared with Rust Complex<T> storage without pack/unpack passes.

This should be implemented in CubeK first and integrated into tenferro separately.

Goal

Expose power-of-two CFFT, RFFT, and IRFFT APIs using an interleaved global-memory ABI for both precisions:

• C32: scalar storage f32, physical order [re0, im0, re1, im1, ...]
• C64: scalar storage f64, physical order [re0, im0, re1, im1, ...]
• RFFT: F32 -> C32 and F64 -> C64
• IRFFT: C32 -> F32 and C64 -> F64
• CFFT: C32 -> C32 and C64 -> C64

The public tensor metadata must retain the logical complex shape. Do not model a complex value by appending a visible trailing dimension of length 2.

Design direction

Use a hybrid representation:

• Global input/output ABI is interleaved.
• On kernel load, separate each complex scalar into real and imaginary values.
• Keep existing shared memory and four-step intermediate scratch split into real/imaginary buffers.
• Re-interleave only when storing the final output.

This avoids standalone pack/unpack kernels and preserves the existing small and four-step algorithms. A full rewrite of shared memory and scratch to interleaved storage is out of scope unless benchmarks demonstrate a benefit.

Add a dedicated interleaved complex binding/layout type that:

• records the logical complex element count separately from the physical scalar count;
• maps logical complex index i to physical scalar offsets 2*i and 2*i + 1;
• validates byte length, alignment, dtype, shape, strides, transform axis, and non-overlapping output requirements;
• accepts caller-owned input/output/scratch bindings, with no host transfers or hidden output allocation.

Precision and backend capability

The API and dispatch must support both precision families without silently narrowing values.

• Generalize kernels, twiddle generation, constants, and normalization away from hard-coded f32.
• Dispatch F32/C32 to f32 kernels and F64/C64 to f64 kernels.
• Before launch, check whether the selected CubeCL runtime/backend supports the required scalar type and operations.
• If native f64 kernels are unavailable, return a typed UnsupportedElementType/capability error. Do not fall back to f32, CPU execution, or host staging.
• Add f64 execution tests on every backend that advertises the capability, plus explicit rejection tests on unsupported backends.

In particular, defining the C64 ABI now is required even where the current WebGPU/Metal path cannot execute f64 shaders yet.

Public API requirements

• Make CFFT, RFFT, and IRFFT caller-owned launch APIs public.
• Replace public assert!/unwrap validation paths with typed errors.
• Expose an explicit normalization/scale policy; do not hard-code normalization inside only one transform.
• Guarantee out-of-place operation. Reject unsupported overlapping bindings with a typed error; in-place CFFT can be added separately after aliasing behavior is designed and tested.
• Preserve arbitrary transform axis/stride/batching behavior supported by BatchSignalLayout.
• Preserve current padding/truncation behavior through explicit signal/spectrum lengths.
• Initial transform sizes remain power-of-two with n_fft >= 2.

Implementation plan

PR 1: ABI, validation, and public surface

• Introduce interleaved complex binding/layout metadata.
• Add dtype pairs F32/C32 and F64/C64.
• Add typed launch/shape/dtype/capability/overlap errors.
• Publicize caller-owned CFFT/RFFT/IRFFT launch entry points.
• Add explicit normalization/scale configuration.
• Add layout-only tests proving Rust-compatible Complex32 and Complex64 byte order and logical-shape handling.

PR 2: Interleaved F32/C32 kernels

• Update small CFFT loads/stores to consume/produce interleaved global buffers while retaining split shared memory.
• Update the four-step path so its first read and final reorder use the interleaved ABI; retain split intermediate scratch.
• Make RFFT write interleaved C32 and IRFFT read interleaved C32.
• Cover small, large, and four-step paths.

PR 3: Generic scalar plumbing and F64/C64 kernels

• Parameterize arithmetic, twiddles, constants, scratch sizing, and scaling by scalar type.
• Add f64 dispatch and backend capability checks.
• Run F64/C64 numerical tests on capable backends and typed rejection tests elsewhere.
• Verify there is no implicit cast to f32 in the kernel path.

PR 4: Integration hardening and benchmarks

• Test axes 0/middle/last, batched and strided layouts, padding/truncation, odd/even spectrum boundaries where applicable, and minimum/large supported sizes.
• Test caller-owned buffers, overlap rejection, invalid physical lengths, wrong dtype pairs, and unsupported backend capabilities.
• Benchmark against the current split-buffer implementation and a pack/unpack baseline.
• Document the ABI and supported backend/type matrix.

Acceptance criteria

• Public interleaved CFFT/RFFT/IRFFT launch APIs exist.
• C32 and C64 use Rust-compatible [re, im] scalar ordering without a visible trailing dimension.
• F32/C32 works for small, large, and four-step paths.
• F64/C64 works on every backend that advertises required f64 support.
• Unsupported f64 backends fail before launch with a typed error.
• No silent precision narrowing, CPU fallback, host staging, or standalone pack/unpack pass.
• Internal split shared memory/scratch is retained unless benchmarks justify changing it.
• Normalization/scale behavior is explicit and tested for forward and inverse transforms.
• Axis, stride, batch, and padding/truncation coverage is retained.
• Public validation paths do not panic.
• ABI documentation and a backend/type support table are included.

Non-goals

• Arbitrary non-power-of-two FFT sizes.
• A general CubeCL complex-number type/lowering.
• Full interleaving of shared memory and four-step scratch.
• CPU FFT implementation or automatic backend fallback.
• Plan caching; the current CubeK plan is launch metadata rather than the process-global cache problem tracked in tenferro.

[shinaoka]

F32/C32 interleaved FFT support has been merged in #7.

Remaining follow-up:

• Add F64/C64 APIs and typed unsupported-capability handling for backends such as Metal that lack native F64 support.

[shinaoka]

CPU backend correctness/CI follow-up merged in #8.

• added final cube synchronization to interleaved CFFT/RFFT/IRFFT shared-memory kernels
• bounded the normal CPU shared-memory path test while retaining true device-limit coverage under heavy
• made unsupported-size validation use checked arithmetic for CPU backends that report system RAM as shared memory
• PR CI passed on stable and previous Rust toolchains

F32/C32 support is now merged and CI-clean. F64/C64 and typed unsupported-capability handling for Metal remain follow-up work in this issue.

[shinaoka]

Scope correction: CubeK does not need F64/C64 support. CubeK remains the Metal/WebGPU F32/C32 FFT provider; F64/C64 are handled by explicitly selected CPU or CUDA providers in tenferro.

The required interleaved C32 CFFT/RFFT/IRFFT implementation is merged in #7, with the CPU correctness follow-up merged in #8. Closing this CubeK task as complete.