Enable half-precision complex Metal FFT kernels - #3981
Merged
Conversation
PhysicistJohn
force-pushed
the
draft/raw-complex-half-fft
branch
from
August 6, 2026 19:12
3c63352 to
e15932f
Compare
This was referenced Aug 6, 2026
PhysicistJohn
force-pushed
the
draft/raw-complex-half-fft
branch
from
August 7, 2026 03:43
e15932f to
72975b0
Compare
zcbenz
approved these changes
Aug 8, 2026
zcbenz
left a comment
Member
There was a problem hiding this comment.
Looks good to me, thanks for splitting the PRs to make them easy to review.
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.
zcbenz
force-pushed
the
draft/raw-complex-half-fft
branch
from
August 8, 2026 01:19
72975b0 to
7479a76
Compare
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.
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.
Implementation
complex32_t = complex_t<half>and assert its four-byte, two-lanelayout.
FFTStorageTraitsspecialization derives the scalar lane fromcomplex_t<T>storage,
complex_t<T>converts to and from the two-lane vector type, and apost_inoverload lifts storage values into lanes.threadgroup-memory sizes.
normalization ahead of reduced-precision accumulation, avoiding overflow of
representable half-complex intermediates. The reordering is gated on the
scalar type, so the float path keeps the existing arithmetic exactly.
Correctness
The float path is bitwise unchanged by construction: every reordering above is
gated on
!metal::is_same_v<scalar_T, float>, and a host-side bitwise harnessconfirms identical outputs against current
mainfor codelets, twiddles,Stockham, four-step (forward and inverse), and the reader/writer paths. All
FFT CTest cases pass.
A half-complex sweep on the isolated pre-integration implementation 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%. Happy to rerun the sweep on this
tree if useful.
Performance against untouched main
Apple M5 Max; untouched
mainatfb5133e1; measurements on the isolatedpre-integration implementation at
3c63352d. Each variant used the samesource-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 half kernels are unchanged by the integration, so the half speedups above
carry over. The FP32 control column reflects the pre-gating implementation, in
which the normalization reorderings also ran on float; in this tree those
reorderings are gated off for float, so the merged FP32 path is bitwise
identical to
mainby construction and performance-neutral, and is at least asfast as the measured control. 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
xcrun metalcompilation offft.metalwith allcomplex32_tinstantiations, plus a compile-only storage regression exercising the half
and bfloat16 lane plumbing.
main.overflows half without the early normalization.
LAPACK linalg cases fail on this machine, identically on untouched
main—a local toolchain issue, unrelated to this diff.)
pre-commit(clang-format) andgit diff --check.symbols.
Binary impact
mlx.metallibThe five threadgroup-size instantiations per planner are byte-identical code
differing only in scratchpad allocation, mirroring the existing float FFT
layout. Measured in isolation, the new half planner code is 182 KB versus
180 KB for the float equivalents (+1%), so the binary impact is the
established static-sizing structure applied to one more storage type, not new
weight from the dtype.