Skip to content

Enable half-precision complex Metal FFT kernels - #3981

Merged
zcbenz merged 2 commits into
ml-explore:mainfrom
PhysicistJohn:draft/raw-complex-half-fft
Aug 8, 2026
Merged

Enable half-precision complex Metal FFT kernels#3981
zcbenz merged 2 commits into
ml-explore:mainfrom
PhysicistJohn:draft/raw-complex-half-fft

Conversation

@PhysicistJohn

@PhysicistJohn PhysicistJohn commented Aug 2, 2026

Copy link
Copy Markdown
Contributor

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

  • Define complex32_t = complex_t<half> and assert its four-byte, two-lane
    layout.
  • Bridge packed complex storage into the FFT reader/writer: an
    FFTStorageTraits specialization derives the scalar lane from complex_t<T>
    storage, complex_t<T> converts to and from the two-lane vector type, and a
    post_in overload lifts storage values into lanes.
  • Instantiate the existing C2C kernel templates at all five
    threadgroup-memory sizes.
  • Compute reciprocal scales in float before narrowing to the storage lane.
  • For reduced lanes only, apply inverse convolution and four-step
    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 harness
confirms identical outputs against current main for 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 main at fb5133e1; measurements on the isolated
pre-integration implementation at 3c63352d. Each variant 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

Plan Main FP32 Candidate FP32 Candidate half Half vs main Half vs candidate FP32
Stockham 0.09141 ms 0.09175 ms 0.04250 ms 2.151x 2.159x
Rader 0.09412 ms 0.09441 ms 0.07392 ms 1.273x 1.277x
Bluestein 0.20906 ms 0.21075 ms 0.17941 ms 1.165x 1.175x
Four-step 0.27790 ms 0.27830 ms 0.12398 ms 2.242x 2.245x

Inverse

Plan Main FP32 Candidate FP32 Candidate half Half vs main Half vs candidate FP32
Stockham 0.09251 ms 0.09299 ms 0.04301 ms 2.151x 2.162x
Rader 0.09837 ms 0.09317 ms 0.07418 ms 1.326x 1.256x
Bluestein 0.22584 ms 0.21198 ms 0.18189 ms 1.242x 1.165x
Four-step 0.27614 ms 0.27642 ms 0.12436 ms 2.220x 2.223x

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 main by construction and performance-neutral, and is at least as
fast 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.fft API results.

Validation

  • Native xcrun metal compilation of fft.metal with all complex32_t
    instantiations, plus a compile-only storage regression exercising the half
    and bfloat16 lane plumbing.
  • Host-side bitwise harness of the float paths against current main.
  • Half-path checks against double-precision references, plus a tone case that
    overflows half without the early normalization.
  • Static Release build and full serial CTest: all FFT cases pass. (Seven
    LAPACK linalg cases fail on this machine, identically on untouched main
    a local toolchain issue, unrelated to this diff.)
  • pre-commit (clang-format) and git diff --check.
  • Exactly 25 intended half-complex C2C entry points; no duplicate or half-real
    symbols.

Binary impact

Artifact Increase
Complete mlx.metallib 987,536 bytes (0.60%)

The 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.

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me, thanks for splitting the PRs to make them easy to review.

PhysicistJohn and others added 2 commits August 8, 2026 10:17
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
zcbenz force-pushed the draft/raw-complex-half-fft branch from 72975b0 to 7479a76 Compare August 8, 2026 01:19
@zcbenz
zcbenz merged commit a88d454 into ml-explore:main Aug 8, 2026
28 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants