tile and tile_grad support bf16 for xpu (#63075)

zhangyk0314 · web-flow · commit 65fae7c0d1fd · 2024-03-29T16:32:13.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -1048,8 +1048,10 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT64,
                      phi::DataType::BOOL,
                      phi::DataType::FLOAT64,
-                     phi::DataType::FLOAT32})},
-      {"tile_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+                     phi::DataType::FLOAT32,
+                     phi::DataType::BFLOAT16})},
+      {"tile_grad",
+       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::BFLOAT16})},
       {"transpose2_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
diff --git a/paddle/phi/kernels/xpu/tile_grad_kernel.cc b/paddle/phi/kernels/xpu/tile_grad_kernel.cc
@@ -83,8 +83,8 @@ void TileGradKernel(const Context& dev_ctx,
     using XPUType = typename XPUTypeTrait<T>::Type;
     // int reduce_sum(Context* ctx, const T* x, T* y, const std::vector<int>&
     // xshape, const std::vector<int>& rdims)
-    const auto* out_data = out_grad.data<XPUType>();
-    auto* x_grad_data = x_grad->data<XPUType>();
+    const auto* out_data = reinterpret_cast<const XPUType*>(out_grad.data<T>());
+    auto* x_grad_data = reinterpret_cast<XPUType*>(x_grad->data<T>());
     int r = xpu::reduce_sum<XPUType>(dev_ctx.x_context(),
                                      out_data,
                                      x_grad_data,
@@ -96,4 +96,9 @@ void TileGradKernel(const Context& dev_ctx,
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(tile_grad, XPU, ALL_LAYOUT, phi::TileGradKernel, float) {}
+PD_REGISTER_KERNEL(tile_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::TileGradKernel,
+                   float,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/tile_kernel.cc b/paddle/phi/kernels/xpu/tile_kernel.cc
@@ -29,6 +29,7 @@ void TileKernel(const Context& dev_ctx,
                 const DenseTensor& x,
                 const IntArray& repeat_times_arr,
                 DenseTensor* out) {
+  using XPUType = typename XPUTypeTrait<T>::Type;
   auto rank = x.dims().size();
   std::vector<int64_t> repeat_times = repeat_times_arr.GetData();
   int repeat_times_size = repeat_times.size();
@@ -123,17 +124,23 @@ void TileKernel(const Context& dev_ctx,
                                  vec_out_dims);
 
   } else {
-    ret = xpu::broadcast<T>(dev_ctx.x_context(),
-                            x.data<T>(),
-                            out->data<T>(),
-                            vec_in_dims,
-                            vec_out_dims);
+    const auto* x_data = reinterpret_cast<const XPUType*>(x.data<T>());
+    auto* out_data = reinterpret_cast<XPUType*>(out->data<T>());
+    ret = xpu::broadcast<XPUType>(
+        dev_ctx.x_context(), x_data, out_data, vec_in_dims, vec_out_dims);
   }
   PADDLE_ENFORCE_XDNN_SUCCESS(ret, "broadcast");
 }
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    tile, XPU, ALL_LAYOUT, phi::TileKernel, bool, float, double, int, int64_t) {
-}
+PD_REGISTER_KERNEL(tile,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::TileKernel,
+                   bool,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::dtype::bfloat16) {}
