From 41dddfd55fc59a0ebae7cda7e759b80011741a78 Mon Sep 17 00:00:00 2001 From: Boris Daskalov Date: Wed, 13 Feb 2019 15:30:39 -0800 Subject: [PATCH] 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 --- aten/src/ATen/mkldnn/Runtime.h | 2 +- tools/build_pytorch_libs.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/aten/src/ATen/mkldnn/Runtime.h b/aten/src/ATen/mkldnn/Runtime.h index c58ef2c..b27b4b5 100644 --- 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 --- 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, -- 2.7.4