make rewrited max_pool2d op can be traced

Author: XiaobingSuper

intel_pytorch_extension_py/ops/pooling.py

@@ -2,7 +2,10 @@
 from torch.autograd import Function
 import torch.nn.functional as F
 import _torch_ipex as core
-from torch.nn.modules.utils import _single
+from torch.nn.modules.utils import _single, _pair
+from typing import List
+
+Vector = List[int]
 
 torch_adaptive_avg_pool2d = torch._C._nn.adaptive_avg_pool2d
 torch_max_pool2d = torch.max_pool2d
@@ -49,14 +52,6 @@ def adaptive_avg_pool2d(input, output_size):
         pass
     return torch_adaptive_avg_pool2d(input, output_size)
 
-def max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode):
-    try:
-        if input.device.type == 'dpcpp' and core.get_auto_dnnl():
-            return MaxPoolingFunction.apply(input, kernel_size, stride, padding, dilation, ceil_mode)
-    except RuntimeError:
-        pass
-    return torch_max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
-
 def max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode):
     try:
         if input.device.type == 'dpcpp' and core.get_auto_dnnl():
@@ -65,6 +60,9 @@ def max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode):
         pass
     return torch_max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode)
 
+def max_pool2d(input, kernel_size: Vector, stride: Vector, padding: Vector, dilation: Vector, ceil_mode: bool):
+    return torch.ops.torch_ipex.max_pool2d(input, _pair(kernel_size), _pair(stride), _pair(padding), _pair(dilation), ceil_mode)
+
 torch._C._nn.adaptive_avg_pool2d = adaptive_avg_pool2d
 torch.max_pool2d = max_pool2d
-torch.max_pool3d = max_pool3d
+torch.max_pool3d = max_pool3d

torch_ipex/csrc/cpu/CustomerOps.h

@@ -53,3 +53,57 @@ class NewLinearOp : public torch::autograd::Function<NewLinearOp> {
       return {grad_input, grad_weight, grad_bias};
     }
 };
+
+class NewMaxPoolingOp : public torch::autograd::Function<NewMaxPoolingOp> {
+  public:
+      static at::Tensor forward(
+        torch::autograd::AutogradContext* ctx,
+        at::Tensor input,
+        at::IntArrayRef kernel_size,
+        at::IntArrayRef stride,
+        at::IntArrayRef padding,
+        at::IntArrayRef dilation,
+        bool ceil_mode) {
+        ctx->saved_data["kernel_size"] = kernel_size;
+        ctx->saved_data["stride"] = stride;
+        ctx->saved_data["padding"] = padding;
+        ctx->saved_data["dilation"] = dilation;
+        ctx->saved_data["ceil_mode"] = ceil_mode;
+        if (torch_ipex::check_auto_dnnl() && input.device().type() == c10::DeviceType::DPCPP) {
+          at::Tensor output = torch_ipex::cpu::AtenIpexCPUDev::dil_max_pooling(input, kernel_size, stride,
+              padding, dilation, ceil_mode);
+          ctx->save_for_backward({input, output});
+          return output;
+        } else {
+          at::Tensor output, indices;
+          std::tie(output, indices) = at::max_pool2d_with_indices(input, kernel_size, stride, padding, dilation, ceil_mode);
+          ctx->save_for_backward({input, indices});
+          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 indices = saved[1];
+
+      at::Tensor grad_output = grad_outputs[0];
+      at::Tensor grad_input;
+      at::IntArrayRef kernel_size = at::IntArrayRef(ctx->saved_data["kernel_size"].toIntVector());
+      at::IntArrayRef stride = at::IntArrayRef(ctx->saved_data["stride"].toIntVector());
+      at::IntArrayRef padding = at::IntArrayRef(ctx->saved_data["padding"].toIntVector());
+      at::IntArrayRef dilation = at::IntArrayRef(ctx->saved_data["dilation"].toIntVector());
+      bool ceil_mode = ctx->saved_data["ceil_mode"].toBool();
+
+      if (torch_ipex::check_auto_dnnl() && input.device().type() == c10::DeviceType::DPCPP) {
+        grad_input = torch_ipex::cpu::AtenIpexCPUDev::dil_max_pooling_backward(
+            grad_output, indices, input, kernel_size, stride, padding, dilation, ceil_mode);
+      } else {
+        grad_input = at::max_pool2d_with_indices_backward(grad_output, input, kernel_size,
+            stride, padding, dilation, ceil_mode, indices);
+      }
+      return {grad_input};
+    }
+};

torch_ipex/csrc/cpu/ExtendOPs.cpp

@@ -468,7 +468,7 @@ at::Tensor AtenIpexTypeExt::adaptive_avg_pool2d_backward(const at::Tensor& grad_
 }
 
 at::Tensor AtenIpexTypeExt::max_pooling(const at::Tensor& input, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {
-    return cpu::AtenIpexCPUDev::dil_max_pooling(input, kernel_size, stride, padding, dilation, ceil_mode);
+    return NewMaxPoolingOp::apply(input, kernel_size, stride, padding, dilation, ceil_mode);
 }
 
 at::Tensor AtenIpexTypeExt::max_pooling_backward(const at::Tensor& grad_output, const at::Tensor& output, const at::Tensor& input, at::IntArrayRef kernel_size, at::IntArrayRef stride, at::IntArrayRef padding, at::IntArrayRef dilation, bool ceil_mode) {

torch_ipex/csrc/cpu/RegisterOps.cpp

@@ -3,9 +3,9 @@
 
 static auto registry =
     torch::RegisterOperators()
-       .op("torch_ipex::linear",
-          [](const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias) {
-          return torch_ipex::AtenIpexTypeExt::linear(input, weight, bias);
+       .op("torch_ipex::linear", &torch_ipex::AtenIpexTypeExt::linear)
+       .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);
         });
 
-