[reland][DNNL]:enable max_pool3d and avg_pool3d (#40996)

XiaobingSuper · facebook-github-bot · commit 2b14f2d368ff · 2020-07-16T10:26:45.000-07:00
Summary: Pull Request resolved: #40996 Test Plan: Imported from OSS Differential Revision: D22440766 Pulled By: VitalyFedyunin fbshipit-source-id: 242711612920081eb4a7e5a7e80bc8b2d4c9f978
diff --git a/aten/src/ATen/native/Pooling.cpp b/aten/src/ATen/native/Pooling.cpp
@@ -154,9 +154,14 @@ Tensor max_pool3d(
     IntArrayRef padding,
     IntArrayRef dilation,
     bool ceil_mode) {
+  if (self.is_mkldnn()) {
+    return at::mkldnn_max_pool3d(
+        self, kernel_size, stride, padding, dilation, ceil_mode);
+  }
   auto output_and_indices = at::max_pool3d_with_indices(
       self, kernel_size, stride, padding, dilation, ceil_mode);
   return std::get<0>(output_and_indices);
 }
+
 } // namespace native
 } // namespace at
diff --git a/aten/src/ATen/native/mkldnn/Pooling.cpp b/aten/src/ATen/native/mkldnn/Pooling.cpp
@@ -17,8 +17,17 @@ Tensor mkldnn_max_pool2d(
     IntArrayRef padding,
     IntArrayRef dilation,
     bool ceil_mode) {
-  AT_ERROR(
-      "mkldnn_max_pool2d: ATen not compiled with MKLDNN support");
+  TORCH_CHECK(false, "mkldnn_max_pool2d: ATen not compiled with MKLDNN support");
+}
+
+Tensor mkldnn_max_pool3d(
+    const Tensor& self,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    bool ceil_mode) {
+  TORCH_CHECK(false, "mkldnn_max_pool3d: ATen not compiled with MKLDNN support");
 }
 
 Tensor mkldnn_avg_pool2d(
@@ -29,7 +38,7 @@ Tensor mkldnn_avg_pool2d(
     bool ceil_mode,
     bool count_include_pad,
     c10::optional<int64_t> divisor_override) {
-  AT_ERROR("mkldnn_avg_pool2d: ATen not compiled with MKLDNN support");
+  TORCH_CHECK(false, "mkldnn_avg_pool2d: ATen not compiled with MKLDNN support");
 }
 
 Tensor& mkldnn_avg_pool2d_out(
@@ -41,19 +50,41 @@ Tensor& mkldnn_avg_pool2d_out(
     bool ceil_mode,
     bool count_include_pad,
     c10::optional<int64_t> divisor_override) {
-  AT_ERROR("mkldnn_avg_pool2d_out: ATen not compiled with MKLDNN support");
+  TORCH_CHECK(false, "mkldnn_avg_pool2d_out: ATen not compiled with MKLDNN support");
+}
+
+Tensor mkldnn_avg_pool3d(
+    const Tensor& self,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    bool ceil_mode,
+    bool count_include_pad,
+    c10::optional<int64_t> divisor_override) {
+  TORCH_CHECK(false, "mkldnn_avg_pool3d: ATen not compiled with MKLDNN support");
+}
+
+Tensor& mkldnn_avg_pool3d_out(
+    Tensor& output,
+    const Tensor& self,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    bool ceil_mode,
+    bool count_include_pad,
+    c10::optional<int64_t> divisor_override) {
+  TORCH_CHECK(false, "mkldnn_avg_pool3d_out: ATen not compiled with MKLDNN support");
 }
 
 Tensor mkldnn_adaptive_avg_pool2d(Tensor const& input, IntArrayRef output_size) {
-  AT_ERROR("mkldnn_adaptive_avg_pool2d: ATen not compiled with MKLDNN support");
+  TORCH_CHECK(false, "mkldnn_adaptive_avg_pool2d: ATen not compiled with MKLDNN support");
 }
 
 Tensor& mkldnn_adaptive_avg_pool2d_out(
     Tensor& output,
     const Tensor& input,
     IntArrayRef output_size) {
-  AT_ERROR(
-      "mkldnn_adaptive_avg_pool2d_out: ATen not compiled with MKLDNN support");
+  TORCH_CHECK(false, "mkldnn_adaptive_avg_pool2d_out: ATen not compiled with MKLDNN support");
 }
 
 } // namespace native
@@ -67,21 +98,22 @@ Tensor& mkldnn_adaptive_avg_pool2d_out(
 namespace at {
 namespace native {
 
-static Tensor _mkldnn_pool2d(
+static Tensor _mkldnn_pooling(
     const Tensor& input,
     IntArrayRef kernel_size,
     IntArrayRef stride,
     IntArrayRef padding,
     IntArrayRef dilation,
     bool ceil_mode,
     ideep::algorithm algo) {
-  auto kernel_size_vec = expand_param_if_needed(kernel_size, "kernel_size", 2);
+  const int64_t dims = input.dim() - 2;
+  auto kernel_size_vec = expand_param_if_needed(kernel_size, "kernel_size", dims);
   if (stride.empty()) stride = kernel_size;
-  auto stride_vec = expand_param_if_needed(stride, "stride", 2);
-  auto padding_vec = expand_param_if_needed(padding, "padding", 2);
+  auto stride_vec = expand_param_if_needed(stride, "stride", dims);
+  auto padding_vec = expand_param_if_needed(padding, "padding", dims);
   auto padding_vec_l = padding_vec;
   auto padding_vec_r = padding_vec;
-  auto dilation_vec = expand_param_if_needed(dilation, "dilation", 2);
+  auto dilation_vec = expand_param_if_needed(dilation, "dilation", dims);
 
   const ideep::tensor& x = itensor_from_mkldnn(input);
   std::vector<int64_t> output_sizes;
@@ -152,7 +184,24 @@ Tensor mkldnn_max_pool2d(
     IntArrayRef padding,
     IntArrayRef dilation,
     bool ceil_mode) {
-  return _mkldnn_pool2d(
+  return _mkldnn_pooling(
+      input,
+      kernel_size,
+      stride,
+      padding,
+      dilation,
+      ceil_mode,
+      ideep::algorithm::pooling_max);
+}
+
+Tensor mkldnn_max_pool3d(
+    const Tensor& input,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    IntArrayRef dilation,
+    bool ceil_mode) {
+  return _mkldnn_pooling(
       input,
       kernel_size,
       stride,
@@ -172,7 +221,7 @@ Tensor mkldnn_avg_pool2d(
     c10::optional<int64_t> divisor_override) {
   TORCH_CHECK(!divisor_override.has_value(),
            "mkldnn_avg_pool2d operator does not support divisor");
-  return _mkldnn_pool2d(
+  return _mkldnn_pooling(
       input,
       kernel_size,
       stride,
@@ -192,28 +241,59 @@ Tensor& mkldnn_avg_pool2d_out(
     bool ceil_mode,
     bool count_include_pad,
     c10::optional<int64_t> divisor_override) {
-  AT_ERROR(
-      "mkldnn_avg_pool2d_out: in-place mkldnn operations are not supported yet");
+  TORCH_CHECK(false, "mkldnn_avg_pool2d_out: in-place mkldnn operations are not supported yet");
+}
+
+Tensor mkldnn_avg_pool3d(
+    const Tensor& input,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    bool ceil_mode,
+    bool count_include_pad,
+    c10::optional<int64_t> divisor_override) {
+  TORCH_CHECK(!divisor_override.has_value(), "mkldnn_avg_pool3d operator does not support divisor");
+  return _mkldnn_pooling(
+      input,
+      kernel_size,
+      stride,
+      padding,
+      /*dilation*/ std::vector<int64_t>{1, 1, 1},
+      ceil_mode,
+      count_include_pad ? ideep::algorithm::pooling_avg_include_padding
+                        : ideep::algorithm::pooling_avg_exclude_padding);
+}
+
+Tensor& mkldnn_avg_pool3d_out(
+    Tensor& output,
+    const Tensor& input,
+    IntArrayRef kernel_size,
+    IntArrayRef stride,
+    IntArrayRef padding,
+    bool ceil_mode,
+    bool count_include_pad,
+    c10::optional<int64_t> divisor_override) {
+  TORCH_CHECK(false, "mkldnn_avg_pool3d_out: in-place mkldnn operations are not supported yet");
 }
 
 Tensor mkldnn_adaptive_avg_pool2d(
     Tensor const& input,
     IntArrayRef output_size) {
-  AT_ASSERTM(input.dim() == 4, "mkldnn_adaptive_avg_pool2d: Expect 2D input");
+  TORCH_CHECK(input.dim() == 4, "mkldnn_adaptive_avg_pool2d: Expect 2D input");
 
   auto output_size_vec =
       expand_param_if_needed(output_size, "output_size", input.dim() - 2);
   std::vector<int64_t> kernel_size(input.dim() - 2);
   for (int64_t i = 2; i < input.dim(); ++i) {
     auto s1 = input.size(i);
     auto s2 = output_size_vec[i - 2];
-    AT_ASSERTM(s2 != 0, "output size can not be zero");
-    AT_ASSERTM(
+    TORCH_CHECK(s2 != 0, "output size can not be zero");
+    TORCH_CHECK(
         s1 % s2 == 0,
         "input size is not divisible by the output size is not supported yet");
     kernel_size[i - 2] = s1 / s2;
   }
-  return _mkldnn_pool2d(
+  return _mkldnn_pooling(
       input,
       kernel_size,
       /*stride*/ kernel_size,
@@ -227,8 +307,7 @@ Tensor& mkldnn_adaptive_avg_pool2d_out(
     Tensor& output,
     const Tensor& input,
     IntArrayRef output_size) {
-  AT_ERROR(
-      "mkldnn_adaptive_avg_pool2d_out: in-place mkldnn operations are not supported yet");
+  TORCH_CHECK(false, "mkldnn_adaptive_avg_pool2d_out: in-place mkldnn operations are not supported yet");
 }
 
 
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -1794,6 +1794,11 @@
   dispatch:
     MkldnnCPU: mkldnn_max_pool2d
 
+- func: mkldnn_max_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> Tensor
+  use_c10_dispatcher: full
+  dispatch:
+    MkldnnCPU: mkldnn_max_pool3d
+
 - func: quantized_max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor
   use_c10_dispatcher: full
   dispatch:
@@ -5851,13 +5856,15 @@
   dispatch:
     CPU: avg_pool3d_out_cpu
     CUDA: avg_pool3d_out_cuda
+    MkldnnCPU: mkldnn_avg_pool3d_out
 
 - func: avg_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor
   use_c10_dispatcher: full
   python_module: nn
   dispatch:
     CPU: avg_pool3d_cpu
     CUDA: avg_pool3d_cuda
+    MkldnnCPU: mkldnn_avg_pool3d
     QuantizedCPU: quantized_avg_pool3d
 
 - func: avg_pool3d_backward.grad_input(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] stride, int[3] padding, bool ceil_mode, bool count_include_pad, int? divisor_override, *, Tensor(a!) grad_input) -> Tensor(a!)
diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py
@@ -254,6 +254,25 @@ def test_max_pool2d_stride_none(self):
 
                 self.assertEqual(y1, y2.to_dense())
 
+    def test_max_pool3d(self):
+        N = torch.randint(3, 10, (1,)).item()
+        C = torch.randint(3, 10, (1,)).item()
+
+        for stride in [1, 2, 3]:
+            for D, H, W in [(64, 64, 64), (35, 39, 35), (16, 19, 20), [7, 8, 9]]:
+                x = torch.randn(N, C, D, H, W, dtype=torch.float32) * 10
+
+                for ceil_mode in [False, True]:
+                    max_pool3d = torch.nn.MaxPool3d(
+                        kernel_size=3 if not ceil_mode else 7,
+                        stride=stride,
+                        padding=1,
+                        ceil_mode=ceil_mode)
+
+                    self.assertEqual(
+                        max_pool3d(x),
+                        max_pool3d(x.to_mkldnn()).to_dense())
+
     def test_avg_pool2d(self):
         N = torch.randint(3, 10, (1,)).item()
         C = torch.randint(3, 10, (1,)).item()
@@ -291,6 +310,22 @@ def test_avg_pool2d_stride_none(self):
 
             self.assertEqual(y1, y2.to_dense())
 
+    def test_avg_pool3d(self):
+        N = torch.randint(3, 10, (1,)).item()
+        C = torch.randint(3, 10, (1,)).item()
+        x = torch.randn(N, C, 64, 64, 64, dtype=torch.float32) * 10
+
+        for count_include_pad in [True, False]:
+            avg_pool3d = torch.nn.AvgPool3d(
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                count_include_pad=count_include_pad)
+
+            self.assertEqual(
+                avg_pool3d(x),
+                avg_pool3d(x.to_mkldnn()).to_dense())
+
     def test_adaptive_avg_pool2d(self):
         N = torch.randint(3, 10, (1,)).item()
         C = torch.randint(3, 10, (1,)).item()
diff --git a/torch/_overrides.py b/torch/_overrides.py
@@ -114,6 +114,7 @@ def get_ignored_functions():
         torch.mkldnn_convolution,
         torch.mkldnn_convolution_backward_weights,
         torch.mkldnn_max_pool2d,
+        torch.mkldnn_max_pool3d,
         torch.ones,
         torch.promote_types,
         torch.rand,
