[XPU] support argmax bf16 (#69056)

skywalker2012 · web-flow · commit 4283d97414ab · 2024-11-06T14:58:19.000+08:00
diff --git a/cmake/external/xpu.cmake b/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)
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/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,
diff --git a/paddle/phi/kernels/xpu/arg_min_max_kernel.cc b/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);
 }
 
diff --git a/test/xpu/test_arg_max_op_xpu.py b/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)}
 

Original file line number	Diff line number	Diff line change
`@@ -193,7 +193,8 @@ PD_REGISTER_KERNEL(argmax,`
`193`	`193`	`float,`
`194`	`194`	`int,`
`195`	`195`	`int64_t,`
`196`		`- phi::dtype::float16) {`
	`196`	`+ phi::dtype::float16,`
	`197`	`+ phi::dtype::bfloat16) {`
`197`	`198`	`kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);`
`198`	`199`	`}`
`199`	`200`