Skip to content

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

Merged
ausimian merged 2 commits into
feat/expr-compilerfrom
feat/expr-compiler-unary-ops
Jun 6, 2026
Merged

feat: lower 19 more unary ops in the native Expr compiler#189
ausimian merged 2 commits into
feat/expr-compilerfrom
feat/expr-compiler-unary-ops

Conversation

@ausimian

@ausimian ausimian commented Jun 6, 2026

Copy link
Copy Markdown
Owner

Summary

Closes the largest cluster on the Expr op-coverage checklist (#188): 19 unary ops that the IR didn't lower yet but Emily.Backend already implements eagerly. Each was a silent-fallback cliff in user code that used trig / numeric-stability / bitwise / complex-component / NaN-detection ops.

  • Direct-mapped (17 ops, new IR opcodes 94–110, bit-identical to the Evaluator)expm1, tan, sinh, cosh, acos, asin, atan, acosh, asinh, atanh, round (decimals=0, matching Backend), bitwise_not, is_nan, is_infinity, conjugate, real, imag. Each routes to the same mx::* primitive as the eager unary NIF (c_src/ops/unary.cpp), so the native single-NIF path and the Evaluator land on identical bits.
  • Composed (2 ops, no new opcode)erfc = 1 - erf and cbrt = sign(x) * abs(x)^(1/3). MLX has no primitive for either; the eager Emily.Backend composes them from existing ops and the IR lowerer now does the same, so the bit pattern stays exact too.

After this PR the op-coverage probe (scripts/expr_op_coverage.exs) drops from 37 → 18 misses; the only remaining unary ops are count_leading_zeros / population_count, which are already documented as "no MLX primitive — unlowerable by design" in #188.

Test plan

  • mix precommit clean (701 tests, 79 properties, 40 doctests, 0 failures).
  • mix run scripts/expr_op_coverage.exs confirms all 19 ops now lower under native_fallback: :raise (down from 37 misses to 18, all 19 targeted ops moved from MISS → OK).
  • New equivalence tests in test/emily/compiler_equivalence_test.exs cover:
    • direct-mapped float trig (expm1, tan, sinh, cosh, atan, asinh);
    • inverse-trig with domain-valid inputs (asin/acos/atanh over |x| < 1, acosh over x >= 1);
    • round over tie-break inputs (-1.5, -0.5, 0.5, 1.5, 2.3, -2.7);
    • composed erfc and cbrt over positive, negative, and zero inputs;
    • bitwise_not over s32/s64/u8/u32;
    • is_nan / is_infinity over a tensor mixing finite values, :infinity, :neg_infinity, and :nan (asserting the {:u, 8} predicate dtype after the trailing coerce);
    • conjugate/real/imag over a {:c, 64} input (asserting the {:f, 32} real-component out.type).
  • Opcode parity test (test/emily/opcode_parity_test.exs) green — kOpcodeCount bumped 94 → 111 in lockstep with the Elixir @opcodes map and the C++ enum.
  • Conformance lanes (:native, :fuse) — out of scope, no new lowerings expected to reach them.

ausimian added 2 commits June 6, 2026 20:29
Extends the IR's `@unary_ops` map (lib/emily/ir.ex) and the C++ Opcode
enum + dispatcher with the missing peers of the original 16-op unary
set, closing the largest cluster on the Expr op-coverage checklist
(#188).

Direct-mapped (route to the same `mx::*` primitive as the eager unary
NIF in c_src/ops/unary.cpp, so bit-identical to the Evaluator):
expm1, tan, sinh, cosh, acos, asin, atan, acosh, asinh, atanh, round
(decimals=0, matching Backend), bitwise_not, is_nan, is_infinity,
conjugate, real, imag.

Composed in the lowerer (no MLX primitive, same composition the eager
Backend uses): erfc = 1 - erf; cbrt = sign(x) * abs(x)^(1/3).

Equivalence tests added to compiler_equivalence_test.exs cover the
direct-mapped float trig + inverse-trig (with domain-valid inputs), the
round tie-break, the composed erfc/cbrt, bitwise_not over integer
dtypes (s32/s64/u8/u32), is_nan/is_infinity across finite + nan + inf
+/- (asserting the {:u, 8} predicate dtype after coerce), and the
complex unary trio over a {:c, 64} input (asserting the {:f, 32}
out.type of real/imag).
scripts/expr_op_coverage.exs JIT-runs a tiny defn over each Nx op via
`Emily.Compiler, native: true, native_fallback: :raise` and reports
which ones don't lower yet, printing both a per-op log and a markdown
checklist suitable for pasting into the op-coverage tracking issue
(#188). Re-run after each lowering PR to keep the checklist honest:

    mix run scripts/expr_op_coverage.exs

The probe walks the public Nx API (not internal Expr nodes), so
composite ops that Nx rewrites before reaching the backend are
correctly reported via their top-level name.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant