Merge pull request #177 from ausimian/feat/expr-compiler-take-along-axis

Lower take_along_axis natively + opcode lockstep guard

Author: ausimian

RELEASE.md

@@ -50,6 +50,14 @@
   selection on the native path. Remaining gaps (`gather`/scatter,
   pooling/`window_*`, cumulative) continue to work via the graceful fallback.
 
+- **`take_along_axis` lowers natively** — `Nx.take_along_axis` (the
+  `Nx.Block.TakeAlongAxis` block) now compiles under the native single-NIF
+  path, mirroring `Emily.Backend.native_take_along_axis/4` (cast indices to
+  s32, then `mlx::core::take_along_axis`) bit-for-bit. This was the last op
+  forcing a fallback in `Bumblebee.Text.question_answering`'s answer-span
+  gather, so a DistilBERT question-answering `Nx.Serving` forward now runs
+  fully native — and fused — under `native_fallback: :raise`.
+
 - **`Bumblebee.Text.generation` compiles fully native — greedy and sampling.**
   The headline result: an end-to-end Bumblebee generation (the transformer
   forward, the `defn while` decode loop, dynamic KV-cache writes, `cumsum`

c_src/emily/opcodes.hpp

@@ -153,9 +153,12 @@ enum class Opcode : int64_t {
   Gather = 79,
   // Stack tensors along a new axis. operands [t0, t1, ...]; iattrs [[axis]]
   Stack = 80,
+  // Gather along one axis with a same-rank s32 index tensor.
+  // operands [input, indices]; iattrs [[axis]]
+  TakeAlongAxis = 81,
 };
 
-inline constexpr int64_t kOpcodeCount = 81;
+inline constexpr int64_t kOpcodeCount = 82;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -503,6 +506,15 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
     int axis = emily::checked_int(scalar_at(iattrs, 0, "take"), "axis");
     return mx::take(in[0], in[1], axis, s);
   }
+  case Opcode::TakeAlongAxis: {
+    if (in.size() != 2) {
+      throw std::invalid_argument("take_along_axis expects 2 operands, got " +
+                                  std::to_string(in.size()));
+    }
+    int axis =
+        emily::checked_int(scalar_at(iattrs, 0, "take_along_axis"), "axis");
+    return mx::take_along_axis(in[0], in[1], axis, s);
+  }
   case Opcode::Concatenate: {
     if (in.empty()) {
       throw std::invalid_argument("concatenate expects >= 1 operand");

lib/emily/ir.ex

@@ -132,7 +132,10 @@ defmodule Emily.IR do
     cummin: 78,
     # multi-axis gather: operands [input, idx0, ...]; iattrs [[axes],[slice_sizes]]
     gather: 79,
-    stack: 80
+    stack: 80,
+    # take_along_axis: gather along one axis with a same-rank s32 index
+    # tensor. operands [input, indices]; iattrs [[axis]].
+    take_along_axis: 81
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -163,6 +166,14 @@ defmodule Emily.IR do
           outputs: [ref()]
         }
 
+  @doc """
+  The opcode name -> wire-value map. Exposed for the opcode-parity test,
+  which checks these stay in lockstep with the `Opcode` enum and
+  `kOpcodeCount` in `c_src/emily/opcodes.hpp`.
+  """
+  @spec opcodes() :: %{atom() => non_neg_integer()}
+  def opcodes, do: @opcodes
+
   @doc "Numeric wire value for an opcode name."
   @spec opcode(atom()) :: non_neg_integer()
   def opcode(name) when is_map_key(@opcodes, name), do: Map.fetch!(@opcodes, name)
@@ -915,6 +926,16 @@ defmodule Emily.IR do
     emit_coerced(state, :take, [ri, rx], [[axis]], t.type)
   end
 
+  # Nx.take_along_axis (Nx.Block.TakeAlongAxis). Mirrors
+  # Emily.Backend.native_take_along_axis/4: cast indices to s32, then
+  # mx::take_along_axis along `axis`.
+  defp lower_block(%Nx.Block.TakeAlongAxis{axis: axis}, [input, indices], _expr, t, state) do
+    {ri, state} = lower_node(input, state)
+    {rx, state} = lower_node(indices, state)
+    {rx, state} = emit(state, :astype, [rx], [[dtype_code({:s, 32})]])
+    emit_coerced(state, :take_along_axis, [ri, rx], [[axis]], t.type)
+  end
+
   # Cumulative families. Like Emily.Backend.block/4, the last-axis case uses
   # the native MLX `cumsum`/`cumprod`/`cummax`/`cummin` kernel; interior axes
   # (which MLX can't always factor) fall back to the block's composed

test/emily/compiler_equivalence_test.exs

@@ -368,6 +368,26 @@ defmodule Emily.CompilerEquivalenceTest do
     end
   end
 
+  describe "take_along_axis" do
+    test "along the last axis matches the Evaluator" do
+      x = et([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]])
+      idx = Nx.tensor([[3, 0], [1, 2], [0, 3]], type: :s64, backend: Emily.Backend)
+      assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 1) end, [x, idx])
+    end
+
+    test "along axis 0 matches" do
+      x = et([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+      idx = Nx.tensor([[0, 1, 0], [1, 0, 1], [0, 0, 1]], type: :s64, backend: Emily.Backend)
+      assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 0) end, [x, idx])
+    end
+
+    test "3-D gather along the last axis (transformer-shaped) matches" do
+      x = Nx.iota({1, 2, 4}, type: :f32, backend: Emily.Backend) |> Nx.divide(8.0)
+      idx = Nx.tensor([[[3, 1], [0, 2]]], type: :s32, backend: Emily.Backend)
+      assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 2) end, [x, idx])
+    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.

test/emily/conformance/distilbert_test.exs

@@ -33,7 +33,8 @@ defmodule Emily.Conformance.DistilbertTest do
 
   use ExUnit.Case, async: false
 
-  import Emily.ConformanceHelper, only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2]
+  import Emily.ConformanceHelper,
+    only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2, mode_test: 3]
 
   alias Emily.Bumblebee.FastKernels
 
@@ -228,15 +229,22 @@ defmodule Emily.Conformance.DistilbertTest do
     # (vocab 30522) with a tiny-random model (1124-row embedding)
     # feeds out-of-range token ids into gather and relies on backend
     # OOB behaviour, which is how we originally hit a :nan score.
-    @tag :distilbert_full
-    test "batched_run drives DistilBERT-QA through Nx.Serving" do
+    # `tag: :distilbert_full, lane_tags: false` keeps all three lanes gated
+    # behind `:distilbert_full` (not the lightweight `:native`), so they run
+    # under `--only distilbert_full` and never bloat `--only native`. The
+    # forward is driven through `Nx.Serving`'s `:defn_options`, which is
+    # where the compiler lanes plug in.
+    mode_test "batched_run drives DistilBERT-QA through Nx.Serving",
+      tag: :distilbert_full,
+      lane_tags: false do
       {:ok, model_info} =
         Bumblebee.load_model({:hf, "distilbert-base-uncased-distilled-squad"})
 
       {:ok, tokenizer} =
         Bumblebee.load_tokenizer({:hf, "distilbert-base-uncased-distilled-squad"})
 
-      serving = Bumblebee.Text.question_answering(model_info, tokenizer)
+      serving =
+        Bumblebee.Text.question_answering(model_info, tokenizer, defn_options: predict_opts)
 
       start_supervised!({Nx.Serving, serving: serving, name: __MODULE__.Serving})
 

test/emily/opcode_parity_test.exs

@@ -0,0 +1,67 @@
+defmodule Emily.OpcodeParityTest do
+  @moduledoc """
+  Guards the one hand-maintained lockstep in the native compiler: the
+  opcode wire values live in **two** places that must agree —
+  `Emily.IR`'s `@opcodes` map (Elixir) and the `Opcode` enum +
+  `kOpcodeCount` in `c_src/emily/opcodes.hpp` (C++). Nothing in the build
+  enforces it: a mismatch compiles fine and only misbehaves at runtime
+  (an instruction dispatches to the wrong MLX op, or `valid_opcode`
+  rejects a real one).
+
+  The check is value-based rather than name-based — both sides must be a
+  contiguous `0..N-1` with N == `kOpcodeCount`. That catches every
+  realistic drift (an opcode added to one side only, a forgotten
+  `kOpcodeCount` bump, a duplicate or gapped number) without depending on
+  the Elixir snake_case names matching the C++ PascalCase ones (they
+  don't always: `negate`/`Negative`, `fast_rms_norm`/`FastRMSNorm`). A
+  name/value *permutation* that keeps both contiguous would slip past
+  here, but that produces wrong results and is caught by
+  `Emily.CompilerEquivalenceTest`.
+  """
+  use ExUnit.Case, async: true
+
+  alias Emily.IR
+
+  @header Path.expand("../../c_src/emily/opcodes.hpp", __DIR__)
+
+  # Parse `kOpcodeCount` and the `Opcode` enum's explicit `Name = N,`
+  # values out of the C++ header.
+  defp parse_header do
+    src = File.read!(@header)
+
+    [_, count] = Regex.run(~r/kOpcodeCount\s*=\s*(\d+)/, src)
+    count = String.to_integer(count)
+
+    [_, body] = Regex.run(~r/enum class Opcode\s*:\s*int64_t\s*\{(.*?)\};/s, src)
+
+    values =
+      ~r/^\s*[A-Za-z]\w*\s*=\s*(\d+)\s*,/m
+      |> Regex.scan(body)
+      |> Enum.map(fn [_, v] -> String.to_integer(v) end)
+
+    %{count: count, enum_values: values}
+  end
+
+  test "C++ Opcode enum is contiguous 0..N-1 and matches kOpcodeCount" do
+    %{count: count, enum_values: values} = parse_header()
+
+    assert length(values) == count,
+           "opcodes.hpp has #{length(values)} enum entries but kOpcodeCount is #{count} — " <>
+             "bump kOpcodeCount when adding an opcode"
+
+    assert Enum.sort(values) == Enum.to_list(0..(count - 1)),
+           "Opcode enum values are not a unique, gap-free 0..#{count - 1}"
+  end
+
+  test "Emily.IR @opcodes stays in lockstep with the C++ enum" do
+    %{count: count} = parse_header()
+    opcodes = IR.opcodes()
+
+    assert map_size(opcodes) == count,
+           "Emily.IR has #{map_size(opcodes)} opcodes but c_src/emily/opcodes.hpp kOpcodeCount " <>
+             "is #{count} — the two must be updated together"
+
+    assert opcodes |> Map.values() |> Enum.sort() == Enum.to_list(0..(count - 1)),
+           "Emily.IR @opcodes values are not a unique, gap-free 0..#{count - 1}"
+  end
+end

test/support/conformance_helper.ex

@@ -96,21 +96,27 @@ defmodule Emily.ConformanceHelper do
       suite's own `:*_full` moduletag; otherwise `--only native` would
       start pulling full-size checkpoints. `--only vit_full` then runs all
       three lanes of that suite.
+
+    * `:tag` — an extra tag stamped on *every* lane. Used by the
+      `Nx.Serving` test, which lives in a `:conformance`-tagged module but
+      must stay gated behind `:distilbert_full` like its eval lane:
+      `tag: :distilbert_full, lane_tags: false`.
   """
   defmacro mode_test(name, opts \\ [], do: body) do
     tag_lanes? = Keyword.get(opts, :lane_tags, true)
+    extra_tag = Keyword.get(opts, :tag)
 
     lanes = [
-      lane(false, name, "", [], body),
+      lane([extra_tag], name, "", [], body),
       lane(
-        tag_lanes? && :native,
+        [extra_tag, tag_lanes? && :native],
         name,
         " [native]",
         [compiler: Emily.Compiler, native: true, native_fallback: :raise],
         body
       ),
       lane(
-        tag_lanes? && :native_compiled,
+        [extra_tag, tag_lanes? && :native_compiled],
         name,
         " [native_compiled]",
         [compiler: Emily.Compiler, native: true, native_fallback: :raise, native_compiled: true],
@@ -124,27 +130,25 @@ defmodule Emily.ConformanceHelper do
   end
 
   # Build one `mode_test` lane: a `test` that binds `predict_opts` for the
-  # body, optionally preceded by `@tag tag`. `tag` is `false` to emit no
-  # lane tag (the `*_full` suites rely on their own `:*_full` moduletag).
-  defp lane(tag, name, suffix, predict_opts, body) do
+  # body, preceded by one `@tag` per entry in `tags` (nil/false entries are
+  # dropped). The `*_full` suites pass `lane_tags: false` to drop the
+  # `:native` / `:native_compiled` tags and rely on their own `:*_full`
+  # moduletag (or an explicit `:tag`) instead.
+  defp lane(tags, name, suffix, predict_opts, body) do
+    tags = Enum.reject(tags, &(&1 in [nil, false]))
+
     name_ast =
       if suffix == "", do: name, else: quote(do: unquote(name) <> unquote(suffix))
 
-    test =
-      quote do
-        test unquote(name_ast) do
-          var!(predict_opts) = unquote(predict_opts)
-          unquote(body)
-        end
-      end
+    tag_attrs = for t <- tags, do: quote(do: @tag(unquote(t)))
+
+    quote do
+      (unquote_splicing(tag_attrs))
 
-    if tag do
-      quote do
-        @tag unquote(tag)
-        unquote(test)
+      test unquote(name_ast) do
+        var!(predict_opts) = unquote(predict_opts)
+        unquote(body)
       end
-    else
-      test
     end
   end