[Dy2St][PIR] Enable more control flow uts (#62811)

SigureMo · web-flow · commit b5b50b8a704d · 2024-03-18T22:12:03.000+08:00
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -3865,7 +3865,6 @@ void SliceArrayDenseInferMeta(const MetaTensor& input,
   if (config.is_runtime) {
     return;
   }
-  // out->set_dims(input.dims());
   out->set_dtype(input.dtype());
   out->set_dims(input.dims());
 }
diff --git a/test/dygraph_to_static/ifelse_simple_func.py b/test/dygraph_to_static/ifelse_simple_func.py
@@ -124,7 +124,7 @@ def dyfunc_with_if_else_early_return1():
         b = paddle.zeros([3, 3])
         return a, b
     a = paddle.zeros([2, 2]) + 1
-    return a, None
+    return a, paddle.zeros([3, 3]) + 1
 
 
 def dyfunc_with_if_else_early_return2():
@@ -138,7 +138,7 @@ def dyfunc_with_if_else_early_return2():
         d = paddle.zeros([3, 3]) + 1
         return c, d
     e = paddle.zeros([2, 2]) + 3
-    return e, None
+    return e, paddle.zeros([3, 3]) + 3
 
 
 def dyfunc_with_if_else_with_list_generator(x):
diff --git a/test/dygraph_to_static/test_cache_program.py b/test/dygraph_to_static/test_cache_program.py
@@ -172,6 +172,7 @@ def sum_under_while(limit):
 
 
 class TestToOutputWithCache(Dy2StTestBase):
+    @test_legacy_and_pt_and_pir
     def test_output(self):
         ret = paddle.jit.to_static(sum_even_until_limit)(80, 10)
         self.assertEqual(ret.numpy(), 30)
diff --git a/test/dygraph_to_static/test_for_enumerate.py b/test/dygraph_to_static/test_for_enumerate.py
@@ -554,6 +554,7 @@ def tearDown(self):
 
     @test_legacy_and_pt_and_pir
     def test_for_zip_error(self):
+        # TODO(pir-save-load): enable PIR test after support PIR save load
         with self.assertRaises(RuntimeError):
             model_path = os.path.join(self.temp_dir.name, 'for_zip_error')
             paddle.jit.save(
@@ -568,6 +569,7 @@ def test_for_zip_error(self):
             )
 
     def test_for_zip(self):
+        # TODO(pir-save-load): enable PIR test after support PIR save load
         model_path = os.path.join(self.temp_dir.name, 'for_zip')
         paddle.jit.save(
             paddle.jit.to_static(
diff --git a/test/dygraph_to_static/test_ifelse.py b/test/dygraph_to_static/test_ifelse.py
@@ -22,6 +22,7 @@
     disable_test_case,
     enable_to_static_guard,
     test_ast_only,
+    test_legacy_and_pir,
     test_legacy_and_pt_and_pir,
     test_pir_only,
 )
@@ -564,8 +565,7 @@ def forward(self, a, b, c):
 
 
 class TestDy2StIfElseBackward(Dy2StTestBase):
-    @test_ast_only
-    @test_pir_only
+    @test_legacy_and_pir
     def test_run_backward(self):
         a = paddle.randn((4, 3), dtype='float32')
         a.stop_gradient = False
diff --git a/test/dygraph_to_static/test_logical.py b/test/dygraph_to_static/test_logical.py
@@ -18,7 +18,11 @@
 import unittest
 
 import numpy as np
-from dygraph_to_static_utils import Dy2StTestBase, enable_to_static_guard
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    enable_to_static_guard,
+    test_legacy_and_pt_and_pir,
+)
 
 import paddle
 from paddle.jit.dy2static.transformers.logical_transformer import (
@@ -186,6 +190,7 @@ class TestLogicalNot(TestLogicalBase):
     def _set_test_func(self):
         self.dygraph_func = test_logical_not
 
+    @test_legacy_and_pt_and_pir
     def test_transformed_result(self):
         dygraph_res = self._run_dygraph()
         static_res = self._run_static()
@@ -201,6 +206,7 @@ class TestLogicalNot2(TestLogicalBase):
     def _set_test_func(self):
         self.dygraph_func = test_logical_not_2
 
+    @test_legacy_and_pt_and_pir
     def test_transformed_result(self):
         dygraph_res = self._run_dygraph()
         static_res = self._run_static()
@@ -250,10 +256,12 @@ def _set_test_func(self):
 
 
 class TestCmpopNodeToStr(Dy2StTestBase):
+    @test_legacy_and_pt_and_pir
     def test_exception(self):
         with self.assertRaises(KeyError):
             cmpop_node_to_str(gast.Or())
 
+    @test_legacy_and_pt_and_pir
     def test_expected_result(self):
         self.assertEqual(cmpop_node_to_str(gast.Eq()), "==")
         self.assertEqual(cmpop_node_to_str(gast.NotEq()), "!=")
diff --git a/test/dygraph_to_static/test_program_translator.py b/test/dygraph_to_static/test_program_translator.py
@@ -20,13 +20,9 @@
 import numpy as np
 from dygraph_to_static_utils import (
     Dy2StTestBase,
-    IrMode,
-    ToStaticMode,
-    disable_test_case,
     enable_to_static_guard,
     test_ast_only,
     test_legacy_and_pt_and_pir,
-    test_legacy_only,
 )
 from ifelse_simple_func import (
     dyfunc_with_if_else_early_return1,
@@ -308,8 +304,7 @@ def test_raise_error(self):
 
 
 class TestIfElseEarlyReturn(Dy2StTestBase):
-    # Why add test_legacy_only? : PIR not support if true and false branch output with different rank
-    @test_legacy_only
+    @test_legacy_and_pt_and_pir
     def test_ifelse_early_return1(self):
         answer = np.zeros([2, 2]) + 1
         static_func = paddle.jit.to_static(dyfunc_with_if_else_early_return1)
@@ -321,7 +316,7 @@ def test_ifelse_early_return1(self):
         elif isinstance(out, tuple):
             np.testing.assert_allclose(answer, out[0].numpy(), rtol=1e-05)
 
-    @disable_test_case((ToStaticMode.AST, IrMode.PT))
+    @test_legacy_and_pt_and_pir
     def test_ifelse_early_return2(self):
         answer = np.zeros([2, 2]) + 3
         static_func = paddle.jit.to_static(dyfunc_with_if_else_early_return2)
diff --git a/test/dygraph_to_static/test_warning.py b/test/dygraph_to_static/test_warning.py
@@ -18,7 +18,7 @@
 from dygraph_to_static_utils import (
     Dy2StTestBase,
     test_ast_only,
-    test_legacy_only,
+    test_legacy_and_pir,
 )
 
 import paddle
@@ -43,9 +43,8 @@ def false_fn():
 
 
 class TestReturnNoneInIfelse(Dy2StTestBase):
-    # Why add test_legacy_only? : PIR not support if true and false branch output with different dtype
-    @test_legacy_only
     @test_ast_only
+    @test_legacy_and_pir
     def test_dy2static_warning(self):
         paddle.disable_static()
         with warnings.catch_warnings(record=True) as w:
@@ -62,6 +61,8 @@ def test_dy2static_warning(self):
                     break
             self.assertTrue(flag)
 
+    # TODO(cleanup-legacy-ir): This case cannot be supported by PIR, we should remove this
+    # in the future.
     def test_cond_warning(self):
         paddle.enable_static()
         with warnings.catch_warnings(record=True) as w:

Original file line number	Diff line number	Diff line change
`@@ -3865,7 +3865,6 @@ void SliceArrayDenseInferMeta(const MetaTensor& input,`
`3865`	`3865`	`if (config.is_runtime) {`
`3866`	`3866`	`return;`
`3867`	`3867`	`}`
`3868`		`- // out->set_dims(input.dims());`
`3869`	`3868`	`out->set_dtype(input.dtype());`
`3870`	`3869`	`out->set_dims(input.dims());`
`3871`	`3870`	`}`