make rewrited AdaptiveAvgPool2d op can be traced

Author: XiaobingSuper

intel_pytorch_extension_py/ops/pooling.py

@@ -7,24 +7,8 @@
 
 Vector = List[int]
 
-torch_adaptive_avg_pool2d = torch._C._nn.adaptive_avg_pool2d
-torch_max_pool2d = torch.max_pool2d
 torch_max_pool3d = torch.max_pool3d
 
-class AdaptiveAvgPool2dFunction(Function):
-    @staticmethod
-    def forward(ctx, input, output_size):
-        output = core.adaptive_avg_pool2d(input, _single(output_size))
-        ctx.save_for_backward(input)
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        (input,) = ctx.saved_tensors
-        grad_output = grad_output.contiguous()
-        grad_input = core.adaptive_avg_pool2d_backward(grad_output, input)
-        return (grad_input, None)
-
 class MaxPoolingFunction(Function):
     @staticmethod
     def forward(ctx, input, kernel_size, stride, padding, dilation, ceil_mode):
@@ -44,13 +28,8 @@ def backward(ctx, grad_output):
         grad_input = core.max_pooling_backward(grad_output, output, input, ctx.kernel_size, ctx.stride, ctx.padding, ctx.dilation, ctx.ceil_mode)
         return (grad_input, None, None, None, None, None)
 
-def adaptive_avg_pool2d(input, output_size):
-    try:
-        if input.device.type == 'dpcpp' and core.get_auto_dnnl():
-            return AdaptiveAvgPool2dFunction.apply(input, output_size)
-    except RuntimeError:
-        pass
-    return torch_adaptive_avg_pool2d(input, output_size)
+def adaptive_avg_pool2d(input, output_size: Vector):
+    return torch.ops.torch_ipex.adaptive_avg_pool2d(input, _pair(output_size))
 
 def max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode):
     try:

torch_ipex/csrc/cpu/CustomOPs.h

@@ -107,3 +107,38 @@ class NewMaxPoolingOp : public torch::autograd::Function<NewMaxPoolingOp> {
       return {grad_input, at::Tensor(), at::Tensor(), at::Tensor(), at::Tensor(), at::Tensor()};
     }
 };
+
+class NewApaptiveAvgPoolingOp : public torch::autograd::Function<NewApaptiveAvgPoolingOp> {
+  public:
+      static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx,
+        at::Tensor input,
+        at::IntArrayRef output_size) {
+        ctx->save_for_backward({input});
+
+        at::Tensor output;
+        if (torch_ipex::check_auto_dnnl() && input.device().type() == c10::DeviceType::DPCPP) {
+          output = torch_ipex::cpu::AtenIpexCPUDev::dil_adaptive_avg_pool2d(input, output_size);
+        } else {
+          output = at::_adaptive_avg_pool2d(input, output_size);
+        }
+        return output;
+      }
+
+    static torch::autograd::tensor_list backward(
+        torch::autograd::AutogradContext* ctx,
+        torch::autograd::tensor_list grad_outputs) {
+      auto saved = ctx->get_saved_variables();
+      at::Tensor input = saved[0];
+
+      at::Tensor grad_output = grad_outputs[0].contiguous();
+      at::Tensor grad_input;
+
+      if (torch_ipex::check_auto_dnnl() && input.device().type() == c10::DeviceType::DPCPP) {
+        grad_input = torch_ipex::cpu::AtenIpexCPUDev::dil_adaptive_avg_pool2d_backward(grad_output, input);
+      } else {
+        grad_input = at::_adaptive_avg_pool2d_backward(grad_output, input);
+      }
+      return {grad_input, at::Tensor()};
+    }
+};

torch_ipex/csrc/cpu/ExtendOPs.cpp

@@ -10,7 +10,7 @@
 #include "xsmm/libxsmm_utils.h"
 #include "../utils.h"
 #include "DevOPs.h"
-#include "CustomOps.h"
+#include "CustomOPs.h"
 
 namespace torch_ipex {
 
@@ -460,7 +460,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> AtenIpexTypeExt::linear_backward(
 }
 
 at::Tensor AtenIpexTypeExt::adaptive_avg_pool2d(at::Tensor const& input, at::IntArrayRef output_size) {
-    return cpu::AtenIpexCPUDev::dil_adaptive_avg_pool2d(input, output_size);
+    return NewApaptiveAvgPoolingOp::apply(input, output_size);
 }
 
 at::Tensor AtenIpexTypeExt::adaptive_avg_pool2d_backward(const at::Tensor& grad_output, const at::Tensor& input) {

torch_ipex/csrc/cpu/RegisterOps.cpp

@@ -7,5 +7,8 @@ static auto registry =
        .op("torch_ipex::max_pool2d", [](const at::Tensor& self, c10::List<int64_t> kernel_size,
           c10::List<int64_t> stride, c10::List<int64_t> padding, c10::List<int64_t> dilation, bool ceil_mode=false){
           return torch_ipex::AtenIpexTypeExt::max_pooling(self, kernel_size.vec(), stride.vec(), padding.vec(), dilation.vec(), ceil_mode);
+        })
+       .op("torch_ipex::adaptive_avg_pool2d", [](const at::Tensor&self, c10::List<int64_t> output_size) {
+          return torch_ipex::AtenIpexTypeExt::adaptive_avg_pool2d(self, output_size.vec());
         });