[Reland] Include AT_PARALLEL_OPENMP/AT_PARALLEL_NATIVE/AT_PARALLEL_NA…

…TIVE_TBB to ATen/Config.h (pytorch#39881) Summary: Pull Request resolved: pytorch#39881 Test Plan: Imported from OSS Differential Revision: D22008317 Pulled By: pbelevich fbshipit-source-id: b25714d8643cf584bb3331d70e44f4df06c1b615
xwang233 · Jun 19, 2020 · 88bceb5 · 88bceb5
1 parent 2dd6736
commit 88bceb5
Show file tree

Hide file tree

Showing 8 changed files with 39 additions and 24 deletions.
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -19,7 +19,6 @@ COMMON_COPTS = [
     "-DTH_HAVE_THREAD",
     "-DUSE_FBGEMM",
     "-DUSE_DISTRIBUTED",
-    "-DAT_PARALLEL_NATIVE=1",
     "-DATEN_THREADING=NATIVE",
     "-DNO_CUDNN_DESTROY_HANDLE",
 ] + if_cuda([
@@ -545,6 +544,9 @@ template_rule(
         "@AT_NNPACK_ENABLED@": "0",
         "@CAFFE2_STATIC_LINK_CUDA_INT@": "0",
         "@USE_BLAS@": "1",
+        "@AT_PARALLEL_OPENMP@": "0",
+        "@AT_PARALLEL_NATIVE@": "1",
+        "@AT_PARALLEL_NATIVE_TBB@": "0",
     },
 )
 

diff --git a/aten/src/ATen/Config.h.in b/aten/src/ATen/Config.h.in
@@ -11,3 +11,6 @@
 #define AT_NNPACK_ENABLED() @AT_NNPACK_ENABLED@
 #define CAFFE2_STATIC_LINK_CUDA() @CAFFE2_STATIC_LINK_CUDA_INT@
 #define AT_BUILD_WITH_BLAS() @USE_BLAS@
+#define AT_PARALLEL_OPENMP @AT_PARALLEL_OPENMP@
+#define AT_PARALLEL_NATIVE @AT_PARALLEL_NATIVE@
+#define AT_PARALLEL_NATIVE_TBB @AT_PARALLEL_NATIVE_TBB@
diff --git a/aten/src/ATen/Parallel.h b/aten/src/ATen/Parallel.h
@@ -1,5 +1,6 @@
 #pragma once
 #include <ATen/ATen.h>
+#include <ATen/Config.h>
 #include <ATen/core/ivalue.h>
 #include <c10/macros/Macros.h>
 

diff --git a/aten/src/ATen/ParallelNative.cpp b/aten/src/ATen/ParallelNative.cpp
@@ -1,3 +1,4 @@
+#include <ATen/Config.h>
 #if AT_PARALLEL_NATIVE
 #include <ATen/Parallel.h>
 #include <ATen/PTThreadPool.h>

diff --git a/aten/src/ATen/ParallelNativeTBB.cpp b/aten/src/ATen/ParallelNativeTBB.cpp
@@ -1,3 +1,4 @@
+#include <ATen/Config.h>
 #if AT_PARALLEL_NATIVE_TBB
 #include <ATen/Parallel.h>
 #include <ATen/PTThreadPool.h>

diff --git a/aten/src/ATen/ParallelOpenMP.cpp b/aten/src/ATen/ParallelOpenMP.cpp
@@ -1,3 +1,4 @@
+#include <ATen/Config.h>
 #if AT_PARALLEL_OPENMP
 #include <ATen/Parallel.h>
 

diff --git a/aten/src/ATen/ParallelThreadPoolNative.cpp b/aten/src/ATen/ParallelThreadPoolNative.cpp
@@ -1,3 +1,4 @@
+#include <ATen/Config.h>
 #if AT_PARALLEL_OPENMP || AT_PARALLEL_NATIVE || AT_PARALLEL_NATIVE_TBB
 #include <ATen/Parallel.h>
 #include <ATen/PTThreadPool.h>

diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
@@ -6,6 +6,34 @@ if(USE_VULKAN)
   include(../cmake/VulkanCodegen.cmake)
 endif()
 
+# ATen parallelism settings
+#  OMP - OpenMP for intra-op, native thread pool for inter-op parallelism
+#  NATIVE - using native thread pool for intra- and inter-op parallelism
+#  TBB - using TBB for intra- and native thread pool for inter-op parallelism
+if(INTERN_BUILD_MOBILE AND NOT BUILD_CAFFE2_MOBILE)
+  set(ATEN_THREADING "NATIVE" CACHE STRING "ATen parallel backend")
+else()
+  set(ATEN_THREADING "OMP" CACHE STRING "ATen parallel backend")
+endif()
+
+set(AT_PARALLEL_OPENMP 0)
+set(AT_PARALLEL_NATIVE 0)
+set(AT_PARALLEL_NATIVE_TBB 0)
+
+message(STATUS "Using ATen parallel backend: ${ATEN_THREADING}")
+if("${ATEN_THREADING}" STREQUAL "OMP")
+  set(AT_PARALLEL_OPENMP 1)
+elseif("${ATEN_THREADING}" STREQUAL "NATIVE")
+  set(AT_PARALLEL_NATIVE 1)
+elseif("${ATEN_THREADING}" STREQUAL "TBB")
+  if(NOT USE_TBB)
+    message(FATAL_ERROR "Using TBB backend but USE_TBB is off")
+  endif()
+  set(AT_PARALLEL_NATIVE_TBB 1)
+else()
+  message(FATAL_ERROR "Unknown ATen parallel backend: ${ATEN_THREADING}")
+endif()
+
 # ---[ Declare source file lists
 
 # ---[ ATen build
@@ -896,29 +924,6 @@ elseif(USE_ROCM)
 endif()
 
 
-# ATen parallelism settings
-#  OMP - OpenMP for intra-op, native thread pool for inter-op parallelism
-#  NATIVE - using native thread pool for intra- and inter-op parallelism
-#  TBB - using TBB for intra- and native thread pool for inter-op parallelism
-if(INTERN_BUILD_MOBILE AND NOT BUILD_CAFFE2_MOBILE)
-  set(ATEN_THREADING "NATIVE" CACHE STRING "ATen parallel backend")
-else()
-  set(ATEN_THREADING "OMP" CACHE STRING "ATen parallel backend")
-endif()
-
-message(STATUS "Using ATen parallel backend: ${ATEN_THREADING}")
-if("${ATEN_THREADING}" STREQUAL "OMP")
-  target_compile_definitions(torch_cpu PUBLIC "-DAT_PARALLEL_OPENMP=1")
-elseif("${ATEN_THREADING}" STREQUAL "NATIVE")
-  target_compile_definitions(torch_cpu PUBLIC "-DAT_PARALLEL_NATIVE=1")
-elseif("${ATEN_THREADING}" STREQUAL "TBB")
-  if(NOT USE_TBB)
-    message(FATAL_ERROR "Using TBB backend but USE_TBB is off")
-  endif()
-  target_compile_definitions(torch_cpu PUBLIC "-DAT_PARALLEL_NATIVE_TBB=1")
-else()
-  message(FATAL_ERROR "Unknown ATen parallel backend: ${ATEN_THREADING}")
-endif()
 set(EXPERIMENTAL_SINGLE_THREAD_POOL "0" CACHE STRING
   "Experimental option to use a single thread pool for inter- and intra-op parallelism")
 if("${EXPERIMENTAL_SINGLE_THREAD_POOL}")