[operator migration] Migrate kernel of unique consecutive op. (#44228)

paddle/fluid/operators/unique_consecutive_op.cc

@@ -12,8 +12,7 @@ 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/fluid/operators/unique_consecutive_op.h"
-
+#include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/framework/op_version_registry.h"
 
 namespace paddle {
@@ -118,11 +117,6 @@ namespace ops = paddle::operators;
 REGISTER_OP_WITHOUT_GRADIENT(unique_consecutive,
                              ops::UniqueConsecutiveOp,
                              ops::UniqueConsecutiveOpMaker);
-REGISTER_OP_CPU_KERNEL(unique_consecutive,
-                       ops::UniqueConsecutiveKernel<phi::CPUContext, float>,
-                       ops::UniqueConsecutiveKernel<phi::CPUContext, double>,
-                       ops::UniqueConsecutiveKernel<phi::CPUContext, int32_t>,
-                       ops::UniqueConsecutiveKernel<phi::CPUContext, int64_t>);
 REGISTER_OP_VERSION(unique_consecutive)
     .AddCheckpoint(
         R"ROC(

paddle/fluid/operators/unique_consecutive_op.h → paddle/phi/kernels/cpu/unique_consecutive_functor.h

-/* Copyright (c) 2019 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. */
+// Copyright (c) 2022 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.
 
 #pragma once
-#include <algorithm>
-#include <cmath>
-#include <numeric>
-#include <set>
-#include <unordered_map>
-#include <utility>
-#include <vector>
 
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/fluid/operators/math/concat_and_split.h"
-#include "paddle/fluid/operators/transpose_op.h"
-#include "paddle/fluid/operators/unique_op.h"
+#include "paddle/fluid/framework/tensor_util.h"
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/unique_functor.h"
+
+namespace phi {
 
-namespace paddle {
-namespace operators {
-template <typename InT, typename IndexT>
-static void UniqueConsecutiveFlattendTensor(
-    const framework::ExecutionContext& context,
-    const framework::Tensor& in,
-    framework::Tensor* out,
+template <typename InT, typename IndexT, typename Context>
+static void UniqueConsecutiveFlattenedTensor(const Context& context,
+                                             const DenseTensor& in,
+                                             DenseTensor* out,
                                              bool return_inverse,
-    bool return_counts) {
+                                             bool return_counts,
+                                             DenseTensor* inverse,
+                                             DenseTensor* count) {
   const InT* in_data = in.data<InT>();
   std::vector<InT> out_vec(in.numel());
   std::vector<IndexT> inverse_vec(in.numel());
@@ -65,27 +60,57 @@ static void UniqueConsecutiveFlattendTensor(
   out_vec.resize(output_size);
 
   out->Resize(phi::make_ddim({output_size}));
-  auto* out_data = out->mutable_data<InT>(context.GetPlace());
+  auto* out_data = context.template Alloc<InT>(out);
   std::copy(out_vec.begin(), out_vec.end(), out_data);
 
   if (return_inverse) {
-    auto* inverse = context.Output<framework::Tensor>("Index");
     inverse->Resize(phi::make_ddim({in.numel()}));
-    auto* inverse_data = inverse->mutable_data<IndexT>(context.GetPlace());
+    auto* inverse_data = context.template Alloc<IndexT>(inverse);
     std::copy(inverse_vec.begin(), inverse_vec.end(), inverse_data);
   }
 
   if (return_counts) {
-    auto* count = context.Output<framework::Tensor>("Counts");
     count->Resize(phi::make_ddim({out->numel()}));
-    auto* counts_data = count->mutable_data<IndexT>(context.GetPlace());
+    auto* counts_data = context.template Alloc<IndexT>(count);
     std::copy(counts_vec.begin(), counts_vec.end(), counts_data);
   }
 }
 
-template <class ForwardIt, typename InT, typename IndexT>
+template <typename Context, typename InT>
+struct UniqueConsecutiveFlattenedTensorFunctor {
+  const Context& ctx_;
+  const DenseTensor& in_;
+  DenseTensor* out_;
+  const bool return_inverse_;
+  const bool return_counts_;
+  DenseTensor* inverse_;
+  DenseTensor* count_;
+
+  UniqueConsecutiveFlattenedTensorFunctor(const Context& context,
+                                          const DenseTensor& in,
+                                          DenseTensor* out,
+                                          bool return_inverse,
+                                          bool return_counts,
+                                          DenseTensor* inverse,
+                                          DenseTensor* count)
+      : ctx_(context),
+        in_(in),
+        out_(out),
+        return_inverse_(return_inverse),
+        return_counts_(return_counts),
+        inverse_(inverse),
+        count_(count) {}
+
+  template <typename IndexT>
+  void apply() const {
+    UniqueConsecutiveFlattenedTensor<InT, IndexT, Context>(
+        ctx_, in_, out_, return_inverse_, return_counts_, inverse_, count_);
+  }
+};
+
+template <typename Context, class ForwardIt, typename InT, typename IndexT>
 static ForwardIt UniqueConsecutiveDimImpl(
-    const framework::ExecutionContext& context,
+    const Context& context,
     ForwardIt first,
     ForwardIt last,
     const std::vector<IndexT>& sorted_indices_vec,
@@ -104,7 +129,7 @@ static ForwardIt UniqueConsecutiveDimImpl(
   while (++first != last) {
     int64_t idx_first = std::distance(begin, first);
     int64_t idx_result = std::distance(begin, result);
-    if (!Equal<InT>(*result, *first)) {
+    if (!phi::funcs::Equal<InT>(*result, *first)) {
       if (++result != first) {
         *result = std::move(*first);
       }
@@ -116,13 +141,15 @@ static ForwardIt UniqueConsecutiveDimImpl(
   return ++result;
 }
 
-template <typename DeviceContext, typename InT, typename IndexT>
-static void UniqueConsecutiveDim(const framework::ExecutionContext& context,
-                                 const framework::Tensor& in,
-                                 framework::Tensor* out,
+template <typename Context, typename InT, typename IndexT>
+static void UniqueConsecutiveDim(const Context& context,
+                                 const DenseTensor& in,
+                                 DenseTensor* out,
                                  bool return_inverse,
                                  bool return_counts,
-                                 int axis) {
+                                 int axis,
+                                 DenseTensor* inverse,
+                                 DenseTensor* count) {
   // transpose tensor: eg. axis=1, [dim0, dim1, dim2] -> [dim1, dim0, dim2]
   std::vector<int> permute(in.dims().size());
   std::iota(permute.begin(), permute.end(), 0);
@@ -131,15 +158,14 @@ static void UniqueConsecutiveDim(const framework::ExecutionContext& context,
   std::vector<int64_t> in_trans_dims_vec(phi::vectorize(in.dims()));
   in_trans_dims_vec[axis] = in.dims()[0];
   in_trans_dims_vec[0] = in.dims()[axis];
-  framework::Tensor in_trans;
-  framework::DDim in_trans_dims = phi::make_ddim(in_trans_dims_vec);
+  DenseTensor in_trans;
+  DDim in_trans_dims = phi::make_ddim(in_trans_dims_vec);
   in_trans.Resize(in_trans_dims);
-  in_trans.mutable_data<InT>(context.GetPlace());
-  auto& dev_ctx = context.template device_context<DeviceContext>();
-  TransCompute<DeviceContext, InT>(
-      in.dims().size(), dev_ctx, in, &in_trans, permute);
+  context.template Alloc<InT>(&in_trans);
+  phi::funcs::TransCompute<Context, InT>(
+      in.dims().size(), context, in, &in_trans, permute);
   // reshape tensor: eg. [dim1, dim0, dim2] -> [dim1, dim0*dim2]
-  framework::DDim in_trans_flat_dims = phi::flatten_to_2d(in_trans_dims, 1);
+  DDim in_trans_flat_dims = phi::flatten_to_2d(in_trans_dims, 1);
   in_trans.Resize(in_trans_flat_dims);
 
   std::vector<IndexT> sorted_indices_vec(in_trans.dims()[0]);
@@ -148,21 +174,21 @@ static void UniqueConsecutiveDim(const framework::ExecutionContext& context,
   const InT* in_trans_data = in_trans.data<InT>();
 
   // sort tensor according to indices
-  framework::Tensor input_sorted;
+  DenseTensor input_sorted;
   input_sorted.Resize(in_trans_dims);
-  input_sorted.mutable_data<InT>(context.GetPlace());
+  context.template Alloc<InT>(&input_sorted);
   InT* input_sorted_data = input_sorted.data<InT>();
   for (size_t i = 0; i < sorted_indices_vec.size(); ++i) {
     memcpy(input_sorted_data + i * col,
            in_trans_data + static_cast<int64_t>(sorted_indices_vec[i]) * col,
            col * sizeof(InT));
   }
-  std::vector<framework::Tensor> input_unbind = Unbind(input_sorted);
+  std::vector<DenseTensor> input_unbind = phi::funcs::Unbind(input_sorted);
   std::vector<IndexT> inverse_vec(sorted_indices_vec.size(), 0);
   std::vector<IndexT> counts_vec(sorted_indices_vec.size(), 0);
-  auto last =
-      UniqueConsecutiveDimImpl<std::vector<framework::Tensor>::iterator, InT>(
-          context,
+  auto last = UniqueConsecutiveDimImpl<Context,
+                                       std::vector<DenseTensor>::iterator,
+                                       InT>(context,
                                             input_unbind.begin(),
                                             input_unbind.end(),
                                             sorted_indices_vec,
@@ -171,117 +197,65 @@ static void UniqueConsecutiveDim(const framework::ExecutionContext& context,
   input_unbind.erase(last, input_unbind.end());
   counts_vec.erase(counts_vec.begin() + input_unbind.size(), counts_vec.end());
 
-  math::ConcatFunctor<DeviceContext, InT> concat_functor;
-  framework::Tensor out_trans;
+  phi::funcs::ConcatFunctor<Context, InT> concat_functor;
+  DenseTensor out_trans;
   std::vector<int64_t> out_trans_dims_vec = in_trans_dims_vec;
   out_trans_dims_vec[0] = input_unbind.size();
   out_trans.Resize(phi::make_ddim(out_trans_dims_vec));
-  out_trans.mutable_data<InT>(context.GetPlace());
+  context.template Alloc<InT>(&out_trans);
   std::swap(out_trans_dims_vec[0], out_trans_dims_vec[axis]);
   out->Resize(phi::make_ddim(out_trans_dims_vec));
-  out->mutable_data<InT>(context.GetPlace());
-  concat_functor(dev_ctx, input_unbind, 0, &out_trans);
-  TransCompute<DeviceContext, InT>(
-      out_trans.dims().size(), dev_ctx, out_trans, out, permute);
+  context.template Alloc<InT>(out);
+  concat_functor(context, input_unbind, 0, &out_trans);
+  phi::funcs::TransCompute<Context, InT>(
+      out_trans.dims().size(), context, out_trans, out, permute);
   if (return_inverse) {
-    auto* inverse = context.Output<framework::Tensor>("Index");
-    framework::TensorFromVector(inverse_vec, context.device_context(), inverse);
+    paddle::framework::TensorFromVector(inverse_vec, context, inverse);
   }
   if (return_counts) {
-    auto* count = context.Output<framework::Tensor>("Counts");
-    framework::TensorFromVector(counts_vec, context.device_context(), count);
+    paddle::framework::TensorFromVector(counts_vec, context, count);
   }
 }
 
-template <typename DeviceContext, typename InT>
-struct UniqueConsecutiveFlattendTensorFunctor {
-  const framework::ExecutionContext& ctx_;
-  const framework::Tensor& in_;
-  framework::Tensor* out_;
-  const bool return_inverse_;
-  const bool return_counts_;
-
-  UniqueConsecutiveFlattendTensorFunctor(
-      const framework::ExecutionContext& context,
-      const framework::Tensor& in,
-      framework::Tensor* out,
-      bool return_inverse,
-      bool return_counts)
-      : ctx_(context),
-        in_(in),
-        out_(out),
-        return_inverse_(return_inverse),
-        return_counts_(return_counts) {}
-
-  template <typename IndexT>
-  void apply() const {
-    UniqueConsecutiveFlattendTensor<InT, IndexT>(
-        ctx_, in_, out_, return_inverse_, return_counts_);
-  }
-};
-
-template <typename DeviceContext, typename InT>
+template <typename Context, typename InT>
 struct UniqueConsecutiveDimFunctor {
-  const framework::ExecutionContext& ctx_;
-  const framework::Tensor& in_;
-  framework::Tensor* out_;
+  const Context& ctx_;
+  const DenseTensor& in_;
+  DenseTensor* out_;
   const int axis_;
   const bool return_inverse_;
   const bool return_counts_;
-  UniqueConsecutiveDimFunctor(const framework::ExecutionContext& context,
-                              const framework::Tensor& in,
-                              framework::Tensor* out,
+  DenseTensor* inverse_;
+  DenseTensor* count_;
+
+  UniqueConsecutiveDimFunctor(const Context& context,
+                              const DenseTensor& in,
+                              DenseTensor* out,
                               const int axis,
                               bool return_inverse,
-                              bool return_counts)
+                              bool return_counts,
+                              DenseTensor* inverse,
+                              DenseTensor* count)
       : ctx_(context),
         in_(in),
         out_(out),
         axis_(axis),
         return_inverse_(return_inverse),
-        return_counts_(return_counts) {}
+        return_counts_(return_counts),
+        inverse_(inverse),
+        count_(count) {}
 
   template <typename IndexT>
   void apply() const {
-    UniqueConsecutiveDim<DeviceContext, InT, IndexT>(
-        ctx_, in_, out_, return_inverse_, return_counts_, axis_);
+    UniqueConsecutiveDim<Context, InT, IndexT>(ctx_,
+                                               in_,
+                                               out_,
+                                               return_inverse_,
+                                               return_counts_,
+                                               axis_,
+                                               inverse_,
+                                               count_);
   }
 };
-template <typename DeviceContext, typename T>
-class UniqueConsecutiveKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& context) const override {
-    auto* x = context.Input<framework::Tensor>("X");
-    auto* out = context.Output<framework::Tensor>("Out");
-    auto data_type = static_cast<framework::proto::VarType::Type>(
-        context.Attr<int>("dtype"));
-    if (data_type == framework::proto::VarType::INT32) {
-      PADDLE_ENFORCE_LE(
-          x->numel(),
-          INT_MAX,
-          platform::errors::InvalidArgument(
-              "The number of elements in Input(X) should be less than or "
-              "equal to INT_MAX, but received num is %d. Please set `dtype` to "
-              "int64.",
-              x->numel()));
-    }
-    std::vector<int> axis_vec = context.Attr<std::vector<int>>("axis");
-    bool return_inverse = context.Attr<bool>("return_inverse");
-    bool return_counts = context.Attr<bool>("return_counts");
 
-    if (axis_vec.empty()) {
-      framework::VisitDataTypeTiny(
-          data_type,
-          UniqueConsecutiveFlattendTensorFunctor<DeviceContext, T>(
-              context, *x, out, return_inverse, return_counts));
-    } else {
-      int axis = axis_vec[0];
-      framework::VisitDataTypeTiny(
-          data_type,
-          UniqueConsecutiveDimFunctor<DeviceContext, T>(
-              context, *x, out, axis, return_inverse, return_counts));
-    }
-  }
-};
-}  // namespace operators
-}  // namespace paddle
+}  // namespace phi

paddle/phi/kernels/cpu/unique_consecutive_kernel.cc

+// Copyright (c) 2022 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/unique_consecutive_kernel.h"
+#include "paddle/phi/kernels/cpu/unique_consecutive_functor.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/errors.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+#include "paddle/fluid/framework/data_type.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void UniqueConsecutiveKernel(const Context& dev_ctx,
+                             const DenseTensor& x,
+                             bool return_inverse,
+                             bool return_counts,
+                             const std::vector<int>& axis,
+                             int dtype,
+                             DenseTensor* out,
+                             DenseTensor* index,
+                             DenseTensor* counts) {
+  auto data_type = static_cast<paddle::framework::proto::VarType::Type>(dtype);
+  if (data_type == paddle::framework::proto::VarType::INT32) {
+    PADDLE_ENFORCE_LE(
+        x.numel(),
+        INT_MAX,
+        phi::errors::InvalidArgument(
+            "The number of elements in Input(X) should be less than or "
+            "equal to INT_MAX, but received num is %d. Please set `dtype` to "
+            "int64.",
+            x.numel()));
+  }
+
+  if (axis.empty()) {
+    paddle::framework::VisitDataTypeTiny(
+        data_type,
+        UniqueConsecutiveFlattenedTensorFunctor<Context, T>(
+            dev_ctx, x, out, return_inverse, return_counts, index, counts));
+  } else {
+    int valid_axis = axis[0];
+    paddle::framework::VisitDataTypeTiny(
+        data_type,
+        UniqueConsecutiveDimFunctor<Context, T>(dev_ctx,
+                                                x,
+                                                out,
+                                                valid_axis,
+                                                return_inverse,
+                                                return_counts,
+                                                index,
+                                                counts));
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(unique_consecutive,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::UniqueConsecutiveKernel,
+                   float,
+                   double,
+                   int32_t,
+                   int64_t) {}

paddle/fluid/operators/unique_consecutive_op.cu → paddle/phi/kernels/gpu/unique_consecutive_functor.h

-/* Copyright (c) 2019 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. */
+// Copyright (c) 2022 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.
+
+#pragma once
 #include <thrust/adjacent_difference.h>
 #include <thrust/device_vector.h>
 #include <thrust/execution_policy.h>
@@ -22,159 +24,51 @@ limitations under the License. */
 #include <iostream>
 #include <vector>
 
-#include "paddle/fluid/framework/tensor_util.h"            // TensorToVector()
-#include "paddle/fluid/operators/unique_consecutive_op.h"  // TransComute()
-
-namespace paddle {
-namespace operators {
-
-using Tensor = framework::Tensor;
-
-// Binary function 'equal_to'
-template <typename InT>
-struct BinaryEqual {
-  int64_t col;
-  const InT* in_trans_data;
-
-  BinaryEqual(int64_t _col, const InT* _in_trans_data)
-      : col(_col), in_trans_data(_in_trans_data) {}
-
-  __device__ bool operator()(int64_t a, int64_t b) const {
-    for (int64_t i = 0; i < col; ++i) {
-      InT lhs = in_trans_data[i + a * col];
-      InT rhs = in_trans_data[i + b * col];
-      if (lhs != rhs) {
-        return false;
-      }
-    }
-    return true;
-  }
-};
+#include "paddle/fluid/framework/tensor_util.h"
 
-// Binary function 'not_equal_to'
-template <typename InT>
-struct BinaryNotEqual {
-  int64_t col;
-  const InT* in_trans_data;
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/unique_functor.h"
 
-  BinaryNotEqual(int64_t _col, const InT* _in_trans_data)
-      : col(_col), in_trans_data(_in_trans_data) {}
-
-  __device__ bool operator()(int64_t a, int64_t b) const {
-    for (int64_t i = 0; i < col; ++i) {
-      InT lhs = in_trans_data[i + a * col];
-      InT rhs = in_trans_data[i + b * col];
-      if (lhs != rhs) {
-        return true;
-      }
-    }
-    return false;
-  }
-};
-
-// index_select() function for Tensor
-template <typename InT, typename IndexT>
-void IndexSelect(const framework::ExecutionContext& context,
-                 const Tensor& input,
-                 const Tensor& index,
-                 Tensor* output,
-                 int dim) {
-  auto input_dim = input.dims();
-  auto input_dim_size = input_dim.size();
-  auto output_dim = output->dims();
-
-  auto slice_size = 1;
-  for (auto i = dim + 1; i < input_dim_size; i++) {
-    slice_size *= input_dim[i];
-  }
-
-  auto input_width = slice_size * input_dim[dim];
-  auto output_width = slice_size * output_dim[dim];
-
-  auto outer_nums = 1;
-  for (auto i = 0; i < dim; i++) {
-    outer_nums *= input_dim[i];
-  }
-
-  auto index_size = index.dims()[0];
-
-  std::vector<InT> input_vec;
-  std::vector<IndexT> index_vec;
-  paddle::framework::TensorToVector(
-      input, context.device_context(), &input_vec);
-  paddle::framework::TensorToVector(
-      index, context.device_context(), &index_vec);
-  std::vector<InT> out_vec(output->numel());
-
-  for (int i = 0; i < index_size; i++) {
-    PADDLE_ENFORCE_GE(
-        index_vec[i],
-        0,
-        platform::errors::InvalidArgument(
-            "Variable value (index) of OP(index_select) "
-            "expected >= 0 and < %ld, but got %ld. Please check input "
-            "value.",
-            input_dim[dim],
-            index_vec[i]));
-    PADDLE_ENFORCE_LT(
-        index_vec[i],
-        input_dim[dim],
-        platform::errors::InvalidArgument(
-            "Variable value (index) of OP(index_select) "
-            "expected >= 0 and < %ld, but got %ld. Please check input "
-            "value.",
-            input_dim[dim],
-            index_vec[i]));
-  }
-
-  for (auto i = 0; i < outer_nums; i++) {
-    auto input_start_offset = i * input_width;
-    auto output_start_offset = i * output_width;
-
-    for (auto j = 0; j < index_size; j++) {
-      IndexT index_value = index_vec[j];
-      for (auto k = 0; k < slice_size; k++) {
-        out_vec[output_start_offset + j * slice_size + k] =
-            input_vec[input_start_offset + index_value * slice_size + k];
-      }
-    }
-  }
-  output->mutable_data<InT>(context.GetPlace());
-  framework::TensorFromVector(out_vec, context.device_context(), output);
-  output->Resize(output_dim);
-}
+namespace phi {
 
 // The core logic of computing Unique Consecutive for a flattend Tensor
-template <typename InT, typename IndexT, typename equal_T, typename not_equal_T>
-static void UniqueConsecutiveFlattendCUDATensor(
-    const framework::ExecutionContext& context,
-    const Tensor& in,
-    Tensor* out,
+template <typename Context,
+          typename InT,
+          typename IndexT,
+          typename equal_T,
+          typename not_equal_T>
+static void UniqueConsecutiveFlattenedCUDATensor(const Context& context,
+                                                 const DenseTensor& in,
+                                                 DenseTensor* out,
                                                  bool return_inverse,
                                                  bool return_counts,
                                                  equal_T equal,
                                                  not_equal_T not_equal,
-    int64_t num_input) {
-  // 0. Prepration
-  Tensor in_hat;
-  framework::TensorCopy(in, context.GetPlace(), &in_hat);
-  auto in_data_hat = in_hat.mutable_data<InT>(context.GetPlace());
-
-  Tensor sorted_indices;
+                                                 int64_t num_input,
+                                                 DenseTensor* inverse,
+                                                 DenseTensor* counts) {
+  // 0. Preparation
+  DenseTensor in_hat;
+  phi::Copy(context, in, context.GetPlace(), false, &in_hat);
+  auto in_data_hat = context.template Alloc<InT>(&in_hat);
+
+  DenseTensor sorted_indices;
   sorted_indices.Resize(phi::make_ddim({num_input}));
-  auto sorted_indices_data =
-      sorted_indices.mutable_data<IndexT>(context.GetPlace());
+  auto sorted_indices_data = context.template Alloc<IndexT>(&sorted_indices);
   thrust::sequence(
       thrust::device, sorted_indices_data, sorted_indices_data + num_input);
   // 1. Calculate op result: 'out'
-  Tensor range;
+  DenseTensor range;
   range.Resize(phi::make_ddim({num_input + 1}));
-  auto range_data_ptr = range.mutable_data<IndexT>(context.GetPlace());
+  auto range_data_ptr = context.template Alloc<IndexT>(&range);
   thrust::sequence(
       thrust::device, range_data_ptr, range_data_ptr + num_input + 1);
-  framework::TensorCopy(in_hat, context.GetPlace(), out);
+  phi::Copy(context, in_hat, context.GetPlace(), false, out);
   int num_out;
-  auto out_data = out->mutable_data<InT>(context.GetPlace());
+  auto out_data = context.template Alloc<InT>(out);
   num_out =
       thrust::unique_by_key(
           thrust::device, out_data, out_data + num_input, range_data_ptr, equal)
@@ -184,12 +78,11 @@ static void UniqueConsecutiveFlattendCUDATensor(
 
   // 2. Calculate inverse index: 'inverse'
   if (return_inverse) {
-    Tensor* inverse = context.Output<Tensor>("Index");
     inverse->Resize(phi::make_ddim({num_input}));
-    auto inverse_data = inverse->mutable_data<IndexT>(context.GetPlace());
-    Tensor inv_loc;
+    auto inverse_data = context.template Alloc<IndexT>(inverse);
+    DenseTensor inv_loc;
     inv_loc.Resize(phi::make_ddim({num_input}));
-    auto inv_loc_data_ptr = inv_loc.mutable_data<IndexT>(context.GetPlace());
+    auto inv_loc_data_ptr = context.template Alloc<IndexT>(&inv_loc);
     thrust::adjacent_difference(thrust::device,
                                 in_data_hat,
                                 in_data_hat + num_input,
@@ -209,9 +102,8 @@ static void UniqueConsecutiveFlattendCUDATensor(
   }
   // 3. Calculate 'counts'
   if (return_counts) {
-    Tensor* counts = context.Output<Tensor>("Counts");
     counts->Resize(phi::make_ddim({num_out}));
-    auto count_data = counts->mutable_data<IndexT>(context.GetPlace());
+    auto count_data = context.template Alloc<IndexT>(counts);
     // init 'count_data' as 0
     thrust::fill(thrust::device, count_data, count_data + num_out, 0);
     thrust::device_ptr<IndexT> range_data_ptr_dev(range_data_ptr);
@@ -223,26 +115,72 @@ static void UniqueConsecutiveFlattendCUDATensor(
   }
 }
 
+// functor for processing a flattend Tensor
+template <typename Context, typename InT>
+struct UniqueConsecutiveFlattenedCUDAFunctor {
+  const Context& ctx_;
+  const DenseTensor& in_;
+  DenseTensor* out_;
+  const bool return_inverse_;
+  const bool return_counts_;
+  DenseTensor* inverse_;
+  DenseTensor* count_;
+
+  UniqueConsecutiveFlattenedCUDAFunctor(const Context& context,
+                                        const DenseTensor& in,
+                                        DenseTensor* out,
+                                        bool return_inverse,
+                                        bool return_counts,
+                                        DenseTensor* inverse,
+                                        DenseTensor* count)
+      : ctx_(context),
+        in_(in),
+        out_(out),
+        return_inverse_(return_inverse),
+        return_counts_(return_counts),
+        inverse_(inverse),
+        count_(count) {}
+
+  template <typename IndexT>
+  void apply() const {
+    UniqueConsecutiveFlattenedCUDATensor<Context, InT, IndexT>(
+        ctx_,
+        in_,
+        out_,
+        return_inverse_,
+        return_counts_,
+        thrust::equal_to<InT>(),
+        thrust::not_equal_to<InT>(),
+        in_.numel(),
+        inverse_,
+        count_);
+  }
+};
+
 // The logic of compute unique with axis required, it's a little different
 // from above function
-template <typename InT, typename IndexT, typename equal_T, typename not_equal_T>
-static void ComputeUniqueConsecutiveDims(
-    const framework::ExecutionContext& context,
-    Tensor* sorted_indices,
+template <typename Context,
+          typename InT,
+          typename IndexT,
+          typename equal_T,
+          typename not_equal_T>
+static void ComputeUniqueConsecutiveDims(const Context& context,
+                                         DenseTensor* sorted_indices,
                                          IndexT* sorted_indices_data,
-    Tensor* out,
+                                         DenseTensor* out,
                                          bool return_inverse,
                                          bool return_counts,
                                          equal_T equal,
                                          not_equal_T not_equal,
-    int64_t row) {
+                                         int64_t row,
+                                         DenseTensor* inverse,
+                                         DenseTensor* counts) {
   // 1. inverse indices: 'inverse'
-  Tensor* inverse = context.Output<Tensor>("Index");
   inverse->Resize(phi::make_ddim({row}));
-  auto inverse_data = inverse->mutable_data<IndexT>(context.GetPlace());
-  Tensor inv_loc;
+  auto inverse_data = context.template Alloc<IndexT>(inverse);
+  DenseTensor inv_loc;
   inv_loc.Resize(phi::make_ddim({row}));
-  auto inv_loc_data_ptr = inv_loc.mutable_data<IndexT>(context.GetPlace());
+  auto inv_loc_data_ptr = context.template Alloc<IndexT>(&inv_loc);
   thrust::adjacent_difference(thrust::device,
                               sorted_indices_data,
                               sorted_indices_data + row,
@@ -261,9 +199,9 @@ static void ComputeUniqueConsecutiveDims(
                   inverse_data);
 
   // 2. sorted indices
-  Tensor range;
+  DenseTensor range;
   range.Resize(phi::make_ddim({row + 1}));
-  auto range_data_ptr = range.mutable_data<IndexT>(context.GetPlace());
+  auto range_data_ptr = context.template Alloc<IndexT>(&range);
   thrust::sequence(thrust::device, range_data_ptr, range_data_ptr + row + 1);
   int num_out;
   num_out = thrust::unique_by_key(thrust::device,
@@ -278,23 +216,135 @@ static void ComputeUniqueConsecutiveDims(
   sorted_indices->Resize(phi::make_ddim({num_out}));
 
   // 3. counts: 'counts'
-  Tensor* counts = context.Output<Tensor>("Counts");
   counts->Resize(phi::make_ddim({num_out}));
-  auto count_data = counts->mutable_data<IndexT>(context.GetPlace());
+  auto count_data = context.template Alloc<IndexT>(counts);
   thrust::fill(thrust::device, count_data, count_data + row, 0);
   thrust::adjacent_difference(
       thrust::device, range_data_ptr + 1, range_data_ptr + row + 1, count_data);
 }
 
+// Binary function 'equal_to'
+template <typename InT>
+struct BinaryEqual {
+  int64_t col;
+  const InT* in_trans_data;
+
+  BinaryEqual(int64_t _col, const InT* _in_trans_data)
+      : col(_col), in_trans_data(_in_trans_data) {}
+
+  __device__ bool operator()(int64_t a, int64_t b) const {
+    for (int64_t i = 0; i < col; ++i) {
+      InT lhs = in_trans_data[i + a * col];
+      InT rhs = in_trans_data[i + b * col];
+      if (lhs != rhs) {
+        return false;
+      }
+    }
+    return true;
+  }
+};
+
+// Binary function 'not_equal_to'
+template <typename InT>
+struct BinaryNotEqual {
+  int64_t col;
+  const InT* in_trans_data;
+
+  BinaryNotEqual(int64_t _col, const InT* _in_trans_data)
+      : col(_col), in_trans_data(_in_trans_data) {}
+
+  __device__ bool operator()(int64_t a, int64_t b) const {
+    for (int64_t i = 0; i < col; ++i) {
+      InT lhs = in_trans_data[i + a * col];
+      InT rhs = in_trans_data[i + b * col];
+      if (lhs != rhs) {
+        return true;
+      }
+    }
+    return false;
+  }
+};
+
+// index_select() function for Tensor
+template <typename Context, typename InT, typename IndexT>
+void IndexSelect(const Context& context,
+                 const DenseTensor& input,
+                 const DenseTensor& index,
+                 DenseTensor* output,
+                 int dim) {
+  auto input_dim = input.dims();
+  auto input_dim_size = input_dim.size();
+  auto output_dim = output->dims();
+
+  auto slice_size = 1;
+  for (auto i = dim + 1; i < input_dim_size; i++) {
+    slice_size *= input_dim[i];
+  }
+
+  auto input_width = slice_size * input_dim[dim];
+  auto output_width = slice_size * output_dim[dim];
+
+  auto outer_nums = 1;
+  for (auto i = 0; i < dim; i++) {
+    outer_nums *= input_dim[i];
+  }
+
+  auto index_size = index.dims()[0];
+
+  std::vector<InT> input_vec;
+  std::vector<IndexT> index_vec;
+  paddle::framework::TensorToVector(input, context, &input_vec);
+  paddle::framework::TensorToVector(index, context, &index_vec);
+  std::vector<InT> out_vec(output->numel());
+
+  for (int i = 0; i < index_size; i++) {
+    PADDLE_ENFORCE_GE(
+        index_vec[i],
+        0,
+        phi::errors::InvalidArgument(
+            "Variable value (index) of OP(index_select) "
+            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "value.",
+            input_dim[dim],
+            index_vec[i]));
+    PADDLE_ENFORCE_LT(
+        index_vec[i],
+        input_dim[dim],
+        phi::errors::InvalidArgument(
+            "Variable value (index) of OP(index_select) "
+            "expected >= 0 and < %ld, but got %ld. Please check input "
+            "value.",
+            input_dim[dim],
+            index_vec[i]));
+  }
+
+  for (auto i = 0; i < outer_nums; i++) {
+    auto input_start_offset = i * input_width;
+    auto output_start_offset = i * output_width;
+
+    for (auto j = 0; j < index_size; j++) {
+      IndexT index_value = index_vec[j];
+      for (auto k = 0; k < slice_size; k++) {
+        out_vec[output_start_offset + j * slice_size + k] =
+            input_vec[input_start_offset + index_value * slice_size + k];
+      }
+    }
+  }
+  context.template Alloc<InT>(output);
+  paddle::framework::TensorFromVector(out_vec, context, output);
+  output->Resize(output_dim);
+}
+
 // Calculate unique consecutive when 'axis' is set
-template <typename DeviceContext, typename InT, typename IndexT>
-static void UniqueConsecutiveDimsCUDATensor(
-    const framework::ExecutionContext& context,
-    const Tensor& in,
-    Tensor* out,
+template <typename Context, typename InT, typename IndexT>
+static void UniqueConsecutiveDimsCUDATensor(const Context& context,
+                                            const DenseTensor& in,
+                                            DenseTensor* out,
                                             bool return_inverse,
                                             bool return_counts,
-    int axis) {
+                                            int axis,
+                                            DenseTensor* inverse,
+                                            DenseTensor* counts) {
   // 1. Transpose & reshape
   // Transpose tensor: eg. axis=1, [dim0, dim1, dim2] -> [dim1, dim0, dim2]
   std::vector<int> permute(in.dims().size());
@@ -304,19 +354,18 @@ static void UniqueConsecutiveDimsCUDATensor(
   std::vector<int64_t> in_trans_dims_vec(phi::vectorize(in.dims()));
   in_trans_dims_vec[axis] = in.dims()[0];
   in_trans_dims_vec[0] = in.dims()[axis];
-  framework::Tensor in_trans;
-  framework::DDim in_trans_dims = phi::make_ddim(in_trans_dims_vec);
+  DenseTensor in_trans;
+  DDim in_trans_dims = phi::make_ddim(in_trans_dims_vec);
   in_trans.Resize(in_trans_dims);
-  in_trans.mutable_data<InT>(context.GetPlace());
-  auto& dev_ctx = context.cuda_device_context();
-  TransCompute<DeviceContext, InT>(in.dims().size(),  // num of dims
-                                   dev_ctx,           // device
+  context.template Alloc<InT>(&in_trans);
+  phi::funcs::TransCompute<Context, InT>(in.dims().size(),  // num of dims
+                                         context,           // device
                                          in,                // original Tensor
                                          &in_trans,  // Tensor after reshape
                                          permute);   // index of axis
 
   // Reshape tensor: eg. [dim1, dim0, dim2] -> [dim1, dim0*dim2]
-  framework::DDim in_trans_flat_dims = phi::flatten_to_2d(in_trans_dims, 1);
+  DDim in_trans_flat_dims = phi::flatten_to_2d(in_trans_dims, 1);
   in_trans.Resize(in_trans_flat_dims);
 
   // now 'in_trans' is 2D
@@ -324,16 +373,15 @@ static void UniqueConsecutiveDimsCUDATensor(
   int64_t row = in_trans.dims()[0];
   const InT* in_trans_data = in_trans.data<InT>();
 
-  Tensor sorted_indices;
+  DenseTensor sorted_indices;
   sorted_indices.Resize(phi::make_ddim({row}));
-  auto sorted_indices_data =
-      sorted_indices.mutable_data<IndexT>(context.GetPlace());
+  auto sorted_indices_data = context.template Alloc<IndexT>(&sorted_indices);
 
   // 2. Calculate 'inverse', 'counts'
   // Init index
   thrust::sequence(
       thrust::device, sorted_indices_data, sorted_indices_data + row);
-  ComputeUniqueConsecutiveDims<InT, IndexT>(
+  ComputeUniqueConsecutiveDims<Context, InT, IndexT>(
       context,
       &sorted_indices,
       sorted_indices_data,
@@ -342,143 +390,70 @@ static void UniqueConsecutiveDimsCUDATensor(
       return_counts,
       BinaryEqual<InT>(col, in_trans_data),
       BinaryNotEqual<InT>(col, in_trans_data),
-      row);
+      row,
+      inverse,
+      counts);
 
   // 3. Select indices and reshape back to get 'out'
-  Tensor out_trans;
+  DenseTensor out_trans;
   std::vector<int64_t> out_trans_dims_vec = in_trans_dims_vec;
   out_trans_dims_vec[0] = sorted_indices.numel();
   out_trans.Resize(phi::make_ddim(out_trans_dims_vec));
-  out_trans.mutable_data<InT>(context.GetPlace());
+  context.template Alloc<InT>(&out_trans);
 
-  IndexSelect<InT, IndexT>(context, in_trans, sorted_indices, &out_trans, 0);
+  IndexSelect<Context, InT, IndexT>(
+      context, in_trans, sorted_indices, &out_trans, 0);
 
   std::swap(out_trans_dims_vec[0], out_trans_dims_vec[axis]);
   out->Resize(phi::make_ddim(out_trans_dims_vec));
-  out->mutable_data<InT>(context.GetPlace());
-  std::vector<framework::Tensor> out_trans_unbind = Unbind(out_trans);
-  math::ConcatFunctor<DeviceContext, InT> concat_functor;
-  concat_functor(dev_ctx, out_trans_unbind, 0, &out_trans);
-  TransCompute<DeviceContext, InT>(
-      out_trans.dims().size(), dev_ctx, out_trans, out, permute);
+  context.template Alloc<InT>(out);
+  std::vector<DenseTensor> out_trans_unbind = phi::funcs::Unbind(out_trans);
+  phi::funcs::ConcatFunctor<Context, InT> concat_functor;
+  concat_functor(context, out_trans_unbind, 0, &out_trans);
+  phi::funcs::TransCompute<Context, InT>(
+      out_trans.dims().size(), context, out_trans, out, permute);
 }
 
-// functor for processing a flattend Tensor
-template <typename DeviceContext, typename InT>
-struct UniqueConsecutiveFlattendCUDAFunctor {
-  const framework::ExecutionContext& ctx_;
-  const Tensor& in_;
-  Tensor* out_;
-  const bool return_inverse_;
-  const bool return_counts_;
-
-  UniqueConsecutiveFlattendCUDAFunctor(
-      const framework::ExecutionContext& context,
-      const Tensor& in,
-      Tensor* out,
-      bool return_inverse,
-      bool return_counts)
-      : ctx_(context),
-        in_(in),
-        out_(out),
-        return_inverse_(return_inverse),
-        return_counts_(return_counts) {}
-
-  template <typename IndexT>
-  void apply() const {
-    UniqueConsecutiveFlattendCUDATensor<InT, IndexT>(
-        ctx_,
-        in_,
-        out_,
-        return_inverse_,
-        return_counts_,
-        thrust::equal_to<InT>(),
-        thrust::not_equal_to<InT>(),
-        in_.numel());
-  }
-};
-
 // functor for processing a multi-dimentional Tensor
-template <typename DeviceContext, typename InT>
+template <typename Context, typename InT>
 struct UniqueConsecutiveDimsCUDAFunctor {
-  const framework::ExecutionContext& ctx_;
-  const Tensor& in_;
-  Tensor* out_;
+  const Context& ctx_;
+  const DenseTensor& in_;
+  DenseTensor* out_;
   const int axis_;
   const bool return_inverse_;
   const bool return_counts_;
+  DenseTensor* inverse_;
+  DenseTensor* count_;
 
-  UniqueConsecutiveDimsCUDAFunctor(const framework::ExecutionContext& context,
-                                   const Tensor& in,
-                                   Tensor* out,
+  UniqueConsecutiveDimsCUDAFunctor(const Context& context,
+                                   const DenseTensor& in,
+                                   DenseTensor* out,
                                    const int axis,
                                    bool return_inverse,
-                                   bool return_counts)
+                                   bool return_counts,
+                                   DenseTensor* inverse,
+                                   DenseTensor* count)
       : ctx_(context),
         in_(in),
         out_(out),
         axis_(axis),
         return_inverse_(return_inverse),
-        return_counts_(return_counts) {}
+        return_counts_(return_counts),
+        inverse_(inverse),
+        count_(count) {}
 
   template <typename IndexT>
   void apply() const {
-    UniqueConsecutiveDimsCUDATensor<DeviceContext, InT, IndexT>(
-        ctx_, in_, out_, return_inverse_, return_counts_, axis_);
-  }
-};
-
-// Unique_Consecutive_op CUDA implementation.
-template <typename InT>
-class UniqueConsecutiveKernel<platform::CUDADeviceContext, InT>
-    : public framework::OpKernel<InT> {
- public:
-  void Compute(const framework::ExecutionContext& context) const override {
-    auto* x = context.Input<framework::Tensor>("X");
-    auto* out = context.Output<framework::Tensor>("Out");
-    auto data_type = static_cast<framework::proto::VarType::Type>(
-        context.Attr<int>("dtype"));
-    if (data_type == framework::proto::VarType::INT32) {
-      PADDLE_ENFORCE_LE(
-          x->numel() + 1,
-          INT_MAX,
-          platform::errors::InvalidArgument(
-              "The number of elements in Input(X) should be less than or "
-              "equal to INT_MAX, but received num is %d. Please set `dtype` to "
-              "int64.",
-              x->numel()));
-    }
-
-    std::vector<int> axis_vec = context.Attr<std::vector<int>>("axis");
-    bool return_inverse = context.Attr<bool>("return_inverse");
-    bool return_counts = context.Attr<bool>("return_counts");
-
-    // if 'axis' is not required, flatten the Tensor.
-    if (axis_vec.empty()) {
-      framework::VisitDataTypeTiny(
-          data_type,
-          UniqueConsecutiveFlattendCUDAFunctor<platform::CUDADeviceContext,
-                                               InT>(
-              context, *x, out, return_inverse, return_counts));
-    } else {
-      // 'axis' is required.
-      int axis = axis_vec[0];
-      framework::VisitDataTypeTiny(
-          data_type,
-          UniqueConsecutiveDimsCUDAFunctor<platform::CUDADeviceContext, InT>(
-              context, *x, out, axis, return_inverse, return_counts));
-    }
+    UniqueConsecutiveDimsCUDATensor<Context, InT, IndexT>(ctx_,
+                                                          in_,
+                                                          out_,
+                                                          return_inverse_,
+                                                          return_counts_,
+                                                          axis_,
+                                                          inverse_,
+                                                          count_);
   }
 };
 
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-
-REGISTER_OP_CUDA_KERNEL(
-    unique_consecutive,
-    ops::UniqueConsecutiveKernel<paddle::platform::CUDADeviceContext, float>,
-    ops::UniqueConsecutiveKernel<paddle::platform::CUDADeviceContext, double>,
-    ops::UniqueConsecutiveKernel<paddle::platform::CUDADeviceContext, int32_t>,
-    ops::UniqueConsecutiveKernel<paddle::platform::CUDADeviceContext, int64_t>);
+}  // namespace phi

paddle/phi/kernels/gpu/unique_consecutive_kernel.cu

+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include "paddle/phi/kernels/unique_consecutive_kernel.h"
+#include "paddle/phi/kernels/gpu/unique_consecutive_functor.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/errors.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+#include "paddle/fluid/framework/data_type.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void UniqueConsecutiveKernel(const Context& dev_ctx,
+                             const DenseTensor& x,
+                             bool return_inverse,
+                             bool return_counts,
+                             const std::vector<int>& axis,
+                             int dtype,
+                             DenseTensor* out,
+                             DenseTensor* index,
+                             DenseTensor* counts) {
+  auto data_type = static_cast<paddle::framework::proto::VarType::Type>(dtype);
+  if (data_type == paddle::framework::proto::VarType::INT32) {
+    PADDLE_ENFORCE_LE(
+        x.numel() + 1,
+        INT_MAX,
+        phi::errors::InvalidArgument(
+            "The number of elements in Input(X) should be less than or "
+            "equal to INT_MAX, but received num is %d. Please set `dtype` to "
+            "int64.",
+            x.numel()));
+  }
+
+  // if 'axis' is not required, flatten the Tensor.
+  if (axis.empty()) {
+    paddle::framework::VisitDataTypeTiny(
+        data_type,
+        UniqueConsecutiveFlattenedCUDAFunctor<Context, T>(
+            dev_ctx, x, out, return_inverse, return_counts, index, counts));
+  } else {
+    // 'axis' is required.
+    int valid_axis = axis[0];
+    paddle::framework::VisitDataTypeTiny(
+        data_type,
+        UniqueConsecutiveDimsCUDAFunctor<Context, T>(dev_ctx,
+                                                     x,
+                                                     out,
+                                                     valid_axis,
+                                                     return_inverse,
+                                                     return_counts,
+                                                     index,
+                                                     counts));
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(unique_consecutive,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::UniqueConsecutiveKernel,
+                   float,
+                   double,
+                   int32_t,
+                   int64_t) {}

paddle/phi/kernels/unique_consecutive_kernel.h

+// Copyright (c) 2022 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.
+
+#pragma once
+
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void UniqueConsecutiveKernel(const Context& dev_ctx,
+                             const DenseTensor& x,
+                             bool return_inverse,
+                             bool return_counts,
+                             const std::vector<int>& axis,
+                             int dtype,
+                             DenseTensor* out,
+                             DenseTensor* index,
+                             DenseTensor* counts);
+
+}  // namespace phi

paddle/phi/ops/compat/unique_consecutive_sig.cc

+// Copyright (c) 2022 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/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature UniqueConsecutiveOpArgumentMapping(
+    const ArgumentMappingContext& ctx) {
+  return KernelSignature("unique_consecutive",
+                         {"X"},
+                         {"return_inverse", "return_counts", "axis", "dtype"},
+                         {"Out", "Index", "Counts"});
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(unique_consecutive,
+                           phi::UniqueConsecutiveOpArgumentMapping);