Lower the FFT op family in the native Expr compiler - #183
Merged
Conversation
Nx.fft / ifft (1-D, trailing axis) and the fft2 / ifft2 / rfft / irfft
blocks now lower under `compiler: Emily.Compiler, native: true` instead
of forcing a graceful fallback to the evaluator. The MLX ops and eager
NIFs already existed (Native.{fftn,ifftn,rfftn,irfftn} ->
mlx::core::fft::*); only the compiler path was missing.
- opcodes.hpp: add Fftn/Ifftn/Rfftn/Irfftn (88-91), bump kOpcodeCount to
92, dispatch each to the matching mx::fft::* with FFTNorm::Backward
(unnormalized), mirroring the eager fft.cpp entry points.
- ir.ex: add the opcodes; lower :fft/:ifft (lower_op) on the trailing
axis (mirroring Emily.Backend.{fft,ifft}/3, which ignore opts[:axis]);
lower the Nx.Block.{FFT2,IFFT2,RFFT,IRFFT} blocks (lower_block) with the
block's sizes/axes. Output coerced to out.type (complex64 for the
forward transforms, real for irfft) like the backend wrap.
- compiler_equivalence_test.exs: native-vs-evaluator bit-identical cases
for fft/ifft (+ explicit length, batched), fft2/ifft2, and rfft/irfft.
This was the op forcing a fallback in a Whisper speech_to_text serving
(the log-mel featurizer's STFT), so that path now compiles fully native.
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.
Nx.fft/ifft(1-D, trailing axis) and thefft2/ifft2/rfft/irfftblocks now lower under
compiler: Emily.Compiler, native: trueinstead offorcing a graceful fallback to the evaluator. Same shape as the
take_along_axisaddition: the MLX ops and eager NIFs already existed(
Native.{fftn,ifftn,rfftn,irfftn}→mlx::core::fft::*); only thecompiler path was missing.
What surfaced this
Running a Whisper
speech_to_textserving withnative: truelogged:That's the log-mel featurizer's STFT. With this change that defn compiles
fully native. (The
whisper_fullconformance native lane never caught it— it tests the encoder/decoder forward on pre-computed features, not the
featurizer.)
Changes
c_src/emily/opcodes.hpp— addFftn/Ifftn/Rfftn/Irfftn(88–91), bump
kOpcodeCountto 92, dispatch each to the matchingmx::fft::*withFFTNorm::Backward(unnormalized), mirroring the eagerfft.cppentry points. operands[input]; iattrs[[sizes...],[axes...]].lib/emily/ir.ex— add the opcodes; lower:fft/:ifft(lower_op)on the trailing axis (mirroring
Emily.Backend.{fft,ifft}/3, which usethe last axis and ignore
opts[:axis], so native stays bit-identical tothe evaluator); lower
Nx.Block.{FFT2,IFFT2,RFFT,IRFFT}(lower_block)with the block's sizes/axes. Output coerced to
out.type(complex64 forthe forward transforms, real for
irfft) like the backendwrap.compiler_equivalence_test.exs— native-vs-evaluator bit-identicalcases for fft/ifft (+ explicit
length, batched), fft2/ifft2, rfft/irfft.Verification
native: true(test config defaultsnative_fallback: :raise, so a fallback would raise — passing provesfull native lowering). Opcode-parity test confirms the IR↔C++ enum
lockstep (88–91 /
kOpcodeCount92).mix precommitgreen: compile (warnings-as-errors), format, credo--strict, docs, 686 tests, 0 failures.
Resolves the
:fftgap noted while trying the native compiler on theWhisper livebook.