Merge pull request #29 from ausimian/m14.5-mlx-from-source

M14.5: worker-thread dispatch for vendored MLX

Author: ausimian

.github/workflows/ci.yml

@@ -23,19 +23,15 @@ jobs:
       MIX_ENV: test
     steps:
       - uses: actions/checkout@v4
+        with:
+          submodules: recursive
 
       - name: Setup Beam
         uses: erlef/setup-beam@v1
         with:
           version-file: .tool-versions
           version-type: strict
 
-      - name: Cache MLX prebuilt
-        uses: actions/cache@v4
-        with:
-          path: ~/Library/Caches/emily
-          key: mlx-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('mix.exs') }}
-
       - name: Cache deps
         uses: actions/cache@v4
         with:
@@ -47,7 +43,7 @@ jobs:
         uses: actions/cache@v4
         with:
           path: _build
-          key: build-${{ runner.os }}-${{ hashFiles('mix.lock', 'c_src/**', 'Makefile', '.tool-versions') }}
+          key: build-${{ runner.os }}-${{ hashFiles('mix.lock', 'c_src/**', 'Makefile', '.tool-versions', '.gitmodules') }}
           restore-keys: build-${{ runner.os }}-
 
       # The Bumblebee cache holds tiny-random HuggingFace fixtures

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "vendor/mlx"]
+	path = vendor/mlx
+	url = https://github.com/ml-explore/mlx.git

Makefile

@@ -1,6 +1,6 @@
 PRIV_DIR       := $(MIX_APP_PATH)/priv
 NIF_SO         := $(PRIV_DIR)/libemily.so
-MLX_STAGE_DIR  := $(PRIV_DIR)/mlx/lib
+METALLIB       := $(PRIV_DIR)/mlx.metallib
 
 BUILD_DIR := $(EMILY_CACHE_DIR)/build-$(EMILY_VERSION)
 
@@ -17,14 +17,18 @@ CXXFLAGS += -I$(ERTS_INCLUDE_DIR) -Ic_src
 # -Werror.
 CXXFLAGS += -isystem $(FINE_INCLUDE_DIR) -isystem $(MLX_INCLUDE_DIR)
 
-LDFLAGS  := -L$(MLX_LIB_DIR) -lmlx -shared
+# Static link: embed libmlx.a into the NIF and link system frameworks
+# that MLX depends on directly (previously resolved transitively through
+# the dylib).
+LDFLAGS  := -shared
+LDFLAGS  += $(MLX_LIB_DIR)/libmlx.a
 
 UNAME_S := $(shell uname -s)
 ifeq ($(UNAME_S),Darwin)
-    LDFLAGS += -undefined dynamic_lookup -flat_namespace -rpath @loader_path/mlx/lib
+    LDFLAGS += -undefined dynamic_lookup -flat_namespace
+    LDFLAGS += -framework Metal -framework Foundation -framework Accelerate
     JOBS    := $(shell sysctl -n hw.ncpu)
 else
-    LDFLAGS += -Wl,-rpath,'$$ORIGIN/mlx/lib'
     JOBS    := $(shell nproc)
 endif
 
@@ -43,7 +47,7 @@ MAKE_JOBS ?= $(JOBS)
 
 .PHONY: all clean bench-native
 
-all: $(NIF_SO)
+all: $(NIF_SO) $(METALLIB)
 
 # ------------------------------------------------------------------
 # bench-native: standalone C++ microbenchmarks under bench/native/.
@@ -61,8 +65,8 @@ $(BENCH_NATIVE_BIN): $(BENCH_NATIVE_SRC) | $(BUILD_DIR)
 	$(CXX) -std=c++17 -O3 -Wall -Wextra \
 	    -isystem $(MLX_INCLUDE_DIR) \
 	    $(BENCH_NATIVE_SRC) \
-	    -L$(MLX_LIB_DIR) -lmlx \
-	    -Wl,-rpath,$(MLX_LIB_DIR) \
+	    $(MLX_LIB_DIR)/libmlx.a \
+	    -framework Metal -framework Foundation -framework Accelerate \
 	    -o $(BENCH_NATIVE_BIN)
 
 bench-native: $(BENCH_NATIVE_BIN)
@@ -79,12 +83,14 @@ $(BUILD_DIR)/%.o: c_src/%.cpp $(HEADERS) | $(BUILD_DIR)
 	@mkdir -p $(dir $@)
 	$(CXX) $(CXXFLAGS) -c $< -o $@
 
-$(MLX_STAGE_DIR): | $(PRIV_DIR)
-	@mkdir -p $(MLX_STAGE_DIR)
-	@cp -a $(MLX_LIB_DIR)/. $(MLX_STAGE_DIR)/
-
-$(NIF_SO): $(OBJECTS) $(MLX_STAGE_DIR) | $(PRIV_DIR)
+$(NIF_SO): $(OBJECTS) | $(PRIV_DIR)
 	$(CXX) $(OBJECTS) -o $(NIF_SO) $(LDFLAGS)
 
+# MLX searches for mlx.metallib colocated with the loaded binary
+# (see vendor/mlx/mlx/backend/metal/device.cpp:load_default_library).
+# Stage the compiled shader library next to the NIF .so in priv/.
+$(METALLIB): $(MLX_LIB_DIR)/mlx.metallib | $(PRIV_DIR)
+	cp $< $@
+
 clean:
-	rm -rf $(BUILD_DIR) $(NIF_SO) $(PRIV_DIR)/mlx
+	rm -rf $(BUILD_DIR) $(NIF_SO) $(METALLIB)

PLAN.md

@@ -696,34 +696,10 @@ Fixed by:
    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
+GPU execution. See M14.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`,
-`determinant`, and friends route through `via_binary` today. Correct,
-BinaryBackend-slow. MLX exposes most natively under `mx::linalg::*`.
-
-- Bind each available `mx::linalg::*` function as a Native NIF.
-- Replace the `via_binary` Backend callbacks with Native dispatch.
-- Document divergences (MLX's pivot strategy may differ from Nx's
-  reference; numerical conditioning thresholds may differ).
-
-**Testing**:
-- Native unit tests against hand-computed references for small
-  matrices (3×3, 4×4) where the answer is checkable.
-- Backend property tests vs. `Nx.BinaryBackend` with shape generators
-  biased toward well-conditioned inputs (random Gaussian → QR →
-  reconstruct). Document the conditioning-bound failure mode for
-  ill-conditioned cases.
-- Existing `via_binary` fallbacks for any op MLX doesn't implement
-  natively; add a fallback-coverage test for the residual.
-
-**Exit:** all `mx::linalg::*`-backed callbacks pass property suite;
-remaining `via_binary` linalg paths documented with rationale.
-
-### M15.5 — MLX upgrade (build from source)
+### M14.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
@@ -754,6 +730,30 @@ and removes the `safe_eval` mutex introduced in the post-M14 fix.
 **Exit:** concurrent soak tests pass without mutex; MLX build-from-source
 documented; stress test confirms concurrent Metal dispatch is stable.
 
+### M15 — Native linalg
+
+`lu`, `svd`, `qr`, `cholesky`, `triangular_solve`, `eigh`,
+`determinant`, and friends route through `via_binary` today. Correct,
+BinaryBackend-slow. MLX exposes most natively under `mx::linalg::*`.
+
+- Bind each available `mx::linalg::*` function as a Native NIF.
+- Replace the `via_binary` Backend callbacks with Native dispatch.
+- Document divergences (MLX's pivot strategy may differ from Nx's
+  reference; numerical conditioning thresholds may differ).
+
+**Testing**:
+- Native unit tests against hand-computed references for small
+  matrices (3×3, 4×4) where the answer is checkable.
+- Backend property tests vs. `Nx.BinaryBackend` with shape generators
+  biased toward well-conditioned inputs (random Gaussian → QR →
+  reconstruct). Document the conditioning-bound failure mode for
+  ill-conditioned cases.
+- Existing `via_binary` fallbacks for any op MLX doesn't implement
+  natively; add a fallback-coverage test for the residual.
+
+**Exit:** all `mx::linalg::*`-backed callbacks pass property suite;
+remaining `via_binary` linalg paths documented with rationale.
+
 ### M16 — Mixed-precision training
 
 bf16 activations + f32 master weights + loss scaling is the standard

RELEASE.md

@@ -2,39 +2,50 @@
 
 ## 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.
 
 ## Added
 
-- M14 — Serving concurrency: stream-per-process. `Emily.Stream` lets
-  each BEAM process use its own Metal command queue for concurrent
-  inference. `Emily.Stream.new/1` creates a stream,
-  `Emily.Stream.with_stream/2` scopes all ops in a block to that
-  stream, and `Emily.Stream.synchronize/1` waits for completion.
-  The stream index is passed explicitly to every op NIF (no
-  thread-local race) via a `-1` sentinel for "use default stream"
-  (backwards-compatible). `Emily.Compiler.__partitions_options__/1`
-  error message now points to `Emily.Stream`.
-  - **New files**: `c_src/stream.cpp` (4 stream management NIFs),
-    `lib/emily/stream.ex` (`Emily.Stream` struct + API),
-    `test/emily/stream_test.exs`, `test/soak/stream_concurrency_test.exs`.
-  - **Modified**: every op NIF gained a trailing `int64_t s` stream
-    parameter; `Emily.Native` stubs, `Emily.Backend`, and all test
-    files updated accordingly.
-  - README now documents both concurrency patterns (stream-per-process
-    and pooled servings).
+- M14.5 — Worker-thread dispatch for vendored MLX. Replaces the
+  stream-index NIF convention (M14) and the `safe_eval` mutex with a
+  proper per-stream dedicated OS thread. Each `WorkerThread` (C++ class
+  in `c_src/emily/worker.hpp`) owns an MLX stream and its Metal
+  `CommandEncoder` on a single thread; NIFs dispatch work via
+  `run_sync` (promise/future, blocks caller for ~1-10 µs). This
+  eliminates the thread-local `CommandEncoder` mismatch that caused
+  SIGABRT/SIGSEGV under BEAM process migration, and removes the global
+  eval mutex that serialised all GPU work.
+  - **MLX built from source.** MLX is vendored as a git submodule
+    (`vendor/mlx`, pinned to commit `8e649be4`). The Makefile builds
+    `libmlx.a` via cmake and statically links it into the NIF. The
+    Metal shader library (`mlx.metallib`) is staged into `priv/` at
+    compile time. No prebuilt download step.
+  - **`Emily.MlxStream`** GenServer owns the default `WorkerThread`
+    resource under the application supervisor. `default_worker/0`
+    caches the worker ref in the process dictionary to avoid per-op
+    GenServer calls. User-created streams (`Emily.Stream.new/1`)
+    allocate their own worker; `enif_monitor_process` stops the thread
+    when the creating process exits.
+  - **NIF signature change**: every op NIF takes a worker ref as its
+    first parameter (after env) instead of a stream index as its last.
+    `Emily.Native`, `Emily.Backend`, and all test files updated
+    accordingly. `c_src/stream.cpp` (old stream management NIFs)
+    deleted.
+  - **Concurrent Metal dispatch.** Multiple `WorkerThread`s =
+    multiple MLX streams = concurrent Metal command queues. Model
+    weights (`mx::array`) are shared across threads (refcounted,
+    thread-safe reads). Shared data must be materialised before
+    passing across workers (lazy tensors are bound to their creating
+    worker's stream).
+  - **New files**: `c_src/emily/worker.hpp` (WorkerThread class),
+    `c_src/worker_nif.cpp` (resource registration with `down`
+    callback), `lib/emily/mlx_stream.ex` (GenServer),
+    `test/soak/eval_concurrency_test.exs`.
+  - **`Emily.Stream`** rewritten: struct holds a `worker` ref instead
+    of an integer index; `with_stream/2` stores it in the process
+    dictionary under `:emily_worker`; `synchronize/1` removed.
 
 - M13 — EXLA gradient conformance. Adds a third gradient oracle —
   EXLA (XLA CPU backend) — to catch bugs where Emily and BinaryBackend

c_src/emily/tensor.hpp

@@ -15,7 +15,6 @@
 #include <mlx/mlx.h>
 
 #include <cstdint>
-#include <mutex>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -60,39 +59,4 @@ unwrap_all(const std::vector<fine::ResourcePtr<Tensor>> &tensors) {
   return out;
 }
 
-// Resolve a stream index from Elixir into an mx::Stream.
-// -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::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