[v1.x][FEATURE] Add MKLDNN Deconvolution 1D and 3D support (#20137) (#20726)

Co-authored-by: Paweł Głomski <[email protected]>

Author: samskalicky

src/operator/nn/deconvolution-inl.h

@@ -219,7 +219,11 @@ class DeconvolutionOp {
     using namespace mshadow::expr;
 
     if (param_.kernel.ndim() > 2) {
-      LOG(FATAL) << "If not using CUDNN, only 1D or 2D Deconvolution is supported";
+      LOG(FATAL) << "Only 1D or 2D Deconvolution is natively supported. "
+                 << ((MXNET_USE_MKLDNN || MXNET_USE_CUDNN)
+                         ? "Fallback to native implementation (if occurred) is therefore "
+                           "impossible for 3D Deconvolution."
+                         : "");
     }
 
     CHECK_EQ(req[deconv::kOut], kWriteTo);

src/operator/nn/deconvolution.cc

@@ -42,9 +42,14 @@ static void DeconvolutionComputeExCPU(const nnvm::NodeAttrs& attrs,
                                       const std::vector<NDArray>& outputs) {
   const DeconvolutionParam& params = nnvm::get<DeconvolutionParam>(attrs.parsed);
   if (SupportMKLDNNDeconv(params, inputs[0])) {
-    MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
-    MKLDNNRun(MKLDNNDeconvolutionForward, attrs, ctx, inputs, req, outputs);
-    MKLDNN_OPCHECK_RUN(DeconvolutionCompute<cpu>, attrs, ctx, inputs, req, outputs);
+    if (params.kernel.ndim() == 3) {
+      // we cannot check the output, as 3D deconvolution is not natively supported yet
+      MKLDNNRun(MKLDNNDeconvolutionForward, attrs, ctx, inputs, req, outputs);
+    } else {
+      MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+      MKLDNNRun(MKLDNNDeconvolutionForward, attrs, ctx, inputs, req, outputs);
+      MKLDNN_OPCHECK_RUN(DeconvolutionCompute<cpu>, attrs, ctx, inputs, req, outputs);
+    }
     return;
   }
   FallBackCompute(DeconvolutionCompute<cpu>, attrs, ctx, inputs, req, outputs);
@@ -57,9 +62,14 @@ static void DeconvolutionGradComputeExCPU(const nnvm::NodeAttrs& attrs,
                                           const std::vector<NDArray>& outputs) {
   const DeconvolutionParam& params = nnvm::get<DeconvolutionParam>(attrs.parsed);
   if (SupportMKLDNNDeconv(params, inputs[0])) {
-    MKLDNN_OPCHECK_INIT(true, outputs.size(), inputs, outputs);
-    MKLDNNRun(MKLDNNDeconvolutionBackward, attrs, ctx, inputs, req, outputs);
-    MKLDNN_OPCHECK_RUN(DeconvolutionGradCompute<cpu>, attrs, ctx, inputs, req, outputs);
+    if (params.kernel.ndim() == 3) {
+      // we cannot check the output, as 3D deconvolution is not natively supported yet
+      MKLDNNRun(MKLDNNDeconvolutionBackward, attrs, ctx, inputs, req, outputs);
+    } else {
+      MKLDNN_OPCHECK_INIT(true, outputs.size(), inputs, outputs);
+      MKLDNNRun(MKLDNNDeconvolutionBackward, attrs, ctx, inputs, req, outputs);
+      MKLDNN_OPCHECK_RUN(DeconvolutionGradCompute<cpu>, attrs, ctx, inputs, req, outputs);
+    }
     return;
   }
   FallBackCompute(DeconvolutionGradCompute<cpu>, attrs, ctx, inputs, req, outputs);
@@ -99,12 +109,12 @@ static bool DeconvolutionShape(const nnvm::NodeAttrs& attrs,
                                mxnet::ShapeVector *in_shape,
                                mxnet::ShapeVector *out_shape) {
   const DeconvolutionParam& param_ = nnvm::get<DeconvolutionParam>(attrs.parsed);
-#if MXNET_USE_CUDNN == 0
+#if MXNET_USE_CUDNN == 0 && MXNET_USE_MKLDNN == 0
   if (param_.kernel.ndim() > 2) {
-    LOG(FATAL) << "If not using CUDNN, only 1D or 2D Deconvolution is supported";
+    LOG(FATAL) << "If not using CUDNN or MKLDNN, only 1D or 2D Deconvolution is supported";
     return false;
   }
-#endif  // CUDNN
+#endif
 
   using namespace mshadow;
   if (!param_.no_bias) {

src/operator/nn/mkldnn/mkldnn_deconvolution.cc

@@ -29,7 +29,8 @@ namespace mxnet {
 namespace op {
 
 bool SupportMKLDNNDeconv(const DeconvolutionParam &params, const NDArray &input) {
-  return params.kernel.ndim() == 2 && input.shape().ndim() == 4 &&
+  return params.kernel.ndim() >= 1 && params.kernel.ndim() <= 3 &&
+         input.shape().ndim() == (params.kernel.ndim() + 2) &&
          (input.dtype() == mshadow::kFloat32 || input.dtype() == mshadow::kBfloat16);
 }
 
@@ -322,10 +323,10 @@ DeconvDescCreator::DeconvDescCreator(const DeconvolutionParam &param, const NDAr
       strides(param.stride.ndim()),
       padding(param.pad.ndim()),
       dilates(param.dilate.ndim()) {
-  // assuming only deconv2D is supported for now
   CHECK_EQ(param.stride.ndim(), param.pad.ndim());
   CHECK_EQ(param.stride.ndim(), param.dilate.ndim());
-  CHECK_EQ(param.stride.ndim(), 2);
+  CHECK_GE(param.stride.ndim(), 1);
+  CHECK_LE(param.stride.ndim(), 3);
   for (int i = 0; i < param.stride.ndim(); ++i) {
     strides[i] = param.stride[i];
     padding[i] = param.pad[i];

tests/python/mkl/test_mkldnn.py

@@ -471,8 +471,8 @@ def check_convolution_training(stype):
 @with_seed()
 def test_Deconvolution():
     def check_Deconvolution_training(stype):
-        for shape in [(3, 3, 10, 10)]: # testing only 2D for now
-            data_tmp = np.random.randint(256, size=shape)
+        for shape in [(3, 3, 10), (3, 3, 10, 10), (3, 3, 10, 10, 10)]:
+            data_tmp = np.random.normal(-0.1, 1, size=shape)
             data = mx.symbol.Variable('data', stype=stype)
 
             if np.array(shape).shape[0] == 3:
@@ -481,6 +481,11 @@ def check_Deconvolution_training(stype):
             elif np.array(shape).shape[0] == 4:
                 test = mx.symbol.Deconvolution(data=data, kernel=(3, 3), stride=(2, 2), num_filter=4)
                 weight_tmp = np.random.normal(-0.1, 0.1, size=(3, 4, 3, 3))
+            elif np.array(shape).shape[0] == 5 and stype == "default":
+                # Unable to test fallback to native implementation for non-default storage types
+                # as 3D deconvolution is not natively supported
+                test = mx.symbol.Deconvolution(data=data, kernel=(3,3,3), stride=(2,2,2), num_filter=4)
+                weight_tmp = np.random.normal(-0.1, 0.1, size=(3, 4, 3, 3, 3))
             else:
                 return 0
             bias_tmp = np.random.normal(0.1, 0.1, size=(4,))