// Copyright (c) 2024 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/funcs/math/beam_search.h"
#include "paddle/phi/backends/gpu/gpu_device_function.h"
#include "paddle/phi/backends/gpu/gpu_launch_config.h"

namespace phi {
namespace math {

struct Triple {
  __device__ __forceinline__ Triple() = default;
  __device__ __forceinline__ Triple(int o, int i, float s)
      : offset(o), id(i), score(s) {}

  __device__ __forceinline__ void set(int o, int i, float s) {
    offset = o;
    id = i;
    score = s;
  }

  __device__ __forceinline__ void operator=(const Triple& in) {
    offset = in.offset;
    id = in.id;
    score = in.score;
  }

  __device__ __forceinline__ bool operator<(const float s) const {
    return score < s;
  }

  __device__ __forceinline__ bool operator<(const Triple& in) const {
    return (score < in.score) || ((score == in.score) && (offset < in.offset));
  }

  int offset;
  int id;
  float score;
};

__device__ __forceinline__ void Insert(Triple* top_beam,
                                       const Triple& p,
                                       int beam_size) {
  if (p < top_beam[beam_size - 1]) {
    return;
  }
  for (int k = beam_size - 2; k >= 0; --k) {
    if (top_beam[k] < p) {
      top_beam[k + 1] = top_beam[k];
    } else {
      top_beam[k + 1] = p;
      return;
    }
  }
  top_beam[0] = p;
}

template <int MaxThreadsPerSeq, bool IsAccumulated = true>
__device__ __forceinline__ int SelectTopBeam(Triple* top_beam,
                                             const int64_t* pre_ids,
                                             const float* pre_scores,
                                             const int64_t* ids,
                                             const float* scores,
                                             const int seq_offset_start,
                                             const int seq_offset_end,
                                             const int seq_width,
                                             int beam_size,
                                             int end_id,
                                             int used_threads) {
  // top_beam is shared memory
  const int tid = threadIdx.x;
  const int tid_of_seq = threadIdx.x % MaxThreadsPerSeq;

  int num_used_threads = used_threads;

  Triple* top_beam_local = top_beam + tid * beam_size;
  if (tid_of_seq < num_used_threads) {
    for (int i = 0; i < beam_size; ++i) {
      top_beam_local[i].set(-1, -1, -INFINITY);
    }

    for (int offset = seq_offset_start; offset < seq_offset_end; ++offset) {
      int pre_id = static_cast<int>(pre_ids[offset]);
      if (pre_id == end_id) {
        if (tid_of_seq == 0) {
          Triple tmp(offset, end_id, pre_scores[offset]);
          Insert(top_beam_local, tmp, beam_size);
        }
      } else {
        int64_t index = static_cast<int64_t>(offset) * seq_width + tid_of_seq;
        if (!IsAccumulated) {
          float pre_score = pre_scores[offset];
          for (int i = tid_of_seq; i < seq_width; i += num_used_threads) {
            float score = pre_score + __logf(scores[index]);
            int id = ids ? static_cast<int>(ids[index]) : i;
            Triple tmp(offset, id, score);
            Insert(top_beam_local, tmp, beam_size);
            index += num_used_threads;
          }
        } else {
          for (int i = tid_of_seq; i < seq_width; i += num_used_threads) {
            int id = ids ? static_cast<int>(ids[index]) : i;
            float score = scores[index];
            Triple tmp(offset, id, score);
            Insert(top_beam_local, tmp, beam_size);
            index += num_used_threads;
          }
        }
      }
    }
  }

  while (num_used_threads > 1) {
    if (num_used_threads > 16) {
      __syncthreads();
    }

    if ((num_used_threads & 0x1) != 0) {
      // If num_used_threads is a odd number, merge local top_beam of thread 0
      // and num_used_threads - 1
      if (tid_of_seq == 0) {
        int index_in_sh = (num_used_threads - 1 + tid) * beam_size;
        for (int i = 0; i < beam_size; i++) {
          Insert(top_beam_local, top_beam[index_in_sh], beam_size);
          index_in_sh++;
        }
      }
    }

    num_used_threads = num_used_threads >> 1;
    if (tid_of_seq < num_used_threads) {
      int index_in_sh = (num_used_threads + tid) * beam_size;
      for (int i = 0; i < beam_size; i++) {
        Insert(top_beam_local, top_beam[index_in_sh], beam_size);
        index_in_sh++;
      }
    }
  }

  if (tid_of_seq == 0) {
    int num_items = 0;
    for (int i = 0; i < beam_size; ++i) {
      num_items =
          (top_beam_local[i].score > -INFINITY) ? num_items + 1 : num_items;
    }
    return num_items;
  }

  return 0;
}

__device__ __forceinline__ bool PruneEndBeams(Triple* top_beam_local,
                                              const int64_t* pre_ids,
                                              const int end_id,
                                              int num_items) {
  bool finish_flag = true;
  for (int i = 0; i < num_items; ++i) {
    int offset = top_beam_local[i].offset;
    if (top_beam_local[i].id != end_id ||
        static_cast<int>(pre_ids[offset]) != end_id) {
      finish_flag = false;
      break;
    }
  }
  return finish_flag;
}

template <bool ReturnParentIdx = false>
__device__ __forceinline__ void WriteBack(int64_t* selected_ids,
                                          float* selected_scores,
                                          int* parent_idx,
                                          size_t* selected_offsets,
                                          Triple* top_beam_local,
                                          const int seq_offset_start,
                                          const int seq_offset_end,
                                          const int selected_seq_start,
                                          const int selected_seq_length) {
  const int tid = threadIdx.x;  // use 1 thread only for each sequence
  int global_index = selected_seq_start;
  for (int global_offset = seq_offset_start; global_offset < seq_offset_end;
       ++global_offset) {
    for (int local_index = 0; local_index < selected_seq_length;
         ++local_index) {
      if (top_beam_local[local_index].offset == global_offset) {
        selected_ids[global_index] =
            static_cast<int64_t>(top_beam_local[local_index].id);
        selected_scores[global_index] = top_beam_local[local_index].score;
        if (ReturnParentIdx) {
          parent_idx[global_index] = static_cast<int>(global_offset);
        }
        global_index++;
      }
    }
    selected_offsets[global_offset + 1] = static_cast<size_t>(global_index);
  }
}

template <int MaxLength, int MaxThreadsPerSeq, int MaxSeqs>
__device__ void BeamSearchDetails(int64_t* selected_ids,
                                  float* selected_scores,
                                  int* parent_idx,
                                  size_t* selected_offsets,
                                  const int64_t* pre_ids,
                                  const float* pre_scores,
                                  const int64_t* ids,
                                  const float* scores,
                                  const int seq_offset_start,
                                  const int seq_offset_end,
                                  const int seq_width,
                                  int beam_size,
                                  int end_id,
                                  bool is_accumulated,
                                  int num_used_threads) {
  __shared__ Triple top_beam[MaxLength];

  int num_items = 0;
  if (is_accumulated) {
    num_items = SelectTopBeam<MaxThreadsPerSeq, true>(top_beam,
                                                      pre_ids,
                                                      pre_scores,
                                                      ids,
                                                      scores,
                                                      seq_offset_start,
                                                      seq_offset_end,
                                                      seq_width,
                                                      beam_size,
                                                      end_id,
                                                      num_used_threads);
  } else {
    num_items = SelectTopBeam<MaxThreadsPerSeq, false>(top_beam,
                                                       pre_ids,
                                                       pre_scores,
                                                       ids,
                                                       scores,
                                                       seq_offset_start,
                                                       seq_offset_end,
                                                       seq_width,
                                                       beam_size,
                                                       end_id,
                                                       num_used_threads);
  }

  const int tid = threadIdx.x;  // use 1 thread only for each sequence
  const int tid_of_seq = tid % MaxThreadsPerSeq;
  if (tid_of_seq == 0) {
    // Use 1 thread for each sequence.
    Triple* top_beam_local = top_beam + tid * beam_size;
    bool finish_flag =
        PruneEndBeams(top_beam_local, pre_ids, end_id, num_items);

    int selected_seq_start = 0;
    int selected_seq_length = finish_flag ? 0 : num_items;

    if (MaxSeqs > 1) {
      const int seq_id = (MaxSeqs > 1) ? tid / MaxThreadsPerSeq : tid;
      __shared__ int shared_mem[MaxSeqs];

      // [0, MaxSeqs - 1], length of each sequences
      shared_mem[seq_id] = selected_seq_length;
      __syncthreads();

      for (int s = 0; s < seq_id; ++s) {
        selected_seq_start += shared_mem[s];
      }

      if (seq_id == 0) {
        selected_offsets[0] = 0;
      }
    } else {
      selected_offsets[0] = 0;
    }

    if (parent_idx) {
      WriteBack<true>(selected_ids,
                      selected_scores,
                      parent_idx,
                      selected_offsets,
                      top_beam_local,
                      seq_offset_start,
                      seq_offset_end,
                      selected_seq_start,
                      selected_seq_length);
    } else {
      WriteBack<false>(selected_ids,
                       selected_scores,
                       parent_idx,
                       selected_offsets,
                       top_beam_local,
                       seq_offset_start,
                       seq_offset_end,
                       selected_seq_start,
                       selected_seq_length);
    }
  }
}

template <int MaxLength, int MaxThreadsPerSeq, int MaxSeqs>
__global__ void BeamSearchKernel(int64_t* selected_ids,
                                 float* selected_scores,
                                 int* parent_idx,
                                 size_t* selected_offsets,
                                 const int64_t* pre_ids,
                                 const float* pre_scores,
                                 const int64_t* ids,
                                 const float* scores,
                                 const size_t* seq_offsets,
                                 const int num_seqs,
                                 const int seq_width,
                                 int beam_size,
                                 int end_id,
                                 bool is_accumulated,
                                 int num_used_threads) {
  const int tid = threadIdx.x;
  const int seq_id = (MaxSeqs > 1) ? tid / MaxThreadsPerSeq : tid;

  int seq_offset_start = static_cast<int>(seq_offsets[seq_id]);
  int seq_offset_end = static_cast<int>(seq_offsets[seq_id + 1]);

  BeamSearchDetails<MaxLength, MaxThreadsPerSeq, MaxSeqs>(selected_ids,
                                                          selected_scores,
                                                          parent_idx,
                                                          selected_offsets,
                                                          pre_ids,
                                                          pre_scores,
                                                          ids,
                                                          scores,
                                                          seq_offset_start,
                                                          seq_offset_end,
                                                          seq_width,
                                                          beam_size,
                                                          end_id,
                                                          is_accumulated,
                                                          num_used_threads);
}

template <int MaxLength, int MaxThreadsPerSeq>
__global__ void BeamSearchKernelSingle(int64_t* selected_ids,
                                       float* selected_scores,
                                       int* parent_idx,
                                       size_t* selected_offsets,
                                       const int64_t* pre_ids,
                                       const float* pre_scores,
                                       const int64_t* ids,
                                       const float* scores,
                                       const int seq_length,
                                       const int seq_width,
                                       int beam_size,
                                       int end_id,
                                       bool is_accumulated,
                                       int num_used_threads) {
  const int seq_offset_start = 0;
  const int seq_offset_end = seq_length;

  BeamSearchDetails<MaxLength, MaxThreadsPerSeq, 1>(selected_ids,
                                                    selected_scores,
                                                    parent_idx,
                                                    selected_offsets,
                                                    pre_ids,
                                                    pre_scores,
                                                    ids,
                                                    scores,
                                                    seq_offset_start,
                                                    seq_offset_end,
                                                    seq_width,
                                                    beam_size,
                                                    end_id,
                                                    is_accumulated,
                                                    num_used_threads);
}

static inline int GetNumUsedThreads(const int max_threads_per_seq,
                                    const int seq_width,
                                    int beam_size) {
  int num_used_threads = (seq_width + beam_size - 1) / beam_size;
  num_used_threads = max_threads_per_seq < num_used_threads
                         ? max_threads_per_seq
                         : num_used_threads;

  num_used_threads =
      num_used_threads > 32
          ? (num_used_threads >> 5) << 5
          : (num_used_threads > 16
                 ? 32
                 : (num_used_threads > 8
                        ? 16
                        : (num_used_threads > 4
                               ? 8
                               : (num_used_threads > 2 ? 4
                                                       : num_used_threads))));
  return num_used_threads;
}

template <typename T>
class BeamSearchFunctor<GPUContext, T> {
 public:
  void operator()(const GPUContext& dev_ctx,
                  const DenseTensor* pre_ids,
                  const DenseTensor* pre_scores,
                  const DenseTensor* ids,
                  const DenseTensor* scores,
                  DenseTensor* selected_ids,
                  DenseTensor* selected_scores,
                  DenseTensor* parent_idx,
                  size_t level,
                  size_t beam_size,
                  int end_id,
                  bool is_accumulated) {
    auto abs_lod = ToAbsOffset(scores->lod());

    const int64_t* pre_ids_data = pre_ids->data<int64_t>();
    const float* pre_scores_data = pre_scores->data<float>();
    const int64_t* ids_data = ids ? ids->data<int64_t>() : nullptr;
    const float* scores_data = scores->data<float>();

    const size_t num_seqs = abs_lod[level].size() - 1;
    size_t seq_width = 1;
    for (int i = 1; i < scores->dims().size(); i++) {
      seq_width *= scores->dims()[i];
    }

    // Reserve a big enough memory.
    auto selected_dims =
        make_ddim({static_cast<int64_t>(num_seqs * beam_size), 1});
    selected_ids->Resize(selected_dims);
    int64_t* selected_ids_data = dev_ctx.template Alloc<int64_t>(selected_ids);
    selected_scores->Resize(selected_dims);
    float* selected_scores_data =
        dev_ctx.template Alloc<float>(selected_scores);
    if (parent_idx != nullptr) {
      parent_idx->Resize({static_cast<int64_t>(num_seqs * beam_size)});
    }
    int* parent_idx_data =
        parent_idx ? dev_ctx.template Alloc<int>(parent_idx) : nullptr;

    LegacyLoD selected_lod(2);
    selected_lod[0].assign(abs_lod[level].begin(), abs_lod[level].end());
    selected_lod[1].resize(scores->dims()[0] + 1);
    phi::MixVector<size_t> mix_vector(&selected_lod[1]);
    phi::MixVector<size_t> mixv_abs(&abs_lod[level]);
    size_t* selected_offsets = mix_vector.CUDAMutableData(dev_ctx.GetPlace());

    if (num_seqs == 1) {
      const int seq_length = static_cast<int>(abs_lod[level][1]);
      const int kMaxThreadsPerSeq = 1024;
      int num_used_threads = GetNumUsedThreads(kMaxThreadsPerSeq,
                                               static_cast<int>(seq_width),
                                               static_cast<int>(beam_size));
      switch (phi::backends::gpu::RoundToPowerOfTwo(beam_size * seq_width)) {
        CUDA_LAUNCH_KERNEL_HELPER(
            BeamSearchKernelSingle<kPowerOfTwoDim, kMaxThreadsPerSeq>
            <<<1, kMaxThreadsPerSeq, 0, dev_ctx.stream()>>>(
                selected_ids_data,
                selected_scores_data,
                parent_idx_data,
                selected_offsets,
                pre_ids_data,
                pre_scores_data,
                ids_data,
                scores_data,
                seq_length,
                static_cast<int>(seq_width),
                static_cast<int>(beam_size),
                static_cast<int>(end_id),
                is_accumulated,
                num_used_threads));
      }
    } else if (num_seqs <= 4) {
      const size_t* seq_offsets = mixv_abs.CUDAData(dev_ctx.GetPlace());
      // Use only 1 block
      const int kMaxThreadsPerSeq = 32;
      const int kMaxSeqs = 4;
      int num_used_threads = GetNumUsedThreads(kMaxThreadsPerSeq,
                                               static_cast<int>(seq_width),
                                               static_cast<int>(beam_size));
      switch (
          phi::backends::gpu::RoundToPowerOfTwo(beam_size * num_seqs * 32)) {
        CUDA_LAUNCH_KERNEL_HELPER(
            BeamSearchKernel<kPowerOfTwoDim, kMaxThreadsPerSeq, kMaxSeqs>
            <<<1, num_seqs * kMaxThreadsPerSeq, 0, dev_ctx.stream()>>>(
                selected_ids_data,
                selected_scores_data,
                parent_idx_data,
                selected_offsets,
                pre_ids_data,
                pre_scores_data,
                ids_data,
                scores_data,
                seq_offsets,
                static_cast<int>(num_seqs),
                static_cast<int>(seq_width),
                static_cast<int>(beam_size),
                end_id,
                is_accumulated,
                num_used_threads));
      }
    } else {
      PADDLE_THROW(common::errors::Unimplemented(
          "Not implemented other number of sequences yet."));
    }

    dev_ctx.Wait();
    mix_vector.CopyToCPU();
    if (!CheckLegacyLoD(selected_lod)) {
      PADDLE_THROW(common::errors::InvalidArgument(
          "lod %s is not right in"
          " beam_search, please check your code.",
          LoDToString(selected_lod)));
    }

    selected_ids->set_lod(selected_lod);
    selected_scores->set_lod(selected_lod);
    if (selected_lod[1].back() < num_seqs * beam_size) {
      auto final_selected_dims =
          make_ddim({static_cast<int64_t>(selected_lod[1].back()), 1});
      selected_ids->Resize(final_selected_dims);
      selected_scores->Resize(final_selected_dims);
      if (parent_idx) {
        parent_idx->Resize({static_cast<int64_t>(selected_lod[1].back())});
      }
    }
  }
};

template class BeamSearchFunctor<GPUContext, int>;
template class BeamSearchFunctor<GPUContext, int64_t>;
template class PADDLE_API BeamSearchFunctor<GPUContext, float>;
template class BeamSearchFunctor<GPUContext, double>;

}  // namespace math
}  // namespace phi