diff --git a/RELEASE.md b/RELEASE.md index 9f37fec..5c0b286 100644 --- a/RELEASE.md +++ b/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 diff --git a/c_src/emily/opcodes.hpp b/c_src/emily/opcodes.hpp index cf13f93..2095bbe 100644 --- a/c_src/emily/opcodes.hpp +++ b/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 &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: { diff --git a/lib/emily/ir.ex b/lib/emily/ir.ex index 8bebfeb..6b6fbdc 100644 --- a/lib/emily/ir.ex +++ b/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,6 +311,10 @@ 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, @@ -294,16 +322,33 @@ defmodule Emily.IR do 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 diff --git a/scripts/expr_op_coverage.exs b/scripts/expr_op_coverage.exs new file mode 100644 index 0000000..9be5c2c --- /dev/null +++ b/scripts/expr_op_coverage.exs @@ -0,0 +1,311 @@ +# scripts/expr_op_coverage.exs +# +# Probe every Nx op through `Emily.Compiler, native: true, +# native_fallback: :raise` and report which ones don't lower yet. +# +# Each probe builds a tiny defn over a representative input. An op the +# IR lowers returns a tensor; an op it doesn't raises `ArgumentError`, +# whose message names what's missing ("does not yet lower op :foo" / +# "does not yet lower the block Foo"). Composite Nx ops that decompose +# to lower-level primitives (e.g. `Nx.square` -> `multiply`) are still +# reported by their top-level name — they pass when every primitive in +# their expansion lowers. +# +# Run: +# mix run scripts/expr_op_coverage.exs +# +# Output: +# - per-op line, sorted into OK / MISS / ERROR / UNSUPPORTED groups +# - a final markdown checklist suitable for pasting into a GitHub issue + +defmodule ExprOpCoverage do + @opts [compiler: Emily.Compiler, native: true, native_fallback: :raise] + + def probe(name, fun, args) do + jit = Nx.Defn.jit(fun, @opts) + _ = apply(jit, args) + {:ok, name} + rescue + e in ArgumentError -> + msg = Exception.message(e) + + cond do + msg =~ "does not yet lower op" -> {:miss, name, extract_op(msg)} + msg =~ "does not yet lower the block" -> {:miss, name, extract_block(msg)} + msg =~ "cannot lower" -> {:unsup, name, msg |> first_line()} + msg =~ "is not supported" -> {:miss, name, msg |> first_line()} + msg =~ "no fallback" -> {:miss, name, msg |> first_line()} + true -> {:error, name, msg |> first_line()} + end + + e -> + {:error, name, Exception.message(e) |> first_line()} + end + + defp extract_op(msg) do + case Regex.run(~r/does not yet lower op (:[a-z_0-9]+)/, msg) do + [_, op] -> "op #{op}" + _ -> first_line(msg) + end + end + + defp extract_block(msg) do + case Regex.run(~r/does not yet lower the block ([A-Za-z0-9_.]+)/, msg) do + [_, mod] -> "block #{mod}" + _ -> first_line(msg) + end + end + + defp first_line(msg), do: msg |> String.split("\n") |> hd() |> String.slice(0, 120) +end + +# ---------- probe definitions ---------- + +t = fn list, opts -> Nx.tensor(list, opts) end +f = fn list -> Nx.tensor(list) end +s32 = fn list -> Nx.tensor(list, type: :s32) end +u8 = fn list -> Nx.tensor(list, type: :u8) end +c = fn list -> Nx.tensor(Enum.map(list, &Complex.new(&1, 0.0))) end + +probes = + [ + # ============ Unary elementwise ============ + {:unary, :exp, fn x -> Nx.exp(x) end, [f.([1.0])]}, + {:unary, :expm1, fn x -> Nx.expm1(x) end, [f.([1.0])]}, + {:unary, :log, fn x -> Nx.log(x) end, [f.([1.0])]}, + {:unary, :log1p, fn x -> Nx.log1p(x) end, [f.([1.0])]}, + {:unary, :sigmoid, fn x -> Nx.sigmoid(x) end, [f.([1.0])]}, + {:unary, :cos, fn x -> Nx.cos(x) end, [f.([1.0])]}, + {:unary, :sin, fn x -> Nx.sin(x) end, [f.([1.0])]}, + {:unary, :tan, fn x -> Nx.tan(x) end, [f.([1.0])]}, + {:unary, :cosh, fn x -> Nx.cosh(x) end, [f.([1.0])]}, + {:unary, :sinh, fn x -> Nx.sinh(x) end, [f.([1.0])]}, + {:unary, :tanh, fn x -> Nx.tanh(x) end, [f.([1.0])]}, + {:unary, :acosh, fn x -> Nx.acosh(x) end, [f.([1.5])]}, + {:unary, :asinh, fn x -> Nx.asinh(x) end, [f.([1.0])]}, + {:unary, :atanh, fn x -> Nx.atanh(x) end, [f.([0.5])]}, + {:unary, :acos, fn x -> Nx.acos(x) end, [f.([0.5])]}, + {:unary, :asin, fn x -> Nx.asin(x) end, [f.([0.5])]}, + {:unary, :atan, fn x -> Nx.atan(x) end, [f.([0.5])]}, + {:unary, :sqrt, fn x -> Nx.sqrt(x) end, [f.([4.0])]}, + {:unary, :rsqrt, fn x -> Nx.rsqrt(x) end, [f.([4.0])]}, + {:unary, :cbrt, fn x -> Nx.cbrt(x) end, [f.([8.0])]}, + {:unary, :negate, fn x -> Nx.negate(x) end, [f.([1.0])]}, + {:unary, :sign, fn x -> Nx.sign(x) end, [f.([1.0])]}, + {:unary, :abs, fn x -> Nx.abs(x) end, [f.([-1.0])]}, + {:unary, :bitwise_not, fn x -> Nx.bitwise_not(x) end, [s32.([1])]}, + {:unary, :is_nan, fn x -> Nx.is_nan(x) end, [f.([1.0])]}, + {:unary, :is_infinity, fn x -> Nx.is_infinity(x) end, [f.([1.0])]}, + {:unary, :conjugate, fn x -> Nx.conjugate(x) end, [c.([1.0])]}, + {:unary, :real, fn x -> Nx.real(x) end, [c.([1.0])]}, + {:unary, :imag, fn x -> Nx.imag(x) end, [c.([1.0])]}, + {:unary, :floor, fn x -> Nx.floor(x) end, [f.([1.5])]}, + {:unary, :ceil, fn x -> Nx.ceil(x) end, [f.([1.5])]}, + {:unary, :round, fn x -> Nx.round(x) end, [f.([1.5])]}, + {:unary, :erf, fn x -> Nx.erf(x) end, [f.([1.0])]}, + {:unary, :erfc, fn x -> Nx.erfc(x) end, [f.([1.0])]}, + {:unary, :erf_inv, fn x -> Nx.erf_inv(x) end, [f.([0.5])]}, + {:unary, :bitcast, fn x -> Nx.bitcast(x, :s32) end, [f.([1.0])]}, + {:unary, :population_count, fn x -> Nx.population_count(x) end, [s32.([1])]}, + {:unary, :count_leading_zeros, fn x -> Nx.count_leading_zeros(x) end, [s32.([1])]}, + + # ============ Binary arithmetic / bitwise ============ + {:binary, :add, fn a, b -> Nx.add(a, b) end, [f.([1.0]), f.([1.0])]}, + {:binary, :subtract, fn a, b -> Nx.subtract(a, b) end, [f.([1.0]), f.([1.0])]}, + {:binary, :multiply, fn a, b -> Nx.multiply(a, b) end, [f.([1.0]), f.([1.0])]}, + {:binary, :divide, fn a, b -> Nx.divide(a, b) end, [f.([1.0]), f.([1.0])]}, + {:binary, :pow, fn a, b -> Nx.pow(a, b) end, [f.([2.0]), f.([3.0])]}, + {:binary, :remainder, fn a, b -> Nx.remainder(a, b) end, [f.([5.0]), f.([3.0])]}, + {:binary, :atan2, fn a, b -> Nx.atan2(a, b) end, [f.([1.0]), f.([1.0])]}, + {:binary, :max, fn a, b -> Nx.max(a, b) end, [f.([1.0]), f.([2.0])]}, + {:binary, :min, fn a, b -> Nx.min(a, b) end, [f.([1.0]), f.([2.0])]}, + {:binary, :quotient, fn a, b -> Nx.quotient(a, b) end, [s32.([5]), s32.([2])]}, + {:binary, :bitwise_and, fn a, b -> Nx.bitwise_and(a, b) end, [s32.([1]), s32.([1])]}, + {:binary, :bitwise_or, fn a, b -> Nx.bitwise_or(a, b) end, [s32.([1]), s32.([1])]}, + {:binary, :bitwise_xor, fn a, b -> Nx.bitwise_xor(a, b) end, [s32.([1]), s32.([1])]}, + {:binary, :left_shift, fn a, b -> Nx.left_shift(a, b) end, [s32.([1]), s32.([2])]}, + {:binary, :right_shift, fn a, b -> Nx.right_shift(a, b) end, [s32.([4]), s32.([1])]}, + + # ============ Compare / logical ============ + {:compare, :equal, fn a, b -> Nx.equal(a, b) end, [f.([1.0]), f.([1.0])]}, + {:compare, :not_equal, fn a, b -> Nx.not_equal(a, b) end, [f.([1.0]), f.([1.0])]}, + {:compare, :less, fn a, b -> Nx.less(a, b) end, [f.([1.0]), f.([2.0])]}, + {:compare, :less_equal, fn a, b -> Nx.less_equal(a, b) end, [f.([1.0]), f.([2.0])]}, + {:compare, :greater, fn a, b -> Nx.greater(a, b) end, [f.([1.0]), f.([2.0])]}, + {:compare, :greater_equal, fn a, b -> Nx.greater_equal(a, b) end, [f.([1.0]), f.([2.0])]}, + {:compare, :logical_and, fn a, b -> Nx.logical_and(a, b) end, [u8.([1]), u8.([0])]}, + {:compare, :logical_or, fn a, b -> Nx.logical_or(a, b) end, [u8.([1]), u8.([0])]}, + {:compare, :logical_xor, fn a, b -> Nx.logical_xor(a, b) end, [u8.([1]), u8.([0])]}, + {:compare, :logical_not, fn x -> Nx.logical_not(x) end, [u8.([0])]}, + + # ============ Reductions ============ + {:reduce, :sum, fn x -> Nx.sum(x) end, [f.([1.0, 2.0])]}, + {:reduce, :product, fn x -> Nx.product(x) end, [f.([1.0, 2.0])]}, + {:reduce, :all, fn x -> Nx.all(x) end, [u8.([1])]}, + {:reduce, :any, fn x -> Nx.any(x) end, [u8.([1])]}, + {:reduce, :reduce_max, fn x -> Nx.reduce_max(x) end, [f.([1.0, 2.0])]}, + {:reduce, :reduce_min, fn x -> Nx.reduce_min(x) end, [f.([1.0, 2.0])]}, + {:reduce, :argmax, fn x -> Nx.argmax(x) end, [f.([1.0, 2.0])]}, + {:reduce, :argmin, fn x -> Nx.argmin(x) end, [f.([1.0, 2.0])]}, + + # ============ Shape ops ============ + {:shape, :reshape, fn x -> Nx.reshape(x, {2, 1}) end, [f.([1.0, 2.0])]}, + {:shape, :squeeze, fn x -> Nx.squeeze(x, axes: [0]) end, [f.([[1.0]])]}, + {:shape, :transpose, fn x -> Nx.transpose(x) end, [f.([[1.0, 2.0]])]}, + {:shape, :as_type, fn x -> Nx.as_type(x, :s32) end, [f.([1.0])]}, + {:shape, :broadcast, fn x -> Nx.broadcast(x, {2, 2}) end, [f.([1.0, 2.0])]}, + {:shape, :pad, fn x -> Nx.pad(x, 0.0, [{1, 1, 0}]) end, [f.([1.0, 2.0])]}, + {:shape, :reverse, fn x -> Nx.reverse(x) end, [f.([1.0, 2.0])]}, + {:shape, :concatenate, fn a, b -> Nx.concatenate([a, b]) end, [f.([1.0]), f.([2.0])]}, + {:shape, :stack, fn a, b -> Nx.stack([a, b]) end, [f.([1.0]), f.([2.0])]}, + + # ============ Linalg core ============ + {:linalg, :dot, fn a, b -> Nx.dot(a, b) end, [f.([1.0, 2.0]), f.([1.0, 2.0])]}, + {:linalg, :conv, fn x, k -> Nx.conv(x, k, strides: 1, padding: :valid) end, + [f.([[[1.0, 2.0, 3.0]]]), f.([[[1.0]]])]}, + + # ============ Selection / indexing ============ + {:select, :select, fn p, t, fa -> Nx.select(p, t, fa) end, + [u8.([1]), f.([1.0]), f.([2.0])]}, + {:select, :clip, fn x -> Nx.clip(x, Nx.tensor(0.0), Nx.tensor(1.0)) end, [f.([0.5])]}, + {:select, :slice, fn x -> Nx.slice(x, [0], [1]) end, [f.([1.0, 2.0])]}, + {:select, :put_slice, fn x, u -> Nx.put_slice(x, [0], u) end, [f.([1.0, 2.0]), f.([3.0])]}, + {:select, :gather, fn x, i -> Nx.gather(x, i) end, [f.([1.0, 2.0]), s32.([[0], [1]])]}, + {:select, :indexed_put, fn x, i, u -> Nx.indexed_put(x, i, u) end, + [f.([1.0, 2.0]), s32.([[0]]), f.([3.0])]}, + {:select, :indexed_add, fn x, i, u -> Nx.indexed_add(x, i, u) end, + [f.([1.0, 2.0]), s32.([[0]]), f.([3.0])]}, + {:select, :take, fn x, i -> Nx.take(x, i) end, [f.([1.0, 2.0]), s32.([0])]}, + {:select, :take_along_axis, fn x, i -> Nx.take_along_axis(x, i) end, + [f.([1.0, 2.0]), s32.([0])]}, + + # ============ Sort / argsort / top_k ============ + {:sort, :sort, fn x -> Nx.sort(x) end, [f.([2.0, 1.0])]}, + {:sort, :argsort, fn x -> Nx.argsort(x) end, [f.([2.0, 1.0])]}, + {:sort, :top_k, fn x -> Nx.top_k(x, k: 1) end, [f.([2.0, 1.0])]}, + + # ============ Creation ============ + {:create, :iota, fn -> Nx.iota({4}) end, []}, + {:create, :eye, fn -> Nx.eye(2) end, []}, + + # ============ FFT family ============ + {:fft, :fft, fn x -> Nx.fft(x, length: 4) end, [f.([1.0, 0.0, 0.0, 0.0])]}, + {:fft, :ifft, fn x -> Nx.ifft(x, length: 4) end, + [Nx.tensor([Complex.new(1.0, 0.0), Complex.new(0.0, 0.0), Complex.new(0.0, 0.0), Complex.new(0.0, 0.0)])]}, + {:fft, :fft2, fn x -> Nx.fft2(x) end, [Nx.iota({2, 2}, type: :f32)]}, + {:fft, :ifft2, fn x -> Nx.ifft2(x) end, + [Nx.tensor([[Complex.new(1.0, 0.0), Complex.new(0.0, 0.0)], [Complex.new(0.0, 0.0), Complex.new(0.0, 0.0)]])]}, + {:fft, :rfft, fn x -> Nx.rfft(x, length: 4) end, [f.([1.0, 0.0, 0.0, 0.0])]}, + {:fft, :irfft, fn x -> Nx.irfft(x, length: 4) end, + [Nx.tensor([Complex.new(1.0, 0.0), Complex.new(0.0, 0.0), Complex.new(0.0, 0.0)])]}, + + # ============ Window (pooling) reductions ============ + {:window, :window_sum, fn x -> Nx.window_sum(x, {2}) end, [f.([1.0, 2.0, 3.0])]}, + {:window, :window_max, fn x -> Nx.window_max(x, {2}) end, [f.([1.0, 2.0, 3.0])]}, + {:window, :window_min, fn x -> Nx.window_min(x, {2}) end, [f.([1.0, 2.0, 3.0])]}, + {:window, :window_product, fn x -> Nx.window_product(x, {2}) end, [f.([1.0, 2.0, 3.0])]}, + {:window, :window_mean, fn x -> Nx.window_mean(x, {2}) end, [f.([1.0, 2.0, 3.0])]}, + + # ============ Cumulative ============ + {:cumulative, :cumulative_sum, fn x -> Nx.cumulative_sum(x) end, [f.([1.0, 2.0])]}, + {:cumulative, :cumulative_product, fn x -> Nx.cumulative_product(x) end, [f.([1.0, 2.0])]}, + {:cumulative, :cumulative_max, fn x -> Nx.cumulative_max(x) end, [f.([1.0, 2.0])]}, + {:cumulative, :cumulative_min, fn x -> Nx.cumulative_min(x) end, [f.([1.0, 2.0])]}, + # interior-axis variants — the IR last-axis fast-path doesn't apply + {:cumulative, :cumulative_sum_axis0, fn x -> Nx.cumulative_sum(x, axis: 0) end, + [f.([[1.0, 2.0], [3.0, 4.0]])]}, + + # ============ Linalg blocks ============ + {:linalg_block, :cholesky, fn x -> Nx.LinAlg.cholesky(x) end, + [f.([[4.0, 2.0], [2.0, 3.0]])]}, + {:linalg_block, :svd, fn x -> Nx.LinAlg.svd(x) end, [f.([[1.0, 0.0], [0.0, 1.0]])]}, + {:linalg_block, :qr, fn x -> Nx.LinAlg.qr(x) end, [f.([[1.0, 0.0], [0.0, 1.0]])]}, + {:linalg_block, :eigh, fn x -> Nx.LinAlg.eigh(x) end, [f.([[2.0, 0.0], [0.0, 1.0]])]}, + {:linalg_block, :lu, fn x -> Nx.LinAlg.lu(x) end, [f.([[1.0, 0.0], [0.0, 1.0]])]}, + {:linalg_block, :determinant, fn x -> Nx.LinAlg.determinant(x) end, + [f.([[1.0, 0.0], [0.0, 1.0]])]}, + {:linalg_block, :solve, fn a, b -> Nx.LinAlg.solve(a, b) end, + [f.([[1.0, 0.0], [0.0, 1.0]]), f.([1.0, 1.0])]}, + {:linalg_block, :triangular_solve, fn a, b -> Nx.LinAlg.triangular_solve(a, b) end, + [f.([[1.0, 0.0], [0.0, 1.0]]), f.([1.0, 1.0])]}, + + # ============ Other Nx.Block ============ + {:block, :all_close, fn a, b -> Nx.all_close(a, b) end, [f.([1.0]), f.([1.0])]}, + {:block, :phase, fn x -> Nx.phase(x) end, [c.([1.0])]} + ] + +# ---------- run ---------- + +# Use the eager Emily.Backend so input materialization works. +Nx.global_default_backend(Emily.Backend) + +results = + Enum.map(probes, fn {cat, name, fun, args} -> + {cat, name, ExprOpCoverage.probe({cat, name}, fun, args)} + end) + +# ---------- print plain log ---------- + +IO.puts("\n========= per-op result (#{length(results)} probes) =========\n") + +grouped = + Enum.reduce(results, %{ok: [], miss: [], unsup: [], error: []}, fn {cat, name, res}, acc -> + tag = + case res do + {:ok, _} -> :ok + {:miss, _, _} -> :miss + {:unsup, _, _} -> :unsup + {:error, _, _} -> :error + end + + Map.update!(acc, tag, &[{cat, name, res} | &1]) + end) + |> Map.new(fn {k, v} -> {k, Enum.reverse(v)} end) + +for {tag, list} <- [{:ok, "OK"}, {:miss, "MISS"}, {:unsup, "UNSUPPORTED"}, {:error, "ERROR"}] do + rows = Map.fetch!(grouped, tag) + IO.puts("--- #{list} (#{length(rows)}) ---") + + for {cat, name, res} <- rows do + detail = + case res do + {:ok, _} -> "" + {_, _, why} -> " — #{why}" + end + + IO.puts(" [#{cat}] #{name}#{detail}") + end + + IO.puts("") +end + +# ---------- print markdown checklist ---------- + +IO.puts("\n========= markdown checklist (paste into the issue) =========\n") + +miss_by_cat = + grouped.miss + |> Enum.group_by(fn {cat, _, _} -> cat end) + |> Enum.sort() + +for {cat, rows} <- miss_by_cat do + IO.puts("### #{cat} (#{length(rows)})") + + for {_, name, {_, _, why}} <- Enum.sort_by(rows, fn {_, n, _} -> n end) do + IO.puts("- [ ] `#{name}` — #{why}") + end + + IO.puts("") +end + +unsup_rows = grouped.unsup + +if unsup_rows != [] do + IO.puts("### Unlowerable by design (#{length(unsup_rows)})") + + for {_, name, {_, _, why}} <- Enum.sort_by(unsup_rows, fn {_, n, _} -> n end) do + IO.puts("- [ ] `#{name}` — #{why}") + end + + IO.puts("") +end diff --git a/test/emily/compiler_equivalence_test.exs b/test/emily/compiler_equivalence_test.exs index b9ccf91..5e24da1 100644 --- a/test/emily/compiler_equivalence_test.exs +++ b/test/emily/compiler_equivalence_test.exs @@ -68,6 +68,80 @@ defmodule Emily.CompilerEquivalenceTest do assert_equiv(op, [x]) end end + + # Direct-mapped unary ops added alongside the Expr op-coverage sweep + # (see #188). Each routes through the same `mx::*` primitive as the + # eager unary NIF, so the native single-NIF path is bit-identical to + # the Evaluator (which dispatches the same op via Emily.Backend). + test "extra direct-mapped float unary ops match the evaluator" do + x = et([0.5, -1.25, 2.0, 0.1]) + + for op <- [ + &Nx.expm1/1, + &Nx.tan/1, + &Nx.sinh/1, + &Nx.cosh/1, + &Nx.atan/1, + &Nx.asinh/1 + ] do + assert_equiv(op, [x]) + end + end + + test "inverse-trig unary ops match the evaluator (domain-restricted inputs)" do + # asin/acos: |x| <= 1; atanh: |x| < 1; acosh: x >= 1. + x = et([0.1, -0.5, 0.25, -0.75]) + + for op <- [&Nx.asin/1, &Nx.acos/1, &Nx.atanh/1] do + assert_equiv(op, [x]) + end + + assert_equiv(&Nx.acosh/1, [et([1.0, 1.5, 2.0, 3.0])]) + end + + test "round (away from zero) matches the evaluator on tie-breaking inputs" do + # Nx.round/1 documents "round away from zero"; MLX's mx::round + # decimals=0 follows the same. The Backend hard-codes decimals=0 + # and the IR dispatcher does too. + assert_equiv(&Nx.round/1, [et([-1.5, -0.5, 0.5, 1.5, 2.3, -2.7])]) + end + + test "composed erfc and cbrt match the evaluator (no MLX primitive)" do + # Both compose from existing ops (erfc: 1 - erf; cbrt: sign * abs^(1/3)) + # the same way Emily.Backend does, so the bit pattern stays exact. + assert_equiv(&Nx.erfc/1, [et([0.0, 0.5, 1.0, -0.5, 2.0])]) + assert_equiv(&Nx.cbrt/1, [et([8.0, -8.0, 0.125, -0.125, 0.0])]) + end + + test "bitwise_not on integer inputs matches the evaluator" do + for type <- [:s32, :s64, :u8, :u32] do + assert_equiv(&Nx.bitwise_not/1, [et([0, 1, 2, 255], type: type)]) + end + end + + test "is_nan / is_infinity on float inputs match the evaluator" do + # nan/inf live alongside finite values so the predicate output (Nx pred == + # {:u, 8}) exercises all three branches; the trailing astype coerces + # MLX's bool to {:u, 8} like every other unary op. + x = et([1.0, :infinity, :neg_infinity, :nan, 0.0]) + + out = assert_equiv(&Nx.is_nan/1, [x]) + assert out.type == {:u, 8} + assert_equiv(&Nx.is_infinity/1, [x]) + end + + test "complex unary ops (conjugate / real / imag) match the evaluator" do + x = et([Complex.new(1.0, 2.0), Complex.new(-3.0, 4.0), Complex.new(0.0, -1.0)]) + + conj = assert_equiv(&Nx.conjugate/1, [x]) + assert conj.type == {:c, 64} + + re = assert_equiv(&Nx.real/1, [x]) + assert re.type == {:f, 32} + + im = assert_equiv(&Nx.imag/1, [x]) + assert im.type == {:f, 32} + end end describe "binary arithmetic" do