[PIR] open test in test_ifelse for PIR (#60372)

chen2016013 · web-flow · commit 7909f768418b · 2023-12-29T10:31:06.000+08:00
* open test

* update

* update

* update as comment
diff --git a/paddle/common/ddim.cc b/paddle/common/ddim.cc
@@ -267,7 +267,7 @@ DDim DDim::transpose(const std::vector<int>& axis) const {
 
 DDim ComputeCompatibleDim(const DDim& dim1, const DDim& dim2) {
   IR_ENFORCE(dim1.size() == dim2.size(),
-             "Does not support rank inconsistency: dim1=%d, dim2=%d",
+             "Does not support rank inconsistency: rank1=%d, rank2=%d",
              dim1.size(),
              dim2.size());
   std::vector<int64_t> result;
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -102,12 +102,18 @@ void IfOp::Build(pir::Builder &builder,             // NOLINT
                               "The dtype in output[%d] of "
                               "true_block&false_block must be equal.",
                               i));
-        PADDLE_ENFORCE_EQ(l_type.data_layout(),
-                          r_type.data_layout(),
-                          phi::errors::PreconditionNotMet(
-                              "The date_layout in output[%d] of "
-                              "true_block&false_block must be equal.",
-                              i));
+        if (l_type.data_layout() != phi::DataLayout::UNDEFINED &&
+            r_type.data_layout() != phi::DataLayout::UNDEFINED) {
+          PADDLE_ENFORCE_EQ(
+              l_type.data_layout(),
+              r_type.data_layout(),
+              phi::errors::PreconditionNotMet(
+                  "The data_layout in output[%d] of "
+                  "true_block (%s) & false_block (%s) must be equal.",
+                  i,
+                  l_type.data_layout(),
+                  r_type.data_layout()));
+        }
         PADDLE_ENFORCE_EQ(l_type.lod(),
                           r_type.lod(),
                           phi::errors::PreconditionNotMet(
diff --git a/test/dygraph_to_static/ifelse_simple_func.py b/test/dygraph_to_static/ifelse_simple_func.py
@@ -60,9 +60,12 @@ def dyfunc_with_if_else2(x, col=100):
         # TODO: Don't support return non-Tensor in Tensor-dependent `if` statement currently.
         #  `x` is Tensor, `col` is not Tensor, and `col` is the return value of `true_fn` after transformed.
         # col = -1
-        col = paddle.tensor.fill_constant(shape=[1], value=-1, dtype="int64")
+        col = paddle.tensor.fill_constant(shape=[], value=-1, dtype="int64")
+    else:
+        col = paddle.tensor.fill_constant(shape=[], value=1, dtype="int64")
     if paddle.mean(x).numpy() > x.numpy()[row][col]:
-        y = paddle.nn.functional.relu(x)
+        x_pow = paddle.pow(x, 2)
+        y = paddle.nn.functional.relu(x_pow)
     else:
         x_pow = paddle.pow(x, 2)
         y = paddle.tanh(x_pow)
@@ -100,9 +103,12 @@ def false_fn_0(q, x, y):
         x = x + 1
         z = x + 2
         q = x + 3
+        m = x + 2
+        n = x + 3
     else:
         y = y + 1
         z = x - 2
+        q = x + 3
         m = x + 2
         n = x + 3
 
@@ -165,6 +171,22 @@ def nested_if_else(x_v):
             tmp = y * w
             y = paddle.nn.functional.relu(tmp)
             if paddle.mean(y).numpy() < batch_size:
+                tmp = paddle.tensor.fill_constant(
+                    y.shape, dtype='float32', value=-1
+                )
+                y = paddle.abs(y)
+            else:
+                tmp = paddle.tensor.fill_constant(
+                    y.shape, dtype='float32', value=-1
+                )
+                y = y - tmp
+        else:
+            tmp = y * w
+            y = paddle.nn.functional.relu(tmp)
+            if paddle.mean(y).numpy() < batch_size:
+                tmp = paddle.tensor.fill_constant(
+                    y.shape, dtype='float32', value=-1
+                )
                 y = paddle.abs(y)
             else:
                 tmp = paddle.tensor.fill_constant(
@@ -173,6 +195,11 @@ def nested_if_else(x_v):
                 y = y - tmp
     else:
         y = x_v - bias
+        w = paddle.tensor.fill_constant([feat_size], dtype='float32', value=10)
+        tmp = y * w
+        y = paddle.nn.functional.relu(tmp)
+        tmp = paddle.tensor.fill_constant(y.shape, dtype='float32', value=-1)
+        y = paddle.abs(y)
     return y
 
 
@@ -223,12 +250,14 @@ def nested_if_else_3(x):
             )
             # `z` is created in above code block.
             z = y + 1
+            out = x - 1
         else:
             res = paddle.tensor.fill_constant(
                 value=3, shape=x.shape, dtype="int32"
             )
             # `out` is a new var.
             out = x + 1
+            z = y - 1
     return res
 
 
@@ -378,7 +407,6 @@ def __init__(self):
 
 def if_tensor_case(x):
     x = base.dygraph.to_variable(x)
-
     mean = paddle.mean(x)
     # It is equivalent to `if mean != 0`
     if mean:
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_pt_and_pir,
     test_legacy_only,
 )
 from ifelse_simple_func import (
@@ -69,13 +70,21 @@ def setUp(self):
         self.error = "Your if/else have different number of return value."
 
     @test_ast_only
+    @test_legacy_and_pt_and_pir
     def test_error(self):
         if self.dyfunc:
             with self.assertRaisesRegex(Dygraph2StaticException, self.error):
                 with enable_to_static_guard(True):
                     self.assertTrue(paddle.jit.to_static(self.dyfunc)(self.x))
 
 
+class TestDy2StIfElseRetInt2(TestDy2staticException):
+    def setUp(self):
+        self.x = np.random.random([5]).astype('float32')
+        self.error = "Your if/else have different number of return value."
+        self.dyfunc = dyfunc_ifelse_ret_int2
+
+
 class TestDygraphIfElse(Dy2StTestBase):
     def setUp(self):
         self.x = np.random.random([10, 16]).astype('float32')
@@ -104,7 +113,6 @@ def setUp(self):
 
     # TODO(dev): fix AST mode
     @disable_test_case((ToStaticMode.AST, IrMode.PT))
-    @disable_test_case((ToStaticMode.AST, IrMode.LEGACY_IR))
     def test_ast_to_func(self):
         self.assertTrue((self._run_dygraph() == self._run_static()).all())
 
@@ -143,6 +151,10 @@ def setUp(self):
         self.x = np.random.random([10, 16]).astype('float32')
         self.dyfunc = dyfunc_with_if_else_with_list_generator
 
+    @test_legacy_and_pt_and_pir
+    def test_ast_to_func(self):
+        self.assertTrue((self._run_dygraph() == self._run_static()).all())
+
 
 class TestDygraphNestedIfElse(Dy2StTestBase):
     def setUp(self):
@@ -172,6 +184,10 @@ def setUp(self):
         self.x = np.random.random([10, 16]).astype('float32')
         self.dyfunc = nested_if_else_2
 
+    @test_legacy_and_pt_and_pir
+    def test_ast_to_func(self):
+        self.assertTrue((self._run_dygraph() == self._run_static()).all())
+
 
 class TestDygraphNestedIfElse3(Dy2StTestBase):
     def setUp(self):
@@ -269,12 +285,20 @@ def setUp(self):
         self.x = np.random.random([10, 16]).astype('float32')
         self.dyfunc = if_with_and_or_2
 
+    @test_legacy_and_pt_and_pir
+    def test_ast_to_func(self):
+        self.assertTrue((self._run_dygraph() == self._run_static()).all())
+
 
 class TestDygraphIfElseWithAndOr3(TestDygraphIfElse):
     def setUp(self):
         self.x = np.random.random([10, 16]).astype('float32')
         self.dyfunc = if_with_and_or_3
 
+    @test_legacy_and_pt_and_pir
+    def test_ast_to_func(self):
+        self.assertTrue((self._run_dygraph() == self._run_static()).all())
+
 
 class TestDygraphIfElseWithAndOr4(TestDygraphIfElse):
     def setUp(self):
@@ -439,11 +463,14 @@ def init_net(self):
 
     def _run(self, mode, to_static):
         with enable_to_static_guard(to_static):
-            net = paddle.jit.to_static(self.Net(mode))
+            if to_static:
+                net = paddle.jit.to_static(self.Net(mode))
+            else:
+                net = self.Net(mode)
             ret = net(self.x, self.y)
-
             return ret.numpy()
 
+    @test_legacy_and_pt_and_pir
     def test_train_mode(self):
         self.assertTrue(
             (
@@ -452,6 +479,7 @@ def test_train_mode(self):
             ).all()
         )
 
+    @test_legacy_and_pt_and_pir
     def test_infer_mode(self):
         self.assertTrue(
             (
@@ -467,6 +495,7 @@ def init_net(self):
 
 
 class TestNewVarCreateInOneBranch(Dy2StTestBase):
+    @test_legacy_and_pt_and_pir
     def test_var_used_in_another_for(self):
         def case_func(training):
             # targets and targets_list is dynamically defined by training
@@ -510,13 +539,6 @@ def test_ast_to_func(self):
         self.assertIsInstance(self.out[1], int)
 
 
-class TestDy2StIfElseRetInt2(TestDy2staticException):
-    def setUp(self):
-        self.x = np.random.random([5]).astype('float32')
-        self.error = "Your if/else have different number of return value."
-        self.dyfunc = dyfunc_ifelse_ret_int2
-
-
 class TestDy2StIfElseRetInt3(TestDy2StIfElseRetInt1):
     def setUp(self):
         self.x = np.random.random([5]).astype('float32')
