Skip to content

Commit e0047df

Browse files
committed
feat: lower pad, eye and triangular_solve in the native Expr compiler
Closes the top-level Expr op cluster on the Expr op-coverage checklist (#188). Two new IR opcodes — `pad` (113) and `linalg_solve_triangular` (114) — plus a dedicated `eye` materialise-as-constant clause. - `pad` lowers to `mx::pad` ("constant" mode) — operands [input, pad_value]; iattrs [axes, lows, highs]. Interior dilation (> 0) raises matching Emily.Backend.pad/4, which has the same MLX constraint. - `eye` follows the same shape-static-creation pattern as iota: Nx.eye on Nx.BinaryBackend at lower time, materialised as a captured const. Nx.eye already handles rank > 2 by broadcasting the identity onto the trailing two axes, matching Emily.Backend.eye/2. - `triangular_solve` decomposes Nx's four `transform_a` × `left_side` combinations into transposes around the bare `linalg_solve_triangular` opcode (operands [a, b]; iattrs [[upper]]), mirroring Emily.Backend.triangular_solve/4. The C++ dispatcher routes `mx::linalg::solve_triangular` to the CPU stream per call (MLX's solver is CPU-only), matching c_src/ops/linalg.cpp. Probe drops from 15 → 12 misses. Tests cover pad on 1-D / 2-D with symmetric and asymmetric padding (plus the interior > 0 raise), eye 2-D and batched (rank-3 with leading batch broadcast), and triangular_solve across all four `transform_a` × `left_side` cases plus the upper-triangular variant.
1 parent 981bac2 commit e0047df

4 files changed

Lines changed: 235 additions & 2 deletions

File tree

RELEASE.md

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,17 @@
6666
primitive) lowers to `(a != 0) != (b != 0)`, mirroring the eager
6767
Backend composition. All three are bit-identical to the Evaluator.
6868

69+
- **`pad`, `eye`, `triangular_solve` lower natively** — closes the
70+
top-level-Expr-op cluster on #188. `pad` is the constant-pad path
71+
(interior dilation raises, same as `Emily.Backend.pad/4`); `eye`
72+
materialises as a captured constant (identity matrix at lower time,
73+
same trick `iota` uses); `triangular_solve` decomposes all four
74+
`transform_a` × `left_side` combinations into transposes around a
75+
bare `linalg_solve_triangular` opcode, mirroring
76+
`Emily.Backend.triangular_solve/4`. The C++ dispatcher routes
77+
`mx::linalg::solve_triangular` to the CPU stream per call (it's
78+
CPU-only), matching the eager NIF.
79+
6980
- **`take_along_axis` lowers natively**`Nx.take_along_axis` (the
7081
`Nx.Block.TakeAlongAxis` block) now compiles under the native single-NIF
7182
path, mirroring `Emily.Backend.native_take_along_axis/4` (cast indices to

c_src/emily/opcodes.hpp

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -207,9 +207,20 @@ enum class Opcode : int64_t {
207207
// matching Emily.Backend.quotient/3).
208208
Arctan2 = 111,
209209
FloorDivide = 112,
210+
// Constant-aware pad: pads `input` with the `pad_value` scalar by `low_pad`
211+
// and `high_pad` on each `axis`. operands [input, pad_value]; iattrs
212+
// [[axes...], [lows...], [highs...]]. MLX has no interior dilation, so
213+
// the lowerer rejects interior > 0 (matches Emily.Backend.pad/4).
214+
Pad = 113,
215+
// Solve A x = b (or x A = b) where A is triangular. The Backend handles
216+
// `transform_a`/`left_side` in the lowerer (via transpose ops on a/b);
217+
// this opcode is the bare kernel call. operands [a, b]; iattrs [[upper]].
218+
// Routed to MLX's CPU stream (`mx::linalg::solve_triangular` is CPU-only),
219+
// mirroring c_src/ops/linalg.cpp's eager NIF.
220+
LinalgSolveTriangular = 114,
210221
};
211222

212-
inline constexpr int64_t kOpcodeCount = 113;
223+
inline constexpr int64_t kOpcodeCount = 115;
213224

214225
// Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
215226
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,
765776
case Opcode::FloorDivide:
766777
need2(in, "floor_divide");
767778
return mx::floor_divide(in[0], in[1], s);
779+
// --- Shape / linalg peers ---
780+
case Opcode::Pad:
781+
// operands [input, pad_value]; iattrs [[axes...], [lows...], [highs...]].
782+
// Mirrors c_src/ops/shape.cpp's pad_nif: `mx::pad` with "constant" mode
783+
// (interior is rejected up-stack by the Elixir lowerer; same constraint
784+
// as the eager NIF).
785+
if (in.size() != 2) {
786+
throw std::invalid_argument("pad expects 2 operands, got " +
787+
std::to_string(in.size()));
788+
}
789+
return mx::pad(in[0], emily::to_int_vec(attr_at(iattrs, 0, "pad")),
790+
emily::to_mlx_shape(attr_at(iattrs, 1, "pad")),
791+
emily::to_mlx_shape(attr_at(iattrs, 2, "pad")), in[1],
792+
"constant", s);
793+
case Opcode::LinalgSolveTriangular: {
794+
// mx::linalg::solve_triangular is CPU-only — eager NIF in
795+
// c_src/ops/linalg.cpp routes to mx::default_stream(cpu) too. The
796+
// dispatcher's `s` arg is the replay stream (typically GPU); override
797+
// here so the solver runs on CPU like the eager path.
798+
need2(in, "linalg_solve_triangular");
799+
bool upper =
800+
scalar_at(iattrs, 0, "linalg_solve_triangular") != 0;
801+
auto cpu = mx::default_stream(mx::Device(mx::Device::DeviceType::cpu));
802+
return mx::linalg::solve_triangular(in[0], in[1], upper, cpu);
803+
}
768804
// --- Scatter (shares the eager index.cpp entry points) ---
769805
case Opcode::Scatter:
770806
case Opcode::ScatterAdd: {

lib/emily/ir.ex

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,15 @@ defmodule Emily.IR do
188188
# so the IR routes Nx :quotient through it (cast both to out.type,
189189
# floor_divide, same bits as the eager Backend).
190190
arctan2: 111,
191-
floor_divide: 112
191+
floor_divide: 112,
192+
# Constant-aware pad (operands [input, pad_value]; iattrs [axes, lows,
193+
# highs]); MLX has no interior dilation so the lowerer rejects
194+
# interior > 0, same as Emily.Backend.pad/4.
195+
pad: 113,
196+
# CPU-only triangular solve (operands [a, b]; iattrs [[upper]]); the
197+
# lowerer handles transform_a/left_side by transposing a/b/output
198+
# around this bare kernel call, mirroring Emily.Backend.triangular_solve/4.
199+
linalg_solve_triangular: 114
192200
}
193201

194202
# Quant mode string -> code; decoded by qmode_from_code in
@@ -829,6 +837,16 @@ defmodule Emily.IR do
829837
|> materialize_const(t.shape, t.type, state)
830838
end
831839

840+
# eye: same shape as iota — a pure creation op with all-static shape /
841+
# type. Materialize as a captured constant via Nx.eye on the host
842+
# backend (which already handles the rank > 2 batch case, identity on
843+
# the trailing two axes — matching Emily.Backend.eye/2's broadcast).
844+
defp lower_op(%T{data: %Nx.Defn.Expr{op: :eye, args: []}} = t, state) do
845+
t.shape
846+
|> Nx.eye(type: t.type, backend: Nx.BinaryBackend)
847+
|> materialize_const(t.shape, t.type, state)
848+
end
849+
832850
# Nx.block node: args [struct, in_args, expr, callback]. Known fused
833851
# structs lower to their fused opcode (matching Emily.Backend.block/4,
834852
# which the Evaluator dispatches through); unknown structs lower by
@@ -945,6 +963,72 @@ defmodule Emily.IR do
945963
emit_coerced(state, op, [rt, rs, ri], [window, strides, pad_lo, pad_hi], t.type)
946964
end
947965

966+
# pad(input, pad_value, padding_config): constant-pad each axis by
967+
# `padding_config = [{lo, hi, interior}, ...]`. Mirrors
968+
# Emily.Backend.pad/4 — MLX has no interior dilation, so interior > 0
969+
# raises (no fallback to interior expansion). The pad_value is an Expr
970+
# scalar tensor and lowers as an operand (input scalars baked at lower
971+
# time end up as a const ref; runtime scalars as inputs — both routes
972+
# work because the pad opcode takes the value as operand[1]).
973+
defp lower_op(
974+
%T{data: %Nx.Defn.Expr{op: :pad, args: [a, pad_value, padding_config]}} = t,
975+
state
976+
) do
977+
lows = Enum.map(padding_config, fn {lo, _, _} -> lo end)
978+
highs = Enum.map(padding_config, fn {_, hi, _} -> hi end)
979+
interiors = Enum.map(padding_config, fn {_, _, interior} -> interior end)
980+
981+
if Enum.any?(interiors, &(&1 > 0)) do
982+
raise ArgumentError,
983+
"Emily Expr compiler does not lower :pad with interior > 0 " <>
984+
"(MLX has no primitive; Emily.Backend.pad/4 also raises)."
985+
end
986+
987+
axes = Enum.to_list(0..(length(lows) - 1)//1)
988+
{ra, state} = lower_node(a, state)
989+
{rp, state} = lower_node(pad_value, state)
990+
emit_coerced(state, :pad, [ra, rp], [axes, lows, highs], t.type)
991+
end
992+
993+
# triangular_solve(a, b, opts): solve A x = b (or x A = b) with A
994+
# triangular. Mirrors Emily.Backend.triangular_solve/4 — the bare kernel
995+
# is operands [a, b] + [[upper]], and the four `transform_a` /
996+
# `left_side` combinations are decomposed here into transposes around
997+
# the kernel call (same Native.transpose sequence the eager Backend
998+
# uses). MLX's mx::linalg::solve_triangular runs on the CPU stream,
999+
# which the C++ dispatcher overrides per call.
1000+
defp lower_op(%T{data: %Nx.Defn.Expr{op: :triangular_solve, args: [a, b, opts]}} = t, state) do
1001+
{ra, state} = lower_node(a, state)
1002+
{rb, state} = lower_node(b, state)
1003+
lower = opts[:lower]
1004+
1005+
{r, state} =
1006+
case {opts[:transform_a], opts[:left_side]} do
1007+
{:none, true} ->
1008+
emit(state, :linalg_solve_triangular, [ra, rb], [[bool_int(not lower)]])
1009+
1010+
{:transpose, true} ->
1011+
{ra_t, state} = emit(state, :transpose, [ra], [mat_transpose_axes(a.shape)])
1012+
emit(state, :linalg_solve_triangular, [ra_t, rb], [[bool_int(lower)]])
1013+
1014+
{:none, false} ->
1015+
{ra_t, state} = emit(state, :transpose, [ra], [mat_transpose_axes(a.shape)])
1016+
{rb_t, state} = emit(state, :transpose, [rb], [mat_transpose_axes(b.shape)])
1017+
{xt, state} = emit(state, :linalg_solve_triangular, [ra_t, rb_t], [[bool_int(lower)]])
1018+
emit(state, :transpose, [xt], [mat_transpose_axes(t.shape)])
1019+
1020+
{:transpose, false} ->
1021+
{rb_t, state} = emit(state, :transpose, [rb], [mat_transpose_axes(b.shape)])
1022+
1023+
{xt, state} =
1024+
emit(state, :linalg_solve_triangular, [ra, rb_t], [[bool_int(not lower)]])
1025+
1026+
emit(state, :transpose, [xt], [mat_transpose_axes(t.shape)])
1027+
end
1028+
1029+
coerce(r, t.type, state)
1030+
end
1031+
9481032
# Nx.reverse along one or more axes (the conv backward flips the kernel).
9491033
# Reversing is order-independent across axes, so chain a single-axis
9501034
# `flip` (mx negative-stride slice) per axis. Empty axes => identity.
@@ -1364,6 +1448,13 @@ defmodule Emily.IR do
13641448

13651449
defp dim_product(axes, shape), do: Enum.reduce(axes, 1, &(elem(shape, &1) * &2))
13661450

1451+
# The "matrix transpose" axes permutation: keep leading batch axes, swap
1452+
# the trailing two. Mirrors Emily.Backend.mat_transpose_axes/1.
1453+
defp mat_transpose_axes(shape) do
1454+
rank = tuple_size(shape)
1455+
Enum.to_list(0..(rank - 3)//1) ++ [rank - 1, rank - 2]
1456+
end
1457+
13671458
# Coerce a ref to `type` (emit an astype). MLX astype to the same dtype
13681459
# is a no-op, so this is safe to apply unconditionally — it mirrors
13691460
# Emily.Backend.wrap/3's coerce and keeps the node's dtype exact.

test/emily/compiler_equivalence_test.exs

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -350,6 +350,101 @@ defmodule Emily.CompilerEquivalenceTest do
350350
assert_equiv(fn t -> Nx.add(t, Nx.iota({4})) end, [x])
351351
assert_equiv(fn t -> Nx.add(t, Nx.iota({4}, type: :f32)) end, [x])
352352
end
353+
354+
# eye is a pure creation op like iota — the IR materialises the
355+
# identity matrix as a captured constant via Nx.eye on the host
356+
# backend (which handles the rank > 2 batch broadcast itself, matching
357+
# Emily.Backend.eye/2). The native and evaluator paths both end up
358+
# comparing against the same materialised bits.
359+
test "eye lowers to a constant and matches (2-D and batched)" do
360+
x = et([[0.0, 0.0], [0.0, 0.0]])
361+
assert_equiv(fn t -> Nx.add(t, Nx.eye(2)) end, [x])
362+
assert_equiv(fn t -> Nx.add(t, Nx.eye({2, 2}, type: :f32)) end, [x])
363+
364+
y = Nx.broadcast(Nx.tensor(0.0, backend: Emily.Backend), {2, 3, 3})
365+
assert_equiv(fn t -> Nx.add(t, Nx.eye({2, 3, 3}, type: :f32)) end, [y])
366+
end
367+
end
368+
369+
describe "pad / triangular_solve" do
370+
# pad with constant value matches Emily.Backend.pad/4 (mx::pad with
371+
# mode "constant", no interior dilation). Cover symmetric and
372+
# asymmetric padding on both 1-D and 2-D inputs.
373+
test "constant-pad on 1-D and 2-D inputs matches the evaluator" do
374+
x = et([1.0, 2.0, 3.0])
375+
assert_equiv(fn t -> Nx.pad(t, 0.0, [{1, 1, 0}]) end, [x])
376+
assert_equiv(fn t -> Nx.pad(t, -1.0, [{2, 0, 0}]) end, [x])
377+
378+
y = et([[1.0, 2.0], [3.0, 4.0]])
379+
assert_equiv(fn t -> Nx.pad(t, 0.0, [{1, 1, 0}, {2, 2, 0}]) end, [y])
380+
assert_equiv(fn t -> Nx.pad(t, 99.0, [{0, 1, 0}, {1, 0, 0}]) end, [y])
381+
end
382+
383+
# Interior padding (interior > 0) has no MLX primitive — both the
384+
# eager Backend and the IR raise on it. The IR's raise must surface
385+
# cleanly (the graceful `:eval` fallback path is what catches it in
386+
# production; here we assert the underlying behaviour).
387+
test "pad with interior > 0 raises in the IR (no fallback)" do
388+
x = et([1.0, 2.0, 3.0])
389+
390+
assert_raise ArgumentError, ~r/interior > 0/, fn ->
391+
run(fn t -> Nx.pad(t, 0.0, [{0, 0, 1}]) end, [x],
392+
compiler: Emily.Compiler,
393+
native: true,
394+
native_fallback: :raise
395+
)
396+
end
397+
end
398+
399+
# triangular_solve over the four `transform_a` × `left_side`
400+
# combinations. The IR composes each case with transposes around the
401+
# bare `linalg_solve_triangular` opcode (operands [a, b]; iattrs
402+
# [[upper]]) — exactly the sequence Emily.Backend.triangular_solve/4
403+
# uses. mx::linalg::solve_triangular is CPU-only; the C++ dispatcher
404+
# overrides the stream per call, mirroring c_src/ops/linalg.cpp.
405+
test "lower-triangular A x = b (default opts) matches the evaluator" do
406+
a = et([[2.0, 0.0], [3.0, 1.0]])
407+
b = et([4.0, 9.0])
408+
assert_equiv(&Nx.LinAlg.triangular_solve/2, [a, b])
409+
end
410+
411+
test "upper-triangular A x = b matches" do
412+
a = et([[2.0, 1.0], [0.0, 3.0]])
413+
b = et([5.0, 6.0])
414+
assert_equiv(fn a, b -> Nx.LinAlg.triangular_solve(a, b, lower: false) end, [a, b])
415+
end
416+
417+
test "transposed A (transform_a: :transpose, left_side: true) matches" do
418+
a = et([[2.0, 0.0], [3.0, 1.0]])
419+
b = et([4.0, 9.0])
420+
421+
assert_equiv(
422+
fn a, b -> Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose) end,
423+
[a, b]
424+
)
425+
end
426+
427+
test "right-side solve x A = b (left_side: false) matches" do
428+
a = et([[2.0, 0.0], [3.0, 1.0]])
429+
b = et([[4.0, 9.0], [1.0, 2.0]])
430+
431+
assert_equiv(
432+
fn a, b -> Nx.LinAlg.triangular_solve(a, b, left_side: false) end,
433+
[a, b]
434+
)
435+
end
436+
437+
test "transposed right-side solve (transform_a: :transpose, left_side: false) matches" do
438+
a = et([[2.0, 0.0], [3.0, 1.0]])
439+
b = et([[4.0, 9.0], [1.0, 2.0]])
440+
441+
assert_equiv(
442+
fn a, b ->
443+
Nx.LinAlg.triangular_solve(a, b, transform_a: :transpose, left_side: false)
444+
end,
445+
[a, b]
446+
)
447+
end
353448
end
354449

355450
describe "fused kernels (Emily.Fast blocks)" do

0 commit comments

Comments
 (0)