Fix sequential MKL case (#22062)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/22062
ghimport-source-id: a30255d7453c4ffecf40215a785c1e06b7296368

Test Plan:
USE_CUDA=0 PARALLEL_BACKEND=OPENMP BLAS=MKL USE_MKLDNN=1 MKL_SEQ=1
MKLDNN_THREADING=SEQ BUILD_BINARY=1 python setup.py develop --cmake

./build/bin/parallel_info

Imported from OSS

Differential Revision: D15938079

Pulled By: ilia-cher

fbshipit-source-id: e7ef0c5bc75ebb845ebe66bf76a4070d45305b35

aten/src/ATen/ParallelOpenMP.cpp

@@ -20,12 +20,14 @@ void init_num_threads() {
   if (nthreads > 0) {
     set_num_threads(nthreads);
   } else {
-#if defined(_OPENMP) && defined(TH_BLAS_MKL)
+#if defined(_OPENMP) && defined(TH_BLAS_MKL) && !defined(TH_BLAS_MKL_SEQ)
     // 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
     // leaks in GCC 5.4)
     omp_set_num_threads(mkl_get_max_threads());
+#elif defined(_OPENMP)
+    omp_set_num_threads(intraop_default_num_threads());
 #endif
   }
 }

cmake/Dependencies.cmake

@@ -171,6 +171,9 @@ if (NOT INTERN_BUILD_MOBILE)
 
   if (MKL_FOUND)
     ADD_DEFINITIONS(-DTH_BLAS_MKL)
+    if ("${MKL_THREADING}" STREQUAL "SEQ")
+      ADD_DEFINITIONS(-DTH_BLAS_MKL_SEQ=1)
+    endif()
     if (MSVC AND MKL_LIBRARIES MATCHES ".*libiomp5md\\.lib.*")
       ADD_DEFINITIONS(-D_OPENMP_NOFORCE_MANIFEST)
       set(AT_MKL_MT 1)

cmake/Modules/FindMKL.cmake

@@ -34,19 +34,20 @@ SET(INTEL_COMPILER_DIR "/opt/intel" CACHE STRING
   "Root directory of the Intel Compiler Suite (contains ipp, mkl, etc.)")
 SET(INTEL_MKL_DIR "/opt/intel/mkl" CACHE STRING
   "Root directory of the Intel MKL (standalone)")
-SET(INTEL_MKL_SEQUENTIAL OFF CACHE BOOL
-  "Force using the sequential (non threaded) libraries")
-SET(INTEL_MKL_TBB OFF CACHE BOOL
-  "Force using TBB library")
+SET(MKL_THREADING "OMP" CACHE STRING "MKL flavor: SEQ, TBB or OMP (default)")
 
-IF (INTEL_MKL_TBB)
-  IF (USE_TBB)
-    message(STATUS "MKL is using TBB")
-  ELSE()
-    message(FATAL_ERROR "INTEL_MKL_TBB is set to true but TBB is not used")
-  ENDIF()
+IF (NOT "${MKL_THREADING}" STREQUAL "SEQ" AND
+    NOT "${MKL_THREADING}" STREQUAL "TBB" AND
+    NOT "${MKL_THREADING}" STREQUAL "OMP")
+  MESSAGE(FATAL_ERROR "Invalid MKL_THREADING (${MKL_THREADING}), should be one of: SEQ, TBB, OMP")
 ENDIF()
 
+IF ("${MKL_THREADING}" STREQUAL "TBB" AND NOT USE_TBB)
+  MESSAGE(FATAL_ERROR "MKL_THREADING is TBB but USE_TBB is turned off")
+ENDIF()
+
+MESSAGE(STATUS "MKL_THREADING = ${MKL_THREADING}")
+
 # Checks
 CHECK_TYPE_SIZE("void*" SIZE_OF_VOIDP)
 IF ("${SIZE_OF_VOIDP}" EQUAL 8)
@@ -59,7 +60,7 @@ ELSE ("${SIZE_OF_VOIDP}" EQUAL 8)
   SET(mkl64s)
 ENDIF ("${SIZE_OF_VOIDP}" EQUAL 8)
 IF(CMAKE_COMPILER_IS_GNUCC)
-  IF (INTEL_MKL_TBB)
+  IF ("${MKL_THREADING}" STREQUAL "TBB")
     SET(mklthreads "mkl_tbb_thread")
     SET(mklrtls "tbb")
   ELSE()
@@ -68,7 +69,7 @@ IF(CMAKE_COMPILER_IS_GNUCC)
   ENDIF()
   SET(mklifaces  "intel" "gf")
 ELSE(CMAKE_COMPILER_IS_GNUCC)
-  IF (INTEL_MKL_TBB)
+  IF ("${MKL_THREADING}" STREQUAL "TBB")
     SET(mklthreads "mkl_tbb_thread")
     SET(mklrtls "tbb")
   ELSE()
@@ -255,7 +256,7 @@ IF (NOT MKL_LIBRARIES)
 ENDIF (NOT MKL_LIBRARIES)
 
 # First: search for parallelized ones with intel thread lib
-IF (NOT INTEL_MKL_SEQUENTIAL)
+IF (NOT "${MKL_THREADING}" STREQUAL "SEQ")
   FOREACH(mklrtl ${mklrtls} "")
     FOREACH(mkliface ${mklifaces})
       FOREACH(mkl64 ${mkl64s} "")
@@ -268,7 +269,7 @@ IF (NOT INTEL_MKL_SEQUENTIAL)
       ENDFOREACH(mkl64)
     ENDFOREACH(mkliface)
   ENDFOREACH(mklrtl)
-ENDIF (NOT INTEL_MKL_SEQUENTIAL)
+ENDIF (NOT "${MKL_THREADING}" STREQUAL "SEQ")
 
 # Second: search for sequential ones
 FOREACH(mkliface ${mklifaces})

cmake/Modules/FindMKLDNN.cmake

@@ -96,6 +96,8 @@ IF(MKL_FOUND)
   SET(MKLDNN_THREADING "OMP:COMP" CACHE STRING "")
 ENDIF(MKL_FOUND)
 
+MESSAGE(STATUS "MKLDNN_THREADING = ${MKLDNN_THREADING}")
+
 IF (MKLDNN_THREADING STREQUAL "TBB")
   IF (USE_TBB)
     message(STATUS "MKL-DNN is using TBB")

setup.py

@@ -105,8 +105,8 @@
 #     then the build will fail if the requested BLAS is not found, otherwise
 #     the BLAS will be chosen based on what is found on your system.
 #
-#   MKL_SEQ=1
-#     chooses a sequential version of MKL library (in case of BLAS=MKL)
+#   MKL_THREADING
+#     MKL flavor: SEQ, TBB or OMP (default)
 #
 #   USE_FBGEMM
 #     Enables use of FBGEMM

tools/setup_helpers/cmake.py

@@ -309,11 +309,9 @@ class CMake:
         if os.getenv('USE_TBB'):
             CMake.defines(args, USE_TBB=check_env_flag('USE_TBB'))
 
-        if os.getenv('MKL_SEQ'):
-            CMake.defines(args, INTEL_MKL_SEQUENTIAL=check_env_flag('MKL_SEQ'))
-
-        if os.getenv('MKL_TBB'):
-            CMake.defines(args, INTEL_MKL_TBB=check_env_flag('MKL_TBB'))
+        mkl_threading = os.getenv('MKL_THREADING')
+        if mkl_threading:
+            CMake.defines(args, MKL_THREADING=mkl_threading)
 
         mkldnn_threading = os.getenv('MKLDNN_THREADING')
         if mkldnn_threading: