Fix stale compile cache entry when released on another thread

mlx/compile.cpp

@@ -3,6 +3,7 @@
 #include <atomic>
 #include <cstdlib>
 #include <map>
+#include <mutex>
 #include <sstream>
 #include <unordered_map>
 #include <unordered_set>
@@ -313,15 +314,19 @@ class CompilerCache {
     std::shared_ptr<void> 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<CacheEntry> find(
       std::uintptr_t fun_id,
       const std::vector<array>& inputs,
       bool shapeless,
       const std::vector<uint64_t>& constants) {
+    apply_pending_erases();
+
     // Find the cache entries for |fun_id|.
-    std::vector<CacheEntry>& entries = cache_[fun_id];
+    std::vector<std::shared_ptr<CacheEntry>>& 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<CacheEntry>(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<std::mutex> 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<std::mutex> 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<std::mutex> lock(caches_mutex());
+    caches().insert(this);
+  }
+
+  ~CompilerCache() {
+    std::lock_guard<std::mutex> 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<std::uintptr_t> fun_ids;
+    {
+      std::lock_guard<std::mutex> 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<CompilerCache*>& caches() {
+    static std::unordered_set<CompilerCache*> caches;
+    return caches;
   }
 
   friend CompilerCache& compiler_cache();
-  std::unordered_map<std::uintptr_t, std::vector<CacheEntry>> cache_;
+  std::unordered_map<std::uintptr_t, std::vector<std::shared_ptr<CacheEntry>>>
+      cache_;
+
+  // Guarded by caches_mutex(), the flag allows checking it without the lock.
+  std::vector<std::uintptr_t> pending_erases_;
+  std::atomic<bool> 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) {

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()

tests/compile_tests.cpp

@@ -6,8 +6,12 @@
 #include "doctest/doctest.h"
 
 #include <cmath>
+#include <condition_variable>
 #include <limits>
+#include <mutex>
+#include <thread>
 
+#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<float>(), 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<array>& in) {
+        return std::vector<array>{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<array>& inputs) {
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      stage = 1;
+    }
+    cv.notify_one();
+    {
+      std::unique_lock<std::mutex> 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<array>{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<float>();
+  });
+
+  {
+    std::unique_lock<std::mutex> lk(mtx);
+    cv.wait(lk, [&stage] { return stage == 1; });
+  }
+  detail::compile_erase(outer_id);
+  {
+    std::lock_guard<std::mutex> 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<array>& inputs) {
+    return std::vector<array>{inputs[0] + 1.0f};
+  };
+  auto add_two = [](const std::vector<array>& inputs) {
+    return std::vector<array>{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<float>();
+    }
+    {
+      std::lock_guard<std::mutex> lk(mtx);
+      stage = 1;
+    }
+    cv.notify_one();
+    {
+      std::unique_lock<std::mutex> 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<float>();
+  });
+
+  {
+    std::unique_lock<std::mutex> lk(mtx);
+    cv.wait(lk, [&stage] { return stage == 1; });
+  }
+  detail::compile_erase(fun_id);
+  {
+    std::lock_guard<std::mutex> lk(mtx);
+    stage = 2;
+  }
+  cv.notify_one();
+  worker.join();
+
+  CHECK_EQ(before, 1.0f);
+  CHECK_EQ(after, 2.0f);
+}