PaddlePaddle
diff --git a/‎paddle/phi/backends/xpu/xpu2_op_list.cc
Copy file name to clipboard
+36-29Lines changed: 36 additions & 29 deletions b/‎paddle/phi/backends/xpu/xpu2_op_list.cc
Copy file name to clipboard
+36-29Lines changed: 36 additions & 29 deletions
diff --git a/‎paddle/phi/backends/xpu/xpu3_op_list.cc
Copy file name to clipboard
+21-7Lines changed: 21 additions & 7 deletions b/‎paddle/phi/backends/xpu/xpu3_op_list.cc
Copy file name to clipboard
+21-7Lines changed: 21 additions & 7 deletions
diff --git a/‎paddle/phi/kernels/xpu/activation_grad_kernel.cc
Copy file name to clipboard
+24-7Lines changed: 24 additions & 7 deletions b/‎paddle/phi/kernels/xpu/activation_grad_kernel.cc
Copy file name to clipboard
+24-7Lines changed: 24 additions & 7 deletions
diff --git a/‎paddle/phi/kernels/xpu/activation_kernel.cc
Copy file name to clipboard
+15-3Lines changed: 15 additions & 3 deletions b/‎paddle/phi/kernels/xpu/activation_kernel.cc
Copy file name to clipboard
+15-3Lines changed: 15 additions & 3 deletions
diff --git a/‎paddle/phi/kernels/xpu/full_kernel.cc
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/phi/kernels/xpu/full_kernel.cc
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/phi/kernels/xpu/reduce_mean_grad_kernel.cc
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions b/‎paddle/phi/kernels/xpu/reduce_mean_grad_kernel.cc
Copy file name to clipboard
+3-3Lines changed: 3 additions & 3 deletions
@@ -429,7 +429,14 @@ XPUOpMap& get_kl2_ops() {
                      phi::DataType::UINT8,
                      phi::DataType::BOOL,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
+      {"fill_constant_batch_size_like",
+       XPUKernelSet({phi::DataType::INT64,
+                     phi::DataType::INT32,
+                     phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"flatten2_grad",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
@@ -471,35 +478,22 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT64,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::FLOAT32,
+                     phi::DataType::BFLOAT16})},
       {"full_batch_size_like",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"full_like",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT32,
                      phi::DataType::FLOAT64,
-                     phi::DataType::FLOAT16})},
-      {"full_batch_size_like",
-       XPUKernelSet({phi::DataType::INT64,
-                     phi::DataType::INT32,
-                     phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT64,
-                     phi::DataType::FLOAT16})},
-      {"full_batch_size_like",
-       XPUKernelSet({phi::DataType::INT64,
-                     phi::DataType::INT32,
-                     phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
-      {"fill_constant_batch_size_like",
-       XPUKernelSet({phi::DataType::INT64,
-                     phi::DataType::INT32,
-                     phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"fused_multi_transformer_xpu",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"fused_multi_transformer_int8_xpu",
@@ -679,7 +673,9 @@ XPUOpMap& get_kl2_ops() {
       {"matmul",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"mean_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"mean",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
@@ -742,8 +738,14 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"pool3d",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
-      {"pow", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
-      {"pow_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"pow",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
+      {"pow_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"pow2_decay_with_linear_warmup", XPUKernelSet({phi::DataType::FLOAT32})},
       {"prior_box", XPUKernelSet({phi::DataType::FLOAT32})},
       {"prelu", XPUKernelSet({phi::DataType::FLOAT32})},
@@ -777,9 +779,7 @@ XPUOpMap& get_kl2_ops() {
       {"reduce_mean_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"reduce_mean",
-       XPUKernelSet({phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16,
-                     phi::DataType::BFLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"reduce_min_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"reduce_min",
        XPUKernelSet({phi::DataType::FLOAT32,
@@ -839,8 +839,14 @@ XPUOpMap& get_kl2_ops() {
       {"roi_align_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"roll", XPUKernelSet({phi::DataType::FLOAT32})},
       {"roll_grad", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"rsqrt", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"rsqrt_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"rsqrt",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
+      {"rsqrt_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"save",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT64,
@@ -853,6 +859,7 @@ XPUOpMap& get_kl2_ops() {
       {"scale",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"scatter",
 
@@ -448,6 +448,12 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT64,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16})},
+      {"fill_constant_batch_size_like",
+       XPUKernelSet({phi::DataType::INT64,
+                     phi::DataType::INT32,
+                     phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"flash_attn_grad",
        XPUKernelSet({phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT32,
@@ -493,16 +499,19 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT8,
                      phi::DataType::FLOAT32})},
       {"flip", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"full_batch_size_like",
+      {"full",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
-      {"fill_constant_batch_size_like",
+                     phi::DataType::BFLOAT16})},
+      {"full_batch_size_like",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"fused_multi_transformer_xpu",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"fused_rotary_position_embedding",
@@ -759,8 +768,11 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"pool3d",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
-      {"pow", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"pow_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"pow", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+      {"pow_grad",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"pow2_decay_with_linear_warmup", XPUKernelSet({phi::DataType::FLOAT32})},
       {"prior_box", XPUKernelSet({phi::DataType::FLOAT32})},
       {"prelu_grad",
@@ -788,7 +800,9 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
       {"reduce_mean_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"reduce_mean",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
 
@@ -465,6 +465,7 @@ void PowGradKernel(const Context& dev_ctx,
                    const DenseTensor& dout,
                    const Scalar& factor,
                    DenseTensor* dx) {
+  using XPUType = typename XPUTypeTrait<T>::Type;
   dev_ctx.template Alloc<T>(dx);
   const T* x_data = x.data<T>();
   const T* y_grad = dout.data<T>();
@@ -483,10 +484,12 @@ void PowGradKernel(const Context& dev_ctx,
   float pow_factor = factor.to<float>();
 
   auto xpu_context = dev_ctx.x_context();
-  // int pow_grad(Context* ctx, const T* x, const T* dy, T* dx, int len, float
-  // factor);
-  int r =
-      xpu::pow_grad(xpu_context, x_data, y_grad, x_grad, x.numel(), pow_factor);
+  int r = xpu::pow_grad(xpu_context,
+                        reinterpret_cast<const XPUType*>(x_data),
+                        reinterpret_cast<const XPUType*>(y_grad),
+                        reinterpret_cast<XPUType*>(x_grad),
+                        x.numel(),
+                        pow_factor);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "pow_grad");
 }
 
@@ -741,13 +744,30 @@ PD_REGISTER_KERNEL(hardswish_grad,
                    phi::HardSwishGradKernel,
                    float,
                    phi::dtype::float16) {}
+
 PD_REGISTER_KERNEL(sigmoid_grad,
                    XPU,
                    ALL_LAYOUT,
                    phi::SigmoidGradKernel,
                    float,
                    phi::dtype::float16) {}
 
+PD_REGISTER_KERNEL(pow_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::PowGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+
+PD_REGISTER_KERNEL(rsqrt_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::RsqrtGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(exp_grad, ExpGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(log_grad, LogGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(leaky_relu_grad, LeakyReluGradKernel)
@@ -759,6 +779,3 @@ PD_REGISTER_ACTIVATION_GRAD_KERNEL(mish_grad, MishGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(softplus_grad, SoftplusGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(sin_grad, SinGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(cos_grad, CosGradKernel)
-PD_REGISTER_ACTIVATION_GRAD_KERNEL(rsqrt_grad, RsqrtGradKernel)
-
-PD_REGISTER_KERNEL(pow_grad, XPU, ALL_LAYOUT, phi::PowGradKernel, float) {}
 
@@ -665,8 +665,21 @@ PD_REGISTER_KERNEL(cos,
                    phi::dtype::float16,
                    phi::dtype::bfloat16) {}
 
-PD_REGISTER_KERNEL(
-    pow, XPU, ALL_LAYOUT, phi::PowKernel, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(pow,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::PowKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+
+PD_REGISTER_KERNEL(rsqrt,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::RsqrtKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 
 #define PD_REGISTER_ACTIVATION_KERNEL(name, func) \
   PD_REGISTER_KERNEL(name, XPU, ALL_LAYOUT, phi::func, float) {}
@@ -676,4 +689,3 @@ PD_REGISTER_ACTIVATION_KERNEL(floor, FloorKernel)
 PD_REGISTER_ACTIVATION_KERNEL(mish, MishKernel)
 PD_REGISTER_ACTIVATION_KERNEL(reciprocal, ReciprocalKernel)
 PD_REGISTER_ACTIVATION_KERNEL(softplus, SoftplusKernel)
-PD_REGISTER_ACTIVATION_KERNEL(rsqrt, RsqrtKernel)
 
@@ -153,7 +153,8 @@ PD_REGISTER_KERNEL(full_batch_size_like,
                    int,
                    int64_t,
                    bool,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
 }
 
 
@@ -26,18 +26,18 @@ template <typename T, typename Context>
 void ReduceMeanGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           const DenseTensor& out_grad,
-                          const IntArray& dims_arr,
+                          const IntArray& dims,
                           bool keep_dim,
                           bool reduce_all,
                           DenseTensor* x_grad) {
   using XPUType = typename XPUTypeTrait<T>::Type;
-  reduce_all = recompute_reduce_all(x, dims_arr, reduce_all);
+  reduce_all = recompute_reduce_all(x, dims, reduce_all);
   dev_ctx.template Alloc<T>(x_grad);
   const XPUType* dy_data = reinterpret_cast<const XPUType*>(out_grad.data<T>());
 
   XPUType* x_data = reinterpret_cast<XPUType*>(x_grad->data<T>());
 
-  auto reduce_dims = dims_arr.GetData();
+  auto reduce_dims = dims.GetData();
 
   std::vector<int> xdims = common::vectorize<int>(x.dims());
   std::vector<int> ydims = common::vectorize<int>(out_grad.dims());
Original file line number	Diff line number	Diff line change
`@@ -153,7 +153,8 @@ PD_REGISTER_KERNEL(full_batch_size_like,`
`153`	`153`	`int,`
`154`	`154`	`int64_t,`
`155`	`155`	`bool,`
`156`		`- phi::dtype::float16) {`
	`156`	`+ phi::dtype::float16,`
	`157`	`+ phi::dtype::bfloat16) {`
`157`	`158`	`kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);`
`158`	`159`	`}`
`159`	`160`