fix bug of substitute dim expr for group (#62941)

zyfncg · web-flow · commit 41dc10408772 · 2024-03-22T15:12:06.000+08:00
diff --git a/paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.cc b/paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.cc
@@ -575,7 +575,7 @@ std::vector<pir::Value> GetMinimalInputs(
       [&](pir::Value input_tensor,
           const std::vector<symbol::DimExpr>& dim_exprs) {
         for (const auto& dim_expr : dim_exprs) {
-          if (dim_expr.isa<int64_t>()) continue;
+          if (!dim_expr.isa<std::string>()) continue;
           if (handled_dim_exprs.insert(dim_expr).second) {
             first_occurred_input_tensors.insert(input_tensor);
           }
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
@@ -116,7 +116,6 @@ void ApplyBuildGroupOpPass(
 
   pass_manager->AddPass(pir::CreateBuildCinnPass());
   if (has_dynamic_shape) {
-    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
     pass_manager->AddPass(cinn::dialect::ir::CreateInsertBroadcastPass());
   }
   pass_manager->Run(program);
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.cc
@@ -670,6 +670,7 @@ CollectSubstituteDimExprMap(
     const GroupPtr& group,
     pir::ShapeConstraintIRAnalysis& shape_analysis) {  // NOLINT
   std::unordered_map<symbol::DimExpr, symbol::DimExpr> dim_expr_map;
+  std::unordered_set<std::string> base_dim_expr_set;
 
   VisitEachInputValue(group, [&](::pir::Value value) {
     if (!shape_analysis.HasShapeOrDataForValue(value)) {
@@ -682,9 +683,33 @@ CollectSubstituteDimExprMap(
         dim_expr_map[dim_expr] =
             symbol::DimExpr(shape_analysis.GetNextSymName());
       }
+      if (dim_expr.isa<std::string>()) {
+        base_dim_expr_set.insert(dim_expr.Get<std::string>());
+      }
     });
   });
 
+  const std::unordered_set<symbol::DimExpr> dim_exprs_no_outer_symbol = [&] {
+    auto HasOuterBasicSymbol = [&](const symbol::DimExpr& dim_expr) {
+      for (const auto& symbol : symbol::CollectDimExprSymbols(dim_expr)) {
+        if (base_dim_expr_set.count(symbol) == 0) {
+          return true;
+        }
+      }
+      return false;
+    };
+    std::unordered_set<symbol::DimExpr> result;
+    for (const auto& kv : dim_expr_map) {
+      if (IsComplicatedDimExpr(kv.first) && !HasOuterBasicSymbol(kv.first)) {
+        result.insert(kv.first);
+      }
+    }
+    return result;
+  }();
+  for (const auto& dim_expr : dim_exprs_no_outer_symbol) {
+    dim_expr_map.erase(dim_expr);
+  }
+
   return dim_expr_map;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -575,7 +575,7 @@ std::vector<pir::Value> GetMinimalInputs(`
`575`	`575`	`[&](pir::Value input_tensor,`
`576`	`576`	`const std::vector<symbol::DimExpr>& dim_exprs) {`
`577`	`577`	`for (const auto& dim_expr : dim_exprs) {`
`578`		`- if (dim_expr.isa<int64_t>()) continue;`
	`578`	`+ if (!dim_expr.isa<std::string>()) continue;`
`579`	`579`	`if (handled_dim_exprs.insert(dim_expr).second) {`
`580`	`580`	`first_occurred_input_tensors.insert(input_tensor);`
`581`	`581`	`}`
Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,6 @@ void ApplyBuildGroupOpPass(`
`116`	`116`
`117`	`117`	`pass_manager->AddPass(pir::CreateBuildCinnPass());`
`118`	`118`	`if (has_dynamic_shape) {`
`119`		`- pass_manager->AddPass(pir::CreateShapeOptimizationPass());`
`120`	`119`	`pass_manager->AddPass(cinn::dialect::ir::CreateInsertBroadcastPass());`
`121`	`120`	`}`
`122`	`121`	`pass_manager->Run(program);`