No.55 complete argsort FP16 test, add argsort BF16 support and test

paddle/phi/kernels/gpu/argsort_grad_kernel.cu

@@ -222,4 +222,5 @@ PD_REGISTER_KERNEL(argsort_grad,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}

paddle/phi/kernels/gpu/argsort_kernel.cu

@@ -40,6 +40,10 @@ namespace detail {
 template <>
 struct radix_key_codec_base<phi::dtype::float16>
     : radix_key_codec_integral<phi::dtype::float16, uint16_t> {};
+
+template <>
+struct radix_key_codec_base<phi::dtype::bfloat16>
+    : radix_key_codec_integral<phi::dtype::bfloat16, uint16_t> {};
 }  // namespace detail
 }  // namespace rocprim
 #else
@@ -48,6 +52,11 @@ namespace cub {
 template <>
 struct NumericTraits<phi::dtype::float16>
     : BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::dtype::float16> {};
+
+template <>
+struct NumericTraits<phi::dtype::bfloat16>
+    : BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::dtype::bfloat16> {
+};
 }  // namespace cub
 #endif
 
@@ -519,6 +528,7 @@ PD_REGISTER_KERNEL(argsort,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
 }

python/paddle/fluid/tests/unittests/test_argsort_op.py

@@ -15,6 +15,7 @@
 import unittest
 
 import numpy as np
+from eager_op_test import OpTest, convert_float_to_uint16
 
 import paddle
 from paddle import fluid
@@ -392,6 +393,7 @@ def setUp(self):
         self.data = np.random.rand(*self.input_shape)
 
     def test_api(self):
+        paddle.enable_static()
         with fluid.program_guard(fluid.Program()):
             input = paddle.static.data(
                 name="input", shape=self.input_shape, dtype="float64"
@@ -513,5 +515,51 @@ def test_fp16(self):
                 out = exe.run(feed={'x': x_np}, fetch_list=[out])
 
 
+@unittest.skipIf(
+    not core.is_compiled_with_cuda()
+    or not core.is_bfloat16_supported(core.CUDAPlace(0)),
+    "core is not compiled with CUDA and not support the bfloat16",
+)
+class TestArgsortBF16OP(OpTest):
+    def setUp(self):
+        self.init()
+        self.op_type = "argsort"
+        self.python_api = paddle.argsort
+        self.public_python_api = paddle.argsort
+        self.dtype = np.uint16
+        self.descending = False
+        self.attrs = {"axis": self.axis, "descending": self.descending}
+        self.x = np.random.rand(*self.input_shape).astype(np.float32)
+        self.sorted_x = np.sort(self.x, kind='heapsort', axis=self.axis).astype(
+            np.float32
+        )
+        self.indices = np.argsort(
+            self.x, kind='heapsort', axis=self.axis
+        ).astype(np.float32)
+        self.inputs = {'X': convert_float_to_uint16(self.x)}
+        self.outputs = {
+            'Out': convert_float_to_uint16(self.sorted_x),
+            "Indices": convert_float_to_uint16(self.indices),
+        }
+
+    def init(self):
+        self.input_shape = [
+            1000,
+        ]
+        self.axis = 0
+
+    def test_check_output(self):
+        place = core.CUDAPlace(0)
+        self.check_output_with_place(place)
+
+    def test_check_grad(self):
+        place = core.CUDAPlace(0)
+        self.check_grad_with_place(
+            place,
+            ['X'],
+            'Out',
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

python/paddle/tensor/search.py

@@ -108,6 +108,7 @@ def argsort(x, axis=-1, descending=False, name=None):
                 'int16',
                 'int32',
                 'int64',
+                'uint16',
                 'uint8',
             ],
             'argsort',