diff --git a/mlx/compile.cpp b/mlx/compile.cpp index 12d7397be4..b4b68acc63 100644 --- a/mlx/compile.cpp +++ b/mlx/compile.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -313,15 +314,19 @@ 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, const std::vector& constants) { + 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]( @@ -347,35 +352,53 @@ 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(); } 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 +407,54 @@ 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_; + 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() { @@ -1119,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) { 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() 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); +}