From 530f6b797e7547a804c5f34997c22204d9520e50 Mon Sep 17 00:00:00 2001
From: HongyuJia <jiahongyu@baidu.com>
Date: Thu, 1 Sep 2022 10:42:07 +0800
Subject: [PATCH] [phi] Migrate coalesce_tensor XPU kernel to PHI (#45579)

* migrate coalesce_tensor_xpu, test=kunlun

* fix xpu SetConstant, test=kunlun

* fix include error, test=kunlun
---
 paddle/fluid/operators/coalesce_tensor_op.cc | 10 ----------
 paddle/phi/kernels/coalesce_tensor_kernel.cc | 17 +++++++++++++++++
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/paddle/fluid/operators/coalesce_tensor_op.cc b/paddle/fluid/operators/coalesce_tensor_op.cc
index bdc0f4799bc..3e77bd91baf 100644
--- a/paddle/fluid/operators/coalesce_tensor_op.cc
+++ b/paddle/fluid/operators/coalesce_tensor_op.cc
@@ -529,16 +529,6 @@ REGISTER_OP_CUDA_KERNEL(
     ops::CoalesceTensorOpKernel<paddle::platform::NPUDeviceContext, double>);
 #endif
 
-#ifdef PADDLE_WITH_XPU
-REGISTER_OP_XPU_KERNEL(
-    coalesce_tensor,
-    ops::CoalesceTensorOpKernel<paddle::platform::XPUDeviceContext,
-                                plat::float16>,
-    ops::CoalesceTensorOpKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::CoalesceTensorOpKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::CoalesceTensorOpKernel<paddle::platform::XPUDeviceContext, double>);
-#endif
-
 #if defined(PADDLE_WITH_ASCEND_CL)
 REGISTER_OP_NPU_KERNEL(
     coalesce_tensor,
diff --git a/paddle/phi/kernels/coalesce_tensor_kernel.cc b/paddle/phi/kernels/coalesce_tensor_kernel.cc
index 41548fa03a9..7c3e3aafd22 100644
--- a/paddle/phi/kernels/coalesce_tensor_kernel.cc
+++ b/paddle/phi/kernels/coalesce_tensor_kernel.cc
@@ -22,6 +22,10 @@
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
+#ifdef PADDLE_WITH_XPU
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/kernels/funcs/math_function_impl.h"
+#endif
 
 namespace phi {
 
@@ -280,3 +284,16 @@ PD_REGISTER_KERNEL(coalesce_tensor,
   kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
 }
 #endif
+
+#ifdef PADDLE_WITH_XPU
+PD_REGISTER_KERNEL(coalesce_tensor,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::CoalesceTensorKernel,
+                   phi::dtype::float16,
+                   int,
+                   float,
+                   double) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}
+#endif
-- 
GitLab