make naive engine thread local.

nickguletskii · Sep 5, 2018 · b9d844e · b9d844e
1 parent 1418034
commit b9d844e
Showing 1 changed file with 20 additions and 3 deletions.
diff --git a/src/engine/engine.cc b/src/engine/engine.cc
@@ -29,6 +29,17 @@
 
 namespace mxnet {
 namespace engine {
+
+inline bool IsNaiveEngine() {
+  const char *type = getenv("MXNET_ENGINE_TYPE");
+  if (type == nullptr) {
+    return false;
+  } else {
+    std::string stype = type;
+    return stype == "NaiveEngine";
+  }
+}
+
 inline Engine* CreateEngine() {
   const char *type = getenv("MXNET_ENGINE_TYPE");
   const bool default_engine = (type == nullptr);
@@ -59,12 +70,18 @@ inline Engine* CreateEngine() {
 }  // namespace engine
 
 std::shared_ptr<Engine> Engine::_GetSharedRef() {
-  static std::shared_ptr<Engine> sptr(engine::CreateEngine());
-  return sptr;
+  static bool is_naive = engine::IsNaiveEngine();
+  if (is_naive) {
+    static thread_local std::shared_ptr<Engine> sptr(engine::CreateEngine());
+    return sptr;
+  } else {
+    static std::shared_ptr<Engine> sptr(engine::CreateEngine());
+    return sptr;
+  }
 }
 
 Engine* Engine::Get() {
-  static Engine *inst = _GetSharedRef().get();
+  Engine *inst = _GetSharedRef().get();
   return inst;
 }
 }  // namespace mxnet