Merge pull request PaddlePaddle#13 from Fridge003/multi-down

skip multi-downstream nodes when doing trivial sink
Fridge003 · Apr 28, 2024 · 2a85bbf · 2a85bbf
2 parents 03d4fa2 + 18a3eb8
commit 2a85bbf
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 39 deletions.
diff --git a/paddle/cinn/operator_fusion/graph_transformer/matcher.h b/paddle/cinn/operator_fusion/graph_transformer/matcher.h
@@ -45,6 +45,18 @@ struct CanFuseRxTMatcher {
   }
 };
 
+struct SinkTrivialMatcher {
+  template <typename T>
+  bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
+    return StmtPatternGraphMatcher<TrivialPattern<T>>()(graph, node) &&
+           node->downstream().size() == 1 &&
+           (std::holds_alternative<ReducePattern<Phrase>>(
+                node->downstream().at(0)->stmt_pattern()) ||
+            std::holds_alternative<TrivialPattern<Phrase>>(
+                node->downstream().at(0)->stmt_pattern()));
+  }
+};
+
 struct CanFuseReduceTreeMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
@@ -133,13 +145,6 @@ struct HorizontalFusionMatcher {
   }
 };
 
-struct NonSinkNodeMatcher {
-  template <typename T>
-  bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
-    return !node->downstream().empty();
-  }
-};
-
 struct IsOutputNodeMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {

diff --git a/paddle/cinn/operator_fusion/graph_transformer/operation.h b/paddle/cinn/operator_fusion/graph_transformer/operation.h
@@ -79,35 +79,48 @@ struct LiftReduceToReduceTreeOperation {
 };
 
 struct MergeTrivialPatternOperation {
-  // TOOD(@wuzhanfei)
   template <typename Phrase>
   void operator()(PatternGraph<Phrase>* graph,
                   PatternNodePtr<Phrase> upstream) {
-    std::vector<PatternNodePtr<Phrase>> fusion_candidate =
-        upstream->downstream();
-    upstream->ClearDownstream();
-    for (const auto& downstream : fusion_candidate) {
-      if (std::holds_alternative<ReducePattern<Phrase>>(
-              downstream->stmt_pattern()) ||
-          std::holds_alternative<TrivialPattern<Phrase>>(
-              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->AddNodeToDownstream(downstream);
-      }
-    }
-    if (upstream->downstream().empty()) {
-      graph->RemoveNode(upstream);
-    }
+    const auto& downstream = node->downstream().at(0);
+    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();
   }
 };
 
+// struct MergeTrivialPatternOperation {
+//   template <typename Phrase>
+//   void operator()(PatternGraph<Phrase>* graph,
+//                   PatternNodePtr<Phrase> upstream) {
+//     std::vector<PatternNodePtr<Phrase>> fusion_candidate =
+//         upstream->downstream();
+//     upstream->ClearDownstream();
+//     for (const auto& downstream : fusion_candidate) {
+//       if (std::holds_alternative<ReducePattern<Phrase>>(
+//               downstream->stmt_pattern()) ||
+//           std::holds_alternative<TrivialPattern<Phrase>>(
+//               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->AddNodeToDownstream(downstream);
+//       }
+//     }
+//     if (upstream->downstream().empty()) {
+//       graph->RemoveNode(upstream);
+//     }
+//   }
+// };
+
 struct LiftToHorizontalFusionPatternOperation {
   template <typename Phrase>
   void operator()(PatternGraph<Phrase>* graph, PatternNodePtr<Phrase> node) {

diff --git a/paddle/cinn/operator_fusion/pattern_graph.cc b/paddle/cinn/operator_fusion/pattern_graph.cc
@@ -28,8 +28,6 @@ template <typename T>
 std::vector<PatternNodePtr<T>> PatternGraph<T>::ClusterOps() {
   VLOG(4) << "[Group Cluster] Initial Condition: " << GraphInfo();
 
-  // TODO(@wuzhanfei) Remove All IsOutputPattern Check
-
   VLOG(4) << "[Group Cluster] Start SinkTrivialPattern";
   SinkTrivialPattern();
   VLOG(4) << "[Group Cluster] After SinkTrivialPattern: " << GraphInfo();
@@ -128,13 +126,10 @@ std::vector<PatternNodePtr<T>> PatternGraph<T>::SortByReverseTopoOrder() {
 
 template <typename T>
 void PatternGraph<T>::SinkTrivialPattern() {
-  // TODO(@wuzhanfei) change sink trivial pattern algorithm, skip pattern with
-  // multi downstream
-  GraphTransformer<
-      NodePattern,
-      T,
-      And<NonSinkNodeMatcher, StmtPatternGraphMatcher<TrivialPattern<T>>>,
-      MergeTrivialPatternOperation>(this);
+  GraphTransformer<NodePattern,
+                   T,
+                   SinkTrivialMatcher,
+                   MergeTrivialPatternOperation>(this);
 }
 
 template <typename T>

diff --git a/paddle/cinn/operator_fusion/pattern_graph.h b/paddle/cinn/operator_fusion/pattern_graph.h
@@ -52,6 +52,7 @@ class PatternGraph {
                               const PatternNodePtr<T>& downstream,
                               MergePatternFn<T> merge_pattern_fn);
   std::vector<PatternNodePtr<T>> SortByTopoOrder();
+  std::vector<PatternNodePtr<T>> SortByReverseTopoOrder();
 
   const PatternNodePtrSet<T>& all_pattern_nodes() const {
     return all_pattern_nodes_;