[mkldnn-v1.0] Enable more checks for MXNET_USE_MKLDNN (#16520)

* open USE_MKLDNN check * trigger ci * ci
apache · Oct 19, 2019 · ac7792d · ac7792d
1 parent f7f4959
commit ac7792d
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 8 deletions.
diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc
@@ -627,7 +627,7 @@ int MXNDArrayGetData(NDArrayHandle handle,
                      void **out_pdata) {
   API_BEGIN();
   NDArray *arr = static_cast<NDArray*>(handle);
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
   if (arr->IsMKLDNNData()) {
     arr->Reorder2DefaultAsync();
     arr->WaitToRead();

diff --git a/src/common/utils.h b/src/common/utils.h
@@ -47,7 +47,7 @@
 #include <limits>
 
 #include "../operator/mxnet_op.h"
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
 #include "../operator/nn/mkldnn/mkldnn_base-inl.h"
 #endif
 
@@ -484,7 +484,7 @@ inline void LogStorageFallback(const nnvm::NodeAttrs& attrs,
     "0 to suppress this warning.";
   os << "\nStorage type fallback detected:\n" << op_str << warning;
   LogOnce(os.str());
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
   if (!MKLDNNEnvSet()) common::LogOnce("MXNET_MKLDNN_ENABLED flag is off. "
                                        "You can re-enable by setting MXNET_MKLDNN_ENABLED=1");
   if (GetMKLDNNCacheSize() != -1) common::LogOnce("MXNET_MKLDNN_CACHE_NUM is set."

diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
@@ -1617,7 +1617,7 @@ void NDArray::Save(dmlc::Stream *strm) const {
   } else {
     this->WaitToRead();
     nd_cpu = *this;
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
     if (nd_cpu.IsMKLDNNData()) {
       nd_cpu = nd_cpu.Reorder2Default();
     }

diff --git a/src/operator/tensor/matrix_op-inl.h b/src/operator/tensor/matrix_op-inl.h
@@ -541,7 +541,7 @@ inline bool SliceForwardInferStorageType(const nnvm::NodeAttrs& attrs,
   }
 
   if (in_stype == kDefaultStorage) {
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
     if (dev_mask == Context::kCPU && MKLDNNEnvSet()
         && SupportMKLDNNSlice(param)) {
       dispatched = storage_type_assign(&out_stype, kDefaultStorage,

diff --git a/src/operator/tensor/matrix_op.cc b/src/operator/tensor/matrix_op.cc
@@ -425,7 +425,7 @@ void SliceExCPU(const nnvm::NodeAttrs& attrs,
   auto in_stype = inputs[0].storage_type();
   if (in_stype == kCSRStorage) {
     SliceCsrImpl<cpu>(param, ctx, inputs[0], req[0], outputs[0]);
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
   } else if (in_stype == kDefaultStorage) {
     if (SupportMKLDNN(inputs[0])) {
       MKLDNNSlice(param, ctx, inputs[0], req[0], outputs[0]);
@@ -488,7 +488,7 @@ Example::
 .set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_slice"})
 .set_attr<FCompute>("FCompute<cpu>", SliceOpForward<cpu>)
 .set_attr<FComputeEx>("FComputeEx<cpu>", SliceExCPU)
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
 .set_attr<bool>("TIsMKLDNN", true)
 #endif
 .add_argument("data", "NDArray-or-Symbol", "Source input")

diff --git a/src/storage/cpu_device_storage.h b/src/storage/cpu_device_storage.h
@@ -53,7 +53,7 @@ class CPUDeviceStorage {
   /*!
    * \brief Alignment of allocation.
    */
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
   // MKLDNN requires special alignment. 64 is used by the MKLDNN library in
   // memory allocation.
   static constexpr size_t alignment_ = kMKLDNNAlign;