diff --git a/RELEASE.md b/RELEASE.md index 6973ede..74a66c5 100644 --- a/RELEASE.md +++ b/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 diff --git a/c_src/emily/opcodes.hpp b/c_src/emily/opcodes.hpp index 53c7a38..0676031 100644 --- a/c_src/emily/opcodes.hpp +++ b/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 &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(op))); diff --git a/lib/emily/ir.ex b/lib/emily/ir.ex index 3f3f331..18f12d4 100644 --- a/lib/emily/ir.ex +++ b/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 diff --git a/test/emily/compiler_equivalence_test.exs b/test/emily/compiler_equivalence_test.exs index 987bdbb..a85621a 100644 --- a/test/emily/compiler_equivalence_test.exs +++ b/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.