Merge pull request #190 from ausimian/feat/expr-compiler-binary-ops

feat: lower atan2, quotient and logical_xor in the native Expr compiler

Author: ausimian

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

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<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: {

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,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,
@@ -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)

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