Merge pull request #191 from ausimian/feat/expr-compiler-pad-eye-trisolve

feat: lower pad, eye and triangular_solve in the native Expr compiler

Author: ausimian

RELEASE.md

@@ -66,6 +66,17 @@
   primitive) lowers to `(a != 0) != (b != 0)`, mirroring the eager
   Backend composition. All three are bit-identical to the Evaluator.
 
+- **`pad`, `eye`, `triangular_solve` lower natively** — closes the
+  top-level-Expr-op cluster on #188. `pad` is the constant-pad path
+  (interior dilation raises, same as `Emily.Backend.pad/4`); `eye`
+  materialises as a captured constant (identity matrix at lower time,
+  same trick `iota` uses); `triangular_solve` decomposes all four
+  `transform_a` × `left_side` combinations into transposes around a
+  bare `linalg_solve_triangular` opcode, mirroring
+  `Emily.Backend.triangular_solve/4`. The C++ dispatcher routes
+  `mx::linalg::solve_triangular` to the CPU stream per call (it's
+  CPU-only), matching the eager NIF.
+
 - **`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

c_src/emily/opcodes.hpp

@@ -207,9 +207,20 @@ enum class Opcode : int64_t {
   // matching Emily.Backend.quotient/3).
   Arctan2 = 111,
   FloorDivide = 112,
+  // Constant-aware pad: pads `input` with the `pad_value` scalar by `low_pad`
+  // and `high_pad` on each `axis`. operands [input, pad_value]; iattrs
+  // [[axes...], [lows...], [highs...]]. MLX has no interior dilation, so
+  // the lowerer rejects interior > 0 (matches Emily.Backend.pad/4).
+  Pad = 113,
+  // Solve A x = b (or x A = b) where A is triangular. The Backend handles
+  // `transform_a`/`left_side` in the lowerer (via transpose ops on a/b);
+  // this opcode is the bare kernel call. operands [a, b]; iattrs [[upper]].
+  // Routed to MLX's CPU stream (`mx::linalg::solve_triangular` is CPU-only),
+  // mirroring c_src/ops/linalg.cpp's eager NIF.
+  LinalgSolveTriangular = 114,
 };
 
-inline constexpr int64_t kOpcodeCount = 113;
+inline constexpr int64_t kOpcodeCount = 115;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -765,6 +776,31 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
   case Opcode::FloorDivide:
     need2(in, "floor_divide");
     return mx::floor_divide(in[0], in[1], s);
+  // --- Shape / linalg peers ---
+  case Opcode::Pad:
+    // operands [input, pad_value]; iattrs [[axes...], [lows...], [highs...]].
+    // Mirrors c_src/ops/shape.cpp's pad_nif: `mx::pad` with "constant" mode
+    // (interior is rejected up-stack by the Elixir lowerer; same constraint
+    // as the eager NIF).
+    if (in.size() != 2) {
+      throw std::invalid_argument("pad expects 2 operands, got " +
+                                  std::to_string(in.size()));
+    }
+    return mx::pad(in[0], emily::to_int_vec(attr_at(iattrs, 0, "pad")),
+                   emily::to_mlx_shape(attr_at(iattrs, 1, "pad")),
+                   emily::to_mlx_shape(attr_at(iattrs, 2, "pad")), in[1],
+                   "constant", s);
+  case Opcode::LinalgSolveTriangular: {
+    // mx::linalg::solve_triangular is CPU-only — eager NIF in
+    // c_src/ops/linalg.cpp routes to mx::default_stream(cpu) too. The
+    // dispatcher's `s` arg is the replay stream (typically GPU); override
+    // here so the solver runs on CPU like the eager path.
+    need2(in, "linalg_solve_triangular");
+    bool upper =
+        scalar_at(iattrs, 0, "linalg_solve_triangular") != 0;
+    auto cpu = mx::default_stream(mx::Device(mx::Device::DeviceType::cpu));
+    return mx::linalg::solve_triangular(in[0], in[1], upper, cpu);
+  }
   // --- Scatter (shares the eager index.cpp entry points) ---
   case Opcode::Scatter:
   case Opcode::ScatterAdd: {

lib/emily/ir.ex

@@ -188,7 +188,15 @@ defmodule Emily.IR do
     # so the IR routes Nx :quotient through it (cast both to out.type,
     # floor_divide, same bits as the eager Backend).
     arctan2: 111,
-    floor_divide: 112
+    floor_divide: 112,
+    # Constant-aware pad (operands [input, pad_value]; iattrs [axes, lows,
+    # highs]); MLX has no interior dilation so the lowerer rejects
+    # interior > 0, same as Emily.Backend.pad/4.
+    pad: 113,
+    # CPU-only triangular solve (operands [a, b]; iattrs [[upper]]); the
+    # lowerer handles transform_a/left_side by transposing a/b/output
+    # around this bare kernel call, mirroring Emily.Backend.triangular_solve/4.
+    linalg_solve_triangular: 114
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -829,6 +837,16 @@ defmodule Emily.IR do
     |> materialize_const(t.shape, t.type, state)
   end
 
+  # eye: same shape as iota — a pure creation op with all-static shape /
+  # type. Materialize as a captured constant via Nx.eye on the host
+  # backend (which already handles the rank > 2 batch case, identity on
+  # the trailing two axes — matching Emily.Backend.eye/2's broadcast).
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :eye, args: []}} = t, state) do
+    t.shape
+    |> Nx.eye(type: t.type, backend: Nx.BinaryBackend)
+    |> materialize_const(t.shape, t.type, state)
+  end
+
   # Nx.block node: args [struct, in_args, expr, callback]. Known fused
   # structs lower to their fused opcode (matching Emily.Backend.block/4,
   # which the Evaluator dispatches through); unknown structs lower by
@@ -945,6 +963,72 @@ defmodule Emily.IR do
     emit_coerced(state, op, [rt, rs, ri], [window, strides, pad_lo, pad_hi], t.type)
   end
 
+  # pad(input, pad_value, padding_config): constant-pad each axis by
+  # `padding_config = [{lo, hi, interior}, ...]`. Mirrors
+  # Emily.Backend.pad/4 — MLX has no interior dilation, so interior > 0
+  # raises (no fallback to interior expansion). The pad_value is an Expr
+  # scalar tensor and lowers as an operand (input scalars baked at lower
+  # time end up as a const ref; runtime scalars as inputs — both routes
+  # work because the pad opcode takes the value as operand[1]).
+  defp lower_op(
+         %T{data: %Nx.Defn.Expr{op: :pad, args: [a, pad_value, padding_config]}} = t,
+         state
+       ) do
+    lows = Enum.map(padding_config, fn {lo, _, _} -> lo end)
+    highs = Enum.map(padding_config, fn {_, hi, _} -> hi end)
+    interiors = Enum.map(padding_config, fn {_, _, interior} -> interior end)
+
+    if Enum.any?(interiors, &(&1 > 0)) do
+      raise ArgumentError,
+            "Emily Expr compiler does not lower :pad with interior > 0 " <>
+              "(MLX has no primitive; Emily.Backend.pad/4 also raises)."
+    end
+
+    axes = Enum.to_list(0..(length(lows) - 1)//1)
+    {ra, state} = lower_node(a, state)
+    {rp, state} = lower_node(pad_value, state)
+    emit_coerced(state, :pad, [ra, rp], [axes, lows, highs], t.type)
+  end
+
+  # triangular_solve(a, b, opts): solve A x = b (or x A = b) with A
+  # triangular. Mirrors Emily.Backend.triangular_solve/4 — the bare kernel
+  # is operands [a, b] + [[upper]], and the four `transform_a` /
+  # `left_side` combinations are decomposed here into transposes around
+  # the kernel call (same Native.transpose sequence the eager Backend
+  # uses). MLX's mx::linalg::solve_triangular runs on the CPU stream,
+  # which the C++ dispatcher overrides per call.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :triangular_solve, args: [a, b, opts]}} = t, state) do
+    {ra, state} = lower_node(a, state)
+    {rb, state} = lower_node(b, state)
+    lower = opts[:lower]
+
+    {r, state} =
+      case {opts[:transform_a], opts[:left_side]} do
+        {:none, true} ->
+          emit(state, :linalg_solve_triangular, [ra, rb], [[bool_int(not lower)]])
+
+        {:transpose, true} ->
+          {ra_t, state} = emit(state, :transpose, [ra], [mat_transpose_axes(a.shape)])
+          emit(state, :linalg_solve_triangular, [ra_t, rb], [[bool_int(lower)]])
+
+        {:none, false} ->
+          {ra_t, state} = emit(state, :transpose, [ra], [mat_transpose_axes(a.shape)])
+          {rb_t, state} = emit(state, :transpose, [rb], [mat_transpose_axes(b.shape)])
+          {xt, state} = emit(state, :linalg_solve_triangular, [ra_t, rb_t], [[bool_int(lower)]])
+          emit(state, :transpose, [xt], [mat_transpose_axes(t.shape)])
+
+        {:transpose, false} ->
+          {rb_t, state} = emit(state, :transpose, [rb], [mat_transpose_axes(b.shape)])
+
+          {xt, state} =
+            emit(state, :linalg_solve_triangular, [ra, rb_t], [[bool_int(not lower)]])
+
+          emit(state, :transpose, [xt], [mat_transpose_axes(t.shape)])
+      end
+
+    coerce(r, t.type, state)
+  end
+
   # Nx.reverse along one or more axes (the conv backward flips the kernel).
   # Reversing is order-independent across axes, so chain a single-axis
   # `flip` (mx negative-stride slice) per axis. Empty axes => identity.
@@ -1364,6 +1448,13 @@ defmodule Emily.IR do
 
   defp dim_product(axes, shape), do: Enum.reduce(axes, 1, &(elem(shape, &1) * &2))
 
+  # The "matrix transpose" axes permutation: keep leading batch axes, swap
+  # the trailing two. Mirrors Emily.Backend.mat_transpose_axes/1.
+  defp mat_transpose_axes(shape) do
+    rank = tuple_size(shape)
+    Enum.to_list(0..(rank - 3)//1) ++ [rank - 1, rank - 2]
+  end
+
   # Coerce a ref to `type` (emit an astype). MLX astype to the same dtype
   # is a no-op, so this is safe to apply unconditionally — it mirrors
   # Emily.Backend.wrap/3's coerce and keeps the node's dtype exact.

test/emily/compiler_equivalence_test.exs

@@ -350,6 +350,101 @@ defmodule Emily.CompilerEquivalenceTest do
       assert_equiv(fn t -> Nx.add(t, Nx.iota({4})) end, [x])
       assert_equiv(fn t -> Nx.add(t, Nx.iota({4}, type: :f32)) end, [x])
     end
+
+    # eye is a pure creation op like iota — the IR materialises the
+    # identity matrix as a captured constant via Nx.eye on the host
+    # backend (which handles the rank > 2 batch broadcast itself, matching
+    # Emily.Backend.eye/2). The native and evaluator paths both end up
+    # comparing against the same materialised bits.
+    test "eye lowers to a constant and matches (2-D and batched)" do
+      x = et([[0.0, 0.0], [0.0, 0.0]])
+      assert_equiv(fn t -> Nx.add(t, Nx.eye(2)) end, [x])
+      assert_equiv(fn t -> Nx.add(t, Nx.eye({2, 2}, type: :f32)) end, [x])
+
+      y = Nx.broadcast(Nx.tensor(0.0, backend: Emily.Backend), {2, 3, 3})
+      assert_equiv(fn t -> Nx.add(t, Nx.eye({2, 3, 3}, type: :f32)) end, [y])
+    end
+  end
+
+  describe "pad / triangular_solve" do
+    # pad with constant value matches Emily.Backend.pad/4 (mx::pad with
+    # mode "constant", no interior dilation). Cover symmetric and
+    # asymmetric padding on both 1-D and 2-D inputs.
+    test "constant-pad on 1-D and 2-D inputs matches the evaluator" do
+      x = et([1.0, 2.0, 3.0])
+      assert_equiv(fn t -> Nx.pad(t, 0.0, [{1, 1, 0}]) end, [x])
+      assert_equiv(fn t -> Nx.pad(t, -1.0, [{2, 0, 0}]) end, [x])
+
+      y = et([[1.0, 2.0], [3.0, 4.0]])
+      assert_equiv(fn t -> Nx.pad(t, 0.0, [{1, 1, 0}, {2, 2, 0}]) end, [y])
+      assert_equiv(fn t -> Nx.pad(t, 99.0, [{0, 1, 0}, {1, 0, 0}]) end, [y])
+    end
+
+    # Interior padding (interior > 0) has no MLX primitive — both the
+    # eager Backend and the IR raise on it. The IR's raise must surface
+    # cleanly (the graceful `:eval` fallback path is what catches it in
+    # production; here we assert the underlying behaviour).
+    test "pad with interior > 0 raises in the IR (no fallback)" do
+      x = et([1.0, 2.0, 3.0])
+
+      assert_raise ArgumentError, ~r/interior > 0/, fn ->
+        run(fn t -> Nx.pad(t, 0.0, [{0, 0, 1}]) end, [x],
+          compiler: Emily.Compiler,
+          native: true,
+          native_fallback: :raise
+        )
+      end
+    end
+
+    # triangular_solve over the four `transform_a` × `left_side`
+    # combinations. The IR composes each case with transposes around the
+    # bare `linalg_solve_triangular` opcode (operands [a, b]; iattrs
+    # [[upper]]) — exactly the sequence Emily.Backend.triangular_solve/4
+    # uses. mx::linalg::solve_triangular is CPU-only; the C++ dispatcher
+    # overrides the stream per call, mirroring c_src/ops/linalg.cpp.
+    test "lower-triangular A x = b (default opts) matches the evaluator" do
+      a = et([[2.0, 0.0], [3.0, 1.0]])
+      b = et([4.0, 9.0])
+      assert_equiv(&Nx.LinAlg.triangular_solve/2, [a, b])
+    end
+
+    test "upper-triangular A x = b matches" do
+      a = et([[2.0, 1.0], [0.0, 3.0]])
+      b = et([5.0, 6.0])
+      assert_equiv(fn a, b -> Nx.LinAlg.triangular_solve(a, b, lower: false) end, [a, b])
+    end
+
+    test "transposed A (transform_a: :transpose, left_side: true) matches" do
+      a = et([[2.0, 0.0], [3.0, 1.0]])
+      b = et([4.0, 9.0])
+
+      assert_equiv(
+        fn a, b -> Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose) end,
+        [a, b]
+      )
+    end
+
+    test "right-side solve x A = b (left_side: false) matches" do
+      a = et([[2.0, 0.0], [3.0, 1.0]])
+      b = et([[4.0, 9.0], [1.0, 2.0]])
+
+      assert_equiv(
+        fn a, b -> Nx.LinAlg.triangular_solve(a, b, left_side: false) end,
+        [a, b]
+      )
+    end
+
+    test "transposed right-side solve (transform_a: :transpose, left_side: false) matches" do
+      a = et([[2.0, 0.0], [3.0, 1.0]])
+      b = et([[4.0, 9.0], [1.0, 2.0]])
+
+      assert_equiv(
+        fn a, b ->
+          Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose, left_side: false)
+        end,
+        [a, b]
+      )
+    end
   end
 
   describe "fused kernels (Emily.Fast blocks)" do