refactor: address CM2 code-review findings

- Raise on unhandled block structs instead of lowering their composed
  expansion. For blocks Emily.Backend.block/4 dispatches through a fused
  / native kernel (SDPAWithSinks, the Nx.Block.LinAlg.* / Take / FFT /
  cumulative families) the composed path would silently diverge from the
  Evaluator; a clear "unsupported" is the honest no-fallback behaviour.
  More fused blocks land with their opcodes.
- Gate the affine packed-shape output-feature formula to mode "affine"
  in quantized_matmul_defn; other layouts derive the feature count from
  the dequantized shape.
- Add Emily.IR.emit_coerced/5 and route the lower_block clauses through
  it (removes the emit-then-coerce footgun).
- C++: fold scalar_attr into scalar_at (one implementation); fix the
  stale qmode_from_code doc comment.

Author: ausimian

c_src/emily/opcodes.hpp

@@ -110,7 +110,7 @@ enum class Opcode : int64_t {
 
 inline constexpr int64_t kOpcodeCount = 61;
 
-// Quant mode code (Emily.IR.mode_code/1) -> MLX mode string.
+// Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
   switch (code) {
   case 0: return "affine";
@@ -160,15 +160,6 @@ inline const std::vector<int64_t> &attr0(const std::vector<std::vector<int64_t>>
   return a[0];
 }
 
-inline int64_t scalar_attr(const std::vector<std::vector<int64_t>> &a,
-                           const char *name) {
-  const auto &v = attr0(a, name);
-  if (v.size() != 1) {
-    throw std::invalid_argument(std::string(name) + " expects one attribute value");
-  }
-  return v[0];
-}
-
 inline const std::vector<int64_t> &
 attr_at(const std::vector<std::vector<int64_t>> &a, std::size_t i,
         const char *name) {
@@ -314,7 +305,7 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
   // --- Cast / shape ---
   case Opcode::Astype:
     return mx::astype(arg1(in, "astype"),
-                      emily::to_mlx_dtype_code(scalar_attr(iattrs, "astype")), s);
+                      emily::to_mlx_dtype_code(scalar_at(iattrs, 0, "astype")), s);
   case Opcode::Reshape:
     return mx::reshape(arg1(in, "reshape"),
                        emily::to_mlx_shape(attr0(iattrs, "reshape")), s);
@@ -374,7 +365,7 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
                                   std::to_string(in.size()));
     }
     auto eps = static_cast<float>(
-        emily::f64_from_bits(scalar_attr(iattrs, "fast_rms_norm")));
+        emily::f64_from_bits(scalar_at(iattrs, 0, "fast_rms_norm")));
     return mx::fast::rms_norm(in[0], std::optional<mx::array>(in[1]), eps, s);
   }
   case Opcode::FastLayerNorm: {
@@ -383,7 +374,7 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
                                   std::to_string(in.size()));
     }
     auto eps = static_cast<float>(
-        emily::f64_from_bits(scalar_attr(iattrs, "fast_layer_norm")));
+        emily::f64_from_bits(scalar_at(iattrs, 0, "fast_layer_norm")));
     return mx::fast::layer_norm(in[0], std::optional<mx::array>(in[1]),
                                 std::optional<mx::array>(in[2]), eps, s);
   }

lib/emily/ir.ex

@@ -501,25 +501,22 @@ defmodule Emily.IR do
   defp lower_block(%FB.RMSNorm{eps: eps}, [x, weight], _expr, t, state) do
     {rx, state} = lower_node(x, state)
     {rw, state} = lower_node(weight, state)
-    {r, state} = emit(state, :fast_rms_norm, [rx, rw], [[float_bits(eps)]])
-    coerce(r, t.type, state)
+    emit_coerced(state, :fast_rms_norm, [rx, rw], [[float_bits(eps)]], t.type)
   end
 
   defp lower_block(%FB.LayerNorm{eps: eps}, [x, weight, bias], _expr, t, state) do
     {rx, state} = lower_node(x, state)
     {rw, state} = lower_node(weight, state)
     {rb, state} = lower_node(bias, state)
-    {r, state} = emit(state, :fast_layer_norm, [rx, rw, rb], [[float_bits(eps)]])
-    coerce(r, t.type, state)
+    emit_coerced(state, :fast_layer_norm, [rx, rw, rb], [[float_bits(eps)]], t.type)
   end
 
   defp lower_block(%FB.RoPE{} = b, [x, offset], _expr, t, state) do
     {rx, state} = lower_node(x, state)
     {ro, state} = lower_node(offset, state)
 
     attrs = [[b.dims], [bool_int(b.traditional)], [float_bits(b.base)], [float_bits(b.scale)]]
-    {r, state} = emit(state, :fast_rope, [rx, ro], attrs)
-    coerce(r, t.type, state)
+    emit_coerced(state, :fast_rope, [rx, ro], attrs, t.type)
   end
 
   defp lower_block(%FB.RoPEWithFreqs{} = b, [x, offset, freqs], _expr, t, state) do
@@ -528,25 +525,29 @@ defmodule Emily.IR do
     {rf, state} = lower_node(freqs, state)
 
     attrs = [[b.dims], [bool_int(b.traditional)], [float_bits(b.scale)]]
-    {r, state} = emit(state, :fast_rope_freqs, [rx, ro, rf], attrs)
-    coerce(r, t.type, state)
+    emit_coerced(state, :fast_rope_freqs, [rx, ro, rf], attrs, t.type)
   end
 
   defp lower_block(%FB.SDPA{scale: scale, causal: causal}, [q, k, v], _expr, t, state) do
     {rq, state} = lower_node(q, state)
     {rk, state} = lower_node(k, state)
     {rv, state} = lower_node(v, state)
-    {r, state} = emit(state, :fast_sdpa, [rq, rk, rv], [[float_bits(scale)], [bool_int(causal)]])
-    coerce(r, t.type, state)
+
+    emit_coerced(
+      state,
+      :fast_sdpa,
+      [rq, rk, rv],
+      [[float_bits(scale)], [bool_int(causal)]],
+      t.type
+    )
   end
 
   defp lower_block(%FB.SDPAWithMask{scale: scale}, [q, k, v, mask], _expr, t, state) do
     {rq, state} = lower_node(q, state)
     {rk, state} = lower_node(k, state)
     {rv, state} = lower_node(v, state)
     {rm, state} = lower_node(mask, state)
-    {r, state} = emit(state, :fast_sdpa_mask, [rq, rk, rv, rm], [[float_bits(scale)]])
-    coerce(r, t.type, state)
+    emit_coerced(state, :fast_sdpa_mask, [rq, rk, rv, rm], [[float_bits(scale)]], t.type)
   end
 
   defp lower_block(%QB.QuantizedMatmul{} = qb, [x, q, s, b], _expr, t, state) do
@@ -562,14 +563,20 @@ defmodule Emily.IR do
       [Map.fetch!(@quant_modes, qb.mode)]
     ]
 
-    {r, state} = emit(state, :quantized_matmul, [rx, rq, rs, rb], attrs)
-    coerce(r, t.type, state)
+    emit_coerced(state, :quantized_matmul, [rx, rq, rs, rb], attrs, t.type)
   end
 
-  # Unknown block struct: lower its pre-composed default expansion (the
-  # `expr` arg) instead of the fused kernel — no runtime fallback.
-  defp lower_block(_struct, _in_args, expr, _t, state) do
-    lower_node(expr, state)
+  # 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
+  # kernel (e.g. SDPAWithSinks, the Nx.Block.LinAlg.* / Take / FFT /
+  # cumulative families) — a worse failure than a clear "unsupported".
+  # Additional fused blocks are added alongside their opcode.
+  defp lower_block(struct, _in_args, _expr, _t, _state) do
+    raise ArgumentError,
+          "Emily Expr compiler does not yet lower the block " <>
+            "#{inspect(struct.__struct__)} (no fallback). Supported: RMSNorm, " <>
+            "LayerNorm, RoPE, RoPEWithFreqs, SDPA, SDPAWithMask, QuantizedMatmul."
   end
 
   defp bool_int(true), do: 1
@@ -592,6 +599,14 @@ defmodule Emily.IR do
     {ref, %{state | instrs: [instr | state.instrs], n_instrs: state.n_instrs + 1}}
   end
 
+  # Emit an instruction then coerce its output to `type` (mirrors
+  # Emily.Backend.wrap/3). The trailing coerce is mandatory on every
+  # value-producing op, so the helper keeps the per-clause tail honest.
+  defp emit_coerced(state, opcode, operands, iattrs, type) do
+    {r, state} = emit(state, opcode, operands, iattrs)
+    coerce(r, type, state)
+  end
+
   # Materialize an Nx tensor (already on a host backend) as a captured
   # const / weight ref, held by the program for its lifetime. `:const`
   # holds materialized literal constants (and iota); `:capture` holds

lib/emily/quantization.ex

@@ -171,12 +171,15 @@ defmodule Emily.Quantization do
   # bits ∈ {2,4,8} the packed last axis unpacks by 32/bits.
   defp qmm_out_features(%QuantizedWeight{value: value}, true), do: elem(Nx.shape(value), 0)
 
-  defp qmm_out_features(%QuantizedWeight{value: value, bits: bits}, false)
+  defp qmm_out_features(%QuantizedWeight{value: value, bits: bits, mode: "affine"}, false)
        when bits in [2, 4, 8] do
+    # Affine packs `32 / bits` lanes per u32 along the last axis.
     shape = Nx.shape(value)
     elem(shape, tuple_size(shape) - 1) * div(32, bits)
   end
 
+  # Other layouts (microscaled, or affine bits 3/6): derive the output
+  # feature count from the dequantized weight's shape.
   defp qmm_out_features(qw, false) do
     shape = Nx.shape(dequantize_defn(qw))
     elem(shape, tuple_size(shape) - 1)