[XPU] add some bf16 ops (#59505)

houj04 · web-flow · commit b10c4ed50971 · 2023-12-01T13:36:55.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -275,7 +275,9 @@ XPUOpMap& get_kl3_ops() {
       {"elementwise_min",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"elementwise_mul_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"elementwise_mul",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
@@ -360,6 +362,7 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT32,
                      phi::DataType::BOOL,
                      phi::DataType::FLOAT16,
+                     phi::DataType::FLOAT64,
                      phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT32})},
       {"fill_diagonal_tensor",
@@ -376,6 +379,7 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::UINT8,
                      phi::DataType::BOOL,
                      phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT64,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16})},
       {"flatten2_grad",
@@ -771,8 +775,13 @@ XPUOpMap& get_kl3_ops() {
       {"sgd_dense_param_sparse_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"silu_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
-      {"silu", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
+      {"silu",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"size",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
@@ -847,9 +856,13 @@ XPUOpMap& get_kl3_ops() {
       {"square",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"squared_l2_norm",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"squared_l2_norm_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"squeeze2_grad",
        XPUKernelSet({phi::DataType::FLOAT64,
                      phi::DataType::INT64,
@@ -944,16 +957,21 @@ XPUOpMap& get_kl3_ops() {
       {"tril_triu",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::BOOL})},
       {"tril",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::BOOL})},
       {"triu",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT16,
-                     phi::DataType::BFLOAT16})},
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::BOOL})},
       {"tril_triu_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/kernels/xpu/activation_grad_kernel.cc b/paddle/phi/kernels/xpu/activation_grad_kernel.cc
@@ -379,10 +379,19 @@ struct XPUSiluGradFunctor : public funcs::BaseActivationFunctor<T> {
     XPUType* x_grad = reinterpret_cast<XPUType*>(dx->data<T>());
 
     if (std::getenv("XPU_PADDLE_ACT_LUT") != nullptr) {
-      int r = xpu::fast_swish_grad(
-          dev_ctx.x_context(), x_data, y_grad, x_grad, dx->numel());
-      PADDLE_ENFORCE_XDNN_SUCCESS(r, "fast_swish_grad");
+      if (!std::is_same<T, ::phi::dtype::bfloat16>::value) {
+        // use fast_swish_grad if NOT bf16
+        int r = xpu::fast_swish_grad(
+            dev_ctx.x_context(), x_data, y_grad, x_grad, dx->numel());
+        PADDLE_ENFORCE_XDNN_SUCCESS(r, "fast_swish_grad");
+      } else {
+        // use plain swish_grad
+        int r = xpu::swish_grad(
+            dev_ctx.x_context(), x_data, y_grad, x_grad, dx->numel());
+        PADDLE_ENFORCE_XDNN_SUCCESS(r, "swish_grad");
+      }
     } else {
+      // use plain swish_grad
       int r = xpu::swish_grad(
           dev_ctx.x_context(), x_data, y_grad, x_grad, dx->numel());
       PADDLE_ENFORCE_XDNN_SUCCESS(r, "swish_grad");
@@ -697,7 +706,8 @@ PD_REGISTER_KERNEL(silu_grad,
                    ALL_LAYOUT,
                    phi::SiluGradKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 
 #define PD_REGISTER_ACTIVATION_GRAD_KERNEL(name, func) \
   PD_REGISTER_KERNEL(name, XPU, ALL_LAYOUT, phi::func, float) {}
diff --git a/paddle/phi/kernels/xpu/activation_kernel.cc b/paddle/phi/kernels/xpu/activation_kernel.cc
@@ -334,10 +334,19 @@ struct XPUSiluFunctor : public funcs::BaseActivationFunctor<T> {
 
     auto xpu_context = dev_ctx.x_context();
     if (std::getenv("XPU_PADDLE_ACT_LUT") != nullptr) {
-      int r = xpu::fast_swish(
-          xpu_context, x_data, y_data, x.numel(), nullptr, nullptr);
-      PADDLE_ENFORCE_XDNN_SUCCESS(r, "fast_swish");
+      if (!std::is_same<T, ::phi::dtype::bfloat16>::value) {
+        // use fast_swish if NOT bf16
+        int r = xpu::fast_swish(
+            xpu_context, x_data, y_data, x.numel(), nullptr, nullptr);
+        PADDLE_ENFORCE_XDNN_SUCCESS(r, "fast_swish");
+      } else {
+        // use plain swish
+        int r = xpu::swish(
+            xpu_context, x_data, y_data, x.numel(), nullptr, nullptr);
+        PADDLE_ENFORCE_XDNN_SUCCESS(r, "swish");
+      }
     } else {
+      // use plain swish
       int r =
           xpu::swish(xpu_context, x_data, y_data, x.numel(), nullptr, nullptr);
       PADDLE_ENFORCE_XDNN_SUCCESS(r, "swish");
@@ -574,8 +583,13 @@ void HardSwishKernel(const Context& dev_ctx,
 
 PD_REGISTER_KERNEL(
     relu, XPU, ALL_LAYOUT, phi::ReluKernel, float, phi::dtype::float16) {}
-PD_REGISTER_KERNEL(
-    silu, XPU, ALL_LAYOUT, phi::SiluKernel, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(silu,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::SiluKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 PD_REGISTER_KERNEL(
     elu, XPU, ALL_LAYOUT, phi::EluKernel, float, phi::dtype::float16) {}
 PD_REGISTER_KERNEL(
diff --git a/paddle/phi/kernels/xpu/elementwise_multiply_grad_kernel.cc b/paddle/phi/kernels/xpu/elementwise_multiply_grad_kernel.cc
@@ -57,4 +57,5 @@ PD_REGISTER_KERNEL(multiply_grad,
                    ALL_LAYOUT,
                    phi::MultiplyGradKernel,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    float) {}
diff --git a/paddle/phi/kernels/xpu/full_kernel.cc b/paddle/phi/kernels/xpu/full_kernel.cc
@@ -120,6 +120,7 @@ PD_REGISTER_KERNEL(full,
                    ALL_LAYOUT,
                    phi::FullKernel,
                    float,
+                   double,
                    int8_t,
                    uint8_t,
                    int16_t,
@@ -134,6 +135,7 @@ PD_REGISTER_KERNEL(full_like,
                    ALL_LAYOUT,
                    phi::FullLikeKernel,
                    float,
+                   double,
                    uint8_t,
                    int16_t,
                    int,
diff --git a/paddle/phi/kernels/xpu/squared_l2_norm_grad_kernel.cc b/paddle/phi/kernels/xpu/squared_l2_norm_grad_kernel.cc
@@ -64,4 +64,5 @@ PD_REGISTER_KERNEL(squared_l2_norm_grad,
                    ALL_LAYOUT,
                    phi::SquaredL2NormGradKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/squared_l2_norm_kernel.cc b/paddle/phi/kernels/xpu/squared_l2_norm_kernel.cc
@@ -59,4 +59,5 @@ PD_REGISTER_KERNEL(squared_l2_norm,
                    ALL_LAYOUT,
                    phi::SquaredL2NormKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/tril_triu_kernel.cc b/paddle/phi/kernels/xpu/tril_triu_kernel.cc
@@ -70,14 +70,24 @@ PD_REGISTER_KERNEL(tril_triu,
                    phi::TrilTriuKernel,
                    int,
                    float,
-                   phi::dtype::float16) {}
-PD_REGISTER_KERNEL(
-    tril, XPU, ALL_LAYOUT, phi::TrilKernel, int, float, phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool) {}
+PD_REGISTER_KERNEL(tril,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::TrilKernel,
+                   int,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool) {}
 PD_REGISTER_KERNEL(triu,
                    XPU,
                    ALL_LAYOUT,
                    phi::TriuKernel,
                    int,
                    float,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   bool) {}
diff --git a/test/xpu/test_activation_op_xpu.py b/test/xpu/test_activation_op_xpu.py
@@ -22,7 +22,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
-from op_test import OpTest
+from op_test import OpTest, convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -97,10 +97,16 @@ def set_case(self):
             self.init_shape()
 
             np.random.seed(1024)
-            x = np.random.uniform(-1, 1, self.shape).astype(self.dtype)
+            x = np.random.uniform(-1, 1, self.shape)
+
+            if self.dtype == np.uint16:
+                # bfloat16 actually
+                new_x = convert_float_to_uint16(x)
+            else:
+                new_x = x.astype(self.dtype)
             out = x / (np.exp(-x) + 1)
 
-            self.inputs = {'X': x}
+            self.inputs = {'X': new_x}
             self.outputs = {'Out': out}
             self.attrs = {'use_xpu': True}
 
@@ -1150,7 +1156,13 @@ def set_case(self):
             self.init_config()
             out = ref_swish(self.x)
 
-            self.inputs = {'X': self.x}
+            if self.dtype == np.uint16:
+                # bfloat16 actually
+                new_x = convert_float_to_uint16(self.x)
+            else:
+                new_x = self.x.astype(self.dtype)
+
+            self.inputs = {'X': new_x}
             self.outputs = {'Out': out}
             self.attrs = {'use_xpu': True}
 
diff --git a/test/xpu/test_squared_l2_norm_op_xpu.py b/test/xpu/test_squared_l2_norm_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,10 +42,25 @@ def setUp(self):
             self.use_mkldnn = False
             self.max_relative_error = 0.05
             self.set_inputs()
-            self.inputs = {'X': self.x}
-            self.outputs = {
-                'Out': np.array([np.square(np.linalg.norm(self.x))])
-            }
+
+            if self.dtype == np.uint16:
+                # bfloat16 actually
+                new_x = convert_float_to_uint16(self.x)
+            else:
+                new_x = self.x.astype(self.dtype)
+
+            out = np.square(np.linalg.norm(self.x))
+
+            if self.dtype == np.uint16:
+                # bfloat16 actually
+                new_out = convert_float_to_uint16(out)
+            else:
+                new_out = out.astype(self.dtype)
+
+            new_out = np.array([new_out])
+
+            self.inputs = {'X': new_x}
+            self.outputs = {'Out': new_out}
 
         def test_check_output(self):
             self.check_output_with_place(self.place)
@@ -53,21 +69,17 @@ def test_check_grad(self):
             self.check_grad_with_place(self.place, ['X'], 'Out')
 
         def set_inputs(self):
-            self.x = np.random.uniform(-1, 1, (13, 19)).astype(self.in_type)
+            self.x = np.random.uniform(-1, 1, (13, 19))
             self.x[np.abs(self.x) < self.max_relative_error] = 0.1
 
     class TestSquaredL2NormOp_1(TestSquaredL2NormOp):
         def set_inputs(self):
-            self.x = np.random.uniform(-0.2, 0.2, (8, 128, 24)).astype(
-                self.in_type
-            )
+            self.x = np.random.uniform(-0.2, 0.2, (8, 128, 24))
             self.x[np.abs(self.x) < self.max_relative_error] = 0.02
 
     class TestSquaredL2NormOp_2(TestSquaredL2NormOp):
         def set_inputs(self):
-            self.x = np.random.uniform(-0.1, 0.1, (2, 128, 256)).astype(
-                self.in_type
-            )
+            self.x = np.random.uniform(-0.1, 0.1, (2, 128, 256))
             self.x[np.abs(self.x) < self.max_relative_error] = 0.01
 
 
