support argmax bf16

Author: skywalker2012

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

@@ -21,6 +21,7 @@
     get_xpu_op_support_types,
 )
 from op_test_xpu import XPUOpTest
+from op_test import convert_float_to_uint16
 
 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)}