[CINN] Fix randomness in fusion (#69547)

huangjiyi · web-flow · commit 7e236b7feda9 · 2024-11-21T14:48:32.000+08:00
diff --git a/paddle/cinn/operator_fusion/graph_transformer/search_algorithm.h b/paddle/cinn/operator_fusion/graph_transformer/search_algorithm.h
@@ -42,7 +42,8 @@ struct SearchAlgorithm<NodePattern, GraphMatcher, GraphOperation> {
       if (GraphMatcher()(*graph_, iter_node) &&
           !visited_nodes.count(iter_node)) {
         visited_nodes.insert(iter_node);
-        VLOG(4) << "Find Matched Node: " << iter_node->id();
+        VLOG(4) << "Find Matched Node: " << iter_node->id() << "(" << iter_node
+                << ")";
         return iter_node;
       }
     }
@@ -76,7 +77,8 @@ struct SearchAlgorithm<NodePairPattern, GraphMatcher, GraphOperation> {
         const auto& pair = std::make_pair(i, j);
         if (GraphMatcher()(*graph_, i, j) && !visited_node_pair.count(pair)) {
           visited_node_pair.insert(pair);
-          VLOG(4) << "Find Matched Node Pair: (" << i << ", " << j << ")";
+          VLOG(4) << "Find Matched Node Pair: (" << i->id() << ", " << j->id()
+                  << ")";
           return pair;
         }
       }
diff --git a/paddle/cinn/operator_fusion/pattern.h b/paddle/cinn/operator_fusion/pattern.h
@@ -25,6 +25,16 @@
 
 namespace cinn::fusion {
 
+enum class PatternType {
+  Trivial = 0,
+  Reduce,
+  ReduceTree,
+  ReduceTreePlusTrivial,
+  ItersPermutation,
+  Horizontal,
+  Unsupport,
+};
+
 struct PatternContent {
   explicit PatternContent(pir::Operation* op) : op(op) {}
   pir::Operation* op;
@@ -43,6 +53,7 @@ struct TrivialPattern {
   std::vector<pir::Operation*> ops() const { return ops_; }
   pir::Operation* sink_op() const { return sink_op_; }
 
+  static PatternType type() { return PatternType::Trivial; }
   static std::string name() { return "Trivial"; }
 
   static std::string UniqueId() {
@@ -67,6 +78,7 @@ struct ReducePattern {
   std::vector<pir::Operation*> ops() const { return ops_; }
   pir::Operation* GetReduceOp() const { return ops_.back(); }
 
+  static PatternType type() { return PatternType::Reduce; }
   static std::string name() { return "Reduce"; }
 
   static std::string UniqueId() {
@@ -108,6 +120,7 @@ struct ReduceTreePattern {
     return result;
   }
 
+  static PatternType type() { return PatternType::ReduceTree; }
   static std::string name() { return "ReduceTree"; }
 
   static std::string UniqueId() {
@@ -174,6 +187,7 @@ struct ReduceTreePlusTrivialPattern {
   }
   std::vector<size_t> fake_reduce_iter_idx;
 
+  static PatternType type() { return PatternType::ReduceTreePlusTrivial; }
   static std::string name() { return "ReduceTreePlusTrivial"; }
 
   static std::string UniqueId() {
@@ -219,6 +233,7 @@ struct ItersPermutationPattern {
   std::vector<pir::Operation*> ops_;
   std::vector<pir::Operation*> ops() const { return ops_; }
 
+  static PatternType type() { return PatternType::ItersPermutation; }
   static std::string name() { return "ItersPermutation"; }
   static std::string UniqueId() {
     static std::atomic<int64_t> counter = 0;
@@ -246,6 +261,7 @@ struct HorizontalFusionPattern {
   std::vector<PaddingStmtPattern> padding_patterns_;
   inline std::vector<pir::Operation*> ops() const;
 
+  static PatternType type() { return PatternType::Horizontal; }
   static std::string name() { return "Horizontal"; }
 
   static std::string UniqueId() {
@@ -268,6 +284,7 @@ struct UnsupportPattern {
   std::vector<pir::Operation*> ops_;
   std::vector<pir::Operation*> ops() const { return ops_; }
 
+  static PatternType type() { return PatternType::Unsupport; }
   static std::string name() { return "Unsupport"; }
 
   static std::string UniqueId() {
@@ -332,6 +349,10 @@ static std::string StmtPatternDebugStr(const StmtPattern& stmt) {
   return ss.str();
 }
 
+static PatternType GetPatternType(const StmtPattern& s) {
+  return std::visit([](const auto& impl) { return impl.type(); }, s);
+}
+
 static std::string GetPatternName(const StmtPattern& s) {
   return std::visit([](const auto& impl) { return impl.name(); }, s);
 }
diff --git a/paddle/cinn/operator_fusion/pattern_graph.cc b/paddle/cinn/operator_fusion/pattern_graph.cc
@@ -299,16 +299,22 @@ PatternGraph::PatternGraph(const std::vector<PatternContent>& contents,
 }
 
 void PatternGraph::RemoveNode(const PatternNodePtr& node) {
-  VLOG(4) << "Start Remove: " << node;
-  if (all_pattern_nodes_.find(node) != all_pattern_nodes_.end()) {
-    VLOG(4) << "Removed! ";
-    all_pattern_nodes_.erase(node);
+  VLOG(4) << "Start Remove: " << node->id() << "(" << node << ")";
+  for (auto it = all_pattern_nodes_.begin(); it != all_pattern_nodes_.end();
+       ++it) {
+    // Here we use traversal instead of count() or find() builtin function
+    // because all_pattern_nodes_ is sorted by node id when initialization
+    // but node id may be changed in copy instruction that may destroy the
+    // order of set.
+    if ((*it)->id() == node->id()) {
+      VLOG(4) << "Removed " << (*it)->id();
+      all_pattern_nodes_.erase(it);
+      break;
+    }
   }
-
   for (const PatternNodePtr& upstream : node->upstream()) {
     upstream->RemoveNodeFromDownstream(node);
   }
-
   for (const PatternNodePtr& downstream : node->downstream()) {
     downstream->RemoveNodeFromUpstream(node);
   }
@@ -321,7 +327,7 @@ void PatternGraph::AppendNode(const PatternNodePtr& node) {
 std::string PatternGraph::GraphInfo() const {
   std::stringstream ss;
   ss << "\n========= GraphInfo ===========";
-  for (const auto& v : SortByTopoOrder()) {
+  for (const auto& v : all_pattern_nodes_) {
     ss << "\n##############################";
     ss << "\n" << v->DebugStr();
     ss << "    IsOutput: " << IsOutputNodeMatcher()(*this, v);
diff --git a/paddle/cinn/operator_fusion/pattern_graph.h b/paddle/cinn/operator_fusion/pattern_graph.h
@@ -20,8 +20,6 @@
 
 namespace cinn::fusion {
 
-using PatternNodePtrSet = std::unordered_set<PatternNodePtr>;
-
 using MergePatternFn =
     std::function<StmtPattern(const StmtPattern&, const StmtPattern&)>;
 class PatternGraph {
diff --git a/paddle/cinn/operator_fusion/pattern_node.h b/paddle/cinn/operator_fusion/pattern_node.h
@@ -70,6 +70,7 @@ struct PatternNode {
   void set_stmt_pattern(const StmtPattern& pattern) { stmt_pattern_ = pattern; }
   const std::vector<PatternNodePtr>& upstream() const { return upstream_; }
   const std::vector<PatternNodePtr>& downstream() const { return downstream_; }
+  PatternType type() const { return GetPatternType(stmt_pattern_); }
   std::string name() const { return GetPatternName(stmt_pattern_); }
   std::string id() const { return GetPatternId(stmt_pattern_); }
   void set_return() const { SetReturnInstr(stmt_pattern_); }
@@ -106,4 +107,16 @@ struct PatternNode {
 };
 
 using PatternNodePtr = std::shared_ptr<PatternNode>;
+
+struct PatternNodeCompare {
+  bool operator()(const PatternNodePtr& lhs, const PatternNodePtr& rhs) const {
+    int lhs_id = std::stoi(
+        lhs->id().substr(lhs->id().find_last_of('_') + 1, std::string::npos));
+    int rhs_id = std::stoi(
+        rhs->id().substr(rhs->id().find_last_of('_') + 1, std::string::npos));
+    return lhs->type() == rhs->type() ? lhs_id < rhs_id
+                                      : lhs->type() < rhs->type();
+  }
+};
+using PatternNodePtrSet = std::set<PatternNodePtr, PatternNodeCompare>;
 }  // namespace cinn::fusion

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,8 @@ struct SearchAlgorithm<NodePattern, GraphMatcher, GraphOperation> {`
`42`	`42`	`if (GraphMatcher()(*graph_, iter_node) &&`
`43`	`43`	`!visited_nodes.count(iter_node)) {`
`44`	`44`	`visited_nodes.insert(iter_node);`
`45`		`- VLOG(4) << "Find Matched Node: " << iter_node->id();`
	`45`	`+ VLOG(4) << "Find Matched Node: " << iter_node->id() << "(" << iter_node`
	`46`	`+ << ")";`
`46`	`47`	`return iter_node;`
`47`	`48`	`}`
`48`	`49`	`}`
`@@ -76,7 +77,8 @@ struct SearchAlgorithm<NodePairPattern, GraphMatcher, GraphOperation> {`
`76`	`77`	`const auto& pair = std::make_pair(i, j);`
`77`	`78`	`if (GraphMatcher()(*graph_, i, j) && !visited_node_pair.count(pair)) {`
`78`	`79`	`visited_node_pair.insert(pair);`
`79`		`- VLOG(4) << "Find Matched Node Pair: (" << i << ", " << j << ")";`
	`80`	`+ VLOG(4) << "Find Matched Node Pair: (" << i->id() << ", " << j->id()`
	`81`	`+ << ")";`
`80`	`82`	`return pair;`
`81`	`83`	`}`
`82`	`84`	`}`