[XPU] Add bf16 Cumsum, update xhpc and bkcl (#68052)

cqulilujia · RuohengMa · web-flow · commit 6a86f64ab1f2 · 2024-09-06T17:57:11.000+08:00
* [PHI] add cumsum bf16

* update XHPC date and refine unittest

---------

Co-authored-by: maruoheng &lt;maruoheng@baidu.com&gt;
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,9 +30,9 @@ if(NOT DEFINED XPU_XRE_BASE_VERSION)
   set(XPU_XRE_BASE_VERSION "4.32.0.1")
 endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "eb35/20240904")
+  set(XPU_XHPC_BASE_DATE "eb35/20240906")
 endif()
-set(XPU_XCCL_BASE_VERSION "1.2.11b")
+set(XPU_XCCL_BASE_VERSION "1.2.11c")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
   set(XPU_XFT_BASE_VERSION "20230602")
 endif()
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -239,7 +239,8 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
                      phi::DataType::INT32,
-                     phi::DataType::INT64})},
+                     phi::DataType::INT64,
+                     phi::DataType::BFLOAT16})},
       {"cumsum_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -248,7 +248,8 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
                      phi::DataType::INT32,
-                     phi::DataType::INT64})},
+                     phi::DataType::INT64,
+                     phi::DataType::BFLOAT16})},
       {"cumsum_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/kernels/xpu/cum_kernel.cc b/paddle/phi/kernels/xpu/cum_kernel.cc
@@ -88,4 +88,5 @@ PD_REGISTER_KERNEL(cumsum,
                    float,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/test/xpu/test_cumsum_op_xpu.py b/test/xpu/test_cumsum_op_xpu.py
@@ -20,6 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -42,13 +43,19 @@ def set_case(self):
             self.op_type = 'cumsum'
             self.init_config()
 
-            self.data = np.random.uniform(-1.0, 1.0, self.input_shape).astype(
-                self.dtype
-            )
+            if self.dtype == np.uint16:
+                self.data = np.random.uniform(
+                    -1.0, 1.0, self.input_shape
+                ).astype('float32')
+                self.inputs = {'X': convert_float_to_uint16(self.data)}
+            else:
+                self.data = np.random.uniform(
+                    -1.0, 1.0, self.input_shape
+                ).astype(self.dtype)
+                self.inputs = {
+                    'X': self.data,
+                }
             reference_out = np.cumsum(self.data, axis=self.axis)
-            self.inputs = {
-                'X': self.data,
-            }
             self.attrs = {
                 'use_xpu': True,
                 'axis': self.axis,
@@ -89,6 +96,11 @@ def init_config(self):
             self.input_shape = [2, 2, 255]
             self.axis = -1
 
+    class XPUTestCumsum5(TestCumsumOPBase):
+        def init_config(self):
+            self.input_shape = [2, 3, 5, 255]
+            self.axis = 2
+
 
 support_types = get_xpu_op_support_types('cumsum')
 for stype in support_types:
