[GitHub] [incubator-mxnet] TaoLv commented on a change in pull request #17884: [MKL-DNN] Integrate Conv3d and Pool3d/1d

[GitBox]

TaoLv commented on a change in pull request #17884: [MKL-DNN] Integrate Conv3d 
and Pool3d/1d
URL: https://github.com/apache/incubator-mxnet/pull/17884#discussion_r396523088
 
 

 ##########
 File path: src/operator/nn/mkldnn/mkldnn_pooling-inl.h
 ##########
 @@ -114,15 +115,21 @@ inline bool SupportMKLDNNPooling(const PoolingParam 
&param,
     return true;
   } else {
     if (param.pool_type == pool_enum::kAvgPooling) {
-      CHECK_EQ(dshape.ndim(), 4);
+      CHECK(dshape.ndim() == 3 || dshape.ndim() == 4 || dshape.ndim() == 5);
       // mkldnn works differently when padding is asymmetric, so let's skip 
this case.
-      if (param.pad[0] == GetPaddingSizeFull(dshape[2], param.pad[0], 
param.pad[0], param.kernel[0],
-                                             param.stride[0]) &&
-          param.pad[1] == GetPaddingSizeFull(dshape[3], param.pad[1], 
param.pad[1], param.kernel[1],
-                                             param.stride[1])) {
-        return true;
+      bool is_symmetric = true;
+      switch (dshape.ndim()) {
+        case 5:
+          is_symmetric = is_symmetric && (param.pad[2] == 
GetPaddingSizeFull(dshape[4],
+                                param.pad[2], param.pad[2], param.kernel[2], 
param.stride[2]));
+        case 4:
+          is_symmetric = is_symmetric && (param.pad[1] == 
GetPaddingSizeFull(dshape[3],
+                                param.pad[1], param.pad[1], param.kernel[1], 
param.stride[1]));
 
 Review comment:
   I see both pad[0] and pad[1] are checked in previous code.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services