[XPU] support argmin bf16 (#68453)

skywalker2012 · web-flow · commit 438e66b9046c · 2024-10-01T08:53:29.000+08:00
* support argmin bf16

* convert_float_to_uint16

* pre-commit modify
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -49,7 +49,9 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16})},
       {"arg_min",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"argsort_grad",
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
diff --git a/paddle/phi/kernels/xpu/arg_min_max_kernel.cc b/paddle/phi/kernels/xpu/arg_min_max_kernel.cc
@@ -196,7 +196,12 @@ PD_REGISTER_KERNEL(argmax,
   kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
 }
 
-PD_REGISTER_KERNEL(
-    argmin, XPU, ALL_LAYOUT, phi::ArgMinKernel, float, phi::dtype::float16) {
+PD_REGISTER_KERNEL(argmin,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ArgMinKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
 }
diff --git a/test/xpu/test_arg_min_op_xpu.py b/test/xpu/test_arg_min_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.argmin(self.x, axis=self.axis)}
 

Original file line number	Diff line number	Diff line change
`@@ -196,7 +196,12 @@ PD_REGISTER_KERNEL(argmax,`
`196`	`196`	`kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);`
`197`	`197`	`}`
`198`	`198`
`199`		`-PD_REGISTER_KERNEL(`
`200`		`- argmin, XPU, ALL_LAYOUT, phi::ArgMinKernel, float, phi::dtype::float16) {`
	`199`	`+PD_REGISTER_KERNEL(argmin,`
	`200`	`+ XPU,`
	`201`	`+ ALL_LAYOUT,`
	`202`	`+ phi::ArgMinKernel,`
	`203`	`+ float,`
	`204`	`+ phi::dtype::float16,`
	`205`	`+ phi::dtype::bfloat16) {`
`201`	`206`	`kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);`
`202`	`207`	`}`