Merge pull request #192 from ausimian/feat/expr-compiler-misc-blocks

feat: lower LogicalNot, AllClose and Phase blocks in the native Expr compiler

Author: ausimian

RELEASE.md

@@ -77,6 +77,17 @@
   `mx::linalg::solve_triangular` to the CPU stream per call (it's
   CPU-only), matching the eager NIF.
 
+- **`Nx.Block.LogicalNot`, `Nx.Block.AllClose`, `Nx.Block.Phase` lower
+  natively** — closes the misc-block cluster on #188. No new opcodes:
+  `LogicalNot` emits the existing `:logical_not` op directly (same as
+  `Emily.Backend.native_logical_not/2`); `AllClose` composes five
+  existing primitives — `astype` → `abs(a - b) <= atol + rtol * abs(b)`
+  → optional `isnan` OR for `equal_nan: true` → reduce-all — exactly
+  matching `Emily.Backend.native_all_close/4`; `Phase` lowers the
+  block's `atan2(imag(t), real(t))` expansion via TopK-style parameter
+  seeding (every primitive in the expansion was already on the native
+  path), 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

lib/emily/ir.ex

@@ -1354,6 +1354,80 @@ defmodule Emily.IR do
     {{:multi_refs, leaf_refs}, state}
   end
 
+  # Nx.logical_not (Nx.Block.LogicalNot). Mirrors
+  # Emily.Backend.native_logical_not/2 — emit the :logical_not opcode
+  # directly (MLX returns bool), trailing coerce produces {:u, 8}.
+  defp lower_block(%Nx.Block.LogicalNot{}, [t], _expr, out, state) do
+    {rt, state} = lower_node(t, state)
+    emit_coerced(state, :logical_not, [rt], [], out.type)
+  end
+
+  # Nx.all_close (Nx.Block.AllClose). Mirrors
+  # Emily.Backend.native_all_close/4 op-for-op so the native and eager
+  # paths land on identical bits: cast both to the merged float type,
+  # compute `abs(a - b) <= atol + rtol * abs(b)`, optionally OR with the
+  # `equal_nan` mask (`isnan(a) AND isnan(b)`), then reduce over every
+  # axis via `:all`.
+  defp lower_block(
+         %Nx.Block.AllClose{equal_nan: equal_nan, rtol: rtol, atol: atol},
+         [a, b],
+         _expr,
+         out,
+         state
+       ) do
+    {ra, state} = lower_node(a, state)
+    {rb, state} = lower_node(b, state)
+
+    merged = Nx.Type.merge(a.type, b.type) |> Nx.Type.to_floating()
+    code = dtype_code(merged)
+    {ra, state} = emit(state, :astype, [ra], [[code]])
+    {rb, state} = emit(state, :astype, [rb], [[code]])
+
+    {diff_raw, state} = emit(state, :subtract, [ra, rb])
+    {diff, state} = emit(state, :abs, [diff_raw])
+
+    {atol_ref, state} = scalar_const(atol, merged, state)
+    {rtol_ref, state} = scalar_const(rtol, merged, state)
+    {abs_b, state} = emit(state, :abs, [rb])
+    {rtol_x_abs_b, state} = emit(state, :multiply, [rtol_ref, abs_b])
+    {tol, state} = emit(state, :add, [atol_ref, rtol_x_abs_b])
+
+    {close, state} = emit(state, :less_equal, [diff, tol])
+
+    {close, state} =
+      if equal_nan do
+        {na, state} = emit(state, :isnan, [ra])
+        {nb, state} = emit(state, :isnan, [rb])
+        {both_nan, state} = emit(state, :logical_and, [na, nb])
+        emit(state, :logical_or, [close, both_nan])
+      else
+        {close, state}
+      end
+
+    axes = Enum.to_list(0..(tuple_size(a.shape) - 1)//1)
+    {result, state} = emit(state, :all, [close], [axes, [0]])
+    coerce(result, out.type, state)
+  end
+
+  # Nx.phase (Nx.Block.Phase) := atan2(imag(t), real(t)). Backend
+  # falls through to the composed expansion (no fused kernel); the IR
+  # does the same here, bound to the real in_args via the TopK-style
+  # parameter seeding (block-local :parameter nodes are FRESH and would
+  # otherwise resolve to the outer function's input slots). All three
+  # primitives in the expansion (atan2/imag/real) already lower
+  # natively, so the result is bit-identical to the Evaluator.
+  defp lower_block(%Nx.Block.Phase{}, [t], expr, _out, state) do
+    {arg_ref, state} = lower_node(t, state)
+
+    seed =
+      expr
+      |> collect_block_params(%{})
+      |> Map.new(fn {id, 0} -> {id, arg_ref} end)
+
+    state = %{state | cache: Map.merge(state.cache, seed)}
+    lower_node(expr, state)
+  end
+
   # Any other block struct raises. Lowering the block's composed
   # expansion would silently diverge from the Evaluator whenever
   # Emily.Backend.block/4 dispatches that struct through a fused / native

test/emily/compiler_equivalence_test.exs

@@ -366,6 +366,73 @@ defmodule Emily.CompilerEquivalenceTest do
     end
   end
 
+  describe "misc Nx.Block lowerings" do
+    # Nx.logical_not (Nx.Block.LogicalNot) — emits the IR :logical_not
+    # opcode directly, mirroring Emily.Backend.native_logical_not/2. The
+    # trailing coerce produces the {:u, 8} predicate dtype.
+    test "logical_not matches the evaluator and produces u8" do
+      out = assert_equiv(&Nx.logical_not/1, [et([0, 1, 2, 0, 3], type: :s32)])
+      assert out.type == {:u, 8}
+
+      assert_equiv(&Nx.logical_not/1, [et([0.0, 1.0, -1.5, 0.0])])
+    end
+
+    # Nx.all_close (Nx.Block.AllClose) — the IR composes the same five
+    # primitive sequence as Emily.Backend.native_all_close/4, so the
+    # output bit pattern is identical to the evaluator. Exercise the
+    # close, far, default-tolerance, custom-tolerance, and equal_nan
+    # branches.
+    test "all_close on close inputs returns 1 (within default tolerance)" do
+      a = et([1.0, 2.0, 3.0])
+      b = et([1.0 + 1.0e-6, 2.0 + 1.0e-6, 3.0 - 1.0e-6])
+      out = assert_equiv(&Nx.all_close/2, [a, b])
+      assert out.type == {:u, 8}
+    end
+
+    test "all_close on far inputs returns 0" do
+      a = et([1.0, 2.0, 3.0])
+      b = et([1.0, 2.0, 99.0])
+      assert_equiv(&Nx.all_close/2, [a, b])
+    end
+
+    test "all_close with custom rtol/atol matches" do
+      a = et([1.0, 2.0, 3.0])
+      b = et([1.1, 2.0, 3.0])
+      assert_equiv(fn a, b -> Nx.all_close(a, b, atol: 0.2) end, [a, b])
+      assert_equiv(fn a, b -> Nx.all_close(a, b, atol: 0.05) end, [a, b])
+    end
+
+    test "all_close with equal_nan: true matches (NaN compares equal to NaN)" do
+      a = et([1.0, :nan, 3.0])
+      b = et([1.0, :nan, 3.0])
+      assert_equiv(fn a, b -> Nx.all_close(a, b, equal_nan: true) end, [a, b])
+      assert_equiv(fn a, b -> Nx.all_close(a, b, equal_nan: false) end, [a, b])
+    end
+
+    # Nx.phase (Nx.Block.Phase) := atan2(imag(t), real(t)). Backend
+    # falls through to the composed expansion (no fused kernel); the IR
+    # lowers the same expansion via the TopK-style param-seeding, so the
+    # result is bit-identical to the evaluator. Exercise the four
+    # quadrants plus the real axis (phase == 0 for real, ±pi for
+    # negative real) and the imaginary axis (±pi/2).
+    test "phase matches the evaluator across the four complex quadrants" do
+      x =
+        et([
+          Complex.new(1.0, 1.0),
+          Complex.new(-1.0, 1.0),
+          Complex.new(-1.0, -1.0),
+          Complex.new(1.0, -1.0),
+          Complex.new(2.0, 0.0),
+          Complex.new(-2.0, 0.0),
+          Complex.new(0.0, 3.0),
+          Complex.new(0.0, -3.0)
+        ])
+
+      out = assert_equiv(&Nx.phase/1, [x])
+      assert out.type == {:f, 32}
+    end
+  end
+
   describe "pad / triangular_solve" do
     # pad with constant value matches Emily.Backend.pad/4 (mx::pad with
     # mode "constant", no interior dilation). Cover symmetric and