rlx-fft

Learned butterfly FFT + spectral pipelines (mel, Welch PSD, top-K Welch peaks), compiled via RLX.

Workspace 0.2.7 — depends on upstream rlx* 0.2.6. Publish tier 2 (scripts/publish.sh), after rlx-cli / rlx-models-core.

cargo run -p rlx-fft --release -- --help

Welch peaks (fast top-K spikes)

Extract top-K frequency spikes (bin, power) without materializing a full Welch PSD. The fast path uses 2 Welch segments (vs 8 in full Welch); an ultra-fast path uses 1 segment for minimum latency.

CLI bench

# Auto strategy (default) — picks fastest path for batch + device
cargo run -p rlx-fft --release -- bench-welch-peaks \
  --n-fft 256 --batch 32 --k 16 --train-steps 0

# Batch sweep + Metal GPU crossover
cargo run -p rlx-fft --features apple-silicon --release -- bench-welch-peaks \
  --n-fft 256 --batch 32,256,1024,4096,8192 --device metal --train-steps 0 --iters 15

# Force a specific strategy (see table below)
cargo run -p rlx-fft --release -- bench-welch-peaks \
  --n-fft 256 --batch 32 --strategy ultra

cargo run -p rlx-fft --features apple-silicon --release -- bench-welch-peaks \
  --n-fft 256 --batch 8192 --device metal --strategy rlx

# K sweep — plot latency vs top-K (JSON rows tagged with batch + k)
cargo run -p rlx-fft --features apple-silicon --release -- bench-welch-peaks \
  --n-fft 256 --batch 8192 --k 4,8,16,32,64 --device metal --train-steps 0 --iters 15 \
  --strategy rlx --json /tmp/welch-k-sweep.json

Sweep output ends with a k crossover table (rustfft / stream / rlx / picker ms per K). Combine with --batch for a full grid, e.g. --batch 32,8192 --k 4,16,64.

Fusion phase bench (IO + latency)

Compare baseline interleaved readback, Phase 1 block layout, and Phase 2 fused Op::WelchPeaks:

cargo run -p rlx-fft --features dev,apple-silicon --release -- bench-fusion-phases \
  --n-fft 256 --batch 8192 --k 16 --device metal --iters 15

# Batch sweep + JSON
cargo run -p rlx-fft --features dev,apple-silicon --release -- bench-fusion-phases \
  --n-fft 256 --batch 32,1024,8192 --k 16 --device metal --iters 15 \
  --json /tmp/fusion-phases.json

Output includes IO profiles (kernel launches, sync points, host readback bytes) and per-phase speedup vs baseline.

# WGPU (Vulkan/Metal/DX12 via wgpu)
cargo run -p rlx-fft --features dev,gpu --release -- bench-fusion-phases \
  --n-fft 256 --batch 8192 --k 16 --device wgpu --iters 15

# CUDA (when NVIDIA toolkit + `rlx-runtime/cuda` available)
cargo run -p rlx-fft --features dev,cuda --release -- bench-fusion-phases \
  --n-fft 256 --batch 8192 --k 16 --device cuda --iters 15

Bench output includes a welch_peaks_picker_<strategy> row using auto or forced selection, e.g.:

[welch-peaks] picker (auto): batch=8192 device=Metal -> rlx_compiled

Strategy picker

Use AutoWelchPeaks in Rust or --strategy on the CLI.

Auto selection (IO-aware picker)

Auto mode estimates each strategy with an Ayala-style latency–bandwidth model (T ≈ L·M + S/W) using graph_io profiles and per-device BackendCostModel (CPU rustfft paths vs fused Op::WelchPeaks on GPU). Fused GPU estimates apply a calibrated compute scale (~7.5× IO-only on Metal, from bench-fusion-phases phase-2); CPU rustfft gets a batch growth penalty when compared on GPU devices. It picks the lowest predicted cost.

Calibrate with bench-fusion-phases (phase-2 fused vs IO-model line; prints suggested fused_io_compute_scale) and bench-welch-peaks (picker vs rustfft crossover). NVIDIA CUDA: ../rlx/rig.sh bench-rlx-fft-welch-peaks windows 256,1024,8192 cuda or quick picker check ../rlx/rig.sh bench-welch-peaks-fft windows 8192 cuda (Mar 2026 rig: CUDA scale 0.43, auto picker → rlx_compiled / FusedOp ~19 ms vs rustfft ~63 ms at batch 8192). Metal/Mlx use unified-memory rustfft penalties. WGPU/Vulkan use native in-arena WelchPeaks at large batch; preliminary WGPU scale ~2.2. welch_peaks_io_fusion_gate / welch_peaks_fusion_gate_breakdown use rlx_compile::IoFusionGate (readback savings minus host-thunk penalty). At compile time, rlx_compile::SelectPeaksOnlyOutputs runs in the fusion pipeline when the backend claims Fft + WelchPeaks: it drops redundant spectrum outputs and promotes peaks-only readback when the IO gate favors fusion (bench-fusion-phases phase3 row). welch_peaks_fusion_gate_breakdown exposes should_fuse_io vs final should_fuse (adds a large-batch compute floor vs block RLX). Auto picker only estimates rlx / learned when both fused_welch_peaks_auto_viable and the IO gate pass (Metal rejects fusion below ~batch 512). At runtime, CompiledRlxWelchPeaksExec::compile_adaptive picks fused Op::WelchPeaks or block FFT + host top-K (rlx_welch_peaks_exec_kind). AutoWelchPeaks::welch_peaks_batch accepts full 8-segment or fast 2-segment layout; use welch_peaks_batch_fast when you already have the fast buffer.

just features=apple-silicon bench-welch-peaks -- --n-fft 256 --batch 1024,8192 --k 16 --device metal
just bench-fusion-phases -- --n-fft 256 --batch 1024,8192 --k 16 --device metal --iters 15

Legacy reference thresholds (used only with RLX_FFT_LEGACY_PICKER): batch ≤ 256 CPU / ≤ 128 GPU → ultra; mid batch → fast; batch ≥ 8192 GPU → rlx; sparse learned gates → learned.

Reference peaks for training/bench error always use full 8-segment Welch; student paths use 1–2 segments.

Rust API

use rlx_fft::{
    AutoWelchPeaks, WelchPeaksPickMode, WelchPeaksStrategy,
    parse_welch_peaks_strategy, pick_welch_peaks_strategy,
};

// Auto (recommended)
let mut picker = AutoWelchPeaks::new(batch, n_fft, k, Some("auto"))?;
println!("strategy: {}", picker.strategy_label());

// Force a strategy
let mut picker = AutoWelchPeaks::with_strategy(
    batch, n_fft, k, Some("metal"), WelchPeaksStrategy::RlxCompiled,
)?;

// Parse CLI-style string
let mode = parse_welch_peaks_strategy("fast")?; // Force(FastStreaming)
let mut picker = AutoWelchPeaks::with_options(
    batch, n_fft, k, Some("cpu"), None, mode,
)?;

// With learned model (for learned strategy or auto sparse-gate path)
let mut picker = AutoWelchPeaks::with_learned(
    batch, n_fft, k, Some("metal"), Some(&model),
)?;

// Full 8-segment layout (e2e / reference pipelines)
let peaks = picker.welch_peaks_batch(&signal)?;

// Fast 2-segment layout (production hot path — no truncate copy)
let fast_signal = fast_params.welch.truncate_batch(&signal, batch, full_frame)?;
let peaks = picker.welch_peaks_batch_fast(&fast_signal)?;

Strategy string aliases (for parse_welch_peaks_strategy / --strategy):

Performance notes (n=256, Apple Silicon reference)

RLX compiled paths need large batch to amortize GPU launch; rustfft wins at small batch.

Training peaks into the learned model

End-to-end training includes a peak-matching loss on the fast 2-segment path. --k / --peak-k sets how many spikes are matched during training and at inference (learned, compiled-learned, and picker learned strategy).

cargo run -p rlx-fft --release -- train-e2e \
  --n-fft 256 --batch 8 --peak-k 16 --peak-weight 2.0 --steps 2000

# bench-e2e: same K for WelchPeaks pipelines + teacher training
cargo run -p rlx-fft --release -- bench-e2e \
  --n-fft 256 --batch 8 --peak-k 8 --train-first --steps 500

At inference, FastLearnedFftModel::welch_peaks_batch accepts any WelchPeakParams::fast_for_n_fft(n_fft, k) — K is not baked into weights, but training with the target K improves peak accuracy.

Tests

just features=apple-silicon test-rlx-fft-welch-peaks
just test-rlx-fft-fusion-gate
cargo test -p rlx-fft welch_peaks_compile::tests --features apple-silicon

Modules