[PIR] Refine fuse linear param grad add pass (#61853)

zhangbo9674 · web-flow · commit 74a8c5c1c4c2 · 2024-02-22T11:17:03.000+08:00
* fix

* add ut
diff --git a/paddle/fluid/pir/transforms/fusion/fused_linear_param_grad_add_pass.cc b/paddle/fluid/pir/transforms/fusion/fused_linear_param_grad_add_pass.cc
@@ -115,8 +115,9 @@ class FusedMatmulGradAddPattern : public paddle::drr::DrrPatternBase {
       auto weight_grad_dims =
           pir::GetShapeFromValue(match_ctx.Tensor("weight_grad"));
       auto dweight_dims = pir::GetShapeFromValue(match_ctx.Tensor("dweight"));
+      auto weight_grad_use_count = match_ctx.Tensor("weight_grad").use_count();
       return (weight_grad_dims == dweight_dims && x_trans == false &&
-              y_trans == false);
+              y_trans == false && weight_grad_use_count == 1);
     });
 
     paddle::drr::ResultPattern res = pat.ResultPattern();
@@ -146,6 +147,73 @@ class FusedMatmulGradAddPattern : public paddle::drr::DrrPatternBase {
   }
 };
 
+// matmul + reshape + reshape + matmul + reshape + add_ -> matmul +
+// fused_liner_param_gard_add
+class FusedMatmulReshapeMatmulAddPattern : public paddle::drr::DrrPatternBase {
+ public:
+  std::string name() const override {
+    return "FusedMatmulReshapeMatmulAddPattern";
+  }
+  void operator()(paddle::drr::DrrPatternContext *ctx) const override {
+    paddle::drr::SourcePattern pat = ctx->SourcePattern();
+    const auto &full_int_array1 =
+        pat.Op(paddle::dialect::FullIntArrayOp::name());
+    const auto &reshape1 = pat.Op(paddle::dialect::ReshapeOp::name());
+    reshape1({&pat.Tensor("x"), &full_int_array1()},
+             {&pat.Tensor("reshape_x"), &pat.Tensor("reshape_x_xshape")});
+
+    const auto &full_int_array2 =
+        pat.Op(paddle::dialect::FullIntArrayOp::name());
+    const auto &reshape2 = pat.Op(paddle::dialect::ReshapeOp::name());
+    reshape2({&pat.Tensor("dy"), &full_int_array2()},
+             {&pat.Tensor("reshape_dy"), &pat.Tensor("reshape_dy_xshape")});
+
+    const auto &matmul = pat.Op(paddle::dialect::MatmulOp::name(),
+                                {{"transpose_x", pat.Attr("trans_x")},
+                                 {"transpose_y", pat.Attr("trans_y")}});
+    pat.Tensor("matmul_out") =
+        matmul(pat.Tensor("reshape_x"), pat.Tensor("reshape_dy"));
+
+    const auto &full_int_array3 =
+        pat.Op(paddle::dialect::FullIntArrayOp::name());
+    const auto &reshape3 = pat.Op(paddle::dialect::ReshapeOp::name());
+    reshape3({&pat.Tensor("matmul_out"), &full_int_array3()},
+             {&pat.Tensor("w_grad"), &pat.Tensor("w_grad_xshape")});
+
+    const auto &add_ = pat.Op(paddle::dialect::Add_Op::name());
+    pat.Tensor("dweight_inplace") =
+        add_(pat.Tensor("dweight"), pat.Tensor("w_grad"));
+
+    pat.RequireNativeCall([&](const paddle::drr::MatchContext &match_ctx) {
+      const auto &x_trans = match_ctx.Attr<bool>("trans_x");
+      const auto &y_trans = match_ctx.Attr<bool>("trans_y");
+      auto w_grad_dims = pir::GetShapeFromValue(match_ctx.Tensor("w_grad"));
+      auto dweight_dims = pir::GetShapeFromValue(match_ctx.Tensor("dweight"));
+      return (w_grad_dims == dweight_dims && x_trans == true &&
+              y_trans == false);
+    });
+
+    paddle::drr::ResultPattern res = pat.ResultPattern();
+    const auto &muti_precision_attr =
+        res.ComputeAttr([](const paddle::drr::MatchContext &match_ctx) -> bool {
+          return !(pir::GetDataTypeFromValue(match_ctx.Tensor("dweight")) ==
+                   pir::GetDataTypeFromValue(match_ctx.Tensor("w_grad")));
+        });
+
+    const auto &fused_linear_param_grad_add =
+        res.Op(paddle::dialect::FusedLinearParamGradAddOp::name(),
+               {{{"multi_precision", muti_precision_attr},
+                 {"has_bias", res.BoolAttr(false)}}});
+
+    fused_linear_param_grad_add(
+        {&res.Tensor("x"),
+         &res.Tensor("dy"),
+         &res.Tensor("dweight"),
+         &res.InputNoneTensor()},
+        {&res.Tensor("dweight_inplace"), &res.Tensor("dbias_out")});
+  }
+};
+
 // matmul + 0 = add_(0,1) -> fused_liner_param_gard_add
 class FusedMatmulAddaPattern : public paddle::drr::DrrPatternBase {
  public:
@@ -362,6 +430,7 @@ class FusedLinearParamGradAddPass : public pir::PatternRewritePass {
     ps.Add(paddle::drr::Create<FusedMatmulAddbPattern>(context));
     ps.Add(paddle::drr::Create<FusedMatmulAddGradAddaPattern>(context));
     ps.Add(paddle::drr::Create<FusedMatmulAddGradAddbPattern>(context));
+    ps.Add(paddle::drr::Create<FusedMatmulReshapeMatmulAddPattern>(context));
 
     return ps;
   }
diff --git a/test/cpp/pir/pattern_rewrite/drr_fuse_linear_param_grad_add_test.cc b/test/cpp/pir/pattern_rewrite/drr_fuse_linear_param_grad_add_test.cc
@@ -154,6 +154,32 @@ void BuildProgram3(pir::Builder &builder) {  // NOLINT
   builder.Build<paddle::dialect::FetchOp>(matmul_op2.out(), "dx", 3);
 }
 
+void BuildProgram4(pir::Builder &builder) {  // NOLINT
+  paddle::dialect::FullOp input_x = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{2, 2048, 5120}, 1.5, phi::DataType::FLOAT16);
+  paddle::dialect::FullOp input_dy = builder.Build<paddle::dialect::FullOp>(
+      std::vector<int64_t>{2, 2048, 6912}, 2.5, phi::DataType::FLOAT16);
+  paddle::dialect::FullOp input_weight_grad =
+      builder.Build<paddle::dialect::FullOp>(
+          std::vector<int64_t>{5120, 6912}, 1.0, phi::DataType::FLOAT16);
+
+  paddle::dialect::ReshapeOp reshape_x =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          input_x.out(), std::vector<int64_t>{4096, 5120});
+  paddle::dialect::ReshapeOp reshape_dy =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          input_dy.out(), std::vector<int64_t>{4096, 6912});
+  paddle::dialect::MatmulOp matmul_op =
+      builder.Build<paddle::dialect::MatmulOp>(
+          reshape_x.out(), reshape_dy.out(), true, false);
+  paddle::dialect::ReshapeOp output_dw =
+      builder.Build<paddle::dialect::ReshapeOp>(
+          matmul_op.out(), std::vector<int64_t>{5120, 6912});
+  paddle::dialect::Add_Op add__op1 = builder.Build<paddle::dialect::Add_Op>(
+      input_weight_grad.out(), output_dw.out());
+  builder.Build<paddle::dialect::FetchOp>(add__op1.out(), "dw", 0);
+}
+
 bool verify_pass(const pir::Program &program) {
   for (auto &op : *(program.block())) {
     if (op.name() == paddle::dialect::FusedLinearParamGradAddOp::name()) {
@@ -238,3 +264,23 @@ TEST(DrrTest, FusedLinearParamGradAdd3) {
   CHECK_EQ(pm.Run(&program), true);
   EXPECT_EQ(verify_pass(program), true);
 }
+
+TEST(DrrTest, FusedMatmulReshapeMatmulAddPattern) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+  ctx->GetOrRegisterDialect<pir::BuiltinDialect>();
+  pir::Program program(ctx);
+  pir::Builder builder = pir::Builder(ctx, program.block());
+  BuildProgram4(builder);
+
+  EXPECT_EQ(program.block()->size(), 12u);
+
+  pir::PassManager pm(ctx);
+  pm.AddPass(pir::CreateFusedLinearParamGradAddPass());
+  pm.EnablePassTiming();
+  pm.EnableIRPrinting();
+
+  CHECK_EQ(pm.Run(&program), true);
+  EXPECT_EQ(verify_pass(program), true);
+  EXPECT_EQ(program.block()->size(), 5u);
+}
