Skip to content

Template Metal C2C FFT scalar lanes - #3969

Open
PhysicistJohn wants to merge 1 commit into
ml-explore:mainfrom
PhysicistJohn:codex/generic-metal-fft
Open

Template Metal C2C FFT scalar lanes#3969
PhysicistJohn wants to merge 1 commit into
ml-explore:mainfrom
PhysicistJohn:codex/generic-metal-fft

Conversation

@PhysicistJohn

@PhysicistJohn PhysicistJohn commented Aug 1, 2026

Copy link
Copy Markdown

Motivation

MLX's Metal C2C FFT core hard-codes float2. This PR is the narrow Metal-kernel foundation for #2112: it removes that source-level coupling so a later, end-to-end reduced-precision complex proposal can be evaluated without rewriting the FFT implementation first.

The target use case is complex I/Q and signal-processing workloads, where storing real and imaginary lanes in FP16 or BF16 halves storage and bytes transferred relative to two FP32 lanes. This PR does not claim that lower precision is appropriate for every workload; public dtype semantics, host dispatch, operation coverage, accuracy policy, and new static instantiations remain separate work.

Existing accelerator-library precedent

This is not a novel representation:

  • NVIDIA cuFFT supports half-complex CUDA_C_16F and bfloat16-complex CUDA_C_16BF transforms through its type-generic planning API.
  • AMD rocFFT exposes rocfft_precision_half together with complex interleaved and planar array layouts.

Those are direct GPU-vendor FFT libraries. This patch only prepares MLX's C2C Metal implementation for equivalent storage lanes; it does not add a public reduced-precision complex type.

Scope

  • Template C2C radix arithmetic, shared memory, and I/O conversion over the scalar lane type.
  • Preserve the current public API, host dispatch, static kernel-instantiation table, and packed float32 RFFT/IRFFT paths.
  • Generalize the C2C four-step second pass so every eligible lane type retains the current float32 indexing, inversion, and normalization behavior.

Binary-size impact

No new static kernels or public exports are instantiated.

  • Local release mlx.metallib: 130,977,720 B → 130,994,264 B
  • Delta: +16,544 B (+0.01263%)
  • metal-nm -g exported-symbol output: exact match to the base build

Validation

  • Native xcrun metal compile and metallib link of the changed fft.metal using MLX's build flags.
  • Native CMake build with deployment target macOS 15.0 (NAX kernels excluded by the existing build gate): ctest 261/261 passed, including all FFT cases.
  • Branch-local Metal Python extension: python -m unittest test_fft -v16 passed, 2 expected skips.
  • Compile-only C2C instantiations for half2 and bfloat2 across Stockham, Rader, Bluestein, and four-step planners.
  • pre-commit (clang-format) and git diff --check passed.

Focused float32 regression checks

Path Correctness Generic / base median
C2C, N=256, batch=4096 Bit-identical 0.2952 / 0.2954 ms (0.9993x)
Four-step second pass Bit-identical 0.2384 / 0.2350 ms (1.0145x)

The second measurement is a focused dispatch-scale check, not a claim of a material performance change. There is no end-to-end reduced-precision performance claim until a future public dtype and dispatch path exists.

@PhysicistJohn
PhysicistJohn marked this pull request as draft August 1, 2026 19:27
@PhysicistJohn
PhysicistJohn force-pushed the codex/generic-metal-fft branch from 4b0e3b7 to 90b7f12 Compare August 1, 2026 19:39
@PhysicistJohn
PhysicistJohn marked this pull request as ready for review August 1, 2026 19:40
@PhysicistJohn
PhysicistJohn marked this pull request as draft August 1, 2026 19:43
Template C2C FFT arithmetic and I/O over their scalar lane while retaining the existing static entry points and packed real paths.

Make the C2C four-step second pass use the same generic storage path so future scalar lanes preserve the existing indexing and scaling behavior.
@PhysicistJohn
PhysicistJohn force-pushed the codex/generic-metal-fft branch from 90b7f12 to 3bac56f Compare August 1, 2026 20:03
@PhysicistJohn PhysicistJohn changed the title Generalize Metal FFT complex scalar handling Template Metal C2C FFT scalar lanes Aug 1, 2026
@PhysicistJohn
PhysicistJohn marked this pull request as ready for review August 1, 2026 20:15
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