[PIR] Avoid replace op which has subblocks in CSE (#65722)

SigureMo · web-flow · commit baea4ab71c1c · 2024-07-05T03:44:59.000+08:00
diff --git a/paddle/fluid/pir/transforms/general/common_subexpression_elimination_pass.cc b/paddle/fluid/pir/transforms/general/common_subexpression_elimination_pass.cc
@@ -325,6 +325,12 @@ struct Expression {
               << " vs rhs [" << rhs << "] " << rhs->name() << " equal";
       return true;
     }
+    if (lhs->num_regions() > 0 || rhs->num_regions() > 0) {
+      VLOG(7) << "[CheckOperationEqual] lhs [" << lhs << "] " << lhs->name()
+              << " vs rhs [" << rhs << "] " << rhs->name()
+              << " has region, which is not supported";
+      return false;
+    }
     if (lhs->name() != rhs->name()) {
       VLOG(7) << "[CheckOperationEqual] lhs [" << lhs << "] " << lhs->name()
               << " vs rhs [" << rhs << "] " << rhs->name() << " name not equal";
@@ -480,6 +486,14 @@ struct CSEAnalyzer {
       return;
     }
 
+    // Handle sub blocks
+    for (auto& region : *op) {
+      for (auto& block : region) {
+        SimplifyBlock(&block, expression_table);
+      }
+    }
+
+    // Handle the operation
     auto expr = expression_table->CreateExpression(op);
     expression_table->RegisiterExpression(expr);
     auto maybe_same_expression = expression_table->Lookup(expr);
@@ -492,12 +506,6 @@ struct CSEAnalyzer {
             std::make_pair(expr.op(), maybe_same_expression.value().op()));
       }
     }
-    // Handle sub blocks
-    for (auto& region : *op) {
-      for (auto& block : region) {
-        SimplifyBlock(&block, expression_table);
-      }
-    }
   }
 
   void SimplifyBlock(pir::Block* block,
diff --git a/test/ir/pir/test_cse_pass.py b/test/ir/pir/test_cse_pass.py
@@ -373,6 +373,28 @@ def test_can_not_replace_used_by_inplace_op(self):
             paddle.base.libpaddle.pir.apply_cse_pass(main_program)
             self.assert_op_count_equal(main_program, {"pd_op.add": 2})
 
+    def test_can_not_replace_op_with_subblocks(self):
+        with program_scope_guard() as main_program:
+            # Inputs
+            x1 = paddle.static.data("x1", [2, 2], dtype="float32")
+            x2 = paddle.static.data("x2", [2, 2], dtype="float32")
+            cond = paddle.static.data("cond", [], dtype="bool")
+            loop_var = paddle.static.data("i", [], dtype="int32")
+
+            def loop_body(loop_var):
+                tmp_var = x1 * x2
+                return [loop_var + 1]
+
+            def get_cond(loop_var):
+                return cond
+
+            paddle.static.nn.while_loop(get_cond, loop_body, [loop_var])
+            paddle.static.nn.while_loop(get_cond, loop_body, [loop_var])
+
+            self.assert_op_count_equal(main_program, {"pd_op.while": 2})
+            paddle.base.libpaddle.pir.apply_cse_pass(main_program)
+            self.assert_op_count_equal(main_program, {"pd_op.while": 2})
+
 
 if __name__ == "__main__":
     unittest.main()
