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: https://github.com/pytorch/pytorch/pull/19014
Differential Revision: D16094090
Pulled By: ezyang
fbshipit-source-id: 3f0b1d1a59a327ea0d1456e2752f2edd78d96ccc
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 1b02218..3771d15 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -136,6 +136,9 @@
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
diff --git a/cmake/Modules/FindMKLDNN.cmake b/cmake/Modules/FindMKLDNN.cmake
index 4994259..bcb16b0 100644
--- a/cmake/Modules/FindMKLDNN.cmake
+++ b/cmake/Modules/FindMKLDNN.cmake
@@ -34,7 +34,6 @@
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 @@
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)
diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake
index c6d4bb7..955e7ef 100644
--- a/cmake/Summary.cmake
+++ b/cmake/Summary.cmake
@@ -101,6 +101,9 @@
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}")
diff --git a/setup.py b/setup.py
index c298719..236e2ff 100644
--- a/setup.py
+++ b/setup.py
@@ -380,6 +380,10 @@
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']:
