Add gpu kernel for argsort op

paddle/fluid/operators/argsort_op.cu

+/* Copyright (c) 2016 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 <thrust/execution_policy.h>
+#include <thrust/sort.h>
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/argsort_op.h"
+#include "paddle/fluid/platform/assert.h"
+#include "paddle/fluid/platform/cuda_device_function.h"
+#include "paddle/fluid/platform/cuda_primitives.h"
+
+namespace paddle {
+namespace operators {
+
+using Tensor = framework::Tensor;
+using platform::PADDLE_CUDA_NUM_THREADS;
+
+template <typename T>
+__global__ void PermuteInData(const T* in, const int64_t* trg_idx, int64_t n,
+                              T* med_out) {
+  int index = threadIdx.x + blockDim.x * blockIdx.x;
+  if (index < n) {
+    med_out[trg_idx[index]] = in[index];
+  }
+}
+
+template <typename T>
+__global__ void Sort(int64_t axis_dim, int64_t groups, T* med_out,
+                     int64_t* med_ids) {
+  int index = threadIdx.x + blockDim.x * blockIdx.x;
+  if (index < groups) {
+    thrust::sort_by_key(thrust::device, med_out + index * axis_dim,
+                        med_out + axis_dim * (1 + index),
+                        med_ids + index * axis_dim);
+  }
+}
+
+template <typename T>
+__global__ void PermuteMediateData(const T* med_out, const int64_t* med_ids,
+                                   const int64_t* trg_idx, int64_t n, T* out,
+                                   int64_t* indices) {
+  int index = threadIdx.x + blockDim.x * blockIdx.x;
+  if (index < n) {
+    out[index] = med_out[trg_idx[index]];
+    indices[index] = med_ids[trg_idx[index]];
+  }
+}
+
+template <typename T>
+class ArgsortOpCUDAKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    auto* input = ctx.Input<Tensor>("X");
+    auto* output = ctx.Output<Tensor>("Out");
+    auto* indices = ctx.Output<Tensor>("Indices");
+    int axis = ctx.Attr<int>("axis");
+
+    auto in_dims = input->dims();
+    axis = (axis == -1) ? (in_dims.size() - 1) : axis;
+
+    const T* in_data = input->data<T>();
+    T* out_data = output->mutable_data<T>(ctx.GetPlace());
+    int64_t* ids_data = indices->mutable_data<int64_t>(ctx.GetPlace());
+
+    int64_t numel = input->numel();
+    int64_t groups = numel / in_dims[axis];
+
+    // Mediate tensor for sorting
+    Tensor mediate_output;
+    T* med_out_data =
+        mediate_output.mutable_data<T>(input->dims(), ctx.GetPlace());
+
+    // The target index of each elemement in mediate tensor
+    std::vector<int64_t> target_idx(numel, 0);
+    // To record the index along the given axis for the data in mediate tensor
+    std::vector<int64_t> mediate_indices(numel, 0);
+    std::vector<int64_t> in_dims_out_axis = vectorize(in_dims);
+    in_dims_out_axis.erase(in_dims_out_axis.begin() + axis);
+    for (int64_t index = 0; index < numel; ++index) {
+      int64_t tmp = index;
+      int64_t pos_in_axis = 0;
+      std::vector<int64_t> shape;
+      for (int64_t j = in_dims.size() - 1; j >= 0; --j) {
+        if (j != axis) {
+          shape.push_back(tmp % in_dims[j]);
+        } else {
+          pos_in_axis = tmp % in_dims[j];
+        }
+        tmp /= in_dims[j];
+      }
+      std::reverse(shape.begin(), shape.end());
+      int64_t group = (shape.size() > 0) ? shape[0] : 0;
+      for (size_t j = 0; j < shape.size() - 1; ++j) {
+        group = group * in_dims_out_axis[j + 1] + shape[j + 1];
+      }
+
+      target_idx[index] = group * in_dims[axis] + pos_in_axis;
+      mediate_indices[target_idx[index]] = pos_in_axis;
+    }
+
+    thrust::device_vector<int64_t> med_ids_dev(mediate_indices.begin(),
+                                               mediate_indices.end());
+    int64_t* med_ids_data = thrust::raw_pointer_cast(med_ids_dev.data());
+    thrust::device_vector<int64_t> trg_idx_dev(target_idx.begin(),
+                                               target_idx.end());
+    int64_t* trg_idx = thrust::raw_pointer_cast(trg_idx_dev.data());
+
+    auto stream = reinterpret_cast<const platform::CUDADeviceContext&>(
+                      ctx.device_context())
+                      .stream();
+    auto num_threads = PADDLE_CUDA_NUM_THREADS;
+
+    PermuteInData<<<(numel - 1) / num_threads + 1, num_threads, 0, stream>>>(
+        in_data, trg_idx, numel, med_out_data);
+
+    Sort<<<(groups - 1) / num_threads + 1, num_threads, 0, stream>>>(
+        in_dims[axis], groups, med_out_data, med_ids_data);
+
+    PermuteMediateData<<<(numel - 1) / num_threads + 1, num_threads, 0,
+                         stream>>>(med_out_data, med_ids_data, trg_idx, numel,
+                                   out_data, ids_data);
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+REGISTER_OP_CUDA_KERNEL(argsort, paddle::operators::ArgsortOpCUDAKernel<float>,
+                        paddle::operators::ArgsortOpCUDAKernel<double>);