Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions RELEASE.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
59 changes: 58 additions & 1 deletion c_src/emily/opcodes.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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: {
Expand Down
80 changes: 78 additions & 2 deletions lib/emily/ir.ex
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 """
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
Loading