xpu set_value_grad support (#71047)

QingshuChen · web-flow · commit 4731cae73d0a · 2025-02-08T17:52:42.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -1145,6 +1145,7 @@ XPUOpMap& get_kl3_ops() {
       {"set_value_with_tensor",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::BOOL,
@@ -1153,6 +1154,7 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::INT64,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT16})},
       {"sgd", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"silu_grad",
diff --git a/paddle/phi/kernels/xpu/set_value_grad_kernel.cc b/paddle/phi/kernels/xpu/set_value_grad_kernel.cc
@@ -427,6 +427,7 @@ PD_REGISTER_KERNEL(set_value_grad,
                    phi::SetValueGradKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {}
 
@@ -436,5 +437,6 @@ PD_REGISTER_KERNEL(set_value_with_scalar_grad,
                    phi::SetValueWithScalarGradKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {}
