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

support LiftToAnchorPattern for reduce tree pattern
Fridge003 · May 8, 2024 · f84c7f4 · f84c7f4
2 parents e10f83a + 0f8b3f5
commit f84c7f4
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 13 deletions.
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/cinn_to_pd_util.cc b/paddle/cinn/hlir/dialect/operator/transforms/cinn_to_pd_util.cc
@@ -37,8 +37,8 @@ pir::Attribute ArrayAttributeToIntArrayAttribute(
 }
 
 const auto& handler_reduce_max_op =
-    [&](::pir::Operation* op,
-        const ::pir::Builder& builder) -> ::pir::Operation* {
+    [](::pir::Operation* op,
+       const ::pir::Builder& builder) -> ::pir::Operation* {
   VLOG(6) << "transform " << op->name() << " from cinn_op to pd_op";
   auto cinn_op = op->dyn_cast<cinn::dialect::ReduceMaxOp>();
   auto attr = cinn_op.attributes();

diff --git a/paddle/cinn/operator_fusion/graph_transformer/matcher.h b/paddle/cinn/operator_fusion/graph_transformer/matcher.h
@@ -88,6 +88,21 @@ struct CanFuseReduceTreeAndTrivialMatcher {
   }
 };
 
+struct LiftToAnchorPatternMatcher {
+  template <typename T>
+  bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
+    bool not_reduce_tree =
+        !StmtPatternGraphMatcher<ReduceTreePattern<T>>()(graph, node) &&
+        !StmtPatternGraphMatcher<ReduceTreePlusTrivialPattern<T>>()(graph,
+                                                                    node);
+    bool reduce_tree_with_single_reduce =
+        StmtPatternGraphMatcher<ReduceTreePattern<T>>()(graph, node) &&
+        std::get<ReduceTreePattern<T>>(node->stmt_pattern()).childs().size() ==
+            0;
+    return not_reduce_tree || reduce_tree_with_single_reduce;
+  }
+};
+
 struct RecomputeNodeMatcher {
   template <typename T>
   bool operator()(const PatternGraph<T>& graph, const PatternNodePtr<T>& node) {
@@ -103,8 +118,11 @@ struct HasUpstreamAnchorMatcher {
                   const PatternNodePtr<T>& upstream,
                   const PatternNodePtr<T>& downstream) {
     return graph.policy_manager()
-        .template GetPolicy<AnchorSearchPolicy>()
-        ->HasUpstreamAnchor(upstream, downstream);
+               .template GetPolicy<GeneralTopoPolicy>()
+               ->CanFuse(upstream, downstream) &&
+           graph.policy_manager()
+               .template GetPolicy<AnchorSearchPolicy>()
+               ->HasUpstreamAnchor(upstream, downstream);
   }
 };
 
@@ -114,8 +132,11 @@ struct HasDownstreamAnchorMatcher {
                   const PatternNodePtr<T>& upstream,
                   const PatternNodePtr<T>& downstream) {
     return graph.policy_manager()
-        .template GetPolicy<AnchorSearchPolicy>()
-        ->HasDownstreamAnchor(upstream, downstream);
+               .template GetPolicy<GeneralTopoPolicy>()
+               ->CanFuse(upstream, downstream) &&
+           graph.policy_manager()
+               .template GetPolicy<AnchorSearchPolicy>()
+               ->HasDownstreamAnchor(upstream, downstream);
   }
 };
 

diff --git a/paddle/cinn/operator_fusion/pattern_graph.cc b/paddle/cinn/operator_fusion/pattern_graph.cc
@@ -178,17 +178,14 @@ void PatternGraph<T>::ReduceTree_Trivial_Fusion() {
 
 template <typename T>
 void PatternGraph<T>::LiftToAnchorPattern() {
-  GraphTransformer<
-      NodePattern,
-      T,
-      And<Not<StmtPatternGraphMatcher<ReduceTreePattern<T>>>,
-          Not<StmtPatternGraphMatcher<ReduceTreePlusTrivialPattern<T>>>>,
-      LiftToAnchorPatternOperation>(this);
+  GraphTransformer<NodePattern,
+                   T,
+                   LiftToAnchorPatternMatcher,
+                   LiftToAnchorPatternOperation>(this);
 }
 
 template <typename T>
 void PatternGraph<T>::AnchorPatternFusion() {
-  // TODO(@wuzhanfei)
   GraphTransformer<ReverseTopoNodePairPattern,
                    T,
                    HasUpstreamAnchorMatcher,