The compiler cache and tracing state are process-global and unsynchronized in 0.31.x — is the 0.32 thread-local redesign the intended thread-safety contract? #4377

Summary

In 0.31.x, three pieces of state that every compile() call touches are process-global
with no synchronization: the compiler cache, the tracing stack, and the counters behind
in_tracing() / retain_graph(). Any binding that lets two threads call a compiled
function has to serialize all of it itself, and the serialization has to span the whole
trace rather than just the cache lookup, because tracing state is global too.

Reading the 0.32.0 and 0.32.1 headers and shipped binaries, it looks like this was
addressed — by making the cache and the tracing state thread_local, and (in 0.32.1) by
putting the per-thread cache behind a shared_ptr so the erase path can hold a weak_ptr
to it. I could not find that written down anywhere, so this issue is mostly asking whether
that is a deliberate, supported thread-safety contract that bindings can rely on, or an
incidental consequence of another change. The rest is the detail I gathered, in case it is
useful.

Versions I read:

• 0.31.1 — full C++ sources, as vendored into mlx-swift (Source/Cmlx/mlx).
  All file:line references below are to this tree unless stated otherwise.
• 0.32.0 and 0.32.1 — headers and libmlx.dylib from the published wheels. I do
  not have the .cpp files for these, so the 0.32 claims come from headers plus the
  symbol table (specifically, which symbols land in __DATA,__thread_vars /
  __DATA,__thread_bss). I have flagged each one with how it was established.

1. What the 0.31.x code does

1.1 CompilerCache has no mutex, and find() hands out a reference held across the whole trace

compile.cpp:298-384. The cache is std::unordered_map<std::uintptr_t, std::vector<CacheEntry>> cache_ (:383), reached through a function-local singleton:

// compile.cpp:386-389
CompilerCache& compiler_cache() {
  static CompilerCache compiler_cache_;
  return compiler_cache_;
}

There is no mutex in the class, and grep -n 'mutex\|lock_guard\|atomic\|thread_local' mlx/compile.cpp mlx/transforms.cpp mlx/transforms_impl.h returns nothing.

find() is a mutating lookup on both levels:

// compile.cpp:313-321
// Returns a reference to a CacheEntry which can be updated
// by the caller to avoid copying large tapes / inputs / outputs
CacheEntry& find(std::uintptr_t fun_id, ...) {
  std::vector<CacheEntry>& entries = cache_[fun_id];   // insert; may rehash
  ...
  entries.push_back(CacheEntry{stream, shapeless});    // :363 — may reallocate
  return entries.back();                               // :364
}

The comment at :313-314 is the important part: the returned reference is deliberately
long-lived. The single caller holds it across the entire compile:

// compile.cpp:1116
auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants);

and then writes and reads through it at :1119 (entry.empty), :1121-1127
(entry.constants, entry.inputs, entry.outputs, entry.extra — filled by
compile_trace, which calls back into user code), :1132 (entry.tape), :1137
(compile_simplify), :1144 (compile_fuse), and finally :1150-1153
(compile_replace(entry.tape, entry.inputs, entry.outputs, ...) and entry.extra).

Erase and clear are bare:

// compile.cpp:367-373
void erase(std::uintptr_t fun_id) { cache_.erase(fun_id); }
void clear() { cache_.clear(); }

Two things follow that are worth separating, because one of them is often assumed away.
A rehash of cache_ caused by another key's insertion does not invalidate the
std::vector<CacheEntry>& — unordered_map rehashing invalidates iterators but not
references to elements. The real hazards are: (i) concurrent unsynchronized modification
of the map is a data race in its own right regardless of reference stability; (ii)
entries.push_back at :363 reallocating the vector invalidates a CacheEntry& another
thread is compiling into; (iii) erase(fun_id) destroys that vector outright, so an
in-flight CacheEntry& becomes a reference to a destroyed object.

One more detail: the cache key includes default_stream(default_device()) (:346), and
at 0.31.x the default stream lives in the process-global Scheduler singleton
(scheduler.h:159, reached via scheduler.cpp:9-15), so it is shared too.

1.2 Tracing state is process-global, not per-thread

// transforms.cpp:45-50
std::vector<std::pair<char, char>>& detail::InTracing::trace_stack() {
  static std::vector<std::pair<char, char>> trace_stack_;   // not thread_local
  return trace_stack_;
}
int detail::InTracing::grad_counter{0};                      // not thread_local
int detail::RetainGraph::tracing_counter{0};                 // not thread_local

InTracing's constructor and destructor push_back/pop_back that vector and adjust
grad_counter (transforms_impl.h:24-32), so two threads entering any transform at the
same time race on a std::vector and on a plain int. The destructor also reads
trace_stack().back().second (transforms_impl.h:30) before popping, which is only
correct if the stack is that thread's own.

in_dynamic_tracing() makes the assumption explicit — transforms_impl.h:37-40 reads
trace_stack().front().first under the comment "compile is always and only the outer-most
transform". With a shared stack, front() may be a different thread's entry.

Everything that reads this state is therefore affected by what other threads are doing.
The complete set in 0.31.1:

And the writers: compile_trace (compile.cpp:397), vjp (transforms.cpp:333), jvp
(:531), vmap_trace (:699), and RetainGraph around each primitive().vjp(...)
(:478).

The fast.cpp:825 row is the one I would point at first, because it needs no data race at
all to misbehave: one thread computing gradients flips is_training for every other
thread's attention call, deterministically, for as long as the gradient trace is alive.

1.3 The compile() deleter erases the cache unlocked, from an arbitrary thread

// compile.cpp:1206-1212
auto pfun = std::shared_ptr<std::function<...>>(
    new std::function<...>{fun},
    [](auto* p) {
      detail::compile_erase(reinterpret_cast<std::uintptr_t>(p));
      delete p;
    });

compile_erase is compiler_cache().erase(fun_id) (:1184-1186). The deleter runs
wherever the last shared_ptr reference happens to drop, which is not necessarily the
thread that compiled, and takes no lock — so it is an unsynchronized unordered_map::erase
against concurrent find(). It is also the one erase site a binding cannot wrap, since it
is internal to mlx::core::compile().

2. What this means for an embedder

The race. Two threads calling different compiled functions for the first time both
enter CompilerCache::find, both execute cache_[fun_id] at :321, and modify the same
unordered_map concurrently — bucket array, size, and node links, with no happens-before
between them. That alone is undefined behavior; in practice it shows up as a lost entry,
a corrupted bucket chain, or a crash inside __hash_table::__do_rehash. Two threads
calling the same compiled function for the first time are worse: both can observe
entry.empty == true at :1119, both set it false, and both then write entry.inputs,
entry.outputs, entry.tape and entry.extra through the same reference while
compile_trace interleaves their graph construction. And because find() returns a
reference the caller keeps for the whole of compile_trace → compile_dfs →
compile_simplify → compile_fuse → compile_replace, a concurrent push_back at :363
or an erase from §1.3 can move or destroy the entry out from under a compile that is
still filling it. The tracing stack races independently and for the same reason. None of
this is a narrow interleaving — the compile path is long, and the window is the whole of
it.

The serialization it forces. Because the tracing flag is global rather than per-thread,
a binding cannot fix this by locking only the cache lookup: while any thread is inside
compile_trace, in_tracing() is true for every thread, so every other thread's eval
must also be excluded or it will retain graphs it should detach (transforms.cpp:144-147)
and may throw (:97-106). The only sound discipline available to a binding is a single
process-wide lock held across the entire compiled call and every eval/async_eval.
That is what mlx-swift does — its evalLock is exactly this — and the cost is that no two
compiled calls, and no compiled call and any eval, can ever overlap, on any stream or
device. That ceiling is the real consequence: not a rare crash, but a structural cap on
concurrent multi-stream use that the binding has no way to lift on its own. It also has a
sharp edge. A companion pull request against ml-explore/mlx-swift describes a permanent
deadlock in that package: the per-CompiledFunction NSLock and the global evalLock
end up acquired in opposite orders on different threads, because compile_trace calls back
into Swift while evalLock is held and the callee is another compiled function that takes
its own instance lock first. Every MLXNN activation is a file-scope compiled function, so
any traced body that calls silu(x) is a nesting site; two threads, one nested and one
plain, wedge on first use. That PR's fix — one global lock order, evalLock outermost —
resolves the deadlock but deliberately keeps the full serialization, precisely because
§1.2 means it cannot be narrowed from the Swift side.

3. What 0.32.0 and 0.32.1 appear to have changed

This is the part I would most like confirmed, since I am reading headers and symbol tables
rather than sources.

0.32.0

transforms_impl.h, diffed against 0.31.1 — two one-word changes and one addition:

-  static int grad_counter;                 (InTracing)
+  static thread_local int grad_counter;
-  static int tracing_counter;              (RetainGraph)
+  static thread_local int tracing_counter;
+  struct InExportTracing { ... static int& counter(); };

trace_stack()'s declaration is unchanged, so the header cannot say whether its local
static became thread-local — but the binary can. In 0.32.0's libmlx.dylib, all of these
are in __DATA,__thread_vars / __DATA,__thread_bss:

__ZZN3mlx4core6detail9InTracing11trace_stackEvE12trace_stack_
__ZN3mlx4core6detail9InTracing12grad_counterE
__ZN3mlx4core6detail11RetainGraph15tracing_counterE
__ZZN3mlx4core6detail15InExportTracing7counterEvE8counter_
__ZZN3mlx4core6detail14compiler_cacheEvE15compiler_cache_

So in 0.32.0 the tracing stack, both counters, and the compiler cache singleton itself
are thread-local. For comparison, the only thread_local anywhere in 0.31.1's core outside
the backends is in the bundled pocketfft.

The same symbol dump shows default_stream_storage and stream_from_thread_local_stream
in __thread_vars as well, i.e. default streams became per-thread too — which matters here
because the cache key includes the default stream (compile.cpp:346).

0.32.1

compile_impl.h changes shape:

-MLX_API void compile_erase(std::uintptr_t fun_id);
-MLX_API bool compile_cache_empty();
-MLX_API void compile_clear_cache();
+// Get the compiler cache of current thread.
+class CompileCache;
+using CompileCacheWeakPtr = std::weak_ptr<CompileCache>;
+MLX_API CompileCacheWeakPtr compile_cache();
+MLX_API void compile_erase(const CompileCacheWeakPtr& cache, std::uintptr_t fun_id);
+MLX_API void compile_clear_cache(const CompileCacheWeakPtr& cache);

and the 0.32.1 symbol table fills in the rest: CompilerCache is renamed
detail::CompileCache; the raw per-thread accessor is detail::compile_cache_unsafe()
whose local static cache is thread-local; the map is now
unordered_map<uintptr_t, shared_ptr<vector<CacheEntry>>>; the cache object is
make_shared<CompileCache>(); and find returns a
tuple<CacheEntry&, shared_ptr<vector<CacheEntry>>> — a reference plus an owning handle
that keeps the entry vector alive for as long as the caller holds it.

Between them those two releases look like they close all three items in §1: the race is
gone because nothing is shared; the reference-across-the-trace problem is handled by
returning the owning shared_ptr alongside the reference; and the deleter problem is
handled by having it capture a weak_ptr to the cache of the thread that compiled, so it
erases from the right cache or from none at all if that thread is gone.

What I would like to know

1. Is this intended and supported? Concretely: may a binding now assume that
   compile() and the transforms are safe to call from multiple threads concurrently, or
   is thread-locality here an implementation detail that could be reverted? Bindings are
   holding process-wide locks on the strength of the old behavior, and will not drop them
   without something to point at.
2. Documentation. There is nothing in the docs stating which entry points are
   thread-safe. Even one paragraph — "the compiler cache and tracing state are per-thread;
   array and the scheduler are not / are" — would let embedders size their locks
   correctly. compile_cache_unsafe() suggests the distinction is already understood
   internally.
3. Cost of per-thread caches. Thread-local means no reuse across threads: N worker
   threads pay N full traces (compile_trace + compile_dfs + compile_simplify +
   compile_fuse) and hold N tapes. For an embedder whose compiled functions are
   process-wide singletons — mlx-swift's MLXNN activations are file-scope globals — that
   is N× first-call latency and N× tape residency. Was a shared cache under a mutex
   considered and rejected on complexity grounds, or is a shared cache still open?
4. Per-thread default streams. Since 0.32 the default stream is thread-local and is
   part of the cache key, so two threads calling the same compiled function land on
   different streams by default. That seems right, but it is a behavior change worth a
   release note.
5. Backport. 0.31.x is what a number of downstreams are pinned to, including
   mlx-swift, which vendors 0.31.1. Is a note on the 0.31 line worth it, or is "upgrade
   to 0.32.1" the answer?

4. Fix options, if any of the above is still open

Recorded mainly for the shared-cache question in (3).

A. A mutex on CompilerCache. The obvious shape, and the one that preserves
cross-thread reuse. The complication is find()'s contract: it returns a CacheEntry&
that the caller holds through the entire compile, so a std::lock_guard inside find()
protects nothing that matters. Three ways out, in increasing order of appeal:

• Return a copy. Defeats the stated purpose of the reference (compile.cpp:313-314) —
  copying tapes, inputs and outputs is exactly what it exists to avoid.
• Return a locked handle. find() returns an RAII object owning the lock; the caller
  holds it to :1153. Correct and small, but it makes the whole compile a critical
  section, so concurrent first-use compiles of different functions serialize — better
  than today, still a bottleneck, and it needs care against re-entrancy (a traced body
  that calls another compiled function; note :1108-1113 already short-circuits that case
  before touching the cache, so a recursive mutex may not even be needed, but the
  interaction should be deliberate rather than incidental).
• Lock only the lookup, make entries independently owned. Guard cache_ with a mutex
  and store shared_ptr<vector<CacheEntry>> as the mapped type, so the map can rehash and
  erase can run while an outstanding entry stays alive. This is, as far as I can tell,
  exactly the ownership model 0.32.1 adopted — just with a mutex over one shared map
  instead of a thread-local map per thread. A per-entry once_flag or state enum would
  then be needed so that two threads racing on the same cold entry produce one trace rather
  than two.

B. thread_local tracing stack. Independent of A and, I would argue, necessary
regardless: without it the serialization in §2 cannot be narrowed even if the cache itself
is made safe, because in_tracing() leaks across threads. The things that assume
process-global tracing are enumerated in the table in §1.2; reading them again with
thread_local in mind, none of them wants the global reading — is_tracer(),
in_grad_tracing() and in_dynamic_tracing() are all asking "am I under a transform?",
and every one of them gets a wrong answer when the answer comes from another thread. The
only care needed is that a transform's trace and the evals it triggers stay on one thread,
which they do: compile_trace (compile.cpp:397), vjp (transforms.cpp:333), jvp
(:531) and vmap_trace (:699) all scope InTracing to a single synchronous call
frame. This is what 0.32.0 did.

C. Lifetime-safe erase. Whatever A resolves to, the deleter at compile.cpp:1209-1212
must take the same lock, or hold a weak handle to the cache it is erasing from. It is the
one path a binding cannot wrap. 0.32.1's weak_ptr<CompileCache> handles it.

D. Can Python users hit any of this? Honestly: mostly no on a GIL build, and the
reason is worth stating because it explains why this surfaced in a Swift binding rather
than in Python. In 0.31.1's bindings, the compile path never releases the GIL —
PyCompiledFun::call_impl calls mx::detail::compile(...)(inputs) directly
(python/src/transforms.cpp:558), and ~PyCompiledFun calls compile_erase
(:580-583) with the GIL held as usual for a destructor. So §1.1 and §1.3 are effectively
serialized by the interpreter. mx.eval and mx.async_eval do release it
(python/src/transforms.cpp:1171, :1191), so §1.2 is reachable: a thread computing
gradients holds grad_counter > 0 across its Python callback, and any other thread that
calls mx.fast.scaled_dot_product_attention in that window gets is_training == true
(fast.cpp:825) and a different kernel. That one needs no data race and no free-threading.
On a free-threaded build (3.13t / 3.14t) the GIL argument disappears entirely and the
compile-path races become reachable directly — so if free-threaded wheels are planned, the
0.32 work is load-bearing for them. I could not check whether such wheels are published;
the wheel I have is cp312-cp312-macosx_26_0_arm64.

5. Reproduction

None of the below has been run — they are sketches for someone with a build. All are
written to fail on 0.31.x and pass on 0.32.x, so they double as a check on §3.

R1 — the cache race (TSan). Two threads, two distinct fun_ids, so both take the
cache_[fun_id] insert path at compile.cpp:321; a fresh input shape each iteration
forces the cold path at :1119 every time.

// build against mlx with -fsanitize=thread
#include <thread>
#include <vector>
#include "mlx/mlx.h"
namespace mx = mlx::core;

std::vector<mx::array> f(const std::vector<mx::array>& x) {
  return {mx::exp(mx::negative(mx::sum(x[0])))};
}
std::vector<mx::array> g(const std::vector<mx::array>& x) {
  return {mx::square(mx::add(x[0], x[0]))};
}

int main() {
  auto cf = mx::compile(f);
  auto cg = mx::compile(g);
  auto run = [](auto fn, float v) {
    for (int i = 1; i <= 128; ++i) {
      auto x = mx::full({i}, v, mx::float32);   // new shape => new CacheEntry
      mx::eval(fn({x}));
    }
  };
  std::thread t1(run, cf, 1.0f), t2(run, cg, 2.0f);
  t1.join(); t2.join();
}

Expected on 0.31.x: TSan races on CompilerCache::cache_ and on
InTracing::trace_stack_; without TSan, intermittent crashes in the hash table. Expected
on 0.32.x: clean.

R2 — the erase/find interleaving. Same shape, but thread B churns compiled-function
lifetimes so compile_erase runs concurrently with thread A's find:

std::thread a([]{ /* R1's `run`, one function, 500 iterations */ });
std::thread b([]{
  for (int i = 0; i < 500; ++i) {
    auto h = mx::compile([](const std::vector<mx::array>& x) {
      return std::vector<mx::array>{mx::abs(x[0])};
    });                                   // lambda => no addressable target
    mx::eval(h({mx::full({4}, 1.0f, mx::float32)}));
  }                                       // dtor => shared_ptr deleter => compile_erase
});

The lambda matters: get_function_address (compile.cpp:288-296) returns 0 for it, which
is what selects the shared_ptr-deleter branch at :1205-1220.

R3 — the tracing flag, no sanitizer needed. This one is a plain wrong answer rather
than UB, and it reproduces from Python on a stock GIL build:

import threading, mlx.core as mx

stop = False

def grads():                       # holds grad_counter > 0 across its Python callback
    f = lambda x: mx.sum(mx.exp(x))
    x = mx.random.normal((256, 256)); mx.eval(x)
    while not stop:
        mx.eval(mx.grad(f)(x))

t = threading.Thread(target=grads); t.start()
# meanwhile, on the main thread, call mx.fast.scaled_dot_product_attention and
# observe `is_training` (fast.cpp:825) flip with the other thread's activity --
# a different kernel is selected and logsumexp is produced, for an inference-only call.
stop = True; t.join()

Expected on 0.31.x: the attention call's behavior depends on whether the other thread is
mid-gradient. Expected on 0.32.x: no dependence, since grad_counter is thread-local.

Appendix — how the 0.32 claims were established

I did not have 0.32 sources, so for the record:

• transforms_impl.h and compile_impl.h were diffed directly, 0.31.1 → 0.32.0 → 0.32.1.
• Thread-locality was read off nm -m libmlx.dylib: a thread_local function-local static
  lands in __DATA,__thread_vars with a $tlv$init counterpart in __DATA,__thread_bss,
  and a non-thread-local one does not. 0.32.0 and 0.32.1 both place trace_stack_,
  grad_counter, tracing_counter, InExportTracing::counter_, the compiler-cache
  singleton, and the default-stream storage there; 0.31.1's sources contain no
  thread_local in these files at all.
• The 0.32.1 cache redesign was read from demangled symbol names —
  detail::CompileCache::find(...), detail::compile_erase(std::weak_ptr<CompileCache> const&, unsigned long), detail::compile_cache_unsafe(), and the
  unordered_map<unsigned long, shared_ptr<vector<CompileCache::CacheEntry>>>
  instantiations — plus the header declarations quoted above. Anything I could not
  establish this way is phrased as a question, not a claim.