PaddlePaddle
diff --git a/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/ast_gen_ius/ast_gen.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+17Lines changed: 17 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
Copy file name to clipboard
+9-1Lines changed: 9 additions & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc
Copy file name to clipboard
+9-1Lines changed: 9 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/add_store_in_fusion_op_pass.cc
Copy file name to clipboard
+122Lines changed: 122 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/add_store_in_fusion_op_pass.cc
Copy file name to clipboard
+122Lines changed: 122 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/add_store_in_fusion_op_pass.h
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/add_store_in_fusion_op_pass.h
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc
Copy file name to clipboard
+15-4Lines changed: 15 additions & 4 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc
Copy file name to clipboard
+15-4Lines changed: 15 additions & 4 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.cc
Copy file name to clipboard
+12-1Lines changed: 12 additions & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/lower_cinn_fusion_op_pass.cc
Copy file name to clipboard
+12-1Lines changed: 12 additions & 1 deletion
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
Copy file name to clipboard
+4-8Lines changed: 4 additions & 8 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
Copy file name to clipboard
+4-8Lines changed: 4 additions & 8 deletions
@@ -244,7 +244,7 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {
       if (FLAGS_group_schedule_tiling_first && is_keep_dim) {
         continue;
       }
-      if (!FLAGS_group_schedule_tiling_first && !FLAGS_cinn_bucket_compile &&
+      if ((!FLAGS_group_schedule_tiling_first || !FLAGS_cinn_bucket_compile) &&
           shape[i] == Expr(1)) {
         continue;
       }
 
@@ -170,6 +170,16 @@ void FusionOp::Print(pir::IrPrinter& printer) {
   os << " \n }";
 }
 
+void YieldStoreOp::Build(pir::Builder& builder,
+                         pir::OperationArgument& argument,
+                         pir::Value x,
+                         pir::Type output_type) {
+  argument.inputs = {x};
+  argument.output_types = {output_type};
+}
+
+void YieldStoreOp::VerifySig() {}
+
 bool ConcatOp::InferSymbolicShape(
     pir::ShapeConstraintIRAnalysis* shape_analysis) {
   VLOG(4) << "Infer symbolic shape for cinn_op.concat";
@@ -501,3 +511,4 @@ IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::FusionOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::ConcatOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::SplitOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::GenerateShapeOp);
+IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::YieldStoreOp);
@@ -82,6 +82,22 @@ class IR_API FusionOp : public pir::Op<FusionOp> {
   void Print(pir::IrPrinter &printer);  // NOLINT
 };
 
+// YieldStoreOp represents a store operation for
+// seperate local variable and ouptut
+class IR_API YieldStoreOp : public pir::Op<YieldStoreOp> {
+ public:
+  using Op::Op;
+  static const char *name() { return "cinn_op.yield_store"; }
+  static constexpr uint32_t attributes_num = 0;
+  static constexpr const char **attributes_name = nullptr;
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value x,
+                    pir::Type output_type);
+
+  void VerifySig();
+};
+
 class IR_API ConcatOp
     : public pir::Op<ConcatOp, paddle::dialect::InferSymbolicShapeInterface> {
  public:
@@ -170,3 +186,4 @@ IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::FusionOp)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::ConcatOp)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::SplitOp)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::GenerateShapeOp);
+IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::YieldStoreOp);
@@ -56,6 +56,7 @@ void OperatorDialect::initialize() {
   RegisterOp<FusionOp>();
   RegisterOp<ConcatOp>();
   RegisterOp<SplitOp>();
+  RegisterOp<YieldStoreOp>();
   RegisterOp<GenerateShapeOp>();
   RegisterAttribute<GroupInfoAttribute>();
   RegisterAttribute<CINNKernelInfoAttribute>();
 
@@ -25,6 +25,7 @@
 
 #include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/dynamic_reshape_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/check_infer_symbolic_pass.h"
@@ -47,6 +48,7 @@
 
 COMMON_DECLARE_bool(print_ir);
 COMMON_DECLARE_bool(check_infer_symbolic);
+PD_DECLARE_bool(group_schedule_tiling_first);
 
 namespace cinn::dialect::ir {
 
@@ -130,6 +132,7 @@ void ApplyGroupOpPass(::pir::Program* program,
 
   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);
@@ -140,7 +143,12 @@ void ApplyDivideGroupOpToFusionOpPass(
     const std::function<std::shared_ptr<pir::PassManager>()>&
         CreatePassManager) {
   std::shared_ptr<pir::PassManager> pass_manager = CreatePassManager();
-  pass_manager->AddPass(cinn::dialect::ir::CreateDivideGroupOpToFusionOpPass());
+  if (FLAGS_group_schedule_tiling_first) {
+    pass_manager->AddPass(cinn::dialect::ir::CreateCinnGroupClusterPass());
+  } else {
+    pass_manager->AddPass(
+        cinn::dialect::ir::CreateDivideGroupOpToFusionOpPass());
+  }
   pass_manager->Run(program);
 }
 
 
@@ -0,0 +1,122 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/cinn/hlir/dialect/operator/transforms/add_store_in_fusion_op_pass.h"
+
+#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
+#include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/pir/include/core/builtin_type_interfaces.h"
+#include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
+#include "paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h"
+
+namespace cinn {
+namespace dialect {
+namespace ir {
+
+class AddYieldStoreInFusionOpPattern
+    : public pir::OpRewritePattern<::pir::YieldOp> {
+ public:
+  using pir::OpRewritePattern<::pir::YieldOp>::OpRewritePattern;
+
+  bool MatchAndRewrite(::pir::YieldOp op,
+                       pir::PatternRewriter& rewriter) const override {
+    for (auto i = 0; i < op->num_operands(); ++i) {
+      if (op->operand_source(i)
+              .defining_op()
+              ->isa<cinn::dialect::ReshapeOp>()) {
+        auto pre_name = op->operand_source(i).defining_op()->name();
+
+        if (op->operand_source(i).use_count() > 1) {
+          continue;
+        }
+
+        if ((pre_name != "cinn_op.reduce_sum") &&
+            (pre_name != "cinn_op.reduce_max")) {
+          auto new_full = rewriter.Build<cinn::dialect::YieldStoreOp>(
+              op->operand_source(i).defining_op()->operand_source(0),
+              op->operand_source(i).type());
+
+          op->operand(i).set_source(new_full.result(0));
+
+          continue;
+        }
+      }
+
+      if (op->operand_source(i).use_count() == 1) {
+        continue;
+      }
+
+      auto new_full = rewriter.Build<cinn::dialect::YieldStoreOp>(
+          op->operand_source(i), op->operand_source(i).type());
+
+      op->operand(i).set_source(new_full.result(0));
+    }
+
+    return true;
+  }
+};
+
+class AddStoreInFusionOpPass : public pir::Pass {
+ public:
+  AddStoreInFusionOpPass()
+      : pir::Pass("add_store_in_fusion_op", /*opt_level=*/1) {}
+
+  bool Initialize(pir::IrContext* context) override {
+    pir::RewritePatternSet ps(context);
+    ps.Add<AddYieldStoreInFusionOpPattern>(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 = 1;
+    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::FusionOp>()) {
+            auto fusion_op = op.dyn_cast<cinn::dialect::FusionOp>();
+            if (fusion_op.GetOperators().size() == 2 &&
+                fusion_op.GetOperators()
+                    .front()
+                    ->isa<cinn::dialect::ReshapeOp>()) {
+              continue;
+            }
+            auto [_, num_rewrites] =
+                pir::ApplyPatternsGreedily(&op, patterns_, cfg);
+            AddStatistics(num_rewrites);
+          }
+        }
+      }
+    }
+  }
+
+  bool CanApplyOn(pir::Operation* op) const override {
+    return op->num_regions() > 0;
+  }
+
+ private:
+  pir::FrozenRewritePatternSet patterns_;
+};
+
+std::unique_ptr<pir::Pass> CreateAddStoreInFusionOpPass() {
+  return std::make_unique<AddStoreInFusionOpPass>();
+}
+
+}  // namespace ir
+}  // namespace dialect
+}  // namespace cinn
@@ -0,0 +1,28 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+#include "paddle/pir/include/pass/pass.h"
+
+namespace cinn {
+namespace dialect {
+namespace ir {
+
+std::unique_ptr<pir::Pass> CreateAddStoreInFusionOpPass();
+
+}  // namespace ir
+}  // namespace dialect
+}  // namespace cinn
@@ -287,10 +287,13 @@ ::pir::GroupOpsVec CloneOps(
     auto new_op = op->Clone(*ir_mapping, clone_options);
     auto& shape_analysis =
         pir::ShapeAnalysisManager::Instance().Get(op->GetParentProgram());
+
     for (size_t i = 0; i < op->num_results(); ++i) {
-      shape_analysis.SetShapeOrDataForValue(
-          new_op->result(i),
-          shape_analysis.GetShapeOrDataForValue(op->result(i)));
+      if (shape_analysis.HasShapeOrDataForValue(op->result(i))) {
+        shape_analysis.SetShapeOrDataForValue(
+            new_op->result(i),
+            shape_analysis.GetShapeOrDataForValue(op->result(i)));
+      }
     }
 
     vec_new_op_list.push_back(new_op);
@@ -398,7 +401,13 @@ bool CanFuse(const GroupClusterNode& first,
 
     if (first.loop_ranges != second.loop_ranges) {
       sch_node->type = hlir::framework::pir::ScheduleAlignType::kBroadcast;
-      sch_node->axis_info = first.reduce_axis;
+      for (auto& d : first.reduce_axis) {
+        if (d < 0) {
+          sch_node->axis_info.push_back(d + first.loop_ranges.size());
+        } else {
+          sch_node->axis_info.push_back(d);
+        }
+      }
       sch_node->factor_info = first.loop_ranges;
     }
     return true;
@@ -531,6 +540,8 @@ void GetClusterNodeBasicInfo(::pir::Operation* op,
     sch_node->axis_info =
         cinn::dialect::ir::GetVectorAttr(op, "broadcast_axes");
     sch_node->factor_info = cinn::dialect::ir::GetVectorAttr(op, "out_shape");
+  } else if (op->name() == "cinn_op.generate_shape") {
+    // do nothing for now
   } else {
     PADDLE_THROW(phi::errors::Unimplemented(
         "only support elementwise, broadcast, reduce type"));
 
@@ -690,11 +690,23 @@ class FusionOpPattern : public pir::OpRewritePattern<cinn::dialect::FusionOp> {
   std::shared_ptr<Group> RebuildGroup(cinn::dialect::FusionOp fusion_op) const {
     auto group = std::make_shared<Group>();
     group->op_pattern_kind = cinn::hlir::framework::OpPatternKind::kElementWise;
+    if (fusion_op.attributes().count("group_info")) {
+      auto attr = fusion_op.attribute("group_info")
+                      .dyn_cast<cinn::dialect::GroupInfoAttribute>()
+                      .data();
+
+      group->op_pattern_kind = attr.op_pattern_kind;
+      group->loop_ranges = attr.loop_ranges;
+
+      group->reduce_axis = attr.reduce_axis;
+      group->alignment_schedule_info = attr.alignment_schedule_info;
+    }
 
     // Rebuild ops of the group
     for (auto op : fusion_op.GetOperators()) {
       if (!op->isa<::pir::YieldOp>()) {
         group->ops.push_back(op);
+
         group->ops_set.insert(op);
         group->op_pattern_kind =
             static_cast<int>(CompatibleInfo::OpKind(*op)) >
@@ -709,7 +721,6 @@ class FusionOpPattern : public pir::OpRewritePattern<cinn::dialect::FusionOp> {
     for (size_t i = 0; i < yield_op->num_operands(); ++i) {
       auto in = yield_op->operand_source(i);
       group->output_values.push_back(in);
-
       group->output_ops.insert(in.defining_op());
     }
 
 
@@ -203,13 +203,15 @@ class ReshapeOpPattern
     auto scale_factor_gen_op = op->operand_source(1).defining_op();
     auto full_op =
         scale_factor_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>();
-    return flag && full_op;
+    auto not_combine_input =
+        op->result(0).use_count() == 1 &&
+        op->result(0).first_use().owner()->name() == "builtin.combine";
+    return flag && full_op && (!not_combine_input);
   }
 
   void Rewrite(paddle::dialect::ReshapeOp op,
                pir::PatternRewriter &rewriter) const override {
     auto scale_factor_gen_op = op->operand_source(1).defining_op();
-
     auto full_op =
         scale_factor_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>();
     // scale is generator by full op
@@ -725,16 +727,10 @@ pir::RewritePatternSet PdOpToCinnOpPass::InitializePatterns(
   ps.Add(paddle::drr::Create<MinOpPattern>(context));
   ps.Add(paddle::drr::Create<ProdOpPattern>(context));
   ps.Add<ReshapeOpPattern>(context);
-  ps.Add<Pool2dOpPattern>(context);
-  ps.Add<ConcatOpPattern>(context);
-  ps.Add<SliceOpPattern>(context);
   ps.Add<PowOpPattern>(context);
-  ps.Add<SplitWithNumOpPattern>(context);
   ps.Add<AddNOpPattern>(context);
-  ps.Add<SplitOpPattern>(context);
   ps.Add<ExpandOpPattern>(context);
   ps.Add<IsCloseOpPattern>(context);
-  // ps.Add(paddle::drr::Create<UniformOpPattern>(context));
 
   return ps;
 }
Original file line number	Diff line number	Diff line change
`@@ -244,7 +244,7 @@ ir::Expr AstGen::Build(const ir::Tensor& tensor, TensorGroup* tensor_group) {`
`244`	`244`	`if (FLAGS_group_schedule_tiling_first && is_keep_dim) {`
`245`	`245`	`continue;`
`246`	`246`	`}`
`247`		`- if (!FLAGS_group_schedule_tiling_first && !FLAGS_cinn_bucket_compile &&`
	`247`	`+ if ((!FLAGS_group_schedule_tiling_first \|\| !FLAGS_cinn_bucket_compile) &&`
`248`	`248`	`shape[i] == Expr(1)) {`
`249`	`249`	`continue;`
`250`	`250`	`}`