Merge pull request #156 from ausimian/feat/expr-compiler-cm7

feat: graceful whole-defn fallback for the native compiler (CM7)

Author: ausimian

RELEASE.md

@@ -21,15 +21,28 @@
   dynamic KV-cache writes (`put_slice` at a runtime offset), container/tuple
   outputs, and `cond` (lowered to a select-chain). DistilBERT and ViT forwards
   run end-to-end under the compiler with `config :emily, :fallback, :raise`.
-  Unsupported constructs — `while` loops and arbitrary BEAM `reduce` functions —
-  raise a clear compile-time error rather than silently falling back, so
-  generation loops stay driven from Elixir.
+  Constructs the IR can't lower yet — `while` loops, arbitrary BEAM `reduce`
+  functions — are handled by the graceful fallback below.
 
   An opt-in compiled eval mode additionally wraps the replay in
   `mlx::core::compile`, fusing the elementwise runs (rms-norm, softmax, SiLU
   gating, residual adds) the replay leaves as separate kernels — measured at
   ~1.5–1.6× over the plain replay on a decode-shaped transformer block.
 
+- **Graceful native fallback — `native_fallback: :eval` (the default).** When a
+  `native: true` defn contains an op or construct the Expr compiler can't lower
+  yet, the *whole* defn now routes through `Nx.Defn.Evaluator` (each op then
+  dispatches through `Emily.Backend`, with its own per-op fallback) and emits a
+  one-shot `[:emily, :compiler, :fallback]` telemetry event — instead of
+  raising. This makes it safe to install the compiler globally on any model:
+
+      Nx.Defn.global_default_options(compiler: Emily.Compiler, native: true)
+
+  Covered subgraphs (e.g. encoder forwards) run fully native; the rest is
+  transparently evaluated. Pass `native_fallback: :raise` (or set
+  `config :emily, native_fallback: :raise`) to fail instead — the conformance
+  suites use this to prove a model lowers fully native.
+
 - `Emily.async_eval/1` (and `Emily.Native.async_eval/2`) schedule evaluation of
   one or more lazy graphs **without blocking on the GPU**, wrapping
   `mlx::core::async_eval`. The work is handed to the device's command queue and

config/test.exs

@@ -9,3 +9,11 @@ config :emily,
   debug_detect_nan_inf: false,
   test_fixture_debug_bounds_check: true,
   test_fixture_debug_detect_nan_inf: true
+
+# Keep the native compiler strict in the test suite: an op the Expr
+# compiler can't lower raises rather than silently falling back to the
+# Evaluator. This preserves the no-fallback conformance gates (CM5) and
+# the unsupported-op assertions. The runtime default is `:eval` (graceful
+# fallback); tests that exercise the fallback pass `native_fallback: :eval`
+# per call.
+config :emily, native_fallback: :raise

lib/emily/compiler.ex

@@ -68,6 +68,19 @@ defmodule Emily.Compiler do
       and Bumblebee passes `:cache` through for its own per-scope
       cache suffixing. Neither is used by the Evaluator walk, but
       rejecting them would break those servings.
+    * `:native` — `true` compiles the traced `Nx.Defn.Expr` to a flat
+      IR and replays the whole graph in a single NIF call per invocation.
+      Defaults to `false`, which runs the op-by-op Evaluator walk.
+    * `:native_fallback` — `:eval` (default) or `:raise`. Controls what
+      happens when `native: true` but the expression contains an op or
+      construct the IR can't lower yet. `:eval` routes the *whole* defn
+      through `Nx.Defn.Evaluator` (each op then dispatches through
+      `Emily.Backend`, with its own per-op `via_binary` fallback) and
+      fires a one-shot `[:emily, :compiler, :fallback]` event, so
+      installing `compiler: Emily.Compiler, native: true` globally is
+      safe on any model. `:raise` re-raises the lowering error instead —
+      use it in CI to prove a model lowers fully native. The per-call
+      option wins over `config :emily, :native_fallback, :eval | :raise`.
 
   Any other option is silently dropped. This matches how
   `Nx.Defn.Evaluator` and EXLA handle their own option lists, and is
@@ -106,15 +119,19 @@ defmodule Emily.Compiler do
     :max_concurrency,
     :batch_keys,
     :cache,
-    :native
+    :native,
+    :native_fallback
   ]
 
   @impl true
   def __jit__(key, vars, fun, args_list, opts) do
     opts = take_known_opts(opts)
 
     if Keyword.get(opts, :native, false) do
-      compile_native(vars, fun).(args_list)
+      case build_native(key, vars, fun, opts) do
+        {:ok, run} -> run.(args_list)
+        :fallback -> Evaluator.__jit__(key, vars, fun, args_list, drop_native_opts(opts))
+      end
     else
       Evaluator.__jit__(key, vars, fun, args_list, opts)
     end
@@ -125,24 +142,73 @@ defmodule Emily.Compiler do
     opts = take_known_opts(opts)
 
     if Keyword.get(opts, :native, false) do
-      compile_native(vars, fun)
+      case build_native(key, vars, fun, opts) do
+        {:ok, run} -> run
+        :fallback -> Evaluator.__compile__(key, vars, fun, drop_native_opts(opts))
+      end
     else
       Evaluator.__compile__(key, vars, fun, opts)
     end
   end
 
-  # The single-NIF compiled path (CM1+): trace the function into an
-  # Nx.Defn.Expr, lower it to a flat IR once, compile it into a `Program`
-  # resource (captured in this closure), and replay the whole graph in
-  # one NIF call per invocation. Op coverage is still partial — an
-  # unsupported op raises in `Emily.IR.lower/1` (no silent fallback).
-  defp compile_native(vars, fun) do
+  # Build the single-NIF native closure for `fun`, or signal `:fallback`.
+  #
+  # The Expr trace (`fun.(vars)`) runs *outside* the rescue so a genuine
+  # caller error surfaces unchanged; only the lowering + program build is
+  # guarded. `Emily.IR.lower/1` raises `ArgumentError` on an op or
+  # construct it can't lower yet. Unless `:native_fallback` is `:raise`,
+  # we emit a one-shot `[:emily, :compiler, :fallback]` event and return
+  # `:fallback`, so the caller routes the whole defn through
+  # `Nx.Defn.Evaluator` (which dispatches each op through `Emily.Backend`,
+  # with its own per-op `via_binary` fallback). This keeps a global
+  # `native: true` install safe on any model.
+  @spec build_native(term(), [Nx.Tensor.t()], fun(), keyword()) ::
+          {:ok, ([term()] -> [Nx.Tensor.t()])} | :fallback
+  defp build_native(key, vars, fun, opts) do
+    # Resolve (and validate) the mode up front so a misconfigured
+    # `:native_fallback` raises on every call — including the happy path —
+    # rather than lying dormant until the first lowering failure.
+    mode = native_fallback_mode(opts)
+
+    # The Expr trace runs outside `lower/3`'s guard so a genuine caller
+    # error surfaces unchanged.
     expr = fun.(vars)
 
     {template, leaves_rev} =
       Composite.traverse(expr, [], fn leaf, acc -> {Nx.to_template(leaf), [leaf | acc]} end)
 
-    program = leaves_rev |> Enum.reverse() |> IR.lower() |> Program.compile()
+    case lower(Enum.reverse(leaves_rev), mode, key) do
+      {:ok, ir} -> {:ok, replay_closure(template, ir)}
+      :fallback -> :fallback
+    end
+  end
+
+  # Lower the output leaves to a flat IR. `Emily.IR.lower/1` is the *only*
+  # guarded step: it raises `ArgumentError` on an op or construct it can't
+  # lower yet, which we turn into a graceful `:fallback` (or re-raise in
+  # `:raise` mode). `Program.compile/1` is deliberately kept outside the
+  # rescue (in `replay_closure/2`) — it raises only on malformed IR, i.e. a
+  # compiler bug, which must surface loudly rather than be masked as an
+  # "unsupported op" fallback.
+  defp lower(leaves, mode, key) do
+    {:ok, IR.lower(leaves)}
+  rescue
+    e in ArgumentError ->
+      case mode do
+        :raise ->
+          reraise(e, __STACKTRACE__)
+
+        :eval ->
+          Emily.Telemetry.compiler_fallback(key, e)
+          :fallback
+      end
+  end
+
+  # The single-NIF compiled path (CM1+): compile the lowered IR into a
+  # `Program` resource (captured in this closure) and replay the whole
+  # graph in one NIF call per invocation.
+  defp replay_closure(template, ir) do
+    program = Program.compile(ir)
 
     fn [params] ->
       worker = Emily.MlxStream.default_worker()
@@ -155,6 +221,32 @@ defmodule Emily.Compiler do
     end
   end
 
+  # Per-call `:native_fallback` opt wins over `config :emily,
+  # :native_fallback`, defaulting to `:eval`. `Keyword.fetch/2` (not `||`)
+  # so an explicit `native_fallback: false` is rejected, not silently
+  # treated as "unset".
+  defp native_fallback_mode(opts) do
+    mode =
+      case Keyword.fetch(opts, :native_fallback) do
+        {:ok, m} -> m
+        :error -> Application.get_env(:emily, :native_fallback, :eval)
+      end
+
+    case mode do
+      m when m in [:eval, :raise] ->
+        m
+
+      other ->
+        raise ArgumentError,
+              "invalid :native_fallback #{inspect(other)}; expected :eval | :raise"
+    end
+  end
+
+  # Strip the Emily-only native knobs before delegating to the Evaluator
+  # — it ignores keys it doesn't consume, but handing it `native: true`
+  # when we've decided *not* to compile natively would be misleading.
+  defp drop_native_opts(opts), do: Keyword.drop(opts, [:native, :native_fallback])
+
   defp native_ref(%T{data: %B{ref: r}}), do: r
   defp native_ref(%T{} = t), do: Nx.backend_transfer(t, B).data.ref
 

lib/emily/telemetry.ex

@@ -46,6 +46,18 @@ defmodule Emily.Telemetry do
   `:fallback` is unset (`true` → `:warn`, `false` → `:silent`).
   Prefer `:fallback` in new code; if both are set, `:fallback` wins.
 
+  ### Native-compiler fallback
+
+  `[:emily, :compiler, :fallback]` — a discrete event (not a span) that
+  fires when a `native: true` defn can't be lowered by the Expr compiler
+  and routes through `Nx.Defn.Evaluator` instead (see
+  `Emily.Compiler`'s `:native_fallback` option). Measurements:
+  `:count` (always `1`). Metadata: `:key` (the JIT key) and `:reason`
+  (the lowering error message, which names the unsupported op or
+  construct). A one-shot `Logger.warning` per distinct `:reason` is also
+  emitted — set `config :emily, :native_fallback, :raise` to fail
+  instead of falling back.
+
   ### Memory stats (poll-driven)
 
   `[:emily, :memory, :stats]` — discrete event, not a span. Call
@@ -137,13 +149,21 @@ defmodule Emily.Telemetry do
   end
 
   defp warn_fallback(op, input_shapes) do
-    key = {op, input_shapes}
+    warn_once(
+      {op, input_shapes},
+      "Emily: #{op} on shapes #{inspect(input_shapes)} fell back to " <>
+        "Nx.BinaryBackend; this path is ~100× slower than native MLX."
+    )
+  end
 
+  # One-shot `Logger.warning` per distinct `key`, deduped via the shared
+  # ETS table so a hot path that repeatedly falls back logs once, not per
+  # call. Keys are namespaced by caller (`{op, shapes}` for the backend
+  # fallback, `{:compiler, reason}` for the native-compiler fallback) so
+  # the two never collide.
+  defp warn_once(key, message) do
     if :ets.insert_new(@dedup_table, {key}) do
-      Logger.warning(
-        "Emily: #{op} on shapes #{inspect(input_shapes)} fell back to " <>
-          "Nx.BinaryBackend; this path is ~100× slower than native MLX."
-      )
+      Logger.warning(message)
     end
 
     :ok
@@ -155,6 +175,32 @@ defmodule Emily.Telemetry do
             "Set `config :emily, fallback: :warn` to log instead, or `:silent` to ignore."
   end
 
+  @doc false
+  # Called from `Emily.Compiler` when a `native: true` defn cannot be
+  # lowered by the Expr compiler and routes through `Nx.Defn.Evaluator`
+  # instead. Fires the discrete `[:emily, :compiler, :fallback]` event
+  # (always) and a one-shot `Logger.warning` per distinct reason, deduped
+  # via the shared ETS table. `reason` is the lowering error message — it
+  # names the unsupported op/construct.
+  @spec compiler_fallback(term(), Exception.t()) :: :ok
+  def compiler_fallback(key, exception) do
+    reason = Exception.message(exception)
+
+    :telemetry.execute(
+      [:emily, :compiler, :fallback],
+      %{count: 1},
+      %{key: key, reason: reason}
+    )
+
+    warn_once(
+      {:compiler, reason},
+      "Emily: native compilation fell back to Nx.Defn.Evaluator — #{reason} " <>
+        "The defn ran op-by-op via the evaluator. Set " <>
+        "`config :emily, native_fallback: :raise` (or pass " <>
+        "`native_fallback: :raise`) to fail instead."
+    )
+  end
+
   @doc false
   # Called exactly once from `Emily.Application.start/2` before the
   # supervisor starts. Safe to call again; an existing table is

test/emily/compiler_control_flow_test.exs

@@ -7,7 +7,10 @@ defmodule Emily.CompilerControlFlowTest do
   use ExUnit.Case, async: true
   import Nx.Defn
 
-  @native [compiler: Emily.Compiler, native: true]
+  # `native_fallback: :raise` is explicit (not relying on config/test.exs)
+  # so the "unsupported control flow raises" assertions below stay a local,
+  # config-independent gate.
+  @native [compiler: Emily.Compiler, native: true, native_fallback: :raise]
   @eval [compiler: Emily.Compiler]
 
   defn if_fn(x) do

test/emily/compiler_fallback_test.exs

@@ -0,0 +1,96 @@
+defmodule Emily.CompilerFallbackTest do
+  @moduledoc """
+  CM7 — graceful whole-defn fallback. When a `native: true` defn hits an
+  op the Expr compiler can't lower, the default (`native_fallback: :eval`)
+  routes the *whole* defn through `Nx.Defn.Evaluator` (each op then
+  dispatches through `Emily.Backend`) and fires a
+  `[:emily, :compiler, :fallback]` telemetry event — rather than raising.
+  `native_fallback: :raise` restores the strict no-fallback behaviour the
+  conformance gates rely on.
+
+  These pass `:native_fallback` per call (the test suite default is
+  `:raise`, set in `config/test.exs`), so no global state is mutated and
+  the cases stay `async`.
+  """
+  use ExUnit.Case, async: true
+  import Nx.Defn
+
+  # `Nx.sort` is not lowered by the Expr compiler yet (CM7) but IS
+  # supported by `Emily.Backend`, so the evaluator path produces a correct
+  # result to compare against.
+  defn(sort_fn(x), do: Nx.sort(x))
+
+  # Fully supported by the IR — must still compile native, no fallback.
+  defn(supported_fn(x), do: Nx.add(Nx.multiply(x, 2.0), 1.0))
+
+  defp t(data), do: Nx.tensor(data, backend: Emily.Backend)
+
+  # Named module-function handler (not an anonymous fn) so `:telemetry`
+  # doesn't log its local-function performance note during the run.
+  @doc false
+  def forward_event(_event, meas, meta, {pid, ref}), do: send(pid, {ref, :event, meas, meta})
+
+  defp attach(event, ref) do
+    id = "cm7-#{inspect(ref)}"
+    :telemetry.attach(id, event, &__MODULE__.forward_event/4, {self(), ref})
+    on_exit(fn -> :telemetry.detach(id) end)
+  end
+
+  describe "native_fallback: :eval (graceful, the runtime default)" do
+    test "an unsupported op falls back to the evaluator and matches it" do
+      x = t([3.0, 1.0, 2.0, 0.0])
+
+      native =
+        Nx.Defn.jit(&sort_fn/1, compiler: Emily.Compiler, native: true, native_fallback: :eval).(
+          x
+        )
+
+      eval = Nx.Defn.jit(&sort_fn/1, compiler: Emily.Compiler).(x)
+
+      assert %Emily.Backend{} = native.data
+      assert Nx.to_binary(native) == Nx.to_binary(eval)
+    end
+
+    test "fires a [:emily, :compiler, :fallback] event naming the op" do
+      ref = make_ref()
+      attach([:emily, :compiler, :fallback], ref)
+
+      Nx.Defn.jit(&sort_fn/1, compiler: Emily.Compiler, native: true, native_fallback: :eval).(
+        t([2.0, 1.0])
+      )
+
+      assert_receive {^ref, :event, %{count: 1}, %{reason: reason}}
+      assert reason =~ "sort"
+    end
+
+    test "a fully supported defn compiles native rather than falling back" do
+      # Asserted under `:raise` (not by checking the absence of a global
+      # telemetry event, which would race with concurrent async modules):
+      # with no fallback available, success proves the defn lowered fully
+      # native instead of silently degrading to the evaluator.
+      x = t([1.0, 2.0, 3.0])
+
+      out =
+        Nx.Defn.jit(&supported_fn/1,
+          compiler: Emily.Compiler,
+          native: true,
+          native_fallback: :raise
+        ).(x)
+
+      assert %Emily.Backend{} = out.data
+
+      assert Nx.to_binary(out) ==
+               Nx.to_binary(Nx.Defn.jit(&supported_fn/1, compiler: Emily.Compiler).(x))
+    end
+  end
+
+  describe "native_fallback: :raise (strict, the conformance-gate mode)" do
+    test "an unsupported op raises rather than falling back" do
+      assert_raise ArgumentError, ~r/sort/, fn ->
+        Nx.Defn.jit(&sort_fn/1, compiler: Emily.Compiler, native: true, native_fallback: :raise).(
+          t([1.0, 2.0])
+        )
+      end
+    end
+  end
+end