Merge pull request #189 from ausimian/feat/expr-compiler-unary-ops

feat: lower 19 more unary ops in the native Expr compiler

Author: ausimian

RELEASE.md

@@ -50,6 +50,14 @@
   selection on the native path. Remaining gaps (`gather`/scatter,
   pooling/`window_*`, cumulative) continue to work via the graceful fallback.
 
+- **19 more unary ops lower natively** — `expm1`, `tan`, `sinh`, `cosh`,
+  `acos`/`asin`/`atan`, `acosh`/`asinh`/`atanh`, `round`, `bitwise_not`,
+  `is_nan`/`is_infinity`, `conjugate`, `real`/`imag` route to the same
+  `mx::*` primitive as the eager unary NIF (bit-identical to the
+  Evaluator); `erfc` and `cbrt` compose from existing ops, mirroring
+  `Emily.Backend`'s eager composition. Closes the largest cluster on the
+  Expr op-coverage checklist (#188).
+
 - **`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

@@ -179,9 +179,31 @@ enum class Opcode : int64_t {
   // idx0, ...] (one s32 index array per scattered axis); iattrs [[axes...]].
   Scatter = 92,    // overwrite (last write wins on duplicates)
   ScatterAdd = 93, // accumulate
+  // Unary elementwise (round 2 — added alongside the @unary_ops expansion
+  // for the missing Nx ops; map to the same mx::* primitives the eager
+  // unary NIFs use, see c_src/ops/unary.cpp).
+  Expm1 = 94,
+  Tan = 95,
+  Sinh = 96,
+  Cosh = 97,
+  Arccos = 98,
+  Arcsin = 99,
+  Arctan = 100,
+  Arccosh = 101,
+  Arcsinh = 102,
+  Arctanh = 103,
+  // Round-half-away-from-zero. Backend hard-codes decimals=0
+  // (Nx.round/1 takes no decimals arg); the dispatcher does too.
+  Round = 104,
+  BitwiseInvert = 105,
+  Isnan = 106,
+  Isinf = 107,
+  Conjugate = 108,
+  Real = 109,
+  Imag = 110,
 };
 
-inline constexpr int64_t kOpcodeCount = 94;
+inline constexpr int64_t kOpcodeCount = 111;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -695,6 +717,41 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
                            emily::to_mlx_shape(attr_at(iattrs, 0, "irfftn")),
                            emily::to_int_vec(attr_at(iattrs, 1, "irfftn")),
                            mx::fft::FFTNorm::Backward, s);
+  // --- Unary elementwise (round 2) ---
+  case Opcode::Expm1:
+    return mx::expm1(arg1(in, "expm1"), s);
+  case Opcode::Tan:
+    return mx::tan(arg1(in, "tan"), s);
+  case Opcode::Sinh:
+    return mx::sinh(arg1(in, "sinh"), s);
+  case Opcode::Cosh:
+    return mx::cosh(arg1(in, "cosh"), s);
+  case Opcode::Arccos:
+    return mx::arccos(arg1(in, "arccos"), s);
+  case Opcode::Arcsin:
+    return mx::arcsin(arg1(in, "arcsin"), s);
+  case Opcode::Arctan:
+    return mx::arctan(arg1(in, "arctan"), s);
+  case Opcode::Arccosh:
+    return mx::arccosh(arg1(in, "arccosh"), s);
+  case Opcode::Arcsinh:
+    return mx::arcsinh(arg1(in, "arcsinh"), s);
+  case Opcode::Arctanh:
+    return mx::arctanh(arg1(in, "arctanh"), s);
+  case Opcode::Round:
+    return mx::round(arg1(in, "round"), /*decimals=*/0, s);
+  case Opcode::BitwiseInvert:
+    return mx::bitwise_invert(arg1(in, "bitwise_invert"), s);
+  case Opcode::Isnan:
+    return mx::isnan(arg1(in, "isnan"), s);
+  case Opcode::Isinf:
+    return mx::isinf(arg1(in, "isinf"), s);
+  case Opcode::Conjugate:
+    return mx::conjugate(arg1(in, "conjugate"), s);
+  case Opcode::Real:
+    return mx::real(arg1(in, "real"), s);
+  case Opcode::Imag:
+    return mx::imag(arg1(in, "imag"), s);
   // --- Scatter (shares the eager index.cpp entry points) ---
   case Opcode::Scatter:
   case Opcode::ScatterAdd: {

lib/emily/ir.ex

@@ -158,7 +158,31 @@ defmodule Emily.IR do
     # idx0, ...] (one s32 index array per scattered axis); iattrs [[axes...]].
     # scatter overwrites (last-write on duplicates); scatter_add accumulates.
     scatter: 92,
-    scatter_add: 93
+    scatter_add: 93,
+    # Unary elementwise (round 2 — the missing Nx ops alongside the
+    # original @unary_ops set). Names mirror the eager unary NIF names
+    # (c_src/ops/unary.cpp) so the IR opcode atom matches the MLX entry
+    # point — Nx renaming (e.g. acos -> arccos, is_nan -> isnan) happens
+    # in @unary_ops below, just like the eager Backend's @renamed_unary.
+    expm1: 94,
+    tan: 95,
+    sinh: 96,
+    cosh: 97,
+    arccos: 98,
+    arcsin: 99,
+    arctan: 100,
+    arccosh: 101,
+    arcsinh: 102,
+    arctanh: 103,
+    # round/0 — Backend hard-codes mx::round's decimals to 0
+    # (Nx.round/1 takes no decimals arg). Dispatcher does the same.
+    round: 104,
+    bitwise_invert: 105,
+    isnan: 106,
+    isinf: 107,
+    conjugate: 108,
+    real: 109,
+    imag: 110
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -287,23 +311,44 @@ defmodule Emily.IR do
   }
 
   # Unary elementwise: no coercion (MLX preserves the dtype Nx expects).
+  # The post-emit `coerce/3` then astype-casts to out.type so MLX ops
+  # whose dtype rule differs from Nx (e.g. `is_nan`/`is_infinity` returning
+  # a bool that Nx wants as {:u, 8}, or `real`/`imag` whose Nx out.type is
+  # the real component) line up — same machinery as the original 16 ops.
   @unary_ops %{
     negate: :negative,
     abs: :abs,
     sign: :sign,
     sqrt: :sqrt,
     rsqrt: :rsqrt,
     exp: :exp,
+    expm1: :expm1,
     log: :log,
     log1p: :log1p,
     sin: :sin,
     cos: :cos,
+    tan: :tan,
     tanh: :tanh,
+    sinh: :sinh,
+    cosh: :cosh,
+    acos: :arccos,
+    asin: :arcsin,
+    atan: :arctan,
+    acosh: :arccosh,
+    asinh: :arcsinh,
+    atanh: :arctanh,
     sigmoid: :sigmoid,
     floor: :floor,
     ceil: :ceil,
+    round: :round,
     erf: :erf,
-    erf_inv: :erf_inv
+    erf_inv: :erf_inv,
+    bitwise_not: :bitwise_invert,
+    is_nan: :isnan,
+    is_infinity: :isinf,
+    conjugate: :conjugate,
+    real: :real,
+    imag: :imag
   }
 
   @doc """
@@ -405,6 +450,30 @@ defmodule Emily.IR do
     emit(state, :astype, [ra], [[dtype_code(t.type)]])
   end
 
+  # erfc(x) := 1 - erf(x). Mirrors Emily.Backend.erfc/2 — MLX has no
+  # erfc primitive, so the eager path also composes from erf + subtract.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :erfc, args: [a]}} = t, state) do
+    {ra, state} = lower_node(a, state)
+    {erf_r, state} = emit(state, :erf, [ra])
+    {one_ref, state} = scalar_const(1.0, t.type, state)
+    {r, state} = emit(state, :subtract, [one_ref, erf_r])
+    coerce(r, t.type, state)
+  end
+
+  # cbrt(x) := sign(x) * abs(x)^(1/3). Mirrors Emily.Backend.cbrt/2 —
+  # MLX has no cbrt primitive. Splitting via sign+abs keeps the negative
+  # branch correct (`x^(1/3)` over negatives lands in complex), matching
+  # the eager path's bit pattern.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :cbrt, args: [a]}} = t, state) do
+    {ra, state} = lower_node(a, state)
+    {sign_r, state} = emit(state, :sign, [ra])
+    {abs_r, state} = emit(state, :abs, [ra])
+    {third, state} = scalar_const(1.0 / 3.0, t.type, state)
+    {pow_r, state} = emit(state, :power, [abs_r, third])
+    {r, state} = emit(state, :multiply, [sign_r, pow_r])
+    coerce(r, t.type, state)
+  end
+
   # bitcast: reinterpret the bytes as out.type (mirrors Emily.Backend.bitcast/2,
   # which calls mx::view). Used by the RNG path to turn random bits into floats.
   defp lower_op(%T{data: %Nx.Defn.Expr{op: :bitcast, args: [a]}} = t, state) do
@@ -1343,6 +1412,13 @@ defmodule Emily.IR do
     {{:const, idx}, %{state | consts: [ref | state.consts], n_consts: idx + 1}}
   end
 
+  # Bake a `{}` scalar of `type` as a captured const operand. Used by the
+  # composite lowerers (erfc, cbrt) whose Backend mirrors build the same
+  # scalar through `scalar_ref/2`.
+  defp scalar_const(value, type, state) do
+    materialize_const(Nx.tensor(value, type: type, backend: Nx.BinaryBackend), {}, type, state)
+  end
+
   defp materialize_capture(tensor, shape, type, state) do
     ref = Emily.Native.from_binary(Nx.to_binary(tensor), Tuple.to_list(shape), type)
     idx = state.n_captures