FFT (extension)

tenferro-fft is the FFT extension package for tenferro. It is an extension crate imported directly alongside tenferro-runtime or tenferro-tensor. Concrete non-AD execution uses TensorFftExt and TensorReadFftExt; eager execution uses EagerTensorFftExt behind autodiff; traced graphs use TracedTensorFftExt.

The current implementation provides one-dimensional CPU transforms backed by RustFFT, an explicitly selected CUDA path backed by cuFFT, and an explicitly selected Apple Metal path backed by CubeK. The public API is ordinary Rust extension-trait methods, so most users do not need to work with the lower-level extension machinery directly.

Setup

When working from a local checkout, use paths that match your project layout. For a scratch crate created directly inside the tenferro-rs checkout, include an empty [workspace] table:

[workspace]

Then add the dependencies:

[dependencies]
num-complex = "0.4"
tenferro-runtime = { path = "../crates/tenferro-runtime" }
tenferro-tensor = { path = "../crates/tenferro-tensor" }
tenferro-cpu = { path = "../crates/tenferro-cpu" }
tenferro-ad = { path = "../crates/tenferro-ad" }
tenferro-fft = { path = "../crates/tenferro-fft", features = ["autodiff"] }

The Apple shared path is not released yet. Until a later release task publishes it, use matching path dependencies from the same tenferro-rs checkout. The workspace pins the reviewed CubeCL and CubeK revisions; applications do not need to declare CubeCL or CubeK directly.

Concrete and graph-only users can omit tenferro-ad and the autodiff feature. Enable tenferro-fft’s autodiff feature when registering FFT AD rules. rustfft is pulled in automatically by tenferro-fft, and the first local build can take a few minutes on a fresh machine.

For the Apple shared CPU/Metal path, also enable the WebGPU feature on both operation and backend crates:

[dependencies]
tenferro-fft = { path = "../crates/tenferro-fft", features = ["webgpu"] }
tenferro-gpu = { path = "../crates/tenferro-gpu", default-features = false, features = ["webgpu"] }
tenferro-cpu = { path = "../crates/tenferro-cpu", default-features = false, features = ["cpu-faer"] }
tenferro-linalg = { path = "../crates/tenferro-linalg", default-features = false, features = ["cpu-faer"] }
tenferro-tensor = { path = "../crates/tenferro-tensor" }

tenferro-linalg and the CPU provider are needed by the Cholesky tutorial; FFT-only applications may omit tenferro-linalg.

For CUDA cuFFT execution from released packages, enable the CUDA feature on both the FFT extension and its provider:

[dependencies]
tenferro-fft = { version = "0.2", features = ["cuda"] }
tenferro-gpu = { version = "0.2", default-features = false, features = ["cuda"] }

The CUDA provider needs a compatible NVIDIA toolkit and driver. Set CUDA_PATH to the toolkit root used by CubeCL/NVRTC and include its lib64 directory (plus any separately installed CUDA vendor-library directories) in LD_LIBRARY_PATH. Set CUBECL_DEBUG_LOG=0 to suppress generated-kernel log output. cuFFT is loaded by tenferro-fft, not tenferro-gpu; use TENFERRO_CUFFT_PATH for an optional ordered, colon-separated list of explicit library paths:

export CUDA_PATH=/usr/local/cuda
export LD_LIBRARY_PATH="$CUDA_PATH/lib64:${LD_LIBRARY_PATH:-}"
export CUBECL_DEBUG_LOG=0
export TENFERRO_CUFFT_PATH=/opt/cuda/lib64/libcufft.so.11:/opt/cuda/lib64/libcufft.so.10

The loader tries non-empty TENFERRO_CUFFT_PATH entries in order, then falls back to libcufft.so.11, libcufft.so.10, and libcufft.so. This covers CUDA 12/cuFFT 11 and CUDA 11/cuFFT 10 installations. A path that cannot be opened is skipped so a later override or default can work; if no candidate can be loaded, or a loaded library cannot provide the required cuFFT symbols, the operation returns a typed provider/load error. This lookup fallback is only for finding cuFFT: it never selects CPU execution or performs an implicit transfer.

Current API

The initial API mirrors the common PyTorch and JAX one-dimensional FFT families:

Operation family	Purpose
`fft`, `ifft`	complex-to-complex transforms; real input may be promoted to complex output
`rfft`, `irfft`	real-to-complex and complex-to-real one-dimensional transforms

Each function accepts an optional transform length n, an axis, and an FftNorm value. Negative axes are normalized relative to the input rank. The normalization modes are:

Mode	Behavior
`FftNorm::Backward`	forward unscaled, inverse scaled by `1 / n`
`FftNorm::Forward`	forward scaled by `1 / n`, inverse unscaled
`FftNorm::Ortho`	forward and inverse scaled by `1 / sqrt(n)`

Backward is the default and matches NumPy, PyTorch, and JAX.

CUDA concrete execution

The CUDA path is explicit: upload a host tensor, run rfft with a CUDA execution session, verify that the result is still CUDA-resident, synchronize, and download only when host assertions are needed. The complete runnable example is the source of truth for this guide:

//! Explicit CUDA cuFFT execution with host/device transfers at visible boundaries.

use num_complex::Complex32;
use tenferro_fft::{FftNorm, TensorFftExt};
use tenferro_gpu::cuda::{
    cuda_devices, download_tensor, gpu_available, upload_tensor, CudaBackend,
};
use tenferro_runtime::BackendSessionHost;
use tenferro_tensor::{DType, MemoryKind, Tensor, TensorRead};

const TUTORIAL_SKIP_MARKER: &str = "TENFERRO_TUTORIAL_SKIP:";

fn skip_or_fail(
    require_cuda: bool,
    reason: impl std::fmt::Display,
) -> Result<(), Box<dyn std::error::Error>> {
    let reason = reason.to_string();
    if require_cuda {
        return Err(std::io::Error::other(format!(
            "CUDA FFT tutorial requires CUDA assertions, but {reason}"
        ))
        .into());
    }

    eprintln!("{TUTORIAL_SKIP_MARKER} CUDA FFT tutorial skipped: {reason}");
    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let require_cuda = std::env::var("TENFERRO_REQUIRE_CUDA")
        .map(|value| value == "1")
        .unwrap_or(false);
    if !gpu_available() {
        return skip_or_fail(require_cuda, "no usable CUDA device is available");
    }

    let devices = match cuda_devices() {
        Ok(devices) => devices,
        Err(error) => {
            return skip_or_fail(
                require_cuda,
                format!("CUDA device enumeration failed: {error}"),
            );
        }
    };
    let Some(device) = devices.into_iter().next() else {
        return skip_or_fail(
            require_cuda,
            "CUDA reported available but device enumeration returned no devices",
        );
    };

    // Keep one backend/runtime for the upload, FFT, synchronization, and download.
    let mut backend = CudaBackend::new(device.id())?;
    let host = Tensor::from_vec_col_major([4], vec![1.0_f32, 2.0, 3.0, 4.0])?;
    let gpu_input = upload_tensor(backend.runtime(), &host)?;

    let input_read = TensorRead::from_tensor(&gpu_input);
    assert_eq!(input_read.backend_family(), Some("cuda"));
    assert_eq!(input_read.placement().memory_kind, MemoryKind::Device);
    let input_domain = input_read
        .allocation_domain()
        .ok_or("uploaded tensor has no CUDA allocation domain")?;

    // FFT execution consumes the already uploaded tensor; it does not transfer it.
    let spectrum = backend
        .with_backend_session(|session| gpu_input.rfft(None, 0, FftNorm::Backward, session))?;

    // Check residency before crossing the explicit device-to-host boundary.
    let spectrum_read = TensorRead::from_tensor(&spectrum);
    assert_eq!(spectrum_read.backend_family(), Some("cuda"));
    assert_eq!(spectrum_read.placement().memory_kind, MemoryKind::Device);
    assert_eq!(spectrum_read.allocation_domain(), Some(input_domain));
    assert_eq!(spectrum.dtype(), DType::C32);
    assert_eq!(spectrum.shape(), &[3]);

    // The cuFFT execution runs inside the credentialed raw session: the bound
    // stream is synchronized before the session returns, while a fresh
    // session-scoped work area and the input/output retention guards pin the
    // allocations. The explicit download below synchronizes stream-managed
    // postprocessing and is the visible device-to-host boundary for the final
    // output.
    let host_spectrum = download_tensor(backend.runtime(), &spectrum)?;
    assert_eq!(host_spectrum.dtype(), DType::C32);
    assert_eq!(host_spectrum.shape(), &[3]);
    let values = host_spectrum.as_slice::<Complex32>()?;
    let expected = [
        Complex32::new(10.0, 0.0),
        Complex32::new(-2.0, 2.0),
        Complex32::new(-2.0, 0.0),
    ];
    assert_eq!(values.len(), expected.len());
    for (index, (actual, expected)) in values.iter().zip(expected).enumerate() {
        assert!(
            (*actual - expected).norm() <= 1.0e-5,
            "rfft value {index} differs: actual {actual:?}, expected {expected:?}"
        );
    }

    Ok(())
}

In ordinary docs CI, the CUDA feature is compile-checked and the tutorial runner may also run this binary without a CUDA device. In that case the binary prints a TENFERRO_TUTORIAL_SKIP: diagnostic and the runner reports the hardware-dependent tutorial as skipped; it does not claim that the CUDA assertions ran. The same explicit skip is used if device enumeration is empty, even if the availability probe was positive. A compile-only --no-run check never attempts CUDA execution.

To require the assertions to execute, set TENFERRO_REQUIRE_CUDA=1. With this strict mode, unavailable CUDA, device-enumeration errors, or an empty device list return failure instead of a skip. A successful command therefore proves that this tutorial reached and completed its CUDA assertions. For example, on an A100 or another configured CUDA host:

TENFERRO_REQUIRE_CUDA=1 \
CUBECL_DEBUG_LOG=0 \
CUDA_PATH=/usr/local/cuda \
LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH:-} \
cargo run -p tenferro-tutorial-code --no-default-features \
  --features cpu-faer,cuda,doc-snippets --bin cuda_fft

CUDA supports one-dimensional transforms with these dtype combinations:

Operation	CUDA input → output
`fft`	`C32 → C32`, `C64 → C64`, `F32 → C32` (full Hermitian), `F64 → C64` (full Hermitian)
`ifft`	`C32 → C32`, `C64 → C64`
`rfft`	`F32 → C32`, `F64 → C64` (one-sided)
`irfft`	`C32 → F32`, `C64 → F64`

The input must already be resident on the exact CudaRuntime borrowed by the execution session. CUDA FFT does not implicitly transfer a host or foreign- runtime tensor and never falls back to CPU RustFFT; upload and download are caller-visible operations. n truncates or zero-pads C2C/R2C input on device, with padding allocated and filled as semantic device zeros; C2R uses n only for the real output length and consumes its validated half-spectrum unchanged. Normalization is applied on device using the same FftNorm rules described above. cuFFT execution enters the credentialed raw session: a fresh session-scoped work area is allocated, the input/output allocations are retained, the vendor call is enqueued, and the bound stream is synchronized before the session returns, so the session’s tensor spans and work area cannot be reclaimed while vendor work is still using them. A failed synchronization intentionally forgets the work area and the input/output retention guards (issue #967 invariant). Subsequent CUDA normalization, Hermitian completion, and axis restoration remain stream-managed; an explicit download synchronizes the final output. Cache eviction or clear retires plans under a context-restoring guard rather than dropping them while queued work may still be active. cuFFT/provider failures are returned as typed errors rather than producing a host result.

Concrete Tensor And TensorRead

Use TensorFftExt when you have an owned compact Tensor and want immediate non-AD execution on an explicit backend. Use TensorReadFftExt when the input is a borrowed view or other read-oriented value. The _read suffix is reserved for that TensorRead surface; compact Tensor inputs use unsuffixed method names.

use num_complex::Complex64;
use tenferro_cpu::CpuBackend;
use tenferro_fft::{FftNorm, TensorFftExt, TensorReadFftExt};
use tenferro_runtime::BackendSessionHost;
use tenferro_tensor::{Tensor, TensorRead, TensorView, TypedTensorView};

let mut backend = CpuBackend::new();
backend.with_backend_session(|session| -> Result<(), tenferro_tensor::Error> {
    let x = Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0])?;
    let full = x.fft(None, -1, FftNorm::Backward, session)?;
    let one_sided = x.rfft(None, -1, FftNorm::Backward, session)?;
    assert_eq!(full.as_slice::<Complex64>()?[0], Complex64::new(10.0, 0.0));
    assert_eq!(one_sided.shape(), &[3]);

    let data = [1.0_f64, 99.0, 2.0, 99.0, 3.0, 99.0, 4.0];
    let view = TypedTensorView::from_slice([4], [2], 0, &data)?;
    let read = TensorRead::from_view(TensorView::F64(view));
    let read_full = read.fft_read(None, -1, FftNorm::Backward, session)?;
    assert_eq!(
        read_full.as_slice::<Complex64>()?[0],
        Complex64::new(10.0, 0.0),
    );
    Ok(())
})?;

TypedTensor<T> wrappers are not part of the current API. FFT operations can change dtype (rfft real to complex, irfft complex to real), so typed return contracts need a separate design.

Eager Tensors

Use EagerTensorFftExt for immediate execution in an EagerRuntime. The methods have the same names and arguments as TracedTensorFftExt, register the FFT execution runtime on demand, and record the existing extension operation when gradients are enabled.

use num_complex::Complex64;
use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
use tenferro_fft::{EagerTensorFftExt, FftNorm};

let x = EagerTensor::from_tensor_in(
    Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0])?,
    EagerRuntime::new()?,
)?;
let spectrum = x.rfft(None, -1, FftNorm::Backward)?;
let restored = spectrum.irfft(Some(4), -1, FftNorm::Backward)?;

assert_eq!(spectrum.shape(), &[3]);
assert_eq!(restored.to_tensor()?.as_slice::<f64>()?, &[1.0, 2.0, 3.0, 4.0]);

Apple shared CPU and Metal execution

AppleContext makes allocation ownership shared, not backend selection implicit. Clone its mutable backend handles and pass the desired backend to each FFT call:

use tenferro_cpu::CpuBackend;
use tenferro_fft::{FftNorm, TensorFftExt};
use tenferro_gpu::apple::AppleContext;
use tenferro_runtime::BackendSessionHost;
use tenferro_tensor::Tensor;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let context = AppleContext::new()?;
    let host = Tensor::from_vec_col_major([8], vec![1.0_f32; 8])?;
    let managed = context.upload_tensor(&host)?;
    let after_creation = context.transfer_stats();

    let mut cpu = context.cpu_backend().clone();
    let cpu_spectrum = cpu
        .with_backend_session(|session| managed.rfft(None, 0, FftNorm::Backward, session))?;

    let mut metal = context.metal_backend().clone();
    let metal_spectrum = metal
        .with_backend_session(|session| managed.rfft(None, 0, FftNorm::Backward, session))?;
    metal.synchronize()?;

    assert_eq!(cpu_spectrum.shape(), metal_spectrum.shape());
    assert_eq!(context.transfer_stats(), after_creation);
    Ok(())
}

RustFFT supports managed F32, F64, C32, and C64 tensors. The initial CubeK Metal implementation supports C32 CFFT/IFFT, F32 one-sided RFFT, and C32 IRFFT for power-of-two transform lengths of at least 2. RFFT/IRFFT may pad or truncate to a supported requested length; CFFT cannot change the input-axis length. Full-spectrum real FFT, F64/C64, non-power-of-two lengths, foreign domains, and ordinary device-local WebGPU buffers return typed errors. There is no automatic RustFFT fallback.

The complete runnable tutorial additionally proves stable input allocation identity across CPU/Metal/CPU use, equal CPU/Metal results, unchanged post-creation transfer counters, C64 CPU success, and typed C64 Metal failure:

cargo test -p tenferro-tutorial-code --no-default-features \
  --features cpu-faer,apple-shared,doc-snippets --test tutorial_binaries

Source: docs/tutorial-code/src/bin/apple_shared_fft.rs.

Traced Graphs

use num_complex::Complex64;
use tenferro_cpu::CpuBackend;
use tenferro_fft::{FftNorm, TracedTensorFftExt};
use tenferro_runtime::{GraphCompiler, Runtime, TracedTensor};

fn cpu_runtime_with_fft() -> Result<Runtime, Box<dyn std::error::Error>> {
    let backend = CpuBackend::new();
    let mut builder = Runtime::builder();
    builder.register_engine(tenferro_cpu::runtime_engine_registration(&backend)?)?;
    builder.install_extension_module(tenferro_fft::extension_module::<CpuBackend>(
        tenferro_cpu::runtime_engine_id()?,
    )?)?;
    Ok(builder.build()?)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let x = TracedTensor::from_vec_col_major(
        vec![4],
        vec![
            Complex64::new(1.0, 0.0),
            Complex64::new(2.0, 0.0),
            Complex64::new(3.0, 0.0),
            Complex64::new(4.0, 0.0),
        ],
    )
    .unwrap();
    let y = x.fft(None, -1, FftNorm::Backward)?;

    let mut compiler = GraphCompiler::new();
    let program = compiler.compile(&y)?;
    let runtime = cpu_runtime_with_fft()?;
    let mut outputs = runtime.run_compiled(&program, &[])?;
    assert_eq!(outputs.len(), 1);
    let out = outputs.remove(0);
    assert_eq!(out.shape(), &[4]);
    assert_eq!(
        out.as_slice::<Complex64>().unwrap(),
        &[
            Complex64::new(10.0, 0.0),
            Complex64::new(-2.0, 2.0),
            Complex64::new(-2.0, 0.0),
            Complex64::new(-2.0, -2.0),
        ],
    );

    Ok(())
}

For real-input transforms, the transformed axis follows the standard half-spectrum shape rule: input length n produces n / 2 + 1 complex values using integer division. When irfft receives n = None, it infers the output length as 2 * (input_len - 1). That matches even-length round trips; for odd original lengths it silently returns one element too short, so pass Some(original_len).

Planned Extensions

The remaining FFT families are planned but not part of the initial API:

Operation family	Purpose
`fftn`, `ifftn`	multidimensional complex transforms
`rfftn`, `irfftn`	multidimensional real/half-spectrum transforms
`fft2`, `ifft2`, `rfft2`, `irfft2`	two-dimensional convenience wrappers

Compatibility Target

The compatibility target is the behavior users expect from:

torch.fft.fft, torch.fft.ifft, torch.fft.rfft, torch.fft.irfft, and their n/2 variants,
jax.numpy.fft.fft, jax.numpy.fft.ifft, jax.numpy.fft.rfft, jax.numpy.fft.irfft, and their n/2 variants.

The extension should normalize axes and lengths before execution, then return results in the same logical axis order as the input. Backend-specific layout or transposition needed to call an FFT implementation should stay inside the extension.

Automatic Differentiation

FFT is linear, so the extension can support AD through registered extension rules. The current package registers JVP/VJP rules for complex-to-complex fft and ifft: the tangent or cotangent is transformed with the same extension op and normalization.

Use AdContext for explicit extension-rule ownership, or import tenferro_ad::TracedTensorAdExt for the compact traced AD method syntax. For eager AD, construct the runtime with EagerRuntime::with_cpu_backend_and_ad_context using the same AdContext.

Real-to-complex and complex-to-real AD are not enabled yet. They require the usual Hermitian symmetry handling so cotangents match the half-spectrum convention; until those rules are implemented and tested, AD through rfft and irfft reports an unsupported operation instead of returning an incorrect gradient.

Status

tenferro-fft currently lives in the top-level tenferro-fft crate. It supports one-dimensional fft, ifft, rfft, and irfft through CPU RustFFT on host or matching Apple managed tensors, CUDA cuFFT on exact CUDA residency, and the narrower CubeK Metal matrix described above. Multidimensional FFT families remain future work; CUDA one-dimensional execution is implemented.

For the general extension mechanism, see Custom Tensor Operations.