Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 56 additions & 0 deletions PLAN.md
Original file line number Diff line number Diff line change
Expand Up @@ -674,6 +674,31 @@ silent about it and ships a tested pattern.
**Exit:** both patterns documented; concurrency soak demonstrates
the streamed path is stable.

**Post-M14 note — eval serialisation:**
MLX is not thread-safe (ml-explore/mlx#2133). The Metal
`CommandEncoder` is shared state — concurrent `mx::eval` calls from
different OS threads crash with `"A command encoder is already
encoding to this command buffer"` (SIGABRT) or SIGSEGV from corrupted
encoder state. The M14 soak tests (8 concurrent workers via
`Task.async_stream`) exposed this by being the first code path to call
`mx::eval` from multiple dirty-CPU scheduler threads simultaneously.

Fixed by:
1. **`emily::safe_eval()`**: a mutex-serialised `mx::eval` wrapper in
`c_src/emily/tensor.hpp`. All eval callsites route through it.
Graph-building ops (regular scheduler) remain lock-free.
2. **Removed `set_default_stream` from `with_stream/2`**: the NIF
mutated MLX thread-local state on BEAM scheduler threads, which
is unreliable since BEAM processes migrate between OS threads.
The process-dictionary-based stream routing was already correct.
3. **Hardened `resolve_stream(-1)`**: the -1 fallback now reads the
device default directly instead of the (potentially corrupted)
thread-local default.

The mutex serialises Metal dispatch at the cost of true concurrent
GPU execution. See M15.5 (MLX upgrade) for the plan to restore
concurrency via MLX's native thread-local `CommandEncoder` support.

### M15 — Native linalg

`lu`, `svd`, `qr`, `cholesky`, `triangular_solve`, `eigh`,
Expand All @@ -698,6 +723,37 @@ BinaryBackend-slow. MLX exposes most natively under `mx::linalg::*`.
**Exit:** all `mx::linalg::*`-backed callbacks pass property suite;
remaining `via_binary` linalg paths documented with rationale.

### M15.5 — MLX upgrade (build from source)

Emily pins MLX 0.25.1 via pre-built binaries from `cocoa-xu/mlx-build`.
MLX gained native thread-safety on `main` in April 2026 (thread-local
`CommandEncoder` ml-explore/mlx#3348, `ThreadLocalStream` C++ API
ml-explore/mlx#3405), but neither fix is in any release yet (latest:
0.31.1). Building from source unblocks true concurrent Metal dispatch
and removes the `safe_eval` mutex introduced in the post-M14 fix.

- **Build MLX from source** (from `main` or 0.32+ when released)
instead of fetching the pre-built 0.25.1 tarball. Extend `mix.exs`
to support an `MLX_SOURCE` env var pointing at a local MLX build.
- **Audit API changes** between 0.25.1 and target version. Emily's
C++ surface is narrow (core ops, linalg, streams, eval, allocator),
but six months of MLX releases may rename or remove functions.
- **Adopt `ThreadLocalStream` C++ API**: each BEAM dirty-CPU thread
gets its own MLX stream automatically, enabling true per-thread
Metal command queues without the eval mutex.
- **Remove `emily::safe_eval` mutex** once native thread-safety is
validated — revert to direct `mx::eval` calls.

**Testing**:
- Amplified stress test (16 workers, 100 iterations) passes without
mutex under the new MLX build.
- Full test suite 20x with zero crashes.
- Benchmark `to_binary` latency under concurrent load: confirm
throughput improves vs. the mutex-serialised path.

**Exit:** concurrent soak tests pass without mutex; MLX build-from-source
documented; stress test confirms concurrent Metal dispatch is stable.

### M16 — Mixed-precision training

bf16 activations + f32 master weights + loss scaling is the standard
Expand Down
10 changes: 10 additions & 0 deletions RELEASE.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,16 @@

## Fixed

- **Fix SIGABRT/SIGSEGV from concurrent `mx::eval` dispatch.** MLX's
Metal `CommandEncoder` is not thread-safe (ml-explore/mlx#2133).
Concurrent `mx::eval` calls from BEAM dirty-CPU scheduler threads
triggered `"A command encoder is already encoding"` assertions or
SIGSEGV from corrupted encoder state. Fixed by serialising all
`mx::eval` calls through `emily::safe_eval()` (mutex in
`c_src/emily/tensor.hpp`). Also removed `set_default_stream` calls
from `with_stream/2` — the NIF mutated MLX thread-local state which
is unreliable under BEAM process migration. Hardened
`resolve_stream(-1)` to avoid reading the thread-local default.
- Relax MNIST convergence canary threshold from 97% to 96% to eliminate
stochastic flaps (observed 96.99% on occasional runs). The test is a
sanity gate, not a performance benchmark.
Expand Down
33 changes: 30 additions & 3 deletions c_src/emily/tensor.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#include <mlx/mlx.h>

#include <cstdint>
#include <mutex>
#include <stdexcept>
#include <string>
#include <vector>
Expand Down Expand Up @@ -60,12 +61,38 @@ unwrap_all(const std::vector<fine::ResourcePtr<Tensor>> &tensors) {
}

// Resolve a stream index from Elixir into an mx::Stream.
// -1 (the sentinel for "no explicit stream") falls through to the
// thread-local default — backwards-compatible with pre-M14 code paths.
// -1 (the sentinel for "no explicit stream") returns the default GPU
// stream (index 0). We intentionally avoid mx::default_stream() here
// because that reads a thread-local which can be corrupted on BEAM
// scheduler threads — the BEAM migrates processes between OS threads,
// so thread-local state is unreliable.
inline mx::Stream resolve_stream(int64_t stream_index) {
if (stream_index < 0)
return mx::default_stream(mx::default_device());
return mx::default_stream(mx::Device(mx::Device::DeviceType::gpu));
return mx::get_stream(static_cast<int>(stream_index));
}

// MLX is not thread-safe (ml-explore/mlx#2133). In particular, the
// Metal CommandEncoder is shared state — concurrent mx::eval calls
// from different OS threads crash with "A command encoder is already
// encoding to this command buffer". BEAM dirty-CPU schedulers are a
// thread pool, so concurrent to_binary / eval NIF calls race.
//
// Serialise all mx::eval calls behind a single mutex until MLX gains
// native thread-safety (expected 0.32+, see ml-explore/mlx#3348).
inline std::mutex &eval_mutex() {
static std::mutex m;
return m;
}

inline void safe_eval(mx::array &a) {
std::lock_guard<std::mutex> lock(eval_mutex());
mx::eval(a);
}

inline void safe_eval(std::initializer_list<mx::array> arrays) {
std::lock_guard<std::mutex> lock(eval_mutex());
mx::eval(arrays);
}

} // namespace emily
4 changes: 2 additions & 2 deletions c_src/emily_nif.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ FINE_NIF(from_binary, 0);
fine::Term to_binary(ErlNifEnv *env, fine::ResourcePtr<Tensor> tensor, int64_t s) {
auto stream = emily::resolve_stream(s);
auto materialized = mx::contiguous(tensor->array, false, stream);
mx::eval(materialized);
emily::safe_eval(materialized);

// Defensive: mx::contiguous is supposed to give a row-contiguous
// layout. If it ever doesn't, we'd be aliasing a strided buffer and
Expand Down Expand Up @@ -113,7 +113,7 @@ FINE_NIF(dtype, 0);
// eval/1 — force evaluation of the lazy graph rooted at this tensor.
// Dirty CPU: waits for MLX to finish.
fine::Ok<> eval(ErlNifEnv *, fine::ResourcePtr<Tensor> tensor) {
mx::eval(tensor->array);
emily::safe_eval(tensor->array);
return fine::Ok<>{};
}
FINE_NIF(eval, ERL_NIF_DIRTY_JOB_CPU_BOUND);
Expand Down
7 changes: 5 additions & 2 deletions c_src/stream.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,11 @@ int64_t new_stream(ErlNifEnv *, fine::Atom device_atom) {
FINE_NIF(new_stream, 0);

// set_default_stream/1 — set the thread-local default stream.
// Belt-and-suspenders for Emily.Stream.with_stream/2; every op NIF
// also receives the stream index explicitly.
//
// WARNING: BEAM processes migrate between OS threads, so thread-local
// state is unreliable. with_stream/2 no longer calls this — it routes
// streams via the process dictionary instead. This NIF is retained for
// advanced use cases but should be avoided in normal code.
fine::Ok<> set_default_stream(ErlNifEnv *, int64_t stream_index) {
mx::set_default_stream(mx::get_stream(static_cast<int>(stream_index)));
return fine::Ok<>{};
Expand Down
18 changes: 5 additions & 13 deletions lib/emily/stream.ex
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,9 @@ defmodule Emily.Stream do
`with_stream/2` stores the stream index in the process dictionary.
`Emily.Backend` reads it via `Process.get(:emily_stream, -1)` and
passes it as an explicit argument to every NIF call. Each NIF
resolves the index to an `mx::Stream` (or falls back to the
thread-local default when the index is -1). This avoids the
thread-local race that would occur if we relied solely on
`set_default_stream` — BEAM processes can migrate between OS
resolves the index to an `mx::Stream` (or falls back to the default
GPU stream when the index is -1). This is process-dictionary based,
not thread-local, because BEAM processes migrate between OS
scheduler threads between NIF calls.

## Concurrent serving patterns
Expand Down Expand Up @@ -73,18 +72,13 @@ defmodule Emily.Stream do
Execute `fun` with the given stream as the default for MLX ops.

Stores the stream index in the process dictionary so that
`Emily.Backend` passes it to every NIF call. Also sets the
thread-local default stream as a belt-and-suspenders measure for
code that calls `Emily.Native` directly.

The previous stream (if any) is restored in an `after` block, so
nesting is safe.
`Emily.Backend` passes it to every NIF call. The previous stream
(if any) is restored in an `after` block, so nesting is safe.
"""
@spec with_stream(t(), (-> result)) :: result when result: var
def with_stream(%__MODULE__{index: index}, fun) when is_function(fun, 0) do
prev = Process.get(:emily_stream)
Process.put(:emily_stream, index)
Emily.Native.set_default_stream(index)

try do
fun.()
Expand All @@ -93,8 +87,6 @@ defmodule Emily.Stream do
nil -> Process.delete(:emily_stream)
idx -> Process.put(:emily_stream, idx)
end

if prev, do: Emily.Native.set_default_stream(prev)
end
end

Expand Down
Loading