fix

taichi-dev · Nov 9, 2022 · 3fcac2c · 3fcac2c
1 parent b5e661d
commit 3fcac2c
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/taichi/program/program.h b/taichi/program/program.h
@@ -6,6 +6,7 @@
 #include <optional>
 #include <atomic>
 #include <stack>
+#include <shared_mutex>
 
 #define TI_RUNTIME_HOST
 #include "taichi/aot/module_builder.h"
@@ -126,9 +127,11 @@ class TI_DLL_EXPORT Program {
   CompileConfig &this_thread_config() {
     auto thread_id = std::this_thread::get_id();
     if (!configs.count(thread_id)) {
-      std::lock_guard<std::mutex> _(compile_config_mut);
+      std::unique_lock<std::shared_mutex> _(compile_config_mut);
       configs[thread_id] = configs[main_thread_id_];
+      return configs[thread_id];
     }
+    std::shared_lock<std::shared_mutex> _(compile_config_mut);
     return configs[thread_id];
   }
 
@@ -397,7 +400,7 @@ class TI_DLL_EXPORT Program {
   std::vector<std::unique_ptr<Ndarray>> ndarrays_;
   std::vector<std::unique_ptr<Texture>> textures_;
 
-  std::mutex compile_config_mut;
+  mutable std::shared_mutex compile_config_mut;
 };
 
 }  // namespace taichi::lang