Move compare_raw_kernel to legacy (#53928)

* Move compare_raw_kernel to legacy

* fix

* Update compare_kernel.cc

* Move compare_raw_kernel to legacy

paddle/phi/kernels/compare_kernel.h

@@ -18,17 +18,11 @@ limitations under the License. */
 
 namespace phi {
 
-#define DECALRE_COMPARE_KERNEL(name)         \
-  template <typename T, typename Context>    \
-  void name##RawKernel(const Context& ctx,   \
-                       const DenseTensor& x, \
-                       const DenseTensor& y, \
-                       int axis,             \
-                       DenseTensor* out);    \
-  template <typename T, typename Context>    \
-  void name##Kernel(const Context& ctx,      \
-                    const DenseTensor& x,    \
-                    const DenseTensor& y,    \
+#define DECALRE_COMPARE_KERNEL(name)      \
+  template <typename T, typename Context> \
+  void name##Kernel(const Context& ctx,   \
+                    const DenseTensor& x, \
+                    const DenseTensor& y, \
                     DenseTensor* out);
 
 DECALRE_COMPARE_KERNEL(LessThan)

paddle/phi/kernels/cpu/compare_kernel.cc

@@ -97,20 +97,6 @@ PD_REGISTER_KERNEL(equal_all,
                      phi::dtype::float16,                 \
                      phi::dtype::bfloat16) {              \
     kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
-  }                                                       \
-  PD_REGISTER_KERNEL(name##_raw,                          \
-                     CPU,                                 \
-                     ALL_LAYOUT,                          \
-                     phi::func##RawKernel,                \
-                     bool,                                \
-                     int16_t,                             \
-                     int,                                 \
-                     int64_t,                             \
-                     float,                               \
-                     double,                              \
-                     phi::dtype::float16,                 \
-                     phi::dtype::bfloat16) {              \
-    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
   }
 PD_REGISTER_COMPARE_KERNEL(less_than, LessThan)
 PD_REGISTER_COMPARE_KERNEL(less_equal, LessEqual)

paddle/phi/kernels/impl/compare_kernel_impl.h

@@ -37,21 +37,13 @@ inline void CompareAllKernelImpl(const Context& ctx,
                                  DenseTensor* out);
 
 #define DEFINE_COMPARE_KERNEL(name, functor, inverse_functor)      \
-  template <typename T, typename Context>                          \
-  void name##RawKernel(const Context& ctx,                         \
-                       const DenseTensor& x,                       \
-                       const DenseTensor& y,                       \
-                       int axis,                                   \
-                       DenseTensor* out) {                         \
-    CompareKernelImpl<T, Context, functor<T>, inverse_functor<T>>( \
-        ctx, x, y, axis, out);                                     \
-  }                                                                \
   template <typename T, typename Context>                          \
   void name##Kernel(const Context& ctx,                            \
                     const DenseTensor& x,                          \
                     const DenseTensor& y,                          \
                     DenseTensor* out) {                            \
-    name##RawKernel<T, Context>(ctx, x, y, -1, out);               \
+    CompareKernelImpl<T, Context, functor<T>, inverse_functor<T>>( \
+        ctx, x, y, -1, out);                                       \
   }
 
 DEFINE_COMPARE_KERNEL(LessThan,

paddle/phi/kernels/kps/compare_kernel.cu

@@ -114,30 +114,6 @@ PD_REGISTER_KERNEL(not_equal, KPS, ALL_LAYOUT, phi::NotEqualKernel, int) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
-PD_REGISTER_KERNEL(
-    less_than_raw, KPS, ALL_LAYOUT, phi::LessThanRawKernel, int) {
-  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
-}
-PD_REGISTER_KERNEL(
-    less_equal_raw, KPS, ALL_LAYOUT, phi::LessEqualRawKernel, int) {
-  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
-}
-PD_REGISTER_KERNEL(
-    greater_than_raw, KPS, ALL_LAYOUT, phi::GreaterThanRawKernel, int) {
-  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
-}
-PD_REGISTER_KERNEL(
-    greater_equal_raw, KPS, ALL_LAYOUT, phi::GreaterEqualRawKernel, int) {
-  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
-}
-PD_REGISTER_KERNEL(equal_raw, KPS, ALL_LAYOUT, phi::EqualRawKernel, int) {
-  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
-}
-PD_REGISTER_KERNEL(
-    not_equal_raw, KPS, ALL_LAYOUT, phi::NotEqualRawKernel, int) {
-  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
-}
-
 #else
 
 PD_REGISTER_KERNEL(equal_all,
@@ -166,20 +142,6 @@ PD_REGISTER_KERNEL(equal_all,
                      phi::dtype::float16,                 \
                      phi::dtype::bfloat16) {              \
     kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
-  }                                                       \
-  PD_REGISTER_KERNEL(name##_raw,                          \
-                     KPS,                                 \
-                     ALL_LAYOUT,                          \
-                     phi::func##RawKernel,                \
-                     bool,                                \
-                     int16_t,                             \
-                     int,                                 \
-                     int64_t,                             \
-                     float,                               \
-                     double,                              \
-                     phi::dtype::float16,                 \
-                     phi::dtype::bfloat16) {              \
-    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
   }
 
 PD_REGISTER_COMPARE_KERNEL(less_than, LessThan)

paddle/phi/kernels/legacy/cpu/compare_kernel.cc

+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/compare_kernel.h"
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/impl/compare_kernel_impl.h"
+
+namespace phi {
+
+template <typename T,
+          typename Context,
+          typename Functor,
+          typename InverseFunctor>
+inline void CompareRawKernelImpl(const Context& ctx,
+                                 const DenseTensor& x,
+                                 const DenseTensor& y,
+                                 int axis,
+                                 DenseTensor* out) {
+  ctx.template Alloc<bool>(out);
+  if (x.dims().size() >= y.dims().size()) {
+    funcs::ElementwiseCompute<Functor, T, bool>(
+        ctx, x, y, Functor(), out, axis);
+  } else {
+    funcs::ElementwiseCompute<InverseFunctor, T, bool>(
+        ctx, x, y, InverseFunctor(), out, axis);
+  }
+}
+
+template <typename T, typename Context>
+void LessThanRawKernel(const Context& ctx,
+                       const DenseTensor& x,
+                       const DenseTensor& y,
+                       int axis,
+                       DenseTensor* out) {
+  CompareRawKernelImpl<T,
+                       Context,
+                       funcs::LessThanFunctor<T>,
+                       funcs::GreaterThanFunctor<T>>(ctx, x, y, axis, out);
+}
+
+template <typename T, typename Context>
+void LessEqualRawKernel(const Context& ctx,
+                        const DenseTensor& x,
+                        const DenseTensor& y,
+                        int axis,
+                        DenseTensor* out) {
+  CompareRawKernelImpl<T,
+                       Context,
+                       funcs::LessEqualFunctor<T>,
+                       funcs::GreaterEqualFunctor<T>>(ctx, x, y, axis, out);
+}
+
+template <typename T, typename Context>
+void GreaterThanRawKernel(const Context& ctx,
+                          const DenseTensor& x,
+                          const DenseTensor& y,
+                          int axis,
+                          DenseTensor* out) {
+  CompareRawKernelImpl<T,
+                       Context,
+                       funcs::GreaterThanFunctor<T>,
+                       funcs::LessThanFunctor<T>>(ctx, x, y, axis, out);
+}
+template <typename T, typename Context>
+void GreaterEqualRawKernel(const Context& ctx,
+                           const DenseTensor& x,
+                           const DenseTensor& y,
+                           int axis,
+                           DenseTensor* out) {
+  CompareRawKernelImpl<T,
+                       Context,
+                       funcs::GreaterEqualFunctor<T>,
+                       funcs::LessEqualFunctor<T>>(ctx, x, y, axis, out);
+}
+template <typename T, typename Context>
+void EqualRawKernel(const Context& ctx,
+                    const DenseTensor& x,
+                    const DenseTensor& y,
+                    int axis,
+                    DenseTensor* out) {
+  CompareRawKernelImpl<T,
+                       Context,
+                       funcs::EqualFunctor<T>,
+                       funcs::EqualFunctor<T>>(ctx, x, y, axis, out);
+}
+template <typename T, typename Context>
+void NotEqualRawKernel(const Context& ctx,
+                       const DenseTensor& x,
+                       const DenseTensor& y,
+                       int axis,
+                       DenseTensor* out) {
+  CompareRawKernelImpl<T,
+                       Context,
+                       funcs::NotEqualFunctor<T>,
+                       funcs::NotEqualFunctor<T>>(ctx, x, y, axis, out);
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(less_than_raw,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::LessThanRawKernel,
+                   bool,
+                   int16_t,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+
+#define PD_REGISTER_COMPARE_RAW_KERNEL(name, func)        \
+  PD_REGISTER_KERNEL(name##_raw,                          \
+                     CPU,                                 \
+                     ALL_LAYOUT,                          \
+                     phi::func##RawKernel,                \
+                     bool,                                \
+                     int16_t,                             \
+                     int,                                 \
+                     int64_t,                             \
+                     float,                               \
+                     double,                              \
+                     phi::dtype::float16,                 \
+                     phi::dtype::bfloat16) {              \
+    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
+  }
+PD_REGISTER_COMPARE_RAW_KERNEL(less_equal, LessEqual)
+PD_REGISTER_COMPARE_RAW_KERNEL(greater_than, GreaterThan)
+PD_REGISTER_COMPARE_RAW_KERNEL(greater_equal, GreaterEqual)
+PD_REGISTER_COMPARE_RAW_KERNEL(equal, Equal)
+PD_REGISTER_COMPARE_RAW_KERNEL(not_equal, NotEqual)

paddle/phi/kernels/legacy/kps/compare_kernel.cu

+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/broadcast_function.h"
+#include "paddle/phi/kernels/impl/compare_kernel_impl.h"
+
+#ifdef PADDLE_WITH_XPU_KP
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#else
+#include <thrust/fill.h>
+
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/compare_kernel.h"
+#include "paddle/phi/kernels/funcs/elementwise_base.h"
+#include "paddle/phi/kernels/gpu/reduce.h"
+#include "paddle/phi/kernels/primitive/functor_primitives.h"
+#endif
+
+namespace phi {
+
+template <typename T>
+struct BitwiseAdd {
+  // Bitwise add operator, returns <tt>a + b</tt>
+  inline T initial() { return static_cast<T>(true); }
+
+  __host__ __device__ __forceinline__ T operator()(const T& a,
+                                                   const T& b) const {
+    return a & b;
+  }
+};
+
+template <typename T,
+          typename Context,
+          typename Functor,
+          typename InverseFunctor>
+inline void CompareCudaRawKernelImpl(const Context& ctx,
+                                     const DenseTensor& x,
+                                     const DenseTensor& y,
+                                     int axis,
+                                     DenseTensor* out) {
+  ctx.template Alloc<bool>(out);
+  std::vector<const DenseTensor*> ins{&x, &y};
+  std::vector<DenseTensor*> outs{out};
+  funcs::BroadcastKernel<bool>(ctx, ins, &outs, Functor(), axis);
+}
+
+template <typename T, typename Context>
+void LessThanRawKernel(const Context& ctx,
+                       const DenseTensor& x,
+                       const DenseTensor& y,
+                       int axis,
+                       DenseTensor* out) {
+  CompareCudaRawKernelImpl<T,
+                           Context,
+                           funcs::LessThanFunctor<T>,
+                           funcs::GreaterThanFunctor<T>>(ctx, x, y, axis, out);
+}
+
+template <typename T, typename Context>
+void LessEqualRawKernel(const Context& ctx,
+                        const DenseTensor& x,
+                        const DenseTensor& y,
+                        int axis,
+                        DenseTensor* out) {
+  CompareCudaRawKernelImpl<T,
+                           Context,
+                           funcs::LessEqualFunctor<T>,
+                           funcs::GreaterEqualFunctor<T>>(ctx, x, y, axis, out);
+}
+
+template <typename T, typename Context>
+void GreaterThanRawKernel(const Context& ctx,
+                          const DenseTensor& x,
+                          const DenseTensor& y,
+                          int axis,
+                          DenseTensor* out) {
+  CompareCudaRawKernelImpl<T,
+                           Context,
+                           funcs::GreaterThanFunctor<T>,
+                           funcs::LessThanFunctor<T>>(ctx, x, y, axis, out);
+}
+template <typename T, typename Context>
+void GreaterEqualRawKernel(const Context& ctx,
+                           const DenseTensor& x,
+                           const DenseTensor& y,
+                           int axis,
+                           DenseTensor* out) {
+  CompareCudaRawKernelImpl<T,
+                           Context,
+                           funcs::GreaterEqualFunctor<T>,
+                           funcs::LessEqualFunctor<T>>(ctx, x, y, axis, out);
+}
+template <typename T, typename Context>
+void EqualRawKernel(const Context& ctx,
+                    const DenseTensor& x,
+                    const DenseTensor& y,
+                    int axis,
+                    DenseTensor* out) {
+  CompareCudaRawKernelImpl<T,
+                           Context,
+                           funcs::EqualFunctor<T>,
+                           funcs::EqualFunctor<T>>(ctx, x, y, axis, out);
+}
+template <typename T, typename Context>
+void NotEqualRawKernel(const Context& ctx,
+                       const DenseTensor& x,
+                       const DenseTensor& y,
+                       int axis,
+                       DenseTensor* out) {
+  CompareCudaRawKernelImpl<T,
+                           Context,
+                           funcs::NotEqualFunctor<T>,
+                           funcs::NotEqualFunctor<T>>(ctx, x, y, axis, out);
+}
+
+}  // namespace phi
+
+#ifdef PADDLE_WITH_XPU_KP
+PD_REGISTER_KERNEL(
+    less_than_raw, KPS, ALL_LAYOUT, phi::LessThanRawKernel, int) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+PD_REGISTER_KERNEL(
+    less_equal_raw, KPS, ALL_LAYOUT, phi::LessEqualRawKernel, int) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+PD_REGISTER_KERNEL(
+    greater_than_raw, KPS, ALL_LAYOUT, phi::GreaterThanRawKernel, int) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+PD_REGISTER_KERNEL(
+    greater_equal_raw, KPS, ALL_LAYOUT, phi::GreaterEqualRawKernel, int) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+PD_REGISTER_KERNEL(equal_raw, KPS, ALL_LAYOUT, phi::EqualRawKernel, int) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+PD_REGISTER_KERNEL(
+    not_equal_raw, KPS, ALL_LAYOUT, phi::NotEqualRawKernel, int) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+
+#else
+
+PD_REGISTER_KERNEL(less_than_raw,
+                   KPS,
+                   ALL_LAYOUT,
+                   phi::LessThanRawKernel,
+                   bool,
+                   int16_t,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+
+#define PD_REGISTER_COMPARE_RAW_KERNEL(name, func)        \
+  PD_REGISTER_KERNEL(name##_raw,                          \
+                     KPS,                                 \
+                     ALL_LAYOUT,                          \
+                     phi::func##RawKernel,                \
+                     bool,                                \
+                     int16_t,                             \
+                     int,                                 \
+                     int64_t,                             \
+                     float,                               \
+                     double,                              \
+                     phi::dtype::float16,                 \
+                     phi::dtype::bfloat16) {              \
+    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
+  }
+
+PD_REGISTER_COMPARE_RAW_KERNEL(less_equal, LessEqual)
+PD_REGISTER_COMPARE_RAW_KERNEL(greater_than, GreaterThan)
+PD_REGISTER_COMPARE_RAW_KERNEL(greater_equal, GreaterEqual)
+PD_REGISTER_COMPARE_RAW_KERNEL(equal, Equal)
+PD_REGISTER_COMPARE_RAW_KERNEL(not_equal, NotEqual)
+
+#endif

paddle/phi/kernels/legacy/xpu/compare_kernel.cc

+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/compare_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/backends/xpu/xpu_header.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename XPUType, typename Context>
+void XPUCompareRawKernelImpl(const Context& dev_ctx,
+                             const DenseTensor& x,
+                             const DenseTensor& y,
+                             DenseTensor* out,
+                             std::function<int(xpu::Context*,
+                                               const XPUType*,
+                                               const XPUType*,
+                                               bool*,
+                                               const std::vector<int>&,
+                                               const std::vector<int>&)> func) {
+  auto x_shape = vectorize<int>(x.dims());
+  auto y_shape = vectorize<int>(y.dims());
+
+  if (x.dims().size() == 0) {
+    x_shape = std::vector<int>({1});
+  }
+  if (y.dims().size() == 0) {
+    y_shape = std::vector<int>({1});
+  }
+
+  auto x_data = reinterpret_cast<const XPUType*>(x.data<T>());
+  auto y_data = reinterpret_cast<const XPUType*>(y.data<T>());
+  auto* out_data = dev_ctx.template Alloc<bool>(out);
+
+  int ret =
+      func(dev_ctx.x_context(), x_data, y_data, out_data, x_shape, y_shape);
+  PADDLE_ENFORCE_XDNN_SUCCESS(ret, "compare op");
+}
+
+#define DEFINE_XPU_COMPARE_RAW_KERNEL(name, functor)                     \
+  template <typename T, typename Context>                                \
+  void name##RawKernel(const Context& dev_ctx,                           \
+                       const DenseTensor& x,                             \
+                       const DenseTensor& y,                             \
+                       int axis,                                         \
+                       DenseTensor* out) {                               \
+    using XPUType = typename XPUTypeTrait<T>::Type;                      \
+    auto f = [](xpu::Context* ctx,                                       \
+                const XPUType* x,                                        \
+                const XPUType* y,                                        \
+                bool* z,                                                 \
+                const std::vector<int>& xshape,                          \
+                const std::vector<int>& yshape) {                        \
+      return functor(ctx, x, y, z, xshape, yshape);                      \
+    };                                                                   \
+    XPUCompareRawKernelImpl<T, XPUType, Context>(dev_ctx, x, y, out, f); \
+  }
+
+DEFINE_XPU_COMPARE_RAW_KERNEL(Equal, xpu::broadcast_equal<XPUType>)
+DEFINE_XPU_COMPARE_RAW_KERNEL(NotEqual, xpu::broadcast_not_equal<XPUType>)
+DEFINE_XPU_COMPARE_RAW_KERNEL(LessThan, xpu::broadcast_less_than<XPUType>)
+DEFINE_XPU_COMPARE_RAW_KERNEL(LessEqual, xpu::broadcast_less_equal<XPUType>)
+DEFINE_XPU_COMPARE_RAW_KERNEL(GreaterThan, xpu::broadcast_greater_than<XPUType>)
+DEFINE_XPU_COMPARE_RAW_KERNEL(GreaterEqual,
+                              xpu::broadcast_greater_equal<XPUType>)
+
+#undef DEFINE_XPU_COMPARE_RAW_KERNEL
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(less_than_raw,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::LessThanRawKernel,
+                   int,
+                   int64_t,
+                   float,
+                   phi::dtype::float16) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
+}
+
+#define PD_REGISTER_COMPARE_RAW_KERNEL(name, func)        \
+  PD_REGISTER_KERNEL(name##_raw,                          \
+                     XPU,                                 \
+                     ALL_LAYOUT,                          \
+                     phi::func##RawKernel,                \
+                     int,                                 \
+                     int64_t,                             \
+                     float,                               \
+                     phi::dtype::float16,                 \
+                     bool) {                              \
+    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
+  }
+
+PD_REGISTER_COMPARE_RAW_KERNEL(less_equal, LessEqual)
+PD_REGISTER_COMPARE_RAW_KERNEL(greater_than, GreaterThan)
+PD_REGISTER_COMPARE_RAW_KERNEL(greater_equal, GreaterEqual)
+PD_REGISTER_COMPARE_RAW_KERNEL(equal, Equal)
+PD_REGISTER_COMPARE_RAW_KERNEL(not_equal, NotEqual)

paddle/phi/kernels/xpu/compare_kernel.cc

@@ -54,11 +54,10 @@ void XPUCompareKernelImpl(const Context& dev_ctx,
 
 #define DEFINE_XPU_COMPARE_KERNEL(name, functor)                      \
   template <typename T, typename Context>                             \
-  void name##RawKernel(const Context& dev_ctx,                        \
-                       const DenseTensor& x,                          \
-                       const DenseTensor& y,                          \
-                       int axis,                                      \
-                       DenseTensor* out) {                            \
+  void name##Kernel(const Context& dev_ctx,                           \
+                    const DenseTensor& x,                             \
+                    const DenseTensor& y,                             \
+                    DenseTensor* out) {                               \
     using XPUType = typename XPUTypeTrait<T>::Type;                   \
     auto f = [](xpu::Context* ctx,                                    \
                 const XPUType* x,                                     \
@@ -69,13 +68,6 @@ void XPUCompareKernelImpl(const Context& dev_ctx,
       return functor(ctx, x, y, z, xshape, yshape);                   \
     };                                                                \
     XPUCompareKernelImpl<T, XPUType, Context>(dev_ctx, x, y, out, f); \
-  }                                                                   \
-  template <typename T, typename Context>                             \
-  void name##Kernel(const Context& dev_ctx,                           \
-                    const DenseTensor& x,                             \
-                    const DenseTensor& y,                             \
-                    DenseTensor* out) {                               \
-    name##RawKernel<T, Context>(dev_ctx, x, y, -1, out);              \
   }
 
 DEFINE_XPU_COMPARE_KERNEL(Equal, xpu::broadcast_equal<XPUType>)
@@ -100,17 +92,6 @@ PD_REGISTER_KERNEL(less_than,
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
-PD_REGISTER_KERNEL(less_than_raw,
-                   XPU,
-                   ALL_LAYOUT,
-                   phi::LessThanRawKernel,
-                   int,
-                   int64_t,
-                   float,
-                   phi::dtype::float16) {
-  kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
-}
-
 #define PD_REGISTER_COMPARE_KERNEL(name, func)            \
   PD_REGISTER_KERNEL(name,                                \
                      XPU,                                 \
@@ -122,19 +103,7 @@ PD_REGISTER_KERNEL(less_than_raw,
                      phi::dtype::float16,                 \
                      bool) {                              \
     kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
-  }                                                       \
-  PD_REGISTER_KERNEL(name##_raw,                          \
-                     XPU,                                 \
-                     ALL_LAYOUT,                          \
-                     phi::func##RawKernel,                \
-                     int,                                 \
-                     int64_t,                             \
-                     float,                               \
-                     phi::dtype::float16,                 \
-                     bool) {                              \
-    kernel->OutputAt(0).SetDataType(phi::DataType::BOOL); \
   }
-
 PD_REGISTER_COMPARE_KERNEL(less_equal, LessEqual)
 PD_REGISTER_COMPARE_KERNEL(greater_than, GreaterThan)
 PD_REGISTER_COMPARE_KERNEL(greater_equal, GreaterEqual)

test/cpp/prim/test_static_prim.cc

@@ -41,6 +41,7 @@ PD_DECLARE_KERNEL(multiply, CPU, ALL_LAYOUT);
 PD_DECLARE_KERNEL(concat, CPU, ALL_LAYOUT);
 PD_DECLARE_KERNEL(less_equal, CPU, ALL_LAYOUT);
 PD_DECLARE_KERNEL(less_than, CPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(less_than_raw, CPU, ALL_LAYOUT);
 PD_DECLARE_KERNEL(equal, CPU, ALL_LAYOUT);
 PD_DECLARE_KERNEL(not_equal, CPU, ALL_LAYOUT);
 PD_DECLARE_KERNEL(greater_equal, CPU, ALL_LAYOUT);
@@ -60,6 +61,7 @@ PD_DECLARE_KERNEL(multiply, KPS, ALL_LAYOUT);
 PD_DECLARE_KERNEL(concat, GPU, ALL_LAYOUT);
 PD_DECLARE_KERNEL(less_equal, KPS, ALL_LAYOUT);
 PD_DECLARE_KERNEL(less_than, KPS, ALL_LAYOUT);
+PD_DECLARE_KERNEL(less_than_raw, KPS, ALL_LAYOUT);
 PD_DECLARE_KERNEL(equal, KPS, ALL_LAYOUT);
 PD_DECLARE_KERNEL(not_equal, KPS, ALL_LAYOUT);
 PD_DECLARE_KERNEL(greater_equal, KPS, ALL_LAYOUT);