Skip to content

Enable half-precision complex Metal FFT kernels - #7

Open
PhysicistJohn wants to merge 1 commit into
draft/reduced-complex-storage-bridgefrom
draft/raw-complex-half-fft
Open

Enable half-precision complex Metal FFT kernels#7
PhysicistJohn wants to merge 1 commit into
draft/reduced-complex-storage-bridgefrom
draft/raw-complex-half-fft

Conversation

@PhysicistJohn

@PhysicistJohn PhysicistJohn commented Aug 2, 2026

Copy link
Copy Markdown
Owner

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

  • Define complex32_t = complex_t<half> and assert its four-byte, two-lane
    layout.
  • Instantiate the existing C2C kernel templates at all five threadgroup-memory
    sizes.
  • Compute reciprocal scales before narrowing to the storage lane.
  • Apply inverse convolution and four-step normalization before reduced-precision
    accumulation, avoiding overflow of finite half-complex results in temporary
    buffers.

Correctness

Candidate FP32 output is bit-identical to untouched main FP32 for the
representative 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 main at fb5133e1; this PR 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 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.fft API results.

Validation

  • Metal macOS 14 compilation with -Wall -Wextra -Werror.
  • Static Release build and full serial CTest: 262/262 on the reviewed
    integration tree.
  • JIT Release build and FFT CTests: 7/7.
  • Exactly 25 intended half-complex C2C entry points; no duplicate or half-real
    symbols.
  • Clang-format dry run and git diff --check.

Binary impact

Artifact Increase
fft.air 745,408 bytes
Complete mlx.metallib 994,384 bytes (0.759%)

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

1 participant