Skip to content

Lower the FFT op family in the native Expr compiler - #183

Merged
ausimian merged 1 commit into
feat/expr-compilerfrom
feat/expr-compiler-fft
Jun 6, 2026
Merged

Lower the FFT op family in the native Expr compiler#183
ausimian merged 1 commit into
feat/expr-compilerfrom
feat/expr-compiler-fft

Conversation

@ausimian

@ausimian ausimian commented Jun 6, 2026

Copy link
Copy Markdown
Owner

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. Same shape as the
take_along_axis addition: the MLX ops and eager NIFs already existed
(Native.{fftn,ifftn,rfftn,irfftn}mlx::core::fft::*); only the
compiler path was missing.

What surfaced this

Running a Whisper speech_to_text serving with native: true logged:

Emily Expr compiler does not yet lower op :fft (no fallback). It will
be added in a later milestone. The defn ran op-by-op via the evaluator.

That's the log-mel featurizer's STFT. With this change that defn compiles
fully native. (The whisper_full conformance native lane never caught it
— it tests the encoder/decoder forward on pre-computed features, not the
featurizer.)

Changes

  • c_src/emily/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. 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 use
    the last axis and ignore opts[:axis], so native stays bit-identical to
    the evaluator); lower Nx.Block.{FFT2,IFFT2,RFFT,IRFFT} (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, rfft/irfft.

Verification

  • New FFT equivalence cases pass under native: true (test config defaults
    native_fallback: :raise, so a fallback would raise — passing proves
    full native lowering). Opcode-parity test confirms the IR↔C++ enum
    lockstep (88–91 / kOpcodeCount 92).
  • mix precommit green: compile (warnings-as-errors), format, credo
    --strict, docs, 686 tests, 0 failures.

Resolves the :fft gap noted while trying the native compiler on the
Whisper livebook.

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.
@ausimian
ausimian merged commit 88d9f51 into feat/expr-compiler Jun 6, 2026
3 checks passed
@ausimian
ausimian deleted the feat/expr-compiler-fft branch June 6, 2026 06:25
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