[XPU] support bf16 reduce_mean, reduce_sum&grad (#59166)

runzhech · web-flow · commit f2ae4249aa9d · 2023-11-22T10:41:08.000+08:00
* support bf16 reduce_mean, reduce_sum&amp;grad  on kl3

* support bf16 reduce_mean, reduce_sum&amp;grad  on kl3

* refresh the CI

* refresh CI
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -24,7 +24,7 @@ set(XPU_XFT_LIB_NAME "libxft.so")
 set(XPU_XPTI_LIB_NAME "libxpti.so")
 
 if(NOT DEFINED XPU_BASE_DATE)
-  set(XPU_BASE_DATE "20231103")
+  set(XPU_BASE_DATE "20231114")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.1.6.1")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -670,20 +670,26 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::INT64})},
       {"reduce_mean_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"reduce_mean",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"reduce_min_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"reduce_min", XPUKernelSet({phi::DataType::FLOAT32})},
       {"reduce_prod",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
-      {"reduce_sum_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"reduce_sum_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"reduce_sum",
        XPUKernelSet({phi::DataType::FLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::INT8,
-                     phi::DataType::FLOAT32})},
+                     phi::DataType::FLOAT32,
+                     phi::DataType::BFLOAT16})},
       {"relu6", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"relu6_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"relu_grad",
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -633,20 +633,26 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT64})},
       {"reduce_mean_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"reduce_mean",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"reduce_min_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"reduce_min", XPUKernelSet({phi::DataType::FLOAT32})},
       {"reduce_prod",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
-      {"reduce_sum_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"reduce_sum_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"reduce_sum",
        XPUKernelSet({phi::DataType::FLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::INT8,
-                     phi::DataType::FLOAT32})},
+                     phi::DataType::FLOAT32,
+                     phi::DataType::BFLOAT16})},
       {"relu6", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"relu6_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"relu_grad",
diff --git a/paddle/phi/core/visit_type.h b/paddle/phi/core/visit_type.h
@@ -437,6 +437,8 @@ namespace phi {
       PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::INT64, int64_t, __VA_ARGS__) \
       PD_PRIVATE_CASE_TYPE(                                                    \
           NAME, ::phi::DataType::FLOAT16, phi::float16, __VA_ARGS__)           \
+      PD_PRIVATE_CASE_TYPE(                                                    \
+          NAME, ::phi::DataType::BFLOAT16, phi::bfloat16, __VA_ARGS__)         \
       PD_PRIVATE_CASE_TYPE(NAME, ::phi::DataType::FLOAT32, float, __VA_ARGS__) \
       default:                                                                 \
         PADDLE_THROW(phi::errors::InvalidArgument(                             \
diff --git a/paddle/phi/kernels/cpu/reduce_sum_grad_kernel.cc b/paddle/phi/kernels/cpu/reduce_sum_grad_kernel.cc
@@ -51,6 +51,7 @@ PD_REGISTER_KERNEL(sum_grad,
                    float,
                    double,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int16_t,
                    int,
                    int64_t,
diff --git a/paddle/phi/kernels/reduce_mean_kernel.cc b/paddle/phi/kernels/reduce_mean_kernel.cc
@@ -69,6 +69,11 @@ PD_REGISTER_KERNEL(
 #endif
 
 #if defined(PADDLE_WITH_XPU)
-PD_REGISTER_KERNEL(
-    mean, XPU, ALL_LAYOUT, phi::MeanKernel, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(mean,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::MeanKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 #endif
diff --git a/paddle/phi/kernels/reduce_sum_kernel.cc b/paddle/phi/kernels/reduce_sum_kernel.cc
@@ -94,6 +94,7 @@ PD_REGISTER_KERNEL(sum,
                    phi::SumKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int8_t,
                    int,
                    int64_t) {
diff --git a/paddle/phi/kernels/xpu/reduce_mean_kernel.cc b/paddle/phi/kernels/xpu/reduce_mean_kernel.cc
@@ -59,6 +59,10 @@ void MeanRawKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    mean_raw, XPU, ALL_LAYOUT, phi::MeanRawKernel, float, phi::dtype::float16) {
-}
+PD_REGISTER_KERNEL(mean_raw,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::MeanRawKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/reduce_sum_grad_kernel.cc b/paddle/phi/kernels/xpu/reduce_sum_grad_kernel.cc
@@ -30,9 +30,9 @@ void ReduceSumGradKernel(const Context& dev_ctx,
   using XPUType = typename XPUTypeTrait<T>::Type;
   reduce_all = recompute_reduce_all(x, dims_arr, reduce_all);
   auto dims = dims_arr.GetData();
-  dev_ctx.template Alloc<XPUType>(x_grad);
-  const auto* out_data = out_grad.data<XPUType>();
-  auto* x_grad_data = x_grad->data<XPUType>();
+  dev_ctx.template Alloc(x_grad, x.dtype());
+  auto* out_data = reinterpret_cast<const XPUType*>(out_grad.data());
+  auto* x_grad_data = reinterpret_cast<XPUType*>(x_grad->data());
   const auto& input_dim_size = x.dims().size();
   std::vector<int> true_dims;
   for (size_t i = 0; i < dims.size(); ++i) {
@@ -70,6 +70,12 @@ void ReduceSumGradKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(sum_grad, XPU, ALL_LAYOUT, phi::ReduceSumGradKernel, float) {
+PD_REGISTER_KERNEL(sum_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ReduceSumGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
 }
diff --git a/paddle/phi/kernels/xpu/reduce_sum_kernel.cc b/paddle/phi/kernels/xpu/reduce_sum_kernel.cc
@@ -43,6 +43,7 @@ PD_REGISTER_KERNEL(sum_raw,
                    phi::SumRawKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int8_t,
                    int,
                    int64_t) {
diff --git a/test/xpu/test_reduce_sum_op_xpu.py b/test/xpu/test_reduce_sum_op_xpu.py
@@ -34,6 +34,7 @@ def __init__(self):
     class XPUTestReduceSumBase(XPUOpTest):
         def setUp(self):
             self.place = paddle.XPUPlace(0)
+            self.dtype = self.in_type
             self.init_case()
             self.set_case()
 
@@ -54,13 +55,14 @@ def set_case(self):
                         axis=self.axis, keepdims=self.attrs['keep_dim']
                     )
                 }
+            if self.dtype == np.uint16:
+                self.outputs['Out'] = self.outputs['Out'].astype(np.uint16)
 
         def init_case(self):
             self.shape = (5, 6, 10)
             self.axis = (0,)
             self.reduce_all = False
             self.keep_dim = False
-            self.dtype = self.in_type
 
         def test_check_output(self):
             self.check_output_with_place(self.place)
