/* Copyright (c) 2023 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/backends/gpu/cuda/cuda_graph_with_memory_pool.h"

namespace phi {
namespace funcs {
namespace sparse {

/* Given the indices of a sparse tensor, return a vector of offsets
for the entries in the equivalent dense tensor. */
template <typename IntT, typename Context>
inline DenseTensor GetOffsets(const Context& dev_ctx,
                              const DenseTensor& indices,
                              const std::vector<IntT>& sizes,
                              const IntT dim) {
#ifdef __HIPCC__
  const auto& policy = thrust::hip::par.on(dev_ctx.stream());
#else
  const auto& policy = thrust::cuda::par.on(dev_ctx.stream());
#endif

  auto ndim = indices.dims()[0];
  auto nnz = indices.dims()[1];
  std::vector<IntT> host_strides(ndim, 1);
  if (ndim > 1) {
    for (IntT i = ndim - 2; i >= 0; i--) {
      host_strides[i] = host_strides[i + 1] * (i + 1 == dim ? 1 : sizes[i + 1]);
    }
  }

  const IntArray strides_shape(vectorize<IntT>(indices.dims()));
  DenseTensor strides = Empty<IntT>(dev_ctx, strides_shape);
  auto strides_ptr = strides.data<IntT>();
#if defined(__NVCC__) || defined(__HIPCC__)
  const IntT* stable_st =
      phi::backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
          host_strides.data(), host_strides.size());
#else
  const IntT* stable_st = host_strides.data();
#endif
  memory_utils::Copy(dev_ctx.GetPlace(),
                     strides_ptr,
                     CPUPlace(),
                     stable_st,
                     sizeof(IntT) * host_strides.size(),
                     dev_ctx.stream());

  DenseTensor offsets = Empty<IntT>(dev_ctx, {nnz});
  auto indices_ptr = indices.data<IntT>();

  thrust::transform(
      policy,
      thrust::make_counting_iterator(IntT(0)),
      thrust::make_counting_iterator(IntT(nnz)),
      thrust::device_ptr<IntT>(offsets.data<IntT>()),
      [strides_ptr, indices_ptr, nnz, dim, ndim] __device__(IntT x) {
        IntT pool_index = 0;
        for (IntT j = 0; j < ndim; j++) {
          if (j != dim) {
            auto indice_cur_ptr = indices_ptr + j * nnz + x;
            auto stride = strides_ptr[j];
            pool_index += stride * (*indice_cur_ptr);
          }
        }
        return pool_index;
      });
  return offsets;
}

/* Return pools of indices that align with the given dimension and the
corresponding max values for each pool. */
template <typename T,
          typename IntT,
          typename Context,
          bool requireMxRows = true>
std::tuple<DenseTensor, DenseTensor, DenseTensor, DenseTensor> ComputePoolMax(
    const Context& dev_ctx,
    const DenseTensor& indices,
    const DenseTensor& values,
    const std::vector<IntT>& sizes,
    IntT nvalues,
    const IntT dim) {
#ifdef __HIPCC__
  const auto& policy = thrust::hip::par.on(dev_ctx.stream());
#else
  const auto& policy = thrust::cuda::par.on(dev_ctx.stream());
#endif
  using thrust_ptr = thrust::device_ptr<IntT>;
  auto nnz = indices.dims()[1];
  DenseTensor offsets =
      funcs::sparse::GetOffsets<IntT, Context>(dev_ctx, indices, sizes, dim);
  auto offsets_ptr = offsets.data<IntT>();

  DenseTensor sorted_indices = Empty<IntT>(dev_ctx, {nnz});
  thrust_ptr sorted_indices_thrust_ptr(sorted_indices.data<IntT>());
  thrust::sequence(
      policy, sorted_indices_thrust_ptr, sorted_indices_thrust_ptr + nnz, 0);

  /* sort indices corresponding to offsets */
  thrust::sort(policy,
               sorted_indices_thrust_ptr,
               sorted_indices_thrust_ptr + nnz,
               [offsets_ptr] __device__(IntT x, IntT y) {
                 return offsets_ptr[x] < offsets_ptr[y];
               });

  DenseTensor pool_sizes = Empty<IntT>(dev_ctx, {nnz});

  /* reduce the elements which are grouped by pool index,
  returns all the pool indexes with unique offset value for each. */
  auto new_end =
      thrust::reduce_by_key(policy,
                            sorted_indices_thrust_ptr,
                            sorted_indices_thrust_ptr + nnz,
                            thrust::make_constant_iterator(IntT(1)),
                            thrust::make_discard_iterator(),
                            thrust_ptr(pool_sizes.data<IntT>()),
                            [offsets_ptr] __device__(IntT x, IntT y) {
                              return offsets_ptr[x] == offsets_ptr[y];
                            });
  auto new_sz =
      thrust::distance(thrust_ptr(pool_sizes.data<IntT>()), new_end.second);
  pool_sizes.Resize({new_sz});

  DenseTensor pool_offsets;
  pool_offsets.Resize({new_sz});
  dev_ctx.template Alloc<T>(&pool_offsets);
  phi::Copy(dev_ctx, pool_sizes, dev_ctx.GetPlace(), false, &pool_offsets);

  /* accumulate value for each pool index */
  thrust_ptr pool_offsets_thrust_ptr(pool_offsets.data<IntT>());
  thrust::exclusive_scan(policy,
                         pool_offsets_thrust_ptr,
                         pool_offsets_thrust_ptr + new_sz,
                         pool_offsets_thrust_ptr);

  DenseTensor mx_buffer;
  if (requireMxRows) {
    mx_buffer = phi::Full<T>(
        dev_ctx, {new_sz * nvalues}, -std::numeric_limits<T>::infinity());
    auto mx_buffer_ptr = mx_buffer.data<T>();

    auto pool_sizes_ptr = pool_sizes.data<IntT>();
    auto sorted_indices_ptr = sorted_indices.data<IntT>();
    auto pool_offsets_ptr = pool_offsets.data<IntT>();
    auto values_ptr = values.data<T>();

    /* calculate max value in each pool. */
    thrust::for_each(policy,
                     thrust::make_counting_iterator(IntT(0)),
                     thrust::make_counting_iterator(IntT(new_sz)),
                     [sorted_indices_ptr,
                      pool_sizes_ptr,
                      pool_offsets_ptr,
                      mx_buffer_ptr,
                      values_ptr,
                      nvalues] __device__(IntT index) {
                       IntT curr_pool_size = pool_sizes_ptr[index];
                       auto mx_row = mx_buffer_ptr + index * nvalues;
                       IntT offset = pool_offsets_ptr[index];
                       for (IntT p = 0; p < curr_pool_size; p++) {
                         IntT i = *(sorted_indices_ptr + offset + p);
                         for (IntT j = 0; j < nvalues; j++) {
                           auto value_tmp = *(values_ptr);
                           mx_row[j] = std::max(mx_row[j], value_tmp);
                         }
                       }
                     });
  }
  return std::make_tuple(sorted_indices, pool_offsets, pool_sizes, mx_buffer);
}

inline int GetNumThreads(int nElem) {
#if defined(PADLDE_WITH_ROCM)
  int threadSizes[5] = {16, 32, 64, 128, 256};
#else
  int threadSizes[5] = {32, 64, 128, 256, 512};
#endif
  for (int i = 0; i != 5; ++i) {
    if (nElem <= threadSizes[i]) {
      return threadSizes[i];
    }
  }
  return threadSizes[4];
}

}  // namespace sparse
}  // namespace funcs
}  // namespace phi