// cudadecoder/cuda-decoder.cc
//
// 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.

#if HAVE_CUDA == 1

#include "cuda-decoder.h"
#include "cuda-decoder-kernels.h"

#include <cuda_runtime_api.h>
#include <nvToolsExt.h>
#include <algorithm>
#include <cfloat>
#include <map>
#include <tuple>

namespace kaldi {
namespace cuda_decoder {
CudaDecoder::CudaDecoder(const CudaFst &fst, const CudaDecoderConfig &config,
                         int32 nlanes, int32 nchannels)
    : fst_(fst),
      nlanes_(nlanes),
      nchannels_(nchannels),
      channel_lock_(nchannels + 1),
      extra_cost_min_delta_(0.0f),
      thread_pool_(NULL),
      n_threads_used_(0),
      n_h2h_task_not_done_(0),
      n_init_decoding_h2h_task_not_done_(0),
      h2h_threads_running_(true) {
  ReadConfig(config);
  // Static asserts on constants
  CheckStaticAsserts();
  // Runtime asserts
  KALDI_ASSERT(nlanes > 0);
  KALDI_ASSERT(nchannels > 0);
  KALDI_ASSERT(nlanes_ <= nchannels_);
  // All GPU work in decoder will be sent to compute_st_
  cudaStreamCreate(&compute_st_);
  // Copies D2H of tokens for storage on host are done on
  // copy_st_, in parallel with compute_st_
  cudaStreamCreate(&copy_st_);
  // For all the allocating/initializing process
  // We create a special channel
  // containing the exact state a channel should have when starting a new decode
  // It contains fst.Start(), the non-emitting tokens created by fst.Start(),
  // and all the data used by the decoder.
  // When calling InitDecoding() on a new channel, we simply clone this special
  // channel into that new channel
  ++nchannels_;                       // adding the special initial channel
  init_channel_id_ = nchannels_ - 1;  // Using last one as init_channel_params
  AllocateHostData();
  AllocateDeviceData();
  AllocateDeviceKernelParams();

  InitDeviceParams();
  InitHostData();
  InitDeviceData();

  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventCreate(&nnet3_done_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventCreate(&d2h_copy_acoustic_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventCreate(&d2h_copy_infotoken_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaEventCreate(&d2h_copy_extra_prev_tokens_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaEventCreate(&concatenated_data_ready_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventCreate(&lane_offsets_ready_evt_));

  ComputeInitialChannel();
  --nchannels_;  // removing the special initial channel from the count

  // Making sure that everything is ready to use
  cudaStreamSynchronize(compute_st_);
  KALDI_DECODER_CUDA_CHECK_ERROR();
}

void CudaDecoder::ReadConfig(const CudaDecoderConfig &cst_config) {
  CudaDecoderConfig config = cst_config;  // deep copy
  // Sets the missing values using other values
  config.ComputeConfig();
  default_beam_ = config.default_beam;
  lattice_beam_ = config.lattice_beam;
  ntokens_pre_allocated_ = config.ntokens_pre_allocated;
  max_active_ = config.max_active;
  aux_q_capacity_ = config.aux_q_capacity;
  main_q_capacity_ = config.main_q_capacity;

  KALDI_ASSERT(default_beam_ >= 0.0f);
  KALDI_ASSERT(lattice_beam_ >= 0.0f);
  KALDI_ASSERT(ntokens_pre_allocated_ >= 0);
  KALDI_ASSERT(max_active_ > 0);
  KALDI_ASSERT(main_q_capacity_ > 0);
  KALDI_ASSERT(aux_q_capacity_ >= main_q_capacity_);
}

void CudaDecoder::AllocateDeviceData() {
  hashmap_capacity_ =
      KALDI_CUDA_DECODER_HASHMAP_CAPACITY_FACTOR * main_q_capacity_;
  d_channels_counters_.Resize(nchannels_, 1);
  d_lanes_counters_.Resize(
      nlanes_ + 1,
      1);  // +1 because we sometimes need last+1 value (for offsets)
  d_main_q_state_and_cost_.Resize(nchannels_, main_q_capacity_);
  d_main_q_info_.Resize(nlanes_, main_q_capacity_);
  d_aux_q_state_and_cost_.Resize(nlanes_, aux_q_capacity_);
  d_aux_q_info_.Resize(nlanes_, aux_q_capacity_);
  d_main_q_degrees_prefix_sum_.Resize(nchannels_, main_q_capacity_);
  d_histograms_.Resize(nlanes_, KALDI_CUDA_DECODER_HISTO_NBINS);
  d_main_q_extra_prev_tokens_prefix_sum_.Resize(nlanes_, main_q_capacity_);
  d_main_q_n_extra_prev_tokens_local_idx_.Resize(nlanes_, main_q_capacity_);

  d_main_q_state_hash_idx_.Resize(nlanes_, main_q_capacity_);
  d_main_q_extra_prev_tokens_.Resize(nlanes_, main_q_capacity_);
  d_main_q_extra_and_acoustic_cost_.Resize(nlanes_, main_q_capacity_);
  d_main_q_block_sums_prefix_sum_.Resize(
      nlanes_, KALDI_CUDA_DECODER_DIV_ROUND_UP(main_q_capacity_,
                                               KALDI_CUDA_DECODER_1D_BLOCK) +
                   1);
  d_main_q_arc_offsets_.Resize(nchannels_, main_q_capacity_);
  d_hashmap_values_.Resize(nlanes_, hashmap_capacity_);
  d_main_q_acoustic_cost_.Resize(nlanes_, main_q_capacity_);
  d_extra_and_acoustic_cost_concat_matrix_.Resize(nlanes_, main_q_capacity_);
  d_acoustic_cost_concat_matrix_.Resize(nlanes_, main_q_capacity_);
  d_infotoken_concat_matrix_.Resize(nlanes_, main_q_capacity_);
  d_extra_prev_tokens_concat_matrix_.Resize(nlanes_, main_q_capacity_);
  // Reusing data from aux_q. Those two are never used at the same time
  // d_list_final_tokens_in_main_q_ is used in GetBestPath.
  // the aux_q is used in AdvanceDecoding
  h_list_final_tokens_in_main_q_.Resize(nlanes_, main_q_capacity_);
  d_extra_prev_tokens_concat_ = d_extra_prev_tokens_concat_matrix_.lane(0);
  d_extra_and_acoustic_cost_concat_ =
      d_extra_and_acoustic_cost_concat_matrix_.lane(0);
  d_acoustic_cost_concat_ = d_acoustic_cost_concat_matrix_.lane(0);
  d_infotoken_concat_ = d_infotoken_concat_matrix_.lane(0);
}

void CudaDecoder::AllocateHostData() {
  channel_to_compute_.resize(nlanes_);
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(
      &h_extra_and_acoustic_cost_concat_,
      nlanes_ * main_q_capacity_ * sizeof(*h_extra_and_acoustic_cost_concat_)));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(
      &h_acoustic_cost_concat_,
      nlanes_ * main_q_capacity_ * sizeof(*h_acoustic_cost_concat_)));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(
      &h_extra_prev_tokens_concat_,
      nlanes_ * main_q_capacity_ * sizeof(*h_extra_prev_tokens_concat_)));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(
      &h_infotoken_concat_,
      nlanes_ * main_q_capacity_ * sizeof(*h_infotoken_concat_)));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaMallocHost(&h_extra_and_acoustic_cost_concat_tmp_,
                     nlanes_ * main_q_capacity_ *
                         sizeof(*h_extra_and_acoustic_cost_concat_tmp_)));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(
      &h_acoustic_cost_concat_tmp_,
      nlanes_ * main_q_capacity_ * sizeof(*h_acoustic_cost_concat_tmp_)));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(
      &h_extra_prev_tokens_concat_tmp_,
      nlanes_ * main_q_capacity_ * sizeof(*h_extra_prev_tokens_concat_tmp_)));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(
      &h_infotoken_concat_tmp_,
      nlanes_ * main_q_capacity_ * sizeof(*h_infotoken_concat_tmp_)));
  h_lanes_counters_.Resize(
      nlanes_ + 1,
      1);  // +1 because we sometimes need last+1 value (for offsets)
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(
      &h_channels_counters_, nchannels_ * sizeof(*h_channels_counters_)));

  h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_.resize(nchannels_);
  h_all_tokens_acoustic_cost_.resize(nchannels_);
  h_all_tokens_extra_prev_tokens_.resize(nchannels_);
  h_all_tokens_info_.resize(nchannels_);
  for (int32 ichannel = 0; ichannel < nchannels_; ++ichannel) {
    h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_[ichannel].reserve(
        ntokens_pre_allocated_);
    h_all_tokens_acoustic_cost_[ichannel].reserve(ntokens_pre_allocated_);
    h_all_tokens_info_[ichannel].reserve(ntokens_pre_allocated_);
  }
  h_main_q_end_lane_offsets_.resize(nlanes_ + 1);
  h_emitting_main_q_end_lane_offsets_.resize(nlanes_ + 1);
  h_n_extra_prev_tokens_lane_offsets_.resize(nlanes_ + 1);
  frame_offsets_.resize(nchannels_);
  num_frames_decoded_.resize(nchannels_, -1);
  lanes2channels_todo_.reserve(nlanes_);

  h_all_argmin_cost_.resize(nchannels_, {-1, 0.0f});
  h_all_final_tokens_list_.resize(nchannels_);
  h_all_has_reached_final_.resize(nchannels_);
}

void CudaDecoder::InitDeviceData() {
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemsetAsync(
      d_channels_counters_.MutableData(), 0,
      nchannels_ * sizeof(*d_channels_counters_.MutableData()), compute_st_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemsetAsync(
      d_lanes_counters_.MutableData(), 0,
      nlanes_ * sizeof(*d_lanes_counters_.MutableData()), compute_st_));
  InitHashmapKernel(KaldiCudaDecoderNumBlocks(hashmap_capacity_, nlanes_),
                    KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,
                    *h_device_params_);
  KALDI_DECODER_CUDA_CHECK_ERROR();
}

void CudaDecoder::InitHostData() {}

void CudaDecoder::AllocateDeviceKernelParams() {
  h_device_params_ = new DeviceParams();
  h_kernel_params_ = new KernelParams();
}

void CudaDecoder::InitDeviceParams() {
  // Setting Kernel Params
  // Sent to cuda kernels by copy
  // Making sure we'll be able to send it to the kernels
  KALDI_ASSERT((sizeof(KernelParams) + sizeof(DeviceParams)) <
               KALDI_CUDA_DECODER_MAX_KERNEL_ARGUMENTS_BYTE_SIZE);

  h_device_params_->d_channels_counters = d_channels_counters_.GetView();
  h_device_params_->d_lanes_counters = d_lanes_counters_.GetView();
  h_device_params_->h_lanes_counters = h_lanes_counters_.GetView();
  h_device_params_->d_main_q_state_and_cost =
      d_main_q_state_and_cost_.GetView();
  h_device_params_->d_main_q_info = d_main_q_info_.GetView();
  h_device_params_->d_aux_q_state_and_cost = d_aux_q_state_and_cost_.GetView();
  h_device_params_->d_main_q_extra_and_acoustic_cost =
      d_main_q_extra_and_acoustic_cost_.GetView();
  h_device_params_->d_main_q_acoustic_cost = d_main_q_acoustic_cost_.GetView();
  h_device_params_->d_aux_q_info = d_aux_q_info_.GetView();
  h_device_params_->d_main_q_degrees_prefix_sum =
      d_main_q_degrees_prefix_sum_.GetView();
  h_device_params_->d_main_q_block_sums_prefix_sum =
      d_main_q_block_sums_prefix_sum_.GetView();
  h_device_params_->d_main_q_state_hash_idx =
      d_main_q_state_hash_idx_.GetView();
  h_device_params_->d_main_q_extra_prev_tokens_prefix_sum =
      d_main_q_extra_prev_tokens_prefix_sum_.GetView();
  h_device_params_->d_main_q_n_extra_prev_tokens_local_idx =
      d_main_q_n_extra_prev_tokens_local_idx_.GetView();
  h_device_params_->d_main_q_extra_prev_tokens =
      d_main_q_extra_prev_tokens_.GetView();
  h_device_params_->d_main_q_arc_offsets = d_main_q_arc_offsets_.GetView();
  h_device_params_->d_hashmap_values = d_hashmap_values_.GetView();
  h_device_params_->d_histograms = d_histograms_.GetView();
  h_device_params_->d_arc_e_offsets = fst_.d_e_offsets_;
  h_device_params_->d_arc_ne_offsets = fst_.d_ne_offsets_;
  h_device_params_->d_arc_pdf_ilabels = fst_.d_arc_pdf_ilabels_;
  h_device_params_->d_arc_weights = fst_.d_arc_weights_;
  h_device_params_->d_arc_nextstates = fst_.d_arc_nextstates_;
  h_device_params_->d_fst_final_costs = fst_.d_final_;
  h_device_params_->default_beam = default_beam_;
  h_device_params_->lattice_beam = lattice_beam_;
  h_device_params_->main_q_capacity = main_q_capacity_;
  h_device_params_->aux_q_capacity = aux_q_capacity_;
  h_device_params_->init_channel_id = init_channel_id_;
  h_device_params_->max_nlanes = nlanes_;
  h_device_params_->nstates = fst_.num_states_;
  h_device_params_->init_state = fst_.Start();
  KALDI_ASSERT(h_device_params_->init_state != fst::kNoStateId);
  h_device_params_->init_cost = StdWeight::One().Value();
  h_device_params_->hashmap_capacity = hashmap_capacity_;
  h_device_params_->max_active = max_active_;
  // For the first static_beam_q_length elements of the queue, we will keep the
  // beam static
  adaptive_beam_static_segment_ =
      aux_q_capacity_ / KALDI_CUDA_DECODER_ADAPTIVE_BEAM_STATIC_SEGMENT;
  // For the last adaptive_beam_q_length elements of the queue, we will decrease
  // the beam, segment by segment
  // For more information, please refer to the definition of GetAdaptiveBeam in
  // cuda-decoder-kernels.cu
  int32 adaptive_beam_q_length =
      (aux_q_capacity_ - adaptive_beam_static_segment_);
  int32 adaptive_beam_bin_width =
      adaptive_beam_q_length / KALDI_CUDA_DECODER_ADAPTIVE_BEAM_NSTEPS;
  h_device_params_->adaptive_beam_static_segment =
      adaptive_beam_static_segment_;
  h_device_params_->adaptive_beam_bin_width = adaptive_beam_bin_width;

  // Reusing aux_q memory to list final states in GetLattice
  // Those cannot be used at the same time
  h_device_params_->h_list_final_tokens_in_main_q =
      h_list_final_tokens_in_main_q_.GetView();
}

CudaDecoder::~CudaDecoder() {
  // Stopping h2h tasks
  h2h_threads_running_ = false;
  n_h2h_main_task_todo_cv_.notify_all();
  for (std::thread &thread : cpu_dedicated_threads_) thread.join();
  cudaStreamDestroy(compute_st_);
  cudaStreamDestroy(copy_st_);

  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaFreeHost(h_channels_counters_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaFreeHost(h_extra_and_acoustic_cost_concat_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaFreeHost(h_acoustic_cost_concat_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaFreeHost(h_extra_prev_tokens_concat_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaFreeHost(h_infotoken_concat_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaFreeHost(h_extra_and_acoustic_cost_concat_tmp_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaFreeHost(h_acoustic_cost_concat_tmp_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaFreeHost(h_extra_prev_tokens_concat_tmp_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaFreeHost(h_infotoken_concat_tmp_));
  // Will call the cudaFrees inside destructors
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventDestroy(nnet3_done_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventDestroy(d2h_copy_acoustic_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventDestroy(d2h_copy_infotoken_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaEventDestroy(d2h_copy_extra_prev_tokens_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(
      cudaEventDestroy(concatenated_data_ready_evt_));
  KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventDestroy(lane_offsets_ready_evt_));

  delete h_kernel_params_;
  delete h_device_params_;
}

void CudaDecoder::ComputeInitialChannel() {
  KALDI_ASSERT(nlanes_ > 0);
  const int32 ilane = 0;
  KALDI_ASSERT(ilane == 0);
  // Following kernels working channel_id
  std::vector<ChannelId> channels = {init_channel_id_};
  SetChannelsInKernelParams(channels);  // not calling LoadChannelsStateToLanes,
                                        // init_channel_id_ is a special case
  h_lanes_counters_.lane(ilane)->channel_to_compute = init_channel_id_;

  cudaMemcpyAsync(d_lanes_counters_.MutableData(), h_lanes_counters_.lane(0),
                  1 * sizeof(*h_lanes_counters_.lane(0)),
                  cudaMemcpyHostToDevice, compute_st_);
  h_lanes_counters_.lane(ilane)->main_q_narcs_and_end.y = 0;

  // Adding the start state to the initial token queue
  InitializeInitialLaneKernel(KaldiCudaDecoderNumBlocks(1, 1),
                              KALDI_CUDA_DECODER_ONE_THREAD_BLOCK, compute_st_,
                              *h_device_params_);

  h_lanes_counters_.lane(ilane)->post_expand_aux_q_end = 1;

  PruneAndPreprocess();
  FinalizeProcessNonEmittingKernel(
      KaldiCudaDecoderNumBlocks(1, 1), KALDI_CUDA_DECODER_LARGEST_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_);

  CopyLaneCountersToHostSync();
  PostProcessingMainQueue();
  CopyLaneCountersToHostSync();

  const int32 main_q_end =
      h_lanes_counters_.lane(ilane)->main_q_narcs_and_end.y;
  KALDI_ASSERT(main_q_end > 0);

  // Moving all data linked to init_channel_id_ to host
  // that data will be cloned to other channels when calling InitDecoding
  CopyMainQueueDataToHost();
  SaveChannelsStateFromLanes();

  KALDI_ASSERT(
      h_channels_counters_[init_channel_id_].prev_main_q_narcs_and_end.x > 0);
  KALDI_ASSERT(
      h_channels_counters_[init_channel_id_].prev_main_q_narcs_and_end.y > 0);
}

void CudaDecoder::InitDecoding(const std::vector<ChannelId> &channels) {
  // Cloning the init_channel_id_ channel into all channels in the channels vec
  const int nlanes_used = channels.size();
  // Getting *h_kernel_params ready to use
  LoadChannelsStateToLanes(channels);
  cudaMemcpyAsync(d_lanes_counters_.MutableData(), h_lanes_counters_.lane(0),
                  nlanes_used_ * sizeof(*h_lanes_counters_.lane(0)),
                  cudaMemcpyHostToDevice, compute_st_);

  // Size of the initial main_q
  ChannelCounters &init_channel_counters =
      h_channels_counters_[init_channel_id_];
  const int32 init_main_q_size =
      init_channel_counters.prev_main_q_narcs_and_end.y;

  KALDI_ASSERT(init_main_q_size > 0);
  // Getting the channels ready to compute new utterances
  InitDecodingOnDeviceKernel(
      KaldiCudaDecoderNumBlocks(init_main_q_size, nlanes_used),
      KALDI_CUDA_DECODER_1D_BLOCK, compute_st_, *h_device_params_,
      *h_kernel_params_);

  {
    std::lock_guard<std::mutex> n_h2h_not_done_lk(
        n_init_decoding_h2h_task_not_done_mutex_);
    n_init_decoding_h2h_task_not_done_ += channels.size();
  }
  for (ChannelId ichannel : channels) {
    ChannelCounters &channel_counters = h_channels_counters_[ichannel];
    channel_counters.prev_main_q_narcs_and_end =
        init_channel_counters.prev_main_q_narcs_and_end;
    channel_counters.prev_main_q_n_extra_prev_tokens =
        init_channel_counters.prev_main_q_n_extra_prev_tokens;
    channel_counters.prev_main_q_global_offset = 0;
    channel_counters.prev_main_q_extra_prev_tokens_global_offset = 0;
    channel_counters.prev_beam = default_beam_;

    int32 n_initial_tokens = h_all_tokens_info_[init_channel_id_].size();
    num_frames_decoded_[ichannel] = 0;
    h_channels_counters_[ichannel] = h_channels_counters_[init_channel_id_];
    h_all_argmin_cost_[ichannel] = {-1, 0.0f};
    frame_offsets_[ichannel].clear();
    frame_offsets_[ichannel].push_back(n_initial_tokens);
    if (thread_pool_)
      thread_pool_->enqueue(THREAD_POOL_HIGH_PRIORITY,
                            &CudaDecoder::InitDecodingH2HCopies, this,
                            ichannel);
    else
      InitDecodingH2HCopies(ichannel);
  }
}

void CudaDecoder::InitDecodingH2HCopies(ChannelId ichannel) {
  // Tokens from initial main_q needed on host
  std::unique_lock<std::mutex> channel_lk(channel_lock_[ichannel]);
  // Deep copy
  h_all_tokens_info_[ichannel] = h_all_tokens_info_[init_channel_id_];
  h_all_tokens_acoustic_cost_[ichannel] =
      h_all_tokens_acoustic_cost_[init_channel_id_];
  h_all_tokens_extra_prev_tokens_[ichannel] =
      h_all_tokens_extra_prev_tokens_[init_channel_id_];
  h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_[ichannel] =
      h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_[init_channel_id_];

  bool all_done;
  {
    std::lock_guard<std::mutex> lk_not_done(
        n_init_decoding_h2h_task_not_done_mutex_);
    all_done = (--n_init_decoding_h2h_task_not_done_ == 0);
  }
  if (all_done) {
    init_decoding_h2h_done_.notify_all();
  }
}

void CudaDecoder::LoadChannelsStateToLanes(
    const std::vector<ChannelId> &channels) {
  // Setting that channels configuration in kernel_params
  SetChannelsInKernelParams(channels);
  KALDI_ASSERT(nlanes_used_ > 0);
  for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {
    const ChannelId ichannel = channel_to_compute_[ilane];
    ChannelCounters &channel_counters = h_channels_counters_[ichannel];
    LaneCounters &lane_counters = *h_lanes_counters_.lane(ilane);
    lane_counters.channel_to_compute = ichannel;
    lane_counters.main_q_narcs_and_end =
        channel_counters.prev_main_q_narcs_and_end;
    lane_counters.main_q_n_extra_prev_tokens =
        channel_counters.prev_main_q_n_extra_prev_tokens;
    int32 int_beam = floatToOrderedIntHost(channel_counters.prev_beam);
    lane_counters.int_beam = int_beam;
    lane_counters.adaptive_int_beam_with_validity_index.x = int_beam;
    lane_counters.adaptive_int_beam_with_validity_index.y =
        adaptive_beam_static_segment_;
    lane_counters.main_q_global_offset =
        channel_counters.prev_main_q_global_offset;
    lane_counters.main_q_extra_prev_tokens_global_offset =
        channel_counters.prev_main_q_extra_prev_tokens_global_offset;

    lane_counters.min_int_cost =
        channel_counters.min_int_cost_and_arg_without_final.x;
    lane_counters.prev_arg_min_int_cost =
        channel_counters.min_int_cost_and_arg_without_final.y;
  }
}

void CudaDecoder::SaveChannelsStateFromLanes() {
  KALDI_ASSERT(nlanes_used_ > 0);
  for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {
    const ChannelId ichannel = channel_to_compute_[ilane];
    ChannelCounters &channel_counters = h_channels_counters_[ichannel];
    LaneCounters &lane_counters = *h_lanes_counters_.lane(ilane);
    channel_counters.prev_main_q_narcs_and_end =
        lane_counters.main_q_narcs_and_end;
    channel_counters.prev_main_q_extra_prev_tokens_global_offset =
        lane_counters.main_q_extra_prev_tokens_global_offset;
    channel_counters.prev_main_q_global_offset =
        lane_counters.main_q_global_offset;
    channel_counters.prev_main_q_n_extra_prev_tokens =
        lane_counters.main_q_n_extra_prev_tokens;
    channel_counters.prev_beam = orderedIntToFloatHost(lane_counters.int_beam);
    channel_counters.min_int_cost_and_arg_without_final = {
        lane_counters.min_int_cost, lane_counters.prev_arg_min_int_cost};
  }
  SaveChannelsStateFromLanesKernel(KaldiCudaDecoderNumBlocks(1, nlanes_used_),
                                   KALDI_CUDA_DECODER_ONE_THREAD_BLOCK,
                                   compute_st_, *h_device_params_,
                                   *h_kernel_params_);

  ResetChannelsInKernelParams();
}

int32 CudaDecoder::GetMaxForAllLanes(
    std::function<int32(const LaneCounters &)> func) {
  int32 max_val = 0;
  for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {
    const int32 val = func(*h_lanes_counters_.lane(ilane));
    max_val = std::max(max_val, val);
  }
  return max_val;
}

void CudaDecoder::CopyLaneCountersToHostAsync() {
  cudaMemcpyAsync(h_lanes_counters_.lane(0), d_lanes_counters_.MutableData(),
                  nlanes_used_ * sizeof(*h_lanes_counters_.lane(0)),
                  cudaMemcpyDeviceToHost, compute_st_);
}

void CudaDecoder::CopyLaneCountersToHostSync() {
  CopyLaneCountersToHostAsync();
  cudaStreamSynchronize(compute_st_);
}

// One sync has to happen between PerformConcatenatedCopy and
// MoveConcatenatedCopyToVector
template <typename T>
void CudaDecoder::MoveConcatenatedCopyToVector(
    const int32 ilane, const int32 ichannel,
    const std::vector<int32> &lanes_offsets, T *h_concat,
    std::vector<std::vector<T>> *vecvec) {
  // Unpacking the concatenated vector into individual channel storage
  int32 beg = lanes_offsets[ilane];
  int32 end = lanes_offsets[ilane + 1];
  auto &vec = (*vecvec)[ichannel];
  vec.insert(vec.end(), h_concat + beg, h_concat + end);
}

void CudaDecoder::ApplyMaxActiveAndReduceBeam(enum QUEUE_ID queue_id) {
  // Checking if we should activate max active for the current frame
  // once it is active, it is active for the whole frame (for all non emitting
  // iterations)
  // If at least one lane queue is bigger than max_active,
  // we'll apply a topk on that queue (k=max_active_)
  bool use_aux_q = (queue_id == AUX_Q);
  ComputeCostsHistogramKernel(KaldiCudaDecoderNumBlocks(nlanes_used_),
                              KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,
                              *h_device_params_, *h_kernel_params_, use_aux_q);

  UpdateBeamUsingHistogramKernel(
      KaldiCudaDecoderNumBlocks(1, nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_, use_aux_q);
}

int32 CudaDecoder::NumFramesToDecode(
    const std::vector<ChannelId> &channels,
    std::vector<CudaDecodableInterface *> &decodables, int32 max_num_frames) {
  int32 nframes_to_decode = INT_MAX;
  // std::vector<int> debug_ntokens;
  // std::vector<int> debug_narcs;
  for (int32 ilane = 0; ilane < nlanes_used_; ++ilane) {
    const ChannelId ichannel = channels[ilane];
    const int32 num_frames_decoded = num_frames_decoded_[ichannel];
    KALDI_ASSERT(num_frames_decoded >= 0 &&
                 "You must call InitDecoding() before AdvanceDecoding()");
    int32 num_frames_ready = decodables[ilane]->NumFramesReady();
    // num_frames_ready must be >= num_frames_decoded, or else
    // the number of frames ready must have decreased (which doesn't
    // make sense) or the decodable object changed between calls
    // (which isn't allowed).
    KALDI_ASSERT(num_frames_ready >= num_frames_decoded);
    int32 channel_nframes_to_decode = num_frames_ready - num_frames_decoded;
    nframes_to_decode = std::min(nframes_to_decode, channel_nframes_to_decode);
  }
  if (max_num_frames >= 0)
    nframes_to_decode = std::min(nframes_to_decode, max_num_frames);

  return nframes_to_decode;
}

void CudaDecoder::ExpandArcsEmitting() {
  ExpandArcsKernel<true>(KaldiCudaDecoderNumBlocks(nlanes_used_),
                         KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,
                         *h_device_params_, *h_kernel_params_);

  // Updating a few counters, like resetting aux_q_end to 0...
  // true is for IS_EMITTING
  PostExpandKernel<true>(KaldiCudaDecoderNumBlocks(1, nlanes_used_),
                         KALDI_CUDA_DECODER_ONE_THREAD_BLOCK, compute_st_,
                         *h_device_params_, *h_kernel_params_);
}

void CudaDecoder::ExpandArcsNonEmitting() {
  // false is for non emitting
  ExpandArcsKernel<false>(KaldiCudaDecoderNumBlocks(nlanes_used_),
                          KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,
                          *h_device_params_, *h_kernel_params_);

  // false is for non emitting
  PostExpandKernel<false>(KaldiCudaDecoderNumBlocks(1, nlanes_used_),
                          KALDI_CUDA_DECODER_ONE_THREAD_BLOCK, compute_st_,
                          *h_device_params_, *h_kernel_params_);
}

void CudaDecoder::PruneAndPreprocess() {
  NonEmittingPreprocessAndContractKernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_);
  PostContractAndPreprocessKernel(KaldiCudaDecoderNumBlocks(1, nlanes_used_),
                                  KALDI_CUDA_DECODER_ONE_THREAD_BLOCK,
                                  compute_st_, *h_device_params_,
                                  *h_kernel_params_);
}

void CudaDecoder::PostProcessingMainQueue() {
  ApplyMaxActiveAndReduceBeam(MAIN_Q);

  FillHashmapWithMainQKernel(KaldiCudaDecoderNumBlocks(nlanes_used_),
                             KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,
                             *h_device_params_, *h_kernel_params_);

  EmittingPreprocessAndListExtraPrevTokensStep1Kernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_);

  EmittingPreprocessAndListExtraPrevTokensStep2Kernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_);

  // Step2 wrote main_q_n_extra_prev_tokens
  // it was the last value missing to compute the lanes offsets
  // doing it now
  ComputeLaneOffsetsKernel(KaldiCudaDecoderNumBlocks(1, 1),  // One CTA
                           KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,
                           *h_device_params_, *h_kernel_params_);
  cudaEventRecord(lane_offsets_ready_evt_, compute_st_);

  EmittingPreprocessAndListExtraPrevTokensStep3Kernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_);

  EmittingPreprocessAndListExtraPrevTokensStep4Kernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_);

  ClearHashmapKernel(KaldiCudaDecoderNumBlocks(nlanes_used_),
                     KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,
                     *h_device_params_, *h_kernel_params_);
}

void CudaDecoder::CopyMainQueueDataToHost() {
  cudaEventRecord(concatenated_data_ready_evt_, compute_st_);
  cudaStreamWaitEvent(copy_st_, concatenated_data_ready_evt_,
                      0);  // the copies on copy_st will wait on compute_st_
  cudaEventSynchronize(
      lane_offsets_ready_evt_);  // we need the total size of each segments
  LaunchD2HCopies();

  // Making sure the previous H2H copies are done
  WaitForInitDecodingH2HCopies();
  WaitForH2HCopies();

  std::swap(h_extra_and_acoustic_cost_concat_tmp_,
            h_extra_and_acoustic_cost_concat_);
  std::swap(h_infotoken_concat_tmp_, h_infotoken_concat_);
  std::swap(h_acoustic_cost_concat_tmp_, h_acoustic_cost_concat_);
  std::swap(h_extra_prev_tokens_concat_tmp_, h_extra_prev_tokens_concat_);
  // Saving the offsets computed previously
  lanes2channels_todo_.clear();
  for (int32 ilane = 0; ilane < (nlanes_used_ + 1); ++ilane) {
    h_emitting_main_q_end_lane_offsets_[ilane] =
        h_lanes_counters_.lane(ilane)->main_q_n_emitting_tokens_lane_offset;
    h_main_q_end_lane_offsets_[ilane] =
        h_lanes_counters_.lane(ilane)->main_q_end_lane_offset;
    h_n_extra_prev_tokens_lane_offsets_[ilane] =
        h_lanes_counters_.lane(ilane)->main_q_n_extra_prev_tokens_lane_offset;
    lanes2channels_todo_.push_back(channel_to_compute_[ilane]);
  }

  LaunchH2HCopies();
}

void CudaDecoder::LaunchD2HCopies() {
  // Last offset = total
  int32 nelements_acoustic_costs = h_lanes_counters_.lane(nlanes_used_)
                                       ->main_q_n_emitting_tokens_lane_offset;
  // Moving the d_concat to h_concat (host), async
  if (nelements_acoustic_costs > 0) {
    KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpyAsync(
        h_acoustic_cost_concat_tmp_, d_acoustic_cost_concat_,
        nelements_acoustic_costs * sizeof(*d_acoustic_cost_concat_),
        cudaMemcpyDeviceToHost, copy_st_));
  }
  cudaEventRecord(d2h_copy_acoustic_evt_, copy_st_);

  int32 nelements_infotoken =
      h_lanes_counters_.lane(nlanes_used_)->main_q_end_lane_offset;
  if (nelements_infotoken > 0) {
    KALDI_DECODER_CUDA_API_CHECK_ERROR(
        cudaMemcpyAsync(h_infotoken_concat_tmp_, d_infotoken_concat_,
                        nelements_infotoken * sizeof(*d_infotoken_concat_),
                        cudaMemcpyDeviceToHost, copy_st_));
  }
  cudaEventRecord(d2h_copy_infotoken_evt_, copy_st_);
  int32 nelements_extra_prev_tokens =
      h_lanes_counters_.lane(nlanes_used_)
          ->main_q_n_extra_prev_tokens_lane_offset;
  if (nelements_extra_prev_tokens > 0) {
    KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpyAsync(
        h_extra_prev_tokens_concat_tmp_, d_extra_prev_tokens_concat_,
        nelements_extra_prev_tokens * sizeof(*d_extra_prev_tokens_concat_),
        cudaMemcpyDeviceToHost, copy_st_));
    KALDI_DECODER_CUDA_API_CHECK_ERROR(
        cudaMemcpyAsync(h_extra_and_acoustic_cost_concat_tmp_,
                        d_extra_and_acoustic_cost_concat_,
                        nelements_extra_prev_tokens *
                            sizeof(*d_extra_and_acoustic_cost_concat_),
                        cudaMemcpyDeviceToHost, copy_st_));
  }
  cudaEventRecord(d2h_copy_extra_prev_tokens_evt_, copy_st_);
}

void CudaDecoder::ConcatenateData() {
  ConcatenateLanesDataKernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_,
      h_device_params_->d_main_q_acoustic_cost, d_acoustic_cost_concat_,
      &d_lanes_counters_.lane(0)->main_q_n_emitting_tokens_lane_offset);
  ConcatenateLanesDataKernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_,
      h_device_params_->d_main_q_info, d_infotoken_concat_,
      &d_lanes_counters_.lane(0)->main_q_end_lane_offset);
  ConcatenateLanesDataKernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_,
      h_device_params_->d_main_q_extra_prev_tokens, d_extra_prev_tokens_concat_,
      &d_lanes_counters_.lane(0)->main_q_n_extra_prev_tokens_lane_offset);
  ConcatenateLanesDataKernel(
      KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
      compute_st_, *h_device_params_, *h_kernel_params_,
      h_device_params_->d_main_q_extra_and_acoustic_cost,
      d_extra_and_acoustic_cost_concat_,
      &d_lanes_counters_.lane(0)->main_q_n_extra_prev_tokens_lane_offset);
}

void CudaDecoder::AdvanceDecoding(
    const std::vector<ChannelId> &channels,
    std::vector<CudaDecodableInterface *> &decodables, int32 max_num_frames) {
  if (channels.size() == 0) return;  // nothing to do
  // Context switch : Loading the channels state in lanes
  LoadChannelsStateToLanes(channels);
  KALDI_ASSERT(nlanes_used_ > 0);

  // We'll decode nframes_to_decode, such as all channels have at least that
  // number
  // of frames available
  int32 nframes_to_decode =
      NumFramesToDecode(channels, decodables, max_num_frames);

  // Looping over the frames that we will compute
  for (int32 iframe = 0; iframe < nframes_to_decode; ++iframe) {
    // Loglikelihoods from the acoustic model
    // Setting the loglikelihoods pointers for that frame
    for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {
      ChannelId ichannel = channel_to_compute_[ilane];
      int32 frame = num_frames_decoded_[ichannel];
      h_lanes_counters_.lane(ilane)->loglikelihoods =
          decodables[ilane]->GetLogLikelihoodsCudaPointer(frame);
    }
    cudaMemcpyAsync(d_lanes_counters_.MutableData(), h_lanes_counters_.lane(0),
                    nlanes_used_ * sizeof(*h_lanes_counters_.lane(0)),
                    cudaMemcpyHostToDevice, compute_st_);
    // compute_st_ will wait for nnet3 to complete
    cudaEventRecord(nnet3_done_evt_, cudaStreamPerThread);
    cudaStreamWaitEvent(compute_st_, nnet3_done_evt_, 0);

    // Estimating cutoff using argmin from last frame
    ResetForFrameAndEstimateCutoffKernel(
        KaldiCudaDecoderNumBlocks(1, nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,
        compute_st_, *h_device_params_, *h_kernel_params_);
    // Reset max active status. If necessary, ApplyMaxActiveAndReduceBeam will
    // switch it back on
    compute_max_active_ = false;

    // Processing emitting arcs. We've done the preprocess stage at the end of
    // the previous frame
    ExpandArcsEmitting();
    // We'll loop until we have a small enough number of non-emitting arcs
    // in the token queue. We'll then break the loop
    for (int i = 0; i < KALDI_CUDA_DECODER_N_NON_EMITTING_MAIN_ITERATIONS;
         ++i) {
      // If one of the aux_q contains more than max_active_ tokens,
      // we'll reduce the beam to only keep max_active_ tokens
      ApplyMaxActiveAndReduceBeam(AUX_Q);
      // Prune the aux_q. Apply the latest beam (using the one from
      // ApplyMaxActiveAndReduceBeam if triggered)
      // move the survival tokens to the main queue
      // and do the preprocessing necessary for the next ExpandArcs
      PruneAndPreprocess();

      // "heavy duty" kernel for non-emitting. The long tail of small
      // non-emitting iterations will be done in
      // FinalizeProcessNonEmittingKernel
      ExpandArcsNonEmitting();
    }
    ApplyMaxActiveAndReduceBeam(AUX_Q);
    PruneAndPreprocess();
    // Finalizing process non emitting. Takes care of the long tail,
    // the final iterations with a small numbers of arcs. Do the work inside a
    // single CTA (per lane),
    FinalizeProcessNonEmittingKernel(KaldiCudaDecoderNumBlocks(1, nlanes_used_),
                                     KALDI_CUDA_DECODER_LARGEST_1D_BLOCK,
                                     compute_st_, *h_device_params_,
                                     *h_kernel_params_);

    // We now have our final token main queues for that frame

    // Post processing the tokens for that frame
    // - do the preprocess necessary for the next emitting expand (will happen
    // with next frame)
    // - if a state S has more than one token associated to it, generate the
    // list of those tokens
    // It allows to backtrack efficiently in GetRawLattice
    // - compute the extra costs
    PostProcessingMainQueue();

    // Waiting on previous d2h before writing on same device memory
    cudaStreamWaitEvent(compute_st_, d2h_copy_extra_prev_tokens_evt_, 0);
    // Concatenating the data that will be moved to host into large arrays
    ConcatenateData();
    // Copying the final lane counters for that frame
    CopyLaneCountersToHostSync();
    CheckOverflow();

    // Moving the data necessary for GetRawLattice/GetBestPath back to host for
    // storage
    CopyMainQueueDataToHost();

    for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {
      const ChannelId ichannel = channel_to_compute_[ilane];
      // We're done processing that frame
      ++num_frames_decoded_[ichannel];
      const int32 main_q_end =
          h_lanes_counters_.lane(ilane)->main_q_narcs_and_end.y;
      // Saving frame offsets for GetRawLattice
      frame_offsets_[ichannel].push_back(frame_offsets_[ichannel].back() +
                                         main_q_end);
    }
  }

  SaveChannelsStateFromLanes();
}

void CudaDecoder::CheckOverflow() {
  for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {
    LaneCounters *lane_counters = h_lanes_counters_.lane(ilane);
    bool q_overflow = lane_counters->q_overflow;
    if (q_overflow != OVERFLOW_NONE) {
      // An overflow was prevented in a kernel
      // The algorithm can still go on but quality of the result can be reduced
      // (less tokens were generated)

      if ((q_overflow & OVERFLOW_MAIN_Q) == OVERFLOW_MAIN_Q) {
        // overflowed main_q
        KALDI_WARN
            << "Preventing overflow of main_q. Continuing "
            << "execution but the quality of the output may be decreased. "
            << "To prevent this from happening, please increase the parameter "
               "--main-q-capacity"
            << " and/or decrease --max-active";
      }
      if ((q_overflow & OVERFLOW_AUX_Q) == OVERFLOW_AUX_Q) {
        // overflowed aux_q
        KALDI_WARN
            << "Preventing overflow of aux_q. Continuing "
            << "execution but the quality of the output may be decreased. "
            << "To prevent this from happening, please increase the parameter "
               "--aux-q-capacity"
            << " and/or decrease --beam";
      }

      KALDI_ASSERT(lane_counters->main_q_narcs_and_end.y < main_q_capacity_);
      KALDI_ASSERT(lane_counters->main_q_narcs_and_end.x >= 0);
      KALDI_ASSERT(lane_counters->main_q_narcs_and_end.y >= 0);
      KALDI_ASSERT(lane_counters->post_expand_aux_q_end < aux_q_capacity_);
      KALDI_ASSERT(lane_counters->post_expand_aux_q_end >= 0);
      KALDI_ASSERT(lane_counters->aux_q_end < aux_q_capacity_);
      KALDI_ASSERT(lane_counters->aux_q_end >= 0);
    }
  }
}

// GetBestCost
// returns the minimum cost among all tokens cost in the current frame
// also returns the index of one token with that min cost
//
// Only called at the end of the computation of one audio file
// not optimized
void CudaDecoder::GetBestCost(const std::vector<ChannelId> &channels,
                              bool use_final_costs,
                              std::vector<std::pair<int32, CostType>> *argmins,
                              std::vector<std::vector<std::pair<int, float>>>
                                  *list_finals_token_idx_and_cost,
                              std::vector<bool> *has_reached_final) {
  if (channels.size() == 0) return;
  // Getting the lanes ready to be used with those channels
  LoadChannelsStateToLanes(channels);
  cudaMemcpyAsync(d_lanes_counters_.MutableData(), h_lanes_counters_.lane(0),
                  nlanes_used_ * sizeof(*h_lanes_counters_.lane(0)),
                  cudaMemcpyHostToDevice, compute_st_);

  auto func_main_q_end = [](const LaneCounters &c) {
    return c.main_q_narcs_and_end.y;
  };
  int32 max_main_q_end = GetMaxForAllLanes(func_main_q_end);

  // Step1 : Finding the best cost in the last token queue, with and without
  // final costs.
  // Also saving the indexes of those min.
  GetBestCostStep1Kernel(
      KaldiCudaDecoderNumBlocks(max_main_q_end, nlanes_used_),
      KALDI_CUDA_DECODER_1D_BLOCK, compute_st_, *h_device_params_,
      *h_kernel_params_, use_final_costs, StdWeight::Zero().Value());

  // Step2: Now that we now what the minimum cost is, we list all tokens within
  // [min_cost; min_cost+lattice_beam]
  // min_cost takes into account the final costs if use_final_costs is true,
  // AND if a final state is is present in the last token queue
  GetBestCostStep2Kernel(
      KaldiCudaDecoderNumBlocks(max_main_q_end, nlanes_used_),
      KALDI_CUDA_DECODER_1D_BLOCK, compute_st_, *h_device_params_,
      *h_kernel_params_, use_final_costs, StdWeight::Zero().Value());

  // Step3 : Moves some data to host. We are moving the data that couldn't be
  // moved
  // directly in step 2, e.g. results of atomics (we don't know which one is
  // last)
  GetBestCostStep3Kernel(
      KaldiCudaDecoderNumBlocks(max_main_q_end, nlanes_used_),
      KALDI_CUDA_DECODER_1D_BLOCK, compute_st_, *h_device_params_,
      *h_kernel_params_);

  // Resetting the datastructures
  argmins->clear();
  has_reached_final->clear();
  list_finals_token_idx_and_cost->clear();
  // list_finals_token_idx_and_cost is a vector<vector<>>
  // Each channel will have its own list of tokens within [best;
  // best+lattice_beam]
  list_finals_token_idx_and_cost->resize(nlanes_used_);
  // Waiting for the copy
  cudaStreamSynchronize(compute_st_);
  for (int32 ilane = 0; ilane < nlanes_used_; ++ilane) {
    int2 minarg = h_lanes_counters_.lane(ilane)->min_int_cost_and_arg;
    // Min cost in that channel last token queue
    CostType min_cost = orderedIntToFloatHost(minarg.x);
    // index of that min cost
    int32 arg = minarg.y;
    // Saving both in output
    argmins->push_back({arg, min_cost});
    // Whether or not the last token queue contains at least one token
    // associated with a final FST state
    has_reached_final->push_back(
        h_lanes_counters_.lane(ilane)->has_reached_final);
    // Number of tokens within [min_cost; min_cost+lattice_beam]
    int n_within_lattice_beam =
        h_lanes_counters_.lane(ilane)->n_within_lattice_beam;
    // Loading those tokens
    (*list_finals_token_idx_and_cost)[ilane].resize(n_within_lattice_beam);
    // Moving to output + int2float conversion
    for (int i = 0; i < n_within_lattice_beam; ++i) {
      int global_idx = h_list_final_tokens_in_main_q_.lane(ilane)[i].x;
      float cost_with_final = orderedIntToFloatHost(
          h_list_final_tokens_in_main_q_.lane(ilane)[i].y);
      (*list_finals_token_idx_and_cost)[ilane][i].first = global_idx;
      (*list_finals_token_idx_and_cost)[ilane][i].second = cost_with_final;
    }
  }
}

void CudaDecoder::GetBestPath(const std::vector<ChannelId> &channels,
                              std::vector<Lattice *> &fst_out_vec,
                              bool use_final_probs) {
  KALDI_ASSERT(channels.size() == fst_out_vec.size());
  nvtxRangePushA("GetBestPath");
  GetBestCost(channels, use_final_probs, &argmins_,
              &list_finals_token_idx_and_cost_, &has_reached_final_);

  std::vector<int32> reversed_path;
  for (int32 ilane = 0; ilane < channels.size(); ++ilane) {
    const ChannelId ichannel = channels[ilane];
    const int32 token_with_best_cost = argmins_[ilane].first;
    std::unique_lock<std::mutex> channel_lk(channel_lock_[ichannel]);
    // If that token in that frame f is available, then all tokens in that frame
    // f are available
    WaitForH2HCopies();
    const bool isfinal = has_reached_final_[ilane];
    TokenId token_idx = token_with_best_cost;

    // Backtracking
    // Going all the way from the token with best cost
    // to the beginning (StartState)
    reversed_path.clear();

    // The first token was inserted at the beginning of the queue
    // it always has index 0
    // We backtrack until that first token
    while (token_idx != 0) {
      InfoToken token = h_all_tokens_info_[ichannel][token_idx];
      // We want an arc with extra_cost == 0
      int32 arc_idx;
      TokenId prev_token_idx;
      if (token.IsUniqueTokenForStateAndFrame()) {
        // If we have only one, it is an arc with extra_cost == 0
        arc_idx = token.arc_idx;
        prev_token_idx = token.prev_token;
      } else {
        // Using the first arc with extra_cost == 0
        int32 offset, size;
        std::tie(offset, size) = token.GetSameFSTStateTokensList();
        bool found_best = false;
        for (auto i = 0; i < size; ++i) {
          CostType arc_extra_cost =
              h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_[ichannel]
                                                                     [offset +
                                                                      i].x;
          // Picking one arc on the best path (extra_cost == 0)
          if (arc_extra_cost == 0.0f) {
            InfoToken list_token =
                h_all_tokens_extra_prev_tokens_[ichannel][offset + i];
            arc_idx = list_token.arc_idx;
            prev_token_idx = list_token.prev_token;
            found_best = true;
            break;
          }
        }
        KALDI_ASSERT(found_best);
      }
      reversed_path.push_back(arc_idx);
      token_idx = prev_token_idx;
    }

    Lattice *fst_out = fst_out_vec[ilane];
    fst_out->DeleteStates();
    // Building the output Lattice
    OutputLatticeState curr_state = fst_out->AddState();
    fst_out->SetStart(curr_state);

    for (int32 i = reversed_path.size() - 1; i >= 1; i--) {
      int32 arc_idx = reversed_path[i];

      LatticeArc arc(fst_.h_arc_id_ilabels_[arc_idx],
                     fst_.h_arc_olabels_[arc_idx],
                     LatticeWeight(fst_.h_arc_weights_[arc_idx], 0),
                     fst_.h_arc_nextstate_[arc_idx]);

      arc.nextstate = fst_out->AddState();
      fst_out->AddArc(curr_state, arc);
      curr_state = arc.nextstate;
    }

    // Adding final cost to final state
    if (isfinal && use_final_probs)
      fst_out->SetFinal(
          curr_state,
          LatticeWeight(fst_.h_final_[fst_.h_arc_nextstate_[reversed_path[0]]],
                        0.0));
    else
      fst_out->SetFinal(curr_state, LatticeWeight::One());

    fst::RemoveEpsLocal(fst_out);
  }
  nvtxRangePop();
}

void CudaDecoder::DebugValidateLattice() {
#if 0
	//validate lattice consistency
	for(int frame=0;frame<nframes;frame++) {
		int token_start=frame_offsets_[ichannel][frame];
		int token_end=(frame+1<nframes) ? frame_offsets_[ichannel][frame+1] : total_ntokens;
		int prev_frame_offset=(frame>0) ? frame_offsets_[ichannel][frame-1] : 0;
		int cur_frame_offset=token_start;
		int next_frame_offset=token_end;

		bool found_zero = false;
		//for each token in frame
		for(int i=token_start;i<token_end;i++) {
			if(i==0) continue;  //initial token skip this...
			InfoToken token=h_all_tokens_info_[ichannel][i];
			KALDI_ASSERT(token.prev_token>=0);

			if(token.IsUniqueTokenForStateAndFrame()) {
				//previous token must be lower than the next frame start
				KALDI_ASSERT(token.prev_token<next_frame_offset);
				//previous token must be larger then previous frame start
				KALDI_ASSERT(token.prev_token>=prev_frame_offset);
			} else {
				int32 offset, size;
				std::tie(offset,size) = token.GetNextStateTokensList();
				KALDI_ASSERT(size>0);
				KALDI_ASSERT(offset>=0 && offset<h_all_tokens_extra_prev_tokens_[ichannel].size());
				for(auto j=0; j<size; ++j) {
					KALDI_ASSERT(offset+j<h_all_tokens_extra_prev_tokens_[ichannel].size());
					InfoToken extra_token=h_all_tokens_extra_prev_tokens_[ichannel][offset+j];
					//previous token must be lower than the next frame start
					KALDI_ASSERT(extra_token.prev_token<next_frame_offset);
					//previous token must be larger then previous frame start
					KALDI_ASSERT(extra_token.prev_token>=prev_frame_offset);
				}
			}
		}
	}
#endif
}

CudaDecoder::LatticeStateInternalId CudaDecoder::GetLatticeStateInternalId(
    int32 total_ntokens, TokenId token_idx, InfoToken token) {
  // If we have a unique token for this (frame,fst_state)
  // Then its ID is a unique ID for (frame,fst_state)
  if (token.IsUniqueTokenForStateAndFrame()) return token_idx;

  // If we have multiple tokens for this (frame,fst_state),
  // let's use the "extra_prev_tokens" offset, which is unique for
  // (frame,fst_state) in that case

  // Adding the total_ntokens offset to avoid collisions with the previous
  // case
  return (total_ntokens + token.prev_token);
}

void CudaDecoder::AddFinalTokensToLattice(
    ChannelId ichannel,
    std::vector<std::pair<TokenId, InfoToken>> *q_curr_frame_todo,
    std::unordered_map<LatticeStateInternalId, RawLatticeState>
        *curr_f_raw_lattice_state,
    Lattice *fst_out) {
  // Total number of tokens for that utterance. Used in
  // GetLatticeStateInternalId
  const int32 total_ntokens = h_all_tokens_info_[ichannel].size();
  // Reading the overall best_cost for that utterance's last frame. Was set by
  // GetBestCost
  const CostType best_cost = h_all_argmin_cost_[ichannel].second;
  // Iterating through tokens associated with a final state in the last frame
  for (auto &p : h_all_final_tokens_list_[ichannel]) {
    // This final token has a final cost of final_token_cost
    CostType final_token_cost = p.second;
    // This token has possibly an extra cost compared to the best
    CostType extra_cost = final_token_cost - best_cost;
    // We only want to keep paths that have a cost within [best;
    // best+lattice_beam]
    if (extra_cost > lattice_beam_) {
      continue;
    }

    const TokenId final_token_idx = p.first;
    InfoToken final_token = h_all_tokens_info_[ichannel][final_token_idx];

    // Internal ID for our lattice_state=(iframe, fst_state)
    LatticeStateInternalId state_internal_id =
        GetLatticeStateInternalId(total_ntokens, final_token_idx, final_token);
    decltype(curr_f_raw_lattice_state->end()) map_it;
    bool inserted;

    // We need to create the fst_lattice_state linked to our internal id in the
    // lattice if it doesn't already exists
    // Inserts only if the key doesn't exist in the map
    std::tie(map_it, inserted) = curr_f_raw_lattice_state->insert(
        {state_internal_id, {FLT_MAX, -1, false}});

    // If we've inserted the element, it means that that state didn't exist in
    // the map
    // Because this is a final state, we need to do a bit of extra work to add
    // the final_cost to it
    if (inserted) {
      // We want to figure out which FST state this token is associated to
      // We don't have that info anymore, it wasn't transfered from the GPU
      // We still need it for final tokens, because we need to know which
      // final cost to add in the lattice.
      // To find that original FST state, we need the id of an arc going to
      // that state,
      // then we'll look in the graph and figure out next_state[arc_idx]
      // we just need a valid arc_idx
      int32 arc_idx;
      if (final_token.IsUniqueTokenForStateAndFrame()) {
        // If unique, we can directly use this arc_idx
        arc_idx = final_token.arc_idx;
      } else {
        // If we have multiple tokens associated to that fst state, just pick
        // the first one
        // from the list
        int32 offset, size;
        std::tie(offset, size) = final_token.GetSameFSTStateTokensList();
        InfoToken prev_token =
            h_all_tokens_extra_prev_tokens_[ichannel][offset];
        arc_idx = prev_token.arc_idx;
      }
      // Creating the state associated with our internal id in the lattice
      OutputLatticeState fst_lattice_final_state = fst_out->AddState();
      map_it->second.fst_lattice_state = fst_lattice_final_state;
      q_curr_frame_todo->push_back({final_token_idx, final_token});

      if (h_all_has_reached_final_[ichannel]) {
        // If we have reached final states, adding the final cost
        // We now have a valid arc_idx. We can read the FST state
        StateId fst_next_state = fst_.h_arc_nextstate_[arc_idx];

        fst_out->SetFinal(fst_lattice_final_state,
                          LatticeWeight(fst_.h_final_[fst_next_state], 0.0));
      } else {
        fst_out->SetFinal(fst_lattice_final_state, LatticeWeight::One());
      }
    }

    map_it->second.token_extra_cost =
        std::min(map_it->second.token_extra_cost, extra_cost);
  }
}

void CudaDecoder::AddArcToLattice(
    int32 list_arc_idx, TokenId list_prev_token_idx, InfoToken list_prev_token,
    int32 curr_frame_offset, CostType acoustic_cost,
    CostType this_arc_prev_token_extra_cost,
    LatticeStateInternalId src_state_internal_id,
    OutputLatticeState fst_lattice_start,
    OutputLatticeState to_fst_lattice_state,
    std::vector<std::pair<TokenId, InfoToken>> *q_curr_frame_todo,
    std::vector<std::pair<TokenId, InfoToken>> *q_prev_frame_todo,
    std::unordered_map<LatticeStateInternalId, RawLatticeState>
        *curr_f_raw_lattice_state,
    std::unordered_map<LatticeStateInternalId, RawLatticeState>
        *prev_f_raw_lattice_state,
    std::unordered_set<int32> *f_arc_idx_added, Lattice *fst_out,
    bool *must_replay_frame) {
  // We will now add this arc to the output lattice
  // We know the destination state of the arc (to_fst_lattice_state)
  // We need to figure out its source
  // And propagate the extra cost from the destination to the source of that arc
  // (we go backward)
  OutputLatticeState from_fst_lattice_state;
  // Having the predecessor in the previous frame
  // <=> that token is associated to an emiting arc
  bool emitting = (list_prev_token_idx < curr_frame_offset);
  // Checking if the source of that arc is the start state (original state at
  // the beginning of the decode)
  if (list_prev_token_idx != 0) {
    // Selecting the right map
    // - emitting arc -> previous frame map
    // - non emitting arc -> same frame map
    auto *extra_cost_map =
        emitting ? prev_f_raw_lattice_state : curr_f_raw_lattice_state;
    decltype(extra_cost_map->end()) from_map_it;
    bool inserted;
    // Attempting to insert the state in the map
    std::tie(from_map_it, inserted) =
        extra_cost_map->insert({src_state_internal_id, {FLT_MAX, -1, false}});
    // If it was inserted, its the first time we insert that key in
    // the map
    // we need to put that state in the todo list to be considered
    // next
    if (inserted) {
      auto *todo_list = emitting ? q_prev_frame_todo : q_curr_frame_todo;
      todo_list->push_back({list_prev_token_idx, list_prev_token});
      from_map_it->second.fst_lattice_state = fst_out->AddState();
    }

    // Updating the source extra cost using that arc
    // for an arc a->b
    // extra_cost(a) = min(extra_cost(a),
    //		extra_cost(b) + arc_extra_cost(a->b))
    CostType prev_token_extra_cost = from_map_it->second.token_extra_cost;
    if (this_arc_prev_token_extra_cost < prev_token_extra_cost) {
      // We found a new min
      CostType diff = (prev_token_extra_cost - this_arc_prev_token_extra_cost);
      // If the change is large enough,
      // and if the state that we're writing to was already closed,
      // then we need to replay that frame.
      // if the source state is already closed it means we've
      // read its extra_cost value. Now we're writing again to it.
      // We have to do the first read again, to get the updated
      // value
      // that's why we're replaying that frame
      // (between frames everything is in topological order)
      if (diff > extra_cost_min_delta_ && from_map_it->second.is_state_closed) {
        *must_replay_frame = true;
      }
      prev_token_extra_cost = this_arc_prev_token_extra_cost;
      from_map_it->second.token_extra_cost = prev_token_extra_cost;
    }

    // Reading the OutputLatticeState of the source state in the output lattice
    from_fst_lattice_state = from_map_it->second.fst_lattice_state;
  } else {
    from_fst_lattice_state =
        fst_lattice_start;  // we simply link it to the source
  }

  // Checking if it's the first time we insert an arc with that
  // arc_idx for that frame.
  // If we're replaying that frame, we don't want duplicates
  bool is_this_arc_new = f_arc_idx_added->insert(list_arc_idx).second;
  if (is_this_arc_new) {
    // The following reads will most likely end up in cache misses
    // we could load everything sooner
    LatticeArc arc(
        fst_.h_arc_id_ilabels_[list_arc_idx], fst_.h_arc_olabels_[list_arc_idx],
        LatticeWeight(fst_.h_arc_weights_[list_arc_idx], acoustic_cost),
        to_fst_lattice_state);
    fst_out->AddArc(from_fst_lattice_state, arc);
  }
}

void CudaDecoder::GetTokenRawLatticeData(
    TokenId token_idx, InfoToken token, int32 total_ntokens,
    std::unordered_map<LatticeStateInternalId, RawLatticeState>
        *curr_f_raw_lattice_state,
    CostType *token_extra_cost, OutputLatticeState *to_fst_lattice_state) {
  LatticeStateInternalId next_state_internal_id =
      GetLatticeStateInternalId(total_ntokens, token_idx, token);
  auto to_map_it = curr_f_raw_lattice_state->find(next_state_internal_id);
  // We know this token exists in the output lattice (because it's in
  // q_curr_frame_todo_)
  KALDI_ASSERT(to_map_it != curr_f_raw_lattice_state->end());

  *token_extra_cost = to_map_it->second.token_extra_cost;
  *to_fst_lattice_state = to_map_it->second.fst_lattice_state;

  // We read the extra cost from lattice_next_state
  // We are now closing the state. If we write to it again, we will have
  // to replay that frame
  // (so that the latest extra_cost value is read)
  to_map_it->second.is_state_closed = true;
}

void CudaDecoder::GetSameFSTStateTokenList(
    ChannelId ichannel, InfoToken token, InfoToken **tok_beg,
    float2 **extra_extra_and_acoustic_cost_beg, int32 *nsame) {
  // We now need to consider all tokens related to that (iframe,
  // fst_state)
  // with fst_state being the state this current token is linked to
  // There's two possibilies:
  // a) only one token is associated with that fst_state in that frame.
  // The necessary information
  // is then stored directly in the token (arc_idx, prev_token)
  // b) multiple tokens are associated with that fst_state in that
  // frame. The token that we have right now
  // only contains information on where to find the list of those
  // tokens. It contains (offset, size)
  //
  // In any cases we consider the list of tokens to process as an array
  // of InfoToken, which will
  // be of size 1 in case a), of size > 1 in case b)
  if (token.IsUniqueTokenForStateAndFrame()) {
    *tok_beg = &token;
    // if we've got only one, extra_cost == 0.0
    *extra_extra_and_acoustic_cost_beg = NULL;
    *nsame = 1;
  } else {
    int32 offset, size;
    std::tie(offset, size) = token.GetSameFSTStateTokensList();
    *tok_beg = &h_all_tokens_extra_prev_tokens_[ichannel][offset];
    *extra_extra_and_acoustic_cost_beg =
        &h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_[ichannel]
                                                                [offset];
    *nsame = size;
  }
}

void CudaDecoder::ConsiderTokenForLattice(
    ChannelId ichannel, int32 iprev, int32 total_ntokens, TokenId token_idx,
    OutputLatticeState fst_lattice_start, InfoToken *tok_beg,
    float2 *extra_extra_and_acoustic_cost_beg, CostType token_extra_cost,
    TokenId list_prev_token_idx, int32 list_arc_idx, InfoToken *list_prev_token,
    CostType *this_arc_prev_token_extra_cost, CostType *acoustic_cost,
    OutputLatticeState *lattice_src_state, bool *keep_arc,
    bool *dbg_found_zero) {
  CostType arc_extra_cost;
  if (extra_extra_and_acoustic_cost_beg) {
    float2 both = extra_extra_and_acoustic_cost_beg[iprev];
    arc_extra_cost = both.x;
    *acoustic_cost = both.y;
  } else {
    // If we have only one token for that (iframe,fst_state),
    // Its arc has an extra_cost of zero (it's the only way to
    // get to that state, so it's the best)
    arc_extra_cost = 0.0f;
    *acoustic_cost = h_all_tokens_acoustic_cost_[ichannel][token_idx];
  }
  // If we use that arc to go to prev_token, prev_token will have the
  // following extra cost
  *this_arc_prev_token_extra_cost = token_extra_cost + arc_extra_cost;
  // We need at least one arc_extra_cost of zero for each (iframe,
  // fst_state)
  // The only use for that boolean is in a KALDI_ASSERT,
  // because if something went wrong in the kernels it's not likely
  // that this property will be verified out of luck
  *dbg_found_zero |= (arc_extra_cost == 0.0f);
  *list_prev_token = h_all_tokens_info_[ichannel][list_prev_token_idx];
  // Source of the arc currently considered
  *lattice_src_state =
      (list_prev_token_idx != 0)
          ? GetLatticeStateInternalId(total_ntokens, list_prev_token_idx,
                                      *list_prev_token)
          : fst_lattice_start;

  // We only keep the arc if, when using that arc, we can end up
  // at the last frame with a cost not worse than (best+lattice_beam)
  // this_arc_prev_token_extra_cost contains the accumulated sums
  // of extra costs (through the cheapest possible way) to the last
  // frame
  *keep_arc = (*this_arc_prev_token_extra_cost < lattice_beam_);
}

void CudaDecoder::SwapPrevAndCurrLatticeMap(
    int32 iframe, bool dbg_found_best_path,
    std::vector<std::pair<TokenId, InfoToken>> *q_curr_frame_todo,
    std::vector<std::pair<TokenId, InfoToken>> *q_prev_frame_todo,
    std::unordered_map<LatticeStateInternalId, RawLatticeState>
        *curr_f_raw_lattice_state,
    std::unordered_map<LatticeStateInternalId, RawLatticeState>
        *prev_f_raw_lattice_state,
    std::unordered_set<int32> *f_arc_idx_added) {
  q_prev_frame_todo->swap(*q_curr_frame_todo);
  q_prev_frame_todo->clear();
  prev_f_raw_lattice_state->swap(*curr_f_raw_lattice_state);
  prev_f_raw_lattice_state->clear();
  f_arc_idx_added->clear();

  KALDI_ASSERT(q_prev_frame_todo->empty());
  if (iframe > 0) {
    KALDI_ASSERT(!q_curr_frame_todo->empty());
    if (!dbg_found_best_path) {
      KALDI_WARN << "Warning didn't find exact best path in GetRawLattice";
    }
  }
}

void CudaDecoder::WaitForH2HCopies() {
  std::unique_lock<std::mutex> lk(n_h2h_task_not_done_mutex_);
  h2h_done_.wait(lk, [this] { return (n_h2h_task_not_done_ == 0); });
}

void CudaDecoder::WaitForInitDecodingH2HCopies() {
  std::unique_lock<std::mutex> lk(n_init_decoding_h2h_task_not_done_mutex_);
  init_decoding_h2h_done_.wait(
      lk, [this] { return (n_init_decoding_h2h_task_not_done_ == 0); });
}

void CudaDecoder::PrepareForGetRawLattice(
    const std::vector<ChannelId> &channels, bool use_final_probs) {
  GetBestCost(channels, use_final_probs, &argmins_,
              &list_finals_token_idx_and_cost_, &has_reached_final_);
  for (LaneId ilane = 0; ilane < channels.size(); ++ilane) {
    ChannelId ichannel = channels[ilane];
    std::lock_guard<std::mutex> channel_lk(channel_lock_[ichannel]);
    h_all_argmin_cost_[ichannel] = argmins_[ilane];
    h_all_final_tokens_list_[ichannel].swap(
        list_finals_token_idx_and_cost_[ilane]);
    h_all_has_reached_final_[ichannel] = has_reached_final_[ilane];
  }
}

void CudaDecoder::ConcurrentGetRawLatticeSingleChannel(const ChannelId ichannel,
                                                       Lattice *fst_out) {
  nvtxRangePushA("GetRawLatticeOneChannel");
  // Allocating the datastructures that we need

  // [prev|curr]_f_raw_lattice_state
  // Used to get information about a lattice state (i.e. a (iframe, fst_state)
  // pair)
  // using its LatticeStateInternalId (its ID inside of the decoder)
  // It gives us the OutputLatticeState (its ID in the output lattice)
  // alongside with the extra_cost of that state in the lattice
  // Those maps are used to build the external lattice using what we know
  // internally
  // Using one map per frame. We always know to which frame a token belongs.
  // Using one big map slows everything down
  std::unordered_map<LatticeStateInternalId, RawLatticeState>
      prev_f_raw_lattice_state, curr_f_raw_lattice_state;
  // We want the unicity of each arc_idx for one frame. Important because we
  // can replay a frame (and possibly add multiple time the same arc)
  std::unordered_set<int32> f_arc_idx_added;
  // When backtracking, we read tokens in the current frame (in
  // q_curr_frame_todo_),
  // we backtrack the associated arc, and we add the predecessor either to
  // q_curr_frame_todo_ (non-emitting arc, same frame)
  // or q_prev_frame_todo_ (emitting arc, source in previous frame)
  std::vector<std::pair<TokenId, InfoToken>> q_curr_frame_todo;
  std::vector<std::pair<TokenId, InfoToken>> q_prev_frame_todo;

  // We need to lock the channel to read argmin
  TokenId best_cost_idx;
  {
    std::lock_guard<std::mutex> channel_lk(channel_lock_[ichannel]);
    h_all_tokens_info_.shrink_to_fit();
    h_all_tokens_acoustic_cost_.shrink_to_fit();
    h_all_tokens_extra_prev_tokens_.shrink_to_fit();
    h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_.shrink_to_fit();
    best_cost_idx = h_all_argmin_cost_[ichannel].first;
  }
  KALDI_ASSERT(
      "You need to call PrepareForGetRawLattice before "
      "ConcurrentGetRawLatticeSingleChannel" &&
      best_cost_idx >= 0);
  const int32 nframes = NumFramesDecoded(ichannel);
  // Making sure that this token is available for this channel.
  // We're going to read storage data from this channel. Locking it
  // If that token in that frame f is available, then all tokens in that frame
  // f are available
  WaitForH2HCopies();
  std::unique_lock<std::mutex> channel_lk(channel_lock_[ichannel]);
  // Total number of tokens generated by the utterance on channel ichannel
  const int32 total_ntokens = h_all_tokens_info_[ichannel].size();

  // Preparing output lattice
  // The start state has to be 0 (cf some asserts somewhere else in Kaldi)
  // Adding it now
  fst_out->DeleteStates();
  OutputLatticeState fst_lattice_start = fst_out->AddState();
  fst_out->SetStart(fst_lattice_start);

  // Adding the best tokens returned by GetBestCost to the lattice
  // We also add them to q_curr_frame_todo, and we'll backtrack from there
  AddFinalTokensToLattice(ichannel, &q_curr_frame_todo,
                          &curr_f_raw_lattice_state, fst_out);

  // We're now going to backtrack frame by frame
  // For each frame we're going to process tokens that need to be inserted
  // into the output lattice
  // and add their predecessors to the todo list
  // iframe == -1 contains the start state and the first non emitting tokens.
  // It is not linked to a real frame
  for (int32 iframe = nframes - 1; iframe >= -1; --iframe) {
    // Tokens for the current frame were inserted after this offset in the
    // token list
    const int32 curr_frame_offset =
        (iframe >= 0) ? frame_offsets_[ichannel][iframe] : 0;

    // bool must_replay_frame
    // In some cases we can update an extra_cost that has already been used
    // For instance we process arcs in that order :
    // 1) a -> b, which updates extra_cost[b] using extra_cost[a]
    // 2) c -> a, which updates extra-cost[a] (using extra_cost[c])
    // because the arcs were not considered in topological order, we need to
    // run
    // again the step 1,
    // to get the correct extra_cost[b] (using the latest extra_cost[a])
    // However, we only re-run the step 1 if the value extra_cost[a] has
    // changed more than extra_cost_min_delta_
    bool must_replay_frame;

    // dbg_found_best_path is used in an useful assert, making sure the best
    // path is still there for each frame
    // if something went wrong in the kernels, it's not likely we respect that
    // property out of luck
    bool dbg_found_best_path = false;
    do {
      must_replay_frame = false;
      // Reading something to do. We are pushing stuff back in
      // q_curr_frame_todo while reading it,
      // so it's important to always read q_curr_frame_todo_.size() directly
      // not using a queue, because we may need to recompute the frame (if
      // must_replay_frame is true)
      for (int32 u = 0; u < q_curr_frame_todo.size(); ++u) {
        TokenId token_idx;
        InfoToken token;
        std::tie(token_idx, token) = q_curr_frame_todo[u];
        KALDI_ASSERT(token_idx >= curr_frame_offset);
        CostType token_extra_cost;
        StateId to_fst_lattice_state;
        // Loading the current extra_cost of that token
        // + its associated state in the lattice
        GetTokenRawLatticeData(token_idx, token, total_ntokens,
                               &curr_f_raw_lattice_state, &token_extra_cost,
                               &to_fst_lattice_state);
        dbg_found_best_path |= (token_extra_cost == 0.0f);

        InfoToken *tok_beg;
        float2 *extra_extra_and_acoustic_cost_beg;
        int32 nsamestate;
        // Getting the list of the tokens linked to the same FST state, in the
        // same frame
        // In the GPU decoder a token is linked to a single arc, but we can
        // generate
        // multiple token for a same fst_nextstate in the same frame.
        // In the CPU decoder we would use the forward_links list to store
        // everything in the same metatoken
        // GetSameFSTStateTokenList returns the list of tokens linked to the
        // same FST state than token
        // (in the current frame)
        GetSameFSTStateTokenList(ichannel, token, &tok_beg,
                                 &extra_extra_and_acoustic_cost_beg,
                                 &nsamestate);

        // dbg_found_zero used for debugging. For each FST state, we have a
        // token with the
        // best cost for that FST state
        // that token has an extra_cost of 0.0f. This is a sanity check
        bool dbg_found_zero = false;
        for (int32 iprev = 0; iprev < nsamestate; ++iprev) {
          InfoToken list_prev_token;
          CostType acoustic_cost, this_arc_prev_token_extra_cost;
          bool keep_arc;
          LatticeStateInternalId src_state_internal_id;
          InfoToken list_token = tok_beg[iprev];
          int32 list_prev_token_idx = list_token.prev_token;
          int32 list_arc_idx = list_token.arc_idx;

          ConsiderTokenForLattice(
              ichannel, iprev, total_ntokens, token_idx, fst_lattice_start,
              tok_beg, extra_extra_and_acoustic_cost_beg, token_extra_cost,
              list_prev_token_idx, list_arc_idx, &list_prev_token,
              &this_arc_prev_token_extra_cost, &acoustic_cost,
              &src_state_internal_id, &keep_arc, &dbg_found_zero);

          if (keep_arc)
            AddArcToLattice(list_arc_idx, list_prev_token_idx, list_prev_token,
                            curr_frame_offset, acoustic_cost,
                            this_arc_prev_token_extra_cost,
                            src_state_internal_id, fst_lattice_start,
                            to_fst_lattice_state, &q_curr_frame_todo,
                            &q_prev_frame_todo, &curr_f_raw_lattice_state,
                            &prev_f_raw_lattice_state, &f_arc_idx_added,
                            fst_out, &must_replay_frame);
        }
        KALDI_ASSERT(dbg_found_zero);
      }

      if (must_replay_frame) {
        // We need to replay the frame. Because all states will be read again,
        // we can reopen them (and they will be closed again when being read
        // from again)
        for (auto it = curr_f_raw_lattice_state.begin();
             it != curr_f_raw_lattice_state.end(); ++it) {
          it->second.is_state_closed = false;
        }
      }
    } while (must_replay_frame);

    // Done processing this frame. Swap the datastructures, move on to
    // previous frame (we go --iframe)
    SwapPrevAndCurrLatticeMap(iframe, dbg_found_best_path, &q_curr_frame_todo,
                              &q_prev_frame_todo, &curr_f_raw_lattice_state,
                              &prev_f_raw_lattice_state, &f_arc_idx_added);
  }
  nvtxRangePop();
}

void CudaDecoder::GetRawLattice(const std::vector<ChannelId> &channels,
                                std::vector<Lattice *> &fst_out_vec,
                                bool use_final_probs) {
  KALDI_ASSERT(channels.size() == fst_out_vec.size());
  // Getting the list of the best costs in the lastest token queue.
  // all costs within [best;best+lattice_beam]
  PrepareForGetRawLattice(channels, use_final_probs);
  for (int32 ilane = 0; ilane < channels.size(); ++ilane) {
    const ChannelId ichannel = channels[ilane];
    Lattice *fst_out = fst_out_vec[ilane];
    ConcurrentGetRawLatticeSingleChannel(ichannel, fst_out);
  }
}

void CudaDecoder::SetChannelsInKernelParams(
    const std::vector<ChannelId> &channels) {
  KALDI_ASSERT(channels.size() <= nchannels_);
  KALDI_ASSERT(channels.size() <= nlanes_);
  for (LaneId lane_id = 0; lane_id < channels.size(); ++lane_id)
    channel_to_compute_[lane_id] = channels[lane_id];
  h_kernel_params_->nlanes_used = channels.size();
  nlanes_used_ = channels.size();
}

void CudaDecoder::ResetChannelsInKernelParams() {
  h_kernel_params_->nlanes_used = 0;
  nlanes_used_ = 0;
}

int32 CudaDecoder::NumFramesDecoded(ChannelId ichannel) const {
  KALDI_ASSERT(ichannel < nchannels_);
  return num_frames_decoded_[ichannel];
}

void CudaDecoder::CheckStaticAsserts() {
  // Checking if all constants look ok

  // We need that because we need to be able to do the scan in one pass in the
  // kernel
  // update_beam_using_histogram_kernel
  KALDI_ASSERT(KALDI_CUDA_DECODER_HISTO_NBINS < KALDI_CUDA_DECODER_1D_BLOCK);
  KALDI_ASSERT(KALDI_CUDA_DECODER_NONEM_LT_MAX_NARCS > 0);
}

void CudaDecoder::LaunchH2HCopies() {
  // Each H2H copy counter
  n_acoustic_h2h_copies_todo_.store(nlanes_used_ - 1);
  n_infotoken_h2h_copies_todo_.store(nlanes_used_ - 1);
  n_extra_prev_tokens_h2h_copies_todo_.store(nlanes_used_ - 1);

  {
    std::lock_guard<std::mutex> n_h2h_not_done_lk(n_h2h_task_not_done_mutex_);
    n_h2h_task_not_done_ += thread_pool_ ? n_threads_used_ : 1;
  }
  {
    std::lock_guard<std::mutex> n_h2h_todo_lk(n_h2h_main_task_todo_mutex_);
    n_h2h_main_task_todo_ = thread_pool_ ? n_threads_used_ : 1;
  }

  // Either do the copy locally or send it to the threadpool
  if (thread_pool_) {
    n_h2h_main_task_todo_cv_.notify_all();
  } else {
    ComputeH2HCopies();
  }
}

void CudaDecoder::ComputeH2HCopiesCPUWorker() {
  // Run by a dedicated CPU thread
  while (h2h_threads_running_) {
    ComputeH2HCopies();
  }
}

void CudaDecoder::ComputeH2HCopies() {
  // Waiting for either something to do or the instruction to stop the threads
  {
    std::unique_lock<std::mutex> n_h2h_lk(n_h2h_main_task_todo_mutex_);
    n_h2h_main_task_todo_cv_.wait(n_h2h_lk, [this] {
      return !h2h_threads_running_ || (n_h2h_main_task_todo_ > 0);
    });
    --n_h2h_main_task_todo_;
  }
  // If we are done, stop the wait and return now. ComputeH2HCopiesCPUWorker
  // will also return,
  // stopping the thread
  if (!h2h_threads_running_) return;
  // Waiting for the D2H copies. This is threadsafe
  // Step 1: acoustic costs
  cudaEventSynchronize(d2h_copy_acoustic_evt_);
  int32 ilane;
  while ((ilane = n_acoustic_h2h_copies_todo_.fetch_sub(1)) >= 0) {
    int32 ichannel = lanes2channels_todo_[ilane];
    // Lock Channel
    std::lock_guard<std::mutex> channel_lk(channel_lock_[ichannel]);
    MoveConcatenatedCopyToVector(
        ilane, ichannel, h_emitting_main_q_end_lane_offsets_,
        h_acoustic_cost_concat_, &h_all_tokens_acoustic_cost_);
    // Adding 0.0f acoustic_costs for non-emittings
    int32 main_q_end = h_main_q_end_lane_offsets_[ilane + 1] -
                       h_main_q_end_lane_offsets_[ilane];
    int32 ntokens_emitting = h_emitting_main_q_end_lane_offsets_[ilane + 1] -
                             h_emitting_main_q_end_lane_offsets_[ilane];
    int32 ntokens_nonemitting = main_q_end - ntokens_emitting;
    auto &vec = h_all_tokens_acoustic_cost_[ichannel];
    vec.insert(vec.end(), ntokens_nonemitting, 0.0f);
  }

  // Step 2: infotoken
  cudaEventSynchronize(d2h_copy_infotoken_evt_);
  while ((ilane = n_infotoken_h2h_copies_todo_.fetch_sub(1)) >= 0) {
    int32 ichannel = lanes2channels_todo_[ilane];
    // Lock Channel
    std::lock_guard<std::mutex> channel_lk(channel_lock_[ichannel]);
    MoveConcatenatedCopyToVector(ilane, ichannel, h_main_q_end_lane_offsets_,
                                 h_infotoken_concat_, &h_all_tokens_info_);
  }

  // Step 3: extra prev tokens
  cudaEventSynchronize(d2h_copy_extra_prev_tokens_evt_);
  while ((ilane = n_extra_prev_tokens_h2h_copies_todo_.fetch_sub(1)) >= 0) {
    int32 ichannel = lanes2channels_todo_[ilane];
    // Lock Channel
    std::lock_guard<std::mutex> channel_lk(channel_lock_[ichannel]);
    MoveConcatenatedCopyToVector(
        ilane, ichannel, h_n_extra_prev_tokens_lane_offsets_,
        h_extra_prev_tokens_concat_, &h_all_tokens_extra_prev_tokens_);
    MoveConcatenatedCopyToVector(
        ilane, ichannel, h_n_extra_prev_tokens_lane_offsets_,
        h_extra_and_acoustic_cost_concat_,
        &h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_);
  }

  // If we're the last cpu thread to complete the current tasks, notify the main
  // thread
  bool all_done;
  {
    std::lock_guard<std::mutex> lk_not_done(n_h2h_task_not_done_mutex_);
    all_done = (--n_h2h_task_not_done_ == 0);
  }
  if (all_done) {
    h2h_done_.notify_all();
  }
}

void CudaDecoder::SetThreadPoolAndStartCPUWorkers(ThreadPool *thread_pool,
                                                  int32 nworkers) {
  KALDI_ASSERT(nworkers > 0);
  n_threads_used_ = nworkers;
  thread_pool_ = thread_pool;
  for (int32 i = 0; i < nworkers; ++i)
    cpu_dedicated_threads_.emplace_back(&CudaDecoder::ComputeH2HCopiesCPUWorker,
                                        this);
}

}  // end namespace cuda_decoder
}  // end namespace kaldi

#endif  // HAVE_CUDA == 1