diff --git a/RELEASE.md b/RELEASE.md index 74a66c5..ce5c2e8 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -67,6 +67,14 @@ `speech_to_text` serving — the log-mel featurizer's STFT — so that path now compiles fully native too. +- **`indexed_put` / `indexed_add` (scatter) lower natively** — both now + compile under the native single-NIF path for MLX-scatter-compatible index + layouts (the same layout the native gather already requires), mirroring + `Emily.Backend`'s scatter — split the index tensor into per-axis s32 + arrays, reshape updates into MLX's layout, then `mx::scatter` (overwrite) + / `mx::scatter_add` (accumulate) bit-for-bit. Layouts MLX can't scatter + still route through the evaluator under `native_fallback: :eval`. + - **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 @@ -177,6 +185,18 @@ ### Fixed +- **A tuple-returning `cond` hard-crashed the native compiler instead of + lowering.** A `cond`/`if` whose branches return a tuple (multi-output) hit a + `FunctionClauseError` in the lowerer — and because that isn't an + `ArgumentError`, it escaped the graceful-fallback rescue and faulted rather + than degrading to the evaluator. It now lowers to one `where`-chain per leaf + (same wholesale-select semantics as a single-output `cond`), projected by + `:elem`; a nested/non-tensor container raises cleanly (graceful fallback). + Surfaced by a Whisper `speech_to_text` serving — with this, plus the native + `fft` and `indexed_put` lowering above, the **full Whisper serving + (featurizer STFT + encoder/decoder + autoregressive decode loop) compiles + fully native end-to-end**, gated by `whisper_full_test.exs`. + - **Dilated window reductions (`window_dilations > 1`) returned wrong values.** `window_sum`/`window_max`/`window_min`/`window_product` with a dilated kernel silently produced garbage for windows past the first stride positions, on both diff --git a/c_src/emily/opcodes.hpp b/c_src/emily/opcodes.hpp index 0676031..cf13f93 100644 --- a/c_src/emily/opcodes.hpp +++ b/c_src/emily/opcodes.hpp @@ -175,9 +175,13 @@ enum class Opcode : int64_t { Ifftn = 89, // complex -> complex (inverse) Rfftn = 90, // real -> complex (half spectrum) Irfftn = 91, // complex half-spectrum -> real + // Scatter (Nx.indexed_put / indexed_add). operands [target, updates, + // idx0, ...] (one s32 index array per scattered axis); iattrs [[axes...]]. + Scatter = 92, // overwrite (last write wins on duplicates) + ScatterAdd = 93, // accumulate }; -inline constexpr int64_t kOpcodeCount = 92; +inline constexpr int64_t kOpcodeCount = 94; // Quant mode code (Emily.IR @quant_modes) -> MLX mode string. inline std::string qmode_from_code(int64_t code) { @@ -691,6 +695,21 @@ inline mx::array dispatch_op(Opcode op, const std::vector &in, emily::to_mlx_shape(attr_at(iattrs, 0, "irfftn")), emily::to_int_vec(attr_at(iattrs, 1, "irfftn")), mx::fft::FFTNorm::Backward, s); + // --- Scatter (shares the eager index.cpp entry points) --- + case Opcode::Scatter: + case Opcode::ScatterAdd: { + if (in.size() < 3) { + throw std::invalid_argument( + "scatter expects >= 3 operands (target, updates, >=1 index), got " + + std::to_string(in.size())); + } + // operands [target, updates, idx0, ...]; the index arrays follow updates. + std::vector indices(in.begin() + 2, in.end()); + auto axes = emily::to_int_vec(attr0(iattrs, "scatter")); + return op == Opcode::Scatter + ? mx::scatter(in[0], indices, in[1], axes, s) + : mx::scatter_add(in[0], indices, in[1], axes, 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 18f12d4..0cd7377 100644 --- a/lib/emily/ir.ex +++ b/lib/emily/ir.ex @@ -153,7 +153,12 @@ defmodule Emily.IR do fftn: 88, ifftn: 89, rfftn: 90, - irfftn: 91 + irfftn: 91, + # Scatter (Nx.indexed_put / indexed_add). operands [target, updates, + # idx0, ...] (one s32 index array per scattered axis); iattrs [[axes...]]. + # scatter overwrites (last-write on duplicates); scatter_add accumulates. + scatter: 92, + scatter_add: 93 } # Quant mode string -> code; decoded by qmode_from_code in @@ -661,6 +666,39 @@ defmodule Emily.IR do coerce(r, t.type, state) end + # indexed_put / indexed_add (Nx scatter). Mirrors Emily.Backend's + # apply_scatter: split the {..., R} index tensor into R per-axis s32 index + # arrays, reshape `updates` into MLX's scatter layout, then mx::scatter + # (overwrite) / mx::scatter_add (accumulate). Only MLX-scatter-compatible + # index layouts lower; others raise (no fallback) — the Evaluator handles + # them via_binary under `native_fallback: :eval`, matching native gather. + # Operands [target, updates, idx0, ...]; iattrs [[axes...]]. + defp lower_op( + %T{data: %Nx.Defn.Expr{op: op, args: [target, indices, updates, opts]}} = t, + state + ) + when op in [:indexed_put, :indexed_add] do + axes = opts[:axes] || Enum.to_list(0..(tuple_size(target.shape) - 1)//1) + indices_shape = Tuple.to_list(indices.shape) + + unless scatter_gather_compatible?(indices_shape, axes) do + raise ArgumentError, + "Emily Expr compiler: #{op} index layout #{inspect(indices_shape)} for axes " <> + "#{inspect(axes)} is not MLX-scatter-compatible (no fallback)." + end + + {rt, state} = lower_node(target, state) + {rx, state} = lower_node(indices, state) + {idx_refs, state} = split_indices_for_gather(rx, indices_shape, length(axes), state) + + {ru, state} = lower_node(updates, state) + updates_shape = updates_shape_for_scatter(indices_shape, target.shape, axes) + {ru, state} = emit(state, :reshape, [ru], [updates_shape]) + + opcode = if op == :indexed_put, do: :scatter, else: :scatter_add + emit_coerced(state, opcode, [rt, ru | idx_refs], [axes], t.type) + end + # put_slice(src, start_indices, slice): write `slice` into `src` at # `start_indices`. Mirrors Emily.Backend.put_slice/4 (cast src + update # to out.type), but supports RUNTIME (tensor) start indices — the decode @@ -832,7 +870,7 @@ defmodule Emily.IR do # gather/index there clamps rather than faults, so the discarded value # never changes the result; a hard-faulting op on a not-taken path would # diverge from the Evaluator's lazy single-branch eval. - defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, last]}} = t, state) do + defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, %T{} = last]}} = t, state) do {last_ref, state} = lower_node(last, state) {result, state} = @@ -846,6 +884,52 @@ defmodule Emily.IR do coerce(result, t.type, state) end + # Multi-output cond: each branch returns a tuple of tensors. Lower to one + # `where`-chain per leaf position — identical wholesale-select semantics to + # the single-output case above (the predicate is a whole-tensor scalar bool; + # every branch is still computed). Returns a `{:multi_refs, [...]}` handle + # that `:elem` projects; sibling `:elem`s share it via lower_node's memo. + # A nested / non-tensor container raises (no fallback path for it yet). + defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, last]}}, state) + when is_tuple(last) do + last_leaves = Tuple.to_list(last) + + unless tensors?(last_leaves) and + Enum.all?(clauses, fn {_p, b} -> is_tuple(b) and tensors?(Tuple.to_list(b)) end) do + raise ArgumentError, + "Emily Expr compiler: cond over a nested / non-tensor container is not " <> + "lowered yet (only a flat tuple of tensors)." + end + + # Lower each predicate (cast to pred) + its branch's leaf refs once, so the + # per-leaf where-chains share them. + {clauses, state} = + Enum.map_reduce(clauses, state, fn {pred, body}, st -> + {pred_ref, st} = lower_node(pred, st) + {pred_ref, st} = emit(st, :astype, [pred_ref], [[dtype_code({:pred, 1})]]) + {body_refs, st} = Enum.map_reduce(Tuple.to_list(body), st, &lower_node/2) + {{pred_ref, body_refs}, st} + end) + + rev = Enum.reverse(clauses) + + {refs, state} = + last_leaves + |> Enum.with_index() + |> Enum.map_reduce(state, fn {leaf, j}, st -> + {last_ref, st} = lower_node(leaf, st) + + {result, st} = + Enum.reduce(rev, {last_ref, st}, fn {pred_ref, body_refs}, {else_ref, st2} -> + emit(st2, :where, [pred_ref, Enum.at(body_refs, j), else_ref]) + end) + + coerce(result, leaf.type, st) + end) + + {{:multi_refs, refs}, state} + end + # attach_token: sequences a token (hooks) before `expr`. With no active # hook the token is a no-op, so pass through to the inner expr. Hooks # would need a callback into Elixir mid-graph (program-split) — deferred. @@ -917,11 +1001,14 @@ defmodule Emily.IR do {{:multi, base, _arity}, state} -> {{:instr, base + i}, state} + {{:multi_refs, refs}, state} -> + {Enum.at(refs, i), state} + {_handle, _state} -> raise ArgumentError, "Emily Expr compiler: :elem projects a tuple-producing op it can't " <> - "lower yet (only `while` produces projectable tuples today; other " <> - "multi-output ops are unsupported)." + "lower yet (only `while` and tuple `cond` produce projectable " <> + "tuples today; other multi-output ops are unsupported)." end end @@ -1119,6 +1206,8 @@ defmodule Emily.IR do defp bool_int(true), do: 1 defp bool_int(false), do: 0 + defp tensors?(list), do: Enum.all?(list, &match?(%T{}, &1)) + defp float_like?({kind, _}) when kind in [:f, :bf, :c], do: true defp float_like?(_), do: false @@ -1285,4 +1374,19 @@ defmodule Emily.IR do rank = tuple_size(input_shape) for i <- 0..(rank - 1)//1, do: if(i in axes_set, do: 1, else: elem(input_shape, i)) end + + # Rewrap Nx's updates shape {batch ++ non_indexed_dims} into MLX's scatter + # layout {batch ++ per_axis_slot}, where per_axis_slot has length + # rank(target) with 1 on indexed axes and target_shape[i] elsewhere. Mirrors + # Emily.Backend.updates_shape_for_scatter/3. + defp updates_shape_for_scatter(indices_shape, target_shape, axes) do + batch = Enum.take(indices_shape, length(indices_shape) - 1) + axes_set = MapSet.new(axes) + rank = tuple_size(target_shape) + + trailing = + for i <- 0..(rank - 1)//1, do: if(i in axes_set, do: 1, else: elem(target_shape, i)) + + batch ++ trailing + end end diff --git a/test/emily/compiler_control_flow_test.exs b/test/emily/compiler_control_flow_test.exs index fbcc366..36e4c3d 100644 --- a/test/emily/compiler_control_flow_test.exs +++ b/test/emily/compiler_control_flow_test.exs @@ -45,6 +45,34 @@ defmodule Emily.CompilerControlFlowTest do acc end + # Tuple-returning cond (multi-output): both branches produce a {vec, vec} + # tuple; `:elem` projects each leaf out of the lowered per-leaf where-chains. + defn tuple_if_fn(x) do + {a, b} = + if Nx.greater(Nx.sum(x), 0) do + {Nx.multiply(x, 2), Nx.add(x, 1)} + else + {Nx.negate(x), Nx.subtract(x, 1)} + end + + Nx.add(a, b) + end + + # Multi-clause tuple cond with leaves of DIFFERENT shapes (vector + scalar), + # so each leaf's where-chain is built and coerced independently. + defn tuple_cond3_fn(x) do + s = Nx.sum(x) + + {vec, scalar} = + cond do + Nx.greater(s, 10) -> {Nx.multiply(x, 10), Nx.reduce_max(x)} + Nx.greater(s, 0) -> {Nx.multiply(x, 2), Nx.product(x)} + true -> {Nx.negate(x), s} + end + + Nx.add(vec, scalar) + end + defp equiv(fun, x) do native = Nx.Defn.jit(fun, @native).(x) eval = Nx.Defn.jit(fun, @eval).(x) @@ -71,6 +99,18 @@ defmodule Emily.CompilerControlFlowTest do equiv(&nested_if_fn/1, Nx.tensor(data, backend: Emily.Backend)) end end + + test "tuple-returning if (multi-output cond) projects each leaf" do + for data <- [[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]] do + equiv(&tuple_if_fn/1, Nx.tensor(data, backend: Emily.Backend)) + end + end + + test "multi-clause tuple cond with mixed-shape leaves matches the evaluator" do + for data <- [[5.0, 4.0, 3.0], [1.0, 1.0, 1.0], [-2.0, -2.0, -2.0]] do + equiv(&tuple_cond3_fn/1, Nx.tensor(data, backend: Emily.Backend)) + end + end end describe "defn while" do diff --git a/test/emily/compiler_equivalence_test.exs b/test/emily/compiler_equivalence_test.exs index a85621a..0b70edf 100644 --- a/test/emily/compiler_equivalence_test.exs +++ b/test/emily/compiler_equivalence_test.exs @@ -395,6 +395,32 @@ defmodule Emily.CompilerEquivalenceTest do end end + describe "scatter (indexed_put / indexed_add)" do + test "indexed_put / indexed_add into a 2-D target match the evaluator" do + target = et([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + # Two distinct cells in the {2,3} grid -> order-independent. + idx = Nx.tensor([[0, 0], [1, 2]], type: :s64, backend: Emily.Backend) + upd = et([10.0, 20.0]) + + assert_equiv(fn t, i, u -> Nx.indexed_put(t, i, u) end, [target, idx, upd]) + assert_equiv(fn t, i, u -> Nx.indexed_add(t, i, u) end, [target, idx, upd]) + end + + test "indexed_add accumulates duplicate indices (same MLX kernel both paths)" do + target = et([0.0, 0.0, 0.0, 0.0]) + idx = Nx.tensor([[1], [1], [3]], type: :s64, backend: Emily.Backend) + upd = et([5.0, 7.0, 2.0]) + assert_equiv(fn t, i, u -> Nx.indexed_add(t, i, u) end, [target, idx, upd]) + end + + test "indexed_put on a partial axis set (axes: [0], whole-row writes) matches" do + target = et([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + idx = Nx.tensor([[0], [2]], type: :s64, backend: Emily.Backend) + upd = et([[10.0, 11.0, 12.0], [13.0, 14.0, 15.0]]) + assert_equiv(fn t, i, u -> Nx.indexed_put(t, i, u, axes: [0]) end, [target, idx, upd]) + end + end + describe "window reductions (pooling forward)" do test "2x2 maxpool / sumpool / minpool (CNN-shaped) match the Evaluator" do # {batch, channels, h, w}; pool only the spatial axes, stride 2. diff --git a/test/emily/conformance/whisper_full_test.exs b/test/emily/conformance/whisper_full_test.exs index 30c3ce3..9a24825 100644 --- a/test/emily/conformance/whisper_full_test.exs +++ b/test/emily/conformance/whisper_full_test.exs @@ -84,6 +84,44 @@ defmodule Emily.Conformance.WhisperFullTest do ) end + test "speech_to_text serving lowers fully native — featurizer + decode loop, no fallback" do + repo = {:hf, "openai/whisper-tiny"} + {:ok, whisper} = Bumblebee.load_model(repo) + {:ok, featurizer} = Bumblebee.load_featurizer(repo) + {:ok, tokenizer} = Bumblebee.load_tokenizer(repo) + {:ok, generation_config} = Bumblebee.load_generation_config(repo) + + # The gate is "does the whole graph lower", not transcription quality — + # cap the decode loop so it stays fast. + generation_config = Bumblebee.configure(generation_config, max_new_tokens: 4) + + # `native_fallback: :raise` makes this a no-fallback gate over the ENTIRE + # serving graph — the mel featurizer's STFT (`fft`), the encoder/decoder + # forward, and the autoregressive decode loop (the multi-output `cond` in + # the encoder attention, `indexed_put` cache writes, dynamic slices). Any + # op the Expr compiler can't lower raises here rather than silently + # degrading. The `mode_test` forward pass above never reaches these: it + # feeds pre-computed mel features through a single `Axon.predict`, so it + # exercises neither the featurizer nor the generation loop. + serving = + Bumblebee.Audio.speech_to_text_whisper(whisper, featurizer, tokenizer, generation_config, + defn_options: [compiler: Emily.Compiler, native: true, native_fallback: :raise] + ) + + # ~1 s of deterministic synthetic audio. The featurizer pads it to + # Whisper's 30 s window, so the encoder still runs the full 1500-position + # path (the shape that surfaced the multi-output cond). + audio = Nx.sin(Nx.iota({16_000}, type: :f32) |> Nx.multiply(0.02)) + + %{chunks: chunks} = Nx.Serving.run(serving, audio) + + # Reaching here is the gate: the full path lowered native with zero + # fallback. The output is only sanity-checked (synthetic audio decodes to + # arbitrary tokens; the transcription itself is not pinned). + assert is_list(chunks) and chunks != [] + assert Enum.all?(chunks, &is_binary(&1.text)) + end + @tag :fast_kernels_full test "Whisper-tiny with fused MLX kernels matches the pinned argmax within widened tolerance" do {:ok, %{model: model, params: params}} =