Skip to content
specux

Plans & autotuning

For hot loops on CUDA with fixed parameters, build the transform once as a plan: the kernel compiles once, the plan runs by raw pointer on the caller’s stream, and the tuning knobs live here; the functional API always uses the auto planner.

specux.stft_plan(n_fft=2048, hop_length=None, win_length=None, window="hann",
                 center=True, output="complex", out_dtype=None, eps=1e-10,
                 dtype="float32", device="cuda",
                 ept=0, fpb=0, tpf=0, core="auto", tune=None)
xt = torch.randn(8, 32768, device="cuda")
plan = specux.stft_plan(n_fft=1024, hop_length=256, output="power")
plan = specux.autotune(xt, plan)     # sweep candidates, keep the fastest
S = plan(xt)                         # microseconds per call from here on

The transform parameters (n_fft through eps) mean what they mean on stft; dtype is the compute/storage precision the plan is specialized for ("float32", "float64", "float16"), and device names the CUDA device ("cuda", "cuda:1"). Plans run forward-only: no gradients flow through a plan call.

Tuning knobs

Leave them at 0 / "auto" unless you’re chasing the last microseconds: the auto planner picks well, and autotune measures instead of guessing.

  • ept: elements per thread (0 = auto planner).
  • fpb: frames per block (0 = auto planner).
  • tpf: threads per frame (0 = auto planner).
  • core: kernel core override: "auto", "ct", "warp", "ept16".
  • tune=: pass a representative tensor to autotune at build time.

Autotuning

specux.autotune(x, plan, iters=50, warmup=10) times every candidate configuration on your actual input, keeps the fastest (the plan is modified in place and returned), and persists the winner to the on-disk tuning cache (SPECUX_CACHE_DIR). The cache key is the GPU plus the plan configuration, never the input shape, so one tune serves every signal length and batch size and survives across processes; still pass an x with a production-like batch, since the candidates are timed on it.

Prefer not to manage plans at all? specux.benchmark(True) (or SPECUX_BENCHMARK=1) makes every configuration autotune-and-cache itself on first use; see Runtime options.

FFT plans

specux.fft_plan(n, real=...) is the FFT-family counterpart: a reusable length-bound handle with forward / inverse / bins. The FFT kernels are fully determined by the length, so there is nothing to autotune; see Plans: fft_plan.