Merge pull request #183 from ausimian/feat/expr-compiler-fft

Lower the FFT op family in the native Expr compiler

Author: ausimian

RELEASE.md

@@ -58,6 +58,15 @@
   gather, so a DistilBERT question-answering `Nx.Serving` forward now runs
   fully native — and fused — under `native_fallback: :raise`.
 
+- **The FFT family lowers natively** — `Nx.fft`/`ifft` (1-D, trailing axis)
+  and the `fft2`/`ifft2`/`rfft`/`irfft` blocks now compile under the native
+  single-NIF path instead of falling back. Each mirrors its eager
+  `Emily.Backend` wrapper, routing to the same `mlx::core::fft::*` kernel
+  (unnormalized `FFTNorm::Backward`) bit-for-bit, including the complex64
+  outputs. This was the op forcing a graceful fallback in a Whisper
+  `speech_to_text` serving — the log-mel featurizer's STFT — so that path
+  now compiles fully native too.
+
 - **Window (pooling) ops lower natively — forward and backward.** The
   forward window family (`window_sum`/`window_max`/`window_min`/
   `window_product`, i.e. average and max pooling), the select-and-scatter

c_src/emily/opcodes.hpp

@@ -168,9 +168,16 @@ enum class Opcode : int64_t {
   // [[window...],[strides...],[pad_lo...],[pad_hi...]] (no dilations)
   WindowScatterMax = 86,
   WindowScatterMin = 87,
+  // FFT family — n-dimensional transforms over the given sizes/axes, with
+  // FFTNorm::Backward (unnormalized), matching Nx + the eager fft NIFs.
+  // operands [input]; iattrs [[sizes...], [axes...]].
+  Fftn = 88,   // complex/real -> complex
+  Ifftn = 89,  // complex -> complex (inverse)
+  Rfftn = 90,  // real -> complex (half spectrum)
+  Irfftn = 91, // complex half-spectrum -> real
 };
 
-inline constexpr int64_t kOpcodeCount = 88;
+inline constexpr int64_t kOpcodeCount = 92;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -663,6 +670,27 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
     }
     return mx::stack(in, emily::checked_int(scalar_at(iattrs, 0, "stack"), "axis"),
                      s);
+  // --- FFT family (shares the eager fft.cpp entry points) ---
+  case Opcode::Fftn:
+    return mx::fft::fftn(arg1(in, "fftn"),
+                         emily::to_mlx_shape(attr_at(iattrs, 0, "fftn")),
+                         emily::to_int_vec(attr_at(iattrs, 1, "fftn")),
+                         mx::fft::FFTNorm::Backward, s);
+  case Opcode::Ifftn:
+    return mx::fft::ifftn(arg1(in, "ifftn"),
+                          emily::to_mlx_shape(attr_at(iattrs, 0, "ifftn")),
+                          emily::to_int_vec(attr_at(iattrs, 1, "ifftn")),
+                          mx::fft::FFTNorm::Backward, s);
+  case Opcode::Rfftn:
+    return mx::fft::rfftn(arg1(in, "rfftn"),
+                          emily::to_mlx_shape(attr_at(iattrs, 0, "rfftn")),
+                          emily::to_int_vec(attr_at(iattrs, 1, "rfftn")),
+                          mx::fft::FFTNorm::Backward, s);
+  case Opcode::Irfftn:
+    return mx::fft::irfftn(arg1(in, "irfftn"),
+                           emily::to_mlx_shape(attr_at(iattrs, 0, "irfftn")),
+                           emily::to_int_vec(attr_at(iattrs, 1, "irfftn")),
+                           mx::fft::FFTNorm::Backward, s);
   }
   throw std::invalid_argument("unknown opcode " +
                               std::to_string(static_cast<int64_t>(op)));

lib/emily/ir.ex

@@ -145,7 +145,15 @@ defmodule Emily.IR do
     # window select-and-scatter (pooling backward). operands
     # [input, source, init]; iattrs [[window],[strides],[pad_lo],[pad_hi]].
     window_scatter_max: 86,
-    window_scatter_min: 87
+    window_scatter_min: 87,
+    # FFT family — n-D transforms. operands [input]; iattrs
+    # [[sizes...],[axes...]]. `fft`/`ifft` (1-D, last axis) and the
+    # `fft2`/`ifft2`/`rfft`/`irfft` blocks all route here. Unnormalized
+    # (`FFTNorm::Backward`) is baked C++-side, matching Nx / the eager NIFs.
+    fftn: 88,
+    ifftn: 89,
+    rfftn: 90,
+    irfftn: 91
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -498,6 +506,19 @@ defmodule Emily.IR do
     coerce(r, t.type, state)
   end
 
+  # 1-D FFT / inverse FFT (Nx.fft / Nx.ifft). Mirrors Emily.Backend.{fft,
+  # ifft}/3: route through the n-D MLX kernel restricted to one axis. The
+  # eager path uses the trailing axis and ignores `opts[:axis]`, so we do
+  # too — keeping native bit-identical to the evaluator. Output is complex
+  # (`Nx.Type.to_complex/1`); the trailing coerce matches the backend `wrap`.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: op, args: [a, opts]}} = t, state)
+       when op in [:fft, :ifft] do
+    {ra, state} = lower_node(a, state)
+    axis = tuple_size(a.shape) - 1
+    opcode = if op == :fft, do: :fftn, else: :ifftn
+    emit_coerced(state, opcode, [ra], [[opts[:length]], [axis]], t.type)
+  end
+
   # Non-batched dot -> tensordot over the contraction axes.
   defp lower_op(%T{data: %Nx.Defn.Expr{op: :dot, args: [a, ca, [], b, cb, []]}} = t, state) do
     {ra, state} = lower_node(a, state)
@@ -1012,10 +1033,35 @@ defmodule Emily.IR do
     end
   end
 
+  # FFT family blocks (Nx.fft2 / ifft2 / rfft / irfft). Each mirrors the
+  # matching Emily.Backend.native_* wrapper: route through the n-D MLX
+  # fft/ifft/rfft/irfft kernel with the block's sizes + axes, then coerce to
+  # out.type (complex for the forward transforms, real for irfft). The
+  # block's `eps` is unused (MLX needs none), as in the eager path.
+  defp lower_block(%Nx.Block.FFT2{lengths: lengths, axes: axes}, [t], _expr, out, state) do
+    {rt, state} = lower_node(t, state)
+    emit_coerced(state, :fftn, [rt], [lengths, axes], out.type)
+  end
+
+  defp lower_block(%Nx.Block.IFFT2{lengths: lengths, axes: axes}, [t], _expr, out, state) do
+    {rt, state} = lower_node(t, state)
+    emit_coerced(state, :ifftn, [rt], [lengths, axes], out.type)
+  end
+
+  defp lower_block(%Nx.Block.RFFT{length: length, axis: axis}, [t], _expr, out, state) do
+    {rt, state} = lower_node(t, state)
+    emit_coerced(state, :rfftn, [rt], [[length], [axis]], out.type)
+  end
+
+  defp lower_block(%Nx.Block.IRFFT{length: length, axis: axis}, [t], _expr, out, state) do
+    {rt, state} = lower_node(t, state)
+    emit_coerced(state, :irfftn, [rt], [[length], [axis]], out.type)
+  end
+
   # Any other block struct raises. Lowering the block's composed
   # expansion would silently diverge from the Evaluator whenever
   # Emily.Backend.block/4 dispatches that struct through a fused / native
-  # kernel (e.g. SDPAWithSinks, the Nx.Block.LinAlg.* / FFT families) — a
+  # kernel (e.g. SDPAWithSinks, the Nx.Block.LinAlg.* families) — a
   # worse failure than a clear "unsupported".
   # Additional fused blocks are added alongside their opcode.
   defp lower_block(struct, _in_args, _expr, _t, _state) do

test/emily/compiler_equivalence_test.exs

@@ -476,6 +476,45 @@ defmodule Emily.CompilerEquivalenceTest do
     end
   end
 
+  describe "fft family (signal transforms)" do
+    test "1-D fft / ifft on the trailing axis match the evaluator" do
+      x = et([1.0, 2.0, 3.0, 4.0])
+      assert_equiv(fn t -> Nx.fft(t) end, [x])
+      assert_equiv(fn t -> Nx.ifft(t) end, [x])
+    end
+
+    test "fft with explicit length (zero-pad / truncate) matches" do
+      x = et([1.0, 2.0, 3.0, 4.0, 5.0])
+      assert_equiv(fn t -> Nx.fft(t, length: 8) end, [x])
+      assert_equiv(fn t -> Nx.fft(t, length: 4) end, [x])
+    end
+
+    test "batched 1-D fft transforms the last axis only" do
+      x = et([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]])
+      assert_equiv(fn t -> Nx.fft(t) end, [x])
+    end
+
+    test "ifft(fft(x)) round-trips through a complex intermediate" do
+      x = et([1.0, -2.0, 3.0, -4.0])
+      out = assert_equiv(fn t -> Nx.ifft(Nx.fft(t)) end, [x])
+      assert out.type == {:c, 64}
+    end
+
+    test "2-D fft2 / ifft2 (Nx.Block) match" do
+      x = et([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]])
+      assert_equiv(fn t -> Nx.fft2(t) end, [x])
+      assert_equiv(fn t -> Nx.ifft2(t) end, [x])
+    end
+
+    test "rfft (real -> half spectrum) and irfft (back to real) match" do
+      x = et([1.0, 2.0, 3.0, 4.0])
+      out = assert_equiv(fn t -> Nx.rfft(t) end, [x])
+      assert out.type == {:c, 64}
+      # irfft takes the complex half-spectrum back to a real signal.
+      assert_equiv(fn t -> Nx.irfft(Nx.rfft(t)) end, [x])
+    end
+  end
+
   describe "dynamic put_slice (KV-cache write)" do
     test "put_slice at a runtime offset matches the Evaluator" do
       # {batch, n_kv_heads, max_len, head_dim} KV buffer; write one token.