Enable half-precision complex Metal FFT kernels - #7
Open
PhysicistJohn wants to merge 1 commit into
Open
Conversation
Instantiate C2C Stockham, Rader, Bluestein, and four-step kernels for packed complex_t<half> storage. Move convolution and four-step inverse scaling ahead of reduced-precision accumulation so representable results do not overflow, and compute reciprocals before narrowing to the storage lane.
PhysicistJohn
marked this pull request as ready for review
August 2, 2026 17:25
This was referenced Aug 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
This is one dependent commit on top of the storage adapter in #6 and the
generic Metal foundations in ml-explore#3969 and ml-explore#3970.
Summary
Enable packed
complex_t<half>C2C FFT kernels for Stockham, Rader,Bluestein, and four-step plans. The patch adds the raw Metal backend capability;
public dtype and host-dispatch changes are outside this diff.
Comparable GPU-library support is established in
NVIDIA cuFFT, which supports
half-precision C2C transforms, and
AMD rocFFT,
which supports FP16 real and complex transforms.
The change is five FFT files with 60 insertions and 33 deletions.
Implementation
complex32_t = complex_t<half>and assert its four-byte, two-lanelayout.
sizes.
accumulation, avoiding overflow of finite half-complex results in temporary
buffers.
Correctness
Candidate FP32 output is bit-identical to untouched
mainFP32 for therepresentative Stockham, Rader, Bluestein, and four-step cases below.
A separate half-complex sweep covered 39 lengths from 2 through 1,048,576.
Every forward and inverse result was finite. Worst normalized RMSE was 0.2585%
forward and 0.4101% round trip; the 1,048,576-point round-trip error was 0.2685%.
Performance against untouched main
Apple M5 Max; untouched
mainatfb5133e1; this PR at3c63352d. Eachvariant used the same source-level algorithm harness, FFT plans, dispatch
geometry, 4,194,304 complex values, input samples, five warmups, 21 timed
samples, and 50 transforms per sample. Six runs covered all six execution
orders of main FP32, candidate FP32, and candidate half. Values below are the
median of the six run medians.
Forward
Inverse
The candidate FP32 control measured 0.992x-1.065x the original FP32 throughput
across these eight cases while remaining bit-identical. A second six-run Rader
and Bluestein measurement reproduced their inverse results.
Packed complex storage and external I/O are four bytes per value instead of
eight. These are direct kernel measurements, not public
mx.fftAPI results.Validation
-Wall -Wextra -Werror.integration tree.
symbols.
git diff --check.Binary impact
fft.airmlx.metallib