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

#pragma once
#include <algorithm>
#include <vector>
#include "gflags/gflags.h"
#include "paddle/fluid/framework/mixed_vector.h"
#include "paddle/fluid/memory/malloc.h"
#include "paddle/fluid/operators/math/concat_and_split.h"
#include "paddle/fluid/platform/cuda_graph_with_memory_pool.h"
#include "paddle/fluid/platform/device/gpu/gpu_primitives.h"

#include "paddle/pten/backends/gpu/gpu_context.h"

namespace pten {

template <typename T>
__global__ void ConcatKernel_(const T** inputs,
                              const int64_t* input_cols,
                              int col_size,
                              const int64_t output_rows,
                              const int64_t output_cols,
                              T* output) {
  int tid_x = blockIdx.x * blockDim.x + threadIdx.x;
  int curr_segment = 0;
  int curr_offset = input_cols[0];
  for (; tid_x < output_cols; tid_x += blockDim.x * gridDim.x) {
    int curr_col_offset = input_cols[curr_segment + 1];
    while (curr_col_offset <= tid_x) {
      curr_offset = curr_col_offset;
      ++curr_segment;
      curr_col_offset = input_cols[curr_segment + 1];
    }

    int local_col = tid_x - curr_offset;
    int segment_width = curr_col_offset - curr_offset;

    const T* input_ptr = inputs[curr_segment];
    int tid_y = blockIdx.y * blockDim.y + threadIdx.y;
    for (; tid_y < output_rows; tid_y += blockDim.y * gridDim.y)
      output[tid_y * output_cols + tid_x] =
          input_ptr[tid_y * segment_width + local_col];
  }
}

template <typename T>
__device__ void ConcatKernelDetail(const T** inputs_data,
                                   const int fixed_in_col,
                                   const int out_rows,
                                   const int out_cols,
                                   T* output_data) {
  int tid_x = blockIdx.x * blockDim.x + threadIdx.x;
  for (; tid_x < out_cols; tid_x += blockDim.x * gridDim.x) {
    int split = tid_x * 1.0 / fixed_in_col;
    int in_offset = tid_x - split * fixed_in_col;
    const T* input_ptr = inputs_data[split];
    int tid_y = blockIdx.y * blockDim.y + threadIdx.y;
    for (; tid_y < out_rows; tid_y += blockDim.y * gridDim.y) {
      output_data[tid_y * out_cols + tid_x] =
          input_ptr[tid_y * fixed_in_col + in_offset];
    }
  }
}

template <typename T>
__global__ void ConcatKernel_(const T* input_addr0,
                              const T* input_addr1,
                              const int64_t fixed_in_col,
                              const int64_t out_rows,
                              const int64_t out_cols,
                              T* output_data) {
  const T* inputs_data[2];
  inputs_data[0] = input_addr0;
  inputs_data[1] = input_addr1;
  ConcatKernelDetail<T>(
      inputs_data, fixed_in_col, out_rows, out_cols, output_data);
}

template <typename T>
__global__ void ConcatKernel_(const T* input_addr0,
                              const T* input_addr1,
                              const T* input_addr2,
                              const int64_t fixed_in_col,
                              const int64_t out_rows,
                              const int64_t out_cols,
                              T* output_data) {
  const T* inputs_data[3];
  inputs_data[0] = input_addr0;
  inputs_data[1] = input_addr1;
  inputs_data[2] = input_addr2;
  ConcatKernelDetail<T>(
      inputs_data, fixed_in_col, out_rows, out_cols, output_data);
}

template <typename T>
__global__ void ConcatKernel_(const T* input_addr0,
                              const T* input_addr1,
                              const T* input_addr2,
                              const T* input_addr3,
                              const int64_t fixed_in_col,
                              const int64_t out_rows,
                              const int64_t out_cols,
                              T* output_data) {
  const T* inputs_data[4];
  inputs_data[0] = input_addr0;
  inputs_data[1] = input_addr1;
  inputs_data[2] = input_addr2;
  inputs_data[3] = input_addr3;
  ConcatKernelDetail<T>(
      inputs_data, fixed_in_col, out_rows, out_cols, output_data);
}

template <typename T>
__global__ void ConcatKernel_(const T** inputs_data,
                              const int in_num,
                              const int64_t fixed_in_col,
                              const int64_t out_rows,
                              const int64_t out_cols,
                              T* output_data) {
  ConcatKernelDetail<T>(
      inputs_data, fixed_in_col, out_rows, out_cols, output_data);
}

template <typename T>
__global__ void SplitKernel(const T* input_data,
                            const int64_t in_row,
                            const int64_t in_col,
                            const int64_t* out_cols,
                            int out_cols_size,
                            T** outputs_data) {
  int tid_x = blockIdx.x * blockDim.x + threadIdx.x;
  int curr_segment = 0;
  int curr_offset = out_cols[0];
  for (; tid_x < in_col; tid_x += blockDim.x * gridDim.x) {
    int curr_col_offset = out_cols[curr_segment + 1];
    while (curr_col_offset <= tid_x) {
      curr_offset = curr_col_offset;
      ++curr_segment;
      curr_col_offset = out_cols[curr_segment + 1];
    }

    int local_col = tid_x - curr_offset;
    int segment_width = curr_col_offset - curr_offset;
    T* output_ptr = outputs_data[curr_segment];
    if (output_ptr != nullptr) {
      int tid_y = blockIdx.y * blockDim.y + threadIdx.y;
      for (; tid_y < in_row; tid_y += blockDim.y * gridDim.y)
        output_ptr[tid_y * segment_width + local_col] =
            input_data[tid_y * in_col + tid_x];
    }
  }
}

template <typename T>
__device__ void SplitKernelDetail(const T* input_data,
                                  const int in_row,
                                  const int in_col,
                                  const int fixed_out_col,
                                  T** outputs_data) {
  int tid_x = blockIdx.x * blockDim.x + threadIdx.x;
  for (; tid_x < in_col; tid_x += blockDim.x * gridDim.x) {
    int split = tid_x / fixed_out_col;
    int in_offset = tid_x - split * fixed_out_col;
    T* output_ptr = outputs_data[split];
    if (output_ptr != nullptr) {
      int tid_y = blockIdx.y * blockDim.y + threadIdx.y;
      for (; tid_y < in_row; tid_y += blockDim.y * gridDim.y)
        output_ptr[tid_y * fixed_out_col + in_offset] =
            input_data[tid_y * in_col + tid_x];
    }
  }
}

template <typename T>
__global__ void SplitKernel(const T* input_data,
                            const int64_t in_row,
                            const int64_t in_col,
                            const int64_t fixed_out_col,
                            T** outputs_data) {
  SplitKernelDetail<T>(input_data, in_row, in_col, fixed_out_col, outputs_data);
}

template <typename T>
__global__ void SplitKernel(const T* input_data,
                            const int64_t in_row,
                            const int64_t in_col,
                            const int64_t fixed_out_col,
                            T* outputs_addr0,
                            T* outputs_addr1) {
  T* outputs_data[2];
  outputs_data[0] = outputs_addr0;
  outputs_data[1] = outputs_addr1;
  SplitKernelDetail<T>(input_data, in_row, in_col, fixed_out_col, outputs_data);
}

template <typename T>
__global__ void SplitKernel(const T* input_data,
                            const int64_t in_row,
                            const int64_t in_col,
                            const int64_t fixed_out_col,
                            T* outputs_addr0,
                            T* outputs_addr1,
                            T* outputs_addr2) {
  T* outputs_data[3];
  outputs_data[0] = outputs_addr0;
  outputs_data[1] = outputs_addr1;
  outputs_data[2] = outputs_addr2;
  SplitKernelDetail<T>(input_data, in_row, in_col, fixed_out_col, outputs_data);
}

template <typename T>
__global__ void SplitKernel(const T* input_data,
                            const int64_t in_row,
                            const int64_t in_col,
                            const int64_t fixed_out_col,
                            T* outputs_addr0,
                            T* outputs_addr1,
                            T* outputs_addr2,
                            T* outputs_addr3) {
  T* outputs_data[4];
  outputs_data[0] = outputs_addr0;
  outputs_data[1] = outputs_addr1;
  outputs_data[2] = outputs_addr2;
  outputs_data[3] = outputs_addr3;
  SplitKernelDetail<T>(input_data, in_row, in_col, fixed_out_col, outputs_data);
}

static inline void GetBlockDims(
    const paddle::platform::CUDADeviceContext& context,
    int64_t num_rows,
    int64_t num_cols,
    dim3* block_dims,
    dim3* grid_dims) {
  // Set the thread block and grid according to CurrentDeviceId
  const int kThreadsPerBlock = 1024;
  int block_cols = kThreadsPerBlock;
  if (num_cols < kThreadsPerBlock) {  // block_cols is aligned by 32.
    block_cols = ((num_cols + 31) >> 5) << 5;
  }
  int block_rows = kThreadsPerBlock / block_cols;
  *block_dims = dim3(block_cols, block_rows, 1);

  int max_threads = context.GetMaxPhysicalThreadCount();
  int64_t max_blocks = std::max(max_threads / kThreadsPerBlock, 1);

  int grid_cols =
      std::min((num_cols + block_cols - 1) / block_cols, max_blocks);
  int grid_rows = std::min(max_blocks / grid_cols,
                           std::max(num_rows / block_rows, (int64_t)1));
  *grid_dims = dim3(grid_cols, grid_rows, 1);
}

/*
 * All tensors' dimension should be the same and the values of
 * each dimension must be the same, except the axis dimension.
 */
template <typename T, typename Context>
void ConcatImpl(const Context& context,
                const std::vector<pten::DenseTensor>& input,
                int axis,
                pten::DenseTensor* output) {
  // TODO(zcd): Add input data validity checking
  int in_num = input.size();
  int64_t in_row = 1;
  auto dim_0 = input[0].dims();
  for (int i = 0; i < axis; ++i) {
    in_row *= dim_0[i];
  }
  int64_t in_col = input[0].numel() / in_row;
  int64_t out_row = in_row, out_col = 0;

  int inputs_col_num = in_num + 1;
  std::vector<const T*> inputs_data_vec(in_num);
  std::vector<int64_t> inputs_col_vec(inputs_col_num);
  const T** inputs_data = inputs_data_vec.data();
  int64_t* inputs_col = inputs_col_vec.data();

// There are some differences between hip runtime and NV runtime.
// In NV, when the pageable memory data less than 64K is transferred from
// hosttodevice, it will be automatically asynchronous.
// However, only pinned memory in hip can copy asynchronously
// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#concurrent-execution-host-device
// 3.2.6.1. Concurrent Execution between Host and Device
// Memory copies from host to device of a memory block of 64 KB or less
#ifdef PADDLE_WITH_HIP
  paddle::memory::AllocationPtr data_alloc, col_alloc;
  // TODO(chentianyu03): try to find a method to remove the Alloc function
  data_alloc = paddle::memory::Alloc(paddle::platform::CUDAPinnedPlace(),
                                     in_num * sizeof(T*));
  inputs_data = reinterpret_cast<const T**>(data_alloc->ptr());
  // TODO(chentianyu03): try to find a method to remove the Alloc function
  col_alloc = paddle::memory::Alloc(paddle::platform::CUDAPinnedPlace(),
                                    inputs_col_num * sizeof(int));
  inputs_col = reinterpret_cast<int64_t*>(col_alloc->ptr());
#endif

  inputs_col[0] = 0;
  bool has_same_shape = true;
  for (int i = 0; i < in_num; ++i) {
    int64_t t_cols = input[i].numel() / in_row;
    if (has_same_shape) {
      if (t_cols != in_col) has_same_shape = false;
    }
    out_col += t_cols;
    inputs_col[i + 1] = out_col;
    inputs_data[i] = input[i].data<T>();
  }

  dim3 block_dims;
  dim3 grid_dims;
  GetBlockDims(context, out_row, out_col, &block_dims, &grid_dims);

  paddle::memory::allocation::AllocationPtr tmp_dev_ins_data;
  const T** dev_ins_data = nullptr;
  if (!has_same_shape || in_num < 2 || in_num > 4) {
    tmp_dev_ins_data = paddle::memory::Alloc(context, in_num * sizeof(T*));
    auto* restored = paddle::platform::RestoreHostMemIfCapturingCUDAGraph(
        inputs_data, in_num);
    paddle::memory::Copy(context.GetPlace(),
                         tmp_dev_ins_data->ptr(),
                         paddle::platform::CPUPlace(),
                         restored,
                         in_num * sizeof(T*),
                         context.stream());
    dev_ins_data = reinterpret_cast<const T**>(tmp_dev_ins_data->ptr());
  }

  if (has_same_shape) {
    if (in_num == 2) {
      ConcatKernel_<<<grid_dims, block_dims, 0, context.stream()>>>(
          inputs_data[0],
          inputs_data[1],
          in_col,
          out_row,
          out_col,
          output->data<T>());
    } else if (in_num == 3) {
      ConcatKernel_<<<grid_dims, block_dims, 0, context.stream()>>>(
          inputs_data[0],
          inputs_data[1],
          inputs_data[2],
          in_col,
          out_row,
          out_col,
          output->data<T>());
    } else if (in_num == 4) {
      ConcatKernel_<<<grid_dims, block_dims, 0, context.stream()>>>(
          inputs_data[0],
          inputs_data[1],
          inputs_data[2],
          inputs_data[3],
          in_col,
          out_row,
          out_col,
          output->data<T>());
    } else {
      ConcatKernel_<<<grid_dims, block_dims, 0, context.stream()>>>(
          dev_ins_data, in_num, in_col, out_row, out_col, output->data<T>());
    }
  } else {
    auto tmp_dev_ins_col_data =
        paddle::memory::Alloc(context, inputs_col_num * sizeof(int64_t));

    auto* restored = paddle::platform::RestoreHostMemIfCapturingCUDAGraph(
        inputs_col, inputs_col_num);
    paddle::memory::Copy(context.GetPlace(),
                         tmp_dev_ins_col_data->ptr(),
                         paddle::platform::CPUPlace(),
                         restored,
                         inputs_col_num * sizeof(int64_t),
                         context.stream());
    int64_t* dev_ins_col_data =
        static_cast<int64_t*>(tmp_dev_ins_col_data->ptr());

    ConcatKernel_<<<grid_dims, block_dims, 0, context.stream()>>>(
        dev_ins_data,
        dev_ins_col_data,
        static_cast<int>(inputs_col_num),
        out_row,
        out_col,
        output->data<T>());
  }

#ifdef PADDLE_WITH_HIP
  // Prevent the pinned memory value from being covered and release the memory
  // after the launch kernel of the stream is executed (reapply pinned memory
  // next time)
  auto* data_alloc_released = data_alloc.release();
  auto* col_alloc_released = col_alloc.release();
  context.AddStreamCallback([data_alloc_released, col_alloc_released] {
    paddle::memory::allocation::Allocator::AllocationDeleter(
        data_alloc_released);
    paddle::memory::allocation::Allocator::AllocationDeleter(
        col_alloc_released);
  });
#endif
}

/*
 * All tensors' dimension should be the same and the values of
 * each dimension must be the same, except the axis dimension.
 */
template <typename T, typename Context>
void SplitImpl(const Context& context,
               const pten::DenseTensor& input,
               const std::vector<const pten::DenseTensor*>& ref_inputs,
               int axis,
               std::vector<pten::DenseTensor*>* outputs) {
  // NOTE(zhiqiu): split a tensor of shape [0,3,4] at axis=1, result in 3
  // tensors of shape [0,1,4]
  if (input.numel() == 0) {
    return;
  }

  // TODO(zcd): Add input data validity checking
  int o_num = outputs->size();
  int64_t out_row = 1;
  auto dim_0 = ref_inputs[0]->dims();
  for (int i = 0; i < axis; ++i) {
    out_row *= dim_0[i];
  }

  int64_t out0_col = ref_inputs[0]->numel() / out_row;
  int64_t in_col = 0, in_row = out_row;
  bool has_same_shape = true;

  int outputs_cols_num = o_num + 1;
  std::vector<T*> outputs_data_vec(o_num);
  std::vector<int64_t> outputs_cols_vec(outputs_cols_num);
  T** outputs_data = outputs_data_vec.data();
  int64_t* outputs_cols = outputs_cols_vec.data();

// There are some differences between hip runtime and NV runtime.
// In NV, when the pageable memory data less than 64K is transferred from
// hosttodevice, it will be automatically asynchronous.
// However, only pinned memory in hip can copy asynchronously
// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#concurrent-execution-host-device
// 3.2.6.1. Concurrent Execution between Host and Device
// Memory copies from host to device of a memory block of 64 KB or less
#ifdef PADDLE_WITH_HIP
  paddle::memory::AllocationPtr data_alloc, cols_alloc;
  // TODO(chentianyu03): try to find a method to remove the Alloc function
  data_alloc = paddle::memory::Alloc(paddle::platform::CUDAPinnedPlace(),
                                     o_num * sizeof(T*));
  outputs_data = reinterpret_cast<T**>(data_alloc->ptr());
  // TODO(chentianyu03): try to find a method to remove the Alloc function
  cols_alloc = paddle::memory::Alloc(paddle::platform::CUDAPinnedPlace(),
                                     (outputs_cols_num) * sizeof(int64_t));
  outputs_cols = reinterpret_cast<int64_t*>(cols_alloc->ptr());
#endif

  outputs_cols[0] = 0;
  for (int i = 0; i < o_num; ++i) {
    int64_t t_col = ref_inputs.at(i)->numel() / out_row;
    if (has_same_shape) {
      if (t_col != out0_col) has_same_shape = false;
    }
    in_col += t_col;
    outputs_cols[i + 1] = in_col;
    if (outputs->at(i) != nullptr) {
      outputs_data[i] = outputs->at(i)->data<T>();
    } else {
      outputs_data[i] = nullptr;
    }
  }

  dim3 block_dims;
  dim3 grid_dims;
  GetBlockDims(context, out_row, in_col, &block_dims, &grid_dims);

  paddle::memory::allocation::AllocationPtr tmp_dev_outs_data;
  T** dev_out_gpu_data = nullptr;
  if (!has_same_shape || o_num < 2 || o_num > 4) {
    // TODO(chentianyu03): try to find a method to remove the Alloc function
    tmp_dev_outs_data = paddle::memory::Alloc(context, o_num * sizeof(T*));
    auto* restored = paddle::platform::RestoreHostMemIfCapturingCUDAGraph(
        outputs_data, o_num);
    paddle::memory::Copy(context.GetPlace(),
                         tmp_dev_outs_data->ptr(),
                         paddle::platform::CPUPlace(),
                         restored,
                         o_num * sizeof(T*),
                         context.stream());
    dev_out_gpu_data = reinterpret_cast<T**>(tmp_dev_outs_data->ptr());
  }

  if (has_same_shape) {
    if (o_num == 2) {
      SplitKernel<<<grid_dims, block_dims, 0, context.stream()>>>(
          input.data<T>(),
          in_row,
          in_col,
          out0_col,
          outputs_data[0],
          outputs_data[1]);
    } else if (o_num == 3) {
      SplitKernel<<<grid_dims, block_dims, 0, context.stream()>>>(
          input.data<T>(),
          in_row,
          in_col,
          out0_col,
          outputs_data[0],
          outputs_data[1],
          outputs_data[2]);
    } else if (o_num == 4) {
      SplitKernel<<<grid_dims, block_dims, 0, context.stream()>>>(
          input.data<T>(),
          in_row,
          in_col,
          out0_col,
          outputs_data[0],
          outputs_data[1],
          outputs_data[2],
          outputs_data[3]);
    } else {
      SplitKernel<<<grid_dims, block_dims, 0, context.stream()>>>(
          input.data<T>(), in_row, in_col, out0_col, dev_out_gpu_data);
    }
  } else {
    auto tmp_dev_ins_col_data =
        // TODO(chentianyu03): try to find a method to remove the Alloc function
        paddle::memory::Alloc(context, outputs_cols_num * sizeof(int64_t));
    auto* restored = paddle::platform::RestoreHostMemIfCapturingCUDAGraph(
        outputs_cols, outputs_cols_num);
    paddle::memory::Copy(context.GetPlace(),
                         tmp_dev_ins_col_data->ptr(),
                         paddle::platform::CPUPlace(),
                         restored,
                         outputs_cols_num * sizeof(int64_t),
                         context.stream());
    int64_t* dev_outs_col_data =
        reinterpret_cast<int64_t*>(tmp_dev_ins_col_data->ptr());

    SplitKernel<<<grid_dims, block_dims, 0, context.stream()>>>(
        input.data<T>(),
        in_row,
        in_col,
        dev_outs_col_data,
        static_cast<int>(outputs_cols_num),
        dev_out_gpu_data);
  }
#ifdef PADDLE_WITH_HIP
  // Prevent the pinned memory value from being covered and release the memory
  // after the launch kernel of the stream is executed (reapply pinned memory
  // next time)
  auto* data_alloc_released = data_alloc.release();
  auto* cols_alloc_released = cols_alloc.release();
  context.AddStreamCallback([data_alloc_released, cols_alloc_released] {
    paddle::memory::allocation::Allocator::AllocationDeleter(
        data_alloc_released);
    paddle::memory::allocation::Allocator::AllocationDeleter(
        cols_alloc_released);
  });
#endif
}

}  // namespace pten