Revive shared mutex lockyouki/static-code-cache

7 files changed, 47 insertions, 50 deletions

diff --git a/src/CodeCache.h b/src/CodeCache.h
index f1aa003..8eb28d8 100644
--- a/src/CodeCache.h
+++ b/src/CodeCache.h
@@ -9,17 +9,17 @@
 #include <future>
 #include <map>
 
-//#if __cplusplus >= 201402L && !defined(__APPLE__)
-//// For C++14, use shared_timed_mutex.
-//// some macOS C++14 compilers don't support shared_timed_mutex.
-//#define FBGEMM_USE_SHARED_TIMED_MUTEX
-//#endif
+#if __cplusplus >= 201402L && !defined(__APPLE__)
+// For C++14, use shared_timed_mutex.
+// some macOS C++14 compilers don't support shared_timed_mutex.
+#define FBGEMM_USE_SHARED_TIMED_MUTEX
+#endif
 
-//#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
-//#include <shared_mutex>
-//#else
-//#include <mutex>
-//#endif
+#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
+#include <shared_mutex>
+#else
+#include <mutex>
+#endif
 
 namespace fbgemm {
 
@@ -31,12 +31,12 @@ namespace fbgemm {
 template <typename KEY, typename VALUE>
 class CodeCache {
  private:
-  static thread_local std::map<KEY, std::shared_future<VALUE>> values_;
-//#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
-//  std::shared_timed_mutex mutex_;
-//#else
-//  std::mutex mutex_;
-//#endif
+  std::map<KEY, std::shared_future<VALUE>> values_;
+#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
+  std::shared_timed_mutex mutex_;
+#else
+  std::mutex mutex_;
+#endif
 
  public:
   CodeCache(const CodeCache&) = delete;
@@ -51,36 +51,36 @@ class CodeCache {
 
     // Check for existance of the key
     {
-//#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
-//      //mutex_.lock_shared();
-//#else
-//      std::unique_lock<std::mutex> lock(mutex_);
-//#endif
+#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
+      mutex_.lock_shared();
+#else
+      std::unique_lock<std::mutex> lock(mutex_);
+#endif
 
       auto it = values_.find(key);
       if (it != values_.end()) {
         returnFuture = it->second;
-//#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
-//        //mutex_.unlock_shared();
-//#endif
+#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
+        mutex_.unlock_shared();
+#endif
       } else {
-//#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
-        //mutex_.unlock_shared();
+#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
+        mutex_.unlock_shared();
 
-        //mutex_.lock();
+        mutex_.lock();
         // Need to look up again because there could be race condition from
         // the time gap between mutex_.unlock_shared() and mutex_.lock()
-        //it = values_.find(key);
-        //if (it == values_.end()) {
-//#endif
+        it = values_.find(key);
+        if (it == values_.end()) {
+#endif
           values_[key] = returnFuture = returnPromise.get_future().share();
           needsToGenerate = true;
-//#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
-        //} else {
-        //  returnFuture = it->second;
-        //}
-        //mutex_.unlock();
-//#endif
+#ifdef FBGEMM_USE_SHARED_TIMED_MUTEX
+        } else {
+          returnFuture = it->second;
+        }
+        mutex_.unlock();
+#endif
       }
     }
 
@@ -94,7 +94,4 @@ class CodeCache {
   }
 };
 
-template <typename KEY, typename VALUE>
-thread_local std::map<KEY, std::shared_future<VALUE>> CodeCache<KEY, VALUE>::values_;
-
 } // namespace fbgemm
\ No newline at end of file
diff --git a/src/GenerateKernel.h b/src/GenerateKernel.h
index 7adf527..c0fece4 100644
--- a/src/GenerateKernel.h
+++ b/src/GenerateKernel.h
@@ -160,26 +160,26 @@ class CodeGenBase {
   int VLEN_; ///< Vector width in elements.
 
   static asmjit::JitRuntime &runtime() {
-    static thread_local asmjit::JitRuntime rt; //< JIT Runtime for asmjit,
+    static asmjit::JitRuntime rt; //< JIT Runtime for asmjit,
                                   // depents on other static
                                   // variables.  Required to prevent
                                   // initialization order fiasco
     return rt;
   }
 
-  //static std::mutex rtMutex_;    ///< Controll access to runtime;
+  static std::mutex rtMutex_;    ///< Controll access to runtime;
 
   // The hash depends on accumulate, mc, nc, ncb, kcb, nr, mr, nr_min
-  static thread_local CodeCache<std::tuple<bool, int, int, int, int, int, int, int>,
+  static CodeCache<std::tuple<bool, int, int, int, int, int, int, int>,
                    jit_micro_kernel_fp>
       codeCache_; ///< JIT Code Cache for reuse.
 };
 
-//template <typename TA, typename TB, typename TC, typename accT>
-//std::mutex CodeGenBase<TA, TB, TC, accT>::rtMutex_;
+template <typename TA, typename TB, typename TC, typename accT>
+std::mutex CodeGenBase<TA, TB, TC, accT>::rtMutex_;
 
 template <typename TA, typename TB, typename TC, typename accT>
-thread_local CodeCache<std::tuple<bool, int, int, int, int, int, int, int>,
+CodeCache<std::tuple<bool, int, int, int, int, int, int, int>,
           typename CodeGenBase<TA, TB, TC, accT>::jit_micro_kernel_fp>
     CodeGenBase<TA, TB, TC, accT>::codeCache_;
 
diff --git a/src/GenerateKernelU8S8S32ACC16.cc b/src/GenerateKernelU8S8S32ACC16.cc
index 1dd581f..205af14 100644
--- a/src/GenerateKernelU8S8S32ACC16.cc
+++ b/src/GenerateKernelU8S8S32ACC16.cc
@@ -342,7 +342,7 @@ CodeGenBase<uint8_t, int8_t, int32_t, int16_t>::getOrCreate<inst_set_t::avx2>(
     jit_micro_kernel_fp fn;
     asmjit::Error err;
     {
-      //std::unique_lock<std::mutex> lock(rtMutex_);
+      std::unique_lock<std::mutex> lock(rtMutex_);
       err = runtime().add(&fn, &code);
     }
     if (err) {
diff --git a/src/GenerateKernelU8S8S32ACC16Avx512.cc b/src/GenerateKernelU8S8S32ACC16Avx512.cc
index 8442365..819f33b 100644
--- a/src/GenerateKernelU8S8S32ACC16Avx512.cc
+++ b/src/GenerateKernelU8S8S32ACC16Avx512.cc
@@ -410,7 +410,7 @@ CodeGenBase<uint8_t, int8_t, int32_t, int16_t>::getOrCreate<inst_set_t::avx512>(
     jit_micro_kernel_fp fn;
     asmjit::Error err;
     {
-      //std::unique_lock<std::mutex> lock(rtMutex_);
+      std::unique_lock<std::mutex> lock(rtMutex_);
       err = runtime().add(&fn, &code);
     }
     if (err) {
diff --git a/src/GenerateKernelU8S8S32ACC32.cc b/src/GenerateKernelU8S8S32ACC32.cc
index b052968..dc9c534 100644
--- a/src/GenerateKernelU8S8S32ACC32.cc
+++ b/src/GenerateKernelU8S8S32ACC32.cc
@@ -330,7 +330,7 @@ CodeGenBase<uint8_t, int8_t, int32_t, int32_t>::getOrCreate<inst_set_t::avx2>(
     jit_micro_kernel_fp fn;
     asmjit::Error err;
     {
-      //std::unique_lock<std::mutex> lock(rtMutex_);
+      std::unique_lock<std::mutex> lock(rtMutex_);
       err = runtime().add(&fn, &code);
     }
     if (err) {
diff --git a/src/GenerateKernelU8S8S32ACC32Avx512.cc b/src/GenerateKernelU8S8S32ACC32Avx512.cc
index bf82bb0..5037292 100644
--- a/src/GenerateKernelU8S8S32ACC32Avx512.cc
+++ b/src/GenerateKernelU8S8S32ACC32Avx512.cc
@@ -422,7 +422,7 @@ CodeGenBase<uint8_t, int8_t, int32_t, int32_t>::getOrCreate<inst_set_t::avx512>(
     jit_micro_kernel_fp fn;
     asmjit::Error err;
     {
-      //std::unique_lock<std::mutex> lock(rtMutex_);
+      std::unique_lock<std::mutex> lock(rtMutex_);
       err = runtime().add(&fn, &code);
     }
     if (err) {
diff --git a/src/GenerateKernelU8S8S32ACC32Avx512VNNI.cc b/src/GenerateKernelU8S8S32ACC32Avx512VNNI.cc
index 816d35b..bd8be1f 100644
--- a/src/GenerateKernelU8S8S32ACC32Avx512VNNI.cc
+++ b/src/GenerateKernelU8S8S32ACC32Avx512VNNI.cc
@@ -415,7 +415,7 @@ CodeGenBase<uint8_t, int8_t, int32_t, int32_t>::getOrCreate<
     jit_micro_kernel_fp fn;
     asmjit::Error err;
     {
-      //std::unique_lock<std::mutex> lock(rtMutex_);
+      std::unique_lock<std::mutex> lock(rtMutex_);
       err = runtime().add(&fn, &code);
     }
     if (err) {