[PIR] pir onednn support pad pool2d (#62085)

wanghuancoder · web-flow · commit 99f6b75e2350 · 2024-02-27T19:03:07.000+08:00
* pir onednn support pad pool2d
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -190,16 +190,22 @@
 
 - op : nearest_interp
 
-# - op : pad
+- op : pad
 
 - op : pad3d
   extra_args :
   data_format_tensors : x
   dynamic_fallback : True
 
-# - op : pool2d
+- op : pool2d
+  extra_args : bool use_quantizer=false, str mkldnn_data_type="float32", bool is_test=false
+  data_format_tensors : x
+  dynamic_fallback : True
 
-# - op : pool2d_grad
+- op : pool2d_grad
+  extra_args : bool use_quantizer=false, str mkldnn_data_type="float32", bool is_test=false
+  data_format_tensors : x, out, out_grad
+  dynamic_fallback : True
 
 - op : prelu
   extra_args : bool is_test=false, str mkldnn_data_type="float32"
diff --git a/paddle/phi/kernels/onednn/pool_grad_kernel.cc b/paddle/phi/kernels/onednn/pool_grad_kernel.cc
@@ -18,6 +18,21 @@
 #include "paddle/phi/core/kernel_registry.h"
 
 namespace phi {
+bool Pool2dGradCheckIfOneDNNSupport(const KernelContext* ctx) {
+  if (ctx->AttrAt<bool>(8) == false) {
+    // adaptive
+    return true;
+  }
+  // oneDNN is supporting only unchangable in size pool window
+  auto src_tz = common::vectorize(ctx->InputAt<phi::DenseTensor>(0).dims());
+  const TensorRef& kernel_size_tmp = ctx->AttrAt<TensorRef>(0);
+  IntArray kernel_size_array = IntArray(*kernel_size_tmp.Get());
+  std::vector<int64_t> kernel_size = kernel_size_array.GetData();
+  // Fast but not exhaustive check
+  return ((src_tz[src_tz.size() - 1] % kernel_size[1] == 0) &&
+          (src_tz[src_tz.size() - 2] % kernel_size[0] == 0));
+}
+
 template <typename T, typename Context>
 void Pool2dGradKernel(const Context& dev_ctx,
                       const DenseTensor& x,
@@ -100,4 +115,5 @@ PD_REGISTER_KERNEL(pool2d_grad,
                    float,
                    phi::dtype::bfloat16) {
   kernel->get_kerneltype_forvar_fn_ = phi::PoolOpGradGetKernelTypeForVar;
+  kernel->check_if_onednn_kernel_support_ = phi::Pool2dGradCheckIfOneDNNSupport;
 }
diff --git a/paddle/phi/kernels/onednn/pool_kernel.cc b/paddle/phi/kernels/onednn/pool_kernel.cc
@@ -18,6 +18,21 @@
 #include "paddle/phi/core/kernel_registry.h"
 
 namespace phi {
+bool Pool2dCheckIfOneDNNSupport(const KernelContext* ctx) {
+  if (ctx->AttrAt<bool>(8) == false) {
+    // adaptive
+    return true;
+  }
+  // oneDNN is supporting only unchangable in size pool window
+  auto src_tz = common::vectorize(ctx->InputAt<phi::DenseTensor>(0).dims());
+  const TensorRef& kernel_size_tmp = ctx->AttrAt<TensorRef>(0);
+  IntArray kernel_size_array = IntArray(*kernel_size_tmp.Get());
+  std::vector<int64_t> kernel_size = kernel_size_array.GetData();
+  // Fast but not exhaustive check
+  return ((src_tz[src_tz.size() - 1] % kernel_size[1] == 0) &&
+          (src_tz[src_tz.size() - 2] % kernel_size[0] == 0));
+}
+
 template <typename T, typename Context>
 void Pool2dKernel(const Context& dev_ctx,
                   const DenseTensor& x,
@@ -104,4 +119,5 @@ PD_REGISTER_KERNEL(pool2d,
                    uint8_t,
                    phi::dtype::bfloat16) {
   kernel->get_kerneltype_forvar_fn_ = phi::PoolOpGetKernelTypeForVar;
+  kernel->check_if_onednn_kernel_support_ = phi::Pool2dCheckIfOneDNNSupport;
 }
diff --git a/test/legacy_test/op_test.py b/test/legacy_test/op_test.py
@@ -2643,7 +2643,7 @@ def _is_skip_name(self, name):
         static_checker.check()
         outs, fetch_list = static_checker.outputs, static_checker.fetch_list
 
-        if check_pir_onednn:
+        if check_pir_onednn and place == base.CPUPlace():
             with pir_executor_guard():
                 pir_onednn_static_checker = StaticChecker(self, self.outputs)
                 pir_onednn_static_checker.check()
@@ -3313,7 +3313,7 @@ def check_grad_with_place(
             atol,
         )
 
-        if check_pir_onednn:
+        if check_pir_onednn and place == base.CPUPlace():
             with pir_executor_guard():
                 self.check_grad_with_place_for_static(
                     user_defined_grads,
diff --git a/test/legacy_test/test_pool2d_op.py b/test/legacy_test/test_pool2d_op.py
@@ -428,11 +428,13 @@ def test_check_output(self):
                 check_dygraph=(not self.use_mkldnn),
                 check_cinn=True,
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
         else:
             self.check_output(
                 check_dygraph=(not self.use_mkldnn),
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
     def test_check_grad(self):
@@ -448,6 +450,7 @@ def test_check_grad(self):
                 check_dygraph=(not self.use_mkldnn),
                 check_cinn=True,
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
         elif self.pool_type != "max":
             self.check_grad(
@@ -456,6 +459,7 @@ def test_check_grad(self):
                 max_relative_error=0.07,
                 check_dygraph=(not self.use_mkldnn),
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
     def init_data_format(self):
@@ -599,6 +603,7 @@ def test_check_output(self):
                         place,
                         check_dygraph=(not self.use_mkldnn),
                         check_cinn=True,
+                        check_pir_onednn=self.check_pir_onednn,
                     )
 
         def test_check_grad(self):
@@ -615,6 +620,7 @@ def test_check_grad(self):
                     'Out',
                     check_dygraph=(not self.use_mkldnn),
                     check_cinn=True,
+                    check_pir_onednn=self.check_pir_onednn,
                 )
 
     cls_name = "{}_{}".format(parent.__name__, "CUDNNFp16Op")
@@ -640,6 +646,7 @@ def test_check_output(self):
                         place,
                         check_dygraph=(not self.use_mkldnn),
                         check_cinn=True,
+                        check_pir_onednn=self.check_pir_onednn,
                     )
 
         def test_check_grad(self):
@@ -656,6 +663,7 @@ def test_check_grad(self):
                     'Out',
                     check_dygraph=(not self.use_mkldnn),
                     check_cinn=True,
+                    check_pir_onednn=self.check_pir_onednn,
                 )
 
     cls_name = "{}_{}".format(parent.__name__, "Fp16Op")
@@ -679,6 +687,7 @@ def test_check_output(self):
                     place,
                     check_dygraph=(not self.use_mkldnn),
                     check_cinn=True,
+                    check_pir_onednn=self.check_pir_onednn,
                 )
 
         def test_check_grad(self):
@@ -690,6 +699,7 @@ def test_check_grad(self):
                     'Out',
                     check_dygraph=(not self.use_mkldnn),
                     check_cinn=True,
+                    check_pir_onednn=self.check_pir_onednn,
                 )
 
     cls_name = "{}_{}".format(parent.__name__, "Bf16Op")
@@ -1025,6 +1035,7 @@ def test_check_grad(self):
                 max_relative_error=1.00,
                 check_cinn=True,
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
         elif self.pool_type == "max":
             self.check_grad(
@@ -1033,6 +1044,7 @@ def test_check_grad(self):
                 max_relative_error=1.00,
                 check_cinn=True,
                 check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
             )
 
 
diff --git a/test/mkldnn/test_pool2d_bf16_mkldnn_op.py b/test/mkldnn/test_pool2d_bf16_mkldnn_op.py
@@ -197,7 +197,7 @@ def setUp(self):
         self.outputs = {'Out': convert_float_to_uint16(output)}
 
     def test_check_output(self):
-        self.check_output_with_place(core.CPUPlace())
+        self.check_output_with_place(core.CPUPlace(), check_pir_onednn=True)
 
     def test_check_grad(self):
         x_grad = pool2d_backward_naive(
@@ -215,7 +215,11 @@ def test_check_grad(self):
         )
         x_grad = x_grad / np.prod(self.outputs['Out'].shape)
         self.check_grad_with_place(
-            core.CPUPlace(), {'X'}, 'Out', user_defined_grads=[x_grad]
+            core.CPUPlace(),
+            {'X'},
+            'Out',
+            user_defined_grads=[x_grad],
+            check_pir_onednn=True,
         )
 
 
@@ -247,7 +251,7 @@ def setUp(self):
         self.outputs = {'Out': convert_float_to_uint16(output)}
 
     def test_check_output(self):
-        self.check_output_with_place(core.CPUPlace())
+        self.check_output_with_place(core.CPUPlace(), check_pir_onednn=True)
 
     def test_check_grad(self):
         pass
diff --git a/test/mkldnn/test_pool2d_int8_mkldnn_op.py b/test/mkldnn/test_pool2d_int8_mkldnn_op.py
@@ -24,6 +24,7 @@
 class TestPool2DMKLDNNInt8_Op(TestPool2D_Op):
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
     def init_data_type(self):
         self.dtype = np.int8
@@ -54,7 +55,10 @@ def setUp(self):
     def test_check_output(self):
         # TODO(wangzhongpu): support mkldnn op in dygraph mode
         self.check_output_with_place(
-            core.CPUPlace(), atol=1e-5, check_dygraph=False
+            core.CPUPlace(),
+            atol=1e-5,
+            check_dygraph=False,
+            check_pir_onednn=True,
         )
 
     def test_check_grad(self):
diff --git a/test/mkldnn/test_pool2d_mkldnn_op.py b/test/mkldnn/test_pool2d_mkldnn_op.py
@@ -30,6 +30,7 @@ def create_test_mkldnn_use_ceil_class(parent):
     class TestMKLDNNPool2DUseCeilCase(parent):
         def init_kernel_type(self):
             self.use_mkldnn = True
+            self.check_pir_onednn = True
 
         def init_ceil_mode(self):
             self.ceil_mode = True
@@ -51,6 +52,7 @@ def create_test_mkldnn_class(parent):
     class TestMKLDNNCase(parent):
         def init_kernel_type(self):
             self.use_mkldnn = True
+            self.check_pir_onednn = True
 
         def init_data_type(self):
             self.dtype = np.float32
@@ -78,6 +80,7 @@ def init_pool_type(self):
 
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
     def init_test_case(self):
         self.ksize = [1, 1]
@@ -128,6 +131,7 @@ def init_shape(self):
 
     def init_kernel_type(self):
         self.use_mkldnn = True
+        self.check_pir_onednn = True
 
     def init_data_type(self):
         self.dtype = np.float32
