[Inference] dce pass should not delete control flow op and inplace op (#69112)

yuanlehome · web-flow · commit a69d69bf305c · 2024-11-04T19:07:23.000-08:00
* dce pass should not delete control flow op and inplace op

* fix ci
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_dialect.h b/paddle/fluid/pir/dialect/operator/ir/op_dialect.h
@@ -46,6 +46,11 @@ inline bool IsCustomOp(pir::Operation* op) {
   return op_name.find("custom_op") != op_name.npos;
 }
 
+inline bool IsInplaceOp(pir::Operation* op) {
+  std::string op_name = op->name();
+  return op_name.back() == '_';
+}
+
 inline bool IsTensorRTOp(pir::Operation* op) {
   std::string op_name = op->name();
   return op_name == "pd_op.tensorrt_engine";
diff --git a/paddle/fluid/pir/transforms/general/dead_code_elimination_pass.cc b/paddle/fluid/pir/transforms/general/dead_code_elimination_pass.cc
@@ -16,6 +16,7 @@
 #include <cstdint>
 
 #include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
 
@@ -58,7 +59,9 @@ class DeadCodeEliminationPass : public pir::Pass {
     for (auto& op : block) {
       if (op.HasTrait<pir::SideEffectTrait>() ||
           op.isa<paddle::dialect::DataOp>() ||
-          paddle::dialect::IsCustomOp(&op)) {
+          op.isa<paddle::dialect::WhileOp>() ||
+          paddle::dialect::IsCustomOp(&op) ||
+          paddle::dialect::IsInplaceOp(&op)) {
         continue;
       }
       if (op.use_empty()) {
