// 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.

#include "paddle/phi/kernels/cum_maxmin_kernel.h"

#include <numeric>

#include "paddle/common/hostdevice.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/math_function.h"

namespace phi {

template <
    typename T1,
    typename T2,
    typename BinaryOperation,
    typename std::enable_if<std::is_floating_point<T1>::value, int>::type = 0>
__device__ void binary_op_update(const T1 lhs,
                                 T1* rhs,
                                 const T2 lhs_idx,
                                 T2* rhs_idx,
                                 BinaryOperation binary_op) {
  if (!isnan(*rhs) && (isnan(lhs) || !binary_op(*rhs, lhs))) {
    *rhs = lhs;
    *rhs_idx = lhs_idx;
  }
}

template <typename T1,
          typename T2,
          typename BinaryOperation,
          typename std::enable_if<std::is_integral<T1>::value, int>::type = 0>
__device__ void binary_op_update(const T1 lhs,
                                 T1* rhs,
                                 const T2 lhs_idx,
                                 T2* rhs_idx,
                                 BinaryOperation binary_op) {
  if (!binary_op(*rhs, lhs)) {
    *rhs = lhs;
    *rhs_idx = lhs_idx;
  }
}

template <
    typename T1,
    typename T2,
    typename BinaryOperation,
    typename std::enable_if<std::is_floating_point<T1>::value, int>::type = 0>
__device__ void binary_op_update_v(const T1 lhs,
                                   T1* rhs,
                                   const T2 lhs_idx,
                                   T2* rhs_idx,
                                   BinaryOperation binary_op) {
  if (isnan(lhs) || (!isnan(*rhs) && binary_op(lhs, *rhs))) {
    *rhs = lhs;
    *rhs_idx = lhs_idx;
  }
}

template <typename T1,
          typename T2,
          typename BinaryOperation,
          typename std::enable_if<std::is_integral<T1>::value, int>::type = 0>
__device__ void binary_op_update_v(const T1 lhs,
                                   T1* rhs,
                                   const T2 lhs_idx,
                                   T2* rhs_idx,
                                   BinaryOperation binary_op) {
  if (binary_op(lhs, *rhs)) {
    *rhs = lhs;
    *rhs_idx = lhs_idx;
  }
}

template <typename T1,
          typename T2,
          int num_threads_x,
          int num_threads_y,
          class BinaryFunction>
__global__ void KernelScanInnerWithIndices(const T1* x_data,
                                           T1* values_data,
                                           T2* indices_data,
                                           int64_t num_rows,
                                           int64_t row_size,
                                           T1 init,
                                           BinaryFunction binary_op) {
  __shared__ T1 vbuf[num_threads_y][2 * num_threads_x];
  __shared__ T2 ibuf[num_threads_y][2 * num_threads_x];
  T1* row_buf = vbuf[threadIdx.y];
  T2* row_idx_buf = ibuf[threadIdx.y];

  for (int64_t block_row =
           static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.y);
       block_row < num_rows;
       block_row += blockDim.y * gridDim.x) {
    int64_t row = block_row + static_cast<int64_t>(threadIdx.y);
    const T1* row_self = x_data + row * row_size;
    T1* row_values = values_data + row * row_size;
    T2* row_indices = indices_data + row * row_size;
    T1 block_total = init;
    T2 block_idx_final = 0;
    // Perform scan on one block at a time, keeping track of the total value of
    // all blocks processed so far.
    for (int64_t block_col = 0; block_col < row_size;
         block_col += 2 * num_threads_x) {
      // Load data into shared memory (two values per thread).
      int64_t col1 = block_col + static_cast<int64_t>(threadIdx.x);
      int64_t col2 =
          block_col + num_threads_x + static_cast<int64_t>(threadIdx.x);
      if (row < num_rows) {
        if (col1 < row_size) {
          row_buf[threadIdx.x] = *reinterpret_cast<const T1*>(&row_self[col1]);
          row_idx_buf[threadIdx.x] = col1;
        } else {
          row_buf[threadIdx.x] = init;
        }

        if (col2 < row_size) {
          row_buf[num_threads_x + threadIdx.x] =
              *reinterpret_cast<const T1*>(&row_self[col2]);
          row_idx_buf[num_threads_x + threadIdx.x] = col2;
        } else {
          row_buf[num_threads_x + threadIdx.x] = init;
        }

        if (threadIdx.x == 0) {
          binary_op_update(block_total,
                           &row_buf[0],
                           block_idx_final,
                           &row_idx_buf[0],
                           binary_op);
        }
      }
      __syncthreads();

      // Parallel reduction (up-sweep).
      for (int s = num_threads_x, d = 1; s >= 1; s >>= 1, d <<= 1) {
        if (row < num_rows && threadIdx.x < s) {
          int offset = (2 * threadIdx.x + 1) * d - 1;
          binary_op_update(row_buf[offset],
                           &row_buf[offset + d],
                           row_idx_buf[offset],
                           &row_idx_buf[offset + d],
                           binary_op);
        }
        __syncthreads();
      }

      // Down-sweep.
      for (int s = 2, d = num_threads_x / 2; d >= 1; s <<= 1, d >>= 1) {
        if (row < num_rows && threadIdx.x < s - 1) {
          int offset = 2 * (threadIdx.x + 1) * d - 1;
          binary_op_update(row_buf[offset],
                           &row_buf[offset + d],
                           row_idx_buf[offset],
                           &row_idx_buf[offset + d],
                           binary_op);
        }
        __syncthreads();
      }

      // Write back to output.
      if (row < num_rows) {
        if (col1 < row_size) {
          row_values[col1] = row_buf[threadIdx.x];
          row_indices[col1] = row_idx_buf[threadIdx.x];
        }
        if (col2 < row_size) {
          row_values[col2] = row_buf[num_threads_x + threadIdx.x];
          row_indices[col2] = row_idx_buf[num_threads_x + threadIdx.x];
        }
      }
      block_total = row_buf[2 * num_threads_x - 1];
      block_idx_final = row_idx_buf[2 * num_threads_x - 1];
      __syncthreads();
    }
  }
}

template <typename T1, typename T2, class BinaryFunction>
__global__ void KernelScanOuterWithIndices(const T1* x_data,
                                           T1* values_data,
                                           T2* indices_data,
                                           const uint32_t num_orows,
                                           const uint32_t num_irows,
                                           const uint32_t row_size,
                                           T1 init,
                                           BinaryFunction binary_op) {
  for (uint32_t orow = blockIdx.x; orow < num_orows; orow += gridDim.x) {
    for (uint32_t irow = blockIdx.y * blockDim.x + threadIdx.x;
         irow < num_irows;
         irow += gridDim.y * blockDim.x) {
      const T1* x = x_data + orow * row_size * num_irows + irow;
      T1* values = values_data + orow * row_size * num_irows + irow;
      T2* indices = indices_data + orow * row_size * num_irows + irow;
      T1 out = init;
      T2 out_idx = 0;

      for (T2 col = 0; col < row_size; ++col) {
        const auto val = *reinterpret_cast<const T1*>(x);
        binary_op_update_v(val, &out, col, &out_idx, binary_op);
        *values = out;
        *indices = out_idx;
        x += num_irows;
        values += num_irows;
        indices += num_irows;
      }
    }
  }
}

template <typename T1, typename T2, typename BinaryFunction, typename Context>
void ScanWithIndicesKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           int axis,
                           T1 init,
                           DenseTensor* out,
                           DenseTensor* indices) {
  if (out && out->numel() == 0) {
    dev_ctx.template Alloc<T1>(out);
    dev_ctx.template Alloc<T2>(indices);
    return;
  }
  dev_ctx.template Alloc<T1>(out);
  dev_ctx.template Alloc<T2>(indices);
  // For 0D Tensor
  if (out->numel() == 1) {
    auto raw_dims = out->dims();
    Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), false, out);
    funcs::SetConstant<Context, T2> set_zero;
    set_zero(dev_ctx, indices, static_cast<T2>(0.0));
    out->Resize(raw_dims);
    indices->Resize(raw_dims);
    return;
  }

  BinaryFunction op;
  auto out_dims = out->dims();
  auto size = x.numel();

  PADDLE_ENFORCE_EQ(
      axis < out_dims.size() && axis >= (0 - out_dims.size()),
      true,
      common::errors::OutOfRange(
          "Attr(axis) is out of range, It's expected "
          "to be in range of [-%d, %d]. But received Attr(axis) = %d.",
          out_dims.size(),
          out_dims.size() - 1,
          axis));
  if (axis < 0) {
    axis += out_dims.size();
  }

  const T1* x_data = x.data<T1>();
  T1* values_data = out->data<T1>();
  T2* indices_data = indices->data<T2>();
  if (axis == out_dims.size() - 1) {
    int ndim = x.dims().size();
    int64_t row_size = x.dims()[ndim - 1];
    int64_t num_rows = x.numel() / row_size;

    dim3 threads(16, 32);
    dim3 grid(std::min(
        dev_ctx.GetCUDAMaxGridDimSize()[0],
        static_cast<unsigned int>(std::ceil(static_cast<float>(num_rows) /
                                            static_cast<float>(threads.y)))));

    KernelScanInnerWithIndices<T1, T2, 16, 32>
        <<<grid, threads, 0, dev_ctx.stream()>>>(
            x_data, values_data, indices_data, num_rows, row_size, init, op);
  } else {
    int64_t row_size = x.dims()[axis];
    auto sizes = vectorize(x.dims());

    const int64_t num_orows =
        std::accumulate(sizes.begin(),
                        sizes.begin() + axis,
                        int64_t(1),
                        [](int64_t a, int64_t b) { return a * b; });
    const int64_t num_irows =
        std::accumulate(sizes.begin() + axis + 1,
                        sizes.end(),
                        int64_t(1),
                        [](int64_t a, int64_t b) { return a * b; });

    dim3 threads(std::min(512, static_cast<int>(num_irows)));
    int64_t maxGridDim = dev_ctx.GetCUDAMaxGridDimSize()[1];
    dim3 grid(std::min(maxGridDim, num_orows),
              std::min(maxGridDim,
                       static_cast<int64_t>(
                           std::ceil(static_cast<double>(num_irows) /
                                     static_cast<double>(threads.x)))));

    KernelScanOuterWithIndices<T1, T2>
        <<<grid, threads, 0, dev_ctx.stream()>>>(x_data,
                                                 values_data,
                                                 indices_data,
                                                 num_orows,
                                                 num_irows,
                                                 row_size,
                                                 init,
                                                 op);
  }
}

template <typename T, typename Context>
void CummaxKernel(const Context& dev_ctx,
                  const DenseTensor& x,
                  int axis,
                  DataType dtype,
                  DenseTensor* out,
                  DenseTensor* indices) {
  T init = std::is_floating_point<T>::value
               ? (-1 * std::numeric_limits<T>::infinity())
               : std::numeric_limits<T>::lowest();
  if (dtype == DataType::INT32) {
    ScanWithIndicesKernel<T, int32_t, std::greater_equal<T>, Context>(
        dev_ctx, x, axis, init, out, indices);
  } else if (dtype == DataType::INT64) {
    ScanWithIndicesKernel<T, int64_t, std::greater_equal<T>, Context>(
        dev_ctx, x, axis, init, out, indices);
  }
}

template <typename T, typename Context>
void CumminKernel(const Context& dev_ctx,
                  const DenseTensor& x,
                  int axis,
                  DataType dtype,
                  DenseTensor* out,
                  DenseTensor* indices) {
  T init = std::is_floating_point<T>::value ? std::numeric_limits<T>::infinity()
                                            : std::numeric_limits<T>::max();
  if (dtype == DataType::INT32) {
    ScanWithIndicesKernel<T, int32_t, std::less_equal<T>, Context>(
        dev_ctx, x, axis, init, out, indices);
  } else if (dtype == DataType::INT64) {
    ScanWithIndicesKernel<T, int64_t, std::less_equal<T>, Context>(
        dev_ctx, x, axis, init, out, indices);
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(cummax,
                   GPU,
                   ALL_LAYOUT,
                   phi::CummaxKernel,
                   float,
                   double,
                   int32_t,
                   int64_t) {}

PD_REGISTER_KERNEL(cummin,
                   GPU,
                   ALL_LAYOUT,
                   phi::CumminKernel,
                   float,
                   double,
                   int32_t,
                   int64_t) {}