[XPU] bfloat16 support for round op (#71228)

houj04 · web-flow · commit 249bb0414a18 · 2025-02-24T18:21:09.000+08:00
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,7 +30,7 @@ set(XPU_XFA_LIB_NAME "libxpu_flash_attention.so")
 set(XPU_XPUDNN_LIB_NAME "libxpu_dnn.so")
 
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "dev/20250213")
+  set(XPU_XHPC_BASE_DATE "dev/20250220")
 endif()
 set(XPU_XCCL_BASE_VERSION "3.0.2.3") # For XRE5
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -515,7 +515,10 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16})},
-      {"round", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+      {"round",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"expand_as_v2",
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
diff --git a/paddle/phi/kernels/xpu/activation_kernel.cc b/paddle/phi/kernels/xpu/activation_kernel.cc
@@ -737,8 +737,13 @@ PD_REGISTER_KERNEL(exp,
                    phi::dtype::float16,
                    phi::dtype::bfloat16) {}
 
-PD_REGISTER_KERNEL(
-    round, XPU, ALL_LAYOUT, phi::RoundKernel, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(round,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::RoundKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(
     tan, XPU, ALL_LAYOUT, phi::TanKernel, float, phi::dtype::float16) {}
diff --git a/test/xpu/test_activation_op_xpu.py b/test/xpu/test_activation_op_xpu.py
@@ -109,18 +109,23 @@ class XPUTestRound(TestActivationOPBase):
         def set_case(self):
             self.op_type = 'round'
 
-            self.init_dtype()
+            self.dtype = self.in_type
             self.set_shape()
             self.set_decimals()
 
             np.random.seed(1024)
-            x = np.random.uniform(-1, 1, self.shape).astype(self.dtype) * 100
+            x = np.random.uniform(-100, 100, self.shape)
+
+            if self.dtype == np.uint16:
+                # bfloat16 actually
+                new_x = convert_float_to_uint16(x)
+            else:
+                new_x = x.astype(self.dtype)
             out = np.round(x, decimals=self.decimals)
 
-            self.inputs = {'X': OpTest.np_dtype_to_base_dtype(x)}
+            self.inputs = {'X': OpTest.np_dtype_to_base_dtype(new_x)}
             self.outputs = {'Out': out}
             self.attrs = {'decimals': self.decimals}
-            self.convert_input_output()
 
         def set_shape(self):
             self.shape = [10, 12]
@@ -131,36 +136,43 @@ def set_decimals(self):
         def test_check_grad(self):
             pass
 
-        def convert_input_output(self):
-            if self.dtype == np.uint16:
-                self.inputs = {'X': convert_float_to_uint16(self.inputs['X'])}
-                self.outputs = {
-                    'Out': convert_float_to_uint16(self.outputs['Out'])
-                }
-
     class XPUTestRound_ZeroDIm(XPUTestRound):
         def set_shape(self):
             self.shape = []
 
     class XPUTestRound_decimals1(XPUTestRound):
-        def init_decimals(self):
+        def set_decimals(self):
             self.decimals = 2
 
         def test_round_api(self):
             with dynamic_guard():
-                x_np = (
-                    np.random.uniform(-1, 1, self.shape).astype(self.dtype)
-                    * 100
+                if self.dtype != np.float32:
+                    # no float16 and bfloat16 on cpu
+                    return
+
+                np.random.seed(1024)
+                x_np = np.random.uniform(-100, 100, [10, 12]).astype(self.dtype)
+
+                x_paddle = paddle.to_tensor(x_np, place=paddle.XPUPlace(0))
+                x_paddle_cpu = paddle.to_tensor(x_np, place=paddle.CPUPlace())
+
+                # round using xpu
+                y = paddle.round(x_paddle, self.decimals)
+                # round using cpu
+                y_cpu = paddle.round(x_paddle_cpu, self.decimals)
+                # round using numpy
+                numpy_result = np.round(x_np, decimals=self.decimals)
+
+                # compare
+                np.testing.assert_allclose(
+                    y.numpy(), y_cpu.numpy(), atol=0, rtol=0
                 )
-                out_expect = np.round(x_np, decimals=self.decimals)
-                x_paddle = paddle.to_tensor(
-                    x_np, dtype=self.dtype, place=self.place
+                np.testing.assert_allclose(
+                    y.numpy(), numpy_result, atol=0, rtol=1e-7
                 )
-                y = paddle.round(x_paddle, decimals=self.decimals)
-                np.testing.assert_allclose(y.numpy(), out_expect, rtol=1e-3)
 
     class TestRound_decimals2(XPUTestRound_decimals1):
-        def init_decimals(self):
+        def set_decimals(self):
             self.decimals = -1
 
 
