PaddlePaddle
diff --git a/‎paddle/fluid/pir/transforms/gpu/horizontal_fuse_pass.cc
Copy file name to clipboard
+162-67Lines changed: 162 additions & 67 deletions b/‎paddle/fluid/pir/transforms/gpu/horizontal_fuse_pass.cc
Copy file name to clipboard
+162-67Lines changed: 162 additions & 67 deletions
@@ -42,12 +42,18 @@ class HorizontalFusePattern : public pir::RewritePattern {
                        1 /*benefit*/,
                        context,
                        {} /*generated_names*/) {}
+
+  // distinguish input is X or Y, where to get weight
+  // 1 means w from Y, 0 means w from x
+  mutable uint32_t w_from_source_index;
+
   bool MatchAndRewrite(
       pir::Operation* op,
       pir::PatternRewriter& rewriter) const override {  // NOLINT
     bool match_flag = false;
     for (size_t i = 0; i < op->num_results(); i++) {
-      if (GetOpCntUseX(op->result(i)) > 1) {
+      // only fuse three or more op
+      if (GetOpCntUseValue(op->result(i)) >= 3) {
         match_flag = true;
         RewriteOpsbyValue(op, &rewriter, i);
         VLOG(4) << "horizontal_fuse_pass applied rewrite on [" << op->name()
@@ -58,17 +64,39 @@ class HorizontalFusePattern : public pir::RewritePattern {
   }
 
  private:
-  bool IsValidOp(pir::Operation* curr_op, const pir::Value& x) const {
-    // x must be the first operand of fc/gemm_epilogue/matmul
-    if (!(curr_op->num_operands() > 0 && curr_op->operand_source(0) == x)) {
-      return false;
+  uint32_t operands_source_index(pir::Operation* op,
+                                 const pir::Value& target_value) const {
+    std::vector<pir::Value> sources = op->operands_source();
+    auto it = std::find(sources.begin(), sources.end(), target_value);
+
+    if (it != sources.end()) {
+      return std::distance(sources.begin(), it);
+    } else {
+      return static_cast<uint32_t>(-1);  // use -1 to indicate not found
     }
+  }
+
+  bool IsValidOp(pir::Operation* curr_op, const pir::Value& val) const {
+    // judge prev_op is a ParameterOp\ConstantTensorOp\DataOp
+    auto areOperandsValid = [this, curr_op, val]() -> bool {
+      uint32_t curr_op_input_index = operands_source_index(curr_op, val);
+      for (uint32_t i = 0; i < curr_op->num_operands(); i++) {
+        auto* prev_op = pir::GetDefiningOpForInput(curr_op, i);
+        if (i != curr_op_input_index) {
+          return prev_op && (prev_op->isa<pir::ParameterOp>() ||
+                             prev_op->isa<pir::ConstantTensorOp>() ||
+                             prev_op->isa<paddle::dialect::DataOp>());
+        }
+      }
+      return false;
+    };
 
     if (curr_op->isa<paddle::dialect::GemmEpilogueOp>() ||
         curr_op->isa<paddle::dialect::FcOp>()) {
-      return true;
+      return areOperandsValid();
     }
 
+    // check Attribute
     if (curr_op->isa<paddle::dialect::MatmulOp>()) {
       if (curr_op->HasAttribute("transpose_x") &&
           curr_op->attribute<pir::BoolAttribute>("transpose_x").data() ==
@@ -80,62 +108,60 @@ class HorizontalFusePattern : public pir::RewritePattern {
               true) {
         return false;
       }
-      return true;
+      return areOperandsValid();
     }
 
     return false;
   }
 
-  bool AreAttributeMapsEqual(const pir::AttributeMap& attrs1,
-                             const pir::AttributeMap& attrs2) const {
-    if (attrs1.size() != attrs2.size()) {
+  bool AreOperationsEqual(const pir::Operation* op1,
+                          const pir::Operation* op2) const {
+    if (op1->name() != op2->name() ||
+        op1->num_operands() != op2->num_operands() ||
+        op1->num_results() != op2->num_results()) {
       return false;
     }
-    pir::AttributeMap attrs2_copy(attrs2);
-    for (const auto& kv : attrs1) {
-      auto it = attrs2_copy.find(kv.first);
-      if (it == attrs2_copy.end() || it->second != kv.second) {
-        return false;
-      }
-      attrs2_copy.erase(it);
+
+    // when pdmodel has sub_graph, the shape of y1 and y2 may be different
+    // Check if dimensions sizes are the same and first dimensions are the same
+    const pir::Value& w1 = op1->operand_source(w_from_source_index);
+    const pir::Value& w2 = op2->operand_source(w_from_source_index);
+    auto w1_dims = pir::GetShapeFromValue(w1);
+    auto w2_dims = pir::GetShapeFromValue(w2);
+
+    // Determine which dimension to compare based on w_from_source_index
+    std::size_t compare_dim = (w_from_source_index == 0) ? 1 : 0;
+
+    // Check if dimensions sizes are different or the specified dimension value
+    // is different
+    if (w1_dims.size() != w2_dims.size() ||
+        w1_dims[compare_dim] != w2_dims[compare_dim]) {
+      return false;
     }
-    return attrs2_copy.empty();
+    return true;
   }
 
-  int GetOpCntUseX(const pir::Value& x) const {
-    // At least two same MatmulOp/GemmEpilogueOp/FcOp using x
-    // MatmulOp/GemmEpilogueOp/FcOp mutually exclusive and don't appear at the
-    // same time. All ops'attrs should also be exactly equal. if not, fusion is
-    // not performed.
+  int GetOpCntUseValue(const pir::Value& val) const {
     int op_cnt_use_x = 0;
     pir::Operation* op_example = nullptr;
-    std::string op_example_name;
-    uint32_t op_example_num_operands = 0;
-    pir::AttributeMap op_example_attrs;
-    for (auto it = x.use_begin(); it != x.use_end(); ++it) {
+
+    for (auto it = val.use_begin(); it != val.use_end(); ++it) {
       pir::Operation* curr_op = it.owner();
-      if (!IsValidOp(curr_op, x)) {
-        return 0;
+      if (!IsValidOp(curr_op, val)) {
+        continue;
       }
       if (op_example == nullptr) {
         op_example = curr_op;
-        op_example_name = curr_op->name();
-        op_example_num_operands = curr_op->num_operands();
-        op_example_attrs = curr_op->attributes();
         op_cnt_use_x++;
+        w_from_source_index = operands_source_index(curr_op, val) == 0 ? 1 : 0;
       } else {
-        if (curr_op->name() == op_example_name &&
-            curr_op->num_operands() == op_example_num_operands &&
-            AreAttributeMapsEqual(op_example_attrs, curr_op->attributes())) {
+        if (AreOperationsEqual(op_example, curr_op)) {
           op_cnt_use_x++;
-        } else {
-          return 0;
         }
       }
     }
     return op_cnt_use_x;
   }
-
   template <typename OpTy, typename... Args>
   pir::Value InsertOpAfter(pir::Operation** prev_op_ptr,
                            pir::PatternRewriter* rewriter,
@@ -146,15 +172,72 @@ class HorizontalFusePattern : public pir::RewritePattern {
     return new_op->result(0);
   }
 
+  template <typename OpTy, typename... Args>
+  pir::Value InsertOpWithOrder(pir::Operation** prev_op,
+                               pir::PatternRewriter* rewriter,
+                               int w_from_source_index,
+                               const pir::Value& val,
+                               const pir::Value& w_qkv,
+                               Args&&... args) const {
+    if (w_from_source_index == 0) {
+      return InsertOpAfter<OpTy>(
+          prev_op, rewriter, w_qkv, val, std::forward<Args>(args)...);
+    } else {
+      return InsertOpAfter<OpTy>(
+          prev_op, rewriter, val, w_qkv, std::forward<Args>(args)...);
+    }
+  }
+
+  pir::Value InsertFusedOp(const std::string& fused_matmul_op_name,
+                           pir::Operation** prev_op,
+                           pir::PatternRewriter* rewriter,
+                           const pir::Value& val,
+                           const pir::Value& w_qkv,
+                           const pir::Value& bias_qkv,
+                           const pir::AttributeMap& fused_matmul_op_attrs,
+                           int w_from_source_index) const {
+    if (fused_matmul_op_name == paddle::dialect::MatmulOp::name()) {
+      return InsertOpWithOrder<paddle::dialect::MatmulOp>(
+          prev_op,
+          rewriter,
+          w_from_source_index,
+          val,
+          w_qkv,
+          fused_matmul_op_attrs);
+    } else if (fused_matmul_op_name ==
+               paddle::dialect::GemmEpilogueOp::name()) {
+      return InsertOpWithOrder<paddle::dialect::GemmEpilogueOp>(
+          prev_op,
+          rewriter,
+          w_from_source_index,
+          val,
+          w_qkv,
+          bias_qkv,
+          fused_matmul_op_attrs);
+    } else {
+      return InsertOpWithOrder<paddle::dialect::FcOp>(prev_op,
+                                                      rewriter,
+                                                      w_from_source_index,
+                                                      val,
+                                                      w_qkv,
+                                                      bias_qkv,
+                                                      fused_matmul_op_attrs);
+    }
+  }
+
   void RewriteOpsbyValue(pir::Operation* op,
                          pir::PatternRewriter* rewriter,
                          size_t idx) const {
-    /// x is used by the op, which belongs to a kind of
+    /// val is used by the op, which belongs to a kind of
     /// MatmulOp/GemmEpilogueOp/FcOp
-    // prepare x
-    pir::Value x = op->result(idx);
-    std::vector<int64_t> x_dims = pir::GetShapeFromValue(x);
-    auto x_last_axis = x_dims.size() - 1;
+
+    // prepare val
+    pir::Value val = op->result(idx);
+    std::vector<int64_t> x_dims = pir::GetShapeFromValue(val);
+
+    int64_t x_last_axis =
+        w_from_source_index == 1 ? x_dims.size() - 1 : x_dims.size() - 2;
+
     // prepare src ops
     std::vector<pir::Operation*> fused_matmul_ops;
     std::string fused_matmul_op_name;
@@ -167,15 +250,25 @@ class HorizontalFusePattern : public pir::RewritePattern {
     // prepare outs
     std::vector<pir::Value> src_outs;
 
-    for (auto it = x.use_begin(); it != x.use_end(); ++it) {
+    for (auto it = val.use_begin(); it != val.use_end(); ++it) {
       pir::Operation* curr_op = it.owner();
-      if (!IsValidOp(curr_op, x)) {
+
+      if (!IsValidOp(curr_op, val)) {
         continue;
       }
       fused_matmul_ops.push_back(curr_op);
-      combine_op_inputs_weight.push_back(curr_op->operand_source(1));
-      auto w_dims = pir::GetShapeFromValue(curr_op->operand_source(1));
-      w_shapes.push_back(w_dims[w_dims.size() - 1]);
+      combine_op_inputs_weight.push_back(
+          curr_op->operand_source(w_from_source_index));
+      auto w_dims =
+          pir::GetShapeFromValue(curr_op->operand_source(w_from_source_index));
+
+      // input is X
+      if (w_from_source_index) {
+        w_shapes.push_back(w_dims[w_dims.size() - 1]);
+      } else {
+        // input is Y
+        w_shapes.push_back(w_dims[0]);
+      }
 
       src_outs.push_back(curr_op->result(0));
 
@@ -187,14 +280,14 @@ class HorizontalFusePattern : public pir::RewritePattern {
         combine_op_inputs_bias.push_back(curr_op->operand_source(2));
       }
     }
-
     /// build new graph
     /// insert new ops
     pir::Operation* prev_op = op;
     auto weight_combined = InsertOpAfter<pir::CombineOp>(
         &prev_op, rewriter, combine_op_inputs_weight);
+
     auto w_qkv = InsertOpAfter<paddle::dialect::ConcatOp>(
-        &prev_op, rewriter, weight_combined, 1);
+        &prev_op, rewriter, weight_combined, w_from_source_index);
 
     pir::Value bias_qkv;
     if (combine_op_inputs_bias.size()) {
@@ -204,38 +297,39 @@ class HorizontalFusePattern : public pir::RewritePattern {
           &prev_op, rewriter, bias_combined, 0);
     }
 
-    pir::Value xw_fused;
-    if (fused_matmul_op_name == paddle::dialect::MatmulOp::name()) {
-      xw_fused = InsertOpAfter<paddle::dialect::MatmulOp>(
-          &prev_op, rewriter, x, w_qkv, fused_matmul_op_attrs);
-    } else if (fused_matmul_op_name ==
-               paddle::dialect::GemmEpilogueOp::name()) {
-      xw_fused = InsertOpAfter<paddle::dialect::GemmEpilogueOp>(
-          &prev_op, rewriter, x, w_qkv, bias_qkv, fused_matmul_op_attrs);
-    } else {
-      xw_fused = InsertOpAfter<paddle::dialect::FcOp>(
-          &prev_op, rewriter, x, w_qkv, bias_qkv, fused_matmul_op_attrs);
-    }
+    pir::Value xw_fused = InsertFusedOp(fused_matmul_op_name,
+                                        &prev_op,
+                                        rewriter,
+                                        val,
+                                        w_qkv,
+                                        bias_qkv,
+                                        fused_matmul_op_attrs,
+                                        w_from_source_index);
 
-    auto xw_splitted = InsertOpAfter<paddle::dialect::SplitOp>(
+    pir::Value xw_splitted = InsertOpAfter<paddle::dialect::SplitOp>(
         &prev_op, rewriter, xw_fused, w_shapes, x_last_axis);
+
     rewriter->SetInsertionPointAfter(prev_op);
+
     auto split_builtin_op = rewriter->Build<pir::SplitOp>(xw_splitted);
+
     std::vector<pir::Value> xw = split_builtin_op.outputs();
 
-    /// replace res Value
+    // replace res Value
     for (size_t k = 0; k < src_outs.size(); k++) {
       rewriter->ReplaceAllUsesWith(src_outs[k], xw[k]);
     }
-    /// del old ops
-    for (auto fused_matmul_op : fused_matmul_ops)
+
+    // delete fused matmul op
+    for (auto fused_matmul_op : fused_matmul_ops) {
       rewriter->EraseOp(fused_matmul_op);
+    }
   }
 };
 
 class HorizontalFusePass : public pir::Pass {
  public:
-  HorizontalFusePass() : pir::Pass("horizontal_fuse_pass", 1) {}
+  HorizontalFusePass() : pir::Pass("horizontal_fuse_pass", 2) {}
 
   bool Initialize(pir::IrContext* context) override {
     pir::RewritePatternSet ps(context);
@@ -247,6 +341,7 @@ class HorizontalFusePass : public pir::Pass {
   void Run(pir::Operation* op) override {
     pir::GreedyRewriteConfig cfg;
     cfg.use_top_down_traversal = true;
+
     cfg.max_iterations = 10;
     auto [_, num_rewrites] = pir::ApplyPatternsGreedily(op, patterns_, cfg);
     AddStatistics(num_rewrites);