[XPU] bind bf16 flatten and uint8 stack (#68418)

RuohengMa · web-flow · commit 431ab811b831 · 2024-09-25T14:02:14.000+08:00
* [XPU] bind bf16 flatten and uint8 stack

* add more support types

* add more types

* refine flatten2 unittest

* remove bool of flatten and add flatten_with_xshape

* add bool back
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -477,41 +477,71 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16})},
       {"flatten2_grad",
-       XPUKernelSet({phi::DataType::INT64,
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT64,
                      phi::DataType::INT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::UINT8,
                      phi::DataType::INT8,
-                     phi::DataType::FLOAT32})},
+                     phi::DataType::INT16,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::BOOL})},
       {"flatten2",
-       XPUKernelSet({phi::DataType::INT64,
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT64,
                      phi::DataType::INT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::UINT8,
                      phi::DataType::INT8,
-                     phi::DataType::FLOAT32})},
+                     phi::DataType::INT16,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::BOOL})},
       {"flatten_contiguous_range_grad",
-       XPUKernelSet({phi::DataType::INT64,
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT64,
                      phi::DataType::INT32,
-                     phi::DataType::INT8,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
-                     phi::DataType::FLOAT32})},
+                     phi::DataType::UINT8,
+                     phi::DataType::INT8,
+                     phi::DataType::INT16,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::BOOL})},
       {"flatten_contiguous_range",
-       XPUKernelSet({phi::DataType::INT64,
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT64,
                      phi::DataType::INT32,
-                     phi::DataType::INT8,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
-                     phi::DataType::FLOAT32})},
+                     phi::DataType::UINT8,
+                     phi::DataType::INT8,
+                     phi::DataType::INT16,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::BOOL})},
       {"flatten_grad",
-       XPUKernelSet({phi::DataType::INT64,
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT64,
                      phi::DataType::INT32,
-                     phi::DataType::INT8,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
-                     phi::DataType::FLOAT32})},
+                     phi::DataType::UINT8,
+                     phi::DataType::INT8,
+                     phi::DataType::INT16,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::BOOL})},
       {"flatten",
-       XPUKernelSet({phi::DataType::INT64,
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::INT64,
                      phi::DataType::INT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::UINT8,
                      phi::DataType::INT8,
-                     phi::DataType::FLOAT32})},
+                     phi::DataType::INT16,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::BOOL})},
       {"flip", XPUKernelSet({phi::DataType::FLOAT32})},
       {"full",
        XPUKernelSet({phi::DataType::INT64,
@@ -1140,13 +1170,21 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT16,
-                     phi::DataType::BFLOAT16})},
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::UINT8,
+                     phi::DataType::INT8,
+                     phi::DataType::INT16,
+                     phi::DataType::BOOL,
+                     phi::DataType::FLOAT64})},
       {"stack_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT16,
-                     phi::DataType::BFLOAT16})},
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::UINT8,
+                     phi::DataType::INT8,
+                     phi::DataType::INT16})},
       {"strided_slice",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
diff --git a/paddle/phi/kernels/cpu/flatten2_grad_kernel.cc b/paddle/phi/kernels/cpu/flatten2_grad_kernel.cc
@@ -25,4 +25,5 @@ PD_REGISTER_KERNEL(flatten2_grad,
                    uint8_t,
                    int,
                    int8_t,
-                   int64_t) {}
+                   int64_t,
+                   bool) {}
diff --git a/paddle/phi/kernels/cpu/flatten2_kernel.cc b/paddle/phi/kernels/cpu/flatten2_kernel.cc
@@ -25,4 +25,5 @@ PD_REGISTER_KERNEL(flatten2,
                    uint8_t,
                    int,
                    int8_t,
-                   int64_t) {}
+                   int64_t,
+                   bool) {}
diff --git a/paddle/phi/kernels/flatten_grad_kernel.cc b/paddle/phi/kernels/flatten_grad_kernel.cc
@@ -46,7 +46,8 @@ PD_REGISTER_KERNEL(flatten_grad,
                    uint8_t,
                    int8_t,
                    int,
-                   int64_t) {}
+                   int64_t,
+                   bool) {}
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
 PD_REGISTER_KERNEL(flatten_grad,
@@ -69,12 +70,16 @@ PD_REGISTER_KERNEL(flatten_grad,
                    XPU,
                    ALL_LAYOUT,
                    phi::FlattenGradKernel,
+                   double,
                    float,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
-                   int8_t,
+                   int64_t,
                    int,
-                   int64_t) {}
+                   int16_t,
+                   int8_t,
+                   uint8_t,
+                   bool) {}
 
 #endif
 
diff --git a/paddle/phi/kernels/flatten_kernel.cc b/paddle/phi/kernels/flatten_kernel.cc
@@ -60,7 +60,8 @@ PD_REGISTER_KERNEL(flatten,
                    int8_t,
                    int16_t,
                    int,
-                   int64_t) {}
+                   int64_t,
+                   bool) {}
 
 PD_REGISTER_KERNEL(flatten_with_xshape,
                    CPU,
@@ -73,7 +74,8 @@ PD_REGISTER_KERNEL(flatten_with_xshape,
                    int8_t,
                    int16_t,
                    int,
-                   int64_t) {}
+                   int64_t,
+                   bool) {}
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
 PD_REGISTER_KERNEL(flatten,
@@ -110,24 +112,31 @@ PD_REGISTER_KERNEL(flatten,
                    XPU,
                    ALL_LAYOUT,
                    phi::FlattenKernel,
+                   double,
                    float,
                    phi::dtype::float16,
-                   int8_t,
-                   int16_t,
+                   phi::dtype::bfloat16,
+                   int64_t,
                    int,
-                   int64_t) {}
+                   int16_t,
+                   int8_t,
+                   uint8_t,
+                   bool) {}
 
 PD_REGISTER_KERNEL(flatten_with_xshape,
                    XPU,
                    ALL_LAYOUT,
                    phi::FlattenWithXShapeKernel,
+                   double,
                    float,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
-                   int8_t,
-                   int16_t,
+                   int64_t,
                    int,
-                   int64_t) {}
+                   int16_t,
+                   int8_t,
+                   uint8_t,
+                   bool) {}
 #endif
 
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
diff --git a/paddle/phi/kernels/xpu/flatten2_grad_kernel.cc b/paddle/phi/kernels/xpu/flatten2_grad_kernel.cc
@@ -22,9 +22,15 @@ PD_REGISTER_KERNEL(flatten2_grad,
                    XPU,
                    ALL_LAYOUT,
                    phi::Flatten2GradKernel,
+                   double,
                    float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   int64_t,
                    int,
+                   int16_t,
                    int8_t,
-                   int64_t) {}
+                   uint8_t,
+                   bool) {}
 
 #endif
diff --git a/paddle/phi/kernels/xpu/flatten2_kernel.cc b/paddle/phi/kernels/xpu/flatten2_kernel.cc
@@ -22,9 +22,15 @@ PD_REGISTER_KERNEL(flatten2,
                    XPU,
                    ALL_LAYOUT,
                    phi::Flatten2Kernel,
+                   double,
                    float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   int64_t,
                    int,
+                   int16_t,
                    int8_t,
-                   int64_t) {}
+                   uint8_t,
+                   bool) {}
 
 #endif
diff --git a/paddle/phi/kernels/xpu/stack_grad_kernel.cc b/paddle/phi/kernels/xpu/stack_grad_kernel.cc
@@ -54,7 +54,10 @@ PD_REGISTER_KERNEL(stack_grad,
                    ALL_LAYOUT,
                    phi::StackGradKernel,
                    float,
-                   int,
-                   int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   int64_t,
+                   int,
+                   int16_t,
+                   int8_t,
+                   uint8_t) {}
diff --git a/paddle/phi/kernels/xpu/stack_kernel.cc b/paddle/phi/kernels/xpu/stack_kernel.cc
@@ -59,8 +59,13 @@ PD_REGISTER_KERNEL(stack,
                    XPU,
                    ALL_LAYOUT,
                    phi::StackKernel,
+                   double,
                    float,
-                   int,
-                   int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   int64_t,
+                   int,
+                   int16_t,
+                   int8_t,
+                   uint8_t,
+                   bool) {}
diff --git a/test/xpu/test_flatten2_op_xpu.py b/test/xpu/test_flatten2_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
@@ -39,9 +40,13 @@ def setUp(self):
             self.dtype = self.in_type
             self.place = paddle.XPUPlace(0)
             self.init_test_case()
-            self.inputs = {
-                "X": np.random.random(self.in_shape).astype(self.dtype)
-            }
+            if self.dtype == np.uint16:
+                data = np.random.random(self.in_shape).astype(np.float32)
+                self.inputs = {"X": convert_float_to_uint16(data)}
+            else:
+                self.inputs = {
+                    "X": np.random.random(self.in_shape).astype(self.dtype)
+                }
             self.init_attrs()
             self.outputs = {
                 "Out": self.inputs["X"].reshape(self.new_shape),
diff --git a/test/xpu/test_flatten_contiguous_range_op_xpu.py b/test/xpu/test_flatten_contiguous_range_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
@@ -44,9 +45,13 @@ def setUp(self):
             self.stop_axis = -1
             self.dtype = self.in_type
             self.init_test_case()
-            self.inputs = {
-                "X": np.random.random(self.in_shape).astype(self.dtype)
-            }
+            if self.dtype == np.uint16:
+                data = np.random.random(self.in_shape).astype(np.float32)
+                self.inputs = {"X": convert_float_to_uint16(data)}
+            else:
+                self.inputs = {
+                    "X": np.random.random(self.in_shape).astype(self.dtype)
+                }
             self.init_attrs()
             self.outputs = {
                 "Out": self.inputs["X"].reshape(self.new_shape),
diff --git a/test/xpu/test_stack_op_xpu.py b/test/xpu/test_stack_op_xpu.py
@@ -21,6 +21,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
@@ -47,9 +48,15 @@ def setUp(self):
             self.dtype = self.in_type
             self.x = []
             for i in range(self.num_inputs):
-                self.x.append(
-                    np.random.random(size=self.input_dim).astype(self.dtype)
-                )
+                if self.dtype == np.uint16:
+                    data = np.random.random(size=self.input_dim).astype(
+                        np.float32
+                    )
+                    self.x.append(convert_float_to_uint16(data))
+                else:
+                    self.x.append(
+                        np.random.random(size=self.input_dim).astype(self.dtype)
+                    )
 
             tmp = []
             x_names = self.get_x_names()
