[XPU] fix registration of elementwise_add (#60252)

houj04 · web-flow · commit 0942bcf511e8 · 2023-12-22T15:05:24.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -256,7 +256,9 @@ XPUOpMap& get_kl3_ops() {
       {"elementwise_add_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
-                     phi::DataType::BFLOAT16})},
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::INT64,
+                     phi::DataType::INT32})},
       {"elementwise_add",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
diff --git a/paddle/phi/kernels/legacy/xpu/elementwise_add_kernel.cc b/paddle/phi/kernels/legacy/xpu/elementwise_add_kernel.cc
@@ -55,6 +55,7 @@ PD_REGISTER_KERNEL(add_raw,
                    ALL_LAYOUT,
                    phi::AddRawKernel,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    float,
                    int,
                    int64_t) {}
diff --git a/paddle/phi/kernels/xpu/elementwise_add_grad_kernel.cc b/paddle/phi/kernels/xpu/elementwise_add_grad_kernel.cc
@@ -104,5 +104,7 @@ PD_REGISTER_KERNEL(add_grad,
                    ALL_LAYOUT,
                    phi::AddGradKernel,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    float,
-                   phi::dtype::bfloat16) {}
+                   int,
+                   int64_t) {}
