[CINN] Fix FuseParallelMatmulPass bug (#64405)

jiahy0825 · web-flow · commit 64d80ed4aab0 · 2024-05-20T14:21:43.000+08:00
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/fuse_parallel_matmul_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/fuse_parallel_matmul_pass.cc
@@ -21,6 +21,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/transforms/sub_graph_detector.h"
 #include "paddle/pir/include/core/builtin_dialect.h"
 #include "paddle/pir/include/pass/pass.h"
 #include "paddle/pir/include/pattern_rewrite/frozen_rewrite_pattern_set.h"
@@ -65,7 +66,7 @@ class MergeParallelMatmulPattern
     };
 
     auto input_x = matmul_op.operand_source(0);
-    const std::vector<pir::Operation*> merge_ops = [&]() {
+    std::vector<pir::Operation*> merge_ops = [&]() {
       std::vector<pir::Operation*> ret;
       std::optional<std::vector<std::int64_t>> pre_dim;
       std::vector<std::int64_t> cur_dim;
@@ -99,6 +100,16 @@ class MergeParallelMatmulPattern
     if (merge_ops.size() <= 1) {
       return false;
     }
+    std::sort(
+        merge_ops.begin(),
+        merge_ops.end(),
+        [&](pir::Operation* a, pir::Operation* b) {
+          int a_distance = std::distance(a->GetParent()->begin(),
+                                         a->operator pir::Block::Iterator());
+          int b_distance = std::distance(b->GetParent()->begin(),
+                                         b->operator pir::Block::Iterator());
+          return a_distance < b_distance;
+        });
 
     const std::vector<pir::Value> combine_ins = [&]() {
       std::vector<pir::Value> ret;
@@ -118,6 +129,18 @@ class MergeParallelMatmulPattern
       }
       return ret;
     }();
+    const std::vector<pir::Value> outputs = [&]() {
+      std::vector<pir::Value> ret;
+      for (pir::Operation* matmul_op : merge_ops) {
+        ret.push_back(matmul_op->result(0));
+      }
+      return ret;
+    }();
+
+    auto* insert_point = FindInsertPoint(merge_ops, outputs);
+    MoveUpstreamOpBeforeGroup(
+        merge_ops, merge_ops.back()->GetParent(), insert_point);
+    rewriter.set_insertion_point(insert_point);
 
     auto combine_out = rewriter.Build<pir::CombineOp>(combine_ins).result(0);
     auto concat_out =
@@ -127,7 +150,6 @@ class MergeParallelMatmulPattern
             .result(0);
 
     for (size_t i = 0; i < merge_ops.size(); ++i) {
-      rewriter.SetInsertionPointAfter(merge_ops[i]);
       auto split_out = rewriter
                            .Build<paddle::dialect::SliceOp>(
                                matmul_out,
diff --git a/paddle/fluid/pir/transforms/sub_graph_detector.cc b/paddle/fluid/pir/transforms/sub_graph_detector.cc
@@ -496,37 +496,6 @@ std::vector<pir::Value> AnalysisOutputs(
 
 namespace {
 
-pir::Operation* FindInsertPoint(const GroupOpsVec& group_ops,
-                                const std::vector<pir::Value>& outputs) {
-  // Regard last op as insert position if there are no downstream ops between in
-  // group_ops.
-  pir::Operation* insert_point_op = group_ops.back();
-  auto begin = group_ops.front()->operator Block::ConstIterator();
-  auto end = ++(group_ops.back()->operator Block::ConstIterator());
-  const std::unordered_set<pir::Value> outputs_set(outputs.begin(),
-                                                   outputs.end());
-  const std::unordered_set<const pir::Operation*> group_ops_set(
-      group_ops.begin(), group_ops.end());
-
-  const auto& IsDownstreamOp = [&](const pir::Operation* op) -> bool {
-    if (group_ops_set.find(op) != group_ops_set.end()) return false;
-    for (auto& value : GetUsedExternalValue(*op)) {
-      if (outputs_set.find(value) != outputs_set.end()) {
-        return true;
-      }
-    }
-    return false;
-  };
-  // Find first downstream op as final insert position.
-  for (; begin != end; ++begin) {
-    if (IsDownstreamOp(begin)) {
-      insert_point_op = begin;
-      break;
-    }
-  }
-  return insert_point_op;
-}
-
 struct IncrementalOrder {
   bool operator()(const pir::Operation* lhs, const pir::Operation* rhs) const {
     CHECK(lhs->GetParent() == rhs->GetParent())
@@ -574,6 +543,7 @@ std::unordered_set<pir::Operation*> GetUpstreamOpsAfterPosition(
   }
   return ops;
 }
+}  // namespace
 
 void MoveUpstreamOpBeforeGroup(const GroupOpsVec& group_ops,
                                pir::Block* block,
@@ -595,7 +565,37 @@ void MoveUpstreamOpBeforeGroup(const GroupOpsVec& group_ops,
     op->MoveTo(block, insert_point_op->operator Block::Iterator());
   }
 }
-}  // namespace
+
+pir::Operation* FindInsertPoint(const GroupOpsVec& group_ops,
+                                const std::vector<pir::Value>& outputs) {
+  // Regard last op as insert position if there are no downstream ops between in
+  // group_ops.
+  pir::Operation* insert_point_op = group_ops.back();
+  auto begin = group_ops.front()->operator Block::ConstIterator();
+  auto end = ++(group_ops.back()->operator Block::ConstIterator());
+  const std::unordered_set<pir::Value> outputs_set(outputs.begin(),
+                                                   outputs.end());
+  const std::unordered_set<const pir::Operation*> group_ops_set(
+      group_ops.begin(), group_ops.end());
+
+  const auto& IsDownstreamOp = [&](const pir::Operation* op) -> bool {
+    if (group_ops_set.find(op) != group_ops_set.end()) return false;
+    for (auto& value : GetUsedExternalValue(*op)) {
+      if (outputs_set.find(value) != outputs_set.end()) {
+        return true;
+      }
+    }
+    return false;
+  };
+  // Find first downstream op as final insert position.
+  for (; begin != end; ++begin) {
+    if (IsDownstreamOp(begin)) {
+      insert_point_op = begin;
+      break;
+    }
+  }
+  return insert_point_op;
+}
 
 void ReplaceWithGroupOp(pir::Block* block,
                         const GroupOpsVec& group_ops) {  // NOLINT
diff --git a/paddle/fluid/pir/transforms/sub_graph_detector.h b/paddle/fluid/pir/transforms/sub_graph_detector.h
@@ -73,4 +73,10 @@ class SubgraphDetector {
 std::vector<pir::Value> AnalysisOutputs(const GroupOpsVec& group_ops);
 void ReplaceWithGroupOp(pir::Block* block, const GroupOpsVec& group_ops);
 
+pir::Operation* FindInsertPoint(const GroupOpsVec& group_ops,
+                                const std::vector<pir::Value>& outputs);
+void MoveUpstreamOpBeforeGroup(const GroupOpsVec& group_ops,
+                               pir::Block* block,
+                               pir::Operation* insert_point_op);
+
 }  // namespace pir
