From 269c88026d457bad62b3f6081bd934e9e8f41527 Mon Sep 17 00:00:00 2001 From: yentur Date: Sun, 9 Aug 2026 14:06:33 +0300 Subject: [PATCH 1/3] Fix stale compile cache entry when released on another thread The compile cache is thread local, but compile_erase runs on whichever thread drops the last reference to a compiled function. When that is not the thread that traced it, the erase hits the wrong cache and the entry stays behind. It does not only leak. fun_id is the address of the callable and the cache matches on it plus shapes, dtypes, stream and constants, so once the address is reused a later compile of an unrelated function can match the stale entry and get the dead function's tape. Register the caches process wide so an erase can reach them all. Threads other than the caller are handed the id instead of having their cache touched from the outside, and they apply it the next time they use the cache. The check for pending ids is a single atomic load, so compiled calls stay lock free. --- mlx/compile.cpp | 64 ++++++++++++++++++++++++++++++++++++ python/tests/test_compile.py | 43 ++++++++++++++++++++++++ 2 files changed, 107 insertions(+) diff --git a/mlx/compile.cpp b/mlx/compile.cpp index 12d7397be4..0fea253a42 100644 --- a/mlx/compile.cpp +++ b/mlx/compile.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -320,6 +321,8 @@ class CompilerCache { const std::vector& inputs, bool shapeless, const std::vector& constants) { + apply_pending_erases(); + // Find the cache entries for |fun_id|. std::vector& entries = cache_[fun_id]; @@ -368,14 +371,32 @@ class CompilerCache { } void erase(std::uintptr_t fun_id) { + // The function may have been traced on other threads, and only the thread + // that owns a cache is allowed to touch it, so hand the id over and let + // them erase it the next time they use their cache. + { + std::lock_guard lock(caches_mutex()); + for (auto* cache : caches()) { + if (cache != this) { + cache->pending_erases_.push_back(fun_id); + cache->has_pending_erases_.store(true, std::memory_order_release); + } + } + } cache_.erase(fun_id); } void clear() { + { + std::lock_guard lock(caches_mutex()); + pending_erases_.clear(); + has_pending_erases_.store(false, std::memory_order_release); + } cache_.clear(); } bool empty() { + apply_pending_erases(); return cache_.empty(); } @@ -384,10 +405,53 @@ class CompilerCache { // Make sure the allocator is fully // initialized before the compiler cache allocator::allocator(); + + std::lock_guard lock(caches_mutex()); + caches().insert(this); + } + + ~CompilerCache() { + std::lock_guard lock(caches_mutex()); + caches().erase(this); + } + + // Erase the ids handed over by other threads. Only the owning thread calls + // this so the common case of nothing pending takes no lock. + void apply_pending_erases() { + if (!has_pending_erases_.load(std::memory_order_acquire)) { + return; + } + std::vector fun_ids; + { + std::lock_guard lock(caches_mutex()); + fun_ids.swap(pending_erases_); + has_pending_erases_.store(false, std::memory_order_release); + } + for (auto fun_id : fun_ids) { + cache_.erase(fun_id); + } + } + + // The caches are thread local, but a compiled function can be released on + // a thread other than the one that traced it. Registering them process wide + // lets such an erase reach the thread holding the entry. Both are used by + // the constructor so they outlive every cache. + static std::mutex& caches_mutex() { + static std::mutex mutex; + return mutex; + } + + static std::unordered_set& caches() { + static std::unordered_set caches; + return caches; } friend CompilerCache& compiler_cache(); std::unordered_map> cache_; + + // Guarded by caches_mutex(), the flag allows checking it without the lock. + std::vector pending_erases_; + std::atomic has_pending_erases_{false}; }; CompilerCache& compiler_cache() { diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 7a2c6b9d0d..090f45f222 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -98,6 +98,49 @@ def worker(): raise errors[0] self.assertEqual(results, [(2.0, 2.0)] * 3) + def test_compile_release_on_another_thread(self): + # A function traced on one thread but released on another must still + # drop its cache entry, otherwise a later compile of the same id gets + # handed the dead function's tape instead of being traced again. + traces = [] + + def fun(x): + traces.append(1) + return x + 1 + + holder = {} + traced = threading.Event() + released = threading.Event() + errors = [] + + def worker(): + try: + holder["fn"] = mx.compile(fun) + mx.eval(holder["fn"](mx.array([1.0]))) + traced.set() + self.assertTrue(released.wait(10)) + # The same callable, so the same id. + fn = mx.compile(fun) + mx.eval(fn(mx.array([1.0]))) + except Exception as e: + errors.append(e) + finally: + traced.set() + + # The tracing thread has to outlive the release, on exit it would tear + # down its cache anyway. + thread = threading.Thread(target=worker) + thread.start() + self.assertTrue(traced.wait(10)) + holder.clear() + gc.collect() + released.set() + thread.join() + + if errors: + raise errors[0] + self.assertEqual(len(traces), 2) + def test_compile_grad(self): def loss_fn(x): return mx.exp(x).sum() From cf6e3ac33f9e850fe269317fb9cee9a86217cb5b Mon Sep 17 00:00:00 2001 From: yentur Date: Sun, 9 Aug 2026 17:25:23 +0300 Subject: [PATCH 2/3] Hold cache entries by shared_ptr across the trace find() handed out a bare reference into the vector inside cache_, and the caller holds it across compile_trace(), which runs user code. A traced body that calls an already compiled function on constant inputs reaches find() again, and the drain there can erase the outer in-flight fun_id and free the entry the caller is still writing to. Store the entries as shared_ptr and let the caller keep one for the whole call, so an erase during a trace, re-entrant or from another thread, only drops the cache slot while the entry in use stays alive. --- mlx/compile.cpp | 32 +++++++++++++++++++------------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/mlx/compile.cpp b/mlx/compile.cpp index 0fea253a42..b4b68acc63 100644 --- a/mlx/compile.cpp +++ b/mlx/compile.cpp @@ -314,9 +314,11 @@ class CompilerCache { std::shared_ptr extra; }; - // Returns a reference to a CacheEntry which can be updated - // by the caller to avoid copying large tapes / inputs / outputs - CacheEntry& find( + // Returns a CacheEntry which can be updated by the caller to avoid copying + // large tapes / inputs / outputs. The caller keeps it alive for the whole + // call, so an erase while the trace is running, from this thread or another + // one, only drops the cache slot and leaves the entry in use intact. + std::shared_ptr find( std::uintptr_t fun_id, const std::vector& inputs, bool shapeless, @@ -324,7 +326,7 @@ class CompilerCache { apply_pending_erases(); // Find the cache entries for |fun_id|. - std::vector& entries = cache_[fun_id]; + std::vector>& entries = cache_[fun_id]; // Compare if 2 arrays have same shape and dtype. auto has_same_shape_and_dtype = [shapeless]( @@ -350,23 +352,23 @@ class CompilerCache { // - Default stream and device match the entry's default stream // - Inputs match i.e. shapes and types must be equal. auto stream = default_stream(default_device()); - for (CacheEntry& entry : entries) { + for (auto& entry : entries) { // Check that the default stream and device match - if (entry.stream != stream) { + if (entry->stream != stream) { continue; } - if (entry.shapeless != shapeless) { + if (entry->shapeless != shapeless) { continue; } // Check the inputs match and return if so - if (has_same_shape_and_dtype(inputs, entry.inputs) && - constants == entry.constants) { + if (has_same_shape_and_dtype(inputs, entry->inputs) && + constants == entry->constants) { return entry; } } // Otherwise append a new cache entry - entries.push_back(CacheEntry{stream, shapeless}); + entries.push_back(std::make_shared(stream, shapeless)); return entries.back(); } @@ -447,7 +449,8 @@ class CompilerCache { } friend CompilerCache& compiler_cache(); - std::unordered_map> cache_; + std::unordered_map>> + cache_; // Guarded by caches_mutex(), the flag allows checking it without the lock. std::vector pending_erases_; @@ -1183,8 +1186,11 @@ ArrayFnWithExtra compile( return fun(inputs); } - // Find a cache entry with the correct inputs - auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants); + // Find a cache entry with the correct inputs, held for the whole call so + // that it survives an erase while the trace below is running + auto entry_ptr = + compiler_cache().find(fun_id, inputs, shapeless, constants); + auto& entry = *entry_ptr; // No matching cache entry existed, so compile if (entry.empty) { From 6242ddb41211e7a0dcce7a5f19b2e35195af0ca4 Mon Sep 17 00:00:00 2001 From: yentur Date: Tue, 11 Aug 2026 12:01:28 +0300 Subject: [PATCH 3/3] Add C++ coverage for erase during trace Two cases in compile_tests.cpp. The first erases the outer id from a second thread while a worker is inside its trace, then has the traced body call an already compiled helper on a constant so the nested find() drains on the worker's own cache. The second traces on a worker that stays alive, erases from another thread, then reuses the id for a different function. Both tests were written by @sashko-zakharchuk and are added here with his agreement in ml-explore/mlx#4096. --- tests/compile_tests.cpp | 130 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 130 insertions(+) diff --git a/tests/compile_tests.cpp b/tests/compile_tests.cpp index 30c2f887ac..02226c126c 100644 --- a/tests/compile_tests.cpp +++ b/tests/compile_tests.cpp @@ -6,8 +6,12 @@ #include "doctest/doctest.h" #include +#include #include +#include +#include +#include "mlx/compile_impl.h" #include "mlx/mlx.h" #include "mlx/primitives.h" @@ -875,3 +879,129 @@ TEST_CASE("test compile throwing first trace does not poison cache") { REQUIRE_EQ(out.size(), 1); CHECK_EQ(out[0].item(), 3.0f); } + +TEST_CASE("test compile erase while a trace is in flight") { + // An entry handed to a caller must outlive an erase that lands while the + // caller is still filling it in. Here a second thread erases the outer id + // while the worker is inside its trace, and the traced body then invokes an + // already-compiled helper on a fresh (non-tracer) constant. That helper's + // find() drains the pending erase on the worker's own cache, which frees the + // outer entry unless the caller is holding it. + auto x = zeros({1}, float32); + eval(x); + + constexpr std::uintptr_t nested_id = 0xf00d; + constexpr std::uintptr_t outer_id = 0xbeef; + + auto nested = detail::compile( + [](const std::vector& in) { + return std::vector{in[0] + 2.0f}; + }, + nested_id); + eval(nested({zeros({1}, float32)})); + + std::mutex mtx; + std::condition_variable cv; + int stage = 0; + + auto traced = [&](const std::vector& inputs) { + { + std::lock_guard lk(mtx); + stage = 1; + } + cv.notify_one(); + { + std::unique_lock lk(mtx); + cv.wait(lk, [&stage] { return stage == 2; }); + } + // Non-tracer input, so the compiled helper runs find() here and drains the + // erase queued for the outer id above. + auto tmp = nested({zeros({1}, float32)}); + return std::vector{inputs[0] + tmp[0]}; + }; + + float result = 0.0f; + std::thread worker([&]() { + auto compiled = detail::compile(traced, outer_id); + auto outputs = compiled({x}); + eval(outputs); + result = outputs[0].item(); + }); + + { + std::unique_lock lk(mtx); + cv.wait(lk, [&stage] { return stage == 1; }); + } + detail::compile_erase(outer_id); + { + std::lock_guard lk(mtx); + stage = 2; + } + cv.notify_one(); + worker.join(); + + CHECK_EQ(result, 2.0f); +} + +TEST_CASE("test compile erase from another thread") { + // The compile cache is thread local, so an erase that reaches only the + // calling thread leaves behind the entry of a function traced elsewhere. + // Since |fun_id| is a reused address, the next compile on that id would then + // be handed the dead function's tape. Trace on a worker thread which stays + // alive, erase from this one, then reuse the id for a different function. + auto x = zeros({1}, float32); + eval(x); + + constexpr std::uintptr_t fun_id = 0xc0ffee; + auto add_one = [](const std::vector& inputs) { + return std::vector{inputs[0] + 1.0f}; + }; + auto add_two = [](const std::vector& inputs) { + return std::vector{inputs[0] + 2.0f}; + }; + + std::mutex mtx; + std::condition_variable cv; + int stage = 0; + float before = 0.0f; + float after = 0.0f; + + std::thread worker([&]() { + { + auto compiled = detail::compile(add_one, fun_id); + auto outputs = compiled({x}); + eval(outputs); + before = outputs[0].item(); + } + { + std::lock_guard lk(mtx); + stage = 1; + } + cv.notify_one(); + { + std::unique_lock lk(mtx); + cv.wait(lk, [&stage] { return stage == 2; }); + } + // Same id, different function. Without the erase reaching this thread the + // cached tape still adds one. + auto compiled = detail::compile(add_two, fun_id); + auto outputs = compiled({x}); + eval(outputs); + after = outputs[0].item(); + }); + + { + std::unique_lock lk(mtx); + cv.wait(lk, [&stage] { return stage == 1; }); + } + detail::compile_erase(fun_id); + { + std::lock_guard lk(mtx); + stage = 2; + } + cv.notify_one(); + worker.join(); + + CHECK_EQ(before, 1.0f); + CHECK_EQ(after, 2.0f); +}