[XPU] support several bf16 ops (#68169)

runzhech · web-flow · commit 4a3929581383 · 2024-09-23T13:56:53.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -313,15 +313,21 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"elementwise_floordiv",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"elementwise_max_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"elementwise_max",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"elementwise_min_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"elementwise_min",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"elementwise_mul_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
@@ -395,6 +401,7 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::BOOL,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32})},
       {"expand_v2",
@@ -1148,7 +1155,10 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::BFLOAT16,
                      phi::DataType::INT16,
                      phi::DataType::INT32})},
-      {"sum", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+      {"sum",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::FLOAT16})},
       {"swiglu",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
@@ -1366,7 +1376,9 @@ XPUOpMap& get_kl3_ops() {
       {"resnet_basic_block_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"resnet_basic_block", XPUKernelSet({phi::DataType::FLOAT32})},
       {"fused_gemm_epilogue",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"fused_gemm_epilogue_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"fused_attention",
diff --git a/paddle/phi/kernels/fusion/xpu/fused_gemm_epilogue_kernel.cc b/paddle/phi/kernels/fusion/xpu/fused_gemm_epilogue_kernel.cc
@@ -84,4 +84,5 @@ PD_REGISTER_KERNEL(fused_gemm_epilogue,
                    ALL_LAYOUT,
                    phi::fusion::FusedGemmEpilogueKernel,
                    float,
+                   phi::dtype::bfloat16,
                    phi::dtype::float16) {}
diff --git a/paddle/phi/kernels/legacy/xpu/elementwise_kernel.cc b/paddle/phi/kernels/legacy/xpu/elementwise_kernel.cc
@@ -121,6 +121,7 @@ PD_REGISTER_KERNEL(floor_divide_raw,
                    ALL_LAYOUT,
                    phi::FloorDivideRawKernel,
                    float,
+                   phi::dtype::bfloat16,
                    phi::dtype::float16,
                    int32_t,
                    int64_t) {}
@@ -130,6 +131,7 @@ PD_REGISTER_KERNEL(maximum_raw,
                    phi::MaximumRawKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int32_t,
                    int64_t) {}
 PD_REGISTER_KERNEL(minimum_raw,
@@ -138,6 +140,7 @@ PD_REGISTER_KERNEL(minimum_raw,
                    phi::MinimumRawKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int32_t,
                    int64_t) {}
 PD_REGISTER_KERNEL(remainder_raw,
diff --git a/paddle/phi/kernels/xpu/add_n_kernel.cc b/paddle/phi/kernels/xpu/add_n_kernel.cc
@@ -160,11 +160,17 @@ void AddNArrayKernel(const Context& dev_ctx,
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    add_n, XPU, ALL_LAYOUT, phi::AddNKernel, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(add_n,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::AddNKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 PD_REGISTER_KERNEL(add_n_array,
                    XPU,
                    ALL_LAYOUT,
                    phi::AddNArrayKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/elementwise_kernel.cc b/paddle/phi/kernels/xpu/elementwise_kernel.cc
@@ -82,6 +82,7 @@ PD_REGISTER_KERNEL(floor_divide,
                    ALL_LAYOUT,
                    phi::FloorDivideKernel,
                    float,
+                   phi::dtype::bfloat16,
                    phi::dtype::float16,
                    int32_t,
                    int64_t) {}
@@ -90,6 +91,7 @@ PD_REGISTER_KERNEL(maximum,
                    ALL_LAYOUT,
                    phi::MaximumKernel,
                    float,
+                   phi::dtype::bfloat16,
                    phi::dtype::float16,
                    int32_t,
                    int64_t) {}
@@ -98,6 +100,7 @@ PD_REGISTER_KERNEL(minimum,
                    ALL_LAYOUT,
                    phi::MinimumKernel,
                    float,
+                   phi::dtype::bfloat16,
                    phi::dtype::float16,
                    int32_t,
                    int64_t) {}
diff --git a/paddle/phi/kernels/xpu/expand_as_kernel.cc b/paddle/phi/kernels/xpu/expand_as_kernel.cc
@@ -127,6 +127,7 @@ PD_REGISTER_KERNEL(expand_as,
                    ALL_LAYOUT,
                    phi::ExpandAsKernel,
                    float,
+                   phi::dtype::bfloat16,
                    phi::dtype::float16,
                    bool,
                    int,
diff --git a/test/xpu/test_elementwise_max_op_xpu.py b/test/xpu/test_elementwise_max_op_xpu.py
@@ -20,7 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
-from op_test import skip_check_grad_ci
+from op_test import convert_float_to_uint16, skip_check_grad_ci
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -37,7 +37,9 @@ class TestElementwiseOp(XPUOpTest):
         def setUp(self):
             self.use_xpu = True
             self.op_type = "elementwise_max"
-            self.dtype = self.in_type
+            self.dtype = (
+                self.in_type if self.in_type != np.uint16 else np.float32
+            )
             self.init_input_output()
             # If x and y have the same value, the max() is not differentiable.
             # So we generate test data by the following method
@@ -47,6 +49,9 @@ def init_input_output(self):
             x = np.random.uniform(0.1, 1, [13, 17]).astype(self.dtype)
             sgn = np.random.choice([-1, 1], [13, 17]).astype(self.dtype)
             y = x + sgn * np.random.uniform(0.1, 1, [13, 17]).astype(self.dtype)
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
             self.outputs = {
                 'Out': np.maximum(self.inputs['X'], self.inputs['Y'])
@@ -91,6 +96,9 @@ class TestElementwiseMaxOp_scalar(TestElementwiseOp):
         def init_input_output(self):
             x = np.random.random_integers(-5, 5, [2, 3, 20]).astype(self.dtype)
             y = np.array([0.5]).astype(self.dtype)
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
             self.outputs = {
                 'Out': np.maximum(self.inputs['X'], self.inputs['Y'])
@@ -101,6 +109,9 @@ def init_input_output(self):
             x = np.random.random((100,)).astype(self.dtype)
             sgn = np.random.choice([-1, 1], (100,)).astype(self.dtype)
             y = x + sgn * np.random.uniform(0.1, 1, (100,)).astype(self.dtype)
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
             self.outputs = {
                 'Out': np.maximum(self.inputs['X'], self.inputs['Y'])
@@ -113,6 +124,9 @@ def init_input_output(self):
             y = x[0, 0, :] + sgn * np.random.uniform(1, 2, (100,)).astype(
                 self.dtype
             )
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
 
             self.outputs = {
@@ -128,6 +142,9 @@ def init_input_output(self):
             y = x + sgn * np.random.uniform(1, 2, (2, 3, 1, 5)).astype(
                 self.dtype
             )
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
 
             self.outputs = {
diff --git a/test/xpu/test_elementwise_min_op_xpu.py b/test/xpu/test_elementwise_min_op_xpu.py
@@ -20,7 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
-from op_test import skip_check_grad_ci
+from op_test import convert_float_to_uint16, skip_check_grad_ci
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -39,13 +39,18 @@ def setUp(self):
             # If x and y have the same value, the min() is not differentiable.
             # So we generate test data by the following method
             # to avoid them being too close to each other.
-            self.dtype = self.in_type
+            self.dtype = (
+                self.in_type if self.in_type != np.uint16 else np.float32
+            )
             self.init_input_output()
 
         def init_input_output(self):
             x = np.random.uniform(0.1, 1, [13, 17]).astype(self.dtype)
             sgn = np.random.choice([-1, 1], [13, 17]).astype(self.dtype)
             y = x + sgn * np.random.uniform(0.1, 1, [13, 17]).astype(self.dtype)
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
             self.outputs = {
                 'Out': np.minimum(self.inputs['X'], self.inputs['Y'])
@@ -90,6 +95,9 @@ class TestElementwiseMinOp_scalar(TestElementwiseOp):
         def init_input_output(self):
             x = np.random.random_integers(-5, 5, [10, 3, 4]).astype(self.dtype)
             y = np.array([0.5]).astype(self.dtype)
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
             self.outputs = {
                 'Out': np.minimum(self.inputs['X'], self.inputs['Y'])
@@ -100,6 +108,9 @@ def init_input_output(self):
             x = np.random.random((100,)).astype(self.dtype)
             sgn = np.random.choice([-1, 1], (100,)).astype(self.dtype)
             y = x + sgn * np.random.uniform(0.1, 1, (100,)).astype(self.dtype)
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
             self.outputs = {
                 'Out': np.minimum(self.inputs['X'], self.inputs['Y'])
@@ -112,6 +123,9 @@ def init_input_output(self):
             y = x[0, 0, :] + sgn * np.random.uniform(1, 2, (100,)).astype(
                 self.dtype
             )
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
             self.outputs = {
                 'Out': np.minimum(
@@ -126,6 +140,9 @@ def init_input_output(self):
             y = x + sgn * np.random.uniform(1, 2, (2, 10, 1, 5)).astype(
                 self.dtype
             )
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
+                y = convert_float_to_uint16(y)
             self.inputs = {'X': x, 'Y': y}
             self.outputs = {
                 'Out': np.minimum(self.inputs['X'], self.inputs['Y'])
