// cudadecoder/cuda-decoder-common.h
  //
  // Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
  // Hugo Braun, Justin Luitjens, Ryan Leary
  //
  // 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.
  
  #ifndef KALDI_CUDA_DECODER_CUDA_DECODER_UTILS_H_
  #define KALDI_CUDA_DECODER_CUDA_DECODER_UTILS_H_
  #include "cudamatrix/cu-device.h"
  #include "util/stl-utils.h"
  
  // A decoder channel is linked to one utterance. Frames
  // from the same must be sent to the same channel.
  //
  // A decoder lane is where the computation actually happens
  // a decoder lane is given a frame and its associated channel
  // and does the actual computation
  //
  // An analogy would be lane -> a core, channel -> a software thread
  
  // Some config parameters can be computed using other parameters
  // (e.g. we can set main_q_capacity using max-active)
  // Those values are the different factors between parameters that we know
  // and parameters we want to set
  #define KALDI_CUDA_DECODER_MAX_ACTIVE_MAIN_Q_CAPACITY_FACTOR 4
  #define KALDI_CUDA_DECODER_AUX_Q_MAIN_Q_CAPACITIES_FACTOR 3
  
  // If we're at risk of filling the tokens queue,
  // the beam is reduced to keep only the best candidates in the
  // remaining space
  // We then slowly put the beam back to its default value
  // beam_next_frame = min(default_beam, RECOVER_RATE * beam_previous_frame)
  #define KALDI_CUDA_DECODER_ADAPTIVE_BEAM_RECOVER_RATE 1.2f
  
  // Defines for the cuda decoder kernels
  // It shouldn't be necessary to change the DIMX of the kernels
  
  // Below that value, we launch the persistent kernel for NonEmitting
  #define KALDI_CUDA_DECODER_NONEM_LT_MAX_NARCS 4096
  
  // We know we will have at least X elements in the hashmap
  // We allocate space for X*KALDI_CUDA_DECODER_HASHMAP_CAPACITY_FACTOR elements
  // to avoid having too much collisions
  #define KALDI_CUDA_DECODER_HASHMAP_CAPACITY_FACTOR 1
  
  // Max size of the total kernel arguments
  // 4kb for compute capability >= 2.0
  #define KALDI_CUDA_DECODER_MAX_KERNEL_ARGUMENTS_BYTE_SIZE (4096)
  
  // When applying the max-active, we need to compute a topk
  // to perform that (soft) topk, we compute a histogram
  // here we define the number of bins in that histogram
  // it has to be less than the number of 1D threads
  #define KALDI_CUDA_DECODER_HISTO_NBINS 255
  
  // Number of "heavy duty" process non emitting kernels
  // If more non emitting iterations are required, those will be done
  // in the one-CTA persistent kernel
  #define KALDI_CUDA_DECODER_N_NON_EMITTING_MAIN_ITERATIONS 2
  
  // Adaptive beam parameters
  // We will decrease the beam when we detect that we are generating too many
  // tokens
  // for the first segment of the aux_q, we don't do anything (keep the original
  // beam)
  // the first segment is made of (aux_q
  // capacity)/KALDI_CUDA_DECODER_ADAPTIVE_BEAM_STATIC_SEGMENT
  // then we will decrease the beam step by step, until 0.
  // we will decrease the beam every m elements, with:
  // x = (aux_q capacity)/KALDI_CUDA_DECODER_ADAPTIVE_BEAM_STATIC_SEGMENT (static
  // segment
  // y = (aux_q capacity) - x
  // m = y / KALDI_CUDA_DECODER_ADAPTIVE_BEAM_NSTEPS
  // For more information, please refer to the definition of GetAdaptiveBeam in
  // cuda-decoder-kernels.cu
  #define KALDI_CUDA_DECODER_ADAPTIVE_BEAM_STATIC_SEGMENT 4
  #define KALDI_CUDA_DECODER_ADAPTIVE_BEAM_NSTEPS 8
  // When applying max_active we don't keep exactly max_active_ tokens,
  // but a bit more. And we can call ApplyMaxActiveAndReduceBeam multiple times
  // in the first frame (the first times as a pre-filter, the last time at the
  // very end of the frame)
  // Because keeping a bit more than max_active_ is expected, we add the tolerance
  // so that we can avoid triggering ApplyMaxActiveAndReduceBeam for just a few
  // tokens above the limit
  // at the end of the frame
  
  #define KALDI_CUDA_DECODER_DIV_ROUND_UP(a, b) ((a + b - 1) / b)
  
  #define KALDI_CUDA_DECODER_ASSERT(val, recoverable)                     \
    {                                                                     \
      if ((val) != true) {                                                \
        throw CudaDecoderException("KALDI_CUDA_DECODER_ASSERT", __FILE__, \
                                   __LINE__, recoverable)                 \
      }                                                                   \
    }
  // Macro for checking cuda errors following a cuda launch or api call
  #ifdef NDEBUG
  #define KALDI_DECODER_CUDA_CHECK_ERROR()
  #else
  #define KALDI_DECODER_CUDA_CHECK_ERROR()                                  \
    {                                                                       \
      cudaError_t e = cudaGetLastError();                                   \
      if (e != cudaSuccess) {                                               \
        throw CudaDecoderException(cudaGetErrorName(e), __FILE__, __LINE__, \
                                   false);                                  \
      }                                                                     \
    }
  #endif
  
  #define KALDI_DECODER_CUDA_API_CHECK_ERROR(e)                             \
    {                                                                       \
      if (e != cudaSuccess) {                                               \
        throw CudaDecoderException(cudaGetErrorName(e), __FILE__, __LINE__, \
                                   false);                                  \
      }                                                                     \
    }
  
  #define KALDI_CUDA_DECODER_1D_KERNEL_LOOP(i, n)                \
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
         i += blockDim.x * gridDim.x)
  
  #define KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(offset, th_idx, n) \
    for (int offset = blockIdx.x * blockDim.x, th_idx = threadIdx.x;        \
         offset < (n); offset += blockDim.x * gridDim.x)
  
  #define KALDI_CUDA_DECODER_IS_LAST_1D_THREAD() (threadIdx.x == (blockDim.x - 1))
  
  #define KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(i, n) \
    for (int i = blockIdx.y; i < (n); i += gridDim.y)
  
  #define KALDI_CUDA_DECODER_DIV_ROUND_UP(a, b) ((a + b - 1) / b)
  
  #define KALDI_CUDA_DECODER_1D_BLOCK 256
  #define KALDI_CUDA_DECODER_LARGEST_1D_BLOCK 1024
  #define KALDI_CUDA_DECODER_ONE_THREAD_BLOCK 1
  #define KALDI_CUDA_DECODER_MAX_CTA_COUNT 4096u
  #define KALDI_CUDA_DECODER_MAX_CTA_PER_LANE 512u
  namespace kaldi {
  namespace cuda_decoder {
  
  // Returning the number of CTAs to launch for (N,M) elements to compute
  // M is usually the batch size
  inline dim3 KaldiCudaDecoderNumBlocks(int N, int M) {
    dim3 grid;
    grid.x = KALDI_CUDA_DECODER_DIV_ROUND_UP(N, KALDI_CUDA_DECODER_1D_BLOCK);
    unsigned int max_CTA_per_lane =
        std::max(KALDI_CUDA_DECODER_MAX_CTA_COUNT / M, 1u);
    grid.x = std::min(grid.x, max_CTA_per_lane);
    grid.y = M;
    return grid;
  }
  
  // Use a fixed number of blocks for nlanes
  // Using the max number of CTAs possible for each lane,
  // according to KALDI_CUDA_DECODER_MAX_CTA_COUNT
  // and KALDI_CUDA_DECODER_MAX_CTA_PER_LANE
  inline dim3 KaldiCudaDecoderNumBlocks(int nlanes) {
    dim3 grid;
    unsigned int n_CTA_per_lane =
        std::max(KALDI_CUDA_DECODER_MAX_CTA_COUNT / nlanes, 1u);
    if (n_CTA_per_lane == 0) n_CTA_per_lane = 1;
    grid.x = std::min(KALDI_CUDA_DECODER_MAX_CTA_PER_LANE, n_CTA_per_lane);
    grid.y = nlanes;
    return grid;
  }
  
  typedef int32 StateId;
  typedef float CostType;
  // IntegerCostType is the type used in the lookup table d_state_best_cost
  // and the d_cutoff
  // We use a 1:1 conversion between CostType <--> IntegerCostType
  // IntegerCostType is used because it triggers native atomic operations
  // (CostType does not)
  typedef int32 IntegerCostType;
  typedef int32 LaneId;
  typedef int32 ChannelId;
  
  // On the device we compute everything by batch
  // Data is stored as 2D matrices (BatchSize, 1D_Size)
  // For example, for the token queue, (BatchSize, max_tokens_per_frame_)
  // DeviceMatrix owns the data but is not used to access it.
  // DeviceMatrix is inherited in DeviceLaneMatrix and DeviceChannelMatrix
  // those two classes do the same thing, except that they belong either to a
  // channel or lane
  // that inheritance is done to clarify the code and help debugging
  //
  // To actually access the data, we should request an view through
  // GetView
  // That view contains both host cuda code to access the data. It does not own
  // the data.
  template <typename T>
  // if necessary, make a version that always use ncols_ as the next power of 2
  class DeviceMatrix {
    T *data_;
    void Allocate() {
      KALDI_ASSERT(nrows_ > 0);
      KALDI_ASSERT(ncols_ > 0);
      KALDI_ASSERT(!data_);
      data_ = static_cast<T *>(CuDevice::Instantiate().Malloc(
          (size_t)nrows_ * ncols_ * sizeof(*data_)));
      KALDI_ASSERT(data_);
    }
    void Free() {
      KALDI_ASSERT(data_);
      CuDevice::Instantiate().Free(data_);
    }
  
   protected:
    int32 ncols_;
    int32 nrows_;
  
   public:
    DeviceMatrix() : data_(NULL), ncols_(0), nrows_(0) {}
  
    virtual ~DeviceMatrix() {
      if (data_) Free();
    }
  
    void Resize(int32 nrows, int32 ncols) {
      if (data_) Free();
      KALDI_ASSERT(nrows > 0);
      KALDI_ASSERT(ncols > 0);
      nrows_ = nrows;
      ncols_ = ncols;
      Allocate();
    }
  
    T *MutableData() {
      KALDI_ASSERT(data_);
      return data_;
    }
    // abstract getInterface...
  };
  
  template <typename T>
  // if necessary, make a version that always use ncols_ as the next power of 2
  class HostMatrix {
    T *data_;
    void Allocate() {
      KALDI_ASSERT(nrows_ > 0);
      KALDI_ASSERT(ncols_ > 0);
      KALDI_ASSERT(!data_);
      cudaMallocHost((void **)&data_, (size_t)nrows_ * ncols_ * sizeof(*data_));
      KALDI_ASSERT(data_);
    }
    void Free() {
      KALDI_ASSERT(data_);
      cudaFreeHost(data_);
    }
  
   protected:
    int32 ncols_;
    int32 nrows_;
  
   public:
    HostMatrix() : data_(NULL), ncols_(0), nrows_(0) {}
  
    virtual ~HostMatrix() {
      if (data_) Free();
    }
  
    void Resize(int32 nrows, int32 ncols) {
      if (data_) Free();
      KALDI_ASSERT(nrows > 0);
      KALDI_ASSERT(ncols > 0);
      nrows_ = nrows;
      ncols_ = ncols;
      Allocate();
    }
  
    T *MutableData() {
      KALDI_ASSERT(data_);
      return data_;
    }
    // abstract getInterface...
  };
  
  // Views of DeviceMatrix
  // Those views are created by either DeviceChannelMatrix or
  // DeviceLaneMatrix
  // We can access the data (the matrix) associated with that
  // Device[Channel|Lane]Matrix without owning that data.
  // Which means that we can pass those views by copy
  // without triggering a cudaFree, for instance.
  // Device[Channel|Lane]Matrix owns the data, [Channel|Lane]MatrixInterface just
  // gives access to it
  // Generating both host and device interfaces
  template <typename T>
  struct LaneMatrixView {
    T *data_;
    int32 ncols_;
    __host__ __device__ __inline__ T *lane(const int32 ilane) {
      return &data_[ilane * ncols_];
    }
  };
  
  template <typename T>
  struct ChannelMatrixView {
    T *data_;
    int32 ncols_;
    __host__ __device__ __inline__ T *channel(const int32 ichannel) {
      return &data_[ichannel * ncols_];
    }
  };
  
  // Specializing DeviceMatrix into lane and channel variants.
  // Helps with code clarity/debugging
  template <typename T>
  class DeviceLaneMatrix : public DeviceMatrix<T> {
   public:
    LaneMatrixView<T> GetView() { return {this->MutableData(), this->ncols_}; }
  
    T *lane(const int32 ilane) {
      return &this->MutableData()[ilane * this->ncols_];
    }
  };
  
  template <typename T>
  class HostLaneMatrix : public HostMatrix<T> {
   public:
    LaneMatrixView<T> GetView() { return {this->MutableData(), this->ncols_}; }
  
    T *lane(const int32 ilane) {
      return &this->MutableData()[ilane * this->ncols_];
    }
  };
  
  template <typename T>
  class DeviceChannelMatrix : public DeviceMatrix<T> {
   public:
    ChannelMatrixView<T> GetView() { return {this->MutableData(), this->ncols_}; }
    T *channel(const int32 ichannel) {
      return &this->MutableData()[ichannel * this->ncols_];
    }
  };
  
  // LaneCounters/ChannelCounters
  // The counters are all the singular values associated to a lane/channel
  // For instance  the main queue size. Or the min_cost of all tokens in that
  // queue
  // LaneCounters are used during computation
  struct LaneCounters {
    // hannel that this lane will compute for the current frame
    ChannelId channel_to_compute;
    // Pointer to the loglikelihoods array for this channel and current frame
    BaseFloat *loglikelihoods;
    // Contains both main_q_end and narcs
    // End index of the main queue
    // only tokens at index i with i < main_q_end
    // are valid tokens
    // Each valid token the subqueue main_q[main_q_local_offset, main_q_end[ has
    // a number of outgoing arcs (out-degree)
    // main_q_narcs is the sum of those numbers
    // We sometime need to update both end and narcs at the same time using a
    // single atomic,
    // which is why they're packed together
    int2 main_q_narcs_and_end;
    // contains the requested queue length which can
    // be larger then the actual queue length in the case of overflow
    int32 main_q_requested;
    int32 aux_q_requested;
    int32 aux_q_end;
    int32 post_expand_aux_q_end;  // used for double buffering
    // Some tokens in the same frame share the same token.next_state
    // main_q_n_extra_prev_tokens is the count of those tokens
    int32 main_q_n_extra_prev_tokens;
    // Number of tokens created during the emitting stage
    int32 main_q_n_emitting_tokens;
    // Depending on the value of the parameter "max_tokens_per_frame"
    // we can end up with an overflow when generating the tokens for a frame
    // We try to prevent this from happening using an adaptive beam
    // If an overflow happens, then the kernels no longer insert any data into
    // the queues and set overflow flag to true.
    // queue length.
    // Even if that flag is set, we can continue the execution (quality
    // of the output can be lowered)
    // We use that flag to display a warning to the user
    int32 q_overflow;
    // ExpandArcs reads the tokens in the index range [main_q_local_offset, end[
    int32 main_q_local_offset;
    // We transfer the tokens back to the host at the end of each frame.
    // Which means that tokens at a frame  n > 0 have an offset compared to to
    // those
    // in frame n-1. main_q_global_offset is the overall offset of the current
    // main_q,
    // since frame 0
    // It is used to set the prev_token index.
    int32 main_q_global_offset;
    // Same thing, but for main_q_n_extra_prev_tokens (those are also transfered
    // back to host)
    int32 main_q_extra_prev_tokens_global_offset;
    // Minimum token for that frame
    IntegerCostType min_int_cost;
    // Current beam. Can be different from default_beam,
    // because of the AdaptiveBeam process, or because of
    // ApplyMaxActiveAndReduceBeam
    IntegerCostType int_beam;
    // Adaptive beam. The validity says until which index this adaptive beam is
    // valid.
    // After that index, we need to lower the adaptive beam
    int2 adaptive_int_beam_with_validity_index;
    // min_cost + beam
    IntegerCostType int_cutoff;
    // The histogram for max_active will be computed between min_histo_cost
    // and max_histo_cost. Set for each frame after emitting stage
    CostType min_histo_cost;
    CostType max_histo_cost;
    CostType histo_bin_width;
    bool compute_max_active;
    // offsets used by concatenate_lanes_data_kernel
    int32 main_q_end_lane_offset;
    int32 main_q_n_emitting_tokens_lane_offset;
    int32 main_q_n_extra_prev_tokens_lane_offset;
  
    // --- Only valid after calling GetBestCost
    // min_cost and its arg. Can be different than min_cost, because we may
    // include final costs
    int2 min_int_cost_and_arg;
    // Number of final tokens with cost < best + lattice_beam
    int32 n_within_lattice_beam;
    int32 has_reached_final;  // if there's at least one final token in the queue
    int32 prev_arg_min_int_cost;
  };
  
  // Channel counters
  // Their job is to save the state of a channel, when this channel is idle
  // The channel counters are loaded into the lane counters during the context
  // switches
  struct ChannelCounters {
    // All the following values are just saved values from LaneCounters
    // from the latest context-switch
    int2 prev_main_q_narcs_and_end;
    int32 prev_main_q_n_extra_prev_tokens;
    int32 prev_main_q_global_offset;
    int32 prev_main_q_extra_prev_tokens_global_offset;
    CostType prev_beam;
  
    // Only valid after calling GetBestCost
    // different than min_int_cost : we include the "final" cost
    int2 min_int_cost_and_arg_with_final;
    int2 min_int_cost_and_arg_without_final;
  };
  
  class CudaDecoderException : public std::exception {
   public:
    CudaDecoderException(const char *str_, const char *file_, int line_,
                         const bool recoverable_)
        : str(str_),
          file(file_),
          line(line_),
          buffer(std::string(file) + ":" + std::to_string(line) + " :" +
                 std::string(str)),
          recoverable(recoverable_) {}
    const char *what() const throw() { return buffer.c_str(); }
  
    const char *str;
    const char *file;
    const int line;
    const std::string buffer;
    const bool recoverable;
  };
  
  // InfoToken contains data that needs to be saved for the backtrack
  // in GetBestPath/GetRawLattice
  // We don't need the token.cost or token.next_state.
  struct __align__(8) InfoToken {
    int32 prev_token;
    int32 arc_idx;
    bool IsUniqueTokenForStateAndFrame() {
      // This is a trick used to save space and PCI-E bandwidth (cf
      // preprocess_in_place kernel)
      // This token is associated with a next_state s, created during the
      // processing of frame f.
      // If we have multiple tokens associated with the state s in the frame f,
      // arc_idx < 0 and -arc_idx is the
      // count of such tokens. We will then have to look at another list to read
      // the actually arc_idx and prev_token values
      // If the current token is the only one, prev_token and arc_idx are valid
      // and can be used directly
      return (arc_idx >= 0);
    }
  
    // Called if this token is linked to others tokens in the same frame (cf
    // comments for IsUniqueTokenForStateAndFrame)
    // return the {offset,size} pair necessary to list those tokens in the
    // extra_prev_tokens list
    // They are stored at offset "offset", and we have "size" of those
    std::pair<int32, int32> GetSameFSTStateTokensList() {
      KALDI_ASSERT(!IsUniqueTokenForStateAndFrame());
  
      return {prev_token, -arc_idx};
    }
  };
  
  // Device function, used to set a in an InfoToken the [offset,size] related to
  // InfoToken.GetSameFSTStateTokensList
  __device__ __inline__ void SetSameFSTStateTokensList(int32 offset, int32 size,
                                                       InfoToken *info_token) {
    // We always have size > 0
    *info_token = {offset, -size};
  }
  
  // Used to store the index in the GPU hashmap of that FST state
  // The hashmap is only generated with the final main queue (post max_active_) of
  // each frame
  // Also stores the information or whether or not the owner of that object is the
  // representative of this FSTState
  typedef int32 FSTStateHashIndex;
  
  // 1:1 Conversion float <---> sortable int
  // We convert floats to sortable ints in order
  // to use native atomics operation
  // Those are the host version, used when we transfer an int from the device
  // and we want to convert it to a float
  // (it was created on device by floatToOrderedInt, we'll use
  // orderedIntToFloatHost on host to convert it back to a float)
  __inline__ int32 floatToOrderedIntHost(float floatVal) {
    int32 intVal;
    // Should be optimized away by compiler
    memcpy(&intVal, &floatVal, sizeof(float));
    return (intVal >= 0) ? intVal : intVal ^ 0x7FFFFFFF;
  }
  
  __inline__ float orderedIntToFloatHost(int32 intVal) {
    intVal = (intVal >= 0) ? intVal : intVal ^ 0x7FFFFFFF;
    float floatVal;
    // Should be optimized away by compiler
    memcpy(&floatVal, &intVal, sizeof(float));
    return floatVal;
  }
  
  // Hashmap value. Used when computing the hashmap in PostProcessingMainQueue
  struct __align__(16) HashmapValueT {
    // Map key : fst state
    int32 key;
    // Number of tokens associated to that state
    int32 count;
    // minimum cost for that state + argmin
    unsigned long long min_and_argmin_int_cost_u64;
  };
  
  enum OVERFLOW_TYPE {
    OVERFLOW_NONE = 0,
    OVERFLOW_MAIN_Q = 1,
    OVERFLOW_AUX_Q = 2
  };
  
  enum QUEUE_ID { MAIN_Q = 0, AUX_Q = 1 };
  
  }  // end namespace cuda_decoder
  }  // end namespace kaldi
  
  #endif  // KALDI_CUDA_DECODER_CUDA_DECODER_UTILS_H_