[XPU] gelu and gelu_grad support bfloat16 type (#71809)

zhangyk0314 · web-flow · commit 943477794eb1 · 2025-03-21T14:27:15.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -762,8 +762,13 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16})},
       {"gelu_grad",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
-      {"gelu", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
+      {"gelu",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"generate_sequence_xpu",
        XPUKernelSet({
            phi::DataType::FLOAT32,
diff --git a/paddle/phi/kernels/xpu/gelu_grad_kernel.cc b/paddle/phi/kernels/xpu/gelu_grad_kernel.cc
@@ -45,4 +45,5 @@ PD_REGISTER_KERNEL(gelu_grad,
                    ALL_LAYOUT,
                    phi::GeluGradKernel,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/gelu_kernel.cc b/paddle/phi/kernels/xpu/gelu_kernel.cc
@@ -38,5 +38,10 @@ void GeluKernel(const Context& dev_ctx,
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    gelu, XPU, ALL_LAYOUT, phi::GeluKernel, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(gelu,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::GeluKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
