[xpu] fix pir:fc_xpu_fuse_pass && add bfloat16 for op :abs (#69614)

TR666 · web-flow · commit 71172e3949d8 · 2024-11-24T22:21:02.000-08:00
* [xpu] fix pir:fc_xpu_fuse_pass &amp;&amp; add bfloat16 for op :abs

* [xpu] fix pir:fc_xpu_fuse_pass &amp;&amp; add bfloat16 for op :abs
diff --git a/paddle/fluid/pir/transforms/xpu/fc_xpu_fuse_pass.cc b/paddle/fluid/pir/transforms/xpu/fc_xpu_fuse_pass.cc
@@ -25,6 +25,32 @@
 #include "paddle/pir/include/pass/pass_registry.h"
 #include "paddle/pir/include/pattern_rewrite/pattern_match.h"
 
+/*
+fuse malmul + add to fc_xpu
+For example:
+graph:
+
+                  x       w
+                    \   /
+                      |
+                     mul
+                      |
+                      |
+           bias ---  add
+                      |
+                      |
+                    output
+------------------------------------------------------
+After the pass is applied:
+                   x      w
+                     \  /
+                      |
+            bias--- fc_xpu
+                      |
+                      |
+                    Output
+*/
+
 namespace {
 
 int ConvertActivationType(const std::string &act_type) {
@@ -59,126 +85,81 @@ int ConvertActivationType(const std::string &act_type) {
 
 class FCXpuFusePattern : public paddle::drr::DrrPatternBase {
  private:
-  bool w_transpose_;
-  bool with_bias_;
+  bool transpose_w_;
 
  public:
-  FCXpuFusePattern(bool w_transpose, bool with_bias)
-      : w_transpose_(w_transpose), with_bias_(with_bias) {}
+  explicit FCXpuFusePattern(bool transpose_w) : transpose_w_(transpose_w) {}
   std::string name() const override { return "FCXpuFusePattern"; }
 
   void operator()(paddle::drr::DrrPatternContext *ctx) const override {
     paddle::drr::SourcePattern pat = ctx->SourcePattern();
     const auto &mul = pat.Op(paddle::dialect::MatmulOp::name(),
                              {{"transpose_x", pat.Attr("transpose_x")},
                               {"transpose_y", pat.Attr("transpose_y")}});
-
     mul({&pat.Tensor("x"), &pat.Tensor("w")}, {&pat.Tensor("mul_out")});
-    if (with_bias_) {
-      const auto &add = pat.Op(paddle::dialect::AddOp::name());
-      add({&pat.Tensor("mul_out"), &pat.Tensor("bias")},
-          {&pat.Tensor("add_out")});
-    }
+
+    const auto &add = pat.Op(paddle::dialect::AddOp::name());
+    add({&pat.Tensor("mul_out"), &pat.Tensor("bias")},
+        {&pat.Tensor("add_out")});
+
+    // Constraints
     pat.AddConstraint([&](const paddle::drr::MatchContext &match_ctx) {
       auto w_shape = pir::GetShapeFromValue(match_ctx.Tensor("w"));
-      if (w_shape.size() != 2) {
+      auto x_shape = pir::GetShapeFromValue(match_ctx.Tensor("x"));
+      auto bias_shape = pir::GetShapeFromValue(match_ctx.Tensor("bias"));
+      if (transpose_w_ != match_ctx.Attr<bool>("transpose_y")) {
         return false;
       }
-      if (with_bias_) {
-        auto x_shape = pir::GetShapeFromValue(match_ctx.Tensor("x"));
-        auto bias_shape = pir::GetShapeFromValue(match_ctx.Tensor("bias"));
-        if (w_shape.back() != bias_shape.back()) {
-          return false;
-        }
-      }
-      if (w_transpose_ != match_ctx.Attr<bool>("transpose_y")) {
-        return false;
-      }
-      return true;
+      return (w_shape.size() == 2 && x_shape.size() >= 2 &&
+              bias_shape.size() == 1);
     });
+
+    // Result pattern
     paddle::drr::ResultPattern res = pat.ResultPattern();
-    // x shape
+
     const auto &in_num_col_dims_attr =
-        res.ComputeAttr([](const paddle::drr::MatchContext &match_ctx) -> int {
+        res.ComputeAttr([&](const paddle::drr::MatchContext &match_ctx) -> int {
           auto x_shape = pir::GetShapeFromValue(match_ctx.Tensor("x"));
           return x_shape.size() - 1;
         });
-    // reshape scale shape
-    const auto &expand_1_shape =
-        res.ComputeAttr([&](const paddle::drr::MatchContext &match_ctx)
-                            -> std::vector<int64_t> {
-          return {static_cast<int64_t>(
-              phi::backends::xpu::get_xpu_max_ptr_size(-1))};
-        });
-
-    const auto &out_dtype_attr = res.ComputeAttr(
-        [](const paddle::drr::MatchContext &match_ctx) -> phi::DataType {
-          auto x_dtype = pir::GetDataTypeFromValue(match_ctx.Tensor("x"));
-          if (x_dtype.isa<pir::Float32Type>()) {
-            return phi::DataType::FLOAT32;
-          } else {
-            return phi::DataType::UNDEFINED;
-          }
-        });
 
-    // find fp32 w_max
-    const auto &max_op1 =
-        res.Op(paddle::dialect::MaxOp::name(),
-               {{"axis", res.VectorInt64Attr(std::vector<int64_t>{})},
-                {"keepdim", res.BoolAttr(false)}});
-    const auto &abs_op = res.Op(paddle::dialect::AbsOp::name());
-    res.Tensor("w_abs") = abs_op(res.Tensor("w"));
-    res.Tensor("filter_fp_max") = max_op1(res.Tensor("w_abs"));
-    // set w_max shape
-    const auto &w_max_shape_attr = res.ComputeAttr(
-        [](const paddle::drr::MatchContext &match_ctx) -> std::vector<int64_t> {
-          return {1};
-        });
-    // full 32767 or 127
-    const auto &full1 = res.Op(paddle::dialect::FullOp::name(),
-                               {{"shape", w_max_shape_attr},
-                                {"value", res.Int32Attr(32767)},
-                                {"dtype", res.DataTypeAttr("float32")},
-                                {"place", res.PlaceAttr("cpu")}});
-    const auto &div = res.Op(paddle::dialect::DivideOp::name());
-    // scale1 = fp32_w_max/32767
-    res.Tensor("scale_fp32_max") = div(res.Tensor("filter_fp_max"), full1());
-    // int16_weight = w / scale1
-    res.Tensor("w_quant_tmp") =
-        div(res.Tensor("w"), res.Tensor("scale_fp32_max"));
-    // cast fp32->int16
-    const auto &cast_op = res.Op(paddle::dialect::CastOp::name(),
-                                 {{"dtype", res.DataTypeAttr("int16")}});
-    res.Tensor("w_quant") = cast_op(res.Tensor("w_quant_tmp"));
-    // expand to 4 or 6 pointer size
-    const auto &expand =
-        res.Op(paddle::dialect::ExpandOp::name(), {{"shape", expand_1_shape}});
-    res.Tensor("res_w_max") = expand(res.Tensor("filter_fp_max"));
-    const auto &transpose_y_attr =
-        res.ComputeAttr([](const paddle::drr::MatchContext &match_ctx) -> bool {
-          return match_ctx.Attr<bool>("transpose_y");
-        });
-    if (!w_transpose_) {
-      // perm
+    if (!transpose_w_) {
+      // prepare weight, transpose it if necessary
       const auto &perm_attr = res.ComputeAttr(
           [](const paddle::drr::MatchContext &match_ctx) -> std::vector<int> {
             auto w_shape = pir::GetShapeFromValue(match_ctx.Tensor("w"));
             if (w_shape.size() == 2) {
               return {1, 0};
-            } else if (w_shape.size() == 3) {
-              return {0, 2, 1};
-            } else if (w_shape.size() == 4) {
-              return {0, 1, 3, 2};
             } else {
               PADDLE_THROW(common::errors::Unimplemented(
                   "Not support convert w_shape.size()(%d).", w_shape.size()));
             }
           });
-      const auto &transpose_y =
+      const auto &transpose_op =
           res.Op(paddle::dialect::TransposeOp::name(), {{"perm", perm_attr}});
-      res.Tensor("w_quant_trans") = transpose_y(res.Tensor("w_quant"));
+      res.Tensor("w_trans") = transpose_op(res.Tensor("w"));
+      VLOG(3) << "transpose weight for fc_xpu op";
     }
 
+    const auto &out_dtype_attr = res.ComputeAttr(
+        [](const paddle::drr::MatchContext &match_ctx) -> phi::DataType {
+          auto x_dtype = pir::GetDataTypeFromValue(match_ctx.Tensor("x"));
+          // 目前仅支持以下几种非量化的情况
+          if (x_dtype.isa<pir::Float32Type>()) {
+            return phi::DataType::FLOAT32;
+          } else if (x_dtype.isa<pir::Float16Type>()) {
+            return phi::DataType::FLOAT16;
+          } else if (x_dtype.isa<pir::BFloat16Type>()) {
+            return phi::DataType::BFLOAT16;
+          } else {
+            return phi::DataType::UNDEFINED;
+          }
+        });
+    // only support float32 bias now
+    const auto &cast_op = res.Op(paddle::dialect::CastOp::name(),
+                                 {{"dtype", res.DataTypeAttr("float32")}});
+    res.Tensor("bias_fp32") = cast_op(res.Tensor("bias"));
+
     const auto &fc_xpu =
         res.Op(paddle::dialect::FcXpuOp::name(),
                {{
@@ -194,15 +175,13 @@ class FCXpuFusePattern : public paddle::drr::DrrPatternBase {
         {
             &res.Tensor("x"),
             &res.InputNoneTensor(),
-            w_transpose_ ? &res.Tensor("w_quant")
-                         : &res.Tensor("w_quant_trans"),
-            &res.Tensor("res_w_max"),
-            with_bias_ ? &res.Tensor("bias") : &res.InputNoneTensor(),
+            transpose_w_ ? &res.Tensor("w") : &res.Tensor("w_trans"),
+            &res.InputNoneTensor(),
+            &res.Tensor("bias_fp32"),
             &res.InputNoneTensor(),
             &res.InputNoneTensor(),
         },
-        {with_bias_ ? &res.Tensor("add_out") : &res.Tensor("mul_out"),
-         &res.Tensor("out_max")});
+        {&res.Tensor("add_out"), &res.Tensor("out_max")});
   }
 };
 
@@ -212,12 +191,8 @@ class FCXpuFusePass : public pir::PatternRewritePass {
 
   pir::RewritePatternSet InitializePatterns(pir::IrContext *context) override {
     pir::RewritePatternSet ps(context);
-    for (auto w_transpose : {true, false}) {
-      for (auto with_bias : {true, false}) {
-        ps.Add(paddle::drr::Create<FCXpuFusePattern>(
-            context, w_transpose, with_bias));
-      }
-    }
+    ps.Add(paddle::drr::Create<FCXpuFusePattern>(context, false));
+    ps.Add(paddle::drr::Create<FCXpuFusePattern>(context, true));
     return ps;
   }
 };
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -26,7 +26,10 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"add_layernorm_xpu",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
-      {"abs", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+      {"abs",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"abs_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"accuracy",
@@ -438,7 +441,9 @@ XPUOpMap& get_kl3_ops() {
       {"fast_layernorm_xpu",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"fc_xpu",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"fill",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
diff --git a/paddle/phi/kernels/fusion/xpu/fc_xpu_kernel.cc b/paddle/phi/kernels/fusion/xpu/fc_xpu_kernel.cc
@@ -28,7 +28,7 @@ void FcXPUKernelImpl(const Context& ctx,
                      const DenseTensor& x,
                      const paddle::optional<DenseTensor>& x_max,
                      const DenseTensor& w,
-                     const DenseTensor& w_max,
+                     const paddle::optional<DenseTensor>& w_max,
                      const paddle::optional<DenseTensor>& bias,
                      const paddle::optional<DenseTensor>& scale_max,
                      const paddle::optional<DenseTensor>& out_max_in,
@@ -51,7 +51,8 @@ void FcXPUKernelImpl(const Context& ctx,
   const float* x_max_data =
       x_max.get_ptr() == nullptr ? nullptr : x_max.get_ptr()->data<float>();
   auto* w_data = reinterpret_cast<const XPUTypeW*>(w.data<T_W>());
-  auto* w_max_data = w_max.data<float>();
+  const float* w_max_data =
+      w_max.get_ptr() == nullptr ? nullptr : w_max.get_ptr()->data<float>();
   const float* bias_data =
       bias.get_ptr() == nullptr ? nullptr : bias.get_ptr()->data<float>();
   auto* out_data =
@@ -126,7 +127,7 @@ void FcXPUKernel(const Context& ctx,
                  const DenseTensor& x,
                  const paddle::optional<DenseTensor>& x_max,
                  const DenseTensor& w,
-                 const DenseTensor& w_max,
+                 const paddle::optional<DenseTensor>& w_max,
                  const paddle::optional<DenseTensor>& bias,
                  const paddle::optional<DenseTensor>& scale_max,
                  const paddle::optional<DenseTensor>& out_max_in,
@@ -249,6 +250,32 @@ void FcXPUKernel(const Context& ctx,
     return;
   }
 
+  if (x.dtype() == DataType::BFLOAT16) {
+    // bfloat16 kernel
+    if (w.dtype() == DataType::BFLOAT16) {
+      if (out_dtype == DataType::BFLOAT16) {
+        FC_XPU_KERNEL_IMPL(phi::dtype::bfloat16,
+                           phi::dtype::bfloat16,
+                           phi::dtype::bfloat16,
+                           float);
+      } else {
+        PADDLE_THROW(common::errors::Unimplemented(
+            "Not support x_dtype is %s, w_dtype is %s and out_dtype is "
+            "%s.",
+            DataTypeToString(x.dtype()),
+            DataTypeToString(w.dtype()),
+            DataTypeToString(out_dtype)));
+      }
+    } else {
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Not support x_dtype is %s, w_dtype is %s and out_dtype is %s.",
+          DataTypeToString(x.dtype()),
+          DataTypeToString(w.dtype()),
+          DataTypeToString(out_dtype)));
+    }
+    return;
+  }
+
   PADDLE_THROW(common::errors::Unimplemented(
       "Not support x_dtype is %s, w_dtype is %s and out_dtype is %s.",
       DataTypeToString(x.dtype()),
@@ -265,4 +292,5 @@ PD_REGISTER_KERNEL(fc_xpu,
                    phi::fusion::FcXPUKernel,
                    float,
                    phi::dtype::float16,
-                   int8_t) {}
+                   int8_t,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/abs_kernel.cc b/paddle/phi/kernels/xpu/abs_kernel.cc
@@ -37,6 +37,7 @@ PD_REGISTER_KERNEL(abs,
                    phi::AbsKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int8_t,
                    int32_t,
                    int64_t) {}
diff --git a/paddle/phi/ops/yaml/fused_ops.yaml b/paddle/phi/ops/yaml/fused_ops.yaml
@@ -185,7 +185,7 @@
   kernel :
     func : fc_xpu
     data_type : x
-  optional : bias, x_max, scale_max, out_max_in
+  optional : bias, x_max, w_max, scale_max, out_max_in
 
 - op : fp8_fp8_half_gemm_fused
   args : (Tensor x, Tensor y, Tensor bias, bool transpose_x = false, bool transpose_y = false, float scale = 1.0f, str output_dtype = "float16", str activation_type = "identity")
