Merge pull request #153 from ausimian/feat/expr-compiler-cm6

CM6: opt-in mx::compile eval mode for the program replay

Author: ausimian

bench/program_compile.exs

@@ -0,0 +1,198 @@
+# CM6 mx::compile secondary-win microbench.
+#
+# The single-NIF replay (CM0) already collapses the ~750 per-op
+# BEAM<->worker round-trips into one. CM6 adds an *opt-in* `:compiled`
+# eval mode that wraps that replay in `mx::compile` (cached per stream,
+# keyed off the Program). `mx::compile` can fuse adjacent elementwise
+# kernels, so the question this bench answers is narrow: on a realistic
+# transformer block, does the compiled wrap buy a meaningful secondary
+# encode win over the already-collapsed sync replay?
+#
+# The prior (Emily's own M6 milestone) predicted ~1.04-1.11x: MLX does
+# not fuse matmul with adjacent elementwise. In practice this bench
+# measures ~1.6x at decode-shaped sizes -- because at small sequence
+# lengths kernel-launch and intermediate-memory overhead dominate, and
+# fusing the elementwise runs (rms-norm, softmax, SiLU gating, residuals)
+# removes them. The speedup is reported, not gated (only the correctness
+# guard halts) -- CM6 is additive, not load-bearing: CM0's single-NIF
+# replay already delivered the main dispatch collapse, and CM3 met the
+# tok/s target without this.
+#
+#   mix run bench/program_compile.exs
+#   mix run bench/program_compile.exs -- --seq 128 --iters 300
+#
+# Both paths run the SAME compiled Program; only the eval mode differs,
+# and a correctness guard asserts they agree within a tight f32 tolerance
+# before timing (mx::compile fuses elementwise runs, which reassociates
+# float arithmetic to within a few ULP -- correct, not bit-exact).
+
+alias Emily.{IR, Native, Program}
+alias Nx.Defn.Composite
+
+{opts, _, _} =
+  OptionParser.parse(System.argv(),
+    strict: [seq: :integer, dmodel: :integer, ffn: :integer, iters: :integer]
+  )
+
+seq = opts[:seq] || 64
+d = opts[:dmodel] || 512
+ffn = opts[:ffn] || 2048
+iters = opts[:iters] || 300
+
+worker = Emily.MlxStream.default_worker()
+
+# --- A realistic transformer block, in Nx primitives only ---
+# RMSNorm -> single-head attention (QKV matmuls, softmax, out proj) ->
+# residual -> RMSNorm -> SwiGLU FFN -> residual. The elementwise runs
+# (rms-norm, softmax, SiLU gating, residual adds) are mx::compile's
+# fusion surface; the matmuls are the floor it can't fuse through.
+eps = 1.0e-6
+scale = 1.0 / :math.sqrt(d)
+
+rms = fn t, g ->
+  ms = Nx.mean(Nx.multiply(t, t), axes: [-1], keep_axes: true)
+  Nx.multiply(Nx.divide(t, Nx.sqrt(Nx.add(ms, eps))), g)
+end
+
+silu = fn z -> Nx.multiply(z, Nx.divide(1.0, Nx.add(1.0, Nx.exp(Nx.negate(z))))) end
+
+softmax = fn s ->
+  m = Nx.reduce_max(s, axes: [-1], keep_axes: true)
+  e = Nx.exp(Nx.subtract(s, m))
+  Nx.divide(e, Nx.sum(e, axes: [-1], keep_axes: true))
+end
+
+# Param order: x, g1, w_q, w_k, w_v, w_o, g2, w_gate, w_up, w_down.
+block = fn [x, g1, w_q, w_k, w_v, w_o, g2, w_gate, w_up, w_down] ->
+  h = rms.(x, g1)
+  q = Nx.dot(h, w_q)
+  k = Nx.dot(h, w_k)
+  v = Nx.dot(h, w_v)
+  scores = Nx.multiply(Nx.dot(q, Nx.transpose(k)), scale)
+  ctx = Nx.dot(softmax.(scores), v)
+  x2 = Nx.add(x, Nx.dot(ctx, w_o))
+
+  f = rms.(x2, g2)
+  gated = Nx.multiply(silu.(Nx.dot(f, w_gate)), Nx.dot(f, w_up))
+  Nx.add(x2, Nx.dot(gated, w_down))
+end
+
+# Concrete weights on Emily.Backend; the param exprs become {:input, i}.
+dev = fn shape ->
+  dims = Tuple.to_list(shape)
+  n = Enum.product(dims)
+  bin = for i <- 1..n, into: <<>>, do: <<:math.sin(i * 0.013) * 0.1::float-32-native>>
+  Native.from_binary(bin, dims, {:f, 32})
+end
+
+shapes = [
+  {seq, d},
+  {d},
+  {d, d},
+  {d, d},
+  {d, d},
+  {d, d},
+  {d},
+  {d, ffn},
+  {d, ffn},
+  {ffn, d}
+]
+
+input_refs = Enum.map(shapes, dev)
+
+# Trace the block into an Expr (params at :root) and lower it once.
+vars =
+  shapes
+  |> Enum.with_index()
+  |> Enum.map(fn {shape, i} ->
+    Nx.Defn.Expr.parameter(Nx.template(shape, {:f, 32}), :root, i)
+  end)
+
+expr = block.(vars)
+{_template, leaves_rev} = Composite.traverse(expr, [], fn leaf, acc -> {leaf, [leaf | acc]} end)
+prog = leaves_rev |> Enum.reverse() |> IR.lower() |> Program.compile()
+
+# Correctness guard: compiled wrap must match the plain replay. On a
+# shallow / integer-stable program the two are bit-identical, but
+# mx::compile fuses elementwise runs (rms-norm, the `* scale`, softmax,
+# SiLU gating), which reassociates f32 arithmetic. Through a deep block
+# that shows up as a last-few-ULP drift -- correct, just not bit-exact.
+# So assert a tight absolute tolerance and report the actual drift.
+tol = 1.0e-5
+[sync_out] = Program.eval(worker, prog, input_refs, mode: :sync)
+[compiled_out] = Program.eval(worker, prog, input_refs, mode: :compiled)
+
+floats = fn ref ->
+  for <<v::float-32-native <- Native.to_binary(worker, ref)>>, do: v
+end
+
+sync_floats = floats.(sync_out)
+compiled_floats = floats.(compiled_out)
+
+# Guard length first: Enum.zip/2 truncates to the shorter list, so a real
+# shape/length divergence between the two modes would otherwise be hidden.
+if length(sync_floats) != length(compiled_floats) do
+  IO.puts(:stderr, "FAIL: compiled output length differs from sync output length")
+  System.halt(1)
+end
+
+max_drift =
+  Enum.zip(sync_floats, compiled_floats)
+  |> Enum.reduce(0.0, fn {a, b}, acc -> max(acc, abs(a - b)) end)
+
+if max_drift > tol do
+  IO.puts(:stderr, "FAIL: compiled vs sync max drift #{max_drift} exceeds tol #{tol}")
+  System.halt(1)
+end
+
+build = fn -> Program.eval(worker, prog, input_refs, mode: :build) end
+
+sync = fn ->
+  [out] = Program.eval(worker, prog, input_refs, mode: :sync)
+  Native.eval(worker, out)
+end
+
+compiled = fn ->
+  [out] = Program.eval(worker, prog, input_refs, mode: :compiled)
+  Native.eval(worker, out)
+end
+
+warmup = fn f -> for _ <- 1..30, do: f.() end
+time = fn f ->
+  {us, _} = :timer.tc(fn -> for _ <- 1..iters, do: f.() end)
+  us / iters
+end
+
+for f <- [build, sync, compiled], do: warmup.(f)
+
+build_us = time.(build)
+sync_us = time.(sync)
+compiled_us = time.(compiled)
+
+IO.puts("""
+
+CM6 mx::compile secondary-win microbench
+  transformer block : seq=#{seq} d_model=#{d} ffn=#{ffn}
+  iterations        : #{iters}
+
+  build only (dispatch, no GPU)  : #{Float.round(build_us, 1)} us/iter
+  sync replay (eval)             : #{Float.round(sync_us, 1)} us/iter
+  compiled replay (mx::compile)  : #{Float.round(compiled_us, 1)} us/iter
+
+  compiled / sync speedup        : #{Float.round(sync_us / compiled_us, 3)}x
+  compiled vs sync max drift     : #{max_drift} (f32 fusion reassociation, tol #{tol})
+""")
+
+cond do
+  sync_us / compiled_us >= 1.09 ->
+    IO.puts("RESULT: meets the ~1.09x M6 prediction -- compiled mode earns its keep.")
+
+  sync_us / compiled_us >= 1.05 ->
+    IO.puts("RESULT: marginal (1.05-1.09x). Opt-in is justified; stays off by default.")
+
+  true ->
+    IO.puts(
+      "RESULT: below 1.05x. As the plan anticipated, mx::compile is not load-bearing; " <>
+        ":compiled stays strictly opt-in and the single-NIF replay (CM0) is the real win."
+    )
+end

c_src/emily/program.hpp

@@ -14,14 +14,23 @@
 
 #include "opcodes.hpp"
 #include "tensor.hpp"
+#include "worker.hpp"
 
 #include <fine.hpp>
+#include <mlx/mlx.h>
 
 #include <cstdint>
+#include <functional>
+#include <map>
+#include <memory>
+#include <mutex>
+#include <utility>
 #include <vector>
 
 namespace emily {
 
+namespace mx = mlx::core;
+
 // Operand references are packed into an int64 by the Elixir lowerer
 // (`Emily.IR.pack_ref/1`): the high bits carry the slot kind, the low
 // bits the index. Unpacked here during compile-time validation and
@@ -57,6 +66,55 @@ class Program {
   std::vector<fine::ResourcePtr<Tensor>> consts;
   std::vector<CompiledInstr> instrs;
   std::vector<int64_t> outputs; // packed refs
+
+  using CompiledFn =
+      std::function<std::vector<mx::array>(const std::vector<mx::array> &)>;
+
+  // One `mx::compile`d replay callable, plus a weak handle to the worker
+  // whose thread-local compiler cache holds its traced graph. The handle
+  // lets `~Program` drop the callable back on that worker thread (see the
+  // destructor) — `mx::compile`'s cache erase is thread-affine.
+  struct CompiledEntry {
+    std::weak_ptr<State> worker;
+    CompiledFn fn;
+  };
+
+  // CM6: opt-in mx::compile cache. One compiled replay callable per stream
+  // index — the compiled graph bakes in the captured weights and the
+  // stream, so it must be keyed by stream and rebuilt if used on a
+  // different one. Built lazily on the first compiled eval (eval_mode 3).
+  // This is the *secondary* encode win; the main dispatch-collapse win is
+  // the single-NIF replay itself.
+  std::mutex compile_mtx;
+  std::map<int, CompiledEntry> compiled;
+
+  Program() = default;
+
+  // Drop each compiled callable on the worker thread that built it. MLX's
+  // compiler cache is `thread_local` and the callable's deleter calls
+  // `compile_erase` wherever it is destroyed; this resource is collected on
+  // a BEAM/GC thread, so destroying the callable here would erase the wrong
+  // thread's cache — leaking the worker's traced graph (and its refs to the
+  // captured weight buffers), and risking a later `fun_id` (a recycled heap
+  // address) colliding with the stale entry. Posting the drop to the worker
+  // makes the erase land on the right cache. If the worker is already gone,
+  // its thread exit destroyed the thread-local cache (and our entry) for us.
+  ~Program() {
+    for (auto &kv : compiled) {
+      CompiledEntry &entry = kv.second;
+      if (!entry.fn) {
+        continue;
+      }
+      if (auto st = entry.worker.lock()) {
+        post_to_worker(*st, [fn = std::move(entry.fn)]() mutable { fn = nullptr; });
+      }
+    }
+  }
+
+  // Movable/copyable would be wrong (std::mutex member), and the explicit
+  // destructor suppresses the implicit moves anyway; spell it out.
+  Program(const Program &) = delete;
+  Program &operator=(const Program &) = delete;
 };
 
 } // namespace emily

c_src/emily/worker.hpp

@@ -68,6 +68,24 @@ inline void signal_stop(State &st) {
   st.cv.notify_all();
 }
 
+// Best-effort: enqueue `fn` to run (and then be destroyed) on the worker
+// thread. Used to drop thread-affine resources — e.g. an `mx::compile`
+// cache entry, whose erase must happen on the thread whose thread-local
+// compiler cache holds it — on the thread that created them. Bypasses the
+// queue cap (teardown is tiny and must not be back-pressured). If the
+// worker is already stopping, does nothing: the thread's exit destroys its
+// thread-local state anyway, so the resource is reclaimed regardless.
+inline void post_to_worker(State &st, std::function<void()> fn) {
+  {
+    std::lock_guard<std::mutex> lock(st.mtx);
+    if (st.stop) {
+      return;
+    }
+    st.queue.push([fn = std::move(fn)](mx::Stream &, bool) mutable { fn(); });
+  }
+  st.cv.notify_one();
+}
+
 // Joins worker threads off the BEAM schedulers. Singleton; its thread is
 // created lazily on first use and joined in the NIF unload callback.
 class Reaper {
@@ -204,6 +222,12 @@ class WorkerThread {
     return state_->queue.size();
   }
 
+  // A non-owning handle to this worker's state, for posting thread-affine
+  // teardown (see `post_to_worker`) from an object that outlives — or is
+  // collected independently of — the worker resource. Weak so the worker
+  // can be reclaimed while the holder lives; `lock()` fails once it is.
+  std::weak_ptr<State> weak_state() const { return state_; }
+
 private:
   static void run(std::shared_ptr<State> st) {
     st->stream = mx::new_stream(mx::Device(mx::Device::DeviceType::gpu));