PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.cc
Copy file name to clipboard
+35-1Lines changed: 35 additions & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.cc
Copy file name to clipboard
+35-1Lines changed: 35 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
Copy file name to clipboard
+37-5Lines changed: 37 additions & 5 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
Copy file name to clipboard
+37-5Lines changed: 37 additions & 5 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/dynamic_reshape_pass.cc
Copy file name to clipboard
+5-26Lines changed: 5 additions & 26 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/dynamic_reshape_pass.cc
Copy file name to clipboard
+5-26Lines changed: 5 additions & 26 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/move_generate_shape_ops_to_prologue_pass.cc
Copy file name to clipboard
+20-10Lines changed: 20 additions & 10 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/move_generate_shape_ops_to_prologue_pass.cc
Copy file name to clipboard
+20-10Lines changed: 20 additions & 10 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/insert_broadcast_pass.cc
Copy file name to clipboard
+6-5Lines changed: 6 additions & 5 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/insert_broadcast_pass.cc
Copy file name to clipboard
+6-5Lines changed: 6 additions & 5 deletions
@@ -24,6 +24,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/ir_tensor.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
+#include "paddle/fluid/pir/transforms/shape_optimization_pass.h"
 #include "paddle/pir/include/core/builtin_type.h"
 #include "paddle/pir/include/core/op_base.h"
 #include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
@@ -104,6 +105,20 @@ void GroupOp::Print(pir::IrPrinter& printer) {
   os << " \n }";
 }
 
+bool GroupOp::InferSymbolicShape(
+    ::pir::ShapeConstraintIRAnalysis* shape_analysis) {
+  ::pir::InferSymExprForBlock(*block(), shape_analysis);
+
+  for (uint32_t rst_idx = 0; rst_idx < num_results(); rst_idx++) {
+    auto inner_yield_value = block()->back().operand_source(rst_idx);
+    const auto& shape =
+        shape_analysis->GetShapeOrDataForValue(inner_yield_value);
+    shape_analysis->SetShapeOrDataForValue(result(rst_idx), shape);
+  }
+
+  return true;
+}
+
 void FusionOp::Build(pir::Builder& builder,
                      pir::OperationArgument& argument,
                      const std::vector<pir::Type>& output_types) {
 
@@ -29,7 +29,8 @@
 namespace cinn {
 namespace dialect {
 
-class IR_API GroupOp : public pir::Op<GroupOp> {
+class IR_API GroupOp
+    : public pir::Op<GroupOp, paddle::dialect::InferSymbolicShapeInterface> {
  public:
   using Op::Op;
   static const char *name() { return "cinn_op.group"; }
@@ -51,6 +52,8 @@ class IR_API GroupOp : public pir::Op<GroupOp> {
   pir::Block *block();
   std::vector<pir::Operation *> GetOperators();
 
+  bool InferSymbolicShape(pir::ShapeConstraintIRAnalysis *shape_analysis);
+
   void VerifySig();
   void Print(pir::IrPrinter &printer);  // NOLINT
 };
 
@@ -15,6 +15,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h"
 
 #include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
+#include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/common/ddim.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
@@ -173,6 +174,23 @@ class AddBroadcastToElementwisePattern : public pir::OpRewritePattern<OPTYPE> {
   }
 };
 
+class DeleteUselessBroadcastPattern
+    : public pir::OpRewritePattern<cinn::dialect::BroadcastOp> {
+ public:
+  using pir::OpRewritePattern<cinn::dialect::BroadcastOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(cinn::dialect::BroadcastOp broadcast,
+                       pir::PatternRewriter& rewriter) const override {
+    if (!broadcast->GetParentOp()->isa<cinn::dialect::FusionOp>()) {
+      rewriter.ReplaceAllUsesWith(broadcast.result(0),
+                                  broadcast->operand_source(0));
+      rewriter.EraseOp(broadcast);
+      return true;
+    }
+    return false;
+  }
+};
+
 class AddBroadcastToElementwisePass : public pir::PatternRewritePass {
  public:
   AddBroadcastToElementwisePass()
@@ -224,14 +242,30 @@ class AddBroadcastToElementwisePass : public pir::PatternRewritePass {
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
-    return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
+    return op->num_regions() > 0;
+  }
+};
+
+class DeleteUselessBroadcastPass : public pir::PatternRewritePass {
+ public:
+  DeleteUselessBroadcastPass()
+      : pir::PatternRewritePass("delete_useless_broadcast_pass", 1) {}
+
+  pir::RewritePatternSet InitializePatterns(pir::IrContext* context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add<DeleteUselessBroadcastPattern>(context);
+    return ps;
   }
 };
 
 std::unique_ptr<pir::Pass> CreateAddBroadcastToElementwisePass() {
   return std::make_unique<AddBroadcastToElementwisePass>();
 }
 
+std::unique_ptr<pir::Pass> CreateDeleteUselessBroadcastPass() {
+  return std::make_unique<DeleteUselessBroadcastPass>();
+}
+
 }  // namespace ir
 }  // namespace dialect
 }  // namespace cinn
@@ -23,6 +23,8 @@ namespace ir {
 
 std::unique_ptr<pir::Pass> CreateAddBroadcastToElementwisePass();
 
+std::unique_ptr<pir::Pass> CreateDeleteUselessBroadcastPass();
+
 }  // namespace ir
 }  // namespace dialect
 }  // namespace cinn
@@ -94,27 +94,56 @@ void ApplyCinnPreprocessPass(
     pass_manager->AddPass(cinn::dialect::ir::CreateSimplifyDimExprPass());
     pass_manager->AddPass(
         cinn::dialect::ir::CreateSubstituteDimExprBasedOnConstraintsPass());
-    pass_manager->AddPass(cinn::dialect::ir::CreateInsertBroadcastPass());
     pass_manager->AddPass(cinn::dialect::ir::CreateConvert0DTo1DPass());
-    pass_manager->AddPass(pir::CreateShapeOptimizationPass());
     pass_manager->AddPass(
         cinn::dialect::ir::CreateFuseShapeOpsIntoGenerateShapeOpPass());
     pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
     pass_manager->AddPass(pir::CreateShapeOptimizationPass());
   }
 
+  pass_manager->Run(program);
+}
+
+void ApplyBuildGroupOpPass(
+    ::pir::Program* program,
+    const std::function<std::shared_ptr<pir::PassManager>()>&
+        CreatePassManager) {
+  std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
   pass_manager->AddPass(pir::CreateBuildCinnPass());
+  if (HasDynamicShape(*program)) {
+    pass_manager->AddPass(cinn::dialect::ir::CreateInsertBroadcastPass());
+  }
+  pass_manager->Run(program);
+}
+
+void ApplyGroupOpPass(::pir::Program* program,
+                      const std::function<std::shared_ptr<pir::PassManager>()>&
+                          CreatePassManager) {
+  std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
+  if (HasDynamicShape(*program)) {
+    pass_manager->AddPass(::pir::CreateShapeOptimizationPass());
+    pass_manager->AddPass(
+        cinn::dialect::ir::CreateFuseShapeOpsIntoGenerateShapeOpPass());
+    pass_manager->AddPass(
+        cinn::dialect::ir::CreateMoveGenerateShapeOpsToProloguePass());
+  }
 
-  pass_manager->AddPass(
-      cinn::dialect::ir::CreateMoveGenerateShapeOpsToProloguePass());
   pass_manager->AddPass(cinn::dialect::ir::CreateDynamicReshapeOpPass());
   pass_manager->AddPass(cinn::dialect::ir::CreateReplaceDynamicExpandOpPass());
-  pass_manager->AddPass(cinn::dialect::ir::CreateDivideGroupOpToFusionOpPass());
   pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
 
   pass_manager->Run(program);
 }
 
+void ApplyDivideGroupOpToFusionOpPass(
+    ::pir::Program* program,
+    const std::function<std::shared_ptr<pir::PassManager>()>&
+        CreatePassManager) {
+  std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
+  pass_manager->AddPass(cinn::dialect::ir::CreateDivideGroupOpToFusionOpPass());
+  pass_manager->Run(program);
+}
+
 void ApplyCinnLowerPass(
     ::pir::Program* program,
     const std::function<std::shared_ptr<pir::PassManager>()>&
@@ -148,6 +177,9 @@ void ApplyCinnPass(::pir::Program* program,
                    const std::function<std::shared_ptr<pir::PassManager>()>&
                        CreatePassManager) {
   ApplyCinnPreprocessPass(program, CreatePassManager);
+  ApplyBuildGroupOpPass(program, CreatePassManager);
+  ApplyGroupOpPass(program, CreatePassManager);
+  ApplyDivideGroupOpToFusionOpPass(program, CreatePassManager);
   ApplyCinnLowerPass(program, CreatePassManager);
 }
 
 
@@ -109,43 +109,22 @@ class DynamicUnsqueezeOpPattern
   }
 };
 
-class DynamicReshapeOpPass : public pir::Pass {
+class DynamicReshapeOpPass : public pir::PatternRewritePass {
  public:
   DynamicReshapeOpPass()
-      : pir::Pass("cinn_dynamic_reshape_op_pass", /*opt_level=*/1) {}
+      : pir::PatternRewritePass("cinn_dynamic_reshape_op_pass", 1) {}
 
-  bool Initialize(pir::IrContext* context) override {
+  pir::RewritePatternSet InitializePatterns(pir::IrContext* context) override {
     pir::RewritePatternSet ps(context);
     ps.Add<DynamicReshapeOpPattern>(context);
     ps.Add<DynamicSqueezeOpPattern>(context);
     ps.Add<DynamicUnsqueezeOpPattern>(context);
-    patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
-    return true;
-  }
-
-  void Run(pir::Operation* op) override {
-    pir::GreedyRewriteConfig cfg;
-    cfg.use_top_down_traversal = true;
-    cfg.max_iterations = 10;
-    for (uint32_t i = 0; i < op->num_regions(); ++i) {
-      for (auto& block : op->region(i)) {
-        for (auto& op : block) {
-          if (op.isa<cinn::dialect::GroupOp>()) {
-            auto [_, num_rewrites] =
-                pir::ApplyPatternsGreedily(&op, patterns_, cfg);
-            AddStatistics(num_rewrites);
-          }
-        }
-      }
-    }
+    return ps;
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
-    return op->num_regions() > 0;
+    return op->isa<cinn::dialect::GroupOp>() && op->num_regions() > 0;
   }
-
- private:
-  pir::FrozenRewritePatternSet patterns_;
 };
 
 std::unique_ptr<pir::Pass> CreateDynamicReshapeOpPass() {
 
@@ -206,7 +206,7 @@ class FuseShapeOpsIntoGenerateShapeOpPass : public pir::PatternRewritePass {
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
-    return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
+    return op->num_regions() > 0;
   }
 };
 
 
@@ -67,22 +67,32 @@ class GroupOpGenerateShapeOpsPattern
   }
 };
 
-class MoveGenerateShapeOpsToProloguePass : public pir::PatternRewritePass {
+class MoveGenerateShapeOpsToProloguePass : public pir::Pass {
  public:
   MoveGenerateShapeOpsToProloguePass()
-      : pir::PatternRewritePass("move_generate_shape_ops_to_prologue", 1) {}
+      : pir::Pass("move_generate_shape_ops_to_prologue", /*opt_level=*/1) {}
 
-  pir::RewritePatternSet InitializePatterns(pir::IrContext* context) override {
-    pir::RewritePatternSet ps(context);
-    ps.Add<GroupOpGenerateShapeOpsPattern>(context);
-    return ps;
+  void Run(pir::Operation* op) override {
+    auto group_op = op->dyn_cast<cinn::dialect::GroupOp>();
+    CHECK(group_op);
+    ::pir::IrContext* ctx = ::pir::IrContext::Instance();
+    pir::ShapeConstraintIRAnalysis& shape_analysis =
+        pir::ShapeAnalysisManager::Instance().Get(group_op->GetParentProgram());
+    ShapeOrDataDimExprsAccessor dim_exprs_accessor{
+        .GetShapeOrDataDimExprs =
+            [&](pir::Value value) -> const symbol::ShapeOrDataDimExprs& {
+          return shape_analysis.GetShapeOrDataForValue(value);
+        },
+        .SetShapeOrDataDimExprs =
+            [&](pir::Value value,
+                const symbol::ShapeOrDataDimExprs& dim_exprs) {
+              shape_analysis.SetShapeOrDataForValue(value, dim_exprs);
+            }};
+    MoveGenerateShapeOpsToPrologue(ctx, group_op.block(), dim_exprs_accessor);
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
-    if (!(op->isa<pir::ModuleOp>() && op->num_regions() > 0)) return false;
-    auto* program = op->GetParentProgram();
-    VLOG(4) << "Before MoveGenerateShapeOpsToProloguePass: " << *program;
-    return true;
+    return op->isa<cinn::dialect::GroupOp>() && op->num_regions() > 0;
   }
 };
 
 
@@ -246,6 +246,11 @@ inline bool horizontal_or_vertical_reduce_relation(
   // check producer has same shape with reducer op.
   auto reduce_shape = ::common::vectorize(GetFirstInputShape(reducer));
   auto reduce_axes = GetVectorAttr(reducer, "dim");
+  if (reduce_axes.empty()) {
+    for (size_t i = 0; i < reduce_shape.size(); ++i) {
+      reduce_axes.push_back(i);
+    }
+  }
 
   for (auto& axis : reduce_axes) {
     // if axis = -1, set as shape.size() - 1
 
@@ -15,6 +15,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/insert_broadcast_pass.h"
 
 #include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
+#include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/common/ddim.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
@@ -51,26 +52,26 @@ bool ProcessOp(pir::Operation* op, pir::PatternRewriter* rewriter) {
   const auto& y_shape = shape_analysis.GetShapeOrDataForValue(y);
   const auto& out_shape = shape_analysis.GetShapeOrDataForValue(op->result(0));
 
-  bool has_insert_broadcast = false;
+  if (x_shape == y_shape) {
+    return false;
+  }
 
   pir::Value output_dim_tensor = GetOutputDimTensor(rewriter, x, y);
   if (x_shape.shape() != out_shape.shape() ||
       x_shape.data() != out_shape.data()) {
-    has_insert_broadcast = true;
     pir::Value broadcasted_x =
         rewriter->Build<paddle::dialect::ExpandOp>(x, output_dim_tensor).out();
     op->operand(0).set_source(broadcasted_x);
     shape_analysis.SetShapeOrDataForValue(broadcasted_x, out_shape);
   }
   if (y_shape.shape() != out_shape.shape() ||
       y_shape.data() != out_shape.data()) {
-    has_insert_broadcast = true;
     pir::Value broadcasted_y =
         rewriter->Build<paddle::dialect::ExpandOp>(y, output_dim_tensor).out();
     op->operand(1).set_source(broadcasted_y);
     shape_analysis.SetShapeOrDataForValue(broadcasted_y, out_shape);
   }
-  return has_insert_broadcast;
+  return true;
 }
 
 }  // namespace
@@ -120,7 +121,7 @@ class InsertBroadcastPass : public pir::PatternRewritePass {
   }
 
   bool CanApplyOn(pir::Operation* op) const override {
-    return op->isa<pir::ModuleOp>() && op->num_regions() > 0;
+    return op->isa<cinn::dialect::GroupOp>() && op->num_regions() > 0;
   }
 };
Original file line number	Diff line number	Diff line change
`@@ -206,7 +206,7 @@ class FuseShapeOpsIntoGenerateShapeOpPass : public pir::PatternRewritePass {`
`206`	`206`	`}`
`207`	`207`
`208`	`208`	`bool CanApplyOn(pir::Operation* op) const override {`
`209`		`- return op->isa<pir::ModuleOp>() && op->num_regions() > 0;`
	`209`	`+ return op->num_regions() > 0;`
`210`	`210`	`}`
`211`	`211`	`};`
`212`	`212`