// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "generators.h"
#include "interface.h"
#include "search.h"
#include "search_cuda.h"
#include "cuda_common.h"
#include "beam_search_scorer_cuda.cuh"
#include "beam_search_scorer_cuda.h"
#include "beam_search_topk.h"
#include <queue>
#include <random>

namespace Generators {

Search_Cuda::Search_Cuda(const GeneratorParams& params)
    : Search{params} {
  auto batch_beam_size = params.BatchBeamSize();
  sequence_lengths_ = params.p_device->Allocate<int32_t>(batch_beam_size);

  eos_seen_buffer_ = CudaMallocArray<bool>(batch_beam_size, &eos_seen_);
  done_cpu_ = CudaMallocHostArray<bool>(1);

  eos_token_ids_ = params.p_device->Allocate<int32_t>(params.config.model.eos_token_id.size());
  copy(std::span<const int32_t>{params.config.model.eos_token_id}, eos_token_ids_.CpuSpan());
  eos_token_ids_.CopyCpuToDevice();

  ResetDone();
}

GreedySearch_Cuda::GreedySearch_Cuda(const GeneratorParams& params)
    : Search_Cuda{params} {
  next_tokens_buffer_ = params.p_device->Allocate<int32_t>(params.search.batch_size);
  next_tokens_buffer_.Zero();
  next_tokens_ = gpu_span<int32_t>(next_tokens_buffer_.Span());
}

void GreedySearch_Cuda::EnsureSamplingData() {
  if (samplingdata_)
    return;

  unsigned long long random_seed = (params_->search.random_seed != -1)
                                       ? params_->search.random_seed
                                       : std::random_device{}();

  // Allocate a single buffer for all sampling data
  size_t sampling_buffer_size = cuda::SamplingData::CalculateTotalSize(params_->search.batch_size, params_->config.model.vocab_size, GetStream());
  sampling_buffer_ = params_->p_device->Allocate<uint8_t>(sampling_buffer_size);

  // Create SamplingData with the externally allocated buffer
  samplingdata_ = std::make_unique<cuda::SamplingData>(random_seed, params_->search.batch_size, params_->config.model.vocab_size, GetStream(),
                                                       sampling_buffer_.Span().data(), sampling_buffer_size);
}

BeamSearch_Cuda::BeamSearch_Cuda(const GeneratorParams& params)
    : Search_Cuda{params} {
  assert(params_->search.num_beams > 1);  // If 1, use GreedySearch
  auto batch_beam_size = params_->BatchBeamSize();
  beam_scorer_ = std::make_unique<BeamSearchScorer_Cuda>(*params_, eos_token_ids_.Span());

  topk_next_tokens_ = CudaMallocArray<int32_t>(2 * batch_beam_size);
  topk_next_indices_ = CudaMallocArray<int32_t>(2 * batch_beam_size);
  topk_next_scores_ = CudaMallocArray<float>(2 * batch_beam_size);
  softmax_buffer_ = CudaMallocArray<float>(batch_beam_size * params_->config.model.vocab_size);

  constexpr size_t max_parts_of_vocab = 128;
  size_t topk_buffer_size = batch_beam_size * (max_parts_of_vocab + 1) * params_->search.num_beams * 2 * 2;
  topk_buffer_ = CudaMallocArray<float>(topk_buffer_size);
  static_assert(sizeof(float) == sizeof(int32_t));  // The topk_buffer assumes these match, fix for float16

  CUDA_CHECK(cudaMemsetAsync(topk_buffer_.get(), 0, topk_buffer_size * sizeof(float), GetStream()));
}

BeamSearch_Cuda::~BeamSearch_Cuda() = default;

void Search_Cuda::ResetDone() {
  *done_cpu_ = false;
  done_pending_ = false;
  CUDA_CHECK(cudaMemsetAsync(eos_seen_.data(), 0, eos_seen_.size_bytes(), GetStream()));
}

DeviceSpan<float> Search_Cuda::GetLogits() const {
  return next_token_scores_;
}

void Search_Cuda::SetLogits(DeviceSpan<float> logits) {
  next_token_scores_ = logits;
}

void Search_Cuda::SaveStateForTransactionImpl(bool checkpoint_local_state) {
  if (!transaction_eos_seen_)
    transaction_eos_seen_ = CudaMallocArray<bool>(eos_seen_.size());

  if (checkpoint_local_state) {
    if (!transaction_sequence_lengths_)
      transaction_sequence_lengths_ = CudaMallocArray<int32_t>(sequence_lengths_.size());
    CUDA_CHECK(cudaMemcpyAsync(transaction_sequence_lengths_.get(), sequence_lengths_.Span().data(),
                               sequence_lengths_.size() * sizeof(int32_t), cudaMemcpyDeviceToDevice, GetStream()));
  }
  transaction_saved_sequence_lengths_ = checkpoint_local_state;
  CUDA_CHECK(cudaMemcpyAsync(transaction_eos_seen_.get(), eos_seen_.data(),
                             eos_seen_.size_bytes(), cudaMemcpyDeviceToDevice, GetStream()));
  transaction_done_ = *done_cpu_;
}

void Search_Cuda::RestoreStateForTransactionImpl() {
  if (transaction_saved_sequence_lengths_) {
    CUDA_CHECK(cudaMemcpyAsync(sequence_lengths_.Span().data(), transaction_sequence_lengths_.get(),
                               sequence_lengths_.size() * sizeof(int32_t), cudaMemcpyDeviceToDevice, GetStream()));
  }
  CUDA_CHECK(cudaMemcpyAsync(eos_seen_.data(), transaction_eos_seen_.get(),
                             eos_seen_.size_bytes(), cudaMemcpyDeviceToDevice, GetStream()));
  transaction_saved_sequence_lengths_ = false;
}

void Search_Cuda::SynchronizeStateForTransactionImpl() {
  CUDA_CHECK(cudaStreamSynchronize(GetStream()));
}

void Search_Cuda::CompleteStateRestoreForTransactionImpl() {
  *done_cpu_ = transaction_done_;
  done_pending_ = false;
}

void GreedySearch_Cuda::SaveStateForTransactionImpl(bool checkpoint_local_state) {
  Search_Cuda::SaveStateForTransactionImpl(checkpoint_local_state);
  if (checkpoint_local_state) {
    EnsureSamplingData();
    if (!transaction_next_tokens_)
      transaction_next_tokens_ = CudaMallocArray<int32_t>(next_tokens_.size());
    if (!transaction_curand_states_)
      transaction_curand_states_ = CudaMallocArray<curandState>(params_->search.batch_size);

    CUDA_CHECK(cudaMemcpyAsync(transaction_next_tokens_.get(), next_tokens_.data(),
                               next_tokens_.size_bytes(), cudaMemcpyDeviceToDevice, GetStream()));
    CUDA_CHECK(cudaMemcpyAsync(transaction_curand_states_.get(), samplingdata_->curand_states,
                               params_->search.batch_size * sizeof(curandState), cudaMemcpyDeviceToDevice, GetStream()));
  }
  transaction_saved_sampling_state_ = checkpoint_local_state;
}

void GreedySearch_Cuda::RestoreStateForTransactionImpl() {
  Search_Cuda::RestoreStateForTransactionImpl();
  if (transaction_saved_sampling_state_) {
    CUDA_CHECK(cudaMemcpyAsync(next_tokens_.data(), transaction_next_tokens_.get(),
                               next_tokens_.size_bytes(), cudaMemcpyDeviceToDevice, GetStream()));
    CUDA_CHECK(cudaMemcpyAsync(samplingdata_->curand_states, transaction_curand_states_.get(),
                               params_->search.batch_size * sizeof(curandState), cudaMemcpyDeviceToDevice, GetStream()));
  }
  transaction_saved_sampling_state_ = false;
}

void GreedySearch_Cuda::CompleteStateRestoreForTransactionImpl() {
  Search_Cuda::CompleteStateRestoreForTransactionImpl();
  completion_pending_ = false;
  external_host_copy_ = false;
}

DeviceSpan<int32_t> GreedySearch_Cuda::GetNextTokens() {
  return next_tokens_buffer_;
}

DeviceSpan<int32_t> BeamSearch_Cuda::GetNextTokens() {
  return beam_scorer_->GetNextTokens();
}

DeviceSpan<int32_t> BeamSearch_Cuda::GetNextIndices() {
  return beam_scorer_->GetNextIndices();
}

void BeamSearch_Cuda::SelectTop() {
  cuda::DispatchBlockwiseSoftmaxForward<true>(GetStream(), softmax_buffer_.get(), next_token_scores_.Span().data(), params_->config.model.vocab_size,
                                              params_->config.model.vocab_size, params_->config.model.vocab_size, params_->BatchBeamSize());

  // Copy next_token_scores to CPU
  auto next_token_scores_cpu = CudaMallocHostArray<float>(params_->BatchBeamSize() * params_->config.model.vocab_size);
  CUDA_CHECK(cudaMemcpyAsync(next_token_scores_cpu.get(), softmax_buffer_.get(), params_->BatchBeamSize() * params_->config.model.vocab_size * sizeof(float), cudaMemcpyDeviceToHost, GetStream()));
  CUDA_CHECK(cudaStreamSynchronize(GetStream()));

  auto beam_scores = beam_scorer_->GetNextScores();

  // Add beam score to next token scores. Corresponding python code is like:
  //    next_token_scores = next_token_scores + beam_scores[:, None].expand_as(next_token_scores)
  cuda::LaunchAddProbsKernel(softmax_buffer_.get(), beam_scores.Span().data(),
                             params_->search.batch_size, params_->search.num_beams, params_->config.model.vocab_size, GetStream());

  if (params_->search.num_beams <= 32) {
    constexpr size_t max_parts_of_vocab = 128;
    size_t candidate_count = params_->BatchBeamSize() * 2 * params_->search.num_beams;
    float* topk_tmp_buffer = topk_buffer_.get();
    float* topk_scores_1st_stage = topk_tmp_buffer;
    int32_t* topk_tokens_1st_stage = reinterpret_cast<int32_t*>(topk_scores_1st_stage + candidate_count * max_parts_of_vocab);
    float* topk_scores_2nd_stage = reinterpret_cast<float*>(topk_tokens_1st_stage + candidate_count * max_parts_of_vocab);
    int32_t* topk_tokens_2nd_stage = reinterpret_cast<int32_t*>(topk_scores_2nd_stage + candidate_count);

    cuda::BeamSearchTopK(softmax_buffer_.get(),
                         params_->search.batch_size,
                         params_->search.num_beams,
                         params_->config.model.vocab_size,
                         2 * params_->search.num_beams,
                         topk_scores_1st_stage,
                         topk_tokens_1st_stage,
                         topk_scores_2nd_stage,
                         topk_tokens_2nd_stage,
                         topk_next_scores_.get(),
                         topk_next_tokens_.get(),
                         topk_next_indices_.get(),
                         GetStream());
  } else
    assert(false);

  CUDA_CHECK(cudaStreamSynchronize(GetStream()));

  size_t size = params_->BatchBeamSize() * 2;
  std::span<float> next_scores{topk_next_scores_.get(), size};
  std::span<int32_t> next_tokens{topk_next_tokens_.get(), size};
  std::span<int32_t> next_indices{topk_next_indices_.get(), size};

#if 0  // TODO(ryanhill): Use logging option
  DumpCudaSpan(std::cout, next_scores);
  DumpCudaSpan(std::cout, next_tokens);
  DumpCudaSpan(std::cout, next_indices);
#endif

  beam_scorer_->Process(sequences_, next_scores, next_tokens, next_indices);
  auto next_tokens_device = beam_scorer_->GetNextTokens();
  next_tokens_ = gpu_span<int32_t>(next_tokens_device.Span());
  sequences_.AfterAppendNextTokens(next_tokens_device, params_->BatchBeamSize());
}

void GreedySearch_Cuda::SampleTopKTopP(int k, float p, float temperature) {
  EnsureSamplingData();
  std::span<float> scores = next_token_scores_.Span();
  assert(scores.size() == params_->search.batch_size * params_->config.model.vocab_size);
  cuda::GetSample(samplingdata_.get(), GetStream(), next_tokens_.data(), scores.data(), int(scores.size() / params_->search.batch_size),
                  params_->search.batch_size, k, p, temperature);

  external_host_copy_ = false;
  LaunchNextTokensTail();

  if (!defer_completion_)
    CompleteGeneration();
}

bool GreedySearch_Cuda::BindNextTokensSlot(DeviceSpan<int32_t> slot) {
  if (params_->BatchBeamSize() != 1 || slot.size() != 1)
    return false;

  // Both spans stay alive through the shared_ptr inside DeviceSpan, so a caller that grows its
  // buffer between steps cannot leave this search writing into freed memory.
  next_tokens_buffer_ = slot;
  next_tokens_ = gpu_span<int32_t>(next_tokens_buffer_.Span());
  return true;
}

void GreedySearch_Cuda::OnNextTokensSampled() {
  external_host_copy_ = true;
  LaunchNextTokensTail();
}

// Everything token selection does after the sampler has written next_tokens_, for both the
// per-search and the batched entry points. Nothing here reads a device result, so the launches
// for a whole batch of searches can be queued before anyone waits.
void GreedySearch_Cuda::LaunchNextTokensTail() {
  // Check for EOS
  assert(next_tokens_.size() == eos_seen_.size());
  cuda::Launch_CheckForEOSAndPad(next_tokens_.data(), static_cast<int>(next_tokens_.size()), eos_seen_.data(), eos_token_ids_.Span().data(), static_cast<int>(eos_token_ids_.Span().size()), params_->config.model.pad_token_id, done_cpu_.get(), GetStream());

  // Append tokens. The launch is unconditional so that no device result has to be read back here:
  // CheckForEOSAndPad has already replaced a finished sequence's token with the pad token, and the
  // host-side length below is only advanced for sequences that had not finished, so a slot written
  // for a finished sequence is never read.
  if (sequences_.GetSequenceLength() < sequences_.max_length_) {
    cuda::Launch_AppendNextTokensToSequences(next_tokens_buffer_.Span(), sequences_.GetSequences().Span(), params_->BatchBeamSize(), sequences_.GetSequenceLength(), sequences_.max_length_, GetStream());
  }

  completion_pending_ = true;
  done_pending_ = true;
}

void GreedySearch_Cuda::CompleteGeneration() {
  if (!completion_pending_)
    return;

  // done_cpu_ lives in pinned host memory written by CheckForEOSAndPad, so it is only meaningful
  // once the stream has drained. When completion is deferred, every search in the batch has already
  // launched its work by the time the first of these calls runs, so only that first call waits.
  // Reading the next tokens back here rides along on the same synchronization, which lets callers
  // use GetNextTokens().CpuSpan() afterwards without another round trip. When the tokens were
  // sampled into a shared buffer the caller has already done that copy for the whole batch.
  if (!external_host_copy_)
    next_tokens_buffer_.CopyDeviceToCpu();
  completion_pending_ = false;
  done_pending_ = false;

  if (!*done_cpu_) {
    sequences_.AfterAppendNextTokens(next_tokens_buffer_, params_->BatchBeamSize());
  }

  MarkDoneAtMaxLength();
}

void GreedySearch_Cuda::AppendTokensToSequences(DeviceSpan<int32_t>& tokens) {
  cuda::Launch_AppendNextTokensToSequences(tokens.Span(), sequences_.GetSequences().Span(),
                                           params_->BatchBeamSize(), sequences_.GetSequenceLength(),
                                           sequences_.max_length_, GetStream());
  sequences_.AfterAppendNextTokens(tokens, params_->BatchBeamSize());
}

void GreedySearch_Cuda::MarkDoneAtMaxLength() {
  if (sequences_.GetSequenceLength() >= params_->search.max_length) {
    if (GetLogItems().enabled && GetLogItems().hit_max_length)
      Log("hit_max_length", "greedy cuda hit");
    *done_cpu_ = true;
  }
}

void GreedySearch_Cuda::CommitToken(int32_t token) {
  // The caller already selected this generated token, so skip sampling but retain generated-token
  // EOS, padding, and max-length behavior. Speculative decoding only calls this with batch_size 1.
  CUDA_CHECK(cudaMemcpyAsync(next_tokens_.data(), &token, sizeof(int32_t), cudaMemcpyHostToDevice, GetStream()));
  external_host_copy_ = false;
  LaunchNextTokensTail();
  CompleteGeneration();
}

bool BeamSearch_Cuda::IsDone() const {
  if (beam_scorer_->IsDoneLater())
    return true;

  if (sequences_.GetSequenceLength() == params_->search.max_length) {
    if (GetLogItems().enabled && GetLogItems().hit_max_length)
      Log("hit_max_length", "beam cuda hit");
    return true;
  }
  return false;
}

void BeamSearch_Cuda::Finalize(size_t num_return_sequences) {
  if (finalized_)
    return;
  beam_scorer_->Finalize(sequences_, num_return_sequences);
  finalized_ = true;
}

DeviceSpan<int32_t> BeamSearch_Cuda::GetSequence(size_t index) {
  Finalize(params_->search.num_return_sequences);
  const size_t batch_id = index / params_->search.num_return_sequences;
  const size_t beam_id = index % params_->search.num_return_sequences;
  return beam_scorer_->GetBeamHypothesis(batch_id, beam_id);
}

DeviceSpan<int32_t> BeamSearch_Cuda::GetSequence(size_t batch_id, size_t beam_id) {
  Finalize(params_->search.num_return_sequences);
  return beam_scorer_->GetBeamHypothesis(batch_id, beam_id);
}

std::span<float> Search_Cuda::GetScores(int batch_beam_index) {
  assert(batch_beam_index >= 0 && batch_beam_index < params_->BatchBeamSize());
  return next_token_scores_.subspan(batch_beam_index * params_->config.model.vocab_size, params_->config.model.vocab_size).Span();
}

std::span<float> Search_Cuda::GetScores() {
  return next_token_scores_.Span();
}

// Set user input tokens (batch_beam_size, sequence_length)
void GreedySearch_Cuda::AppendTokens(DeviceSpan<int32_t>& next_tokens) {
  // Caller-provided prompt/continuation tokens are appended as input. Unlike CommitToken, they do
  // not run generated-token EOS/padding handling, and done state is reset when capacity remains.
  ResetDone();
  AppendTokensToSequences(next_tokens);
  MarkDoneAtMaxLength();
  if (*done_cpu_)
    return;

  ResetDone();
}

void BeamSearch_Cuda::AppendTokens(DeviceSpan<int32_t>& next_tokens) {
  auto next_tokens_gpu = next_tokens.Span();
  cuda::Launch_ExpandInputSequences(next_tokens_gpu, sequences_.GetNextSequences().Span(), params_->search.batch_size, params_->search.num_beams, sequences_.max_length_, GetStream());
  cuda::Launch_ExpandInputSequences(next_tokens_gpu, sequences_.GetSequences().Span(), params_->search.batch_size, params_->search.num_beams, sequences_.max_length_, GetStream());
  sequences_.AfterAppendNextTokens(next_tokens, params_->search.batch_size);  // next_tokens is batch_size
  CUDA_CHECK(cudaStreamSynchronize(GetStream()));
}

void GreedySearch_Cuda::RewindTo(size_t index) {
  ResetDone();
  if (index > 0)
    cuda::Launch_GetLastTokens(next_tokens_.data(), sequences_.GetSequences().Span().data(), static_cast<int>(params_->BatchBeamSize()), static_cast<int>(index), sequences_.max_length_, GetStream());
  else
    CUDA_CHECK(cudaMemsetAsync(next_tokens_.data(), 0, params_->search.batch_size * sizeof(int32_t), GetStream()));
  sequences_.RewindTo(index);
}

void Search_Cuda::ApplyMinLength(int min_length) {
  if (sequences_.GetSequenceLength() >= min_length)
    return;

  for (auto eos_token_id : params_->config.model.eos_token_id)
    cuda::LaunchSetScoreProcessor(GetScores().data(), params_->BatchBeamSize(), params_->config.model.vocab_size, eos_token_id, std::numeric_limits<float>::lowest(), GetStream());
}

void Search_Cuda::ApplyRepetitionPenalty(float penalty) {
  if (penalty == 1.0f)
    return;

  cuda::LaunchRepetitionPenaltyProcessor(sequences_.GetSequences().Span().data(),
                                         GetScores().data(), params_->search.batch_size, params_->search.num_beams, params_->config.model.vocab_size,
                                         params_->search.max_length, GetSequenceLength(), penalty, GetStream());
}

}  // namespace Generators
