diff --git a/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc
index 606d07fd598268..9c0a2e4501a729 100644
--- a/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc
+++ b/paddle/cinn/hlir/dialect/operator/transforms/cinn_group_cluster_pass.cc
@@ -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).
diff --git a/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc b/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
index 9bc206c53a234c..910bb59f4a3ad4 100644
--- a/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
+++ b/paddle/cinn/hlir/framework/pir/trivial_op_impl.cc
@@ -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;
diff --git a/paddle/cinn/operator_fusion/group_cluster.h b/paddle/cinn/operator_fusion/group_cluster.h
index aa545699a0d4d5..649a2a6a7dcf9f 100644
--- a/paddle/cinn/operator_fusion/group_cluster.h
+++ b/paddle/cinn/operator_fusion/group_cluster.h
@@ -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;
diff --git a/paddle/cinn/operator_fusion/pattern_graph.cc b/paddle/cinn/operator_fusion/pattern_graph.cc
index 73008c4ec49524..a8ab68cf809b36 100644
--- a/paddle/cinn/operator_fusion/pattern_graph.cc
+++ b/paddle/cinn/operator_fusion/pattern_graph.cc
@@ -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,28 +220,22 @@ 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) {
@@ -249,8 +243,16 @@ PatternNodePtr<T> PatternGraph<T>::MergeNode(
     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);
diff --git a/paddle/cinn/operator_fusion/pattern_graph.h b/paddle/cinn/operator_fusion/pattern_graph.h
index 589235d8d76a8c..e6ba1342623491 100644
--- a/paddle/cinn/operator_fusion/pattern_graph.h
+++ b/paddle/cinn/operator_fusion/pattern_graph.h
@@ -17,11 +17,15 @@
 #include "paddle/cinn/operator_fusion/policy/policy_manager.h"
 #include "paddle/cinn/operator_fusion/policy/relative_judge_policy.h"
 #include "paddle/cinn/operator_fusion/utils.h"
+#include "paddle/common/enforce.h"
 
 namespace cinn::fusion {
 
 template <typename T>
 using PatternNodePtrSet = std::unordered_set<PatternNodePtr<T>>;
+template <typename T>
+using MergePatternFn =
+    std::function<StmtPattern<T>(const StmtPattern<T>&, const StmtPattern<T>&)>;
 
 template <typename T>
 class PatternGraph {
@@ -43,10 +47,18 @@ class PatternGraph {
   void AppendNode(const PatternNodePtr<T>& node);
   std::string GraphInfo() const;
   PatternNodePtr<T> MergeNode(const PatternNodePtr<T>& upstream,
-                              const PatternNodePtr<T>& downstream);
+                              const PatternNodePtr<T>& downstream,
+                              MergePatternFn<T> merge_pattern_fn);
   std::vector<PatternNodePtr<T>> SortByTopoOrder();
 
- public:
+  const PatternNodePtrSet<T>& all_pattern_nodes() const {
+    return all_pattern_nodes_;
+  }
+  const std::vector<pir::Value>& outputs() const { return outputs_; }
+  const PolicyManager<T>& policy_manager() const { return policy_manager_; }
+  const PolicyManager<T>& topo_manager() const { return topo_manager_; }
+
+ private:
   PatternNodePtrSet<T> all_pattern_nodes_;
   std::vector<pir::Value> outputs_;
   PolicyManager<T> policy_manager_;
@@ -79,7 +91,7 @@ struct SearchAlgorithm<NodePattern, Phrase, GraphMatcher, GraphOperation> {
   }
 
   PatternNodePtr<Phrase> FindMatchedNode() {
-    for (PatternNodePtr<Phrase> iter_node : graph_->all_pattern_nodes_) {
+    for (PatternNodePtr<Phrase> iter_node : graph_->all_pattern_nodes()) {
       if (GraphMatcher()(*graph_, iter_node) &&
           !visited_nodes.count(iter_node)) {
         visited_nodes.insert(iter_node);
@@ -113,8 +125,8 @@ struct SearchAlgorithm<NodePairPattern, Phrase, GraphMatcher, GraphOperation> {
   }
   std::optional<std::pair<PatternNodePtr<Phrase>, PatternNodePtr<Phrase>>>
   FindMatchedPair() {
-    for (PatternNodePtr<Phrase> i : graph_->all_pattern_nodes_) {
-      for (PatternNodePtr<Phrase> j : graph_->all_pattern_nodes_) {
+    for (PatternNodePtr<Phrase> i : graph_->all_pattern_nodes()) {
+      for (PatternNodePtr<Phrase> j : graph_->all_pattern_nodes()) {
         if (i == j) continue;
         const auto& pair = std::make_pair(i, j);
         if (GraphMatcher()(*graph_, i, j) && !visited_node_pair.count(pair)) {
@@ -142,9 +154,14 @@ struct SearchAlgorithm<NodePairPattern, Phrase, GraphMatcher, GraphOperation> {
 struct MergeReduceTreeOperation {
   template <typename Phrase>
   void operator()(PatternGraph<Phrase>* graph, PatternNodePtr<Phrase> node) {
-    CHECK_EQ(node->downstream_.size(), 1);
-    auto downstream = node->downstream_.at(0);
-    auto merged_node = graph->MergeNode(node, downstream);
+    PADDLE_ENFORCE_EQ(
+        node->downstream().size(),
+        1,
+        phi::errors::PreconditionNotMet(
+            "The downstream of the ReduceTree node should be 1, but got %d.",
+            node->downstream().size()));
+    auto downstream = node->downstream().at(0);
+    auto merged_node = graph->MergeNode(node, downstream, MergePattern<Phrase>);
     graph->RemoveNode(downstream);
     graph->RemoveNode(node);
     VLOG(4) << "MergeReduceTreeOperation: \nupstream " << node->DebugStr()
@@ -156,13 +173,25 @@ struct MergeReduceTreeOperation {
 struct MergeReduceTreeAndTrivialOperation {
   template <typename Phrase>
   void operator()(PatternGraph<Phrase>* graph, PatternNodePtr<Phrase> node) {
-    CHECK_EQ(node->downstream_.size(), 1);
-    auto downstream = node->downstream_.at(0);
+    PADDLE_ENFORCE_EQ(
+        node->downstream().size(),
+        1,
+        phi::errors::PreconditionNotMet(
+            "The downstream of the ReduceTree node should be 1, but got %d.",
+            node->downstream().size()));
+    auto downstream = node->downstream().at(0);
     auto fake_reduce_iter_idx =
-        graph->policy_manager_.GetFakeReduceIterIdx(node, downstream);
-    PatternNodePtr<Phrase> merged_node = graph->MergeNode(node, downstream);
-    std::get<ReduceTreePlusTrivialPattern<Phrase>>(merged_node->stmt_pattern_)
-        .fake_reduce_iter_idx = fake_reduce_iter_idx;
+        graph->policy_manager().GetFakeReduceIterIdx(node, downstream);
+    const auto merge_pattern_fn = [&fake_reduce_iter_idx](
+                                      const StmtPattern<Phrase>& first,
+                                      const StmtPattern<Phrase>& secend) {
+      auto rt_pattern = std::get<ReduceTreePlusTrivialPattern<Phrase>>(
+          MergePattern<Phrase>(first, secend));
+      rt_pattern.fake_reduce_iter_idx = fake_reduce_iter_idx;
+      return rt_pattern;
+    };
+    PatternNodePtr<Phrase> merged_node =
+        graph->MergeNode(node, downstream, merge_pattern_fn);
     graph->RemoveNode(downstream);
     graph->RemoveNode(node);
     VLOG(4) << "MergeReduceTreeAndTrivialOperation: \nupstream "
@@ -174,8 +203,8 @@ struct MergeReduceTreeAndTrivialOperation {
 struct LiftReduceToReduceTreeOperation {
   template <typename Phrase>
   void operator()(PatternGraph<Phrase>* graph, PatternNodePtr<Phrase> node) {
-    const auto& reduce_pattern = ToReducePattern<Phrase>(node->stmt_pattern_);
-    node->stmt_pattern_ = ReduceTreePattern<Phrase>({}, reduce_pattern);
+    const auto& reduce_pattern = ToReducePattern<Phrase>(node->stmt_pattern());
+    node->set_stmt_pattern(ReduceTreePattern<Phrase>({}, reduce_pattern));
     VLOG(4) << "LiftReduceToReduceTreeOperation: \nnode " << node->DebugStr();
   }
 };
@@ -185,24 +214,25 @@ struct MergeTrivialPatternOperation {
   void operator()(PatternGraph<Phrase>* graph,
                   PatternNodePtr<Phrase> upstream) {
     std::vector<PatternNodePtr<Phrase>> fusion_candidate =
-        upstream->downstream_;
-    upstream->downstream_.clear();
+        upstream->downstream();
+    upstream->ClearDownstream();
     for (const auto& downstream : fusion_candidate) {
       if (std::holds_alternative<ReducePattern<Phrase>>(
-              downstream->stmt_pattern_) ||
+              downstream->stmt_pattern()) ||
           std::holds_alternative<TrivialPattern<Phrase>>(
-              downstream->stmt_pattern_)) {
-        auto merged_node = graph->MergeNode(upstream, downstream);
+              downstream->stmt_pattern())) {
+        auto merged_node =
+            graph->MergeNode(upstream, downstream, MergePattern<Phrase>);
         graph->RemoveNode(downstream);
         VLOG(4) << "MergeTrivialPatternOperation: \nupstream "
                 << upstream->DebugStr() << "\ndownstream "
                 << downstream->DebugStr() << "\nmerged "
                 << merged_node->DebugStr();
       } else {
-        upstream->downstream_.push_back(downstream);
+        upstream->AddNodeToDownstream(downstream);
       }
     }
-    if (upstream->downstream_.empty()) {
+    if (upstream->downstream().empty()) {
       graph->RemoveNode(upstream);
     }
   }
@@ -210,8 +240,34 @@ struct MergeTrivialPatternOperation {
 
 struct LiftToHorizontalFusionPatternOperation {
   template <typename Phrase>
-  void operator()(PatternGraph<Phrase>* graph, PatternNodePtr<Phrase> i) {
-    i->stmt_pattern_ = HorizontalFusionPattern<Phrase>({i->stmt_pattern_});
+  void operator()(PatternGraph<Phrase>* graph, PatternNodePtr<Phrase> node) {
+    node->set_stmt_pattern(
+        HorizontalFusionPattern<Phrase>({node->stmt_pattern()}));
+  }
+};
+
+struct HorizontalFusionOperation {
+  template <typename Phrase>
+  void operator()(PatternGraph<Phrase>* graph,
+                  const PatternNodePtr<Phrase>& i,
+                  const PatternNodePtr<Phrase>& j) {
+    PADDLE_ENFORCE_EQ(
+        GetPatternName(i->stmt_pattern()),
+        HorizontalFusionPattern<Phrase>::name(),
+        phi::errors::PreconditionNotMet(
+            "The pattern of the first node should be HorizontalFusionPattern, "
+            "but got %s.",
+            GetPatternName(i->stmt_pattern())));
+    PADDLE_ENFORCE_EQ(
+        GetPatternName(j->stmt_pattern()),
+        HorizontalFusionPattern<Phrase>::name(),
+        phi::errors::PreconditionNotMet(
+            "The pattern of the second node should be HorizontalFusionPattern, "
+            "but got %s.",
+            GetPatternName(j->stmt_pattern())));
+    graph->MergeNode(i, j, MergePattern<Phrase>);
+    graph->RemoveNode(i);
+    graph->RemoveNode(j);
   }
 };
 
@@ -229,17 +285,18 @@ template <typename StmtPattern>
 struct StmtPatternGraphMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
-    return GetPatternName(node->stmt_pattern_) == StmtPattern::name();
+    return GetPatternName(node->stmt_pattern()) == StmtPattern::name();
   }
 };
 
 struct CanFuseRxTMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
-    return (std::holds_alternative<ReduceTreePattern<T>>(node->stmt_pattern_) &&
-            !node->downstream_.empty() &&
-            std::holds_alternative<TrivialPattern<T>>(
-                node->downstream_.at(0)->stmt_pattern_));
+    return (
+        std::holds_alternative<ReduceTreePattern<T>>(node->stmt_pattern()) &&
+        !node->downstream().empty() &&
+        std::holds_alternative<TrivialPattern<T>>(
+            node->downstream().at(0)->stmt_pattern()));
   }
 };
 
@@ -247,10 +304,10 @@ struct CanFuseReduceTreeMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
     return StmtPatternGraphMatcher<ReduceTreePattern<T>>()(graph, node) &&
-           !node->downstream_.empty() &&
+           !node->downstream().empty() &&
            std::holds_alternative<ReduceTreePattern<T>>(
-               node->downstream_.at(0)->stmt_pattern_) &&
-           graph.policy_manager_.CanFuse(node, node->downstream_.at(0));
+               node->downstream().at(0)->stmt_pattern()) &&
+           graph.policy_manager().CanFuse(node, node->downstream().at(0));
   }
 };
 
@@ -258,10 +315,10 @@ struct CanFuseReduceTreeAndTrivialMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
     return StmtPatternGraphMatcher<ReduceTreePattern<T>>()(graph, node) &&
-           !node->downstream_.empty() &&
+           !node->downstream().empty() &&
            std::holds_alternative<TrivialPattern<T>>(
-               node->downstream_.at(0)->stmt_pattern_) &&
-           graph.policy_manager_.CanFuse(node, node->downstream_.at(0));
+               node->downstream().at(0)->stmt_pattern()) &&
+           graph.policy_manager().CanFuse(node, node->downstream().at(0));
   }
 };
 
@@ -276,45 +333,32 @@ struct HorizontalFusionConstrain {
     if (!StmtPatternGraphMatcher<HorizontalFusionPattern<T>>()(graph, second)) {
       return false;
     }
-    const auto& first_dim = first->sink_op_->result(0)
+    const auto& first_dim = first->sink_op()
+                                ->result(0)
                                 .type()
                                 .template dyn_cast<pir::DenseTensorType>()
                                 .dims();
-    const auto& second_dim = second->sink_op_->result(0)
+    const auto& second_dim = second->sink_op()
+                                 ->result(0)
                                  .type()
                                  .template dyn_cast<pir::DenseTensorType>()
                                  .dims();
-    return graph.topo_manager_.CanFuse(first, second) &&
+    return graph.topo_manager().CanFuse(first, second) &&
            first_dim == second_dim;
   }
 };
 
-struct HorizontalFusionOperation {
-  template <typename T>
-  void operator()(PatternGraph<T>* graph,
-                  const PatternNodePtr<T>& i,
-                  const PatternNodePtr<T>& j) {
-    CHECK(GetPatternName(i->stmt_pattern_) ==
-          HorizontalFusionPattern<T>::name());
-    CHECK(GetPatternName(j->stmt_pattern_) ==
-          HorizontalFusionPattern<T>::name());
-    graph->MergeNode(i, j);
-    graph->RemoveNode(i);
-    graph->RemoveNode(j);
-  }
-};
-
 struct NonSinkNodeMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
-    return !node->downstream_.empty();
+    return !node->downstream().empty();
   }
 };
 
 struct IsOutputNodeMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
-    bool res = IsAnyFirstInSecond(node->sink_op_->results(), graph.outputs_);
+    bool res = IsAnyFirstInSecond(node->sink_op()->results(), graph.outputs());
     return res;
   }
 };
@@ -331,7 +375,7 @@ template <int N>
 struct DownstreamSmallerThan {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
-    return node->downstream_.size() < N;
+    return node->downstream().size() < N;
   }
 };
 
diff --git a/paddle/cinn/operator_fusion/pattern_node.h b/paddle/cinn/operator_fusion/pattern_node.h
index d6c9f8202669ef..459522b8341de2 100644
--- a/paddle/cinn/operator_fusion/pattern_node.h
+++ b/paddle/cinn/operator_fusion/pattern_node.h
@@ -23,15 +23,27 @@ namespace cinn::fusion {
 template <typename T>
 struct PatternNode {
   using PatternNodePtr = std::shared_ptr<PatternNode<T>>;
+  using MergePatternFn = std::function<StmtPattern<T>(const StmtPattern<T>&,
+                                                      const StmtPattern<T>&)>;
 
   explicit PatternNode(const PatternContent<T>& content)
       : sink_op_(content.op), stmt_pattern_(ConvertToStmtPattern<T>(content)) {}
 
   explicit PatternNode(PatternNodePtr fused_up_node,
-                       PatternNodePtr fused_down_node)
+                       PatternNodePtr fused_down_node,
+                       MergePatternFn merge_pattern_fn)
       : sink_op_(fused_down_node->sink_op_),
-        stmt_pattern_(MergePattern<T>(fused_up_node->stmt_pattern_,
-                                      fused_down_node->stmt_pattern_)) {}
+        stmt_pattern_(merge_pattern_fn(fused_up_node->stmt_pattern_,
+                                       fused_down_node->stmt_pattern_)) {
+    // Update the upstream & downstream
+    ExtendVector(&upstream_, fused_up_node->upstream());
+    ExtendVector(&upstream_, fused_down_node->upstream());
+    RemoveFromVector(&upstream_, fused_up_node);
+
+    ExtendVector(&downstream_, fused_up_node->downstream());
+    ExtendVector(&downstream_, fused_down_node->downstream());
+    RemoveFromVector(&downstream_, fused_down_node);
+  }
 
   std::string DebugStr() const {
     std::stringstream ss;
@@ -47,6 +59,27 @@ struct PatternNode {
     return ss.str();
   }
 
+  pir::Operation* sink_op() const { return sink_op_; }
+  const StmtPattern<T>& stmt_pattern() const { return stmt_pattern_; }
+  void set_stmt_pattern(const StmtPattern<T>& pattern) {
+    stmt_pattern_ = pattern;
+  }
+  const std::vector<PatternNodePtr>& upstream() const { return upstream_; }
+  const std::vector<PatternNodePtr>& downstream() const { return downstream_; }
+  void AddNodeToUpstream(PatternNodePtr node) { upstream_.push_back(node); }
+  void AddNodeToDownstream(PatternNodePtr node) { downstream_.push_back(node); }
+  void RemoveNodeFromUpstream(PatternNodePtr node) {
+    RemoveFromVector(&upstream_, node);
+  }
+  void RemoveNodeFromDownstream(PatternNodePtr node) {
+    RemoveFromVector(&downstream_, node);
+  }
+  void ClearUpstream() { upstream_.clear(); }
+  void ClearDownstream() { downstream_.clear(); }
+  void UniqueUpstream() { upstream_ = UniqueVectorBySet(upstream_); }
+  void UniqueDownstream() { downstream_ = UniqueVectorBySet(downstream_); }
+
+ private:
   StmtPattern<T> stmt_pattern_;
   pir::Operation* sink_op_;
 
diff --git a/paddle/cinn/operator_fusion/policy/general_topo_policy.cc b/paddle/cinn/operator_fusion/policy/general_topo_policy.cc
index 53d54b8fa0f65e..e4cca9804a79f2 100644
--- a/paddle/cinn/operator_fusion/policy/general_topo_policy.cc
+++ b/paddle/cinn/operator_fusion/policy/general_topo_policy.cc
@@ -24,7 +24,7 @@ template <typename T>
 bool IsDownstreamNode(const PatternNodePtr<T> start,
                       const PatternNodePtr<T> target) {
   if (start == target) return true;
-  for (const auto& down_node : start->downstream_) {
+  for (const auto& down_node : start->downstream()) {
     if (IsDownstreamNode(down_node, target)) return true;
   }
   return false;
@@ -33,7 +33,7 @@ bool IsDownstreamNode(const PatternNodePtr<T> start,
 template <typename T>
 bool IsIndirectDownstreamNode(const PatternNodePtr<T> start,
                               const PatternNodePtr<T> target) {
-  for (const auto& node : start->downstream_) {
+  for (const auto& node : start->downstream()) {
     if (node == target) continue;
     if (IsDownstreamNode(node, target)) return true;
   }
diff --git a/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc b/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc
index 954593778a7b7f..626f54c215b6ea 100644
--- a/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc
+++ b/paddle/cinn/operator_fusion/policy/relative_judge_policy.cc
@@ -110,12 +110,12 @@ template <typename T>
 bool RelativeJudgePolicy<T>::ReduceTreeGrownCanMerge(
     const PatternNodePtr<T>& upstream, const PatternNodePtr<T>& downstream) {
   const auto& upstream_tree =
-      std::get<ReduceTreePattern<T>>(upstream->stmt_pattern_);
-  VLOG(4) << "upstream->stmt_pattern_:"
+      std::get<ReduceTreePattern<T>>(upstream->stmt_pattern());
+  VLOG(4) << "upstream->stmt_pattern():"
           << OpsDebugStr(GetOpsInPattern<T>(upstream_tree));
   const auto& downstream_tree =
-      std::get<ReduceTreePattern<T>>(downstream->stmt_pattern_);
-  VLOG(4) << "downstream->stmt_pattern_"
+      std::get<ReduceTreePattern<T>>(downstream->stmt_pattern());
+  VLOG(4) << "downstream->stmt_pattern()"
           << OpsDebugStr(GetOpsInPattern<T>(downstream_tree));
   const auto& maybe_downstream_op = GetDownstreamFromCandidate(
       upstream_tree.GetRootPattern(), downstream_tree.FlattenReducePattern());
@@ -202,7 +202,7 @@ std::vector<ValueDim> RelativeJudgePolicy<T>::getUpstreamReduceDims(
     ShardableAxesInfoManager& axes_info) {  // NOLINT
   const auto& split_reduce_input_dims_result =
       SplitReduceInputDimsIfRelatedWithNonReduceAxis(
-          axes_info.GetSignature(upstream->sink_op_), upstream->sink_op_);
+          axes_info.GetSignature(upstream->sink_op()), upstream->sink_op());
   return split_reduce_input_dims_result.non_related;
 }
 
@@ -213,11 +213,11 @@ std::vector<ValueDim> RelativeJudgePolicy<T>::getDownstreamUnrelatedDims(
     ShardableAxesInfoManager& axes_info) {  // NOLINT
   const auto& split_reduce_output_dims_result =
       SplitReduceOutputDimsIfRelatedWithNonReduceAxis(
-          axes_info.GetSignature(upstream->sink_op_), upstream->sink_op_);
+          axes_info.GetSignature(upstream->sink_op()), upstream->sink_op());
   const auto& upstream_non_reduce_dims =
       split_reduce_output_dims_result.related;
   const auto& split_trivial_dims_result = SplitDimsWithRelationship(
-      GetAllValueDimFromValue(downstream->sink_op_->result(0)),
+      GetAllValueDimFromValue(downstream->sink_op()->result(0)),
       upstream_non_reduce_dims);
   VLOG(4) << split_trivial_dims_result.DebugStr();
   return split_trivial_dims_result.non_related;
@@ -237,8 +237,10 @@ bool RelativeJudgePolicy<T>::ReducePlusTrivialCanMerge(
   return res;
 }
 
-static std::vector<ValueDim> GatherDimsExcept(
-    const std::vector<ValueDim>& dims, const std::vector<size_t>& except) {
+namespace {
+
+std::vector<ValueDim> GatherDimsExcept(const std::vector<ValueDim>& dims,
+                                       const std::vector<size_t>& except) {
   std::vector<ValueDim> result;
   for (size_t i = 0; i < dims.size(); i++) {
     if (std::find(except.begin(), except.end(), i) == except.end()) {
@@ -248,7 +250,7 @@ static std::vector<ValueDim> GatherDimsExcept(
   return result;
 }
 
-static symbol::DimExpr GetProductDimExprForValueDims(
+symbol::DimExpr GetProductDimExprForValueDims(
     const std::vector<ValueDim>& dims) {
   if (dims.empty()) {
     return 0;
@@ -262,8 +264,8 @@ static symbol::DimExpr GetProductDimExprForValueDims(
   return shape_analysis.GetProductDimExpr(dims[0].v_, dim_idx);
 }
 
-static bool IsProductSmallerOrEqual(const std::vector<ValueDim>& first,
-                                    const std::vector<ValueDim>& second) {
+bool IsProductSmallerOrEqual(const std::vector<ValueDim>& first,
+                             const std::vector<ValueDim>& second) {
   if (first.empty()) return true;
   const auto& first_product = GetProductDimExprForValueDims(first);
   const auto& second_product = GetProductDimExprForValueDims(second);
@@ -279,15 +281,17 @@ static bool IsProductSmallerOrEqual(const std::vector<ValueDim>& first,
   return shape_analysis.IsEqual(first_product, second_product);
 }
 
+}  // namespace
+
 template <typename T>
 bool RelativeJudgePolicy<T>::IsFlattenDimSmaller(
     const PatternNodePtr<T>& upstream, const PatternNodePtr<T>& downstream) {
   const auto& fakes = GetFakeReduceIterIdx(upstream, downstream);
   VLOG(4) << "IsFlattenDimSmaller: fake is " << utils::Join(fakes, ",");
   const auto& downstream_free_dims = GatherDimsExcept(
-      GetAllValueDimFromValue(downstream->sink_op_->result(0)), fakes);
+      GetAllValueDimFromValue(downstream->sink_op()->result(0)), fakes);
   const auto& upstream_free_dims =
-      GetAllValueDimFromValue(upstream->sink_op_->result(0));
+      GetAllValueDimFromValue(upstream->sink_op()->result(0));
 
   bool res = IsProductSmallerOrEqual(downstream_free_dims, upstream_free_dims);
   VLOG(4) << "IsFlattenDimSmaller: " << res;
@@ -297,12 +301,13 @@ bool RelativeJudgePolicy<T>::IsFlattenDimSmaller(
 template <typename T>
 bool RelativeJudgePolicy<T>::CanFuse(const PatternNodePtr<T>& upstream,
                                      const PatternNodePtr<T>& downstream) {
-  if (std::holds_alternative<ReduceTreePattern<T>>(upstream->stmt_pattern_) &&
-      std::holds_alternative<TrivialPattern<T>>(downstream->stmt_pattern_)) {
+  if (std::holds_alternative<ReduceTreePattern<T>>(upstream->stmt_pattern()) &&
+      std::holds_alternative<TrivialPattern<T>>(downstream->stmt_pattern())) {
     return ReducePlusTrivialCanMerge(upstream, downstream);
   }
-  if (std::holds_alternative<ReduceTreePattern<T>>(upstream->stmt_pattern_) &&
-      std::holds_alternative<ReduceTreePattern<T>>(downstream->stmt_pattern_)) {
+  if (std::holds_alternative<ReduceTreePattern<T>>(upstream->stmt_pattern()) &&
+      std::holds_alternative<ReduceTreePattern<T>>(
+          downstream->stmt_pattern())) {
     return ReduceTreeGrownCanMerge(upstream, downstream);
   }
   return true;  // other case.
@@ -311,15 +316,15 @@ bool RelativeJudgePolicy<T>::CanFuse(const PatternNodePtr<T>& upstream,
 template <typename T>
 std::vector<size_t> RelativeJudgePolicy<T>::GetFakeReduceIterIdx(
     const PatternNodePtr<T>& upstream, const PatternNodePtr<T>& downstream) {
-  if (!std::holds_alternative<ReduceTreePattern<T>>(upstream->stmt_pattern_) &&
-      !std::holds_alternative<TrivialPattern<T>>(downstream->stmt_pattern_)) {
+  if (!std::holds_alternative<ReduceTreePattern<T>>(upstream->stmt_pattern()) &&
+      !std::holds_alternative<TrivialPattern<T>>(downstream->stmt_pattern())) {
     PADDLE_THROW("Illegal Call GetFakeReduceIterIdx");
   }
 
   // TODO(xiongkun): replace after fix bug in relation that if has multi path in
   // graph const auto& split_reduce_dims_result =
   // SplitReduceInputDimsIfRelatedWithNonReduceAxis(
-  // axes_info_.GetSignature(upstream->sink_op_), upstream->sink_op_);
+  // axes_info_.GetSignature(upstream->sink_op()), upstream->sink_op());
 
   // const auto& upstream_reduce_dims = split_reduce_dims_result.non_related;
   // const auto& upstream_non_reduce_dims = split_reduce_dims_result.related;
@@ -327,12 +332,12 @@ std::vector<size_t> RelativeJudgePolicy<T>::GetFakeReduceIterIdx(
 
   const auto& split_reduce_input_dims_result =
       SplitReduceInputDimsIfRelatedWithNonReduceAxis(
-          axes_info_.GetSignature(upstream->sink_op_), upstream->sink_op_);
+          axes_info_.GetSignature(upstream->sink_op()), upstream->sink_op());
   VLOG(4) << split_reduce_input_dims_result.DebugStr();
   const auto& upstream_reduce_dims = split_reduce_input_dims_result.non_related;
   const auto& split_reduce_output_dims_result =
       SplitReduceOutputDimsIfRelatedWithNonReduceAxis(
-          axes_info_.GetSignature(upstream->sink_op_), upstream->sink_op_);
+          axes_info_.GetSignature(upstream->sink_op()), upstream->sink_op());
   VLOG(4) << split_reduce_input_dims_result.DebugStr();
   const auto& upstream_non_reduce_dims =
       split_reduce_output_dims_result.related;
@@ -340,7 +345,7 @@ std::vector<size_t> RelativeJudgePolicy<T>::GetFakeReduceIterIdx(
   // =======================
 
   const auto& split_trivial_dims_result = SplitDimsWithRelationship(
-      GetAllValueDimFromValue(downstream->sink_op_->result(0)),
+      GetAllValueDimFromValue(downstream->sink_op()->result(0)),
       upstream_non_reduce_dims);
 
   const auto& trivial_reorder_dims = split_trivial_dims_result.non_related;
diff --git a/paddle/cinn/operator_fusion/policy/shardable_axes_policy.cc b/paddle/cinn/operator_fusion/policy/shardable_axes_policy.cc
index 24ffa6d862c863..4b8b758f449cd5 100644
--- a/paddle/cinn/operator_fusion/policy/shardable_axes_policy.cc
+++ b/paddle/cinn/operator_fusion/policy/shardable_axes_policy.cc
@@ -62,9 +62,9 @@ bool ShardableAxesRRFusePolicy<T>::ReduceTreeGrownCanMerge(
     return false;
   }
   const auto& upstream_tree =
-      std::get<ReduceTreePattern>(upstream->stmt_pattern_);
+      std::get<ReduceTreePattern>(upstream->stmt_pattern());
   const auto& downstream_tree =
-      std::get<ReduceTreePattern>(downstream->stmt_pattern_);
+      std::get<ReduceTreePattern>(downstream->stmt_pattern());
   const auto& maybe_downstream_op = GetDownstreamFromCandidate(
       upstream_tree.GetRootPattern(), downstream_tree.reduce_patterns_);
   if (!maybe_downstream_op.has_value()) {