[CINN] Refine operator-fusion code (#63579)

SigureMo · web-flow · commit 73ff39586f0d · 2024-04-17T15:38:15.000+08:00
* [CINN] Refine `operator-fusion` code

* cleanup and add comment

* refine comment

* cleanup
diff --git a/paddle/cinn/api/op_topo_pattern.h b/paddle/cinn/api/op_topo_pattern.h
diff --git a/paddle/cinn/operator_fusion/backend/pattern_fuser.cc b/paddle/cinn/operator_fusion/backend/pattern_fuser.cc
@@ -185,7 +185,7 @@ struct FusionOpGetter {
 // tmp transform for reduce_tree and reduce_tree_trivial.
 std::vector<FusionOp> GetFusionOpFromPattern(
     const StmtPattern<BackendStage>& pattern) {
-  return std::visit(FusionOpGetter(), pattern);
+  return std::visit(FusionOpGetter(), pattern.variant());
 }
 
 struct FusionOp2Expr {
diff --git a/paddle/cinn/operator_fusion/pattern.h b/paddle/cinn/operator_fusion/pattern.h
@@ -84,10 +84,18 @@ template <typename T>
 class HorizontalFusionPattern {};
 
 template <typename T>
-using StmtPattern = std::variant<TrivialPattern<T>,
-                                 ReducePattern<T>,
-                                 ReduceTreePattern<T>,
-                                 ReduceTreePlusTrivialPattern<T>,
-                                 HorizontalFusionPattern<T>,
-                                 UnsupportPattern<T>>;
+using StmtPatternBase = std::variant<TrivialPattern<T>,
+                                     ReducePattern<T>,
+                                     ReduceTreePattern<T>,
+                                     ReduceTreePlusTrivialPattern<T>,
+                                     HorizontalFusionPattern<T>,
+                                     UnsupportPattern<T>>;
+
+template <typename T>
+struct StmtPattern final : public StmtPatternBase<T> {
+  using StmtPatternBase<T>::StmtPatternBase;
+  const StmtPatternBase<T>& variant() const {
+    return static_cast<const StmtPatternBase<T>&>(*this);
+  }
+};
 }  // namespace cinn::fusion
diff --git a/paddle/cinn/operator_fusion/pattern_fuser.h b/paddle/cinn/operator_fusion/pattern_fuser.h
@@ -31,6 +31,9 @@
 #include "paddle/cinn/operator_fusion/pattern.h"
 #include "paddle/cinn/operator_fusion/utils.h"
 
+// This file is the protocol of the pattern fuser. Please implement
+// ConvertToStmtPattern and MergePatternImpl in the specializations.
+
 namespace cinn::fusion {
 
 template <typename T>
@@ -40,17 +43,16 @@ ReducePattern<T> ToReducePattern(const StmtPattern<T>& second) {
 
 template <typename T>
 std::string GetPatternName(const StmtPattern<T>& s) {
-  return std::visit([](const auto& impl) { return impl.name(); }, s);
+  return std::visit([](const auto& impl) { return impl.name(); }, s.variant());
 }
 
 template <typename T>
-StmtPattern<T> ConvertToStmtPattern(const PatternContent<T>& content) {
-  CHECK(false) << "Please specialization!";
-}
+StmtPattern<T> ConvertToStmtPattern(const PatternContent<T>& content);
 
 template <typename T>
 std::vector<pir::Operation*> GetOpsInPattern(const StmtPattern<T>& pattern) {
-  return std::visit([](const auto& impl) { return impl.ops(); }, pattern);
+  return std::visit([](const auto& impl) { return impl.ops(); },
+                    pattern.variant());
 }
 
 template <typename T>
@@ -159,27 +161,19 @@ StmtPattern<T> MergePatternImpl(const ReduceTreePattern<T>& upstream,
 
 template <typename T>
 StmtPattern<T> MergePatternImpl(const ReduceTreePattern<T>& first,
-                                const TrivialPattern<T>& second) {
-  CHECK(false) << "Please specialization!";
-}
+                                const TrivialPattern<T>& second);
 
 template <typename T>
 StmtPattern<T> MergePatternImpl(const TrivialPattern<T>& first,
-                                const ReducePattern<T>& second) {
-  CHECK(false) << "Please specialization!";
-}
+                                const ReducePattern<T>& second);
 
 template <typename T>
 StmtPattern<T> MergePatternImpl(const TrivialPattern<T>& first,
-                                const TrivialPattern<T>& second) {
-  CHECK(false) << "Please specialization!";
-}
+                                const TrivialPattern<T>& second);
 
 template <typename T>
 StmtPattern<T> MergePatternImpl(const HorizontalFusionPattern<T>& first,
-                                const HorizontalFusionPattern<T>& second) {
-  CHECK(false) << "Please specialization!";
-}
+                                const HorizontalFusionPattern<T>& second);
 
 template <typename T>
 StmtPattern<T> MergePattern(const StmtPattern<T>& first,
@@ -208,7 +202,7 @@ StmtPattern<T> MergePattern(const StmtPattern<T>& first,
                      << "X" << GetPatternName(second);
       },
   };
-  return std::visit(PatternMatch, first, second);
+  return std::visit(PatternMatch, first.variant(), second.variant());
 }
 
 }  // namespace cinn::fusion
diff --git a/paddle/cinn/operator_fusion/pattern_graph.h b/paddle/cinn/operator_fusion/pattern_graph.h
@@ -33,10 +33,8 @@ class PatternGraph {
 
   std::vector<PatternNodePtr<T>> ClusterOps();
 
- public:
   void SinkTrivialPattern();
   void HorizontalFusion();
-  void FuseReducePattern();
   void ReduceLiftReduceTree();
   void ReduceTreeGrown();
   void ReduceTree_Trivial_Fusion();

Original file line number	Diff line number	Diff line change
`@@ -185,7 +185,7 @@ struct FusionOpGetter {`
`185`	`185`	`// tmp transform for reduce_tree and reduce_tree_trivial.`
`186`	`186`	`std::vector<FusionOp> GetFusionOpFromPattern(`
`187`	`187`	`const StmtPattern<BackendStage>& pattern) {`
`188`		`- return std::visit(FusionOpGetter(), pattern);`
	`188`	`+ return std::visit(FusionOpGetter(), pattern.variant());`
`189`	`189`	`}`
`190`	`190`
`191`	`191`	`struct FusionOp2Expr {`