Crate tenferro_fft 

FFT extension operations for tenferro.

This crate is an out-of-tree ExtensionOp package with an explicit FftBackend capability. [tenferro_cpu::CpuBackend] implements the capability through RustFFT. With the webgpu feature, tenferro_gpu::WebGpuBackend executes C32 CFFT, F32 one-sided RFFT, and C32-to-F32 IRFFT through CubeK on its existing WebGPU placement. That first GPU path supports power-of-two lengths only; unsupported operations and dtypes return an error and never fall back to CPU or transfer tensor data. On macOS, tenferro_gpu::AppleContext pairs that Metal backend with a domain-bound CPU RustFFT backend. Backend choice remains explicit, while matching managed tensors can be used without an intervening download. Concrete non-AD execution uses TensorFftExt and TensorReadFftExt. Eager execution uses EagerTensorFftExt when autodiff is enabled, and traced graph construction uses TracedTensorFftExt.

Examples

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

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).unwrap();

let mut compiler = GraphCompiler::new();
let program = compiler.compile(&y).unwrap();
let backend = CpuBackend::new();
let engine_id = tenferro_cpu::runtime_engine_id().unwrap();
let mut builder = Runtime::builder();
builder
    .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
    .unwrap();
builder
    .install_extension_module(tenferro_fft::extension_module::<CpuBackend>(engine_id).unwrap())
    .unwrap();
let runtime = builder.build().unwrap();
let out = runtime.run_compiled(&program, &[]).unwrap().pop().unwrap();
assert_eq!(out.shape(), &[4]);
assert_eq!(out.as_slice::<Complex64>().unwrap()[0], Complex64::new(10.0, 0.0));

use num_complex::Complex32;
use tenferro_cpu::with_cpu_exec_session;
use tenferro_fft::{FftNorm, TensorFftExt};
use tenferro_gpu::{with_webgpu_exec_session, AppleContext};
use tenferro_tensor::{BackendSessionHost, Tensor};

if let Ok(context) = AppleContext::new() {
    let host = Tensor::from_vec_col_major(
        vec![4],
        vec![Complex32::new(1.0, 0.0); 4],
    ).unwrap();
    let input = context.upload_tensor(&host).unwrap();
    let after_creation = context.transfer_stats();
    let mut cpu = context.cpu_backend().clone();
    let cpu_output = cpu
        .with_backend_session(|session| {
            with_cpu_exec_session(session, |exec_session| {
                input.fft(None, 0, FftNorm::Backward, exec_session)
            })
            .expect("CpuBackend must expose a CPU execution session")
        })
        .unwrap();
    let mut metal = context.metal_backend().clone();
    let output = metal
        .with_backend_session(|session| {
            with_webgpu_exec_session(session, |exec_session| {
                input.fft(None, 0, FftNorm::Backward, exec_session)
            })
            .expect("WebGpuBackend must expose a WebGPU execution session")
        })
        .unwrap();
    metal
        .with_backend_session(|session| {
            with_webgpu_exec_session(session, |exec_session| {
                exec_session.runtime().synchronize()
            })
            .expect("WebGpuBackend must expose a WebGPU execution session")
        })
        .unwrap();
    assert_eq!(output.shape(), &[4]);
    assert_eq!(cpu_output.shape(), output.shape());
    assert_eq!(context.transfer_stats(), after_creation);
}

use num_complex::Complex64;
use tenferro_cpu::{with_cpu_exec_session, CpuBackend};
use tenferro_fft::{FftNorm, TensorFftExt};
use tenferro_tensor::{BackendSessionHost, Tensor};

let x = Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
let mut backend = CpuBackend::new();
let out = backend
    .with_backend_session(|session| {
        with_cpu_exec_session(session, |exec_session| {
            x.fft(None, -1, FftNorm::Backward, exec_session)
        })
        .expect("CpuBackend must expose a CPU execution session")
    })
    .unwrap();

assert_eq!(out.as_slice::<Complex64>().unwrap()[0], Complex64::new(10.0, 0.0));

Structs

FftExecutionCache

Execution-cache state supplied to an FftBackend.

FftExecutor

Reusable concrete FFT executor with an explicitly owned backend-neutral cache.

FftPlanCache

Bounded, caller-owned typed cache for backend FFT plans and workspaces.

FftPlanSpec

Validated, backend-neutral description of one FFT request.

Enums

FftNorm

FFT normalization convention.

FftOperation

One-dimensional FFT operation requested from an FftBackend.

Constants

DEFAULT_FFT_PLAN_CACHE_CAPACITY

Default number of typed entries retained by a caller-owned FftPlanCache.

FFT_EXTENSION_FAMILY_ID

Extension family id used by the tenferro FFT extension.

FFT_PLAN_CACHE_NAME

Runtime cache namespace used for private RustFFT plans.

Traits

FftBackend

Explicit backend capability required by concrete and traced FFT execution.

TensorFftExt

Backend-explicit FFT methods for concrete Tensor values.

TensorReadFftExt

Backend-explicit FFT methods for read-only tensor inputs.

TracedTensorFftExt

FFT extension methods for [TracedTensor].

Functions

extension_module

Build this extension module for one runtime engine.

fft_plan_cache_selector

Select the private CPU RustFFT plan entries in an extension runtime cache.