// 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/argsort_kernel.h"

#include <thrust/copy.h>
#include <thrust/execution_policy.h>
#include <thrust/sequence.h>
#include <thrust/sort.h>
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/gpu/gpu_info.h"
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/blas/blas.h"
#include "paddle/phi/kernels/funcs/cub.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/primitive/functor_primitives.h"
#include "paddle/phi/kernels/transpose_kernel.h"

#ifdef __HIPCC__
#include <rocprim/config.hpp>
#if defined(ROCPRIM_VERSION) && ROCPRIM_VERSION >= 400000
// rocPRIM 4.x (ROCm 7.0+) replaces detail::radix_key_codec_base
// with traits::define for non-builtin / wrapper types.
namespace rocprim {
namespace traits {
template <>
struct define<phi::float16> {
  using float_bit_mask =
      float_bit_mask::values<uint16_t, 0x8000, 0x7C00, 0x03FF>;
};
template <>
struct define<phi::bfloat16> {
  using float_bit_mask =
      float_bit_mask::values<uint16_t, 0x8000, 0x7F80, 0x007F>;
};
}  // namespace traits
}  // namespace rocprim
#else
namespace rocprim {
namespace detail {
template <>
struct radix_key_codec_base<phi::float16>
    : radix_key_codec_integral<phi::float16, uint16_t> {};

template <>
struct radix_key_codec_base<phi::bfloat16>
    : radix_key_codec_integral<phi::bfloat16, uint16_t> {};

#if HIP_VERSION >= 50400000
template <>
struct float_bit_mask<phi::float16> : float_bit_mask<rocprim::half> {};

template <>
struct float_bit_mask<phi::bfloat16> : float_bit_mask<rocprim::bfloat16> {};
#endif
}  // namespace detail
}  // namespace rocprim
#endif  // ROCPRIM_VERSION
#else
// set cub base traits in order to handle float16
namespace cub {
template <>
struct NumericTraits<phi::float16>
    : BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::float16> {};

template <>
struct NumericTraits<phi::bfloat16>
    : BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::bfloat16> {};
}  // namespace cub

#endif

namespace phi {

// Iter for move to next row
struct SegmentOffsetIter {
  EIGEN_DEVICE_FUNC
  explicit SegmentOffsetIter(int num_cols) : num_cols_(num_cols) {}

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int operator()(int idx) const {
    return idx * num_cols_;
  }

  int num_cols_;
};

template <typename T, typename IndType>
__global__ void merge_kernel(const T* A,
                             size_t sizeA,
                             const T* B,
                             size_t sizeB,
                             const IndType* ids_A,
                             const IndType* ids_B,
                             T* out,
                             IndType* out_ids,
                             bool descending) {
  int64_t thread =
      static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
  int64_t num_per_thread = (sizeA + sizeB + thread) / thread;
  for (int64_t offset = 0; offset < num_per_thread; offset++) {
    size_t idx =
        static_cast<size_t>(blockIdx.x) * static_cast<size_t>(blockDim.x) +
        static_cast<size_t>(threadIdx.x) + offset * thread;
    size_t total = sizeA + sizeB;
    if (idx >= total) return;
    size_t left = (idx > sizeB) ? idx - sizeB : 0;
    size_t right = (idx < sizeA) ? idx : sizeA;
    while (left < right) {
      size_t mid = (left + right) / 2;
      size_t b_idx = idx - mid;

      T A_mid, B_bidx;
      if (descending) {
        A_mid = (mid >= sizeA) ? std::numeric_limits<T>::lowest() : A[mid];
        B_bidx = (b_idx >= sizeB) ? std::numeric_limits<T>::lowest() : B[b_idx];
      } else {
        A_mid = (mid >= sizeA) ? std::numeric_limits<T>::max() : A[mid];
        B_bidx = (b_idx >= sizeB) ? std::numeric_limits<T>::max() : B[b_idx];
      }

      if (descending ? (A_mid >= B_bidx) : (A_mid <= B_bidx))
        left = mid + 1;
      else
        right = mid;
    }

    size_t a_idx = left;
    size_t b_idx = idx - a_idx;
    if (a_idx >= sizeA) {
      if (descending ? (A[sizeA - 1] < B[b_idx]) : (A[sizeA - 1] > B[b_idx])) {
        out[idx] = A[sizeA - 1];
        out_ids[idx] = ids_A[sizeA - 1];
      } else {
        out[idx] = B[b_idx];
        out_ids[idx] = ids_B[b_idx];
      }
    } else if (b_idx >= sizeB) {
      out[idx] = A[a_idx];
      out_ids[idx] = ids_A[a_idx];
    } else {
      if (descending ? (A[a_idx] >= B[b_idx]) : (A[a_idx] <= B[b_idx])) {
        out[idx] = A[a_idx];
        out_ids[idx] = ids_A[a_idx];
      } else if (descending ? (a_idx > 0 && (A[a_idx - 1] < B[b_idx]))
                            : (a_idx > 0 && (A[a_idx - 1] > B[b_idx]))) {
        out[idx] = A[a_idx - 1];
        out_ids[idx] = ids_A[a_idx - 1];
      } else {
        out[idx] = B[b_idx];
        out_ids[idx] = ids_B[b_idx];
      }
    }
  }
}

template <typename T>
static __global__ void FillIndex(T* indices, T num_rows, T num_cols) {
  int col_id = threadIdx.x;
  int row_id = blockIdx.x;

  for (T j = row_id; j < num_rows; j += gridDim.x) {
    for (T i = col_id; i < num_cols; i += blockDim.x) {
      indices[j * num_cols + i] = i;
    }
  }
}

#define CUB_ARGSORT_WRAPPER(func, ...)                                        \
  {                                                                           \
    size_t temp_storage_bytes = 0;                                            \
    PADDLE_ENFORCE_GPU_SUCCESS(                                               \
        func(nullptr, temp_storage_bytes, __VA_ARGS__));                      \
    DenseTensor temp_storage;                                                 \
    int64_t temp_size = static_cast<int64_t>(temp_storage_bytes);             \
    PADDLE_ENFORCE_GT(                                                        \
        temp_size,                                                            \
        0,                                                                    \
        common::errors::InvalidArgument(                                      \
            "Argsort temp storage size is %d, but should be greater than 0.", \
            temp_size));                                                      \
    temp_storage.Resize({temp_size});                                         \
    dev_ctx.template Alloc<uint8_t>(&temp_storage);                           \
    PADDLE_ENFORCE_GPU_SUCCESS(                                               \
        func(temp_storage.data<uint8_t>(), temp_storage_bytes, __VA_ARGS__)); \
  }

#define PREDICATE_CUB_ARGSORT(predicate, if_func, else_func, ...) \
  if (predicate)                                                  \
    CUB_ARGSORT_WRAPPER(if_func, __VA_ARGS__)                     \
  else                                                            \
    CUB_ARGSORT_WRAPPER(else_func, __VA_ARGS__)

// Sort by flag descending, True: descending. False: Ascending.
// Default is false.
template <typename T, typename IndType>
void ArgFullSort(const GPUContext& dev_ctx,
                 const DenseTensor* input,
                 DenseTensor* output,
                 DenseTensor* indices,
                 const int64_t num_rows,
                 const int64_t num_cols,
                 const bool descending) {
  PADDLE_ENFORCE_LE_INT_MAX(num_cols, "num_cols");

  auto cu_stream = dev_ctx.stream();
  auto ComputeBlockSize = [](IndType col) {
    if (col > 512)
      return 1024;
    else if (col > 256 && col <= 512)
      return 512;
    else if (col > 128 && col <= 256)
      return 256;
    else
      return 128;
  };
  const int block_size = ComputeBlockSize(num_cols);
  const int64_t maxGridDimX = dev_ctx.GetCUDAMaxGridDimSize()[0];

  const T* inp = input->data<T>();
  IndType* sorted_indices_ptr = indices->data<IndType>();

  // create iter for counting input
  cub::CountingInputIterator<IndType> counting_iter(0);
  // segment_offset is used for move to next row
  cub::TransformInputIterator<IndType,
                              SegmentOffsetIter,
                              cub::CountingInputIterator<IndType>>
      segment_offsets_t(counting_iter, SegmentOffsetIter(num_cols));

  // num_rows is the total segments to be sorted
  constexpr int64_t max_elements = 1 << 30;
  const int64_t total_elements = num_cols * num_rows;
  const int64_t segment_size = num_cols;
  const int64_t element_per_call = std::min(max_elements, total_elements);

  // make sure element_per_call >= segment_size
  const int64_t adjusted_elements_per_call =
      std::max(max_elements, segment_size);

  // make sure batch size is the multiple of segment_size
  const int64_t batch_size =
      (adjusted_elements_per_call / segment_size) * segment_size;
  int64_t offset = 0;
  DenseTensor input_indices;

  T* sorted_out_ptr = sorted_out_ptr = output->data<T>();
  IndType* ind_ptr = nullptr;

  while (offset < total_elements) {
    const int64_t n_elements = std::min(batch_size, total_elements - offset);
    const int64_t n_segments = n_elements / segment_size;

    // allocate a temporary storage for input indices, with shape:
    // [num_segments = n_elements / segment_size, segment_size]
    // will be de-allocated once the sort is done, to save memory and
    // avoid repeated allocation and deallocation
    if (input_indices.initialized()) {
      ind_ptr = input_indices.data<IndType>();
    } else {
      input_indices.Resize({n_segments, segment_size});
      ind_ptr = dev_ctx.template Alloc<IndType>(&input_indices);
    }
    const int64_t grid_size = std::min(n_segments, maxGridDimX);
    // Init a index array
    FillIndex<<<grid_size, block_size, 0, cu_stream>>>(
        ind_ptr, n_segments, segment_size);

    PREDICATE_CUB_ARGSORT(descending,
                          cub::DeviceSegmentedRadixSort::SortPairsDescending,
                          cub::DeviceSegmentedRadixSort::SortPairs,
                          inp + offset,
                          sorted_out_ptr + offset,
                          ind_ptr,
                          sorted_indices_ptr + offset,
                          n_elements,
                          n_segments,
                          segment_offsets_t,
                          segment_offsets_t + 1,
                          0,
                          sizeof(T) * 8,
                          cu_stream);
    offset += n_elements;
  }
}
template <typename T, typename IndType>
void PerSort(const GPUContext& dev_ctx,
             T* out_data,
             int64_t* ids_data,
             IndType start,
             IndType end,
             bool stable,
             bool descending) {
#ifdef PADDLE_WITH_CUDA
  phi::memory_utils::ThrustAllocator<cudaStream_t> allocator(dev_ctx.GetPlace(),
                                                             dev_ctx.stream());
  const auto& exec_policy = thrust::cuda::par(allocator).on(dev_ctx.stream());
#else
  phi::memory_utils::ThrustAllocator<hipStream_t> allocator(dev_ctx.GetPlace(),
                                                            dev_ctx.stream());
  const auto& exec_policy = thrust::hip::par(allocator).on(dev_ctx.stream());
#endif
  if (stable) {
    if (descending) {
      thrust::stable_sort_by_key(exec_policy,
                                 out_data + start,
                                 out_data + end,
                                 ids_data + start,
                                 thrust::greater<T>());
    } else {
      thrust::stable_sort_by_key(
          exec_policy, out_data + start, out_data + end, ids_data + start);
    }
    return;
  } else {
    thrust::sort_by_key(
        exec_policy, out_data + start, out_data + end, ids_data + start);
    if (descending) {
      thrust::reverse(exec_policy, out_data + start, out_data + end);
      thrust::reverse(exec_policy, ids_data + start, ids_data + end);
    }
    return;
  }
}

template <typename T, typename Context>
void ArgsortKernel(const Context& dev_ctx,
                   const DenseTensor& input,
                   int axis,
                   bool descending,
                   bool stable,
                   DenseTensor* output,
                   DenseTensor* indices) {
  auto in_dims = input.dims();
  auto rank = in_dims.size();

  if (input.numel() == 0) {
    output->Resize(in_dims);
    indices->Resize(in_dims);
    dev_ctx.template Alloc<T>(output);
    dev_ctx.template Alloc<int64_t>(indices);
    return;
  }

  axis = (axis < 0) ? (in_dims.size() + axis) : axis;
  const T* in_data = input.data<T>();
  auto size = input.numel();

  if (rank == 0) {
    dev_ctx.template Alloc<T>(output);
    dev_ctx.template Alloc<int64_t>(indices);
    Copy<Context>(dev_ctx, input, dev_ctx.GetPlace(), false, output);
    funcs::set_constant(dev_ctx, indices, static_cast<int64_t>(0));
    return;
  }

  // Use thrust for parallel acceleration when the input size is equal to the
  // length of the 'axis' dimension.
  // Compared to the following 'Special case for full sort', ascending sort is
  // 34 times faster and descending sort is 31 times faster.
  if (size == in_dims[axis]) {
    T* out_data = dev_ctx.template Alloc<T>(output);
    int64_t* ids_data = dev_ctx.template Alloc<int64_t>(indices);
#ifdef PADDLE_WITH_CUDA
    phi::memory_utils::ThrustAllocator<cudaStream_t> allocator(
        dev_ctx.GetPlace(), dev_ctx.stream());
    const auto& exec_policy = thrust::cuda::par(allocator).on(dev_ctx.stream());
#else
    phi::memory_utils::ThrustAllocator<hipStream_t> allocator(
        dev_ctx.GetPlace(), dev_ctx.stream());
    const auto& exec_policy = thrust::hip::par(allocator).on(dev_ctx.stream());
#endif
    auto cu_stream = dev_ctx.stream();
    thrust::sequence(exec_policy, ids_data, ids_data + size);
    thrust::copy(exec_policy, in_data, in_data + size, out_data);
    const int64_t per_number = (1LL << 31) - 1;
    int64_t start = 0;
    int64_t end = std::min(start + per_number, size);
    if (end == size) {
      PerSort<T, int64_t>(
          dev_ctx, out_data, ids_data, start, end, stable, descending);
    } else {
      // Sorting the segments and then merging them
      DenseTensor temp;
      DenseTensor ids;
      temp.Resize(in_dims);
      ids.Resize(in_dims);
      T* temp_data = dev_ctx.template Alloc<T>(&temp);
      int64_t* temp_ids = dev_ctx.template Alloc<int64_t>(&ids);

      while (start != size) {
        PerSort<T, int64_t>(
            dev_ctx, out_data, ids_data, start, end, stable, descending);
        if (start != 0) {
          auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, end);
          merge_kernel<<<config.block_per_grid.x,
                         config.thread_per_block.x,
                         0,
                         cu_stream>>>(out_data,
                                      start,
                                      out_data + start,
                                      end - start,
                                      ids_data,
                                      ids_data + start,
                                      temp_data,
                                      temp_ids,
                                      descending);
          thrust::copy(exec_policy, temp_ids, temp_ids + end, ids_data);
          thrust::copy(exec_policy, temp_data, temp_data + end, out_data);
        }
        start = end;
        end = std::min(start + per_number, size);
      }
    }
    return;
  }

  // Special case for full sort, speedup ~190x.
  if (axis == -1 || axis + 1 == in_dims.size()) {
    const int64_t input_height =
        common::product(slice_ddim(in_dims, 0, in_dims.size() - 1));
    const int64_t input_width = in_dims[in_dims.size() - 1];
    dev_ctx.template Alloc<int64_t>(indices);
    dev_ctx.template Alloc<T>(output);
    ArgFullSort<T, int64_t>(dev_ctx,
                            &input,
                            output,
                            indices,
                            input_height,
                            input_width,
                            descending);
  } else {
    // if not full sort, do transpose first
    std::vector<int> trans;
    for (int i = 0; i < axis; i++) {
      trans.push_back(i);
    }
    trans.push_back(in_dims.size() - 1);
    for (int i = axis + 1; i < in_dims.size() - 1; i++) {
      trans.push_back(i);
    }
    trans.push_back(axis);
    DDim trans_dims(in_dims);
    for (int i = 0; i < trans.size(); i++) {
      trans_dims[i] = in_dims[trans[i]];
    }

    DenseTensor trans_inp;
    trans_inp.Resize(trans_dims);
    T* trans_inp_data = dev_ctx.template Alloc<T>(&trans_inp);
    // Do transpose
    TransposeKernel<T, Context>(dev_ctx, input, trans, &trans_inp);

    const int64_t input_height =
        common::product(slice_ddim(trans_dims, 0, trans_dims.size() - 1));
    const int64_t input_width = trans_dims[trans_dims.size() - 1];

    DenseTensor tmp_out;
    tmp_out.Resize(trans_dims);
    dev_ctx.template Alloc<T>(&tmp_out);

    DenseTensor tmp_indices;
    // temp indices for sorting
    tmp_indices.Resize(trans_dims);
    dev_ctx.template Alloc<int64_t>(&tmp_indices);

    ArgFullSort<T, int64_t>(dev_ctx,
                            &trans_inp,
                            &tmp_out,
                            &tmp_indices,
                            input_height,
                            input_width,
                            descending);
    // delay output allocation until after transpose, to avoid
    // allocating too much memory
    dev_ctx.template Alloc<T>(output);
    dev_ctx.template Alloc<int64_t>(indices);
    // transpose back
    TransposeKernel<T, Context>(dev_ctx, tmp_out, trans, output);
    TransposeKernel<int64_t, Context>(dev_ctx, tmp_indices, trans, indices);
    return;
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(argsort,
                   GPU,
                   ALL_LAYOUT,
                   phi::ArgsortKernel,
                   float,
                   double,
                   int,
                   int64_t,
                   uint8_t,
                   int16_t,
                   phi::float16,
                   phi::bfloat16) {
  kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
}