Make mkldnn Stream object thread_local and enable mkldnn thread-safe (#17022)

Summary:
This PR fixes following issue: https://github.com/pytorch/pytorch/issues/16828

It is a combination of two things:
1) MKLDNN streams are not thread-safe but are currently shared between different threads. This change makes them thread_local
2) By default MKLDNN primitives can share global memory and can't be invoked from multiple threads. This PR enables the MKLDNN_ENABLE_CONCURRENT_EXEC cmake configuration option that makes them thread-safe.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17022

Differential Revision: D14069052

Pulled By: ezyang

fbshipit-source-id: f8f7fcb86c40f5d751fb35dfccc2f802b6e137c6

diff --git a/aten/src/ATen/mkldnn/Runtime.h b/aten/src/ATen/mkldnn/Runtime.h
index c58ef2c..b27b4b5 100644 (file)
--- a/aten/src/ATen/mkldnn/Runtime.h
+++ b/aten/src/ATen/mkldnn/Runtime.h
@@ -29,7 +29,7 @@ private:
 // Stream singleton
 struct Stream {
   static Stream& Instance() {
-    static Stream myInstance;
+    static thread_local Stream myInstance;
     return myInstance;
   };
   stream& get_stream() {

diff --git a/tools/build_pytorch_libs.py b/tools/build_pytorch_libs.py
index fea64e7..fc60213 100644 (file)
--- a/tools/build_pytorch_libs.py
+++ b/tools/build_pytorch_libs.py
@@ -200,6 +200,9 @@ def run_cmake(version,
     if USE_GLOO_IBVERBS:
         cmake_defines(cmake_args, USE_IBVERBS="1", USE_GLOO_IBVERBS="1")
 
+    if USE_MKLDNN:
+        cmake_defines(cmake_args, MKLDNN_ENABLE_CONCURRENT_EXEC="ON")
+
     expected_wrapper = '/usr/local/opt/ccache/libexec'
     if IS_DARWIN and os.path.exists(expected_wrapper):
         cmake_defines(cmake_args,