feat: lower the Nx.Block.LinAlg.* family in the native Expr compiler

Closes the LinAlg-block cluster on the Expr op-coverage checklist
(#188) — seven new lowerings: Cholesky, Solve, QR, Eigh, LU, SVD, and
Determinant.

Two new infrastructure pieces unlock multi-output and nested-block
lowering:

- replay_program in c_src/program.cpp grows a switch for the four
  multi-output linalg opcodes (LinalgQR / LinalgEigh / LinalgLU /
  LinalgSVD) that pushes N values per instruction, the same shape
  `while` already uses. compile_program's per-instruction arity
  calculation extends to match.
- The block-parameter collector goes "shallow": when it encounters a
  nested `:block` Expr node it walks only the block's `in_args` (which
  reference outer-scope tensors), not the inner `expr` whose fresh
  inner block-params live in a different scope. Determinant's
  expansion calls Nx.LinAlg.lu internally, so without this fix the
  outer block's parameter seeding would leak into LU's inner
  parameters and the binding would mis-resolve.

Single-output ops (LinalgCholesky, LinalgSolve) dispatch through the
standard `dispatch_op` path; multi-output ops use emit_multi/6 with
arity 2 (QR, Eigh) or 3 (LU, SVD), returning {:multi_refs, ...} so
the existing `:elem` projection works unchanged.

LU's s32 perm vector is post-processed into a permutation matrix via
`take(eye(n), perm, 0)`, mirroring Emily.Backend.native_lu/3. SVD in
`full_matrices?: false` mode is sliced from the full output to the
requested shape via the same `maybe_slice_svd/4` shape comparison as
Emily.Backend. Determinant has no fused kernel — both paths run the
block's composed 2x2 / 3x3 / NxN expansion via TopK-style parameter
seeding, bit-identical to the Evaluator. All seven dispatch
mx::linalg::* on the CPU stream per call (MLX's linalg is CPU-only),
matching c_src/ops/linalg.cpp's eager NIFs.

Tests in compiler_equivalence_test.exs cover each block on small
inputs (SPD for Cholesky, well-conditioned for Solve, mixed shapes
for QR/SVD, symmetric for Eigh, square for LU/Determinant). SVD's
sign convention is mathematically non-unique and MLX's S also drifts
by ~1 ULP across separate calls (the Backend property tests
acknowledge this at backend_test.exs:771), so SVD tests assert
shapes + singular values within tolerance instead of bit-exact;
every other LinAlg block matches the evaluator bit-for-bit.

Probe drops from 9 → 2 misses; only `count_leading_zeros` and
`population_count` remain (no MLX primitive — by design).

Author: ausimian

RELEASE.md

@@ -88,6 +88,28 @@
   seeding (every primitive in the expansion was already on the native
   path), bit-identical to the Evaluator.
 
+- **`Nx.LinAlg.cholesky` / `solve` / `qr` / `eigh` / `lu` / `svd` /
+  `determinant` lower natively** — closes the LinAlg-block cluster on
+  #188. Two new infrastructure pieces enable this: a small extension to
+  `replay_program` that pushes N values per instruction for multi-output
+  linalg ops (mirroring how `while` already pushes one value per
+  loop-carried state element), and a "shallow" parameter collector that
+  stops at nested `:block` nodes so the outer block's seeding doesn't
+  leak across scope boundaries (Determinant's expansion calls
+  `Nx.LinAlg.lu` internally — without the fix the IR would try to bind
+  LU's fresh inner block-params from the outer scope). Cholesky and
+  Solve are single-output; QR/Eigh/LU/SVD are multi-output. LU's
+  `s32` perm vector is post-processed into a permutation matrix via
+  `take(eye(n), perm, 0)` exactly like `Emily.Backend.native_lu/3`; SVD
+  in `full_matrices?: false` mode is sliced from the full output the
+  same way `Emily.Backend.maybe_slice_svd/4` does (the Gram thin-tall
+  workaround for Issue #84 stays on the eager path for now);
+  Determinant has no fused kernel — both paths run the block's composed
+  2x2 / 3x3 / NxN expansion, via the same TopK-style parameter seeding.
+  All seven dispatch to the CPU stream per call (MLX's `linalg::*` is
+  CPU-only); the C++ side overrides the replay stream just like the
+  eager linalg NIFs.
+
 - **`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

@@ -218,9 +218,32 @@ enum class Opcode : int64_t {
   // Routed to MLX's CPU stream (`mx::linalg::solve_triangular` is CPU-only),
   // mirroring c_src/ops/linalg.cpp's eager NIF.
   LinalgSolveTriangular = 114,
+  // CPU-only linalg decompositions / solvers (mx::linalg::*). Each
+  // routes the call to the CPU stream like the eager NIFs in
+  // c_src/ops/linalg.cpp. Single-output (Cholesky, Solve) dispatches
+  // through the standard `dispatch_op` path; multi-output (QR/Eigh/LU/SVD)
+  // is special-cased in replay_program — the dispatcher returns one
+  // array; multi-output instructions push N values into the program
+  // values vector, the same shape `while` uses.
+  // operands [a]; iattrs [[upper]]
+  LinalgCholesky = 115,
+  // operands [a, b]; no iattrs
+  LinalgSolve = 116,
+  // operands [a]; arity 2 (Q, R); no iattrs
+  LinalgQR = 117,
+  // operands [a]; arity 2 (vals, vecs); uplo hard-coded "L" to match
+  // Emily.Backend.native_eigh/3.
+  LinalgEigh = 118,
+  // operands [a]; arity 3 (perm[s32], L, U); the IR post-processes the
+  // perm vector into a permutation matrix via take(eye, perm, 0),
+  // mirroring Emily.Backend.native_lu/3.
+  LinalgLU = 119,
+  // operands [a]; arity 3 (U, S, V); `full_matrices` hard-coded `true` —
+  // the IR slices for the thin case, matching Emily.Backend.native_svd/3.
+  LinalgSVD = 120,
 };
 
-inline constexpr int64_t kOpcodeCount = 115;
+inline constexpr int64_t kOpcodeCount = 121;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -801,6 +824,27 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
     auto cpu = mx::default_stream(mx::Device(mx::Device::DeviceType::cpu));
     return mx::linalg::solve_triangular(in[0], in[1], upper, cpu);
   }
+  case Opcode::LinalgCholesky: {
+    // CPU-only, mirrors c_src/ops/linalg.cpp's eager NIF.
+    bool upper = scalar_at(iattrs, 0, "linalg_cholesky") != 0;
+    auto cpu = mx::default_stream(mx::Device(mx::Device::DeviceType::cpu));
+    return mx::linalg::cholesky(arg1(in, "linalg_cholesky"), upper, cpu);
+  }
+  case Opcode::LinalgSolve: {
+    need2(in, "linalg_solve");
+    auto cpu = mx::default_stream(mx::Device(mx::Device::DeviceType::cpu));
+    return mx::linalg::solve(in[0], in[1], cpu);
+  }
+  // Multi-output linalg (QR/Eigh/LU/SVD) is special-cased in
+  // replay_program (it pushes N values per instruction instead of one).
+  // Reaching these here would be a compiler bug — surface it loudly
+  // rather than silently fall through to "unknown opcode".
+  case Opcode::LinalgQR:
+  case Opcode::LinalgEigh:
+  case Opcode::LinalgLU:
+  case Opcode::LinalgSVD:
+    throw std::invalid_argument(
+        "multi-output linalg op is handled in replay_program");
   // --- Scatter (shares the eager index.cpp entry points) ---
   case Opcode::Scatter:
   case Opcode::ScatterAdd: {

c_src/program.cpp

@@ -242,6 +242,51 @@ std::vector<mx::array> replay_program(const Program &prog,
       continue;
     }
 
+    // Multi-output linalg decompositions (QR/Eigh/LU/SVD) follow the
+    // same shape as `while` above: one instruction, N values pushed.
+    // mx::linalg::* are CPU-only, mirroring c_src/ops/linalg.cpp's
+    // eager NIFs (the dispatcher's `s` arg is the replay stream, often
+    // GPU — override per call).
+    if (instr.opcode == Opcode::LinalgQR ||
+        instr.opcode == Opcode::LinalgEigh ||
+        instr.opcode == Opcode::LinalgLU ||
+        instr.opcode == Opcode::LinalgSVD) {
+      auto a = resolve(instr.operands.at(0));
+      auto cpu = mx::default_stream(mx::Device(mx::Device::DeviceType::cpu));
+      switch (instr.opcode) {
+      case Opcode::LinalgQR: {
+        auto [q, r] = mx::linalg::qr(a, cpu);
+        values.push_back(std::move(q));
+        values.push_back(std::move(r));
+        break;
+      }
+      case Opcode::LinalgEigh: {
+        auto [vals, vecs] = mx::linalg::eigh(a, "L", cpu);
+        values.push_back(std::move(vals));
+        values.push_back(std::move(vecs));
+        break;
+      }
+      case Opcode::LinalgLU: {
+        auto result = mx::linalg::lu(a, cpu);
+        values.push_back(std::move(result[0]));
+        values.push_back(std::move(result[1]));
+        values.push_back(std::move(result[2]));
+        break;
+      }
+      case Opcode::LinalgSVD: {
+        auto result = mx::linalg::svd(a, /*compute_uv=*/true, cpu);
+        values.push_back(std::move(result[0]));
+        values.push_back(std::move(result[1]));
+        values.push_back(std::move(result[2]));
+        break;
+      }
+      default:
+        // Unreachable — the outer `if` already filtered.
+        throw std::logic_error("unreachable multi-output linalg switch");
+      }
+      continue;
+    }
+
     std::vector<mx::array> operands;
     operands.reserve(instr.operands.size());
     for (auto r : instr.operands) {
@@ -318,8 +363,15 @@ compile_program(ErlNifEnv *, int64_t n_inputs,
     }
 
     Opcode op = static_cast<Opcode>(opcodes[i]);
+    // Multi-output instructions: `while` mirrors its operand count
+    // (one final-state value per loop-carried state element); the
+    // linalg decompositions have fixed arity (Q,R / vals,vecs / P,L,U
+    // / U,S,V). Everything else is single-output.
     int64_t out_count =
-        (op == Opcode::While) ? static_cast<int64_t>(operands[i].size()) : 1;
+        (op == Opcode::While) ? static_cast<int64_t>(operands[i].size())
+        : (op == Opcode::LinalgQR || op == Opcode::LinalgEigh) ? 2
+        : (op == Opcode::LinalgLU || op == Opcode::LinalgSVD)  ? 3
+                                                                : 1;
 
     std::vector<fine::ResourcePtr<Program>> sub;
     if (i < subprograms.size()) {

lib/emily/ir.ex

@@ -196,7 +196,19 @@ defmodule Emily.IR do
     # 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
+    linalg_solve_triangular: 114,
+    # CPU-only linalg decompositions / solvers (mx::linalg::*). Each
+    # routes through MLX's CPU stream like the eager NIFs in
+    # c_src/ops/linalg.cpp. Single-output: Cholesky, Solve. Multi-output:
+    # QR (Q,R), Eigh (vals,vecs), LU (perm[s32],L,U), SVD (U,S,V) — the
+    # IR emits them via emit_multi (one instr, N reserved output slots)
+    # the same way `while` does, and replay_program pushes N values.
+    linalg_cholesky: 115,
+    linalg_solve: 116,
+    linalg_qr: 117,
+    linalg_eigh: 118,
+    linalg_lu: 119,
+    linalg_svd: 120
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -1428,6 +1440,119 @@ defmodule Emily.IR do
     lower_node(expr, state)
   end
 
+  # ---- LinAlg blocks (Nx.Block.LinAlg.*) ----
+  #
+  # Each mirrors the matching Emily.Backend.native_* helper. Single-output
+  # ops (Cholesky, Solve) emit a single instruction; multi-output ops
+  # (QR/Eigh/LU/SVD) use emit_multi/6 to reserve their output slots,
+  # same machinery as `:while`. The CPU stream is forced inside the C++
+  # dispatcher because mx::linalg::* is CPU-only.
+
+  # Nx.LinAlg.cholesky — Backend hard-codes upper: false (lower).
+  defp lower_block(%Nx.Block.LinAlg.Cholesky{}, [t], _expr, out, state) do
+    {rt, state} = lower_node(t, state)
+    emit_coerced(state, :linalg_cholesky, [rt], [[0]], out.type)
+  end
+
+  # Nx.LinAlg.solve — single-output, no opts.
+  defp lower_block(%Nx.Block.LinAlg.Solve{}, [a, b], _expr, out, state) do
+    {ra, state} = lower_node(a, state)
+    {rb, state} = lower_node(b, state)
+    emit_coerced(state, :linalg_solve, [ra, rb], [], out.type)
+  end
+
+  # Nx.LinAlg.qr (mode: :reduced). Backend's `:complete` mode falls back
+  # via_binary; here we don't lower `:complete` (the catch-all raises),
+  # matching the no-fallback contract.
+  defp lower_block(%Nx.Block.LinAlg.QR{mode: :reduced}, [t], expr, _out, state)
+       when is_tuple(expr) do
+    {rt, state} = lower_node(t, state)
+    {base, state} = emit_multi(state, :linalg_qr, [rt], [], [], 2)
+    [q_leaf, r_leaf] = Tuple.to_list(expr)
+    {q_ref, state} = coerce({:instr, base + 0}, q_leaf.type, state)
+    {r_ref, state} = coerce({:instr, base + 1}, r_leaf.type, state)
+    {{:multi_refs, [q_ref, r_ref]}, state}
+  end
+
+  # Nx.LinAlg.eigh — uplo hard-coded "L" (lower) inside the dispatcher,
+  # matching Emily.Backend.native_eigh/3.
+  defp lower_block(%Nx.Block.LinAlg.Eigh{}, [t], expr, _out, state) when is_tuple(expr) do
+    {rt, state} = lower_node(t, state)
+    {base, state} = emit_multi(state, :linalg_eigh, [rt], [], [], 2)
+    [vals_leaf, vecs_leaf] = Tuple.to_list(expr)
+    {vals_ref, state} = coerce({:instr, base + 0}, vals_leaf.type, state)
+    {vecs_ref, state} = coerce({:instr, base + 1}, vecs_leaf.type, state)
+    {{:multi_refs, [vals_ref, vecs_ref]}, state}
+  end
+
+  # Nx.LinAlg.lu — MLX returns (perm[s32], L, U); Backend post-processes
+  # perm into a permutation matrix via `take(eye(n), perm, 0)`. The IR
+  # mirrors that here.
+  defp lower_block(%Nx.Block.LinAlg.LU{}, [t], expr, _out, state) when is_tuple(expr) do
+    {rt, state} = lower_node(t, state)
+    {base, state} = emit_multi(state, :linalg_lu, [rt], [], [], 3)
+    [p_leaf, l_leaf, u_leaf] = Tuple.to_list(expr)
+
+    # perm is a 1-D s32 vector (or batched). Use the leaf p_leaf's
+    # shape to size the eye matrix used as the row-selection source.
+    p_rank = tuple_size(p_leaf.shape)
+    n = elem(p_leaf.shape, p_rank - 1)
+
+    {eye_ref, state} =
+      materialize_const(
+        Nx.eye({n, n}, type: p_leaf.type, backend: Nx.BinaryBackend),
+        {n, n},
+        p_leaf.type,
+        state
+      )
+
+    {perm_s32, state} = emit(state, :astype, [{:instr, base + 0}], [[dtype_code({:s, 32})]])
+    {p_ref, state} = emit(state, :take, [eye_ref, perm_s32], [[0]])
+    {p_ref, state} = coerce(p_ref, p_leaf.type, state)
+    {l_ref, state} = coerce({:instr, base + 1}, l_leaf.type, state)
+    {u_ref, state} = coerce({:instr, base + 2}, u_leaf.type, state)
+    {{:multi_refs, [p_ref, l_ref, u_ref]}, state}
+  end
+
+  # Nx.LinAlg.svd — MLX's full SVD (`compute_uv: true`, U is m×m, V is
+  # n×n, S is length min(m,n)). For thin SVD (`full_matrices?: false`),
+  # slice U and V down to the requested output shape, mirroring
+  # Emily.Backend.maybe_slice_svd/4. The Gram-matrix workaround for very
+  # tall 2-D matrices (Issue #84) is left to the eager path; the IR
+  # equivalent can land alongside it.
+  defp lower_block(%Nx.Block.LinAlg.SVD{}, [t], expr, _out, state) when is_tuple(expr) do
+    {rt, state} = lower_node(t, state)
+    {base, state} = emit_multi(state, :linalg_svd, [rt], [], [], 3)
+    [u_leaf, s_leaf, v_leaf] = Tuple.to_list(expr)
+
+    rank = tuple_size(t.shape)
+    m = elem(t.shape, rank - 2)
+    n = elem(t.shape, rank - 1)
+
+    {u_ref, state} = maybe_slice_svd_ir({:instr, base + 0}, u_leaf.shape, {m, m}, state)
+    {u_ref, state} = coerce(u_ref, u_leaf.type, state)
+    {s_ref, state} = coerce({:instr, base + 1}, s_leaf.type, state)
+    {v_ref, state} = maybe_slice_svd_ir({:instr, base + 2}, v_leaf.shape, {n, n}, state)
+    {v_ref, state} = coerce(v_ref, v_leaf.type, state)
+    {{:multi_refs, [u_ref, s_ref, v_ref]}, state}
+  end
+
+  # Nx.LinAlg.determinant — Backend has no fused kernel; the eager path
+  # runs the block's composed expansion (2x2 / 3x3 / NxN). The IR does
+  # the same here via TopK-style param seeding so the result is
+  # bit-identical to the evaluator.
+  defp lower_block(%Nx.Block.LinAlg.Determinant{}, [t], expr, _out, state) do
+    {arg_ref, state} = lower_node(t, state)
+
+    seed =
+      expr
+      |> collect_block_params(%{})
+      |> Map.new(fn {id, 0} -> {id, arg_ref} end)
+
+    state = %{state | cache: Map.merge(state.cache, seed)}
+    lower_node(expr, state)
+  end
+
   # Any other block struct raises. Lowering the block's composed
   # expansion would silently diverge from the Evaluator whenever
   # Emily.Backend.block/4 dispatches that struct through a fused / native
@@ -1492,13 +1617,22 @@ defmodule Emily.IR do
   defp tensors?(list), do: Enum.all?(list, &match?(%T{}, &1))
 
   # Collect `{parameter_id => position}` for every block-local `:parameter`
-  # reachable from `node`. Used to bind a block expansion's fresh parameters
-  # to the real in_args (see the Nx.Block.TopK lowering). Walks only the small
-  # expansion graph — parameters are leaves, so it never descends into the
-  # in_args' own (possibly large) graph.
+  # reachable from `node` *in the current block scope*. Used to bind a
+  # block expansion's fresh parameters to the real in_args (see the
+  # Nx.Block.TopK / Nx.Block.LinAlg.Determinant lowerings). For nested
+  # `:block` nodes (e.g. determinant's expansion calls Nx.LinAlg.lu),
+  # walk only the in_args — which reference outer-scope tensors — and
+  # NOT the nested `expr`, whose fresh inner block params live in a
+  # different scope and would otherwise leak in as spurious positions.
   defp collect_block_params(%T{data: %Nx.Defn.Expr{op: :parameter, id: id, args: [pos]}}, acc),
     do: Map.put_new(acc, id, pos)
 
+  defp collect_block_params(
+         %T{data: %Nx.Defn.Expr{op: :block, args: [_struct, in_args, _expr, _cb]}},
+         acc
+       ),
+       do: Enum.reduce(in_args, acc, &collect_block_params/2)
+
   defp collect_block_params(%T{data: %Nx.Defn.Expr{args: args}}, acc),
     do: Enum.reduce(args, acc, &collect_block_params/2)
 
@@ -1529,6 +1663,23 @@ defmodule Emily.IR do
     Enum.to_list(0..(rank - 3)//1) ++ [rank - 1, rank - 2]
   end
 
+  # Slice U/V from the SVD's full output down to the thin-SVD shape if the
+  # block requested it. Mirrors Emily.Backend.maybe_slice_svd/4 — when the
+  # trailing two dims match the full shape (`full_last2`), no-op; otherwise
+  # slice from origin to `out_shape` with stride 1.
+  defp maybe_slice_svd_ir(ref, out_shape, full_last2, state) do
+    rank = tuple_size(out_shape)
+
+    if {elem(out_shape, rank - 2), elem(out_shape, rank - 1)} == full_last2 do
+      {ref, state}
+    else
+      starts = List.duplicate(0, rank)
+      strides = List.duplicate(1, rank)
+      stops = Tuple.to_list(out_shape)
+      emit(state, :slice, [ref], [starts, stops, strides])
+    end
+  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.