[CINN / Fusion] Support fusion of Pattern with multi downstream

paddle/cinn/hlir/dialect/operator/ir/attribute_storage.h

@@ -20,6 +20,7 @@
 #include <vector>
 #include "paddle/cinn/hlir/framework/op.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
+#include "paddle/cinn/operator_fusion/fusion_tracker/tracker.h"
 #include "paddle/pir/include/core/attribute_base.h"
 #include "paddle/pir/include/core/operation.h"
 #include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
@@ -100,6 +101,26 @@ struct GroupInfoAttributeStorage : public pir::AttributeStorage {
   ParamKey data_;
 };
 
+struct FusionTrackerPtrAttributeStorage : public pir::AttributeStorage {
+  using ParamKey = cinn::fusion::FusionTrackerPtr;
+  explicit FusionTrackerPtrAttributeStorage(const ParamKey& key) : data_(key) {}
+
+  static FusionTrackerPtrAttributeStorage* Construct(const ParamKey& key) {
+    return new FusionTrackerPtrAttributeStorage(key);
+  }
+
+  static std::size_t HashValue(const ParamKey& key) {
+    return std::hash<ParamKey>()(key);
+  }
+
+  bool operator==(const ParamKey& key) const { return data_ == key; }
+
+  const ParamKey& GetAsKey() const { return data_; }
+
+ private:
+  ParamKey data_;
+};
+
 struct CINNKernelInfoAttributeStorage : public pir::AttributeStorage {
   using ParamKey = cinn::hlir::framework::pir::CINNKernelInfo;
   explicit CINNKernelInfoAttributeStorage(const ParamKey& key) : data_(key) {}

paddle/cinn/hlir/dialect/operator/ir/manual_op.cc

@@ -37,7 +37,8 @@ namespace dialect {
 using DenseTensorType = paddle::dialect::DenseTensorType;
 
 const char* GroupOp::attributes_name[GroupOp::attributes_num] = {"group_info"};
-const char* FusionOp::attributes_name[GroupOp::attributes_num] = {"group_info"};
+const char* FusionOp::attributes_name[FusionOp::attributes_num] = {
+    "group_info", "fusion_tracker"};
 const char* ConcatOp::attributes_name[ConcatOp::attributes_num] = {"axis"};
 const char* SplitOp::attributes_name[SplitOp::attributes_num] = {
     "num_or_sections", "axis"};
@@ -146,13 +147,17 @@ void FusionOp::Build(pir::Builder& builder,
 void FusionOp::Build(pir::Builder& builder,             // NOLINT
                      pir::OperationArgument& argument,  // NOLINT
                      const std::vector<pir::Type>& output_types,
-                     const cinn::dialect::GroupInfo& group_info) {
+                     const cinn::dialect::GroupInfo& group_info,
+                     const cinn::fusion::FusionTrackerPtr& tracker) {
   argument.AddRegion(nullptr);
   argument.output_types = output_types;
 
   argument.AddAttribute("group_info",
                         cinn::dialect::GroupInfoAttribute::get(
                             pir::IrContext::Instance(), group_info));
+  argument.AddAttribute("fusion_tracker",
+                        cinn::dialect::FusionTrackerPtrAttribute::get(
+                            pir::IrContext::Instance(), tracker));
 }
 
 pir::Block* FusionOp::block() {

paddle/cinn/hlir/dialect/operator/ir/manual_op.h

@@ -67,7 +67,7 @@ class IR_API FusionOp
  public:
   using Op::Op;
   static const char *name() { return "cinn_op.fusion"; }
-  static constexpr uint32_t attributes_num = 1;
+  static constexpr uint32_t attributes_num = 2;
   static const char *attributes_name[attributes_num];
   static void Build(pir::Builder &builder,             // NOLINT
                     pir::OperationArgument &argument,  // NOLINT
@@ -76,7 +76,8 @@ class IR_API FusionOp
   static void Build(pir::Builder &builder,             // NOLINT
                     pir::OperationArgument &argument,  // NOLINT
                     const std::vector<pir::Type> &output_types,
-                    const cinn::dialect::GroupInfo &group_info);
+                    const cinn::dialect::GroupInfo &group_info,
+                    const cinn::fusion::FusionTrackerPtr &tracker);
 
   pir::Block *block();
   pir::Block *block() const;

paddle/cinn/hlir/dialect/operator/ir/op_attribute.cc

@@ -16,6 +16,11 @@
 
 namespace cinn {
 namespace dialect {
+
+const cinn::fusion::FusionTrackerPtr &FusionTrackerPtrAttribute::data() const {
+  return storage()->GetAsKey();
+}
+
 const GroupInfo &GroupInfoAttribute::data() const {
   return storage()->GetAsKey();
 }
@@ -29,3 +34,4 @@ CINNKernelInfoAttribute::data() const {
 
 IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::GroupInfoAttribute)
 IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::CINNKernelInfoAttribute)
+IR_DEFINE_EXPLICIT_TYPE_ID(cinn::dialect::FusionTrackerPtrAttribute)

paddle/cinn/hlir/dialect/operator/ir/op_attribute.h

@@ -14,10 +14,26 @@
 
 #pragma once
 #include "paddle/cinn/hlir/dialect/operator/ir/attribute_storage.h"
+#include "paddle/cinn/operator_fusion/fusion_tracker/tracker.h"
 #include "paddle/pir/include/core/attribute_base.h"
 
 namespace cinn {
 namespace dialect {
+class FusionTrackerPtrAttribute : public pir::Attribute {
+ public:
+  using Attribute::Attribute;
+
+  DECLARE_ATTRIBUTE_UTILITY_FUNCTOR(FusionTrackerPtrAttribute,
+                                    FusionTrackerPtrAttributeStorage);
+
+  bool operator<(const FusionTrackerPtrAttribute& right) const {
+    return storage() < right.storage();
+  }
+
+  static std::string name() { return "fusion_tracker"; }
+
+  const cinn::fusion::FusionTrackerPtr& data() const;
+};
 
 class GroupInfoAttribute : public pir::Attribute {
  public:
@@ -56,3 +72,4 @@ class CINNKernelInfoAttribute : public pir::Attribute {
 
 IR_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::GroupInfoAttribute)
 IR_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::CINNKernelInfoAttribute)
+IR_DECLARE_EXPLICIT_TYPE_ID(cinn::dialect::FusionTrackerPtrAttribute)

paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc

@@ -61,6 +61,7 @@ void OperatorDialect::initialize() {
   RegisterOp<GenerateXShapeOp>();
   RegisterAttribute<GroupInfoAttribute>();
   RegisterAttribute<CINNKernelInfoAttribute>();
+  RegisterAttribute<FusionTrackerPtrAttribute>();
 }
 
 void OperatorDialect::PrintType(pir::Type type, std::ostream &os) const {}
@@ -81,6 +82,10 @@ void OperatorDialect::PrintAttribute(pir::Attribute attr,
 
     os << "(" << cinn_kernel_info.data().fn_ptr;
     os << ')';
+  } else if (attr.isa<FusionTrackerPtrAttribute>()) {
+    auto tracker = attr.dyn_cast<FusionTrackerPtrAttribute>();
+    os << "(" << tracker;
+    os << ')';
   } else {
     PADDLE_THROW(::common::errors::Unimplemented(
         "cinn dialect only support GroupInfo and CINNKernelInfo"));

paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.cc

@@ -28,6 +28,7 @@
 #include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/accuracy_check_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/add_broadcast_to_elementwise_pass.h"
+#include "paddle/cinn/hlir/dialect/operator/transforms/add_store_in_fusion_op_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/add_store_in_group_op_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/conv2d_transpose_filter_pass.h"

paddle/cinn/hlir/dialect/operator/transforms/add_store_in_fusion_op_pass.cc

@@ -33,10 +33,6 @@ class AddYieldStoreInFusionOpPattern
   bool MatchAndRewrite(::pir::YieldOp op,
                        pir::PatternRewriter& rewriter) const override {
     for (auto i = 0; i < op->num_operands(); ++i) {
-      if (op->operand_source(i).use_count() == 1) {
-        continue;
-      }
-
       rewriter.SetInsertionPointAfter(op->operand_source(i).defining_op());
       auto store_op = rewriter.Build<cinn::dialect::YieldStoreOp>(
           op->operand_source(i), op->operand_source(i).type());
@@ -69,13 +65,6 @@ class AddStoreInFusionOpPass : public pir::Pass {
       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);

paddle/cinn/hlir/dialect/operator/transforms/attr_adt_type_id.h

@@ -56,6 +56,7 @@ namespace cinn::dialect {
 
 class GroupInfoAttribute;
 class CINNKernelInfoAttribute;
+class FusionTrackerPtrAttribute;
 
 }  // namespace cinn::dialect
 
@@ -86,7 +87,8 @@ class UnclassifiedAttribute {};
   __macro(paddle::dialect::PlaceAttribute)          \
   __macro(paddle::dialect::DataLayoutAttribute)     \
   __macro(cinn::dialect::GroupInfoAttribute)        \
-  __macro(cinn::dialect::CINNKernelInfoAttribute)
+  __macro(cinn::dialect::CINNKernelInfoAttribute)   \
+  __macro(cinn::dialect::FusionTrackerPtrAttribute)
 // clang-format on
 
 using AttrAdtTypeIdBase = ::common::AdtBaseTypeId<

paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc

@@ -33,7 +33,8 @@
 #include "paddle/cinn/hlir/dialect/operator/ir/manual_op.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h"
 #include "paddle/cinn/hlir/framework/pir/utils.h"
-#include "paddle/cinn/operator_fusion/group_cluster.h"
+#include "paddle/cinn/operator_fusion/cluster_interface.h"
+#include "paddle/cinn/operator_fusion/fusion_tracker/tracker.h"
 #include "paddle/common/ddim.h"
 #include "paddle/common/flags.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
@@ -109,6 +110,8 @@ struct GroupClusterNode {
   std::unordered_set<::pir::Value> GetOutsideInput() const {
     return GetListOutsideInput(ops);
   }
+
+  cinn::fusion::FusionTrackerPtr tracker;
 };
 
 std::vector<::pir::Value> GenerateOutputValue(
@@ -199,7 +202,7 @@ ::pir::GroupOpsVec CloneOps(
   return vec_new_op_list;
 }
 
-::pir::Operation* ReplaceWithGroupOp(
+::pir::Operation* ReplaceWithFusionOp(
     pir::PatternRewriter* rewriter,
     const ::pir::GroupOpsVec& group_ops,
     const GroupClusterNode& node,
@@ -218,8 +221,8 @@ ::pir::Operation* ReplaceWithGroupOp(
   // step 2: Replace the old op with GroupOp.
 
   auto output_types = BuildOutType(output_value);
-  auto new_fusion_op =
-      rewriter->Build<cinn::dialect::FusionOp>(output_types, group_info);
+  auto new_fusion_op = rewriter->Build<cinn::dialect::FusionOp>(
+      output_types, group_info, node.tracker);
   pir::Block* fusion_block = new_fusion_op.block();
 
   for (auto op : vec_new_op_list) {
@@ -245,32 +248,42 @@ ::pir::Operation* ReplaceWithGroupOp(
 }
 
 std::vector<GroupClusterNode> GroupSplit(cinn::dialect::GroupOp group_op) {
-  std::function<cinn::fusion::FrontendContent(pir::Operation*)> func =
-      [](pir::Operation* op) { return cinn::fusion::FrontendContent(op); };
+  std::function<cinn::fusion::PatternContent(pir::Operation*)> func =
+      [](pir::Operation* op) { return cinn::fusion::PatternContent(op); };
   const auto& contents = cinn::fusion::MapVector(group_op.GetOperators(), func);
   auto cluster_result = cinn::fusion::ClusterOps(contents, {});
-  std::vector<std::vector<pir::Operation*>> result;
-  std::transform(
-      cluster_result.begin(),
-      cluster_result.end(),
-      std::back_inserter(result),
-      [](const cinn::fusion::PatternNodePtr<cinn::fusion::FrontendStage> node) {
-        return cinn::fusion::GetOpsInPattern(node->stmt_pattern());
-      });
+  std::vector<std::vector<pir::Operation*>> op_sets;
+  std::vector<cinn::fusion::FusionTrackerPtr> trackers;
+  std::transform(cluster_result.begin(),
+                 cluster_result.end(),
+                 std::back_inserter(op_sets),
+                 [](const cinn::fusion::PatternNodePtr node) {
+                   return cinn::fusion::GetOpsInPattern(node->stmt_pattern());
+                 });
+  std::transform(cluster_result.begin(),
+                 cluster_result.end(),
+                 std::back_inserter(trackers),
+                 [](const cinn::fusion::PatternNodePtr node)
+                     -> cinn::fusion::FusionTrackerPtr {
+                   return cinn::fusion::GetFusionTracker(node->stmt_pattern());
+                 });
 
   // Each stmts corresponds to each fusion op(cluster node).
   // Concat all the ops of patterns in the stmts, and make them the op list of
   // cluster node.
   VLOG(4) << "Start Creating Cluster Nodes!";
   std::vector<GroupClusterNode> output_cluster_nodes;
-  for (const auto& op_set : result) {
+  for (int i = 0; i < op_sets.size(); i++) {
+    auto op_set = op_sets[i];
     GroupClusterNode cluster_node;
     for (const auto* op : op_set) {
       cluster_node.ops.push_back(const_cast<pir::Operation*>(op));
       auto op_kind = cinn::hlir::framework::pir::CompatibleInfo::OpKind(*op);
       cluster_node.group_kind =
           cluster_node.group_kind > op_kind ? cluster_node.group_kind : op_kind;
     }
+    // Deep copy trackers to avoid shared tracker conflict in different node
+    cluster_node.tracker = trackers[i]->Clone();
     output_cluster_nodes.push_back(cluster_node);
   }
   VLOG(4) << "Finished Creating Cluster Nodes!";
@@ -320,6 +333,22 @@ std::unordered_map<::pir::Value, size_t> BuildValueOrderByYieldOp(
   return all_output_values;
 }
 
+void UpdateTracker(std::vector<pir::Operation*> uniq_ops,
+                   fusion::FusionTrackerPtr tracker) {
+  std::map<pir::Operation*, int> op2idx;
+  for (int i = 0; i < uniq_ops.size(); ++i) {
+    op2idx[uniq_ops[i]] = i;
+  }
+  for (const auto& t : tracker->instructions_) {
+    if (t->type() == fusion::T_InitPattern) {
+      auto init_instr =
+          cinn::fusion::dynamic_cast_instr_with_err<fusion::InitPatternInstr>(
+              t);
+      init_instr->set_idx(op2idx[init_instr->op_]);
+    }
+  }
+}
+
 }  // namespace
 
 class CinnGroupClusterPattern
@@ -349,7 +378,9 @@ class CinnGroupClusterPattern
       VLOG(4) << "cluster node output size: " << output_values.size();
       auto uniq_ops = SortByOriginalOrderAndUniq(group_op, node.ops);
 
-      auto new_group_op = ReplaceWithGroupOp(
+      UpdateTracker(uniq_ops, node.tracker);
+
+      auto new_group_op = ReplaceWithFusionOp(
           &rewriter, uniq_ops, node, output_values, &ir_mapping);
 
       // TODO(Hongqing-work): delete this after fix bug of

paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.cc

@@ -73,6 +73,7 @@ void SetBroadcastLeafGroup(
           "ops:%d.",
           origin_group->ops().size(),
           new_group->ops().size()));
+
   UpdateGroupShapeExprs(
       new_group, origin_group, value_dim_exprs_list, value_to_dim_expr_idx);
 
@@ -85,6 +86,7 @@ void SetBroadcastLeafGroup(
                                              new_group->GetShapeOrDataExprs(v));
     }
   }
+
   group_list->emplace_back(new_group);
 }
 

paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/utils.cc

@@ -15,6 +15,7 @@
 #include "paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/utils.h"
 
 #include "paddle/cinn/adt/generate_map_expr.h"
+#include "paddle/cinn/hlir/dialect/operator/ir/attribute_storage.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/generate_shape_util.h"
 #include "paddle/cinn/hlir/dialect/operator/ir/op_attribute.h"
 #include "paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/broadcast_with_cf.h"
@@ -107,6 +108,7 @@ OpLoweringGroupPtr BuildOpLoweringGroup(pir::Operation* fusion_op_ptr) {
                           : group_op_kind;
     }
   }
+
   PADDLE_ENFORCE_GT(fusion_op.attributes().count("group_info"),
                     0UL,
                     ::common::errors::InvalidArgument(
@@ -117,7 +119,12 @@ OpLoweringGroupPtr BuildOpLoweringGroup(pir::Operation* fusion_op_ptr) {
                         .data();
 
   const auto& fn_name = attr.fn_name;
-  auto group = std::make_shared<OpLoweringGroup>(ops, fn_name);
+  auto group = std::make_shared<OpLoweringGroup>(
+      ops,
+      fn_name,
+      fusion_op_ptr->attribute("fusion_tracker")
+          .dyn_cast<cinn::dialect::FusionTrackerPtrAttribute>()
+          .data());
 
   group_op_kind =
       static_cast<int>(attr.op_pattern_kind) > static_cast<int>(group_op_kind)

paddle/cinn/hlir/dialect/operator/transforms/pir_to_py_code_converter.cc

@@ -1085,6 +1085,12 @@ struct PirToPyCodeConverterHelper {
       ss << "self." << name << "()";
       return ss.str();
     }
+    std::string operator()(TypeId<cinn::dialect::FusionTrackerPtrAttribute>) {
+      const auto& name = cinn::dialect::FusionTrackerPtrAttribute::name();
+      std::stringstream ss;
+      ss << "self." << name << "()";
+      return ss.str();
+    }
     std::string operator()(TypeId<UnclassifiedAttribute>) {
       return "self.UnclassifiedAttribute()";
     }

paddle/cinn/hlir/framework/pir/op_lowering_group.cc

@@ -133,7 +133,8 @@ std::shared_ptr<OpLoweringGroup> OpLoweringGroup::Clone(
     const std::string name_suffix) const {
   const auto new_fn_name = this->fn_name_ + "_cloned_" + name_suffix;
   // Construct Base information for new Group
-  auto new_group = std::make_shared<OpLoweringGroup>(this->ops_, new_fn_name);
+  auto new_group = std::make_shared<OpLoweringGroup>(
+      this->ops_, new_fn_name, this->fusion_tracker_ptr);
 
   new_group->output_ops_ = this->output_ops_;
   new_group->output_values_ = this->output_values_;