Add USE_MKLDNN_CBLAS build option. (#19014)

Summary:
MKL-DNN is the main library for computation when we use ideep device. It can use kernels implemented by different algorithms (including JIT, CBLAS, etc.) for computation. We add the "USE_MKLDNN_CBLAS" (default OFF) build option so that users can decide whether to use CBLAS computation methods or not.
Pull Request resolved: pytorch/pytorch#19014

Differential Revision: D16094090

Pulled By: ezyang

fbshipit-source-id: 3f0b1d1a59a327ea0d1456e2752f2edd78d96ccc

CMakeLists.txt

@@ -136,6 +136,9 @@ cmake_dependent_option(
   USE_MKLDNN "Use MKLDNN. Only available on x86 and x86_64." ON
   "CPU_INTEL" OFF)
 set(MKLDNN_ENABLE_CONCURRENT_EXEC ${USE_MKLDNN})
+cmake_dependent_option(
+    USE_MKLDNN_CBLAS "Use CBLAS in MKLDNN" OFF
+    "USE_MKLDNN" OFF)
 option(USE_DISTRIBUTED "Use distributed" ON)
 cmake_dependent_option(
     USE_MPI "Use MPI for Caffe2. Only available if USE_DISTRIBUTED is on." ON

cmake/Modules/FindMKLDNN.cmake

@@ -34,7 +34,6 @@ IF (NOT IDEEP_INCLUDE_DIR OR NOT MKLDNN_INCLUDE_DIR)
   RETURN()
 ENDIF(NOT IDEEP_INCLUDE_DIR OR NOT MKLDNN_INCLUDE_DIR)
 LIST(APPEND MKLDNN_INCLUDE_DIR ${IDEEP_INCLUDE_DIR})
-
 IF(MKL_FOUND)
   # Append to mkldnn dependencies
   LIST(APPEND MKLDNN_LIBRARIES ${MKL_LIBRARIES})
@@ -125,7 +124,11 @@ IF(NOT TARGET mkldnn)
   RETURN()
 ENDIF(NOT TARGET mkldnn)
 IF(MKL_FOUND)
-  TARGET_COMPILE_DEFINITIONS(mkldnn PRIVATE -DUSE_MKL)
+  set (USE_MKL_CBLAS -DUSE_MKL)
+  IF(USE_MKLDNN_CBLAS)
+    LIST(APPEND USE_MKL_CBLAS -DUSE_CBLAS)
+  ENDIF(USE_MKLDNN_CBLAS)
+  TARGET_COMPILE_DEFINITIONS(mkldnn PRIVATE USE_MKL_CBLAS)
 ENDIF(MKL_FOUND)
 IF(NOT APPLE AND CMAKE_COMPILER_IS_GNUCC)
   TARGET_COMPILE_OPTIONS(mkldnn PRIVATE -Wno-maybe-uninitialized)

cmake/Summary.cmake

@@ -101,6 +101,9 @@ function (caffe2_print_configuration_summary)
   message(STATUS "  USE_METAL             : ${USE_METAL}")
   message(STATUS "  USE_MKL               : ${CAFFE2_USE_MKL}")
   message(STATUS "  USE_MKLDNN            : ${CAFFE2_USE_MKLDNN}")
+  if(${CAFFE2_USE_MKLDNN})
+    message(STATUS "  USE_MKLDNN_CBLAS      : ${USE_MKLDNN_CBLAS}")
+  endif()
   message(STATUS "  USE_NCCL              : ${USE_NCCL}")
   if(${USE_NCCL})
     message(STATUS "    USE_SYSTEM_NCCL     : ${USE_SYSTEM_NCCL}")

setup.py

@@ -380,6 +380,10 @@ def run(self):
             report('-- Not using CUDA')
         if cmake_cache_vars['USE_MKLDNN']:
             report('-- Using MKLDNN')
+            if cmake_cache_vars['USE_MKLDNN_CBLAS']:
+                report('-- Using CBLAS in MKLDNN')
+            else:
+                report('-- Not using CBLAS in MKLDNN')
         else:
             report('-- Not using MKLDNN')
         if cmake_cache_vars['USE_NCCL'] and cmake_cache_vars['USE_SYSTEM_NCCL']: