From 805b1b6eb19f26f6ce56470866dcd739cc909a36 Mon Sep 17 00:00:00 2001
From: Dan Yeh <chenchen.yeh0423@gmail.com>
Date: Tue, 23 Sep 2025 17:44:03 +0200
Subject: [PATCH] add mutex and guard `cache_` access

---
 mlx/compile.cpp | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mlx/compile.cpp b/mlx/compile.cpp
index 91743ec04b..fbe83e9b85 100644
--- a/mlx/compile.cpp
+++ b/mlx/compile.cpp
@@ -1,6 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 #include <cstdlib>
 #include <map>
+#include <mutex>
 #include <sstream>
 #include <unordered_map>
 #include <unordered_set>
@@ -305,6 +306,7 @@ class CompilerCache {
       const std::vector<array>& inputs,
       bool shapeless,
       const std::vector<uint64_t>& constants) {
+    std::lock_guard<std::mutex> guard(mtx_);
     // Find the cache entries for |fun_id|.
     std::vector<CacheEntry>& entries = cache_[fun_id];
 
@@ -353,10 +355,12 @@ class CompilerCache {
   }
 
   void erase(std::uintptr_t fun_id) {
+    std::lock_guard<std::mutex> guard(mtx_);
     cache_.erase(fun_id);
   }
 
   void clear() {
+    std::lock_guard<std::mutex> guard(mtx_);
     cache_.clear();
   }
 
@@ -368,6 +372,7 @@ class CompilerCache {
   }
 
   friend CompilerCache& compiler_cache();
+  std::mutex mtx_;
   std::unordered_map<std::uintptr_t, std::vector<CacheEntry>> cache_;
 };