Merge pull request #151 from ausimian/feat/expr-compiler-cm4

feat: control flow — :cond (where-chain) + honest deferrals (CM4)

Author: ausimian

lib/emily/ir.ex

@@ -195,6 +195,7 @@ defmodule Emily.IR do
 
   alias Emily.Fast.Block, as: FB
   alias Emily.Quantization.Block, as: QB
+  alias Nx.Defn.Tree
   alias Nx.Tensor, as: T
 
   # Nx Expr op -> IR opcode. Arithmetic/bitwise cast both operands to the
@@ -522,6 +523,76 @@ defmodule Emily.IR do
     lower_block(struct, in_args, expr, t, state)
   end
 
+  # cond: raw args [clauses, last], clauses = [{pred, body}, ...]. Lower to
+  # a select chain `where(p1, b1, where(p2, b2, ... last))`. ALL branches
+  # are evaluated (Nx branches are side-effect-free and shape-compatible);
+  # the result value matches the Evaluator's chosen branch exactly — only
+  # the cost differs (not-taken branches are computed and discarded by the
+  # elementwise select). The predicate is a whole-tensor scalar bool, so
+  # `where` selects a branch wholesale.
+  #
+  # Caveat: a not-taken branch is still computed. On MLX an out-of-bounds
+  # 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
+    {last_ref, state} = lower_node(last, state)
+
+    {result, state} =
+      Enum.reduce(Enum.reverse(clauses), {last_ref, state}, fn {pred, body}, {else_ref, st} ->
+        {pred_ref, st} = lower_node(pred, st)
+        {body_ref, st} = lower_node(body, st)
+        {pred_ref, st} = emit(st, :astype, [pred_ref], [[dtype_code({:pred, 1})]])
+        emit(st, :where, [pred_ref, body_ref, else_ref])
+      end)
+
+    coerce(result, t.type, 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.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :attach_token, args: [token, expr]}}, state) do
+    if Tree.has_hooks?(token, %{}) do
+      raise ArgumentError,
+            "Emily Expr compiler does not support hooks under native compilation " <>
+              "(they require a mid-graph callback into Elixir)."
+    end
+
+    lower_node(expr, state)
+  end
+
+  # reduce / window_reduce with a user-supplied BEAM reducer cannot be
+  # compiled — the reducer would have to run on the host mid-graph. The
+  # fixed-identity aggregates (sum/product/max/min) are separate ops and
+  # already lower natively; only an arbitrary reducer reaches here.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: op}}, _state) when op in [:reduce, :window_reduce] do
+    raise ArgumentError,
+          "Emily Expr compiler cannot lower #{inspect(op)} with an arbitrary " <>
+            "reducer function (it would require a host callback mid-graph; no " <>
+            "fallback). Use the native aggregates (sum/product/reduce_max/" <>
+            "reduce_min) where possible."
+  end
+
+  # while is deferred to a follow-up: the single-NIF replay has no loop
+  # construct, so a data-dependent while needs static-trip unrolling or a
+  # worker-side synced loop. defn while is not used by the core transformer
+  # forwards (decode/generation loops run in Elixir today).
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :while}}, _state) do
+    raise ArgumentError,
+          "Emily Expr compiler does not yet lower defn `while` loops (deferred — " <>
+            "the single-NIF replay has no loop construct)."
+  end
+
+  # :elem is a tuple projection, emitted for any tuple-returning expression
+  # (defn `while`, multi-output ops). Deferred alongside the constructs
+  # that produce surviving tuples.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :elem}}, _state) do
+    raise ArgumentError,
+          "Emily Expr compiler does not yet lower :elem (tuple projection) — it " <>
+            "arises from defn `while` and multi-output ops, which are deferred."
+  end
+
   defp lower_op(%T{data: %Nx.Defn.Expr{op: op}}, _state) do
     raise ArgumentError,
           "Emily Expr compiler does not yet lower op #{inspect(op)} " <>

test/emily/compiler_control_flow_test.exs

@@ -0,0 +1,88 @@
+defmodule Emily.CompilerControlFlowTest do
+  @moduledoc """
+  CM4 — `Nx.Defn` control-flow constructs compile single-NIF, bit-identical
+  to the Evaluator. `:cond` lowers to a `where`-chain (all branches
+  evaluated; the value matches the chosen branch exactly).
+  """
+  use ExUnit.Case, async: true
+  import Nx.Defn
+
+  @native [compiler: Emily.Compiler, native: true]
+  @eval [compiler: Emily.Compiler]
+
+  defn if_fn(x) do
+    if Nx.greater(Nx.sum(x), 0) do
+      Nx.multiply(x, 2)
+    else
+      Nx.negate(x)
+    end
+  end
+
+  defn cond3_fn(x) do
+    s = Nx.sum(x)
+
+    cond do
+      Nx.greater(s, 10) -> Nx.multiply(x, 10)
+      Nx.greater(s, 0) -> Nx.multiply(x, 2)
+      true -> Nx.negate(x)
+    end
+  end
+
+  defn nested_if_fn(x) do
+    y = if Nx.greater(Nx.sum(x), 0), do: Nx.add(x, 1), else: Nx.subtract(x, 1)
+    if Nx.greater(Nx.reduce_max(y), 5), do: Nx.divide(y, 2), else: y
+  end
+
+  defn while_fn(x) do
+    {_i, acc} =
+      while {i = 0, acc = x}, i < 5 do
+        {i + 1, Nx.multiply(acc, 2)}
+      end
+
+    acc
+  end
+
+  defp equiv(fun, x) do
+    native = Nx.Defn.jit(fun, @native).(x)
+    eval = Nx.Defn.jit(fun, @eval).(x)
+    assert Nx.to_binary(native) == Nx.to_binary(eval)
+    native
+  end
+
+  describe ":cond / if" do
+    test "two-branch if selects the right branch" do
+      for data <- [[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0], [0.0, 0.0, 0.0]] do
+        equiv(&if_fn/1, Nx.tensor(data, backend: Emily.Backend))
+      end
+    end
+
+    test "multi-clause cond picks the first matching predicate" do
+      # sums: 12 (>10), 3 (0..10), -6 (else) exercise each clause.
+      for data <- [[5.0, 4.0, 3.0], [1.0, 1.0, 1.0], [-2.0, -2.0, -2.0]] do
+        equiv(&cond3_fn/1, Nx.tensor(data, backend: Emily.Backend))
+      end
+    end
+
+    test "nested conds compose" do
+      for data <- [[3.0, 3.0, 3.0], [-1.0, 0.0, 1.0], [10.0, 10.0, 10.0]] do
+        equiv(&nested_if_fn/1, Nx.tensor(data, backend: Emily.Backend))
+      end
+    end
+  end
+
+  describe "unsupported control flow raises (no silent fallback)" do
+    test "arbitrary reduce/2 fn raises a clear error" do
+      f = fn x -> Nx.reduce(x, 0.0, fn a, b -> Nx.add(a, b) end) end
+
+      assert_raise ArgumentError, ~r/arbitrary reducer/, fn ->
+        Nx.Defn.jit(f, @native).(Nx.tensor([1.0, 2.0, 3.0], backend: Emily.Backend))
+      end
+    end
+
+    test "defn while raises (deferred)" do
+      assert_raise ArgumentError, ~r/while/, fn ->
+        Nx.Defn.jit(&while_fn/1, @native).(Nx.tensor([1.0, 2.0], backend: Emily.Backend))
+      end
+    end
+  end
+end