PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/operator_fusion/group_cluster.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/cinn/operator_fusion/group_cluster.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/cinn/operator_fusion/pattern_graph.cc
Copy file name to clipboard
+34-32Lines changed: 34 additions & 32 deletions b/‎paddle/cinn/operator_fusion/pattern_graph.cc
Copy file name to clipboard
+34-32Lines changed: 34 additions & 32 deletions
@@ -846,7 +846,7 @@ std::vector<GroupClusterNode> NewOpMergeWithOp(
       cluster_result.end(),
       std::back_inserter(result),
       [](const cinn::fusion::PatternNodePtr<cinn::fusion::FrontendStage> node) {
-        return cinn::fusion::GetOpsInPattern(node->stmt_pattern_);
+        return cinn::fusion::GetOpsInPattern(node->stmt_pattern());
       });
 
   // Each stmts corresponds to each fusion op(cluster node).
 
@@ -586,7 +586,7 @@ std::vector<ir::Expr> OperationFusion(
   CHECK(fusion_nodes.size() == 1)
       << "Only support one fusion node in backend now.";
 
-  const auto& output = GetExprFromPattern(fusion_nodes[0]->stmt_pattern_);
+  const auto& output = GetExprFromPattern(fusion_nodes[0]->stmt_pattern());
   VLOG(4) << "Fusion Result: output size is " << output.size();
   for (const auto& expr : output) {
     VLOG(4) << expr;
 
@@ -82,7 +82,7 @@ inline std::vector<fusion::PatternNodePtr<T>> ClusterOps(
   for (const auto& node : result) {
     VLOG(4) << "\n"
             << node->DebugStr() << "\n"
-            << fusion::StmtPatternDebugStr(node->stmt_pattern_);
+            << fusion::StmtPatternDebugStr(node->stmt_pattern());
   }
 
   return result;
 
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/cinn/operator_fusion/pattern_graph.h"
+#include <functional>
 #include "paddle/cinn/operator_fusion/backend/pattern.h"
 #include "paddle/cinn/operator_fusion/backend/pattern_fuser.h"
 #include "paddle/cinn/operator_fusion/frontend/pattern.h"
@@ -58,7 +59,7 @@ std::vector<PatternNodePtr<T>> PatternGraph<T>::SortByTopoOrder() {
   std::list<PatternNodePtr<T>> topo_queue;
   std::map<PatternNodePtr<T>, int> degree;
   for (const auto& node : all_pattern_nodes_) {
-    degree[node] = node->upstream_.size();
+    degree[node] = node->upstream().size();
     if (degree[node] == 0) {
       topo_queue.push_back(node);
     }
@@ -67,7 +68,7 @@ std::vector<PatternNodePtr<T>> PatternGraph<T>::SortByTopoOrder() {
     PatternNodePtr<T> node = topo_queue.front();
     topo_queue.pop_front();
     res.push_back(node);
-    for (const auto& downstream_op : node->downstream_) {
+    for (const auto& downstream_op : node->downstream()) {
       degree[downstream_op] = degree[downstream_op] - 1;
       if (degree[downstream_op] == 0) {
         topo_queue.push_back(downstream_op);
@@ -145,7 +146,6 @@ PatternGraph<T>::PatternGraph(const std::vector<PatternContent<T>>& contents,
     PatternNodePtr<T> node = std::make_shared<PatternNode<T>>(content);
     op_to_node_map[content.op] = node;
     all_pattern_nodes_.emplace(node);
-    node->sink_op_ = content.op;
   }
 
   for (const auto& content : contents) {
@@ -156,7 +156,7 @@ PatternGraph<T>::PatternGraph(const std::vector<PatternContent<T>>& contents,
       ::pir::Operation* input_op = content.op->operand_source(i).defining_op();
       if (op_to_node_map.find(input_op) != op_to_node_map.end()) {
         PatternNodePtr<T> upstream_node = op_to_node_map[input_op];
-        cur_node->upstream_.push_back(upstream_node);
+        cur_node->AddNodeToUpstream(upstream_node);
       }
     }
 
@@ -169,15 +169,15 @@ PatternGraph<T>::PatternGraph(const std::vector<PatternContent<T>>& contents,
         ::pir::Operation* output_op = consumer_it->owner();
         if (op_to_node_map.find(output_op) != op_to_node_map.end()) {
           PatternNodePtr<T> downstream_node = op_to_node_map[output_op];
-          cur_node->downstream_.push_back(downstream_node);
+          cur_node->AddNodeToDownstream(downstream_node);
         }
       }
     }
 
     // unique all upstream / downstream node.
     // c = a + a ; then add will have 2 same upstream.
-    cur_node->downstream_ = UniqueVectorBySet(cur_node->downstream_);
-    cur_node->upstream_ = UniqueVectorBySet(cur_node->upstream_);
+    cur_node->UniqueUpstream();
+    cur_node->UniqueDownstream();
   }
 
   VLOG(4) << "PatternGraph Created, pattern node size: "
@@ -192,12 +192,12 @@ void PatternGraph<T>::RemoveNode(const PatternNodePtr<T>& node) {
     all_pattern_nodes_.erase(node);
   }
 
-  for (PatternNodePtr<T>& upstream : node->upstream_) {
-    RemoveFromVector(&upstream->downstream_, node);
+  for (const PatternNodePtr<T>& upstream : node->upstream()) {
+    upstream->RemoveNodeFromDownstream(node);
   }
 
-  for (PatternNodePtr<T>& downstream : node->downstream_) {
-    RemoveFromVector(&downstream->upstream_, node);
+  for (const PatternNodePtr<T>& downstream : node->downstream()) {
+    downstream->RemoveNodeFromUpstream(node);
   }
 }
 
@@ -220,37 +220,39 @@ std::string PatternGraph<T>::GraphInfo() const {
 
 template <typename T>
 PatternNodePtr<T> PatternGraph<T>::MergeNode(
-    const PatternNodePtr<T>& upstream, const PatternNodePtr<T>& downstream) {
+    const PatternNodePtr<T>& upstream,
+    const PatternNodePtr<T>& downstream,
+    MergePatternFn<T> merge_pattern_fn) {
   PatternNodePtr<T> merged_node =
-      std::make_shared<PatternNode<T>>(upstream, downstream);
+      std::make_shared<PatternNode<T>>(upstream, downstream, merge_pattern_fn);
 
-  // deal with the reference.
-  ExtendVector(&merged_node->upstream_, upstream->upstream_);
-  ExtendVector(&merged_node->upstream_, downstream->upstream_);
-  RemoveFromVector(&merged_node->upstream_, upstream);
-
-  ExtendVector(&merged_node->downstream_, upstream->downstream_);
-  ExtendVector(&merged_node->downstream_, downstream->downstream_);
-  RemoveFromVector(&merged_node->downstream_, downstream);
-
-  for (const auto& upstream_node : merged_node->upstream_) {
-    upstream_node->downstream_.push_back(merged_node);
-    RemoveFromVector(&upstream_node->downstream_, upstream);
-    RemoveFromVector(&upstream_node->downstream_, downstream);
+  // Update upstream and downstream nodes.
+  for (const auto& upstream_node : merged_node->upstream()) {
+    upstream_node->AddNodeToDownstream(merged_node);
+    upstream_node->RemoveNodeFromDownstream(upstream);
+    upstream_node->RemoveNodeFromDownstream(downstream);
   }
-  for (const auto& downstream_node : merged_node->downstream_) {
-    downstream_node->upstream_.push_back(merged_node);
-    RemoveFromVector(&downstream_node->downstream_, upstream);
-    RemoveFromVector(&downstream_node->downstream_, downstream);
+  for (const auto& downstream_node : merged_node->downstream()) {
+    downstream_node->AddNodeToUpstream(merged_node);
+    downstream_node->RemoveNodeFromDownstream(upstream);
+    downstream_node->RemoveNodeFromDownstream(downstream);
   }
 
   const auto vec_unique = [](const std::vector<PatternNodePtr<T>>& vec) {
     auto set = std::unordered_set(vec.begin(), vec.end());
     return set.size() == vec.size();
   };
 
-  CHECK(vec_unique(merged_node->upstream_));
-  CHECK(vec_unique(merged_node->downstream_));
+  PADDLE_ENFORCE_EQ(
+      vec_unique(merged_node->upstream()),
+      true,
+      phi::errors::PreconditionNotMet(
+          "The upstream nodes of the merged node are not unique."));
+  PADDLE_ENFORCE_EQ(
+      vec_unique(merged_node->downstream()),
+      true,
+      phi::errors::PreconditionNotMet(
+          "The downstream nodes of the merged node are not unique."));
 
   // deal with the graph storage.
   AppendNode(merged_node);
Original file line number	Diff line number	Diff line change
`@@ -82,7 +82,7 @@ inline std::vector<fusion::PatternNodePtr<T>> ClusterOps(`
`82`	`82`	`for (const auto& node : result) {`
`83`	`83`	`VLOG(4) << "\n"`
`84`	`84`	`<< node->DebugStr() << "\n"`
`85`		`- << fusion::StmtPatternDebugStr(node->stmt_pattern_);`
	`85`	`+ << fusion::StmtPatternDebugStr(node->stmt_pattern());`
`86`	`86`	`}`
`87`	`87`
`88`	`88`	`return result;`