diff --git a/RELEASE.md b/RELEASE.md index 5c0b286..b7423bd 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -58,6 +58,14 @@ `Emily.Backend`'s eager composition. Closes the largest cluster on the Expr op-coverage checklist (#188). +- **`atan2`, `quotient`, `logical_xor` lower natively** — closes the + binary/compare cluster on the Expr op-coverage checklist (#188). + `atan2` slots into the @arith_binary mapping (cast both to out.type, + then `mx::arctan2`); `quotient` routes through `mx::floor_divide` the + same way `Emily.Backend.quotient/3` does; `logical_xor` (no MLX + primitive) lowers to `(a != 0) != (b != 0)`, mirroring the eager + Backend composition. All three are bit-identical to the Evaluator. + - **`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 2095bbe..40ae73e 100644 --- a/c_src/emily/opcodes.hpp +++ b/c_src/emily/opcodes.hpp @@ -201,9 +201,15 @@ enum class Opcode : int64_t { Conjugate = 108, Real = 109, Imag = 110, + // Binary arithmetic peers of the @arith_binary cluster. arctan2 is the + // direct Backend mapping (atan2: arctan2); floor_divide is the integer + // engine behind Nx.quotient (the lowerer routes quotient -> floor_divide + // matching Emily.Backend.quotient/3). + Arctan2 = 111, + FloorDivide = 112, }; -inline constexpr int64_t kOpcodeCount = 111; +inline constexpr int64_t kOpcodeCount = 113; // Quant mode code (Emily.IR @quant_modes) -> MLX mode string. inline std::string qmode_from_code(int64_t code) { @@ -752,6 +758,13 @@ inline mx::array dispatch_op(Opcode op, const std::vector &in, return mx::real(arg1(in, "real"), s); case Opcode::Imag: return mx::imag(arg1(in, "imag"), s); + // --- Binary arithmetic peers --- + case Opcode::Arctan2: + need2(in, "arctan2"); + return mx::arctan2(in[0], in[1], s); + case Opcode::FloorDivide: + need2(in, "floor_divide"); + return mx::floor_divide(in[0], in[1], 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 6b6fbdc..8608388 100644 --- a/lib/emily/ir.ex +++ b/lib/emily/ir.ex @@ -182,7 +182,13 @@ defmodule Emily.IR do isinf: 107, conjugate: 108, real: 109, - imag: 110 + imag: 110, + # Binary arithmetic peers. arctan2 is the direct atan2 lowering; + # floor_divide is the integer engine Emily.Backend.quotient/3 calls, + # so the IR routes Nx :quotient through it (cast both to out.type, + # floor_divide, same bits as the eager Backend). + arctan2: 111, + floor_divide: 112 } # Quant mode string -> code; decoded by qmode_from_code in @@ -281,6 +287,9 @@ defmodule Emily.IR do # Nx Expr op -> IR opcode. Arithmetic/bitwise cast both operands to the # node's out.type before the op (see backend.ex @renamed_arith_binary). + # Nx :quotient routes here too — Emily.Backend.quotient/3 is + # `floor_divide(astype(a, out.type), astype(b, out.type))`, which is + # exactly this clause with the floor_divide opcode. @arith_binary %{ add: :add, subtract: :subtract, @@ -288,8 +297,10 @@ defmodule Emily.IR do divide: :divide, pow: :power, remainder: :remainder, + atan2: :arctan2, min: :minimum, max: :maximum, + quotient: :floor_divide, bitwise_and: :bitwise_and, bitwise_or: :bitwise_or, bitwise_xor: :bitwise_xor, @@ -434,6 +445,21 @@ defmodule Emily.IR do emit(state, :astype, [r], [[dtype_code(t.type)]]) end + # logical_xor(a, b) := (a != 0) != (b != 0). MLX has no logical_xor + # primitive; mirrors Emily.Backend.logical_xor/3 — three not_equal + # calls (each operand vs a per-dtype zero, then the two booleans), with + # the final coerce to out.type ({:u, 8}) matching the eager wrap. + defp lower_op(%T{data: %Nx.Defn.Expr{op: :logical_xor, args: [a, b]}} = t, state) do + {ra, state} = lower_node(a, state) + {rb, state} = lower_node(b, state) + {za, state} = scalar_const(0, a.type, state) + {zb, state} = scalar_const(0, b.type, state) + {ma, state} = emit(state, :not_equal, [ra, za]) + {mb, state} = emit(state, :not_equal, [rb, zb]) + {r, state} = emit(state, :not_equal, [ma, mb]) + coerce(r, t.type, state) + end + defp lower_op(%T{data: %Nx.Defn.Expr{op: op, args: [a]}} = t, state) when is_map_key(@unary_ops, op) do {ra, state} = lower_node(a, state) diff --git a/test/emily/compiler_equivalence_test.exs b/test/emily/compiler_equivalence_test.exs index 5e24da1..e68325e 100644 --- a/test/emily/compiler_equivalence_test.exs +++ b/test/emily/compiler_equivalence_test.exs @@ -170,6 +170,35 @@ defmodule Emily.CompilerEquivalenceTest do assert_equiv(&Nx.add/2, [a, b]) end + # Binary arithmetic peers added with the Expr op-coverage sweep (#188). + # atan2 is a direct @arith_binary mapping (atan2: :arctan2); quotient + # routes through floor_divide the same way Emily.Backend.quotient/3 does. + test "atan2 matches the evaluator across the four quadrants" do + # Quadrant-covering pairs (y, x): include the axes so the result + # lands on the ±0, ±pi, ±pi/2 boundaries the eager Backend produces. + y = et([1.0, 1.0, -1.0, -1.0, 0.0, 0.0, 2.0, -2.0]) + x = et([1.0, -1.0, -1.0, 1.0, 1.0, -1.0, 0.0, 0.0]) + assert_equiv(&Nx.atan2/2, [y, x]) + end + + test "quotient matches the evaluator for signed and unsigned integer dtypes" do + # Backend casts both operands to out.type then calls floor_divide. + # Quotient is integer-only in Nx, so we exercise the s32/s64/u8/u32 + # paths; the native lane matches the evaluator bit-for-bit because + # both end up at the same mx::floor_divide kernel. + for type <- [:s32, :s64] do + a = et([7, -7, 10, -10, 4, -4], type: type) + b = et([2, 2, 3, 3, 5, 5], type: type) + assert_equiv(&Nx.quotient/2, [a, b]) + end + + for type <- [:u8, :u32] do + a = et([7, 10, 4, 255], type: type) + b = et([2, 3, 5, 2], type: type) + assert_equiv(&Nx.quotient/2, [a, b]) + end + end + test "scalar constant operand (materialized capture)" do x = et([1.0, 2.0, 3.0]) assert_equiv(fn t -> Nx.add(t, 1.5) end, [x]) @@ -192,6 +221,22 @@ defmodule Emily.CompilerEquivalenceTest do assert out.type == {:u, 8} end end + + # logical_xor: MLX has no primitive, so both paths run + # `(a != 0) != (b != 0)` — Emily.Backend.logical_xor/3 eagerly, the + # IR via a dedicated composite clause that emits the same three + # not_equal ops. Inputs span the four truth-table corners across + # float and integer dtypes; the trailing coerce produces {:u, 8}. + test "logical_xor matches across float and integer dtypes" do + a = et([1.0, 0.0, 1.0, 0.0]) + b = et([1.0, 1.0, 0.0, 0.0]) + out = assert_equiv(&Nx.logical_xor/2, [a, b]) + assert out.type == {:u, 8} + + ai = et([1, 0, 1, 0], type: :s32) + bi = et([1, 1, 0, 0], type: :s32) + assert_equiv(&Nx.logical_xor/2, [ai, bi]) + end end describe "cast / shape" do