Phase 1: async substrate + eval canary

c_src/emily/async.hpp

@@ -0,0 +1,143 @@
+// Async NIF machinery. See docs/planning/async-worker-exploration.md
+// for the design rationale.
+//
+// The pattern: a NIF captures the caller PID via enif_self, mints a
+// fresh ref in a process-independent env, enqueues a task onto the
+// target WorkerThread, and returns the ref synchronously. The worker
+// thread runs the task, builds the reply term in the msg_env, and
+// posts {ref, {:ok, payload}} or {ref, {:error, reason}} back to
+// the caller via enif_send. The Elixir side awaits the reply with a
+// pattern-match receive in `Emily.Native.Async.call/1`.
+//
+// Three invariants kept by this code:
+//
+// 1. enif_self must be called on the scheduler thread. The NIF
+//    captures the ErlNifPid by value into the lambda; the worker
+//    thread (non-scheduler) must never call enif_self itself.
+//
+// 2. enif_send with a non-NULL msg_env does not transfer env
+//    ownership. The worker calls enif_free_env after the send
+//    unconditionally — success invalidates the env but the caller
+//    still owns the env object.
+//
+// 3. Resource terms built via enif_make_resource(msg_env, ptr)
+//    internally bump the resource's refcount. ResourcePtrs captured
+//    by the lambda release their refs on lambda exit, but the term
+//    held by msg_env keeps the resource alive until the term is
+//    delivered + GC'd on the receiver or msg_env is freed (if the
+//    receiver PID is dead).
+
+#pragma once
+
+#include "worker.hpp"
+
+#include <fine.hpp>
+#include <mlx/mlx.h>
+
+#include <cstring>
+#include <exception>
+#include <stdexcept>
+#include <utility>
+
+namespace emily {
+
+namespace mx = mlx::core;
+
+namespace __async {
+
+// Build a binary term in msg_env from a null-terminated C string.
+inline ERL_NIF_TERM make_binary_from_cstr(ErlNifEnv *msg_env, const char *s) {
+  size_t len = std::strlen(s);
+  ERL_NIF_TERM term;
+  unsigned char *data = enif_make_new_binary(msg_env, len, &term);
+  std::memcpy(data, s, len);
+  return term;
+}
+
+// Build an error reply term in msg_env, classifying the exception
+// type. Matches fine::nif_impl's sync catch ladder so the Elixir
+// side can raise the same exception classes:
+//   std::invalid_argument -> {:argument, message}
+//   std::runtime_error    -> {:runtime, message}
+//   std::exception        -> {:runtime, message}
+//   ... (any)             -> :unknown
+inline ERL_NIF_TERM
+error_reason_from_current_exception(ErlNifEnv *msg_env) {
+  try {
+    throw;  // re-raise the current exception to classify it
+  } catch (const std::invalid_argument &e) {
+    return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "argument"),
+                            make_binary_from_cstr(msg_env, e.what()));
+  } catch (const std::runtime_error &e) {
+    return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "runtime"),
+                            make_binary_from_cstr(msg_env, e.what()));
+  } catch (const std::exception &e) {
+    return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "runtime"),
+                            make_binary_from_cstr(msg_env, e.what()));
+  } catch (...) {
+    return enif_make_atom(msg_env, "unknown");
+  }
+}
+
+}  // namespace __async
+
+// Run `build_payload` on the worker thread of `w` and post the
+// result back to the caller PID as a message. Returns a fresh ref
+// synchronously; the caller awaits the reply via
+// `Emily.Native.Async.call/1`.
+//
+// `build_payload` signature:
+//     (mx::Stream &stream, ErlNifEnv *msg_env) -> ERL_NIF_TERM
+//
+// The returned ERL_NIF_TERM is the *payload* term; this helper wraps
+// it as `{ref, {:ok, payload}}`. Exceptions thrown by the lambda are
+// caught and posted as `{ref, {:error, reason}}` where `reason` is
+// `{:argument | :runtime, binary}` or `:unknown`.
+template <typename BuildPayload>
+fine::Term async_reply(ErlNifEnv *env,
+                       fine::ResourcePtr<WorkerThread> w,
+                       BuildPayload &&build_payload) {
+  ErlNifPid caller;
+  enif_self(env, &caller);
+
+  // Mint the ref in a durable (process-independent) env so it
+  // survives the NIF return. Copy into the caller's env for the
+  // synchronous return value.
+  ErlNifEnv *msg_env = enif_alloc_env();
+  ERL_NIF_TERM ref_in_msg = enif_make_ref(msg_env);
+  ERL_NIF_TERM ref_to_return = enif_make_copy(env, ref_in_msg);
+
+  try {
+    w->run_async([msg_env, ref_in_msg, caller,
+                  build_payload = std::forward<BuildPayload>(build_payload)]
+                 (mx::Stream &s) mutable {
+      ERL_NIF_TERM reply;
+      try {
+        ERL_NIF_TERM payload = build_payload(s, msg_env);
+        ERL_NIF_TERM ok_tuple = enif_make_tuple2(
+            msg_env, enif_make_atom(msg_env, "ok"), payload);
+        reply = enif_make_tuple2(msg_env, ref_in_msg, ok_tuple);
+      } catch (...) {
+        reply = enif_make_tuple2(
+            msg_env, ref_in_msg,
+            enif_make_tuple2(
+                msg_env, enif_make_atom(msg_env, "error"),
+                __async::error_reason_from_current_exception(msg_env)));
+      }
+
+      // enif_send invalidates msg_env on success but does not take
+      // ownership of the env object itself; free it unconditionally.
+      enif_send(nullptr, &caller, msg_env, reply);
+      enif_free_env(msg_env);
+    });
+  } catch (...) {
+    // Enqueue failed (worker stopped). Reclaim the env and rethrow
+    // so fine::nif_impl surfaces the exception synchronously.
+    enif_free_env(msg_env);
+    throw;
+  }
+
+  return fine::Term(ref_to_return);
+}
+
+}  // namespace emily

c_src/emily/worker.hpp

@@ -3,7 +3,9 @@
 // MLX uses thread-local CommandEncoders — a stream's encoder only
 // exists on the thread that created it. BEAM processes migrate
 // between OS threads, so we pin each MLX stream to a dedicated
-// thread and dispatch work to it via run_sync (promise/future).
+// thread and dispatch work to it via run_sync (promise/future) or
+// run_async (fire-and-forget, with the task posting its own reply
+// via enif_send — see emily/async.hpp).
 
 #pragma once
 
@@ -32,6 +34,29 @@ class WorkerThread {
 
   ~WorkerThread() { stop(); }
 
+  // Enqueue a task without blocking. The caller is responsible for
+  // whatever side-effect the task performs (typically enif_send back
+  // to a caller PID captured at NIF entry; see emily/async.hpp).
+  // Exceptions thrown by the task are swallowed — the task owns
+  // error propagation because there is no future to carry an
+  // exception through.
+  template <typename F>
+  void run_async(F &&f) {
+    {
+      std::lock_guard<std::mutex> lock(mtx_);
+      if (stop_)
+        throw std::runtime_error("worker thread has been stopped");
+      queue_.push([f = std::forward<F>(f), this]() mutable {
+        try {
+          f(stream_);
+        } catch (...) {
+          // Swallow — the task owns error propagation.
+        }
+      });
+    }
+    cv_.notify_one();
+  }
+
   template <typename F>
   auto run_sync(F &&f) -> decltype(f(std::declval<mx::Stream &>())) {
     using R = decltype(f(std::declval<mx::Stream &>()));

c_src/emily_nif.cpp

@@ -3,6 +3,7 @@
 // Op NIFs live in c_src/ops/*.cpp; they share the Tensor resource
 // defined here via emily/tensor.hpp.
 
+#include "emily/async.hpp"
 #include "emily/tensor.hpp"
 #include "emily/worker.hpp"
 
@@ -105,14 +106,24 @@ std::tuple<fine::Atom, int64_t> dtype(ErlNifEnv *, fine::ResourcePtr<Tensor> ten
 }
 FINE_NIF(dtype, 0);
 
-// eval/2 — force evaluation of the lazy graph rooted at this tensor.
-// Dirty CPU: waits for MLX to finish.
-fine::Ok<> eval(ErlNifEnv *, fine::ResourcePtr<WorkerThread> w, fine::ResourcePtr<Tensor> tensor) {
-  w->run_sync([&](mx::Stream &) {
-    mx::eval(tensor->array);
-  });
-  return fine::Ok<>{};
+// eval_nif/2 — force evaluation of the lazy graph rooted at this
+// tensor. Async: the NIF enqueues the eval onto the worker and
+// returns a ref synchronously; the worker posts {ref, {:ok, :ok}}
+// back once eval completes. The Elixir wrapper `Emily.Native.eval/2`
+// awaits via `Emily.Native.Async.call/1`.
+//
+// Runs on a regular scheduler — enqueueing is sub-microsecond and
+// the scheduler is never blocked on MLX work.
+fine::Term eval_nif(ErlNifEnv *env,
+                    fine::ResourcePtr<WorkerThread> w,
+                    fine::ResourcePtr<Tensor> tensor) {
+  return emily::async_reply(
+      env, w,
+      [tensor](mx::Stream &, ErlNifEnv *msg_env) {
+        mx::eval(tensor->array);
+        return enif_make_atom(msg_env, "ok");
+      });
 }
-FINE_NIF(eval, ERL_NIF_DIRTY_JOB_CPU_BOUND);
+FINE_NIF(eval_nif, 0);
 
 FINE_INIT("Elixir.Emily.Native");

lib/emily/native.ex

@@ -9,6 +9,8 @@ defmodule Emily.Native do
   # thread that owns the MLX stream. Core NIFs (from_binary, shape,
   # dtype) and memory introspection NIFs don't take a worker.
 
+  alias Emily.Native.Async
+
   @on_load :__on_load__
   @compile {:autoload, false}
 
@@ -38,8 +40,12 @@ defmodule Emily.Native do
   @spec dtype(tensor()) :: dtype()
   def dtype(_tensor), do: nif()
 
+  @doc false
+  @spec eval_nif(worker(), tensor()) :: reference()
+  def eval_nif(_w, _tensor), do: nif()
+
   @spec eval(worker(), tensor()) :: :ok
-  def eval(_w, _tensor), do: nif()
+  def eval(w, tensor), do: Async.call(eval_nif(w, tensor))
 
   # --- Worker ------------------------------------------------------
 

lib/emily/native/async.ex

@@ -0,0 +1,42 @@
+defmodule Emily.Native.Async do
+  @moduledoc false
+  # Helper for awaiting async NIF replies.
+  #
+  # Async NIFs return a ref synchronously and dispatch the actual
+  # work onto a worker thread that posts
+  # `{ref, {:ok, result}}` or `{ref, {:error, reason}}` back to the
+  # caller PID via `enif_send`. `call/1` awaits that message.
+  #
+  # Error reasons mirror `fine::nif_impl`'s sync catch ladder:
+  #
+  #   {:argument, binary} -> ArgumentError
+  #   {:runtime, binary}  -> RuntimeError
+  #   :unknown            -> RuntimeError
+  #
+  # See `c_src/emily/async.hpp` and
+  # `docs/planning/async-worker-exploration.md`.
+
+  @doc """
+  Await the reply posted by the worker thread for `ref`.
+
+  Blocks the calling process on `receive/1`, not any BEAM scheduler
+  — the scheduler can run other work while the worker executes the
+  op.
+  """
+  @spec call(reference()) :: term()
+  def call(ref) do
+    receive do
+      {^ref, {:ok, result}} ->
+        result
+
+      {^ref, {:error, {:argument, message}}} ->
+        raise ArgumentError, message
+
+      {^ref, {:error, {:runtime, message}}} ->
+        raise RuntimeError, message
+
+      {^ref, {:error, :unknown}} ->
+        raise RuntimeError, "unknown exception thrown within NIF"
+    end
+  end
+end

test/emily/async_eval_test.exs

@@ -0,0 +1,102 @@
+defmodule Emily.AsyncEvalTest do
+  use ExUnit.Case, async: false
+
+  # Regression tests for the async `Emily.Native.eval/2` path.
+  # `eval/2` internally dispatches to `eval_nif/2` (which returns a
+  # ref) and awaits the worker's reply via `Emily.Native.Async.call/1`.
+  #
+  # Higher-level tests in backend_test.exs exercise eval transitively.
+  # These tests pin specific properties of the async path:
+  # - The caller's mailbox stays empty around an eval.
+  # - A process killed mid-eval does not leak, and the worker is
+  #   still usable by other callers afterward.
+  # - Latency on an already-resident tensor is small (plumbing
+  #   overhead, not a pessimisation versus the prior sync path).
+
+  describe "round-trip" do
+    test "returns :ok and leaves an empty mailbox" do
+      stream = Emily.Stream.new(:gpu)
+      worker = stream.worker
+
+      t = Emily.Native.zeros(worker, [4, 4], {:f, 32})
+      :ok = Emily.Native.eval(worker, t)
+
+      assert {:message_queue_len, 0} = Process.info(self(), :message_queue_len)
+    end
+
+    test "thousands of evals from one process drain cleanly" do
+      stream = Emily.Stream.new(:gpu)
+      worker = stream.worker
+
+      t = Emily.Native.zeros(worker, [4, 4], {:f, 32})
+
+      for _ <- 1..1_000 do
+        :ok = Emily.Native.eval(worker, t)
+      end
+
+      assert {:message_queue_len, 0} = Process.info(self(), :message_queue_len)
+    end
+  end
+
+  describe "robustness" do
+    test "caller killed mid-eval; worker remains usable" do
+      stream = Emily.Stream.new(:gpu)
+      worker = stream.worker
+
+      # A doomed process fires off many evals then exits. The worker
+      # will deliver replies to a dead PID; enif_send silently drops
+      # them. The worker must still process subsequent work from
+      # live callers.
+      {pid, mon} =
+        spawn_monitor(fn ->
+          t = Emily.Native.zeros(worker, [64, 64], {:f, 32})
+
+          for _ <- 1..100 do
+            _ref = Emily.Native.eval_nif(worker, t)
+          end
+
+          exit(:normal)
+        end)
+
+      receive do
+        {:DOWN, ^mon, :process, ^pid, _} -> :ok
+      after
+        5_000 -> flunk("doomed process did not exit")
+      end
+
+      # Let the worker drain any queued evals from the dead caller.
+      Process.sleep(100)
+
+      # The worker should still respond to fresh evals.
+      t = Emily.Native.ones(worker, [4, 4], {:f, 32})
+      :ok = Emily.Native.eval(worker, t)
+    end
+  end
+
+  describe "latency" do
+    test "eval of a resident tensor completes in under 1 ms on average" do
+      stream = Emily.Stream.new(:gpu)
+      worker = stream.worker
+
+      t = Emily.Native.zeros(worker, [4, 4], {:f, 32})
+      :ok = Emily.Native.eval(worker, t)
+
+      n = 1_000
+
+      {us, :ok} =
+        :timer.tc(fn ->
+          for _ <- 1..n, do: :ok = Emily.Native.eval(worker, t)
+          :ok
+        end)
+
+      per_op_us = us / n
+
+      # Budget: 1 ms per eval on a cached 4×4 tensor is very
+      # generous — observed locally is ~5-20 µs. A regression that
+      # pushes this over 1000 µs indicates something is queuing or
+      # scheduling incorrectly.
+      assert per_op_us < 1_000,
+             "eval averaged #{Float.round(per_op_us, 2)} µs per call"
+    end
+  end
+end