From eec9c9cbe7e841a9b3fef301d8dc9518d6db2452 Mon Sep 17 00:00:00 2001
From: Yihua Xu <yihua.xu@intel.com>
Date: Fri, 15 Nov 2019 13:54:38 +0800
Subject: [PATCH] Fix jit tls issue (#21151)

---
 paddle/fluid/operators/jit/helper.cc      |  6 ++++++
 paddle/fluid/operators/jit/helper.h       | 17 +++++++++++++++--
 paddle/fluid/operators/jit/kernel_pool.cc |  6 ++++++
 paddle/fluid/operators/jit/kernel_pool.h  | 14 ++++++++++++--
 4 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/paddle/fluid/operators/jit/helper.cc b/paddle/fluid/operators/jit/helper.cc
index f868c847bd8..992228adb50 100644
--- a/paddle/fluid/operators/jit/helper.cc
+++ b/paddle/fluid/operators/jit/helper.cc
@@ -22,6 +22,12 @@ namespace paddle {
 namespace operators {
 namespace jit {
 
+std::unordered_map<std::string, std::shared_ptr<void>>& GetFuncCacheMap() {
+  static thread_local std::unordered_map<std::string, std::shared_ptr<void>>
+      g_func_cache_map;
+  return g_func_cache_map;
+}
+
 #define ONE_CASE(key) \
   case key:           \
     return #key
diff --git a/paddle/fluid/operators/jit/helper.h b/paddle/fluid/operators/jit/helper.h
index 1ac5318d461..9a2447a98fb 100644
--- a/paddle/fluid/operators/jit/helper.h
+++ b/paddle/fluid/operators/jit/helper.h
@@ -15,6 +15,7 @@
 #pragma once
 
 #include <iostream>
+#include <memory>
 #include <string>
 #include <unordered_map>
 #include <utility>  // for std::move
@@ -175,13 +176,25 @@ typename KernelTuple::func_type GetDefaultBestFunc(
   return funcs[0];
 }
 
+extern std::unordered_map<std::string, std::shared_ptr<void>>&
+GetFuncCacheMap();
+
 template <typename KernelTuple, typename PlaceType>
 class KernelFuncs {
  public:
   KernelFuncs() = default;
   static KernelFuncs& Cache() {
-    static thread_local KernelFuncs<KernelTuple, PlaceType> g_func_cache;
-    return g_func_cache;
+    auto& func_cache_map = GetFuncCacheMap();
+    std::string key = typeid(KernelFuncs<KernelTuple, PlaceType>).name();
+    auto iter = func_cache_map.find(key);
+    if (iter != func_cache_map.end()) {
+      return *(KernelFuncs<KernelTuple, PlaceType>*)(iter->second.get());
+    } else {
+      std::shared_ptr<void> cache =
+          std::make_shared<KernelFuncs<KernelTuple, PlaceType>>();
+      func_cache_map.emplace(key, cache);
+      return *(KernelFuncs<KernelTuple, PlaceType>*)(cache.get());
+    }
   }
 
   // the exposed interface to use
diff --git a/paddle/fluid/operators/jit/kernel_pool.cc b/paddle/fluid/operators/jit/kernel_pool.cc
index bc98c644fbe..8eef8e47474 100644
--- a/paddle/fluid/operators/jit/kernel_pool.cc
+++ b/paddle/fluid/operators/jit/kernel_pool.cc
@@ -21,6 +21,12 @@ namespace paddle {
 namespace operators {
 namespace jit {
 
+std::unordered_map<std::string, std::shared_ptr<void>>& GetJITCodesMap() {
+  static thread_local std::unordered_map<std::string, std::shared_ptr<void>>
+      g_jit_codes_map;
+  return g_jit_codes_map;
+}
+
 JitCodeCreatorPool& JitCodeCreatorPool::Instance() {
   static JitCodeCreatorPool g_creator_pool;
   return g_creator_pool;
diff --git a/paddle/fluid/operators/jit/kernel_pool.h b/paddle/fluid/operators/jit/kernel_pool.h
index 04710a54ac9..548c8704126 100644
--- a/paddle/fluid/operators/jit/kernel_pool.h
+++ b/paddle/fluid/operators/jit/kernel_pool.h
@@ -28,6 +28,8 @@ namespace paddle {
 namespace operators {
 namespace jit {
 
+extern std::unordered_map<std::string, std::shared_ptr<void>>& GetJITCodesMap();
+
 template <KernelType KT>
 class JitCodePool {
   typedef std::unique_ptr<GenBase> GenBasePtr;
@@ -36,8 +38,16 @@ class JitCodePool {
  public:
   JitCodePool() = default;
   static JitCodePool& Instance() {
-    static thread_local JitCodePool<KT> g_jit_codes;
-    return g_jit_codes;
+    auto& jit_codes_map = GetJITCodesMap();
+    std::string key = typeid(JitCodePool<KT>).name();
+    auto iter = jit_codes_map.find(key);
+    if (iter != jit_codes_map.end()) {
+      return *(JitCodePool<KT>*)(iter->second.get());
+    } else {
+      std::shared_ptr<void> cache = std::make_shared<JitCodePool<KT>>();
+      jit_codes_map.emplace(key, cache);
+      return *(JitCodePool<KT>*)(cache.get());
+    }
   }
 
   const JitCodeMap& AllKernels() { return codes_; }
-- 
GitLab