// 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 <assert.h>
#include "paddle/phi/kernels/funcs/aligned_vector.h"
#include "paddle/phi/kernels/primitive/functor_primitives.h"

#define CUDART_INF __longlong_as_double(0x7ff0000000000000ULL)
#define CUDART_INF_F __int_as_float(0x7f800000)

namespace phi {
namespace funcs {

constexpr int kWarpSize = 32;

template <typename T>
__inline__ __device__ T Inf();

template <>
__inline__ __device__ float Inf<float>() {
  return CUDART_INF_F;
}

template <>
__inline__ __device__ double Inf<double>() {
  return CUDART_INF;
}

template <typename T,
          template <typename>
          class Functor,
          int ThreadGroupWidth = kWarpSize>
__inline__ __device__ T WarpAllReduce(T val) {
  for (int mask = ThreadGroupWidth / 2; mask > 0; mask /= 2) {
#if PADDLE_WITH_HIP
    val = Functor<T>()(val, __shfl_xor(0xffffffff, val, mask));
#else
    val = Functor<T>()(val, __shfl_xor_sync(0xffffffff, val, mask));
#endif
  }
  return val;
}

#if PADDLE_WITH_HIP
inline void GetNumBlocks(int64_t block_size,
                         int64_t max_blocks,
                         int64_t waves,
                         int* num_blocks) {
  int dev;
  PADDLE_ENFORCE_GPU_SUCCESS(hipGetDevice(&dev));
  int sm_count;
  PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceGetAttribute(
      &sm_count, hipDeviceAttributeMultiprocessorCount, dev));
  int tpm;
  PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceGetAttribute(
      &tpm, hipDeviceAttributeMaxThreadsPerMultiProcessor, dev));
  *num_blocks = std::max<int>(
      1, std::min<int64_t>(max_blocks, sm_count * tpm / block_size * waves));
}
#else
inline void GetNumBlocks(int64_t block_size,
                         int64_t max_blocks,
                         int64_t waves,
                         int* num_blocks) {
  int dev;
  PADDLE_ENFORCE_GPU_SUCCESS(cudaGetDevice(&dev));
  int sm_count;
  PADDLE_ENFORCE_GPU_SUCCESS(
      cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev));
  int tpm;
  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceGetAttribute(
      &tpm, cudaDevAttrMaxThreadsPerMultiProcessor, dev));
  *num_blocks = std::max<int>(
      1, std::min<int64_t>(max_blocks, sm_count * tpm / block_size * waves));
}
#endif

template <typename T,
          typename SourceType,
          typename Context,
          int VecSize,
          int ColsPerThread,
          int RowsPerThread,
          int ThreadGroupWidth,
          bool NeedPadding>
__global__ void LogsumexpWarpImpl(const Context& dev_ctx,
                                  const int64_t num_row,
                                  const int64_t num_col,
                                  const SourceType* in,
                                  SourceType* out) {
  static_assert(ColsPerThread % VecSize == 0, "");
  static_assert(ThreadGroupWidth <= kWarpSize, "");
  static_assert(kWarpSize % ThreadGroupWidth == 0, "");
  constexpr int num_read = ColsPerThread / VecSize;
  assert(num_col <= ColsPerThread * ThreadGroupWidth);
  const int group_id = blockIdx.x * blockDim.y + threadIdx.y;
  const int num_thread_group = gridDim.x * blockDim.y;
  const int thread_id = threadIdx.x;
  const int step = num_thread_group * RowsPerThread;

  using LoadType = AlignedVector<SourceType, VecSize>;
  using StoreType = AlignedVector<SourceType, RowsPerThread>;

  LoadType load_vec;
  StoreType store_vec;

  T buffer[RowsPerThread][ColsPerThread];

  for (int64_t cur_row = group_id * RowsPerThread; cur_row < num_row;
       cur_row += step) {
    T thread_max[RowsPerThread];
// Read data
#pragma unroll
    for (int row_id = 0; row_id < RowsPerThread; row_id++) {
      thread_max[row_id] = -Inf<T>();
      T* row_buffer = buffer[row_id];
#pragma unroll
      for (int read_id = 0; read_id < num_read; read_id++) {
        const int offset = read_id * VecSize;
        const int cur_col = (read_id * ThreadGroupWidth + thread_id) * VecSize;
        if (!NeedPadding || cur_col < num_col) {
          int64_t load_offset = ((cur_row + row_id) * num_col + cur_col);
          Load<SourceType, VecSize>(in + load_offset, &load_vec);
#pragma unroll
          for (int i = 0; i < VecSize; i++) {
            row_buffer[offset + i] = static_cast<T>(load_vec[i]);
            thread_max[row_id] =
                max(thread_max[row_id], row_buffer[offset + i]);
          }
        } else {
#pragma unroll
          for (int i = 0; i < VecSize; i++) {
            row_buffer[offset + i] = -Inf<T>();
          }
        }
      }
    }
    T warp_max[RowsPerThread];
// Get warp max
#pragma unroll
    for (int row_id = 0; row_id < RowsPerThread; row_id++) {
      warp_max[row_id] = WarpAllReduce<T, kps::MaxFunctor, ThreadGroupWidth>(
          thread_max[row_id]);
    }
    T thread_sum[RowsPerThread];
// Calculate
#pragma unroll
    for (int row_id = 0; row_id < RowsPerThread; row_id++) {
      thread_sum[row_id] = 0;
      T* row_buffer = buffer[row_id];
#pragma unroll
      for (int i = 0; i < ColsPerThread; i++) {
        thread_sum[row_id] += exp(row_buffer[i] - warp_max[row_id]);
      }
    }
// Get warp sum and write
#pragma unroll
    for (int row_id = 0; row_id < RowsPerThread; row_id++) {
      T res = log(WarpAllReduce<T, kps::AddFunctor, ThreadGroupWidth>(
          thread_sum[row_id]));
      store_vec[row_id] = static_cast<SourceType>(res + warp_max[row_id]);
    }
    if (thread_id == 0 && cur_row < num_row) {
      Store<SourceType, RowsPerThread>(store_vec, out + cur_row);
    }
  }
}

template <typename T,
          typename SourceType,
          typename Context,
          int VecSize,
          int ColsPerThread,
          int RowsPerThread,
          int ThreadGroupWidth,
          bool NeedPadding>
#if PADDLE_WITH_HIP
inline hipError_t LaunchLogsumexpWarp(const Context& dev_ctx,
                                      const int64_t num_row,
                                      const int64_t num_col,
                                      const SourceType* in,
                                      SourceType* out) {
#else
inline cudaError_t LaunchLogsumexpWarp(const Context& dev_ctx,
                                       const int64_t num_row,
                                       const int64_t num_col,
                                       const SourceType* in,
                                       SourceType* out) {
#endif
  constexpr int block_size = 128;
  constexpr int waves = 32;
  static_assert(block_size % ThreadGroupWidth == 0, "");
  constexpr int thread_groups_per_block = block_size / ThreadGroupWidth;
  dim3 block_dim(ThreadGroupWidth, thread_groups_per_block);
  const int64_t num_blocks =
      (num_row / RowsPerThread + thread_groups_per_block - 1) /
      thread_groups_per_block;
  int grid_dim_x;
  { GetNumBlocks(block_size, num_blocks, waves, &grid_dim_x); }
  LogsumexpWarpImpl<T,
                    SourceType,
                    Context,
                    VecSize,
                    ColsPerThread,
                    RowsPerThread,
                    ThreadGroupWidth,
                    NeedPadding>
      <<<grid_dim_x, block_dim, 0, dev_ctx.stream()>>>(
          dev_ctx, num_row, num_col, in, out);
#if PADDLE_WITH_HIP
  return hipPeekAtLastError();
#else
  return cudaPeekAtLastError();
#endif
}

template <typename T,
          typename SourceType,
          typename Context,
          int VecSize,
          int ColsPerThread,
          int RowsPerThread,
          int ThreadGroupWidth>
#if PADDLE_WITH_HIP
inline hipError_t DispatchLogsumexpWarpWithPadding(const Context& dev_ctx,
                                                   const int64_t num_row,
                                                   const int64_t num_col,
                                                   const SourceType* in,
                                                   SourceType* out) {
#else
inline cudaError_t DispatchLogsumexpWarpWithPadding(const Context& dev_ctx,
                                                    const int64_t num_row,
                                                    const int64_t num_col,
                                                    const SourceType* in,
                                                    SourceType* out) {
#endif
  if (num_col == ColsPerThread * ThreadGroupWidth) {
    return LaunchLogsumexpWarp<T,
                               SourceType,
                               Context,
                               VecSize,
                               ColsPerThread,
                               RowsPerThread,
                               ThreadGroupWidth,
                               false>(dev_ctx, num_row, num_col, in, out);
  } else {
    return LaunchLogsumexpWarp<T,
                               SourceType,
                               Context,
                               VecSize,
                               ColsPerThread,
                               RowsPerThread,
                               ThreadGroupWidth,
                               true>(dev_ctx, num_row, num_col, in, out);
  }
}

template <typename T, typename SourceType, typename Context, int VecSize>
#if PADDLE_WITH_HIP
typename std::enable_if<VecSize == 1, hipError_t>::type
DispatchLogsumexpWarpCols(const Context& dev_ctx,
                          const int64_t num_row,
                          const int64_t num_col,
                          const SourceType* in,
                          SourceType* out) {
#else
typename std::enable_if<VecSize == 1, cudaError_t>::type
DispatchLogsumexpWarpCols(const Context& dev_ctx,
                          const int64_t num_row,
                          const int64_t num_col,
                          const SourceType* in,
                          SourceType* out) {
#endif
  if (num_col <= 0) {
#if PADDLE_WITH_HIP
    return hipErrorInvalidValue;
#else
    return cudaErrorInvalidValue;
#endif
  }
#define HANDLE_THREAD_GROUP(thread_group_width)                    \
  if (num_col <= (thread_group_width)*VecSize) {                   \
    if (num_row % 2 == 0) {                                        \
      return DispatchLogsumexpWarpWithPadding<T,                   \
                                              SourceType,          \
                                              Context,             \
                                              VecSize,             \
                                              VecSize,             \
                                              2,                   \
                                              thread_group_width>( \
          dev_ctx, num_row, num_col, in, out);                     \
    } else {                                                       \
      return DispatchLogsumexpWarpWithPadding<T,                   \
                                              SourceType,          \
                                              Context,             \
                                              VecSize,             \
                                              VecSize,             \
                                              1,                   \
                                              thread_group_width>( \
          dev_ctx, num_row, num_col, in, out);                     \
    }                                                              \
  }
  HANDLE_THREAD_GROUP(1)
  HANDLE_THREAD_GROUP(2)
  HANDLE_THREAD_GROUP(4)
  HANDLE_THREAD_GROUP(8)
  HANDLE_THREAD_GROUP(16)
  HANDLE_THREAD_GROUP(32)
#undef HANDLE_ROWS
// if num_col > 32
#define HANDLE_COL(col)                                 \
  if (num_col <= (col)*kWarpSize) {                     \
    return DispatchLogsumexpWarpWithPadding<T,          \
                                            SourceType, \
                                            Context,    \
                                            VecSize,    \
                                            col,        \
                                            1,          \
                                            kWarpSize>( \
        dev_ctx, num_row, num_col, in, out);            \
  }

  HANDLE_COL(2)
  HANDLE_COL(3)
  HANDLE_COL(4)
  HANDLE_COL(5)
  HANDLE_COL(6)
  HANDLE_COL(7)
  HANDLE_COL(8)
  HANDLE_COL(9)
  HANDLE_COL(10)
  HANDLE_COL(11)
  HANDLE_COL(12)
  HANDLE_COL(13)
  HANDLE_COL(14)
  HANDLE_COL(15)
  HANDLE_COL(16)
  HANDLE_COL(17)
  HANDLE_COL(18)
  HANDLE_COL(19)
  HANDLE_COL(20)
  HANDLE_COL(21)
  HANDLE_COL(22)
  HANDLE_COL(23)
  HANDLE_COL(24)
  HANDLE_COL(25)
  HANDLE_COL(26)
  HANDLE_COL(27)
  HANDLE_COL(28)
  HANDLE_COL(29)
  HANDLE_COL(30)
  HANDLE_COL(31)
  HANDLE_COL(32)
#undef HANDLE_COL
#if PADDLE_WITH_HIP
  return hipErrorInvalidValue;
#else
  return cudaErrorInvalidValue;
#endif
}

template <typename T, typename SourceType, typename Context, int VecSize>
#if PADDLE_WITH_HIP
typename std::enable_if<VecSize == 2, hipError_t>::type
DispatchLogsumexpWarpCols(const Context& dev_ctx,
                          const int64_t num_row,
                          const int64_t num_col,
                          const SourceType* in,
                          SourceType* out) {
#else
typename std::enable_if<VecSize == 2, cudaError_t>::type
DispatchLogsumexpWarpCols(const Context& dev_ctx,
                          const int64_t num_row,
                          const int64_t num_col,
                          const SourceType* in,
                          SourceType* out) {
#endif
  if (num_col <= 0) {
#if PADDLE_WITH_HIP
    return hipErrorInvalidValue;
#else
    return cudaErrorInvalidValue;
#endif
  }
#define HANDLE_THREAD_GROUP(thread_group_width)                    \
  if (num_col <= (thread_group_width)*VecSize) {                   \
    if (num_row % 2 == 0) {                                        \
      return DispatchLogsumexpWarpWithPadding<T,                   \
                                              SourceType,          \
                                              Context,             \
                                              VecSize,             \
                                              VecSize,             \
                                              2,                   \
                                              thread_group_width>( \
          dev_ctx, num_row, num_col, in, out);                     \
    } else {                                                       \
      return DispatchLogsumexpWarpWithPadding<T,                   \
                                              SourceType,          \
                                              Context,             \
                                              VecSize,             \
                                              VecSize,             \
                                              1,                   \
                                              thread_group_width>( \
          dev_ctx, num_row, num_col, in, out);                     \
    }                                                              \
  }
  HANDLE_THREAD_GROUP(1)
  HANDLE_THREAD_GROUP(2)
  HANDLE_THREAD_GROUP(4)
  HANDLE_THREAD_GROUP(8)
  HANDLE_THREAD_GROUP(16)
  HANDLE_THREAD_GROUP(32)
#undef HANDLE_THREAD_GROUP
// if num_col > 32
#define HANDLE_COL(col)                                 \
  if (num_col <= (col)*kWarpSize) {                     \
    return DispatchLogsumexpWarpWithPadding<T,          \
                                            SourceType, \
                                            Context,    \
                                            VecSize,    \
                                            col,        \
                                            1,          \
                                            kWarpSize>( \
        dev_ctx, num_row, num_col, in, out);            \
  }

  HANDLE_COL(4)
  HANDLE_COL(6)
  HANDLE_COL(8)
  HANDLE_COL(10)
  HANDLE_COL(12)
  HANDLE_COL(14)
  HANDLE_COL(16)
  HANDLE_COL(18)
  HANDLE_COL(20)
  HANDLE_COL(22)
  HANDLE_COL(24)
  HANDLE_COL(26)
  HANDLE_COL(28)
  HANDLE_COL(30)
  HANDLE_COL(32)
#undef HANDLE_COL
#if PADDLE_WITH_HIP
  return hipErrorInvalidValue;
#else
  return cudaErrorInvalidValue;
#endif
}

template <typename T, typename SourceType, typename Context>
#if PADDLE_WITH_HIP
inline hipError_t DispatchLogsumexpWarp(const Context& dev_ctx,
                                        const int64_t num_row,
                                        const int64_t num_col,
                                        const SourceType* in,
                                        SourceType* out) {
#else
inline cudaError_t DispatchLogsumexpWarp(const Context& dev_ctx,
                                         const int64_t num_row,
                                         const int64_t num_col,
                                         const SourceType* in,
                                         SourceType* out) {
#endif
  // dispatch logsumexp warp with vecsize
  if (num_col % 2 == 0) {
    return DispatchLogsumexpWarpCols<T, SourceType, Context, 2>(
        dev_ctx, num_row, num_col, in, out);
  } else {
    return DispatchLogsumexpWarpCols<T, SourceType, Context, 1>(
        dev_ctx, num_row, num_col, in, out);
  }
}
}  // namespace funcs
}  // namespace phi