// Copyright (c) 2019 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 "ctc_helper.h"
#include "gpu_ctc_kernels.h"
#include "reduce.h"

template <typename ProbT>
class GpuCTC {
 public:
  GpuCTC(int alphabet_size,
         int minibatch,
         void* workspace,
         CUstream stream,
         int blank_label)
      : out_dim_(alphabet_size),
        minibatch_(minibatch),
        gpu_workspace_(workspace),
        stream_(stream),
        blank_label_(blank_label){};

  // Noncopyable
  GpuCTC(const GpuCTC&) = delete;
  GpuCTC& operator=(const GpuCTC&) = delete;

  ctcStatus_t cost_and_grad(const ProbT* const activations,
                            ProbT* grads,
                            ProbT* costs,
                            const int* const flat_labels,
                            const int* const label_lengths,
                            const int* const input_lengths);

  ctcStatus_t score_forward(const ProbT* const activations,
                            ProbT* costs,
                            const int* const flat_labels,
                            const int* const label_lengths,
                            const int* const input_lengths);

 private:
  template <int NT, int VT>
  ctcStatus_t launch_alpha_beta_kernels(const ProbT* const probs,
                                        ProbT* grads,
                                        bool compute_alpha,
                                        bool compute_beta);

  ctcStatus_t launch_gpu_kernels(const ProbT* const probs,
                                 ProbT* grads,
                                 size_t config,
                                 bool launch_alpha,
                                 bool launch_beta);

  ctcStatus_t setup_gpu_metadata(const int* const flat_labels,
                                 const int* const label_lengths,
                                 const int* const input_lengths);

  ctcStatus_t create_metadata_and_choose_config(const int* const label_lengths,
                                                const int* const flat_labels,
                                                const int* const input_lengths,
                                                size_t& best_config);

  ctcStatus_t compute_probs(const ProbT* const activations);

  ctcStatus_t compute_cost_and_score(const ProbT* const activations,
                                     ProbT* grads,
                                     ProbT* costs,
                                     const int* const flat_labels,
                                     const int* const label_lengths,
                                     const int* const input_lengths,
                                     bool compute_alpha,
                                     bool compute_betas_and_grad);

  int out_dim_;  // Number of characters plus blank
  int minibatch_;

  int S_;
  int T_;

  int activation_cols_;  // Number of columns in activations

  CUstream stream_;
  int blank_label_;

  void* gpu_workspace_;  // Buffer for all temporary GPU memory
  int* utt_length_;      // T
  int* label_sizes_;     // L
  int* repeats_;         // repeats_
  int* label_offsets_;
  int* labels_without_blanks_;
  int* labels_with_blanks_;
  ProbT* alphas_;
  ProbT* nll_forward_;
  ProbT* nll_backward_;
  ProbT* denoms_;  // Temporary storage for denoms for softmax
  ProbT* probs_;   // Temporary storage for probabilities (softmax output)
};

template <typename ProbT>
ctcStatus_t GpuCTC<ProbT>::setup_gpu_metadata(const int* const flat_labels,
                                              const int* const label_lengths,
                                              const int* const input_lengths) {
  size_t gpu_bytes_used = 0;

  nll_forward_ = reinterpret_cast<ProbT*>(static_cast<char*>(gpu_workspace_) +
                                          gpu_bytes_used);
  gpu_bytes_used += minibatch_ * sizeof(ProbT);

  nll_backward_ = reinterpret_cast<ProbT*>(static_cast<char*>(gpu_workspace_) +
                                           gpu_bytes_used);
  gpu_bytes_used += minibatch_ * sizeof(ProbT);

  repeats_ = reinterpret_cast<int*>(static_cast<char*>(gpu_workspace_) +
                                    gpu_bytes_used);
  gpu_bytes_used += minibatch_ * sizeof(int);

  label_offsets_ = reinterpret_cast<int*>(static_cast<char*>(gpu_workspace_) +
                                          gpu_bytes_used);
  gpu_bytes_used += minibatch_ * sizeof(int);

  // This is the max of all S and T for all valid examples in the minibatch.
  // A valid example is one for which L + repeats <= T
  S_ = 0;
  T_ = 0;

  // This is the max of all timesteps, valid or not. Needed to compute offsets
  int Tmax = 0;

  // This is the max of all labels, valid or not. Needed to compute offsets
  int Lmax = 0;
  int total_label_length = 0;

  constexpr int cpu_buffer_size = 64;
  int repeats[cpu_buffer_size];
  int label_offsets[cpu_buffer_size];

  const int num_passes = ctc_helper::div_up(minibatch_, cpu_buffer_size);

  cudaError_t cuda_status;

  for (int pass = 0; pass < num_passes; ++pass) {
    const int start_idx = pass * cpu_buffer_size;
    const int end_idx = std::min(minibatch_, (pass + 1) * cpu_buffer_size);

    for (int j = start_idx; j < end_idx; ++j) {
      const int L = label_lengths[j];
      const int local_T = input_lengths[j];
      const int* label_ptr = &(flat_labels[total_label_length]);

      label_offsets[j % cpu_buffer_size] = total_label_length;
      total_label_length += L;

      int repeat_counter = 0;

      for (int i = 1; i < L; ++i)
        repeat_counter += (label_ptr[i] == label_ptr[i - 1]);

      repeats[j % cpu_buffer_size] = repeat_counter;
      const bool valid_label = ((L + repeat_counter) <= local_T);

      // Only update S and T if label is valid
      S_ = (valid_label) ? std::max(S_, L) : S_;
      T_ = (valid_label) ? std::max(T_, local_T) : T_;

      Tmax = std::max(Tmax, local_T);
      Lmax = std::max(Lmax, L);
    }

    cuda_status = cudaMemcpyAsync(&(repeats_[start_idx]),
                                  repeats,
                                  (end_idx - start_idx) * sizeof(int),
                                  cudaMemcpyHostToDevice,
                                  stream_);
    if (cuda_status != cudaSuccess) return CTC_STATUS_MEMOPS_FAILED;

    cuda_status = cudaMemcpyAsync(&(label_offsets_[start_idx]),
                                  label_offsets,
                                  (end_idx - start_idx) * sizeof(int),
                                  cudaMemcpyHostToDevice,
                                  stream_);
    if (cuda_status != cudaSuccess) return CTC_STATUS_MEMOPS_FAILED;
  }

  S_ = 2 * S_ + 1;
  const int Smax = 2 * Lmax + 1;

  activation_cols_ = minibatch_ * Tmax;

  // Allocate memory for T
  utt_length_ = reinterpret_cast<int*>(static_cast<char*>(gpu_workspace_) +
                                       gpu_bytes_used);
  gpu_bytes_used += minibatch_ * sizeof(int);

  cuda_status = cudaMemcpyAsync(utt_length_,
                                input_lengths,
                                minibatch_ * sizeof(int),
                                cudaMemcpyHostToDevice,
                                stream_);
  if (cuda_status != cudaSuccess) return CTC_STATUS_MEMOPS_FAILED;

  label_sizes_ = reinterpret_cast<int*>(static_cast<char*>(gpu_workspace_) +
                                        gpu_bytes_used);
  gpu_bytes_used += minibatch_ * sizeof(int);
  cuda_status = cudaMemcpyAsync(label_sizes_,
                                label_lengths,
                                minibatch_ * sizeof(int),
                                cudaMemcpyHostToDevice,
                                stream_);
  if (cuda_status != cudaSuccess) return CTC_STATUS_MEMOPS_FAILED;

  labels_without_blanks_ = reinterpret_cast<int*>(
      static_cast<char*>(gpu_workspace_) + gpu_bytes_used);
  gpu_bytes_used += Lmax * minibatch_ * sizeof(int);
  cuda_status = cudaMemcpyAsync(labels_without_blanks_,
                                flat_labels,
                                total_label_length * sizeof(int),
                                cudaMemcpyHostToDevice,
                                stream_);
  if (cuda_status != cudaSuccess) return CTC_STATUS_MEMOPS_FAILED;

  labels_with_blanks_ = reinterpret_cast<int*>(
      static_cast<char*>(gpu_workspace_) + gpu_bytes_used);
  gpu_bytes_used += Smax * minibatch_ * sizeof(int);

  alphas_ = reinterpret_cast<ProbT*>(static_cast<char*>(gpu_workspace_) +
                                     gpu_bytes_used);
  gpu_bytes_used += (S_ * T_) * minibatch_ * sizeof(ProbT);

  denoms_ = reinterpret_cast<ProbT*>(static_cast<char*>(gpu_workspace_) +
                                     gpu_bytes_used);
  gpu_bytes_used += activation_cols_ * sizeof(ProbT);

  probs_ = reinterpret_cast<ProbT*>(static_cast<char*>(gpu_workspace_) +
                                    gpu_bytes_used);
  gpu_bytes_used += out_dim_ * activation_cols_ * sizeof(ProbT);

  return CTC_STATUS_SUCCESS;
}

template <typename ProbT>
template <int NT, int VT>
ctcStatus_t GpuCTC<ProbT>::launch_alpha_beta_kernels(const ProbT* const probs,
                                                     ProbT* grads,
                                                     bool compute_alpha,
                                                     bool compute_beta) {
  // One thread block per utterance
  const int grid_size = minibatch_;

  // The data is laid out so that the next timestep is minibatch entries
  // away
  const int stride = minibatch_;

  if (compute_alpha)
    compute_alpha_kernel<ProbT, NT, VT><<<grid_size, NT, 0, stream_>>>(
        probs,
        label_sizes_,
        utt_length_,
        repeats_,
        labels_without_blanks_,
        label_offsets_,
        labels_with_blanks_,
        alphas_,
        nll_forward_,
        stride,
        out_dim_,
        S_,
        T_,
        blank_label_);

  if (compute_beta) {
    compute_betas_and_grad_kernel<ProbT, NT, VT><<<grid_size, NT, 0, stream_>>>(
        probs,
        label_sizes_,
        utt_length_,
        repeats_,
        labels_with_blanks_,
        alphas_,
        nll_forward_,
        nll_backward_,
        grads,
        stride,
        out_dim_,
        S_,
        T_,
        blank_label_);

    cudaStreamSynchronize(stream_);
  }

  cudaError_t err = cudaGetLastError();
  if (err != cudaSuccess) return CTC_STATUS_EXECUTION_FAILED;

  return CTC_STATUS_SUCCESS;
}

template <typename ProbT>
ctcStatus_t GpuCTC<ProbT>::create_metadata_and_choose_config(
    const int* const flat_labels,
    const int* const label_lengths,
    const int* const input_lengths,
    size_t& best_config) {
  // Setup the metadata for GPU
  ctcStatus_t status =
      setup_gpu_metadata(flat_labels, label_lengths, input_lengths);
  if (status != CTC_STATUS_SUCCESS) return status;

  constexpr int num_configs = 12;

  int config_NT[num_configs] = {
      32, 64, 128, 64, 128, 32, 64, 128, 64, 128, 128, 128};
  int config_VT[num_configs] = {1, 1, 1, 3, 2, 9, 6, 4, 9, 6, 9, 10};

  best_config = 0;

  for (int i = 0; i < num_configs; ++i) {
    if ((config_NT[i] * config_VT[i]) >= S_)
      break;
    else
      best_config++;
  }

  if (best_config >= num_configs) return CTC_STATUS_UNKNOWN_ERROR;

  return CTC_STATUS_SUCCESS;
}

template <typename ProbT>
ctcStatus_t GpuCTC<ProbT>::launch_gpu_kernels(
    const ProbT* const probs, ProbT* grads, size_t config, bool l_a, bool l_b) {
  switch (config) {
    case 0: {
      return launch_alpha_beta_kernels<32, 1>(probs, grads, l_a, l_b);
    }
    case 1: {
      return launch_alpha_beta_kernels<64, 1>(probs, grads, l_a, l_b);
    }
    case 2: {
      return launch_alpha_beta_kernels<128, 1>(probs, grads, l_a, l_b);
    }
    case 3: {
      return launch_alpha_beta_kernels<64, 3>(probs, grads, l_a, l_b);
    }
    case 4: {
      return launch_alpha_beta_kernels<128, 2>(probs, grads, l_a, l_b);
    }
    case 5: {
      return launch_alpha_beta_kernels<32, 9>(probs, grads, l_a, l_b);
    }
    case 6: {
      return launch_alpha_beta_kernels<64, 6>(probs, grads, l_a, l_b);
    }
    case 7: {
      return launch_alpha_beta_kernels<128, 4>(probs, grads, l_a, l_b);
    }
    case 8: {
      return launch_alpha_beta_kernels<64, 9>(probs, grads, l_a, l_b);
    }
    case 9: {
      return launch_alpha_beta_kernels<128, 6>(probs, grads, l_a, l_b);
    }
    case 10: {
      return launch_alpha_beta_kernels<128, 9>(probs, grads, l_a, l_b);
    }
    case 11: {
      return launch_alpha_beta_kernels<128, 10>(probs, grads, l_a, l_b);
    }
  }

  return CTC_STATUS_EXECUTION_FAILED;
}

template <typename ProbT>
ctcStatus_t GpuCTC<ProbT>::compute_probs(const ProbT* const activations) {
  cudaError_t cuda_status;
  cuda_status = cudaMemcpyAsync(probs_,
                                activations,
                                activation_cols_ * out_dim_ * sizeof(ProbT),
                                cudaMemcpyDeviceToDevice,
                                stream_);
  if (cuda_status != cudaSuccess) return CTC_STATUS_MEMOPS_FAILED;

  // Numerically stable SM
  ctcStatus_t ctc_status =
      reduce_max(probs_, denoms_, out_dim_, activation_cols_, 1, stream_);
  if (ctc_status != CTC_STATUS_SUCCESS) return ctc_status;

  // Kernel launch to subtract maximum
  const int NT = 128;
  const int VT = 1;
  const int NV = NT * VT;
  const int num_elements = out_dim_ * activation_cols_;
  const int grid_size = ctc_helper::div_up(num_elements, NV);

  prepare_stable_SM_kernel<ProbT, VT><<<grid_size, NT, 0, stream_>>>(
      ctc_helper::identity<ProbT>(), probs_, denoms_, out_dim_, num_elements);

  // Reduce along columns to calculate denominator
  ctc_status =
      reduce_exp(probs_, denoms_, out_dim_, activation_cols_, 1, stream_);
  if (ctc_status != CTC_STATUS_SUCCESS) return ctc_status;

  // Kernel launch to calculate probabilities
  compute_probs_kernel<ProbT, VT><<<grid_size, NT, 0, stream_>>>(
      ctc_helper::exponential<ProbT>(),
      probs_,
      denoms_,
      out_dim_,
      num_elements);

  truncate_probs_kernel<ProbT, VT><<<grid_size, NT, 0, stream_>>>(probs_,
                                                                  num_elements);

  return CTC_STATUS_SUCCESS;
}

template <typename ProbT>
ctcStatus_t GpuCTC<ProbT>::compute_cost_and_score(
    const ProbT* const activations,
    ProbT* grads,
    ProbT* costs,
    const int* const flat_labels,
    const int* const label_lengths,
    const int* const input_lengths,
    bool compute_alpha,
    bool compute_betas_and_grad) {
  size_t best_config;
  ctcStatus_t status = create_metadata_and_choose_config(
      flat_labels, label_lengths, input_lengths, best_config);
  if (status != CTC_STATUS_SUCCESS) return status;

  status = compute_probs(activations);
  if (status != CTC_STATUS_SUCCESS) return status;

  launch_gpu_kernels(
      probs_, grads, best_config, compute_alpha, compute_betas_and_grad);

  cudaError_t cuda_status_mem, cuda_status_sync;
  cuda_status_mem = cudaMemcpyAsync(costs,
                                    nll_forward_,
                                    sizeof(ProbT) * minibatch_,
                                    cudaMemcpyDeviceToHost,
                                    stream_);
  cuda_status_sync = cudaStreamSynchronize(stream_);
  if (cuda_status_mem != cudaSuccess || cuda_status_sync != cudaSuccess)
    return CTC_STATUS_MEMOPS_FAILED;

  return CTC_STATUS_SUCCESS;
}

template <typename ProbT>
ctcStatus_t GpuCTC<ProbT>::cost_and_grad(const ProbT* const activations,
                                         ProbT* grads,
                                         ProbT* costs,
                                         const int* const flat_labels,
                                         const int* const label_lengths,
                                         const int* const input_lengths) {
  if (activations == nullptr || grads == nullptr || costs == nullptr ||
      flat_labels == nullptr || label_lengths == nullptr ||
      input_lengths == nullptr)
    return CTC_STATUS_INVALID_VALUE;

  return compute_cost_and_score(activations,
                                grads,
                                costs,
                                flat_labels,
                                label_lengths,
                                input_lengths,
                                true,
                                true);
}

template <typename ProbT>
ctcStatus_t GpuCTC<ProbT>::score_forward(const ProbT* const activations,
                                         ProbT* costs,
                                         const int* const flat_labels,
                                         const int* const label_lengths,
                                         const int* const input_lengths) {
  if (activations == nullptr || costs == nullptr || flat_labels == nullptr ||
      label_lengths == nullptr || input_lengths == nullptr)
    return CTC_STATUS_INVALID_VALUE;

  return compute_cost_and_score(activations,
                                nullptr,
                                costs,
                                flat_labels,
                                label_lengths,
                                input_lengths,
                                true,
                                false);
}
