Skip to content

Commit 06d39a4

Browse files
committed
Phase 1: async substrate + eval canary
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).
1 parent d766f0d commit 06d39a4

6 files changed

Lines changed: 339 additions & 10 deletions

File tree

c_src/emily/async.hpp

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
// Async NIF machinery. See docs/planning/async-worker-exploration.md
2+
// for the design rationale.
3+
//
4+
// The pattern: a NIF captures the caller PID via enif_self, mints a
5+
// fresh ref in a process-independent env, enqueues a task onto the
6+
// target WorkerThread, and returns the ref synchronously. The worker
7+
// thread runs the task, builds the reply term in the msg_env, and
8+
// posts {ref, {:ok, payload}} or {ref, {:error, reason}} back to
9+
// the caller via enif_send. The Elixir side awaits the reply with a
10+
// pattern-match receive in `Emily.Native.Async.call/1`.
11+
//
12+
// Three invariants kept by this code:
13+
//
14+
// 1. enif_self must be called on the scheduler thread. The NIF
15+
// captures the ErlNifPid by value into the lambda; the worker
16+
// thread (non-scheduler) must never call enif_self itself.
17+
//
18+
// 2. enif_send with a non-NULL msg_env does not transfer env
19+
// ownership. The worker calls enif_free_env after the send
20+
// unconditionally — success invalidates the env but the caller
21+
// still owns the env object.
22+
//
23+
// 3. Resource terms built via enif_make_resource(msg_env, ptr)
24+
// internally bump the resource's refcount. ResourcePtrs captured
25+
// by the lambda release their refs on lambda exit, but the term
26+
// held by msg_env keeps the resource alive until the term is
27+
// delivered + GC'd on the receiver or msg_env is freed (if the
28+
// receiver PID is dead).
29+
30+
#pragma once
31+
32+
#include "worker.hpp"
33+
34+
#include <fine.hpp>
35+
#include <mlx/mlx.h>
36+
37+
#include <cstring>
38+
#include <exception>
39+
#include <stdexcept>
40+
#include <utility>
41+
42+
namespace emily {
43+
44+
namespace mx = mlx::core;
45+
46+
namespace __async {
47+
48+
// Build a binary term in msg_env from a null-terminated C string.
49+
inline ERL_NIF_TERM make_binary_from_cstr(ErlNifEnv *msg_env, const char *s) {
50+
size_t len = std::strlen(s);
51+
ERL_NIF_TERM term;
52+
unsigned char *data = enif_make_new_binary(msg_env, len, &term);
53+
std::memcpy(data, s, len);
54+
return term;
55+
}
56+
57+
// Build an error reply term in msg_env, classifying the exception
58+
// type. Matches fine::nif_impl's sync catch ladder so the Elixir
59+
// side can raise the same exception classes:
60+
// std::invalid_argument -> {:argument, message}
61+
// std::runtime_error -> {:runtime, message}
62+
// std::exception -> {:runtime, message}
63+
// ... (any) -> :unknown
64+
inline ERL_NIF_TERM
65+
error_reason_from_current_exception(ErlNifEnv *msg_env) {
66+
try {
67+
throw; // re-raise the current exception to classify it
68+
} catch (const std::invalid_argument &e) {
69+
return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "argument"),
70+
make_binary_from_cstr(msg_env, e.what()));
71+
} catch (const std::runtime_error &e) {
72+
return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "runtime"),
73+
make_binary_from_cstr(msg_env, e.what()));
74+
} catch (const std::exception &e) {
75+
return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "runtime"),
76+
make_binary_from_cstr(msg_env, e.what()));
77+
} catch (...) {
78+
return enif_make_atom(msg_env, "unknown");
79+
}
80+
}
81+
82+
} // namespace __async
83+
84+
// Run `build_payload` on the worker thread of `w` and post the
85+
// result back to the caller PID as a message. Returns a fresh ref
86+
// synchronously; the caller awaits the reply via
87+
// `Emily.Native.Async.call/1`.
88+
//
89+
// `build_payload` signature:
90+
// (mx::Stream &stream, ErlNifEnv *msg_env) -> ERL_NIF_TERM
91+
//
92+
// The returned ERL_NIF_TERM is the *payload* term; this helper wraps
93+
// it as `{ref, {:ok, payload}}`. Exceptions thrown by the lambda are
94+
// caught and posted as `{ref, {:error, reason}}` where `reason` is
95+
// `{:argument | :runtime, binary}` or `:unknown`.
96+
template <typename BuildPayload>
97+
fine::Term async_reply(ErlNifEnv *env,
98+
fine::ResourcePtr<WorkerThread> w,
99+
BuildPayload &&build_payload) {
100+
ErlNifPid caller;
101+
enif_self(env, &caller);
102+
103+
// Mint the ref in a durable (process-independent) env so it
104+
// survives the NIF return. Copy into the caller's env for the
105+
// synchronous return value.
106+
ErlNifEnv *msg_env = enif_alloc_env();
107+
ERL_NIF_TERM ref_in_msg = enif_make_ref(msg_env);
108+
ERL_NIF_TERM ref_to_return = enif_make_copy(env, ref_in_msg);
109+
110+
try {
111+
w->run_async([msg_env, ref_in_msg, caller,
112+
build_payload = std::forward<BuildPayload>(build_payload)]
113+
(mx::Stream &s) mutable {
114+
ERL_NIF_TERM reply;
115+
try {
116+
ERL_NIF_TERM payload = build_payload(s, msg_env);
117+
ERL_NIF_TERM ok_tuple = enif_make_tuple2(
118+
msg_env, enif_make_atom(msg_env, "ok"), payload);
119+
reply = enif_make_tuple2(msg_env, ref_in_msg, ok_tuple);
120+
} catch (...) {
121+
reply = enif_make_tuple2(
122+
msg_env, ref_in_msg,
123+
enif_make_tuple2(
124+
msg_env, enif_make_atom(msg_env, "error"),
125+
__async::error_reason_from_current_exception(msg_env)));
126+
}
127+
128+
// enif_send invalidates msg_env on success but does not take
129+
// ownership of the env object itself; free it unconditionally.
130+
enif_send(nullptr, &caller, msg_env, reply);
131+
enif_free_env(msg_env);
132+
});
133+
} catch (...) {
134+
// Enqueue failed (worker stopped). Reclaim the env and rethrow
135+
// so fine::nif_impl surfaces the exception synchronously.
136+
enif_free_env(msg_env);
137+
throw;
138+
}
139+
140+
return fine::Term(ref_to_return);
141+
}
142+
143+
} // namespace emily

c_src/emily/worker.hpp

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@
33
// MLX uses thread-local CommandEncoders — a stream's encoder only
44
// exists on the thread that created it. BEAM processes migrate
55
// between OS threads, so we pin each MLX stream to a dedicated
6-
// thread and dispatch work to it via run_sync (promise/future).
6+
// thread and dispatch work to it via run_sync (promise/future) or
7+
// run_async (fire-and-forget, with the task posting its own reply
8+
// via enif_send — see emily/async.hpp).
79

810
#pragma once
911

@@ -32,6 +34,29 @@ class WorkerThread {
3234

3335
~WorkerThread() { stop(); }
3436

37+
// Enqueue a task without blocking. The caller is responsible for
38+
// whatever side-effect the task performs (typically enif_send back
39+
// to a caller PID captured at NIF entry; see emily/async.hpp).
40+
// Exceptions thrown by the task are swallowed — the task owns
41+
// error propagation because there is no future to carry an
42+
// exception through.
43+
template <typename F>
44+
void run_async(F &&f) {
45+
{
46+
std::lock_guard<std::mutex> lock(mtx_);
47+
if (stop_)
48+
throw std::runtime_error("worker thread has been stopped");
49+
queue_.push([f = std::forward<F>(f), this]() mutable {
50+
try {
51+
f(stream_);
52+
} catch (...) {
53+
// Swallow — the task owns error propagation.
54+
}
55+
});
56+
}
57+
cv_.notify_one();
58+
}
59+
3560
template <typename F>
3661
auto run_sync(F &&f) -> decltype(f(std::declval<mx::Stream &>())) {
3762
using R = decltype(f(std::declval<mx::Stream &>()));

c_src/emily_nif.cpp

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
// Op NIFs live in c_src/ops/*.cpp; they share the Tensor resource
44
// defined here via emily/tensor.hpp.
55

6+
#include "emily/async.hpp"
67
#include "emily/tensor.hpp"
78
#include "emily/worker.hpp"
89

@@ -105,14 +106,24 @@ std::tuple<fine::Atom, int64_t> dtype(ErlNifEnv *, fine::ResourcePtr<Tensor> ten
105106
}
106107
FINE_NIF(dtype, 0);
107108

108-
// eval/2 — force evaluation of the lazy graph rooted at this tensor.
109-
// Dirty CPU: waits for MLX to finish.
110-
fine::Ok<> eval(ErlNifEnv *, fine::ResourcePtr<WorkerThread> w, fine::ResourcePtr<Tensor> tensor) {
111-
w->run_sync([&](mx::Stream &) {
112-
mx::eval(tensor->array);
113-
});
114-
return fine::Ok<>{};
109+
// eval_nif/2 — force evaluation of the lazy graph rooted at this
110+
// tensor. Async: the NIF enqueues the eval onto the worker and
111+
// returns a ref synchronously; the worker posts {ref, {:ok, :ok}}
112+
// back once eval completes. The Elixir wrapper `Emily.Native.eval/2`
113+
// awaits via `Emily.Native.Async.call/1`.
114+
//
115+
// Runs on a regular scheduler — enqueueing is sub-microsecond and
116+
// the scheduler is never blocked on MLX work.
117+
fine::Term eval_nif(ErlNifEnv *env,
118+
fine::ResourcePtr<WorkerThread> w,
119+
fine::ResourcePtr<Tensor> tensor) {
120+
return emily::async_reply(
121+
env, w,
122+
[tensor](mx::Stream &, ErlNifEnv *msg_env) {
123+
mx::eval(tensor->array);
124+
return enif_make_atom(msg_env, "ok");
125+
});
115126
}
116-
FINE_NIF(eval, ERL_NIF_DIRTY_JOB_CPU_BOUND);
127+
FINE_NIF(eval_nif, 0);
117128

118129
FINE_INIT("Elixir.Emily.Native");

lib/emily/native.ex

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ defmodule Emily.Native do
99
# thread that owns the MLX stream. Core NIFs (from_binary, shape,
1010
# dtype) and memory introspection NIFs don't take a worker.
1111

12+
alias Emily.Native.Async
13+
1214
@on_load :__on_load__
1315
@compile {:autoload, false}
1416

@@ -38,8 +40,12 @@ defmodule Emily.Native do
3840
@spec dtype(tensor()) :: dtype()
3941
def dtype(_tensor), do: nif()
4042

43+
@doc false
44+
@spec eval_nif(worker(), tensor()) :: reference()
45+
def eval_nif(_w, _tensor), do: nif()
46+
4147
@spec eval(worker(), tensor()) :: :ok
42-
def eval(_w, _tensor), do: nif()
48+
def eval(w, tensor), do: Async.call(eval_nif(w, tensor))
4349

4450
# --- Worker ------------------------------------------------------
4551

lib/emily/native/async.ex

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
defmodule Emily.Native.Async do
2+
@moduledoc false
3+
# Helper for awaiting async NIF replies.
4+
#
5+
# Async NIFs return a ref synchronously and dispatch the actual
6+
# work onto a worker thread that posts
7+
# `{ref, {:ok, result}}` or `{ref, {:error, reason}}` back to the
8+
# caller PID via `enif_send`. `call/1` awaits that message.
9+
#
10+
# Error reasons mirror `fine::nif_impl`'s sync catch ladder:
11+
#
12+
# {:argument, binary} -> ArgumentError
13+
# {:runtime, binary} -> RuntimeError
14+
# :unknown -> RuntimeError
15+
#
16+
# See `c_src/emily/async.hpp` and
17+
# `docs/planning/async-worker-exploration.md`.
18+
19+
@doc """
20+
Await the reply posted by the worker thread for `ref`.
21+
22+
Blocks the calling process on `receive/1`, not any BEAM scheduler
23+
— the scheduler can run other work while the worker executes the
24+
op.
25+
"""
26+
@spec call(reference()) :: term()
27+
def call(ref) do
28+
receive do
29+
{^ref, {:ok, result}} ->
30+
result
31+
32+
{^ref, {:error, {:argument, message}}} ->
33+
raise ArgumentError, message
34+
35+
{^ref, {:error, {:runtime, message}}} ->
36+
raise RuntimeError, message
37+
38+
{^ref, {:error, :unknown}} ->
39+
raise RuntimeError, "unknown exception thrown within NIF"
40+
end
41+
end
42+
end

test/emily/async_eval_test.exs

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
defmodule Emily.AsyncEvalTest do
2+
use ExUnit.Case, async: false
3+
4+
# Regression tests for the async `Emily.Native.eval/2` path.
5+
# `eval/2` internally dispatches to `eval_nif/2` (which returns a
6+
# ref) and awaits the worker's reply via `Emily.Native.Async.call/1`.
7+
#
8+
# Higher-level tests in backend_test.exs exercise eval transitively.
9+
# These tests pin specific properties of the async path:
10+
# - The caller's mailbox stays empty around an eval.
11+
# - A process killed mid-eval does not leak, and the worker is
12+
# still usable by other callers afterward.
13+
# - Latency on an already-resident tensor is small (plumbing
14+
# overhead, not a pessimisation versus the prior sync path).
15+
16+
describe "round-trip" do
17+
test "returns :ok and leaves an empty mailbox" do
18+
stream = Emily.Stream.new(:gpu)
19+
worker = stream.worker
20+
21+
t = Emily.Native.zeros(worker, [4, 4], {:f, 32})
22+
:ok = Emily.Native.eval(worker, t)
23+
24+
assert {:message_queue_len, 0} = Process.info(self(), :message_queue_len)
25+
end
26+
27+
test "thousands of evals from one process drain cleanly" do
28+
stream = Emily.Stream.new(:gpu)
29+
worker = stream.worker
30+
31+
t = Emily.Native.zeros(worker, [4, 4], {:f, 32})
32+
33+
for _ <- 1..1_000 do
34+
:ok = Emily.Native.eval(worker, t)
35+
end
36+
37+
assert {:message_queue_len, 0} = Process.info(self(), :message_queue_len)
38+
end
39+
end
40+
41+
describe "robustness" do
42+
test "caller killed mid-eval; worker remains usable" do
43+
stream = Emily.Stream.new(:gpu)
44+
worker = stream.worker
45+
46+
# A doomed process fires off many evals then exits. The worker
47+
# will deliver replies to a dead PID; enif_send silently drops
48+
# them. The worker must still process subsequent work from
49+
# live callers.
50+
{pid, mon} =
51+
spawn_monitor(fn ->
52+
t = Emily.Native.zeros(worker, [64, 64], {:f, 32})
53+
54+
for _ <- 1..100 do
55+
_ref = Emily.Native.eval_nif(worker, t)
56+
end
57+
58+
exit(:normal)
59+
end)
60+
61+
receive do
62+
{:DOWN, ^mon, :process, ^pid, _} -> :ok
63+
after
64+
5_000 -> flunk("doomed process did not exit")
65+
end
66+
67+
# Let the worker drain any queued evals from the dead caller.
68+
Process.sleep(100)
69+
70+
# The worker should still respond to fresh evals.
71+
t = Emily.Native.ones(worker, [4, 4], {:f, 32})
72+
:ok = Emily.Native.eval(worker, t)
73+
end
74+
end
75+
76+
describe "latency" do
77+
test "eval of a resident tensor completes in under 1 ms on average" do
78+
stream = Emily.Stream.new(:gpu)
79+
worker = stream.worker
80+
81+
t = Emily.Native.zeros(worker, [4, 4], {:f, 32})
82+
:ok = Emily.Native.eval(worker, t)
83+
84+
n = 1_000
85+
86+
{us, :ok} =
87+
:timer.tc(fn ->
88+
for _ <- 1..n, do: :ok = Emily.Native.eval(worker, t)
89+
:ok
90+
end)
91+
92+
per_op_us = us / n
93+
94+
# Budget: 1 ms per eval on a cached 4×4 tensor is very
95+
# generous — observed locally is ~5-20 µs. A regression that
96+
# pushes this over 1000 µs indicates something is queuing or
97+
# scheduling incorrectly.
98+
assert per_op_us < 1_000,
99+
"eval averaged #{Float.round(per_op_us, 2)} µs per call"
100+
end
101+
end
102+
end

0 commit comments

Comments
 (0)