remove kBFloat16 as its not supported int 1.6

apache · May 3, 2020 · 5bf9c76 · 5bf9c76
1 parent 679ae04
commit 5bf9c76
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 6 deletions.
diff --git a/src/operator/nn/mkldnn/mkldnn_act.cc b/src/operator/nn/mkldnn/mkldnn_act.cc
@@ -51,7 +51,7 @@ bool SupportMKLDNNAct(const ActivationParam& param, const NDArray &input) {
   // MKL-DNN Activation supports 1d, 2d, 3d, 4d and 5d data layout
   if ((input.shape().ndim() < 1) ||
       (input.shape().ndim() > 5) ||
-      !(input.dtype() == mshadow::kFloat32 || input.dtype() == mshadow::kBfloat16))
+      !(input.dtype() == mshadow::kFloat32))
     return false;
   return SupportMKLDNNAct(param);
 }
@@ -65,7 +65,7 @@ bool SupportMKLDNNLeakyRelu(const LeakyReLUParam& param, const NDArray &input) {
   // MKL-DNN Activation supports 1d, 2d, 3d, 4d and 5d data layout
   if ((input.shape().ndim() < 1) ||
       (input.shape().ndim() > 5) ||
-      !(input.dtype() == mshadow::kFloat32 || input.dtype() == mshadow::kBfloat16))
+      !(input.dtype() == mshadow::kFloat32))
     return false;
   return SupportMKLDNNLeakyRelu(param);
 }

diff --git a/src/operator/nn/mkldnn/mkldnn_base-inl.h b/src/operator/nn/mkldnn/mkldnn_base-inl.h
@@ -129,7 +129,7 @@ static inline bool SupportMKLDNN(int dtype, const mxnet::TShape &shape) {
     // MKLDNN currently does not support 0-dim Tensor and 0-size Tensor
     return false;
   }
-  return (dtype == mshadow::kFloat32 || dtype == mshadow::kBfloat16) &&
+  return (dtype == mshadow::kFloat32) &&
                     (ndim == 1 || ndim == 2 || ndim == 4);
 }
 

diff --git a/src/operator/nn/mkldnn/mkldnn_pooling-inl.h b/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
@@ -104,7 +104,7 @@ inline bool SupportMKLDNNPooling(const PoolingParam &param,
   const auto dtype = input.dtype();
 
   if (!(SupportStorageMKLDNN(input.storage_type()) && (ndim == 3 || ndim == 4 || ndim == 5) &&
-       (dtype == mshadow::kFloat32 || dtype == mshadow::kBfloat16)))
+       (dtype == mshadow::kFloat32)))
     return false;
 
   if (!SupportMKLDNNPooling(param))

diff --git a/src/operator/subgraph/mkldnn/mkldnn_subgraph_base-inl.h b/src/operator/subgraph/mkldnn/mkldnn_subgraph_base-inl.h
@@ -29,8 +29,7 @@ static inline bool SupportMKLDNNAttr(const std::shared_ptr<NodeAttr>& node_attr)
   if (node_attr) {
     int ndim = node_attr->ishape[0].ndim();
     return (node_attr->dispatch_mode == DispatchMode::kFComputeEx) &&
-           (node_attr->itype[0] == mshadow::kFloat32 ||
-            node_attr->itype[0] == mshadow::kBfloat16) &&
+           (node_attr->itype[0] == mshadow::kFloat32) &&
            (ndim == 1 || ndim == 2 || ndim == 4 || ndim == 5);
   } else {
     return true;