[PIR] Open pir tests for cond/if op (#60903)

chen2016013 · web-flow · commit 4760488c7545 · 2024-01-18T15:56:16.000+08:00
* open pir tests for cond/if op

* fix
diff --git a/test/dygraph_to_static/test_return.py b/test/dygraph_to_static/test_return.py
@@ -26,6 +26,7 @@
 import paddle
 from paddle.base import core
 from paddle.jit.dy2static.utils import Dygraph2StaticException
+from paddle.pir_utils import test_with_pir_api
 
 SEED = 2020
 np.random.seed(SEED)
@@ -261,6 +262,7 @@ def _run(self):
         return res
 
     def _test_value_impl(self):
+        paddle.disable_static()
         with enable_to_static_guard(False):
             dygraph_res = self._run()
         static_res = self._run()
@@ -277,6 +279,7 @@ def _test_value_impl(self):
             self.assertEqual(dygraph_res, static_res)
 
     @test_ast_only
+    @test_with_pir_api
     def test_transformed_static_result(self):
         self.init_dygraph_func()
         if hasattr(self, "error"):
@@ -307,6 +310,7 @@ def _run(self):
         return res
 
     def _test_value_impl(self):
+        paddle.disable_static()
         with enable_to_static_guard(False):
             dygraph_res = self._run()
         static_res = self._run()
@@ -360,6 +364,7 @@ def _run(self):
         return res
 
     def _test_value_impl(self):
+        paddle.disable_static()
         with enable_to_static_guard(False):
             dygraph_res = self._run()
         static_res = self._run()
@@ -408,6 +413,7 @@ def _run(self):
         return res
 
     def _test_value_impl(self):
+        paddle.disable_static()
         with enable_to_static_guard(False):
             dygraph_res = self._run()
         static_res = self._run()
diff --git a/test/legacy_test/test_math_op_patch.py b/test/legacy_test/test_math_op_patch.py
@@ -19,6 +19,8 @@
 
 import paddle
 from paddle import base
+from paddle.framework import in_pir_mode
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestMathOpPatches(unittest.TestCase):
@@ -231,11 +233,13 @@ def test_equal(self):
         self.assertEqual(c.dtype, base.core.VarDesc.VarType.BOOL)
 
     @prog_scope()
+    @test_with_pir_api
     def test_equal_and_cond(self):
         a = paddle.static.data(name="a", shape=[-1, 1], dtype='float32')
-        a.desc.set_need_check_feed(False)
         b = paddle.static.data(name="b", shape=[-1, 1], dtype='float32')
-        b.desc.set_need_check_feed(False)
+        if not in_pir_mode():
+            a.desc.set_need_check_feed(False)
+            b.desc.set_need_check_feed(False)
         one = paddle.ones(shape=[1], dtype='int32')
         zero = paddle.zeros(shape=[1], dtype='int32')
         cond = one == zero
@@ -245,8 +249,9 @@ def test_equal_and_cond(self):
         exe = base.Executor(place)
         a_np = np.array([3, 4, 10, 14, 9, 18]).astype('float32')
         b_np = np.array([3, 4, 11, 15, 8, 18]).astype('float32')
+
         (c_np,) = exe.run(
-            base.default_main_program(),
+            paddle.static.default_main_program(),
             feed={"a": a_np, "b": b_np},
             fetch_list=[c],
         )
diff --git a/test/standalone_executor/test_standalone_controlflow.py b/test/standalone_executor/test_standalone_controlflow.py
@@ -18,7 +18,7 @@
 
 import paddle
 from paddle.base import core
-from paddle.base.framework import Program, program_guard
+from paddle.pir_utils import test_with_pir_api
 
 paddle.enable_static()
 
@@ -42,19 +42,17 @@ def _get_feed(self):
     def build_program(self):
         def true_func():
             return paddle.tensor.fill_constant(
-                shape=[1, 2], dtype='int32', value=1
-            ), paddle.tensor.fill_constant(
-                shape=[2, 3], dtype='bool', value=True
-            )
+                shape=[1, 2], dtype='float32', value=1
+            ), paddle.tensor.fill_constant(shape=[2, 3], dtype='int64', value=1)
 
         def false_func():
             return paddle.tensor.fill_constant(
                 shape=[3, 4], dtype='float32', value=3
             ), paddle.tensor.fill_constant(shape=[4, 5], dtype='int64', value=2)
 
-        main_program = Program()
-        startup_program = Program()
-        with program_guard(main_program, startup_program):
+        main_program = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program, startup_program):
             x = paddle.tensor.fill_constant(
                 shape=[1], dtype='float32', value=0.1
             )
@@ -84,10 +82,10 @@ def run_dygraph_once(self, feed):
         if x < y:
             out = [
                 paddle.tensor.fill_constant(
-                    shape=[1, 2], dtype='int32', value=1
+                    shape=[1, 2], dtype='float32', value=1
                 ).numpy(),
                 paddle.tensor.fill_constant(
-                    shape=[2, 3], dtype='bool', value=True
+                    shape=[2, 3], dtype='int64', value=1
                 ).numpy(),
             ]
         else:
@@ -111,6 +109,7 @@ def run_new_executor(self, feed):
         out = self._run(feed)
         return out
 
+    @test_with_pir_api
     def test_with_feed(self):
         feed = self._get_feed()
         paddle.enable_static()
diff --git a/test/standalone_executor/test_standalone_multiply_write.py b/test/standalone_executor/test_standalone_multiply_write.py
@@ -17,7 +17,8 @@
 from test_standalone_controlflow import TestCompatibility
 
 import paddle
-from paddle.base.framework import Program
+
+# from paddle.base.framework import Program
 
 paddle.enable_static()
 
@@ -28,8 +29,8 @@ def _get_feed(self):
         return None
 
     def build_program(self):
-        main_program = Program()
-        startup_program = Program()
+        main_program = paddle.static.Program()
+        startup_program = paddle.static.Program()
         with paddle.static.program_guard(main_program, startup_program):
             out = paddle.full((1,), 1)
             inp1 = paddle.full((1,), 2)
