AT_MKL_SEQUENTIAL and build changes (pytorch#70259)

Summary:
Re-land of  pytorch#69419

Pull Request resolved: pytorch#70259

Test Plan: Imported from OSS

Reviewed By: malfet

Differential Revision: D33246757

Pulled By: ngimel

fbshipit-source-id: 738f8558d4cad6752be14108f9931ec3514f6682

BUILD.bazel

@@ -318,6 +318,7 @@ header_template_rule(
     substitutions = {
         "@AT_MKLDNN_ENABLED@": "1",
         "@AT_MKL_ENABLED@": "1",
+        "@AT_MKL_SEQUENTIAL@": "0",
         "@AT_FFTW_ENABLED@": "0",
         "@AT_POCKETFFT_ENABLED@": "0",
         "@AT_NNPACK_ENABLED@": "0",

aten/src/ATen/Config.h.in

@@ -8,6 +8,7 @@
 
 #define AT_MKLDNN_ENABLED() @AT_MKLDNN_ENABLED@
 #define AT_MKL_ENABLED() @AT_MKL_ENABLED@
+#define AT_MKL_SEQUENTIAL() @AT_MKL_SEQUENTIAL@
 #define AT_FFTW_ENABLED() @AT_FFTW_ENABLED@
 #define AT_POCKETFFT_ENABLED() @AT_POCKETFFT_ENABLED@
 #define AT_NNPACK_ENABLED() @AT_NNPACK_ENABLED@

aten/src/ATen/ParallelOpenMP.cpp

@@ -31,7 +31,7 @@ void init_num_threads() {
   if (nthreads > 0) {
     set_num_threads(nthreads);
   } else {
-#if defined(_OPENMP) && AT_MKL_ENABLED() && !defined(TH_BLAS_MKL_SEQ)
+#if defined(_OPENMP) && AT_MKL_ENABLED() && !AT_MKL_SEQUENTIAL()
     // If we are using MKL an OpenMP make sure the number of threads match.
     // Otherwise, MKL and our OpenMP-enabled functions will keep changing the
     // size of the OpenMP thread pool, resulting in worse performance (and memory

cmake/Dependencies.cmake

@@ -247,6 +247,7 @@ endif()
 
 if(NOT INTERN_BUILD_MOBILE)
   set(AT_MKL_ENABLED 0)
+  set(AT_MKL_SEQUENTIAL 0)
   set(AT_MKL_MT 0)
   set(USE_BLAS 1)
   if(NOT (ATLAS_FOUND OR BLIS_FOUND OR GENERIC_BLAS_FOUND OR MKL_FOUND OR OpenBLAS_FOUND OR VECLIB_FOUND OR FlexiBLAS_FOUND))
@@ -258,9 +259,8 @@ if(NOT INTERN_BUILD_MOBILE)
   endif()
 
   if(MKL_FOUND)
-    add_definitions(-DTH_BLAS_MKL)
     if("${MKL_THREADING}" STREQUAL "SEQ")
-      add_definitions(-DTH_BLAS_MKL_SEQ=1)
+      set(AT_MKL_SEQUENTIAL 1)
     endif()
     if(MSVC AND MKL_LIBRARIES MATCHES ".*libiomp5md\\.lib.*")
       add_definitions(-D_OPENMP_NOFORCE_MANIFEST)