decagon-pytorch/src/torch_stablesort/torch_stablesort_cuda.cu

#pragma once

#include <thrust/sort.h>
#include <thrust/device_ptr.h>
#include <thrust/execution_policy.h>

#include "dispatch.h"
#include "torch_stablesort_cuda.h"

template<bool descending, typename T>
struct stable_sort_impl_cuda {
  std::vector<torch::Tensor> operator()(
    torch::Tensor input,
    int dim,
    torch::optional<std::tuple<torch::Tensor, torch::Tensor>> out
  ) const {

    if (input.is_sparse())
      throw std::runtime_error("Sparse tensors are not supported");

    if (input.device().type() != torch::DeviceType::CUDA)
      throw std::runtime_error("Only CUDA tensors are supported");

    if (out != torch::nullopt)
      throw std::runtime_error("out argument is not supported");

    auto x = input.clone();

    if (dim != -1)
      x = torch::transpose(x, dim, -1);

    auto x_sizes = x.sizes();

    x = x.view({ -1, x.size(-1) }).contiguous();

    auto x_outer_stride = x.stride(-2);
    auto x_inner_stride = x.stride(-1);
    auto n_cols = x.size(1);
    auto n_rows = x.size(0);
    auto px = x.data_ptr<T>();

    assert(x_inner_stride == 1);

    auto y = torch::repeat_interleave(
      torch::arange(0, n_cols, 1, torch::TensorOptions()
        .dtype(torch::kInt64)
          .device(x.device())).view({ 1, -1 }),
      n_rows,
      0 /* dim */
    );

    auto y_outer_stride = y.stride(-2);
    auto y_inner_stride = y.stride(-1);
    auto py = y.data_ptr<int64_t>();

    assert(y_inner_stride == 1);

    #define NUM_STREAMS 16
    cudaStream_t streams[NUM_STREAMS];
    for(int i = 0; i < NUM_STREAMS; i++)
        assert(cudaStreamCreate(&streams[i]) == cudaSuccess);

    thrust::host_vector<int64_t> row_indices(n_rows);
    thrust::sequence(row_indices.begin(), row_indices.end());
    thrust::for_each(thrust::host, row_indices.begin(), row_indices.end(),
      [&streams, py, y_outer_stride, px, x_outer_stride, x_inner_stride, n_cols](int64_t i) {
        auto ind_beg = thrust::device_pointer_cast(py + i * y_outer_stride);

        auto val_beg = thrust::device_pointer_cast(px + i * x_outer_stride);
        auto val_end = thrust::device_pointer_cast(px + i * x_outer_stride +
            n_cols * x_inner_stride);

        if (descending)
          thrust::stable_sort_by_key(thrust::cuda::par.on(streams[i % NUM_STREAMS]), val_beg, val_end, ind_beg,
            thrust::greater<T>());
        else
          thrust::stable_sort_by_key(thrust::cuda::par.on(streams[i % NUM_STREAMS]), val_beg, val_end, ind_beg);
      });
    cudaDeviceSynchronize();

    for(int i = 0; i < NUM_STREAMS; i++)
        assert(cudaStreamDestroy(streams[i]) == cudaSuccess);

    x = x.view(x_sizes);
    y = y.view(x_sizes);

    x = (dim == -1) ?
      x :
      torch::transpose(x, dim, -1).contiguous();

    y = (dim == -1) ?
      y :
      torch::transpose(y, dim, -1).contiguous();

    return { x, y };
  }
};

template <typename T>
struct stable_sort_impl_desc_cuda: stable_sort_impl_cuda<true, T> {};

template <typename T>
struct stable_sort_impl_asc_cuda: stable_sort_impl_cuda<false, T> {};

std::vector<torch::Tensor> dispatch_cuda(torch::Tensor input,
  int dim,
  bool descending,
  torch::optional<std::tuple<torch::Tensor, torch::Tensor>> out) {

  if (descending)
    return dispatch<stable_sort_impl_desc_cuda, std::vector<torch::Tensor>>(
        input, dim, out);
  else
    return dispatch<stable_sort_impl_asc_cuda, std::vector<torch::Tensor>>(
        input, dim, out);
}