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 @@ -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
Expand Down
15 changes: 14 additions & 1 deletion c_src/emily/opcodes.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -752,6 +758,13 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &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: {
Expand Down
28 changes: 27 additions & 1 deletion lib/emily/ir.ex
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -281,15 +287,20 @@ 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,
multiply: :multiply,
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,
Expand Down Expand Up @@ -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)
Expand Down
45 changes: 45 additions & 0 deletions test/emily/compiler_equivalence_test.exs
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand All @@ -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
Expand Down