Make mx.compile cache erasing thread safe

mlx/compile.cpp

@@ -3,6 +3,7 @@
 #include <atomic>
 #include <cstdlib>
 #include <map>
+#include <shared_mutex>
 #include <sstream>
 #include <unordered_map>
 #include <unordered_set>
@@ -298,7 +299,7 @@ std::uintptr_t get_function_address(const std::function<T(U...)>& fun) {
   return reinterpret_cast<std::uintptr_t>(*fun_ptr);
 }
 
-class CompilerCache {
+class CompileCache {
  public:
   struct CacheEntry {
     CacheEntry(Stream stream, bool shapeless)
@@ -313,15 +314,38 @@ 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(
+  CompileCache() {
+    // Make sure the allocator is fully initialized before the compiler cache.
+    allocator::allocator();
+  }
+
+  // Returns a reference to a CacheEntry which can be updated by the caller to
+  // avoid copying large tapes / inputs / outputs, with the shared_ptr of
+  // entries to avoid getting erased during compilation.
+  std::tuple<CacheEntry&, std::shared_ptr<std::vector<CacheEntry>>> find(
       std::uintptr_t fun_id,
       const std::vector<array>& inputs,
       bool shapeless,
       const std::vector<uint64_t>& constants) {
-    // Find the cache entries for |fun_id|.
-    std::vector<CacheEntry>& entries = cache_[fun_id];
+    // Find the cache entries for |fun_id| in a thread-safe way.
+    auto entries_ptr = [&]() {
+      // Lookup with shared lock.
+      {
+        std::shared_lock lock(mutex_);
+        auto it = cache_.find(fun_id);
+        if (it != cache_.end()) {
+          return it->second;
+        }
+      }
+      // Insertion with exclusive lock.
+      std::unique_lock lock(mutex_);
+      auto& ptr = cache_[fun_id];
+      if (!ptr) {
+        ptr = std::make_shared<std::vector<CacheEntry>>();
+      }
+      return ptr;
+    }();
+    auto& entries = *entries_ptr;
 
     // Compare if 2 arrays have same shape and dtype.
     auto has_same_shape_and_dtype = [shapeless](
@@ -359,40 +383,37 @@ class CompilerCache {
       // Check the inputs match and return if so
       if (has_same_shape_and_dtype(inputs, entry.inputs) &&
           constants == entry.constants) {
-        return entry;
+        return {entry, std::move(entries_ptr)};
       }
     }
     // Otherwise append a new cache entry
     entries.push_back(CacheEntry{stream, shapeless});
-    return entries.back();
+    return {entries.back(), std::move(entries_ptr)};
   }
 
   void erase(std::uintptr_t fun_id) {
+    std::unique_lock lock(mutex_);
     cache_.erase(fun_id);
   }
 
   void clear() {
+    std::unique_lock lock(mutex_);
     cache_.clear();
   }
 
-  bool empty() {
-    return cache_.empty();
-  }
-
  private:
-  CompilerCache() {
-    // Make sure the allocator is fully
-    // initialized before the compiler cache
-    allocator::allocator();
-  }
-
-  friend CompilerCache& compiler_cache();
-  std::unordered_map<std::uintptr_t, std::vector<CacheEntry>> cache_;
+  // The cache may get its key erased from a separate thread, but its value is
+  // only added and modified in the thread of creation.
+  // Put value in a shared_ptr to avoid race condition when erasing happened
+  // during compilation for the same function.
+  std::unordered_map<std::uintptr_t, std::shared_ptr<std::vector<CacheEntry>>>
+      cache_;
+  std::shared_mutex mutex_;
 };
 
-CompilerCache& compiler_cache() {
-  static thread_local CompilerCache compiler_cache_;
-  return compiler_cache_;
+std::shared_ptr<CompileCache>& compile_cache_unsafe() {
+  static thread_local auto cache = std::make_shared<CompileCache>();
+  return cache;
 }
 
 std::tuple<std::vector<array>, std::vector<array>, std::shared_ptr<void>>
@@ -1120,7 +1141,9 @@ ArrayFnWithExtra compile(
     }
 
     // Find a cache entry with the correct inputs
-    auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants);
+    auto [entry, entries_ptr] =
+        compile_cache_unsafe()->find(fun_id, inputs, shapeless, constants);
+    static_assert(std::is_reference_v<decltype(entry)>);
 
     // No matching cache entry existed, so compile
     if (entry.empty) {
@@ -1192,16 +1215,20 @@ std::function<std::vector<array>(const std::vector<array>&)> compile(
   };
 }
 
-void compile_erase(std::uintptr_t fun_id) {
-  detail::compiler_cache().erase(fun_id);
+CompileCacheWeakPtr compile_cache() {
+  return compile_cache_unsafe();
 }
 
-void compile_clear_cache() {
-  detail::compiler_cache().clear();
+void compile_erase(const CompileCacheWeakPtr& cache, std::uintptr_t fun_id) {
+  if (auto p = cache.lock()) {
+    p->erase(fun_id);
+  }
 }
 
-bool compile_cache_empty() {
-  return detail::compiler_cache().empty();
+void compile_clear_cache(const CompileCacheWeakPtr& cache) {
+  if (auto p = cache.lock()) {
+    p->clear();
+  }
 }
 
 } // namespace detail
@@ -1221,8 +1248,8 @@ std::function<std::vector<array>(const std::vector<array>&)> compile(
     auto pfun = std::shared_ptr<
         std::function<std::vector<array>(const std::vector<array>&)>>(
         new std::function<std::vector<array>(const std::vector<array>&)>{fun},
-        [](auto* p) {
-          detail::compile_erase(reinterpret_cast<std::uintptr_t>(p));
+        [cache = detail::compile_cache()](auto* p) {
+          detail::compile_erase(cache, reinterpret_cast<std::uintptr_t>(p));
           delete p;
         });
     fun_id = reinterpret_cast<std::uintptr_t>(pfun.get());

mlx/compile_impl.h

@@ -27,15 +27,19 @@ MLX_API ArrayFnWithExtra compile(
     bool shapeless,
     std::vector<uint64_t> constants);
 
-// Erase cached compile functions
-MLX_API void compile_erase(std::uintptr_t fun_id);
+// Get the compiler cache of current thread.
+class CompileCache;
+using CompileCacheWeakPtr = std::weak_ptr<CompileCache>;
+MLX_API CompileCacheWeakPtr compile_cache();
+
+// Erase cached compile function.
+MLX_API void compile_erase(
+    const CompileCacheWeakPtr& cache,
+    std::uintptr_t fun_id);
 
 // Clear the compiler cache causing a recompilation of all compiled functions
 // when called again.
-MLX_API void compile_clear_cache();
-
-// Return true if the cache is empty.
-MLX_API bool compile_cache_empty();
+MLX_API void compile_clear_cache(const CompileCacheWeakPtr& cache);
 
 bool compile_available_for_device(const Device& device);
 

mlx/scheduler.cpp

@@ -3,6 +3,7 @@
 #include "mlx/scheduler.h"
 #include "mlx/backend/cpu/eval.h"
 #include "mlx/backend/gpu/eval.h"
+#include "mlx/compile_impl.h"
 #include "mlx/utils.h"
 
 namespace mlx::core {
@@ -27,6 +28,7 @@ void synchronize() {
 }
 
 void clear_streams() {
+  detail::compile_clear_cache(detail::compile_cache());
   cpu::clear_streams();
   gpu::clear_streams();
 }

python/src/random.cpp

@@ -64,6 +64,10 @@ PyKeySequence& default_key() {
   return ks;
 }
 
+void reset_random_state() {
+  default_key().reset();
+}
+
 // A process-global sentinel for `mx.random.state`. Since it is the same object
 // on every thread, capturing it (e.g. with `mx.compile`) is thread-independent;
 // the pytree traversal in trees.cpp resolves it to the calling thread's key.

python/src/random.h

@@ -9,6 +9,9 @@
 namespace mx = mlx::core;
 namespace nb = nanobind;
 
+// Clear the `mx.random.state` python object in current thread.
+void reset_random_state();
+
 // The process-global `mx.random.state` sentinel.
 nb::object random_state_sentinel();
 

python/src/stream.cpp

@@ -9,6 +9,7 @@
 
 #include "mlx/stream.h"
 #include "mlx/utils.h"
+#include "python/src/random.h"
 
 namespace mx = mlx::core;
 namespace nb = nanobind;
@@ -137,7 +138,10 @@ void init_stream(nb::module_& m) {
       R"pbdoc(Make a new stream that will be unique per thread.)pbdoc");
   m.def(
       "clear_streams",
-      &mx::clear_streams,
+      []() {
+        reset_random_state();
+        mx::clear_streams();
+      },
       R"pbdoc(Destroy all streams created in current thread.)pbdoc");
 
   nb::class_<PyStreamContext>(m, "StreamContext", R"pbdoc(

python/src/transforms.cpp

@@ -406,29 +406,13 @@ auto py_vmap(
   };
 }
 
-void ensure_compile_cache_cleanup() {
-  // Make sure each thread using mx.compile would clear its compile cache
-  // before python interpreter exits.
-  struct ThreadCleanup {
-    ~ThreadCleanup() {
-      if (!mx::detail::compile_cache_empty()) {
-        nb::gil_scoped_acquire gil;
-        mx::detail::compile_clear_cache();
-      }
-    }
-  };
-  static thread_local auto clear_cache = []() {
-    mx::detail::compile_clear_cache();
-    return ThreadCleanup{};
-  }();
-}
-
 struct PyCompiledFun {
   nb::callable fun;
   std::uintptr_t fun_id;
   nb::object captured_inputs;
   nb::object captured_outputs;
   bool shapeless;
+  mx::detail::CompileCacheWeakPtr cache;
 
   // Data to attach to the compiled function that contains the python output
   // structure and the number of arrays in said structure.
@@ -456,15 +440,16 @@ struct PyCompiledFun {
   PyCompiledFun& operator=(PyCompiledFun&& other) = delete;
   PyCompiledFun(PyCompiledFun&& other)
       : fun(std::move(other.fun)),
-        fun_id(reinterpret_cast<std::uintptr_t>(fun.ptr())) {
+        fun_id(reinterpret_cast<std::uintptr_t>(fun.ptr())),
+        captured_inputs(std::move(other.captured_inputs)),
+        captured_outputs(std::move(other.captured_outputs)),
+        shapeless(other.shapeless),
+        cache(other.cache) {
     other.fun_id = 0;
-    captured_inputs = std::move(other.captured_inputs);
-    captured_outputs = std::move(other.captured_outputs);
-    shapeless = other.shapeless;
   };
 
   nb::object call_impl(const nb::args& args, const nb::kwargs& kwargs) {
-    ensure_compile_cache_cleanup();
+    cache = mx::detail::compile_cache();
 
     // Flat array inputs
     std::vector<mx::array> inputs;
@@ -599,7 +584,7 @@ struct PyCompiledFun {
   ~PyCompiledFun() {
     nb::gil_scoped_acquire gil;
 
-    mx::detail::compile_erase(fun_id);
+    mx::detail::compile_erase(cache, fun_id);
     fun.reset();
     captured_inputs.reset();
     captured_outputs.reset();
@@ -1553,10 +1538,4 @@ void init_transforms(nb::module_& m) {
           A callable that recomputes intermediate states during gradient
           computation.
       )pbdoc");
-
-  // Ensure the main thread cleanup will happen before the interpreter goes
-  // away. As a result if the other threads join the main thread we should have
-  // a clean tear-down.
-  auto atexit = nb::module_::import_("atexit");
-  atexit.attr("register")(nb::cpp_function(&mx::detail::compile_clear_cache));
 }

python/tests/test_compile.py

@@ -87,6 +87,7 @@ def worker():
                 results.append((y.item(), z.item()))
             except Exception as e:
                 errors.append(e)
+            mx.clear_streams()
 
         for _ in range(3):
             thread = threading.Thread(target=worker)
@@ -98,6 +99,50 @@ 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()
+            mx.clear_streams()
+
+        # 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()
@@ -449,6 +494,7 @@ def test_compile_rng_across_threads(self):
 
         def grab():
             state_from_thread["s"] = mx.random.state
+            mx.clear_streams()
 
         t = threading.Thread(target=grab)
         t.start()
@@ -482,6 +528,7 @@ def worker():
                     results["seed_changes"] = not bool(
                         mx.allclose(c, e, 1e-2, 1e-2).item()
                     )
+                mx.clear_streams()
 
             t = threading.Thread(target=worker)
             t.start()