Add mkldnn support for adaptive_avg_pool2d

bddppq · bddppq · commit b6f75bc4a7a6 · 2019-04-30T11:11:30.000-07:00
diff --git a/aten/src/ATen/native/AdaptiveAveragePooling.cpp b/aten/src/ATen/native/AdaptiveAveragePooling.cpp
@@ -262,9 +262,11 @@ namespace {
     return output;
   }
 
-  Tensor adaptive_avg_pool2d(
-    at::Tensor const& input,
-    IntArrayRef output_size){
+  Tensor adaptive_avg_pool2d(at::Tensor const& input, IntArrayRef output_size) {
+    if (input.is_mkldnn()) {
+      return at::mkldnn_adaptive_avg_pool2d(input, output_size);
+    }
+
     if (output_size[0] == 1 && output_size[1] == 1) {
 //in this case, adaptive pooling is just computing mean over hw dimensions, which can be done more efficiently
        int64_t mean_size = input.size(-1) * input.size(-2);
diff --git a/aten/src/ATen/native/mkldnn/Pooling.cpp b/aten/src/ATen/native/mkldnn/Pooling.cpp
@@ -41,6 +41,19 @@ Tensor& mkldnn_avg_pool2d_out(
     bool count_include_pad) {
   AT_ERROR("mkldnn_avg_pool2d_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");
+}
+
+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");
+}
+
 } // namespace native
 } // namespace at
 
@@ -118,7 +131,7 @@ Tensor mkldnn_avg_pool2d(
       kernel_size,
       stride,
       padding,
-      std::vector<int64_t>{1, 1},
+      /*dilation*/ std::vector<int64_t>{1, 1},
       ceil_mode,
       count_include_pad ? ideep::algorithm::pooling_avg_include_padding
                         : ideep::algorithm::pooling_avg_exclude_padding);
@@ -136,6 +149,42 @@ Tensor& mkldnn_avg_pool2d_out(
       "mkldnn_avg_pool2d_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");
+
+  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 (size_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(
+        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(
+      input,
+      kernel_size,
+      /*stride*/ kernel_size,
+      /*padding*/ {0, 0},
+      /*dilation*/ {1, 1},
+      /*ceil_mode*/ false,
+      /*algo*/ ideep::algorithm::pooling_avg);
+}
+
+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");
+}
+
+
 } // namespace native
 } // namespace at
 
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -3634,10 +3634,16 @@
   dispatch:
     CPU: adaptive_avg_pool2d_out_cpu
     CUDA: adaptive_avg_pool2d_out_cuda
+    MkldnnCPU: mkldnn_adaptive_avg_pool2d_out
 
 - func: adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor
   python_module: nn
 
+- func: mkldnn_adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor
+  dispatch:
+    MkldnnCPU: mkldnn_adaptive_avg_pool2d
+  requires_tensor: True
+
 - func: _adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor
   dispatch:
     CPU: adaptive_avg_pool2d_cpu
diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py
@@ -148,6 +148,17 @@ def test_avg_pool2d(self):
                 avg_pool2d(x),
                 avg_pool2d(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()
+        x = torch.randn(N, C, 224, 224, dtype=torch.float32) * 100
+
+        adaptive_avg_pool2d = torch.nn.AdaptiveAvgPool2d(7)
+
+        self.assertEqual(
+            adaptive_avg_pool2d(x),
+            adaptive_avg_pool2d(x.to_mkldnn()).to_dense())
+
     def test_batch_norm2d(self):
         N = torch.randint(3, 10, (1,)).item()
         C = torch.randint(3, 100, (1,)).item()
