[XPU] add some bf16 ops for kl3 (#59263)

houj04 · web-flow · commit 723f2d3ce9b4 · 2023-11-23T11:02:40.000+08:00
diff --git a/paddle/fluid/imperative/gradient_accumulator.cc b/paddle/fluid/imperative/gradient_accumulator.cc
@@ -275,6 +275,9 @@ void TensorAdd(const VarType& src, VarType* dst) {
       XPUTensorAddFunctor<platform::float16>(place, src_tensor, dst_tensor);
     } else if (data_type == framework::DataTypeTrait<double>::DataType()) {
       XPUTensorAddFunctor<double>(place, src_tensor, dst_tensor);
+    } else if (data_type ==
+               framework::DataTypeTrait<platform::bfloat16>::DataType()) {
+      XPUTensorAddFunctor<platform::bfloat16>(place, src_tensor, dst_tensor);
     } else {
       PADDLE_THROW(platform::errors::Unimplemented(
           "Gradient accumulation of data type (%s) on place (%s) is not "
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -31,7 +31,10 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"accuracy", XPUKernelSet({phi::DataType::FLOAT32})},
       {"adadelta", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"adamw", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+      {"adamw",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"adam", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"adam_dense_param_sparse_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
@@ -176,10 +179,13 @@ XPUOpMap& get_kl3_ops() {
       {"coalesce_tensor",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"concat_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"concat",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT64,
                      phi::DataType::BOOL,
                      phi::DataType::INT8,
@@ -243,10 +249,13 @@ XPUOpMap& get_kl3_ops() {
       {"einsum", XPUKernelSet({phi::DataType::FLOAT32})},
       {"einsum_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"elementwise_add_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"elementwise_add",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"elementwise_div_grad",
@@ -271,6 +280,7 @@ XPUOpMap& get_kl3_ops() {
       {"elementwise_mul",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
       {"elementwise_pow",
@@ -348,6 +358,7 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT32,
                      phi::DataType::BOOL,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT32})},
       {"fill_diagonal_tensor",
        XPUKernelSet({phi::DataType::INT64,
@@ -363,7 +374,8 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::UINT8,
                      phi::DataType::BOOL,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"flatten2_grad",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
@@ -424,7 +436,8 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"gather",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
@@ -531,7 +544,9 @@ XPUOpMap& get_kl3_ops() {
       {"logical_xor", XPUKernelSet({phi::DataType::BOOL})},
       {"lookup_table_v2_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"lookup_table_v2",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"masked_select",
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
@@ -670,13 +685,15 @@ XPUOpMap& get_kl3_ops() {
       {"reshape2_grad",
        XPUKernelSet({phi::DataType::FLOAT64,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::BOOL,
                      phi::DataType::FLOAT32})},
       {"reshape2",
        XPUKernelSet({phi::DataType::FLOAT64,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::BOOL,
@@ -702,6 +719,7 @@ XPUOpMap& get_kl3_ops() {
       {"scale",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"scatter",
@@ -917,7 +935,8 @@ XPUOpMap& get_kl3_ops() {
       {"triu",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"tril_triu_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/backends/xpu/xpu_info.cc b/paddle/phi/backends/xpu/xpu_info.cc
@@ -219,7 +219,7 @@ int get_xpu_max_ptr_size(int dev_id) {
       break;
     default:
       PADDLE_THROW(phi::errors::InvalidArgument(
-          "Only support get max ptr size of XPU1 or XPU2."));
+          "Only support get max ptr size of XPU1, XPU2 or XPU3."));
       break;
   }
   return max_ptr_size;
diff --git a/paddle/phi/kernels/xpu/adamw_kernel.cc b/paddle/phi/kernels/xpu/adamw_kernel.cc
@@ -31,16 +31,17 @@ float GetAbsMax(const Context& dev_ctx,
                 const float* input,
                 float* buffer_xpu,
                 int64_t numel) {
-  float buffer_cpu[6];
+  int max_ptr_size = phi::backends::xpu::get_xpu_max_ptr_size(-1);
+  float buffer_cpu[12];  // 12 is enough even for XPU3
   // int findmax(Context* ctx, const T* x, float* maxptr, int64_t len);
   int r = xpu::findmax<float>(dev_ctx.x_context(), input, buffer_xpu, numel);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "findmax");
   memory_utils::Copy(CPUPlace(),
                      static_cast<void*>(buffer_cpu),
                      dev_ctx.GetPlace(),
                      static_cast<void*>(buffer_xpu),
-                     sizeof(float) * 6);
-  float* max_value = std::max_element(buffer_cpu, buffer_cpu + 6);
+                     sizeof(float) * max_ptr_size);
+  float* max_value = std::max_element(buffer_cpu, buffer_cpu + max_ptr_size);
   return *max_value;
 }
 
@@ -357,8 +358,13 @@ void AdamwDenseKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    adamw, XPU, ALL_LAYOUT, phi::AdamwDenseKernel, float, phi::dtype::float16) {
+PD_REGISTER_KERNEL(adamw,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::AdamwDenseKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   // Skip beta1_pow, beta2_pow, skip_update data transform
   kernel->InputAt(5).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
diff --git a/paddle/phi/kernels/xpu/concat_grad_kernel.cc b/paddle/phi/kernels/xpu/concat_grad_kernel.cc
@@ -102,4 +102,5 @@ PD_REGISTER_KERNEL(concat_grad,
                    ALL_LAYOUT,
                    phi::ConcatGradKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/concat_kernel.cc b/paddle/phi/kernels/xpu/concat_kernel.cc
@@ -118,6 +118,7 @@ PD_REGISTER_KERNEL(concat,
                    float,
                    double,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int64_t,
                    int,
                    int8_t,
diff --git a/paddle/phi/kernels/xpu/elementwise_add_grad_kernel.cc b/paddle/phi/kernels/xpu/elementwise_add_grad_kernel.cc
@@ -99,6 +99,10 @@ void AddGradKernel(const Context& dev_ctx,
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    add_grad, XPU, ALL_LAYOUT, phi::AddGradKernel, phi::dtype::float16, float) {
-}
+PD_REGISTER_KERNEL(add_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::AddGradKernel,
+                   phi::dtype::float16,
+                   float,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/elementwise_add_kernel.cc b/paddle/phi/kernels/xpu/elementwise_add_kernel.cc
@@ -80,6 +80,7 @@ PD_REGISTER_KERNEL(add,
                    ALL_LAYOUT,
                    phi::AddKernel,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    float,
                    int,
                    int64_t) {}
diff --git a/paddle/phi/kernels/xpu/elementwise_multiply_kernel.cc b/paddle/phi/kernels/xpu/elementwise_multiply_kernel.cc
@@ -49,6 +49,7 @@ PD_REGISTER_KERNEL(multiply,
                    ALL_LAYOUT,
                    phi::MultiplyKernel,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    float,
                    int,
                    int64_t) {}
diff --git a/paddle/phi/kernels/xpu/embedding_kernel.cc b/paddle/phi/kernels/xpu/embedding_kernel.cc
@@ -117,4 +117,5 @@ PD_REGISTER_KERNEL(embedding,
                    ALL_LAYOUT,
                    phi::EmbeddingKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/full_kernel.cc b/paddle/phi/kernels/xpu/full_kernel.cc
@@ -126,7 +126,8 @@ PD_REGISTER_KERNEL(full,
                    int,
                    int64_t,
                    bool,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(full_like,
                    XPU,
@@ -138,7 +139,8 @@ PD_REGISTER_KERNEL(full_like,
                    int,
                    int64_t,
                    bool,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
 }
 
diff --git a/paddle/phi/kernels/xpu/gather_nd_kernel.cc b/paddle/phi/kernels/xpu/gather_nd_kernel.cc
@@ -137,4 +137,5 @@ PD_REGISTER_KERNEL(gather_nd,
                    float,
                    int64_t,
                    int,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/scale_kernel.cc b/paddle/phi/kernels/xpu/scale_kernel.cc
@@ -57,5 +57,6 @@ PD_REGISTER_KERNEL(scale,
                    phi::ScaleKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {}
diff --git a/paddle/phi/kernels/xpu/tril_triu_kernel.cc b/paddle/phi/kernels/xpu/tril_triu_kernel.cc
@@ -73,5 +73,11 @@ PD_REGISTER_KERNEL(tril_triu,
                    phi::dtype::float16) {}
 PD_REGISTER_KERNEL(
     tril, XPU, ALL_LAYOUT, phi::TrilKernel, int, float, phi::dtype::float16) {}
-PD_REGISTER_KERNEL(
-    triu, XPU, ALL_LAYOUT, phi::TriuKernel, int, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(triu,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::TriuKernel,
+                   int,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}

Original file line number	Diff line number	Diff line change
`@@ -219,7 +219,7 @@ int get_xpu_max_ptr_size(int dev_id) {`
`219`	`219`	`break;`
`220`	`220`	`default:`
`221`	`221`	`PADDLE_THROW(phi::errors::InvalidArgument(`
`222`		`- "Only support get max ptr size of XPU1 or XPU2."));`
	`222`	`+ "Only support get max ptr size of XPU1, XPU2 or XPU3."));`
`223`	`223`	`break;`
`224`	`224`	`}`
`225`	`225`	`return max_ptr_size;`