[PIR]Migrate MaxPool2D/AvgPool2D into pir (#58114)

0x45f · web-flow · commit d40e7688e7e7 · 2023-10-17T14:29:33.000+08:00
diff --git a/python/paddle/nn/functional/pooling.py b/python/paddle/nn/functional/pooling.py
@@ -15,7 +15,11 @@
 import numpy as np
 
 from paddle import _C_ops, _legacy_C_ops, in_dynamic_mode
-from paddle.base.framework import Variable, in_dygraph_mode
+from paddle.base.framework import (
+    Variable,
+    in_dygraph_mode,
+    in_dynamic_or_pir_mode,
+)
 
 from ...base.data_feeder import check_type, check_variable_and_dtype
 from ...base.layer_helper import LayerHelper
@@ -372,7 +376,7 @@ def avg_pool2d(
         padding, 2, channel_last, ceil_mode=ceil_mode
     )
 
-    if in_dygraph_mode():
+    if in_dynamic_or_pir_mode():
         output = _C_ops.pool2d(
             x,
             kernel_size,
@@ -1254,7 +1258,7 @@ def max_pool2d(
             "When setting return_mask to true, data_format must be set to NCHW in API:max_pool2d"
         )
 
-    if in_dygraph_mode():
+    if in_dynamic_or_pir_mode():
         if return_mask:
             output = _C_ops.max_pool2d_with_index(
                 x, kernel_size, stride, padding, False, False
diff --git a/python/paddle/utils/layers_utils.py b/python/paddle/utils/layers_utils.py
@@ -74,7 +74,9 @@ def convert_to_list(value, n, name, dtype=int):
                 + str(value)
             )
         for single_value in value_list:
-            assert not isinstance(single_value, Variable), (
+            assert not isinstance(
+                single_value, (Variable, paddle.pir.OpResult)
+            ), (
                 "Required numerical type with '%s', but received Tensor."
                 % dtype
             )
diff --git a/test/legacy_test/test_pool2d_api.py b/test/legacy_test/test_pool2d_api.py
@@ -25,6 +25,7 @@
 from paddle import base
 from paddle.base import core
 from paddle.nn.functional import avg_pool2d, max_pool2d
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestPool2D_API(unittest.TestCase):
@@ -52,7 +53,7 @@ def check_avg_static_results(self, place):
 
             exe = base.Executor(place)
             fetches = exe.run(
-                base.default_main_program(),
+                paddle.static.default_main_program(),
                 feed={"input": input_np},
                 fetch_list=[result],
             )
@@ -144,7 +145,7 @@ def check_max_static_results(self, place):
 
             exe = base.Executor(place)
             fetches = exe.run(
-                base.default_main_program(),
+                paddle.static.default_main_program(),
                 feed={"input": input_np},
                 fetch_list=[result],
             )
@@ -360,8 +361,6 @@ def test_pool2d(self):
         for place in self.places:
             self.check_max_dygraph_results(place)
             self.check_avg_dygraph_results(place)
-            self.check_max_static_results(place)
-            self.check_avg_static_results(place)
             self.check_max_dygraph_stride_is_none(place)
             self.check_avg_dygraph_stride_is_none(place)
             self.check_max_dygraph_padding(place)
@@ -370,6 +369,14 @@ def test_pool2d(self):
             self.check_max_dygraph_ceilmode_results(place)
             self.check_max_dygraph_nhwc_results(place)
 
+    @test_with_pir_api
+    def test_pool2d_static(self):
+        paddle.enable_static()
+        for place in self.places:
+            self.check_max_static_results(place)
+            self.check_avg_static_results(place)
+        paddle.disable_static()
+
 
 class TestPool2DError_API(unittest.TestCase):
     def test_error_api(self):

Original file line number	Diff line number	Diff line change
`@@ -74,7 +74,9 @@ def convert_to_list(value, n, name, dtype=int):`
`74`	`74`	`+ str(value)`
`75`	`75`	`)`
`76`	`76`	`for single_value in value_list:`
`77`		`- assert not isinstance(single_value, Variable), (`
	`77`	`+ assert not isinstance(`
	`78`	`+ single_value, (Variable, paddle.pir.OpResult)`
	`79`	`+ ), (`
`78`	`80`	`"Required numerical type with '%s', but received Tensor."`
`79`	`81`	`% dtype`
`80`	`82`	`)`