Reuse Stockham twiddles in large batched Bluestein FFTs

[PhysicistJohn]

Proposed changes

Fused Bluestein FFTs run the same Stockham plan twice and currently recompute its radix twiddles in each pass. For complex64 transforms using the 4096-point internal plan with total batch at least 1024, this change dispatches one private Metal kernel to generate a 584-value twiddle table, then reuses it in both internal transforms. Smaller batches, other plans, and real transforms retain dynamic twiddles. There are no public API changes.

Results

The implementation was benchmarked against main at bb6d960f on an Apple M5 Max using static Metal. Across 20 selected end-to-end fft and ifft cells, the median of the per-cell candidate/main ratios was 1.03856x (3.856%). Each full-matrix cell used six balanced fresh-process rounds, 10 warmups, and nine samples of 30 transforms. Fresh 12- or 24-round targeted reruns placed every initially noisy candidate/main 95% lower bound above 1.0. The targeted 2047-point, batch-2048 cases remained near-neutral at 1.0056x forward and 1.0028x inverse. An unchanged 4096-point Stockham control measured 0.9984x, with a 95% interval of [0.9817, 1.0186].

All 48 deterministic accuracy comparisons passed. Relative to the benchmark baseline, the selected path adds 16,384 bytes of measured peak allocation, 38,432 bytes to mlx.metallib, 4,080 bytes to libmlx.a, and one private Metal kernel entry point.

Validation

• Resealed on exact current main at 7729d587
• Static and JIT Metal builds passed with warnings treated as errors
• Static FFT module: 16 passed, 2 expected skips
• JIT FFT module: 16 passed, 2 expected skips
• In both build-directory native runs, the candidate reproduced the same six failing linalg cases and 17 failed assertions as the exact-main baseline; there were no candidate-only failures
• Pre-commit on all six changed files and diff check passed

Checklist

• I have read the contributing guide
• I have run pre-commit on all changed files
• I have added tests for the selected and threshold-control paths
• Documentation is not needed because there is no public API change

[zcbenz]

Precomputing twiddles should help but doing it in host seems to be eating a lot of performance. Can you try doing it with a metal kernel? We can't really merge changes like this.

[PhysicistJohn]

Thanks, Cheng. Agreed. I moved the table generation onto Metal: one private kernel writes the 584 Stockham twiddles, which the selected fused Bluestein path reuses in both internal passes. Against bb6d960f, the median candidate/main ratio across 20 selected end-to-end cases was 1.03856x; targeted reruns cleared the initially noisy lower bounds, with the 2047-point boundary remaining near-neutral. I also resealed the exact diff on current main at 7729d587: static and JIT Werror builds passed, both FFT runs were 16 passed with 2 expected skips, and the native runs reproduced only the same six linalg cases and 17 failed assertions as baseline. All 48 accuracy comparisons passed.

[zcbenz]

Nice improvement, thanks!