[xpu]:support equal int64 and transpose int64;test=develop

Author: gitliuyf

lite/kernels/x86/transpose_compute.cc

@@ -34,3 +34,25 @@ REGISTER_LITE_KERNEL(transpose2,
     .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kX86))})
     .BindOutput("XShape", {LiteType::GetTensorTy(TARGET(kX86))})
     .Finalize();
+
+REGISTER_LITE_KERNEL(transpose,
+                     kX86,
+                     kFloat,
+                     kNCHW,
+                     paddle::lite::kernels::x86::TransposeCompute<int64_t>,
+                     int64)
+    .BindInput("X", {LiteType::GetTensorTy(TARGET(kX86), PRECISION(kInt64))})
+    .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kX86), PRECISION(kInt64))})
+    .Finalize();
+
+REGISTER_LITE_KERNEL(transpose2,
+                     kX86,
+                     kFloat,
+                     kNCHW,
+                     paddle::lite::kernels::x86::Transpose2Compute<int64_t>,
+                     int64)
+    .BindInput("X", {LiteType::GetTensorTy(TARGET(kX86), PRECISION(kInt64))})
+    .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kX86), PRECISION(kInt64))})
+    .BindOutput("XShape",
+                {LiteType::GetTensorTy(TARGET(kX86), PRECISION(kInt64))})
+    .Finalize();

lite/kernels/xpu/compare_compute.cc

@@ -34,6 +34,18 @@ struct LessThanFunctor {
   }
 };
 
+template <typename T>
+struct EqualFunctor {
+  inline int operator()(xdnn::Context* ctx,
+                        const T* x,
+                        const T* y,
+                        bool* z,
+                        const std::vector<int>& xshape,
+                        const std::vector<int>& yshape) const {
+    return xdnn::broadcast_equal<T>(ctx, x, y, z, xshape, yshape);
+  }
+};
+
 template <PrecisionType PType, class T, class Functor>
 void CompareCompute<PType, T, Functor>::Run() {
   auto& param = this->template Param<operators::CompareParam>();
@@ -152,3 +164,63 @@ REGISTER_LITE_KERNEL(less_than, kXPU, kFloat, kAny, less_than_int64, int64)
                                        DATALAYOUT(kAny))})
     .BindPaddleOpVersion("less_than", 1)
     .Finalize();
+
+using equal_float = paddle::lite::kernels::xpu::CompareCompute<
+    PRECISION(kFloat),
+    float,
+    paddle::lite::kernels::xpu::EqualFunctor<float>>;
+REGISTER_LITE_KERNEL(equal, kXPU, kFloat, kAny, equal_float, def)
+    .BindInput("X",
+               {LiteType::GetTensorTy(TARGET(kXPU),
+                                      PRECISION(kFloat),
+                                      DATALAYOUT(kAny))})
+    .BindInput("Y",
+               {LiteType::GetTensorTy(TARGET(kXPU),
+                                      PRECISION(kFloat),
+                                      DATALAYOUT(kAny))})
+    .BindOutput("Out",
+                {LiteType::GetTensorTy(TARGET(kXPU),
+                                       PRECISION(kBool),
+                                       DATALAYOUT(kAny))})
+    .BindPaddleOpVersion("equal", 1)
+    .Finalize();
+
+using equal_int32 = paddle::lite::kernels::xpu::CompareCompute<
+    PRECISION(kFloat),
+    int,
+    paddle::lite::kernels::xpu::EqualFunctor<int>>;
+REGISTER_LITE_KERNEL(equal, kXPU, kFloat, kAny, equal_int32, int32)
+    .BindInput("X",
+               {LiteType::GetTensorTy(TARGET(kXPU),
+                                      PRECISION(kInt32),
+                                      DATALAYOUT(kAny))})
+    .BindInput("Y",
+               {LiteType::GetTensorTy(TARGET(kXPU),
+                                      PRECISION(kInt32),
+                                      DATALAYOUT(kAny))})
+    .BindOutput("Out",
+                {LiteType::GetTensorTy(TARGET(kXPU),
+                                       PRECISION(kBool),
+                                       DATALAYOUT(kAny))})
+    .BindPaddleOpVersion("equal", 1)
+    .Finalize();
+
+using euqal_int64 = paddle::lite::kernels::xpu::CompareCompute<
+    PRECISION(kFloat),
+    int64_t,
+    paddle::lite::kernels::xpu::EqualFunctor<int64_t>>;
+REGISTER_LITE_KERNEL(equal, kXPU, kFloat, kAny, euqal_int64, int64)
+    .BindInput("X",
+               {LiteType::GetTensorTy(TARGET(kXPU),
+                                      PRECISION(kInt64),
+                                      DATALAYOUT(kAny))})
+    .BindInput("Y",
+               {LiteType::GetTensorTy(TARGET(kXPU),
+                                      PRECISION(kInt64),
+                                      DATALAYOUT(kAny))})
+    .BindOutput("Out",
+                {LiteType::GetTensorTy(TARGET(kXPU),
+                                       PRECISION(kBool),
+                                       DATALAYOUT(kAny))})
+    .BindPaddleOpVersion("equal", 1)
+    .Finalize();

lite/kernels/xpu/transpose_compute.cc

@@ -22,7 +22,8 @@ namespace lite {
 namespace kernels {
 namespace xpu {
 
-void TransposeCompute::Run() {
+template <class T>
+void TransposeCompute<T>::Run() {
   auto& param = this->Param<param_t>();
   auto& ctx = this->ctx_->As<XPUContext>();
   auto x = param.x;
@@ -38,10 +39,10 @@ void TransposeCompute::Run() {
   for (int i = 0; i < ndims; ++i) {
     x_shape_host[i] = x_dims[i];
   }
-  int r =
-      xdnn::transpose<float>(ctx.GetRawContext(),
-                             x->data<float>(),
-                             param.output->mutable_data<float>(TARGET(kXPU)),
+
+  int r = xdnn::transpose<T>(ctx.GetRawContext(),
+                             x->data<T>(),
+                             param.output->mutable_data<T>(TARGET(kXPU)),
                              x_shape_host,
                              axis);
   CHECK_EQ(r, 0);
@@ -56,7 +57,7 @@ REGISTER_LITE_KERNEL(transpose,
                      kXPU,
                      kFloat,
                      kNCHW,
-                     paddle::lite::kernels::xpu::TransposeCompute,
+                     paddle::lite::kernels::xpu::TransposeCompute<float>,
                      def)
     .BindInput("X", {LiteType::GetTensorTy(TARGET(kXPU))})
     .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kXPU))})
@@ -66,17 +67,18 @@ REGISTER_LITE_KERNEL(transpose2,
                      kXPU,
                      kFloat,
                      kNCHW,
-                     paddle::lite::kernels::xpu::TransposeCompute,
+                     paddle::lite::kernels::xpu::TransposeCompute<float>,
                      def)
     .BindInput("X", {LiteType::GetTensorTy(TARGET(kXPU))})
     .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kXPU))})
     .BindOutput("XShape", {LiteType::GetTensorTy(TARGET(kHost))})
     .Finalize();
+
 REGISTER_LITE_KERNEL(transpose2,
                      kXPU,
                      kFloat,
                      kNCHW,
-                     paddle::lite::kernels::xpu::TransposeCompute,
+                     paddle::lite::kernels::xpu::TransposeCompute<int64_t>,
                      def_int64)
     .BindInput("X", {LiteType::GetTensorTy(TARGET(kXPU), PRECISION(kInt64))})
     .BindOutput("Out", {LiteType::GetTensorTy(TARGET(kXPU), PRECISION(kInt64))})

lite/kernels/xpu/transpose_compute.h

@@ -21,6 +21,7 @@ namespace lite {
 namespace kernels {
 namespace xpu {
 
+template <class T>
 class TransposeCompute : public KernelLite<TARGET(kXPU), PRECISION(kFloat)> {
  public:
   using param_t = operators::TransposeParam;