[XPU] add bfloat16 support for CastKernel (#58456)

houj04 · web-flow · commit 64211bade87b · 2023-10-30T10:46:53.000+08:00
* [XPU] add bfloat16 support for CastKernel

* fix unittest.
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -139,6 +139,7 @@ XPUOpMap& get_kl2_ops() {
       {"cast",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT64,
                      phi::DataType::BOOL,
                      phi::DataType::INT8,
@@ -874,6 +875,7 @@ XPUOpMap& get_kl2_ops() {
       {"transfer_dtype",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT64,
                      phi::DataType::BOOL,
                      phi::DataType::UINT8,
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -138,6 +138,7 @@ XPUOpMap& get_kl3_ops() {
       {"cast",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT64,
                      phi::DataType::BOOL,
                      phi::DataType::INT8,
@@ -850,6 +851,7 @@ XPUOpMap& get_kl3_ops() {
       {"transfer_dtype",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT64,
                      phi::DataType::BOOL,
                      phi::DataType::UINT8,
diff --git a/paddle/phi/kernels/xpu/cast_kernel.cc b/paddle/phi/kernels/xpu/cast_kernel.cc
@@ -33,11 +33,41 @@ void CastXPUKernelImpl(const Context& dev_ctx,
     return;
   }
 
+  if (std::is_same<InT, OutT>::value) {
+    int ret = xpu::copy(dev_ctx.x_context(),
+                        reinterpret_cast<const int8_t*>(in_data),
+                        reinterpret_cast<int8_t*>(out_data),
+                        x.numel() * phi::SizeOf(x.dtype()));
+    PADDLE_ENFORCE_XDNN_SUCCESS(ret, "copy");
+    return;
+  }
+
+  if (std::is_same<InT, dtype::bfloat16>::value &&
+          !std::is_same<OutT, float>::value ||
+      !std::is_same<InT, float>::value &&
+          std::is_same<OutT, dtype::bfloat16>::value) {
+    // bfloat -> non float, or non float -> bfloat, use float buffer
+    xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
+    float* cast_buffer = RAII_GUARD.alloc_l3_or_gm<float>(numel);
+    // step 1: InT to float
+    int r = xpu::cast<XPUInT, float>(dev_ctx.x_context(),
+                                     reinterpret_cast<const XPUInT*>(in_data),
+                                     cast_buffer,
+                                     numel);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
+    // step 2: float to OutT
+    r = xpu::cast<float, XPUOutT>(dev_ctx.x_context(),
+                                  cast_buffer,
+                                  reinterpret_cast<XPUOutT*>(out_data),
+                                  numel);
+    PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
+    return;
+  }
+
   int r = xpu::cast<XPUInT, XPUOutT>(dev_ctx.x_context(),
                                      reinterpret_cast<const XPUInT*>(in_data),
                                      reinterpret_cast<XPUOutT*>(out_data),
                                      numel);
-
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
 }
 
@@ -56,6 +86,9 @@ void CastKernel(const Context& dev_ctx,
     case DataType::FLOAT16:
       CastXPUKernelImpl<T, dtype::float16, Context>(dev_ctx, x, out);
       break;
+    case DataType::BFLOAT16:
+      CastXPUKernelImpl<T, dtype::bfloat16, Context>(dev_ctx, x, out);
+      break;
     case DataType::INT64:
       CastXPUKernelImpl<T, int64_t, Context>(dev_ctx, x, out);
       break;
@@ -85,6 +118,7 @@ PD_REGISTER_KERNEL(cast,
                    int32_t,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int64_t,
                    bool,
                    int8_t,
diff --git a/test/xpu/test_cast_op_xpu.py b/test/xpu/test_cast_op_xpu.py
@@ -20,6 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_float_to_uint16, convert_uint16_to_float
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -31,6 +32,7 @@
     'int64': int(core.VarDesc.VarType.INT64),
     'float32': int(core.VarDesc.VarType.FP32),
     'float16': int(core.VarDesc.VarType.FP16),
+    'bfloat16': int(core.VarDesc.VarType.BF16),
     'bool': int(core.VarDesc.VarType.BOOL),
     'int8': int(core.VarDesc.VarType.INT8),
     'uint8': int(core.VarDesc.VarType.UINT8),
@@ -48,6 +50,7 @@ def dynamic_create_class(self):
         classes = []
         for out_type in {
             'float16',
+            'bfloat16',
             'float32',
             'int32',
             'int64',
@@ -71,8 +74,18 @@ def setUp(self):
                 else self.out_typename
             )
 
-            self.inputs = {'X': ipt.astype(in_typename)}
-            self.outputs = {'Out': ipt.astype(in_typename).astype(out_typename)}
+            if in_typename == "bfloat16":
+                ipt_x = convert_float_to_uint16(ipt)
+            else:
+                ipt_x = ipt.astype(in_typename)
+
+            if out_typename == "bfloat16":
+                opt = convert_uint16_to_float(convert_float_to_uint16(ipt_x))
+            else:
+                opt = ipt_x.astype(out_typename)
+
+            self.inputs = {'X': ipt_x}
+            self.outputs = {'Out': opt}
             self.attrs = {
                 'in_dtype': typeid_dict[in_typename],
                 'out_dtype': typeid_dict[out_typename],
