Fix SupportMKLDNN function for Convolution and Reshape

apache · Sep 14, 2021 · 6cfdf0d · 6cfdf0d
1 parent e09c119
commit 6cfdf0d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/src/operator/nn/mkldnn/mkldnn_convolution.cc b/src/operator/nn/mkldnn/mkldnn_convolution.cc
@@ -37,11 +37,10 @@ namespace op {
 DMLC_REGISTER_PARAMETER(MKLDNNConvParam);
 
 bool SupportMKLDNNConv(const ConvolutionParam& params, const NDArray& input) {
-  if ((params.kernel.ndim() != 1) && (params.kernel.ndim() != 2) && (params.kernel.ndim() != 3))
+  if (params.kernel.ndim() > 3 || params.kernel.ndim() == 0)
     return false;
   return IsMKLDNNType(input.dtype()) &&
-         ((input.shape().ndim() == 3) || (input.shape().ndim() == 4) ||
-          (input.shape().ndim() == 5));
+         input.shape().ndim() >= 3 && input.shape().ndim() <= 5;
 }
 
 std::shared_ptr<mkldnn::convolution_forward::primitive_desc> GetConvFwdImpl(

diff --git a/src/operator/nn/mkldnn/mkldnn_reshape.cc b/src/operator/nn/mkldnn/mkldnn_reshape.cc
@@ -36,8 +36,8 @@ namespace op {
 bool SupportMKLDNNReshape(const NDArray& input, const NDArray& output) {
   const int input_ndims  = input.shape().ndim();
   const int output_ndims = output.shape().ndim();
-  return input_ndims >= 1 && input_ndims <= 6 && output_ndims >= 1 && output_ndims <= 6 &&
-         IsMKLDNNType(input.dtype()) && input.shape().Size() > 0;
+  return input.shape().Size() > 0 && input_ndims >= 1 && input_ndims <= 6 && output_ndims >= 1 &&
+         output_ndims <= 6 && IsMKLDNNType(input.dtype());
 }
 
 MKLDNNReshapeFwd::MKLDNNReshapeFwd(const OpReqType& req,