From 06d39a4184f57ead0c4c378ed8a0e9bbf476f59c Mon Sep 17 00:00:00 2001 From: ausimian Date: Sun, 19 Apr 2026 09:32:40 +1000 Subject: [PATCH] Phase 1: async substrate + eval canary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces the async NIF machinery and converts `eval/2` as the first canary. No behaviour change from the caller's perspective — `Emily.Native.eval/2` still blocks until the tensor is materialised and returns `:ok`. Internally the NIF now enqueues the eval onto the worker thread, returns a fresh ref, and awaits the worker's reply via `Emily.Native.Async.call/1` on the caller's mailbox. No BEAM scheduler thread is blocked during MLX work. The four spikes on the exploration plan (A-D) validated the plumbing before this landed; see docs/planning/async-worker-exploration.md. New: - `c_src/emily/async.hpp` — `async_reply` helper that mints a ref, captures the caller PID, enqueues onto the worker, and posts `{ref, {:ok, payload}}` or `{ref, {:error, reason}}` back via `enif_send`. Exception classification mirrors `fine::nif_impl`'s catch ladder (`std::invalid_argument` → `:argument`, `std::runtime_error` → `:runtime`, anything else → `:unknown`). - `lib/emily/native/async.ex` — `call/1` awaits the reply and re-raises errors as `ArgumentError` / `RuntimeError` to match the sync path's semantics. Changed: - `c_src/emily/worker.hpp` — added `run_async` (non-blocking enqueue, task owns error propagation). `run_sync` is unchanged; other NIFs continue to use it. - `c_src/emily_nif.cpp` — `eval` renamed to `eval_nif` and converted to async via `async_reply`. Drops the `ERL_NIF_DIRTY_JOB_CPU_BOUND` flag (no longer dirty-appropriate; enqueue is sub-microsecond and the work runs off-scheduler on the worker thread). - `lib/emily/native.ex` — `eval/2` wraps `eval_nif/2` with `Async.call/1`. Public signature and return type unchanged. Regression tests in `test/emily/async_eval_test.exs` cover: mailbox hygiene (empty after 1k evals), caller killed mid-flight (worker remains usable), latency on a resident 4x4 tensor (averages <20 us locally, budgeted under 1 ms). Full test suite passes (425 tests, 0 failures). --- c_src/emily/async.hpp | 143 +++++++++++++++++++++++++++++++++ c_src/emily/worker.hpp | 27 ++++++- c_src/emily_nif.cpp | 27 +++++-- lib/emily/native.ex | 8 +- lib/emily/native/async.ex | 42 ++++++++++ test/emily/async_eval_test.exs | 102 +++++++++++++++++++++++ 6 files changed, 339 insertions(+), 10 deletions(-) create mode 100644 c_src/emily/async.hpp create mode 100644 lib/emily/native/async.ex create mode 100644 test/emily/async_eval_test.exs diff --git a/c_src/emily/async.hpp b/c_src/emily/async.hpp new file mode 100644 index 0000000..8cecb53 --- /dev/null +++ b/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 +#include + +#include +#include +#include +#include + +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 +fine::Term async_reply(ErlNifEnv *env, + fine::ResourcePtr 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(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 diff --git a/c_src/emily/worker.hpp b/c_src/emily/worker.hpp index c57fc57..ecb3c77 100644 --- a/c_src/emily/worker.hpp +++ b/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 + void run_async(F &&f) { + { + std::lock_guard lock(mtx_); + if (stop_) + throw std::runtime_error("worker thread has been stopped"); + queue_.push([f = std::forward(f), this]() mutable { + try { + f(stream_); + } catch (...) { + // Swallow — the task owns error propagation. + } + }); + } + cv_.notify_one(); + } + template auto run_sync(F &&f) -> decltype(f(std::declval())) { using R = decltype(f(std::declval())); diff --git a/c_src/emily_nif.cpp b/c_src/emily_nif.cpp index 94bb21d..81ce701 100644 --- a/c_src/emily_nif.cpp +++ b/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 dtype(ErlNifEnv *, fine::ResourcePtr 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 w, fine::ResourcePtr 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 w, + fine::ResourcePtr 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"); diff --git a/lib/emily/native.ex b/lib/emily/native.ex index e82bf29..2ab104e 100644 --- a/lib/emily/native.ex +++ b/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 ------------------------------------------------------ diff --git a/lib/emily/native/async.ex b/lib/emily/native/async.ex new file mode 100644 index 0000000..4130e48 --- /dev/null +++ b/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 diff --git a/test/emily/async_eval_test.exs b/test/emily/async_eval_test.exs new file mode 100644 index 0000000..ce9867b --- /dev/null +++ b/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