support argmax bf16

Author: skywalker2012

cmake/external/xpu.cmake

@@ -31,7 +31,7 @@ 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/20241024")
+  set(XPU_XHPC_BASE_DATE "eb35/20241104")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.2.11e")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)

paddle/phi/backends/xpu/xpu3_op_list.cc

@@ -58,7 +58,8 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"arg_min",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,

paddle/phi/kernels/xpu/arg_min_max_kernel.cc

@@ -193,7 +193,8 @@ PD_REGISTER_KERNEL(argmax,
                    float,
                    int,
                    int64_t,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
 }
 

test/xpu/test_arg_max_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
@@ -41,8 +42,17 @@ def setUp(self):
             self.dtype = self.in_type
             self.initTestCase()
 
-            self.x = (np.random.random(self.dims)).astype(self.dtype)
-            self.inputs = {'X': self.x}
+            self.x = (np.random.random(self.dims)).astype(
+                self.dtype if self.dtype != np.uint16 else np.float32
+            )
+
+            self.inputs = {
+                'X': (
+                    self.x
+                    if self.dtype != np.uint16
+                    else convert_float_to_uint16(self.x)
+                )
+            }
             self.attrs = {'axis': self.axis, 'use_xpu': True}
             self.outputs = {'Out': np.argmax(self.x, axis=self.axis)}