Merge pull request #160 from ausimian/feat/expr-compiler-cm12

feat: lower Nx.Random natively (bitcast, erf_inv, dynamic slice) (CM12)

Author: ausimian

RELEASE.md

@@ -61,6 +61,16 @@
   opt-in `mx::compile` eval mode degrades to a sync replay for while-containing
   programs, which it can't trace.)
 
+- **`Nx.Random` compiles native.** The PRNG surface (`split`, `uniform`,
+  `normal`, `randint`, `gumbel`, `choice`) now lowers under the native
+  compiler, so sampling-based generation runs on the single-NIF path and a
+  PRNG key threads through a decode loop as ordinary carried state. This
+  needed three primitives: `bitcast` (random bits → float), `erf_inv`
+  (`normal`), and a **dynamic-start `slice`** — threefry indexes its rotation
+  table by the loop counter, a genuine runtime-start slice (the eager backend
+  materialises the index to a host int; the compiled replay threads it as a
+  runtime `s32` instead, same result).
+
 - **`Emily.Generation` — a model-agnostic decode-loop driver.** JIT-compiles a
   caller-supplied **shape-stable** per-token forward (`fn token, offset, cache,
   params -> {logits, cache} end`) with the native single-NIF compiler and drives

c_src/emily/opcodes.hpp

@@ -133,9 +133,18 @@ enum class Opcode : int64_t {
   // `replay_program` (it needs the subprograms + multi-output), never via
   // `dispatch_op`.
   While = 71,
+  // Reinterpret the bytes as another dtype (Nx.bitcast). operands [a];
+  // iattrs [[dtype_code]]
+  Bitcast = 72,
+  // Inverse error function (for Nx.Random.normal). operands [a]
+  ErfInv = 73,
+  // Slice with runtime (dynamic) start indices, stride 1 (the threefry/RNG
+  // path indexes by a loop counter). operands [a, start(s32 [naxes])];
+  // iattrs [[axes...], [slice_sizes...]]
+  DynSlice = 74,
 };
 
-inline constexpr int64_t kOpcodeCount = 72;
+inline constexpr int64_t kOpcodeCount = 75;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -541,6 +550,21 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
     // Multi-output + carries subprograms; handled directly in
     // replay_program, never dispatched here.
     throw std::invalid_argument("while is handled in replay_program");
+  case Opcode::Bitcast:
+    return mx::view(arg1(in, "bitcast"),
+                    emily::to_mlx_dtype_code(scalar_at(iattrs, 0, "bitcast")), s);
+  case Opcode::ErfInv:
+    return mx::erfinv(arg1(in, "erf_inv"), s);
+  case Opcode::DynSlice:
+    // operands [a, start(s32 [naxes])]; iattrs [[axes...], [slice_sizes...]].
+    // Stride-1 dynamic slice (mx::slice's dynamic-start overload).
+    if (in.size() != 2) {
+      throw std::invalid_argument("dyn_slice expects 2 operands, got " +
+                                  std::to_string(in.size()));
+    }
+    return mx::slice(in[0], in[1],
+                     emily::to_int_vec(attr_at(iattrs, 0, "dyn_slice")),
+                     emily::to_mlx_shape(attr_at(iattrs, 1, "dyn_slice")), s);
   }
   throw std::invalid_argument("unknown opcode " +
                               std::to_string(static_cast<int64_t>(op)));

lib/emily/ir.ex

@@ -120,7 +120,11 @@ defmodule Emily.IR do
     flip: 70,
     # control flow: operands = initial loop-carried state; iattrs [[arity]];
     # subprograms [condition, body]. Multi-output (produces `arity` values).
-    while: 71
+    while: 71,
+    # RNG / dynamic indexing primitives
+    bitcast: 72,
+    erf_inv: 73,
+    dyn_slice: 74
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -256,7 +260,8 @@ defmodule Emily.IR do
     sigmoid: :sigmoid,
     floor: :floor,
     ceil: :ceil,
-    erf: :erf
+    erf: :erf,
+    erf_inv: :erf_inv
   }
 
   @doc """
@@ -358,6 +363,13 @@ defmodule Emily.IR do
     emit(state, :astype, [ra], [[dtype_code(t.type)]])
   end
 
+  # bitcast: reinterpret the bytes as out.type (mirrors Emily.Backend.bitcast/2,
+  # which calls mx::view). Used by the RNG path to turn random bits into floats.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :bitcast, args: [a]}} = t, state) do
+    {ra, state} = lower_node(a, state)
+    emit(state, :bitcast, [ra], [[dtype_code(t.type)]])
+  end
+
   defp lower_op(%T{data: %Nx.Defn.Expr{op: :reshape, args: [a]}} = t, state) do
     {ra, state} = lower_node(a, state)
     emit(state, :reshape, [ra], [Tuple.to_list(t.shape)])
@@ -510,24 +522,50 @@ defmodule Emily.IR do
     coerce(r, t.type, state)
   end
 
-  # slice(t, starts, lengths, strides): static integer starts only. Nx
-  # passes scalar-tensor starts for dynamic slicing; those depend on
-  # runtime values and are deferred (the decode offset becomes a runtime
-  # input in CM3). Stops = starts + lengths (see Emily.Backend.slice/5).
+  # slice(t, starts, lengths, strides). Nx passes starts as integers (static
+  # slice) or scalar tensors (dynamic slice — e.g. threefry indexing a
+  # rotation table by the loop counter). Static starts -> mx::slice with
+  # integer bounds. Dynamic starts -> mx::slice's dynamic-start overload via
+  # the `dyn_slice` opcode (stride 1 only; the eager backend materialises the
+  # start to a host int, but the compiled replay can't, so it threads the
+  # start as a runtime s32 array).
   defp lower_op(
          %T{data: %Nx.Defn.Expr{op: :slice, args: [a, starts, lengths, strides]}} = t,
          state
        ) do
-    unless Enum.all?(starts, &is_integer/1) do
-      raise ArgumentError,
-            "Emily Expr compiler: dynamic (tensor) slice start indices are not yet " <>
-              "supported (they require a runtime input). Got: #{inspect(starts)}"
-    end
-
     {ra, state} = lower_node(a, state)
-    stops = Enum.zip_with(starts, lengths, fn st, l -> st + l end)
-    {r, state} = emit(state, :slice, [ra], [starts, stops, strides])
-    coerce(r, t.type, state)
+
+    if Enum.all?(starts, &is_integer/1) do
+      stops = Enum.zip_with(starts, lengths, fn st, l -> st + l end)
+      {r, state} = emit(state, :slice, [ra], [starts, stops, strides])
+      coerce(r, t.type, state)
+    else
+      unless Enum.all?(strides, &(&1 == 1)) do
+        raise ArgumentError,
+              "Emily Expr compiler: dynamic (tensor) slice start indices are only " <>
+                "supported with unit strides. Got strides: #{inspect(strides)}"
+      end
+
+      # Build the [ndim] s32 start array from the mixed int / scalar-tensor
+      # starts (same machinery as the dynamic put_slice write). Each runtime
+      # start is clamped to `[0, dim - length]` — MLX's dynamic slice reads
+      # out of bounds, whereas Nx (XLA semantics) clamps the start so the
+      # window stays in range; clamp is a no-op for the in-bounds starts the
+      # threefry/RNG path produces.
+      dims = Tuple.to_list(a.shape)
+
+      {start_refs, state} =
+        [starts, lengths, dims]
+        |> Enum.zip()
+        |> Enum.map_reduce(state, fn {start, length, dim}, state ->
+          {r, state} = lower_start_index(start, state)
+          clamp_start(r, dim - length, state)
+        end)
+
+      {start_arr, state} = emit(state, :concatenate, start_refs, [[0]])
+      axes = Enum.to_list(0..(length(starts) - 1)//1)
+      emit_coerced(state, :dyn_slice, [ra, start_arr], [axes, lengths], t.type)
+    end
   end
 
   # put_slice(src, start_indices, slice): write `slice` into `src` at
@@ -945,4 +983,28 @@ defmodule Emily.IR do
     {r, state} = emit(state, :astype, [r], [[dtype_code({:s, 32})]])
     emit(state, :reshape, [r], [[1]])
   end
+
+  # Clamp a runtime s32 `[1]` dynamic-slice start to `[0, hi]`
+  # (hi = dim - length), matching Nx/XLA dynamic-slice semantics — MLX's
+  # dynamic slice would otherwise read out of bounds. Both bounds are static.
+  defp clamp_start(ref, hi, state) do
+    {lo_c, state} =
+      materialize_const(
+        Nx.tensor([0], type: :s32, backend: Nx.BinaryBackend),
+        {1},
+        {:s, 32},
+        state
+      )
+
+    {hi_c, state} =
+      materialize_const(
+        Nx.tensor([hi], type: :s32, backend: Nx.BinaryBackend),
+        {1},
+        {:s, 32},
+        state
+      )
+
+    {r, state} = emit(state, :maximum, [ref, lo_c])
+    emit(state, :minimum, [r, hi_c])
+  end
 end

test/emily/compiler_rng_test.exs

@@ -0,0 +1,97 @@
+defmodule Emily.CompilerRngTest do
+  @moduledoc """
+  CM12 — `Nx.Random` lowers under the native compiler. threefry2x32 is pure
+  tensor ops (bitwise/shift/reshape/concat) plus a `while` whose body indexes
+  a rotation table by the loop counter — a *dynamic* (runtime-start) slice —
+  and `uniform`/`gumbel` turn random bits into floats via `bitcast`. With
+  those wired (CM12), the whole RNG surface compiles, and a PRNG key threads
+  through a generation-style loop as ordinary carried state.
+
+  Determinism makes the evaluator the oracle: for a fixed seed the native
+  single-NIF path must produce bit-identical draws. Run under
+  `native_fallback: :raise` so any un-lowered op fails rather than passing.
+  """
+  use ExUnit.Case, async: true
+  import Nx.Defn
+
+  @native [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+  @eval [compiler: Emily.Compiler]
+
+  defp key(seed \\ 7), do: Nx.Random.key(seed) |> Nx.backend_copy(Emily.Backend)
+
+  defp equiv(fun, args) do
+    native = apply(Nx.Defn.jit(fun, @native), args)
+    eval = apply(Nx.Defn.jit(fun, @eval), args)
+    assert Nx.to_binary(native) == Nx.to_binary(eval)
+    native
+  end
+
+  # A PRNG key threaded through a decode-style loop: each step splits the key,
+  # draws a uniform with one half, keeps the other, and accumulates. Nests
+  # threefry's own `while` (and its dynamic rotation-table slice) inside the
+  # outer loop body.
+  defn sample_loop(key) do
+    {_i, _key, acc} =
+      while {i = 0, key = key, acc = Nx.tensor(0.0, type: :f32)}, i < 5 do
+        keys = Nx.Random.split(key)
+        {u, _} = Nx.Random.uniform(keys[1])
+        {i + 1, keys[0], acc + u}
+      end
+
+    acc
+  end
+
+  describe "Nx.Random primitives: native == evaluator (fixed seed)" do
+    test "uniform", do: equiv(fn k -> Nx.Random.uniform(k, shape: {8}) |> elem(0) end, [key()])
+    test "normal", do: equiv(fn k -> Nx.Random.normal(k, shape: {8}) |> elem(0) end, [key()])
+
+    test "randint",
+      do: equiv(fn k -> Nx.Random.randint(k, 0, 100, shape: {8}) |> elem(0) end, [key()])
+
+    test "gumbel", do: equiv(fn k -> Nx.Random.gumbel(k, shape: {8}) |> elem(0) end, [key()])
+    test "split", do: equiv(fn k -> Nx.Random.split(k, parts: 4) end, [key()])
+  end
+
+  describe "dynamic-start slice (the threefry enabler)" do
+    test "slice with a runtime start index matches the evaluator" do
+      t = Nx.tensor([10.0, 20.0, 30.0, 40.0, 50.0], backend: Emily.Backend)
+      start = Nx.tensor(1, type: :s32, backend: Emily.Backend)
+      out = equiv(fn t, s -> Nx.slice(t, [s], [3]) end, [t, start])
+      assert Nx.to_flat_list(out) == [20.0, 30.0, 40.0]
+    end
+
+    test "bitcast (random bits -> float) matches the evaluator" do
+      bits = Nx.tensor([0, 1_065_353_216, 1_073_741_824], type: :u32, backend: Emily.Backend)
+      equiv(fn b -> Nx.bitcast(b, :f32) end, [bits])
+    end
+
+    test "out-of-bounds runtime start clamps to [0, dim - length] (Nx semantics)" do
+      # Oracle is Nx (BinaryBackend), which clamps the start to dim-length;
+      # MLX's dynamic slice would read out of bounds without the clamp.
+      data = [10.0, 20.0, 30.0, 40.0, 50.0]
+      fun = fn t, s -> Nx.slice(t, [s], [3]) end
+
+      native =
+        Nx.Defn.jit(fun, @native).(
+          Nx.tensor(data, backend: Emily.Backend),
+          Nx.tensor(4, type: :s32, backend: Emily.Backend)
+        )
+
+      canonical = Nx.slice(Nx.tensor(data), [Nx.tensor(4)], [3])
+      assert Nx.to_flat_list(native) == Nx.to_flat_list(canonical)
+      assert Nx.to_flat_list(native) == [30.0, 40.0, 50.0]
+    end
+  end
+
+  describe "key-splitting sampling while" do
+    test "native matches the evaluator bit-for-bit" do
+      equiv(&sample_loop/1, [key()])
+    end
+
+    test "draws actually depend on the seed (real RNG, not a constant)" do
+      a = Nx.Defn.jit(&sample_loop/1, @native).(key(1))
+      b = Nx.Defn.jit(&sample_loop/1, @native).(key(2))
+      refute Nx.to_binary(a) == Nx.to_binary(b)
+    end
+  end
+end