[DanielMercado60660]

Summary

Two threads calling compiled functions can deadlock permanently the first time a
trace runs. The per-instance CompiledFunction.lock and the process-global
evalLock are acquired in opposite orders on different threads. This makes
evalLock the outermost lock so there is only one order.

It is not a rare interleaving: with a cold compiler cache and one thread inside
compile_trace, any second thread entering a compiled function is enough. The
included reproduction wedges on round 1, every run, on both .cpu and .gpu.

While establishing that rule the PR also closes the one place that broke it:
CompiledFunction.deinit mutated the compiler cache without holding evalLock.
That is a data race rather than a deadlock, and it predates both commits below,
but it is the same cache and the same lock, so it is fixed here as its own
commit — see The compiler-cache erase in deinit.

The lock graph

CompiledFunction.call takes the per-instance lock, and then innerCall
takes the global evalLock and holds it through mlx_closure_apply. That gives
lock → evalLock. evalLock genuinely has to span that region: the
compiler cache it protects is a process-global singleton with no mutex
(mlx::core::detail::CompilerCache, compile.cpp), and the tracing flag is not
even per-thread — detail::InTracing::trace_stack() is a plain function-local
static, not thread_local (transforms.cpp), so while any thread traces,
in_tracing() is true for the whole process. Narrowing evalLock is not an
option.

But mlx_closure_apply on a cold cache entry runs compile_trace, which calls
back into the Swift body being compiled — and that body routinely calls other
compiled functions. Every MLXNN activation is a file-scope CompiledFunction,
so silu(x) inside any traced body enters compiledSilu.call and takes a second
per-instance lock while evalLock is held. That is evalLock → lock. The
two orders are inverted and close a cycle:

thread H (nested)                          thread W (plain)
  call(outer)                                call(shared)
    acquire L_cf(outer)                        acquire L_cf(shared)
    acquire evalLock
    mlx_closure_apply -> compile_trace
      -> back into Swift, body calls shared
    want L_cf(shared)  <-- held by W           want evalLock  <-- held by H

Note the per-instance lock is taken unconditionally, before C++ gets a chance
to short-circuit: detail::compile returns fun(inputs) immediately when any
input is_tracer(), so the nested call would not have touched the cache at all.
The Swift-side lock is the sole cause of the cycle.

This is first-use-only. compile_trace runs only when the cache entry is empty,
so once every trace is warm the window closes — and any new shape, dtype, stream
or constant set re-opens it.

Provenance

Neither commit is wrong on its own; the pair is.

• 1f1f9220 — "CompiledFunction is not thread safe (CompiledFunction is not thread safe #226)", Apr 2025 — added the
  per-instance NSLock to guard the observed state. Outer lock.
• b058edaa — "guard compile and vmap tracing calls (guard compile and vmap tracing calls #339)", Jan 2026 — added
  evalLock.lock() inside innerCall because tracing mutates global state.
  Inner lock. This created the inversion.

Transforms+Vmap.swift has the same shape — evalLock held across a trace that
calls back into Swift — so it is on the evalLock → lock side too, and is
consistent with the ordering after this change.

Reproduction

New executable target CompileLockRepro (Source/CompileLockRepro). It is a
separate executable on purpose: a wedged round cannot be recovered and would take
a shared test process down with it.

It is deterministic rather than a race that has to be lost:

• fresh compile() results per round, so the compiler cache is cold and the
  callback into Swift actually happens;
• a semaphore signalled from inside the traced body, so the second thread is
  released only once the first is provably inside compile_trace holding
  evalLock;
• a usleep in that same traced body holding the window open for 200 ms;
• threads owned by the repro, so nothing depends on test-runner parallelism;
• a 10 s per-round deadline, so a hang is a fast hard failure rather than a
  stuck CI job. Exit 1 means deadlock, exit 2 means the trace never ran (e.g.
  MLX_DISABLE_COMPILE), so a pass cannot be a false negative.

swift build --product CompileLockRepro
.build/debug/CompileLockRepro          # --device cpu|gpu

SwiftPM (command line) does not build the Metal shaders and the MLX scheduler
opens a GPU stream when constructed, so a swift build product needs an
mlx.metallib beside the binary even for --device cpu; one built by
xcodebuild -project xcode/MLX.xcodeproj -scheme MLX works.

Before:

compile-lock repro: device=Device(gpu, 0), 20 rounds, 10s deadline per round
FAIL: round 1: thread H did not complete within 10s.
  ...
  enteredTrace=true
  This is the lock-order inversion; the process is deadlocked.

sample(1) of the wedged process, 1572/1572 samples in identical frames:

Thread H-nested-compiled-call
  CompiledFunction.call                Transforms+Compile.swift:40   holds L_cf(outer)
  CompiledFunction.innerCall           Transforms+Compile.swift:112  holds evalLock (:89)
  mlx_closure_apply -> detail::compile compile.cpp:1180 -> :1126
  detail::compile_trace                compile.cpp:408
  [back into Swift] inner (tracers:)   Transforms+Compile.swift:70
  <traced body> -> CompiledFunction.call  Transforms+Compile.swift:40
  NSLocking.withLock -> __psynch_mutexwait

Thread W-plain-compiled-call
  CompiledFunction.call                Transforms+Compile.swift:40   holds L_cf(shared)
  CompiledFunction.innerCall           Transforms+Compile.swift:89
  -[NSRecursiveLock lock] -> __psynch_mutexwait

After: PASS: 20 rounds completed, no deadlock (exit 0), on .cpu and .gpu.

The fix

Four commits.

1. Take evalLock outermost in CompiledFunction.call and drop the
acquisition inside innerCall. One global order, evalLock → lock, so the
inversion cannot form. evalLock is recursive, so the nested acquisition from
inside a trace is free.

The cost is small: innerCall already acquired evalLock before doing any real
work and held it through mlx_closure_apply — that is, across the evaluation
itself — so this widens the held region only by the state gathering and closure
creation at the top of innerCall.

lock also becomes an NSRecursiveLock, so a compiled function whose body
re-enters that same compiled function on one thread no longer self-deadlocks on
a non-recursive NSLock.

2. Skip the per-instance lock when the function has no state. The
_updateInternal swaps in innerCall are the only mutation of anything
reachable from self — every other stored property is let and id is written
once at init — so with inputs and outputs both empty there is nothing for
lock to guard, and evalLock (held by commit 1) already serializes every
caller. This is the common case: every compiled function in MLXNN is stateless.
Optional; the deadlock is fixed without it. Happy to drop it if you would rather
keep the change minimal.

3. Check the ordering in debug builds. A passing repro only shows the race
did not fire. evalLock is recursive, so a successful try() cannot tell "this
thread owns it" from "nobody owns it"; ownership is recorded explicitly instead.
withEvalLock(_:) is now the single choke point for evalLock and keeps a
per-thread depth (pthread-specific, so the bookkeeping needs no lock of its own —
it runs while holding evalLock). CompiledFunction acquires its lock through
withInstanceLock, which requires that depth to be non-zero. All of it compiles
out of release builds.

Moving the per-instance lock back outside evalLock trips it on the first
compiled call, deterministically. Verified by doing exactly that locally:

MLX/Transforms+Eval.swift:66: Assertion failed: CompiledFunction.lock was
acquired without holding evalLock; that is the lock-order inversion, see
CompiledFunction.call

4. Erase the compiler cache under evalLock in deinit. See below.

The compiler-cache erase in deinit

CompiledFunction.deinit called mlx_detail_compile_erase(id) with no lock
held. That is compiler_cache().erase(fun_id) — a std::unordered_map::erase on
the same unguarded process-global cache that innerCall reads and inserts into
through CompilerCache::find, whose first statement is cache_[fun_id] (itself
an insert that can rehash) and whose slow path is entries.push_back. Worse,
find hands back a reference to an entry that the caller then fills in and
reads from across the entire trace and compile_replace. An unlocked erase can
rehash the map out from under a concurrent find, or destroy the very entry
another thread is compiling into.

This is a pre-existing data race — it is older than either commit in
Provenance — and it is not the deadlock. It is in this PR because
it is the same cache and the same lock: with commits 1–3 the rule is "every use
of the compiler cache happens under evalLock", and deinit was the one site
that broke it. Leaving it out would land a debug ordering checker that
deliberately does not check one of the two sites that mutate the cache.

It is commit 4 and touches nothing the first three touch beyond the same file, so
it can be dropped or split without disturbing them.

deinit runs on the releasing thread — what can that thread be holding?

Taking a lock in a deinit deserves the enumeration, since deinit fires
wherever the last reference happens to be dropped:

• A thread holding no MLX lock. It waits for the eval or compiled call in
  flight and then erases. Bounded — evalLock is only ever held for the duration
  of one call — but it is new blocking; see the note on cost below.
• A thread already inside a compiled call or an eval. A traced body that
  drops the last reference to some other compiled function is the realistic
  case. That thread already holds evalLock, which is an NSRecursiveLock, so
  the acquisition is free and cannot self-deadlock.
• A thread holding some CompiledFunction's per-instance lock. After
  commit 1 that implies it holds evalLock too, so this is the previous case. No
  new edge is added to the lock graph.
• A back-end worker thread, while an evalLock holder waits for it. This
  would be a genuine deadlock, and it does not arise: nothing in mlx-swift hands
  a CompiledFunction reference to one. The only strong reference is the closure
  returned by compile(...), and the only other retain is the
  ClosureCaptureState built by new_mlx_closure in innerCall, whose
  mlx_closure_free runs in that same function's defer, on the calling thread,
  with evalLock held. Nothing Swift-owned is stored in a CacheEntry — on the
  mlx_detail_compile path entry.extra is always nullptr — nor in the tape.
• Process teardown. Unchanged in kind: a thread wedged holding evalLock
  already hangs atexit, which is why the repro uses _exit. MLXNN's
  file-scope compiled functions are Swift globals and are not destroyed at exit.

The residual hazard is on the caller's side, and it is worth stating: after this
change, releasing a compiled function is a call into MLX. Dropping one while
holding a lock that an evalLock holder needs is an inversion, exactly as
calling eval() there would be. That is the existing discipline for this library,
but a nil assignment advertises it much less loudly than eval() does.

Why the erase stays synchronous

The obvious way to avoid the blocking would be to hand the erase to a queue. That
is not safe here: id is the object's own address
(Unmanaged.passUnretained(self).toOpaque()), so once the object is freed a new
CompiledFunction can be allocated at the same address and receive the same id.
A deferred erase could delete the new function's cache entries. Doing it in
deinit, before the memory is released, is what keeps the id unambiguous.

Note also that detail::compile_erase has a second caller in compile.cpp — the
shared_ptr deleter in the public mlx::core::compile() — which is likewise
unlocked and can run on any thread. mlx-swift never reaches it (it calls only
mlx_detail_compile), so CompiledFunction.deinit is the only erase site this
package can fix from Swift. If you would rather have the lock inside
CompilerCache itself, that is a change on the C++ side and would cover both.

Tests

Tests/MLXTests/CompileLockOrderingTests.swift — none of these can hang, all are
safe in the shared test process:

• the ownership bookkeeping can report "not held", and is per thread — otherwise
  a violation count of zero would prove nothing;

• a stateful compiled function really reaches the checked acquisition (the
  counters also assert the check ran at all, so a pass is not vacuous);

• the nested case, one compiled function called from inside another's trace;

• the erase interleaving: one thread is parked inside compile_trace holding
  evalLock — it signals a semaphore from inside the traced body and then holds
  the window open — while a second thread drops the last reference to a warmed
  compiled function at that instant. The test asserts that the second thread
  observes the traced body to have finished by the time its erase returns,
  which is what "serialized against find" means, and that exactly one erase
  happened at that point. Reverting deinit to the unlocked form fails it on the
  intended assertion, deterministically:

  CompileLockOrderingTests.swift:252: error: testDeinitEraseSerializesAgainstATraceInFlight
  XCTAssertFalse failed - the compiler-cache erase completed while another thread
  was still inside compile_trace -- the erase is not serialized against find/insert
  

• the erase under load: 200 deinit-erases against 200 cold compiles on two
  owned threads, covering the cache_[fun_id] insert-and-rehash window as well
  as tracing, asserting the results stay correct and that the erase path ran.

What that pair does and does not establish, plainly: mutual exclusion here is
structural, not statistical. Every find/insert is inside withEvalLock (from
call, the only caller of innerCall), and after commit 4 every erase is too;
the tests show the erase path really is reached with evalLock held under the
exact interleaving that used to race, and that adding the acquisition does not
wedge. The load test is a stress loop, so a green run of it alone is not proof
— which is why its assertions are the counters rather than the absence of a
crash.

No sanitizer run is included. TSan works on this toolchain, but swift test --sanitize=thread rebuilds the ~300 vendored MLX C++ translation units under it,
and at the vendored 0.31.1 it would report the pre-existing unsynchronized
globals in transforms.cpp (trace_stack, grad_counter) throughout, which
buries rather than sharpens the signal for this one line. That reasoning stops
holding after a bump to mlx 0.32.x, where those globals are thread_local: with
the background noise gone a TSan run becomes considerably more attractive, and
is the natural way to re-check this lane against a bumped core.

swift test: 541 → 547 tests, 0 failures before and after (541 existing on upstream main plus
the 6 new). CompileLockRepro: fails on round 1 before, 20/20 rounds after, on
.cpu and .gpu.

There were previously no concurrency tests in the package
(grep -rn concurrentPerform Tests/ Source/ is empty), which is presumably why
this went unnoticed.

Not fixed here

By design rather than by bug: because evalLock covers mlx_closure_apply,
every compiled call serializes against every other compiled call and every
eval/asyncEval in the process. This PR does not change that — it only moves
the acquisition a few statements earlier — but it does put a ceiling on
concurrent multi-stream use that is worth being explicit about. Commit 4 extends
the same serialization to the release of a compiled function.

The package vendors mlx 0.31.1, where the compiler cache and the tracing state
are process-global: CompilerCache has no mutex and
detail::InTracing::trace_stack() is a plain function-local static rather than
thread_local. evalLock is a Swift-side workaround for both, which is why it
has to span the whole trace rather than just the evaluation. That is no longer
the upstream shape. mlx 0.32.0 made the tracing stack, InTracing::grad_counter,
RetainGraph::tracing_counter and the compiler-cache singleton itself
thread_local, and 0.32.1 redesigned the cache — detail::CompileCache, an
unordered_map<uintptr_t, shared_ptr<vector<CacheEntry>>>, a find that returns
the entry reference together with an owning shared_ptr that keeps it alive
across the trace, and compile_erase(const weak_ptr<CompileCache>&, uintptr_t)
so a deleter erases from the cache of the thread that compiled, or from none at
all if that thread is gone. After a core bump to 0.32.x, evalLock would no
longer need to cover the trace for correctness and could shrink back to guarding
evaluation.

None of that supersedes this PR. The lock-order inversion fixed here is entirely
Swift-side — two Swift locks taken in opposite orders on different threads — and
it stands at any core version: a per-thread cache removes the C++ data race but
not the cycle. As context rather than a request, one cost of the thread-local
design is that nothing is reused across threads, so N worker threads pay N traces
and hold N tapes — which is why a shared cache under a mutex would still interest
embedders whose compiled functions are process-wide singletons, as every MLXNN
activation is.

Noted separately while reading the same package, and not fixed in this PR:
StreamOrDevice.stream(_:) (Source/MLX/Stream.swift:54-56) discards its
argument — public static func stream(_ stream: Stream) -> StreamOrDevice { StreamOrDevice(Device.defaultStream()) }
— so every caller that names a stream explicitly silently gets the device's
default stream instead. It is upstream-report shaped: a one-line fix
(StreamOrDevice(stream)), no relationship to the lock inversion, and its own
behavioural blast radius, since any code that has been relying on the observed
behaviour would start running where it asked to. Worth filing beside this PR
rather than inside it.

Related issue on ml-explore/mlx: ml-explore/mlx#4377

[davidkoski]

Nice find!

[davidkoski]

I wonder if this could be done in swift testing with the subprocess capability? Perhaps an exercise for the future.

[davidkoski]

Changes look good, thank you for finding this!