diff --git a/RELEASE.md b/RELEASE.md index 67e309e..49db2e3 100644 --- a/RELEASE.md +++ b/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` diff --git a/c_src/emily/opcodes.hpp b/c_src/emily/opcodes.hpp index a4e758b..c40cdaa 100644 --- a/c_src/emily/opcodes.hpp +++ b/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 &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"); diff --git a/lib/emily/ir.ex b/lib/emily/ir.ex index 34b5b34..b819bb4 100644 --- a/lib/emily/ir.ex +++ b/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 diff --git a/test/emily/compiler_equivalence_test.exs b/test/emily/compiler_equivalence_test.exs index 45a7467..3b3ccee 100644 --- a/test/emily/compiler_equivalence_test.exs +++ b/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. diff --git a/test/emily/conformance/distilbert_test.exs b/test/emily/conformance/distilbert_test.exs index 4929472..83f050a 100644 --- a/test/emily/conformance/distilbert_test.exs +++ b/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] + import Emily.ConformanceHelper, + only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2, mode_test: 3] alias Emily.Bumblebee.FastKernels @@ -52,7 +53,7 @@ defmodule Emily.Conformance.DistilbertTest do :ok end - test ":base" do + mode_test ":base" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertModel"}) @@ -63,7 +64,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.hidden_state) == {1, 10, 32} @@ -75,7 +76,7 @@ defmodule Emily.Conformance.DistilbertTest do ) end - test ":for_masked_language_modeling" do + mode_test ":for_masked_language_modeling" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertForMaskedLM"}) @@ -86,7 +87,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 10, 1124} @@ -98,7 +99,7 @@ defmodule Emily.Conformance.DistilbertTest do ) end - test ":for_sequence_classification" do + mode_test ":for_sequence_classification" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-DistilBertForSequenceClassification"} @@ -111,14 +112,14 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 2} assert_all_close(outputs.logits, Nx.tensor([[-0.0047, -0.0103]])) end - test ":for_token_classification" do + mode_test ":for_token_classification" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-DistilBertForTokenClassification"} @@ -131,7 +132,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 10, 2} @@ -141,7 +142,7 @@ defmodule Emily.Conformance.DistilbertTest do ) end - test ":for_question_answering" do + mode_test ":for_question_answering" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-DistilBertForQuestionAnswering"} @@ -154,7 +155,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.start_logits) == {1, 10} assert Nx.shape(outputs.end_logits) == {1, 10} @@ -170,7 +171,7 @@ defmodule Emily.Conformance.DistilbertTest do ) end - test ":for_multiple_choice" do + mode_test ":for_multiple_choice" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-DistilBertForMultipleChoice"} @@ -183,7 +184,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 1} @@ -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}) diff --git a/test/emily/conformance/modernbert_test.exs b/test/emily/conformance/modernbert_test.exs index 99f9b40..b6ef5f1 100644 --- a/test/emily/conformance/modernbert_test.exs +++ b/test/emily/conformance/modernbert_test.exs @@ -33,7 +33,7 @@ defmodule Emily.Conformance.ModernBertTest do @moduletag :conformance @moduletag capture_log: true - test "ModernBert :base forward on Emily.Backend" do + mode_test "ModernBert :base forward on Emily.Backend" do spec = Bumblebee.configure(ModernBert, architecture: :base, @@ -47,7 +47,10 @@ defmodule Emily.Conformance.ModernBertTest do ) model = ModernBert.model(spec) - {init_fn, predict_fn} = Axon.build(model) + # Init on the evaluator (params are random-init, mode-irrelevant); + # gate only the forward pass under `predict_opts`. + {init_fn, _} = Axon.build(model) + {_, predict_fn} = Axon.build(model, predict_opts) input_template = %{ "input_ids" => Nx.template({1, 8}, :s64), diff --git a/test/emily/conformance/nomic_embeddings_test.exs b/test/emily/conformance/nomic_embeddings_test.exs index 7e74c80..4003f7e 100644 --- a/test/emily/conformance/nomic_embeddings_test.exs +++ b/test/emily/conformance/nomic_embeddings_test.exs @@ -27,7 +27,7 @@ defmodule Emily.Conformance.NomicEmbeddingsTest do @moduletag :conformance @moduletag capture_log: true - test "NomicBert :base forward runs end-to-end on Emily.Backend" do + mode_test "NomicBert :base forward runs end-to-end on Emily.Backend" do spec = Bumblebee.configure(NomicBert, architecture: :base, @@ -41,7 +41,10 @@ defmodule Emily.Conformance.NomicEmbeddingsTest do ) model = NomicBert.model(spec) - {init_fn, predict_fn} = Axon.build(model) + # Init on the evaluator (params are random-init, mode-irrelevant); + # gate only the forward pass under `predict_opts`. + {init_fn, _} = Axon.build(model) + {_, predict_fn} = Axon.build(model, predict_opts) input_template = %{ "input_ids" => Nx.template({1, 8}, :s64), diff --git a/test/emily/conformance/smollm3_test.exs b/test/emily/conformance/smollm3_test.exs index fdc7002..fca1179 100644 --- a/test/emily/conformance/smollm3_test.exs +++ b/test/emily/conformance/smollm3_test.exs @@ -30,7 +30,7 @@ defmodule Emily.Conformance.SmolLm3Test do @moduletag :conformance @moduletag capture_log: true - test "SmolLm3 :for_causal_language_modeling forward on Emily.Backend" do + mode_test "SmolLm3 :for_causal_language_modeling forward on Emily.Backend" do spec = Bumblebee.configure(SmolLm3, architecture: :for_causal_language_modeling, @@ -44,7 +44,10 @@ defmodule Emily.Conformance.SmolLm3Test do ) model = SmolLm3.model(spec) - {init_fn, predict_fn} = Axon.build(model) + # Init on the evaluator (params are random-init, mode-irrelevant); + # gate only the forward pass under `predict_opts`. + {init_fn, _} = Axon.build(model) + {_, predict_fn} = Axon.build(model, predict_opts) input_template = %{"input_ids" => Nx.template({1, 8}, :s64)} params = init_fn.(input_template, Axon.ModelState.empty()) diff --git a/test/emily/conformance/vit_full_test.exs b/test/emily/conformance/vit_full_test.exs index accb739..eabbb1e 100644 --- a/test/emily/conformance/vit_full_test.exs +++ b/test/emily/conformance/vit_full_test.exs @@ -30,7 +30,8 @@ defmodule Emily.Conformance.VitFullTest do @moduletag capture_log: true @moduletag timeout: 600_000 - test "google/vit-base-patch16-224 forward pass matches pinned logits slice" do + mode_test "google/vit-base-patch16-224 forward pass matches pinned logits slice", + lane_tags: false do {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "google/vit-base-patch16-224"}) @@ -42,7 +43,7 @@ defmodule Emily.Conformance.VitFullTest do "pixel_values" => Nx.broadcast(Nx.tensor(0.5, type: :f32), {1, 224, 224, 3}) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 1000} diff --git a/test/emily/conformance/vit_test.exs b/test/emily/conformance/vit_test.exs index baa39f7..5e38938 100644 --- a/test/emily/conformance/vit_test.exs +++ b/test/emily/conformance/vit_test.exs @@ -31,7 +31,7 @@ defmodule Emily.Conformance.VitTest do @moduletag capture_log: true @moduletag timeout: 120_000 - test ":base" do + mode_test ":base" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-ViTModel"}) @@ -41,7 +41,7 @@ defmodule Emily.Conformance.VitTest do "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.hidden_state) == {1, 226, 32} assert Nx.shape(outputs.pooled_state) == {1, 32} @@ -59,7 +59,7 @@ defmodule Emily.Conformance.VitTest do ) end - test ":for_image_classification" do + mode_test ":for_image_classification" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-ViTForImageClassification"} @@ -71,7 +71,7 @@ defmodule Emily.Conformance.VitTest do "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 2} @@ -81,7 +81,7 @@ defmodule Emily.Conformance.VitTest do ) end - test ":for_masked_image_modeling" do + mode_test ":for_masked_image_modeling" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-ViTForMaskedImageModeling"} @@ -93,7 +93,7 @@ defmodule Emily.Conformance.VitTest do "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.pixel_values) == {1, 30, 30, 3} diff --git a/test/emily/conformance/whisper_full_test.exs b/test/emily/conformance/whisper_full_test.exs index 78e48c5..30c3ce3 100644 --- a/test/emily/conformance/whisper_full_test.exs +++ b/test/emily/conformance/whisper_full_test.exs @@ -31,7 +31,7 @@ defmodule Emily.Conformance.WhisperFullTest do @moduletag capture_log: true @moduletag timeout: 600_000 - test "openai/whisper-tiny forward pass matches pinned logits slice" do + mode_test "openai/whisper-tiny forward pass matches pinned logits slice", lane_tags: false do {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "openai/whisper-tiny"}) @@ -58,7 +58,7 @@ defmodule Emily.Conformance.WhisperFullTest do "decoder_attention_mask" => decoder_attention_mask } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 6, 51_865} diff --git a/test/emily/conformance/whisper_test.exs b/test/emily/conformance/whisper_test.exs index 57f7d15..4bfb3ff 100644 --- a/test/emily/conformance/whisper_test.exs +++ b/test/emily/conformance/whisper_test.exs @@ -31,7 +31,7 @@ defmodule Emily.Conformance.WhisperTest do @moduletag capture_log: true @moduletag timeout: 300_000 - test ":base" do + mode_test ":base" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-WhisperModel"}) @@ -43,7 +43,7 @@ defmodule Emily.Conformance.WhisperTest do "decoder_attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.hidden_state) == {1, 8, 16} @@ -55,7 +55,7 @@ defmodule Emily.Conformance.WhisperTest do ) end - test ":for_conditional_generation" do + mode_test ":for_conditional_generation" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-WhisperForConditionalGeneration"} @@ -69,7 +69,7 @@ defmodule Emily.Conformance.WhisperTest do "decoder_attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 8, 50_257} diff --git a/test/emily/opcode_parity_test.exs b/test/emily/opcode_parity_test.exs new file mode 100644 index 0000000..d7231a5 --- /dev/null +++ b/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 diff --git a/test/support/conformance_helper.ex b/test/support/conformance_helper.ex index 0ec9a52..51ddbad 100644 --- a/test/support/conformance_helper.ex +++ b/test/support/conformance_helper.ex @@ -29,7 +29,8 @@ defmodule Emily.ConformanceHelper do defmacro __using__(_opts) do quote do - import Emily.ConformanceHelper, only: [assert_all_close: 2, assert_all_close: 3] + import Emily.ConformanceHelper, + only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2, mode_test: 3] setup do Nx.default_backend(Emily.Backend) @@ -38,6 +39,119 @@ defmodule Emily.ConformanceHelper do end end + @doc """ + Define a conformance test in three lanes from a single body. + + Expands to three `ExUnit` tests that share `body` but bind a different + `predict_opts` keyword list: + + * the default lane binds `predict_opts` to `[]` — the evaluator path + Bumblebee/Axon use out of the box (the existing "eval'd" mode); + * the native lane binds `predict_opts` to + `[compiler: Emily.Compiler, native: true, native_fallback: :raise]` + and is additionally tagged `:native`; + * the fusion lane adds `native_compiled: true` (wrapping the replay in + `mx::compile`) and is additionally tagged `:native_compiled`. + + The module is already tagged `:conformance`, so the native and fusion + lanes carry that tag too: `mix test --only conformance` runs all three, + while `mix test --only native` / `mix test --only native_compiled` run + one lane each. Because every lane resolves the same HuggingFace repos, + whichever runs first reads from `~/.cache/bumblebee` for the rest — the + download is paid once. + + `mx::compile` reassociates f32, so the fusion lane's logits are not + bit-identical to the evaluator's; it shares the same reference and + tolerance as the other lanes (these tiny-random forwards drift well + within `assert_all_close`'s default), and `assert_finite!`-style smoke + tests are robust to the drift outright. + + The body must thread `predict_opts` into whatever drives the forward + pass so the two lanes assert against the *identical* reference and + cannot drift apart in maintenance: + + mode_test ":base" do + {:ok, %{model: model, params: params}} = Bumblebee.load_model(...) + outputs = Axon.predict(model, params, inputs, predict_opts) + assert_all_close(outputs.hidden_state, ...) + end + + For `Axon.build`-driven tests, build `init_fn` on the evaluator (params + are random-init, mode-irrelevant) and only `predict_fn` under + `predict_opts`, so the native lane gates the forward pass alone: + + {init_fn, _} = Axon.build(model) + {_, predict_fn} = Axon.build(model, predict_opts) + + `native_fallback: :raise` makes the native lane a no-fallback gate: an + op that does not lower fails the test rather than silently degrading to + the evaluator, so a red native lane is a concrete op-coverage gap. + + ## Options + + * `:lane_tags` (default `true`) — when `false`, the native and fusion + lanes are emitted *without* the cross-cutting `:native` / + `:native_compiled` tags. The heavyweight `*_full` suites pass + `lane_tags: false` so their compiler lanes stay gated behind the + 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([extra_tag], name, "", [], body), + lane( + [extra_tag, tag_lanes? && :native], + name, + " [native]", + [compiler: Emily.Compiler, native: true, native_fallback: :raise], + body + ), + lane( + [extra_tag, tag_lanes? && :native_compiled], + name, + " [native_compiled]", + [compiler: Emily.Compiler, native: true, native_fallback: :raise, native_compiled: true], + body + ) + ] + + quote do + (unquote_splicing(lanes)) + end + end + + # Build one `mode_test` lane: a `test` that binds `predict_opts` for the + # 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)) + + tag_attrs = for t <- tags, do: quote(do: @tag(unquote(t))) + + quote do + (unquote_splicing(tag_attrs)) + + test unquote(name_ast) do + var!(predict_opts) = unquote(predict_opts) + unquote(body) + end + end + end + @doc """ Assert that every element of `left` agrees with `right` within `atol + rtol * |right|`. diff --git a/test/test_helper.exs b/test/test_helper.exs index 58b751d..2205023 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -41,10 +41,29 @@ # the MLX `mx::fast::*` kernels via `Emily.Fast`. Run explicitly: # # mix test --only fast_kernels_full +# +# `:native` and `:native_compiled` are the expression-compiler lanes of +# the tiny-random conformance suites: every `mode_test` (see +# `Emily.ConformanceHelper`) re-runs the forward pass under +# `compiler: Emily.Compiler, native: true, native_fallback: :raise` +# (`:native`) and again with `native_compiled: true` wrapping the replay +# in `mx::compile` (`:native_compiled`), so the same PyTorch reference +# slice validates the evaluator, the native-compiled, and the fused +# paths. Those tests carry `:conformance` too, so `--only conformance` +# runs all three lanes; select one lane alone with: +# +# mix test --only native +# mix test --only native_compiled +# +# Listed in the default exclude defensively — every such test is already +# `:conformance`-tagged, but this keeps a future `:native`-only or +# `:native_compiled`-only test out of the default suite. ExUnit.start( max_cases: System.schedulers_online(), exclude: [ :conformance, + :native, + :native_compiled, :vit_full, :whisper_full, :distilbert_full,