Fix ShapeOrDataDimExpr simplify unwork (#62376)

JiaWenxuan · web-flow · commit 7129945f12c0 · 2024-03-07T16:07:35.000+08:00
* update test case

* fix

* fix concat op infer symbolic

* fix some bugs

* fix some bugs

* fix some bugs

* fix some bugs

* fix some bugs
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
@@ -94,9 +94,6 @@ void ApplyCinnPreprocessPass(
   if (has_dynamic_shape) {
     pass_manager->AddPass(cinn::dialect::ir::CreateConvert0DTo1DPass());
     pass_manager->AddPass(pir::CreateShapeOptimizationPass());
-    pass_manager->AddPass(cinn::dialect::ir::CreateSimplifyDimExprPass());
-    pass_manager->AddPass(
-        cinn::dialect::ir::CreateSubstituteDimExprBasedOnConstraintsPass());
     pass_manager->AddPass(cinn::dialect::ir::CreateConvert0DTo1DPass());
     pass_manager->AddPass(
         cinn::dialect::ir::CreateFuseShapeOpsIntoGenerateShapeOpPass());
@@ -130,6 +127,9 @@ void ApplyGroupOpPass(::pir::Program* program,
         cinn::dialect::ir::CreateFuseShapeOpsIntoGenerateShapeOpPass());
     pass_manager->AddPass(
         cinn::dialect::ir::CreateMoveGenerateShapeOpsToProloguePass());
+    pass_manager->AddPass(
+        cinn::dialect::ir::CreateSubstituteDimExprBasedOnConstraintsPass());
+    pass_manager->AddPass(cinn::dialect::ir::CreateSimplifyDimExprPass());
   }
 
   pass_manager->AddPass(cinn::dialect::ir::CreateDynamicReshapeOpPass());
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/simplify_dim_expr_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/simplify_dim_expr_pass.cc
@@ -28,11 +28,14 @@ namespace ir {
 namespace {
 
 template <typename DoEachT>
-void VisitEachOp(pir::ModuleOp module_op, const DoEachT& DoEach) {
-  for (uint32_t i = 0; i < module_op->num_regions(); i++) {
-    for (pir::Block& block : module_op->region(i)) {
-      for (pir::Operation& op : block) {
-        DoEach(op);
+void VisitEachOp(pir::Operation* op, const DoEachT& DoEach) {
+  for (uint32_t i = 0; i < op->num_regions(); i++) {
+    for (pir::Block& block : op->region(i)) {
+      for (pir::Operation& sub_op : block) {
+        DoEach(sub_op);
+        if (sub_op.num_regions() > 0) {
+          VisitEachOp(&sub_op, DoEach);
+        }
       }
     }
   }
@@ -90,24 +93,36 @@ symbol::ShapeOrDataDimExprs SimplifyShapeOrData(
   return std::visit(lambdas, shape_or_data.variant());
 }
 
-void SimplifyDimExpr(pir::ModuleOp module_op) {
+void SimplifyDimExpr(pir::Operation* module_op) {
   VLOG(4) << "SimplifyDimExpr start";
-  pir::ShapeConstraintIRAnalysis shape_analysis =
-      pir::ShapeAnalysisManager::Instance().Get(module_op.program());
+  pir::ShapeConstraintIRAnalysis* shape_analysis =
+      &pir::ShapeAnalysisManager::Instance().Get(
+          module_op->dyn_cast<pir::ModuleOp>().program());
+
   VisitEachOp(module_op, [&](pir::Operation& op) {
     VisitEachValue(op, [&](pir::Value value) {
-      if (!shape_analysis.HasShapeOrDataForValue(value)) {
+      if (!shape_analysis->HasShapeOrDataForValue(value)) {
         VLOG(4) << "SimplifyDimExpr: shape_analysis can't find ShapeOrData for "
                    "value of the op:"
                 << op.name();
       } else {
         const symbol::ShapeOrDataDimExprs& shape_or_data =
-            shape_analysis.GetShapeOrDataForValue(value);
+            shape_analysis->GetShapeOrDataForValue(value);
+        VLOG(8) << op.name() << "     origin_shape_or_data: " << shape_or_data;
         symbol::ShapeOrDataDimExprs simplified_shape_or_data =
             SimplifyShapeOrData(shape_or_data);
-        shape_analysis.SetShapeOrDataForValue(value, simplified_shape_or_data);
+        VLOG(8) << op.name()
+                << " simplified_shape_or_data: " << simplified_shape_or_data;
+        shape_analysis->SetShapeOrDataForValue(value, simplified_shape_or_data);
       }
     });
+    if (op.num_results() > 0) {
+      pir::shape::SetShapeAttrForOp(
+          &op, shape_analysis->GetShapeOrDataForValue(op.result(0)));
+    } else {
+      pir::shape::SetShapeAttrForOp(
+          &op, shape_analysis->GetShapeOrDataForValue(op.operand_source(0)));
+    }
     // TODO(JiaWenxuan): simplify the attribute "sym_shape_str" of the op
   });
   VLOG(4) << "SimplifyDimExpr end";
@@ -117,10 +132,7 @@ class SimplifyDimExprPass : public pir::Pass {
  public:
   SimplifyDimExprPass() : pir::Pass("simplify_dim_expr_pass", 1) {}
 
-  void Run(pir::Operation* op) override {
-    pir::ModuleOp module_op = op->dyn_cast<pir::ModuleOp>();
-    SimplifyDimExpr(module_op);
-  }
+  void Run(pir::Operation* op) override { SimplifyDimExpr(op); }
 
   bool CanApplyOn(pir::Operation* op) const override {
     return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/substitute_dim_expr_based_on_constraints_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/substitute_dim_expr_based_on_constraints_pass.cc
@@ -18,6 +18,7 @@
 
 #include "paddle/cinn/common/dim_expr_util.h"
 #include "paddle/cinn/common/union_find.h"
+#include "paddle/pir/include/dialect/shape/ir/shape_attribute.h"
 
 namespace cinn {
 namespace dialect {
@@ -26,11 +27,14 @@ namespace ir {
 namespace {
 
 template <typename DoEachT>
-void VisitEachOp(pir::ModuleOp module_op, const DoEachT& DoEach) {
-  for (uint32_t i = 0; i < module_op->num_regions(); i++) {
-    for (pir::Block& block : module_op->region(i)) {
-      for (pir::Operation& op : block) {
-        DoEach(op);
+void VisitEachOp(pir::Operation* op, const DoEachT& DoEach) {
+  for (uint32_t i = 0; i < op->num_regions(); i++) {
+    for (pir::Block& block : op->region(i)) {
+      for (pir::Operation& sub_op : block) {
+        DoEach(sub_op);
+        if (sub_op.num_regions() > 0) {
+          VisitEachOp(&sub_op, DoEach);
+        }
       }
     }
   }
@@ -133,25 +137,39 @@ std::unordered_map<symbol::DimExpr, symbol::DimExpr> GetDimExprSubstitution(
   return substitution_pattern;
 }
 
-void SubstituteDimExprBasedOnConstraints(pir::ModuleOp module_op) {
+void SubstituteDimExprBasedOnConstraints(pir::Operation* module_op) {
   VLOG(4) << "SubstituteDimExprBasedOnConstraints start";
-  pir::ShapeConstraintIRAnalysis shape_analysis =
-      pir::ShapeAnalysisManager::Instance().Get(module_op.program());
+  pir::ShapeConstraintIRAnalysis* shape_analysis =
+      &pir::ShapeAnalysisManager::Instance().Get(
+          module_op->dyn_cast<pir::ModuleOp>().program());
   const std::unordered_map<symbol::DimExpr, symbol::DimExpr>&
-      substitution_pattern = GetDimExprSubstitution(&shape_analysis);
+      substitution_pattern = GetDimExprSubstitution(shape_analysis);
+
   VisitEachOp(module_op, [&](pir::Operation& op) {
     VisitEachValue(op, [&](pir::Value value) {
-      if (!shape_analysis.HasShapeOrDataForValue(value)) {
+      if (!shape_analysis->HasShapeOrDataForValue(value)) {
         VLOG(4) << "Can not find ShapeOrData for value of op(" << op.name()
                 << ") in shape_analysis";
       } else {
         const symbol::ShapeOrDataDimExprs& origin_shape_or_data =
-            shape_analysis.GetShapeOrDataForValue(value);
+            shape_analysis->GetShapeOrDataForValue(value);
+        VLOG(8) << op.name()
+                << "      origin_shape_or_data: " << origin_shape_or_data;
         const symbol::ShapeOrDataDimExprs& substituted_shape_or_data =
             SubstituteShapeOrData(origin_shape_or_data, substitution_pattern);
-        shape_analysis.SetShapeOrDataForValue(value, substituted_shape_or_data);
+        VLOG(8) << op.name()
+                << " substituted_shape_or_data: " << substituted_shape_or_data;
+        shape_analysis->SetShapeOrDataForValue(value,
+                                               substituted_shape_or_data);
       }
     });
+    if (op.num_results() > 0) {
+      pir::shape::SetShapeAttrForOp(
+          &op, shape_analysis->GetShapeOrDataForValue(op.result(0)));
+    } else {
+      pir::shape::SetShapeAttrForOp(
+          &op, shape_analysis->GetShapeOrDataForValue(op.operand_source(0)));
+    }
     // TODO(JiaWenxuan): substitute the attribute "sym_shape_str" of the op
   });
   VLOG(4) << "SubstituteDimExprBasedOnConstraints end";
@@ -163,8 +181,7 @@ class SubstituteDimExprBasedOnConstraintsPass : public pir::Pass {
       : pir::Pass("substitute_dim_expr_based_on_constraints_pass", 1) {}
 
   void Run(pir::Operation* op) override {
-    pir::ModuleOp module_op = op->dyn_cast<pir::ModuleOp>();
-    SubstituteDimExprBasedOnConstraints(module_op);
+    SubstituteDimExprBasedOnConstraints(op);
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
