// cudadecoder/cuda-decoder-kernels.cu
  //
  // 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.
  
  #include <cub/cub.cuh>
  #include "cuda-decoder-kernels.h"
  #include "cuda-decoder-kernels-utils.h"
  
  namespace kaldi {
  namespace cuda_decoder {
  
  // Initialize the hashmap with NO_VAL
  // Called in InitDeviceData, when building the CudaDecoder object
  __global__ void init_hashmap_kernel(DeviceParams cst_dev_params) {
    const int max_nlanes = cst_dev_params.max_nlanes;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, max_nlanes) {
      const int capacity = cst_dev_params.hashmap_capacity;
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(idx, capacity) {
        cst_dev_params.d_hashmap_values.lane(ilane)[idx] =
            KALDI_CUDA_DECODER_HASHMAP_NO_VAL;
      }
    }
  }
  
  // Initialize initial channel on  device
  // Called by ComputeInitialChannel
  // It is NOT called in InitDecoding
  // In InitDecoding we will clone the initial channel into the channel we called
  // InitDecoding on
  // Here we are actually creating this initial channel
  // we do that once in the CudaDecoder constructor.
  //
  // The initial channel is the state of a channel when
  // it will start decoding a new utterance
  // thread (1, 1, 1)
  // blocks(1, 1, 1);
  __global__ void initialize_initial_lane_kernel(DeviceParams cst_dev_params) {
    const int init_ichannel = cst_dev_params.init_channel_id;
    const int init_ilane = 0;
    ChannelCounters *init_channel_counters =
        cst_dev_params.d_channels_counters.channel(init_ichannel);
    LaneCounters *lane_counters =
        cst_dev_params.d_lanes_counters.lane(init_ilane);
  
    // Making the data look like an ExpandArcsEmitting just executed,
    // and put the StartState in the aux_q. We will then pick up a normal
    // execution from there
    // (calling PruneAndPreprocess, then ExpandArcsNonEmitting..)
    lane_counters->aux_q_end = 0;
    lane_counters->aux_q_requested = 0;
    lane_counters->post_expand_aux_q_end = 1;
    lane_counters->main_q_global_offset = 0;
    lane_counters->main_q_local_offset = 0;
    lane_counters->main_q_n_extra_prev_tokens = 0;
    lane_counters->int_cutoff = INT_MAX;
    lane_counters->main_q_n_emitting_tokens = 0;  // all non emitting
    lane_counters->int_beam = floatToOrderedInt(cst_dev_params.default_beam);
    lane_counters->main_q_narcs_and_end = {0, 0};
    lane_counters->main_q_requested = 0;
    lane_counters->prev_arg_min_int_cost = 0;
    const StateId init_state = cst_dev_params.init_state;
    const CostType init_cost = cst_dev_params.init_cost;
    IntegerCostType int_init_cost = floatToOrderedInt(init_cost);
    cst_dev_params.d_aux_q_state_and_cost.lane(init_ilane)[0] = {init_state,
                                                                 int_init_cost};
    lane_counters->min_int_cost = int_init_cost;
    CostType cutoff = orderedIntToFloat(int_init_cost);
    lane_counters->int_cutoff =
        floatToOrderedInt(cutoff + cst_dev_params.default_beam);
    cst_dev_params.d_aux_q_info.lane(init_ilane)[0] = {INT_MIN, -1};
  }
  
  // Called by InitDecoding
  // Called when some channels will start decoding a new utterance
  // do everything that's needed to do on the device to start decoding a new
  // utterance with those channels
  // It clones the initial channel (created in initialize_initial_lane_kernel)
  // into the channels we want to InitDecoding on
  __global__ void init_decoding_on_device_kernel(DeviceParams cst_dev_params,
                                                 KernelParams params) {
    const int init_ichannel = cst_dev_params.init_channel_id;
  
    const ChannelCounters *init_channel_counters =
        cst_dev_params.d_channels_counters.channel(init_ichannel);
    const int32 init_main_q_end =
        init_channel_counters->prev_main_q_narcs_and_end.y;
    const int32 nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(idx, init_main_q_end) {
        const LaneCounters *lane_counters =
            cst_dev_params.d_lanes_counters.lane(ilane);
        const int32 ichannel = lane_counters->channel_to_compute;
        cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[idx] =
            cst_dev_params.d_main_q_state_and_cost.channel(init_ichannel)[idx];
        cst_dev_params.d_main_q_degrees_prefix_sum.channel(ichannel)[idx] =
            cst_dev_params.d_main_q_degrees_prefix_sum.channel(
                init_ichannel)[idx];
        cst_dev_params.d_main_q_arc_offsets.channel(ichannel)[idx] =
            cst_dev_params.d_main_q_arc_offsets.channel(init_ichannel)[idx];
        if (idx == 0) {
          ChannelCounters *channel_counters =
              cst_dev_params.d_channels_counters.channel(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 = cst_dev_params.default_beam;
        }
      }
    }
  }
  
  // Context switch : load
  // Called by LoadChannelsStateToLanes
  // THREADS : (1, 1, 1)
  // BLOCKS : (1, nlanes_used, 1)
  __global__ void load_channels_state_in_lanes_kernel(DeviceParams cst_dev_params,
                                                      KernelParams params) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      const ChannelCounters *channel_counters =
          cst_dev_params.d_channels_counters.channel(ichannel);
      int2 main_q_narcs_and_end = channel_counters->prev_main_q_narcs_and_end;
      lane_counters->main_q_narcs_and_end = main_q_narcs_and_end;
      lane_counters->main_q_n_extra_prev_tokens =
          channel_counters->prev_main_q_n_extra_prev_tokens;
      CostType beam = channel_counters->prev_beam;
      IntegerCostType int_beam = floatToOrderedInt(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 =
          cst_dev_params.adaptive_beam_static_segment;
      lane_counters->main_q_global_offset =
          channel_counters
              ->prev_main_q_global_offset;  // we'll update it after emitting
      lane_counters->main_q_extra_prev_tokens_global_offset =
          channel_counters->prev_main_q_extra_prev_tokens_global_offset;
    }
  }
  
  // Context switch : store
  // Called by SaveChannelsStateFromLanes
  // THREADS : (1, 1, 1)
  // BLOCKS : (1, nchannel_to_compute, 1)
  __global__ void save_channels_state_from_lanes_kernel(
      DeviceParams cst_dev_params, KernelParams params) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      const LaneCounters *lane_counters =
          cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      ChannelCounters *channel_counters =
          cst_dev_params.d_channels_counters.channel(ichannel);
      channel_counters->prev_main_q_global_offset =
          lane_counters->main_q_global_offset;
      channel_counters->prev_main_q_extra_prev_tokens_global_offset =
          lane_counters->main_q_extra_prev_tokens_global_offset;
      channel_counters->prev_main_q_narcs_and_end =
          lane_counters->main_q_narcs_and_end;
      channel_counters->prev_main_q_n_extra_prev_tokens =
          lane_counters->main_q_n_extra_prev_tokens;
      channel_counters->prev_beam = orderedIntToFloat(lane_counters->int_beam);
    }
  }
  
  // compute_lane_offsets_kernel
  // the kernel concatenate_lanes_data concatenates multiple array into a single
  // continuous array
  // compute_lane_offsets_kernel computes the offset of each array into this
  // continous array
  // This kernel is 1D : the lanes are on the X dimension, because we want to
  // compute the offset of those lanes
  __global__ void compute_lane_offsets_kernel(DeviceParams cst_dev_params,
                                              KernelParams params) {
    typedef cub::BlockScan<int4, KALDI_CUDA_DECODER_1D_BLOCK> BlockScan;
    __shared__ typename BlockScan::TempStorage temp_storage;
  
    const int nlanes = params.nlanes_used;
    int4 sum_so_far = {0, 0, 0, 0};
    KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(
        block_offset, thread_idx,
        nlanes + 1) {  // +1 because we are doing an exclusive sum, and we want
                       // all the values
      int32 ilane = block_offset + thread_idx;
      int4 zero4 = {0, 0, 0, 0};
      int4 lane_offsets = zero4;
      if (ilane < nlanes) {  // nlanes, not nlanes+1, because we cannot read +1
                             // values (undefined)
        LaneCounters *d_lane_counters =
            cst_dev_params.d_lanes_counters.lane(ilane);
        int32 main_q_end = d_lane_counters->main_q_narcs_and_end.y;
        int32 n_emitting_tokens = d_lane_counters->main_q_n_emitting_tokens;
        int32 main_q_n_extra_prev_tokens =
            d_lane_counters->main_q_n_extra_prev_tokens;
        lane_offsets = {main_q_end, n_emitting_tokens, main_q_n_extra_prev_tokens,
                        0};
      }
      int4 block_aggregate;
      BlockScan(temp_storage)
          .ExclusiveScan(lane_offsets, lane_offsets, zero4, PlusPlusPlusPlus(),
                         block_aggregate);
      PlusPlusPlusPlus pppp;
      lane_offsets = pppp(lane_offsets, sum_so_far);
      sum_so_far = pppp(sum_so_far, block_aggregate);
      if (ilane < (nlanes + 1)) {  // nlanes+1, to write the output
        LaneCounters *d_lane_counters =
            cst_dev_params.d_lanes_counters.lane(ilane);
        LaneCounters *h_lane_counters =
            cst_dev_params.h_lanes_counters.lane(ilane);
        h_lane_counters->main_q_end_lane_offset =
            d_lane_counters->main_q_end_lane_offset = lane_offsets.x;
        h_lane_counters->main_q_n_emitting_tokens_lane_offset =
            d_lane_counters->main_q_n_emitting_tokens_lane_offset =
                lane_offsets.y;
        h_lane_counters->main_q_n_extra_prev_tokens_lane_offset =
            d_lane_counters->main_q_n_extra_prev_tokens_lane_offset =
                lane_offsets.z;
      }
      __syncthreads();  // reusing temp_storage
    }
  }
  
  // concatenate_lanes_data
  // Called by PerformConcatenatedCopy
  // Creates a concatenate array into concat,
  // by concatenating all the arrays src.lane(ilane)
  // for ilane=0..params.nlanes_used
  // Used to prepare data for copy to Host. We want to avoid small Device2Host
  // copies.
  template <typename T>
  __global__ void concatenate_lanes_data_kernel(DeviceParams cst_dev_params,
                                                KernelParams params,
                                                LaneMatrixView<T> src, T *concat,
                                                int32 *lane_offsets) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      const int32 stride =
          sizeof(LaneCounters) / sizeof(int32);  // offsets are in LaneCounters
      int32 beg = *(lane_offsets + ilane * stride);
      int32 end = *(lane_offsets + (ilane + 1) * stride);
      int32 vec_size = end - beg;
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(idx, vec_size) {
        T d = src.lane(ilane)[idx];
        concat[beg + idx] = d;
      }
    }
  }
  
  // nonemitting_preprocess_and_contract_kernel
  // Called from PruneAndPreprocess
  // This kernels prune the aux_q, move the survival tokens to the main_q,
  // and add the preprocessing information necessary for the next ExpandArcs
  // (the expand that follows PruneAndPreprocess is always non-emitting)
  // It prunes the tokens using the cutoff, and prepare the data necessary for
  // ExpandArcs:
  // d_main_q_degrees_prefix_sum, d_main_q_arc_offsets_
  // The prefix sum is done in one-pass here, using a trick (we compute the prefix
  // sum
  // as we fill the main_q)
  __global__ void nonemitting_preprocess_and_contract_kernel(
      DeviceParams cst_dev_params, KernelParams params) {
    typedef cub::BlockScan<int2, KALDI_CUDA_DECODER_1D_BLOCK> BlockScan;
    __shared__ typename BlockScan::TempStorage sh_temp_storage;
    // We need to move the survival tokens to the main_q
    //
    // sh_main_q_global_block_offset has two purposes :
    // (1) to know where to store the survival tokens in the main_q
    // (2) to perform the prefix sum degrees (of the survival tokens)
    __shared__ int2 sh_main_q_global_block_offset;
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 aux_q_end = lane_counters->post_expand_aux_q_end;
      const IntegerCostType int_cutoff = lane_counters->int_cutoff;
      // Keeping whole CTA alive. We'll use __syncthreads()
      KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(block_offset, thread_idx,
                                                     aux_q_end) {
        const int32 aux_q_idx = block_offset + thread_idx;
        const int32 ichannel = lane_counters->channel_to_compute;
        int32 degree = 0;
        int32 arc_start = -1;
        StateId token_state;
        IntegerCostType token_int_cost;
        // We've kept the whole CTA alive. Now we keep only those will a valid
        // token
        if (aux_q_idx < aux_q_end) {
          const int2 both =
              cst_dev_params.d_aux_q_state_and_cost.lane(ilane)[aux_q_idx];
          token_state = both.x;
          token_int_cost = both.y;
  
          if (token_int_cost < int_cutoff) {
            // We'll keep that token. Loading its arc degree/csr offset now.
            arc_start = cst_dev_params.d_arc_ne_offsets[token_state];
            const int32 arc_end =
                cst_dev_params.d_arc_ne_offsets[token_state + 1];
            degree = arc_end - arc_start;
          }
        }
  
        // If we've set a different arc_start,
        // this thread has a valid unpruned token
        int32 is_pruned = (arc_start == -1);
  
        // We now know which tokens will be moved to the main_q, the remaining
        // will be pruned
        // we now compute a prefix sum inside the CUDA block to determine the
        // local indexes of the unpruned tokens
        // the first unpruned token will have a index of 0, the second 1, ...
        // We also need to compute the prefix sum of the arc degrees
        // we start by doing a local prefix sum inside the CUDA block
        int2 block_prefix_sum_narcs_and_end = {degree, (is_pruned ? 0 : 1)};
        const int2 zero2 = {0, 0};
  
        // Computing the prefix sum (exclusive)
        BlockScan(sh_temp_storage)
            .ExclusiveScan(block_prefix_sum_narcs_and_end,
                           block_prefix_sum_narcs_and_end, zero2, PlusPlus());
  
        if (KALDI_CUDA_DECODER_IS_LAST_1D_THREAD()) {
          // This conditional branch is entered by the last thread
          // Because it is the last, the prefix_sum of that thread contains the
          // sum of all elements
  
          // We also add the value from this thread - the prefix sum is exclusive
          // For the sum, we want it inclusive
          int2 block_sum = block_prefix_sum_narcs_and_end;
          block_sum.x += degree;
          block_sum.y += is_pruned ? 0 : 1;
  
          // Doing two things at the same time :
          // requesting a spot in the main_q to store the survival tokens from
          // this CTA
          // We also increment the narcs value. atomic64.x will contain the number
          // of
          // arcs in the main_q up until the atomic64.y index
          // That's all we need to finish our prefix sum. We add this global
          // offset.
  
          // First atomic to check if we are not overflowing main_q.
          int block_offset =
              atomicAdd(&lane_counters->main_q_requested, block_sum.y);
  
          // Verify that we do not overflow
          if (block_offset + block_sum.y < cst_dev_params.main_q_capacity) {
            // we don't overflow we can safely grab a spot in the main_q
            sh_main_q_global_block_offset =
                atomicAddI2(&lane_counters->main_q_narcs_and_end, block_sum);
          } else {
            // our update would overflow
            lane_counters->q_overflow |= OVERFLOW_MAIN_Q;  // for the host
            sh_main_q_global_block_offset.y =
                cst_dev_params.main_q_capacity;  // used as flag to broadcast the
                                                 // information in the CTA
          }
        }
  
        // Syncing because :
        // - Broadcasting sh_main_q_global_block_offset
        // - We may reuse sh_temp_storage (cf CUB doc)
        __syncthreads();
  
        // Checking if we are overflowing the main_q
        // All threads are executing the next line
        if (sh_main_q_global_block_offset.y == cst_dev_params.main_q_capacity)
          goto end_lane;  // done for this lane
  
        // If we are executing the following lines it means that we are not
        // overflowing the queue
        // We then continue what we were doing
        if (!is_pruned) {
          bool moving_emitting_tokens = (lane_counters->main_q_local_offset == 0);
          // we will move our unpruned token to the main_q, at index main_q_idx
          InfoToken tok_info = cst_dev_params.d_aux_q_info.lane(ilane)[aux_q_idx];
          const int32 main_q_idx =
              sh_main_q_global_block_offset.y + block_prefix_sum_narcs_and_end.y;
          CostType acoustic_cost = 0.0f;
          if (moving_emitting_tokens && tok_info.arc_idx != -1) {
            const int32 arc_ilabel =
                cst_dev_params.d_arc_pdf_ilabels[tok_info.arc_idx];
            acoustic_cost = -lane_counters->loglikelihoods[arc_ilabel];
          }
          cst_dev_params.d_main_q_info.lane(ilane)[main_q_idx] = tok_info;
  
          // Moving the token to the main q
          cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[main_q_idx] = {
              token_state, token_int_cost};
          cst_dev_params.d_main_q_acoustic_cost.lane(ilane)[main_q_idx] =
              acoustic_cost;
          // Saving the global prefix sum
          const int32 prefix_sum_narcs =
              sh_main_q_global_block_offset.x + block_prefix_sum_narcs_and_end.x;
          cst_dev_params.d_main_q_degrees_prefix_sum.channel(
              ichannel)[main_q_idx] = prefix_sum_narcs;
          // Saving the CSR arc offset for that token's state
          // it will be used by the expand kernel, and avoid doing a new random
          // memory access in the expand kernel
          cst_dev_params.d_main_q_arc_offsets.channel(ichannel)[main_q_idx] =
              arc_start;
        }
      }
  
    end_lane:;  // empty statement
    }
  }
  
  // GetAdaptiveBeam is used in ExpandArcs
  // When we generate new tokens by traversing arcs, 
  // we can end up creating a lot of tokens, if the current frame 
  // generated loglikelihoods too uniform for instance (we don't have
  // any good tokens that will reduce the cutoff, so we end up generating
  // a lot of tokens)
  // To avoid overflowing the aux_q, we apply a decreasing beam.
  // With aux_q_end being the current aux_q size, we have a decrease function f, with
  // adaptive_beam = f(aux_q_end)
  // f is a decreasing piecewise constant function
  // Please note that when processing tokens, we usually have dozens of thousands of threads
  // generating tokens. Those are already in flight, and will not reload the beam immediatly.
  // It means that we need to start reducing the beam as soon as we detect that we are generating more tokens than
  // expected. 
  // We can configure the function f using KALDI_CUDA_DECODER_ADAPTIVE_BEAM_STATIC_SEGMENT
  // and KALDI_CUDA_DECODER_ADAPTIVE_BEAM_NSTEPS.
  // We will use default_beam for the first max_tokens_per_frame/KALDI_CUDA_DECODER_ADAPTIVE_BEAM_STATIC_SEGMENT
  // tokens in the aux_q.
  // Once we reach that number, we will decrease the adaptive beam linearly from default_beam to 0,
  // using KALDI_CUDA_DECODER_ADAPTIVE_BEAM_NSTEPS steps
  //
  // x-axis : aux_q_end. How much tokens are already in the aux_q
  // y-axis : adaptive_beam = f(aux_q_end)
  // default_beam _| ________________
  //               |               /\ _________
  //               |                |          _________
  //            0 _|   static_segment                   _________
  //               |________________________________________________
  //               |                                             |     
  //   aux_q_end=  0                                    max_tokens_per_frame
  // We have :     
  // static_segment = max_tokens_per_frame/KALDI_CUDA_DECODER_ADAPTIVE_BEAM_STATIC_SEGMENT
  // and KALDI_CUDA_DECODER_ADAPTIVE_BEAM_NSTEPS = 3
  __device__ void UpdateAdaptiveBeam(const DeviceParams &cst_dev_params,
                                     const int aux_q_index_block_offset,
                                     IntegerCostType min_int_cost,
                                     int2 *adaptive_int_beam_with_validity_index,
                                     LaneCounters *lane_counters) {
    int32 beam_valid_until_idx = adaptive_int_beam_with_validity_index->y;
    if (aux_q_index_block_offset < beam_valid_until_idx) return;  // nothing to do
  
    CostType beam = orderedIntToFloat(adaptive_int_beam_with_validity_index->x);
    while (aux_q_index_block_offset >= beam_valid_until_idx) {
      beam /= 2;
      beam_valid_until_idx += cst_dev_params.adaptive_beam_bin_width;
    }
  
    IntegerCostType new_int_cutoff = (min_int_cost < INT_MAX)
        ? floatToOrderedInt(orderedIntToFloat(min_int_cost) + beam)
        : INT_MAX;
    IntegerCostType int_beam = floatToOrderedInt(beam);
    adaptive_int_beam_with_validity_index->x = int_beam;
    adaptive_int_beam_with_validity_index->y = beam_valid_until_idx;
    // We can have races between the two atomics
    // However the worst than can happen is a CTA might delay updating the beam
    // This is not a critical bug. However, once we have a floatToOrderedInt
    // that is generating unsigned ints, we could merge the two atomics into a
    // single atomic64
    atomicMin(&lane_counters->adaptive_int_beam_with_validity_index.x, int_beam);
    atomicMax(&lane_counters->adaptive_int_beam_with_validity_index.y,
              beam_valid_until_idx);
    atomicMin(&lane_counters->int_cutoff, new_int_cutoff);
  }
  
  // One CTA / lane
  __global__ void reset_for_frame_and_estimate_cutoff_kernel(
      DeviceParams cst_dev_params, KernelParams params) {
    typedef cub::BlockReduce<CostType, KALDI_CUDA_DECODER_1D_BLOCK> BlockReduce;
    __shared__ typename BlockReduce::TempStorage temp_storage;
  
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      ChannelCounters *channel_counters =
          cst_dev_params.d_channels_counters.channel(ichannel);
      if (threadIdx.x == 0) {
        const CostType current_beam = orderedIntToFloat(lane_counters->int_beam);
        // Do some initialization
        lane_counters->q_overflow = OVERFLOW_NONE;
        lane_counters->main_q_n_emitting_tokens = INT_MAX;
        lane_counters->int_cutoff = INT_MAX;
        lane_counters->min_int_cost = INT_MAX;
        lane_counters->q_overflow = OVERFLOW_NONE;
        lane_counters->aux_q_requested = 0;
        lane_counters->main_q_requested = 0;
        lane_counters->main_q_local_offset = 0;
        lane_counters->compute_max_active =
            false;  // will be set to true if necessary
        channel_counters->min_int_cost_and_arg_with_final.x =
            INT_MAX;  // it will be set with atomicMins
        const CostType new_beam =
            fmin(cst_dev_params.default_beam,
                 current_beam * KALDI_CUDA_DECODER_ADAPTIVE_BEAM_RECOVER_RATE);
        lane_counters->int_beam = floatToOrderedInt(new_beam);
      }
      const int32 prev_arg_min = lane_counters->prev_arg_min_int_cost;
      int2 both =
          cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[prev_arg_min];
      int32 int_cost = both.y;
      CostType previous_cost = orderedIntToFloat(int_cost);
      const int32 prev_arg_min_state = both.x;
      int32 arc_start = cst_dev_params.d_arc_e_offsets[prev_arg_min_state];
      int32 arc_end = cst_dev_params.d_arc_e_offsets[prev_arg_min_state + 1];
      int32 narcs = arc_end - arc_start;
      // no loop - we only process the first KALDI_CUDA_DECODER_1D_BLOCK arcs
      // we just want an estimate
      CostType total_cost = FLT_MAX;
      if (threadIdx.x < narcs) {
        int32 iarc = arc_start + threadIdx.x;
        CostType arc_fixed_cost = cst_dev_params.d_arc_weights[iarc];
        const int32 arc_ilabel = cst_dev_params.d_arc_pdf_ilabels[iarc];
        CostType acoustic_cost = -lane_counters->loglikelihoods[arc_ilabel];
        total_cost = previous_cost + arc_fixed_cost +
                     acoustic_cost;  // +0.0f, best prev cost is normalized to 0
      }
  
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(bin_id, KALDI_CUDA_DECODER_HISTO_NBINS) { 
        cst_dev_params.d_histograms.lane(ilane)[bin_id] = 0; // reset for this frame
      }
  
      CostType min = BlockReduce(temp_storage).Reduce(total_cost, cub::Min());
      if (narcs > 0 && threadIdx.x == 0) {
        // narcs > 0 to have at least one valid element in the reduce
        CostType new_cutoff = min + orderedIntToFloat(lane_counters->int_beam);
        IntegerCostType new_int_cutoff = floatToOrderedInt(new_cutoff);
        lane_counters->int_cutoff = new_int_cutoff;
        lane_counters->min_int_cost = floatToOrderedInt(min);
      }
    }
  }
  // ExpandArc kernel
  // This kernel does the actual work of traversing arcs
  //
  // Pseudo code :
  // for all token tok in main_q[main_q_offset...end]:
  //      u = tok.next_state
  //      for all arc a(u->v) in the FST:
  //          v_cost = tok.cost + a.cost + accoustic_cost
  //
  //          if v_cost < cutoff and v_cost < best_state_cost[v]
  //              generate token associated to v, add to aux_q
  //              if necessary update cutoff
  //              if aux_q is getting full, reduce beam
  //
  // For more information please refer to http://kaldi-asr.org/doc/decoders.html
  //
  // ExpandArc rely on some preprocessed data to be able to function
  // for instance, it needs the prefix sum of the arc degree of all token.state in
  // the main_q
  // We need to call a Preprocess kernel before ExpandArc
  //
  // ExpandArc is used for both emitting and nonemitting phases
  // Differences between emitting and nonemitting :
  //      1) params.d_q_arc_offset contains offsets to either emitting or
  //      nonemitting arcs.
  //         It is transparent for this kernel. The differentiation was done in
  //         the Preprocess kernel,
  //         which is responsible for filling the params.d_q_arc_offset array
  //      2) Computation of the acoustic cost. If nonemitting, it is equal to 0.
  //      If emitting, we need
  //         to use values from the acoustic model (through the d_loglikelihoods
  //         array)
  //
  // Note : ExpandArc is not the only kernel able to traverse arcs.
  // FinalizeProcessNonemitting contains a simplified version of expand for only
  // one CUDA block
  template <bool IS_EMITTING>
  __global__ void expand_arcs_kernel(DeviceParams cst_dev_params,
                                     KernelParams params) {
    // BlockScan that we will use to compute token indexes in the output queue,
    // and to find the min cost in the block
    typedef cub::BlockScan<int2, KALDI_CUDA_DECODER_1D_BLOCK> BlockScan;
    __shared__ typename BlockScan::TempStorage sh_temp_storage_scan;
  
    // This kernel writes the new token to the output queue aux_q
    // We will request a spot to store all the new tokens created by threads in
    // this CUDA block
    // sh_aux_q_index_block_offset indicates where to store them in the aux_q
    // tokens created in this CUDA block will be store in :
    // aux_q[sh_aux_q_index_block_offset], aux_q[sh_aux_q_index_block_offset + 1],
    __shared__ int32 sh_aux_q_index_block_offset;
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 main_q_offset = lane_counters->main_q_local_offset;
      const int32 main_q_end = lane_counters->main_q_narcs_and_end.y;
      const int32 total_narcs = lane_counters->main_q_narcs_and_end.x;
      KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(block_offset, thread_idx,
                                                     total_narcs) {
        int2 adaptive_int_beam_with_validity_index =
            lane_counters->adaptive_int_beam_with_validity_index;
        const int32 ichannel = lane_counters->channel_to_compute;
        // Important : this thread is not responsible for a token in the input
        // queue main_q
        // but for an arc, going out of a token in the main_q
        // The main_q contains in total total_narcs
        // and this thread will compute the main_q_arc_index-th arc of the main_q
        // For instance, first thread in the grid with threadIdx.x == 0 and
        // blockIdx.x == 0
        // will process the first arc of the token in main_q[main_q_offset + 0]
        // (if that token has at least one arc)
        //
        // This insure a perfect one thread = one arc load balancing
        // but we have work to do to know exactly which arc is the
        // main_q_arc_index-th arc
        // (what's its source ? its destination ? its arc_idx the FST CSR ?)
        int32 main_q_arc_index = block_offset + thread_idx;
        // We'll need those variables later in the kernel
        // we declare them outside of the "valid_input" scope
        // to be able to access them later
        int32 main_q_idx;
        int32 arc_idx;
        StateId arc_next_state;
        IntegerCostType int_total_cost = INT_MAX;
        if (main_q_arc_index < total_narcs) {
          // Current thread must take care of main_q_arc_index-th arc
          // we need to now what's the source of that arc
          // ie which token.state in main_q does it start from ?
          // We use a binary search in the prefix sum of the token's degree to get
          // that information
          //
          // Example : main_q contains 3 tokens
          // - First token is associated to a state which has 3 outgoing arc
          // - Second token is associated to a state which has 0 outgoing arc
          // - Third token is associated to a state which has 2 outgoing arc
          //
          // We store the degrees in an array :
          // [3, 0, 2]
          //
          // We then compute the exclusive prefix sum of that array :
          // [0, 3, 3, 5]
          //
          // In total, we have 5 arcs in the main_q. ExpandArc will use 5 threads.
          //
          // Let's say we are the fifth thread in ExpandArc.
          // we have threadIdx.x == 4, and blockIdx.x == 0
          // it gives us main_q_arc_index == 4
          // From there we have no idea what we're supposed to do next, we need to
          // have information about the
          // arc that we're supposed to traverse
          //
          // To do that, we look for the maximum index maxle_i in the prefix sum
          // array such prefix_sum[i] <= 4
          //
          // [0, 3, 3, 5]
          //          |
          //         here
          // maxle_i = 2
          // it means that our source token is at index 2 in the main_q
          // and we are computing the arc at index (main_q_arc_index -
          // prefix_sum[maxle_i]) of that token
          // ie the arc at index (4-3) = 1, the second arc of the second token in
          // main_q
  
          // Searching for the source of the arc that we will process
          // (main_q_arc_index)
          // we could preprocess the search in the preprocess kernels - for now
          // this kernel is fast enough
          const int32 *degrees_prefix_sum =
              cst_dev_params.d_main_q_degrees_prefix_sum.channel(ichannel);
          main_q_idx = binsearch_maxle(degrees_prefix_sum, main_q_arc_index,
                                       main_q_offset, main_q_end - 1);
  
          // state_first_arc_idx_in_main_q
          // d_main_q_degrees_prefix_sum contains the prefix sum of the
          // degrees of all tokens in the main_q
          // d_main_q_degrees_prefix_sum[main_q_idx] contains the number of arc
          // in the main_q until that token
          const int32 state_first_arc_idx_in_main_q =
              degrees_prefix_sum[main_q_idx];
  
          // arc_offset_start is the offset in the CSR, to find the arcs
          // related to the state main_q_state_[main_q_idx]
          // it was set by the preprocess kernel
          const int32 arc_offset_start =
              cst_dev_params.d_main_q_arc_offsets.channel(ichannel)[main_q_idx];
  
          // local_arc_index is the arc index for that state
          // if local_arc_index == 2, we will process the second arc
          // of state main_q_state_[main_q_idx]
          const int32 local_arc_index =
              main_q_arc_index - state_first_arc_idx_in_main_q;
  
          // corresponding arc_idx in the FST
          arc_idx = arc_offset_start + local_arc_index;
  
          // Destination of that arc
          arc_next_state = cst_dev_params.d_arc_nextstates[arc_idx];
  
          // Building the total cost incrementally
          // we'll add the acoustic cost and the old token's cost
          const CostType arc_fixed_cost = cst_dev_params.d_arc_weights[arc_idx];
          const CostType prev_token_cost = orderedIntToFloat(
              cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[main_q_idx]
                  .y);
          CostType total_cost = prev_token_cost + arc_fixed_cost;
          const int32 prev_state =
              cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[main_q_idx]
                  .x;
          if (IS_EMITTING) {
            const int32 arc_ilabel = cst_dev_params.d_arc_pdf_ilabels[arc_idx];
            CostType acoustic_cost = -lane_counters->loglikelihoods[arc_ilabel];
            total_cost += acoustic_cost;
          }
          int_total_cost = floatToOrderedInt(total_cost);
  
          // If the total_cost is too large compared to our cutoff (beam search)
          // then let's drop it
          const IntegerCostType int_cutoff = lane_counters->int_cutoff;
          if (int_total_cost >= int_cutoff) int_total_cost = INT_MAX;
        }
  
        // If int_total_cost < INT_MAX, it means that :
        // - this thread had a valid input (main_q_arc_index < total_narcs)
        // - the total_cost of the generated token is < cutoff
        // We will then add that new token in the output queue, aux_q
        // We need to know where to put that token in the aux_q
        // we'll first compute its index inside the CUDA block
        // the first valid output token in the CUDA block will have index 0,
        // the second index 1... We compute that using a prefix sum
        //
        // We also need to find the overall min cost in the CUDA block
        // a prefix sum is a scan operation, and a min a reduce operation
        // we can perform a reduce operation using a scan (using the last value)
        // we compute the prefix sum and the min in one scan, using the data
        // struct CostTypeAndInt
        const int32 has_successor = (int_total_cost < INT_MAX) ? 1 : 0;
  
        int2 int_cost_and_index = {int_total_cost, has_successor};
        BlockScan(sh_temp_storage_scan)
            .InclusiveScan(int_cost_and_index, int_cost_and_index, MinPlus());
        if (KALDI_CUDA_DECODER_IS_LAST_1D_THREAD()) {
          // We are in a divergent branch
          // This is the last thread. The last value of the inclusive scan is the
          // total
          const int32 total_successors_in_block = int_cost_and_index.y;
          // Requesting a spot of size total_successors_in_block in the aux_q
  
          // note:  using 2 atomics here to avoid adding another kernel
          // first request more space
          const int aux_q_index_block_offset = atomicAdd(
              &lane_counters->aux_q_requested, total_successors_in_block);
  
          // check for overflow in aux_q
          // We try to prevent an overflow from happening using an adaptive beam
          // (cf GetAdaptiveBeam)
          if (aux_q_index_block_offset + total_successors_in_block <
              cst_dev_params.aux_q_capacity) {
            // no overflow
  
            // grab the aux_q offset
            sh_aux_q_index_block_offset =
                atomicAdd(&lane_counters->aux_q_end, total_successors_in_block);
  
            // We are not overflowing the queue, updating the global values
              IntegerCostType global_min_int_cost = lane_counters->min_int_cost;
              IntegerCostType local_min_int_cost = int_cost_and_index.x;
              // if we found a lower min_cost, update the global value
              if (local_min_int_cost < global_min_int_cost) {
                global_min_int_cost = local_min_int_cost;
                atomicMin(&lane_counters->min_int_cost, global_min_int_cost);
                CostType beam =
                    orderedIntToFloat(adaptive_int_beam_with_validity_index.x);
                IntegerCostType new_int_cutoff = floatToOrderedInt(
                    orderedIntToFloat(local_min_int_cost) + beam);
                atomicMin(&lane_counters->int_cutoff, new_int_cutoff);
              }
              int32 beam_valid_until_idx =
                  adaptive_int_beam_with_validity_index.y;
              if (aux_q_index_block_offset >= beam_valid_until_idx) {
                // This beam is no longer valid. Updating it
                UpdateAdaptiveBeam(
                    cst_dev_params, aux_q_index_block_offset, global_min_int_cost,
                    &adaptive_int_beam_with_validity_index, lane_counters);
              }
          } else {
            // sh_aux_q_index_block_offset is in shared memory
            // its value is currently invalid (overflow)
            // we set it to a special value and use it as a flag to broadcast
            // the fact that we have an overflow and that all threads should exit
            sh_aux_q_index_block_offset = cst_dev_params.aux_q_capacity;
  
            // Setting the flag for the host. It will be used to print a warning
            // to stderr
            lane_counters->q_overflow |= OVERFLOW_AUX_Q;
  
            // We do not jump to end_lane now, because only
            // the first thread (threadIdx.x == 0) is executing this
            // We wait until the end of the divergent branch
          }
        }
  
        // Sync'ing for two reasons :
        // - Broadcasting sh_aux_q_index_block_offset
        // - reusing sh_temp_storage (cf CUB's doc)
        __syncthreads();
        // The only case where we can have that condition met,
        // is if we detected an overflow if the previous lines
        if (sh_aux_q_index_block_offset == cst_dev_params.aux_q_capacity)
          goto end_lane;  // done for this lane
        //
        // If we're executing the following lines it means everything
        // is valid and we are not overflowing the aux_q
        //
        int_cost_and_index.y -= has_successor;  // we want the exclusive sum now
        const int32 aux_q_block_index = int_cost_and_index.y;
        const int32 aux_q_index = sh_aux_q_index_block_offset + aux_q_block_index;
        if (has_successor) {
          // We save the new token to the aux_q
          cst_dev_params.d_aux_q_state_and_cost.lane(ilane)[aux_q_index] = {
              arc_next_state, int_total_cost};
          // Index of the parent token
          // the parent is the token used as input (source of arc)
          // that parent is at index main_q_idx in the GPU memory
          // However, the main_q is emptied before processing a new frame
          // we need to add the offset related to the previous frames index
          // we add cst_dev_params.main_q_global_offset
          const int32 prev_token =
              lane_counters->main_q_global_offset + main_q_idx;
          assert(main_q_idx >= 0 && main_q_idx < cst_dev_params.main_q_capacity);
          cst_dev_params.d_aux_q_info.lane(ilane)[aux_q_index] = {prev_token,
                                                                  arc_idx};
        }
      }
    end_lane:;  // ";" is an empty statement
    }
  }
  
  // post_expand_kernel
  // Called after expand_arcs_kernel
  // Takes care of what needs to be done after an expand_arcs_kernel
  // execution. Mostly resetting the beam (if adaptive beam was triggered,
  // the max_active_ kernels will take care of selecting a good beam),
  // resetting the number of arcs in the main_q (we've processed them all),
  // etc.
  // Threads (1,1,1)
  // Blocks (1, nlanes_used, 1)
  template <bool IS_EMITTING>
  __global__ void post_expand_kernel(DeviceParams cst_dev_params,
                                     KernelParams params) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      LaneCounters *h_lane_counters = cst_dev_params.h_lanes_counters.lane(ilane);
      const int prev_main_q_end = lane_counters->main_q_narcs_and_end.y;
      const int prev_n_extra_prev_tokens =
          lane_counters->main_q_n_extra_prev_tokens;
      const int aux_q_end = lane_counters->aux_q_end;
      CostType min_cost = orderedIntToFloat(lane_counters->min_int_cost);
      // The next step is the contracting step from aux_q to main_q
      // It will need the aux_q_end value. But it will also empty the aux_q
      // We're resetting aux_q_end to 0 now, but we're saving its old value
      // in another place
      lane_counters->post_expand_aux_q_end = aux_q_end;
      h_lane_counters->post_expand_aux_q_end = aux_q_end;       // pinned memory
      h_lane_counters->q_overflow = lane_counters->q_overflow;  // pinned memory
      lane_counters->aux_q_end = 0;
      lane_counters->aux_q_requested = 0;
      // We are done processing those arcs
      lane_counters->main_q_narcs_and_end.x = 0;
      // Resetting the adaptive beam
      lane_counters->adaptive_int_beam_with_validity_index.x =
          lane_counters->int_beam;
      lane_counters->adaptive_int_beam_with_validity_index.y =
          cst_dev_params.adaptive_beam_static_segment;
      CostType beam = orderedIntToFloat(lane_counters->int_beam);
      lane_counters->int_cutoff = floatToOrderedInt(min_cost + beam);
      // If the adaptive beam kicked in, we want to reset the beam
      // the max-active process will take care of selecting the right beam
      if (IS_EMITTING) {
        // the main_q contains the tokens from the previous frame
        // after emitting, we won't use them anymore to create new tokens
        // we reset the main_q
        lane_counters->main_q_narcs_and_end = {0, 0};
        lane_counters->main_q_requested = 0;
        // The main_q was flushed - we need to update the global_offset
        lane_counters->main_q_global_offset += prev_main_q_end;
        if (threadIdx.x == 0 && blockIdx.x == 0)
          lane_counters->main_q_extra_prev_tokens_global_offset +=
              prev_n_extra_prev_tokens;
        // Moving local offset. Tokens created by last expand
        // will be pruned, and survivals will be moved at the end
        // of the main q. Those tokens will be placed after local_offset
        lane_counters->main_q_requested = 0;
        CostType min_cost = orderedIntToFloat(lane_counters->min_int_cost);
        lane_counters->min_histo_cost = min_cost;
        lane_counters->max_histo_cost = min_cost + beam;
        lane_counters->histo_bin_width = beam / (KALDI_CUDA_DECODER_HISTO_NBINS-1);
      } else {
        lane_counters->main_q_local_offset = prev_main_q_end;
        // reset requested to end of queue
        lane_counters->main_q_requested = prev_main_q_end;
      }
    }
  }
  
  __global__ void post_contract_and_preprocess_kernel(DeviceParams cst_dev_params,
                                                      KernelParams params) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      LaneCounters *h_lane_counters = cst_dev_params.h_lanes_counters.lane(ilane);
      int2 main_q_narcs_and_end = lane_counters->main_q_narcs_and_end;
      h_lane_counters->main_q_narcs_and_end =
          main_q_narcs_and_end;                                 // pinned memory
      h_lane_counters->q_overflow = lane_counters->q_overflow;  // pinned memory
      atomicMin(&lane_counters->main_q_n_emitting_tokens, main_q_narcs_and_end.y);
    }
  }
  
  // Meta-kernel (merging preprocess and expand) but only works with 1 CUDA block
  // Used to avoid calling multiple main kernels (such as expand_arcs_kernel)
  // for the tail of non emitting (lots of iterations with small number of arcs)
  //
  // Code is greatly simplified because we use only one CTA / lane
  //
  // Repeat until new queue empty:
  // 1) Preprocess
  // 2) Expand arcs
  //
  // The preprocess stage is not done on the first iteration, because it was
  // already done by the ProcessAndContract kernel. We always call
  // PruneAndPreprocess before calling FinalizeProcessNonemitting
  //
  // At the end, this kernel finalize the computation for current frame,
  // so that it's ready for next ProcessEmitting
  //
  // This kernel works, but can be greatly simplified now.
  __launch_bounds__(KALDI_CUDA_DECODER_LARGEST_1D_BLOCK, 1) __global__
      void finalize_process_non_emitting_kernel(DeviceParams cst_dev_params,
                                                KernelParams params) {
    typedef cub::BlockScan<int2, KALDI_CUDA_DECODER_LARGEST_1D_BLOCK>
        Int2BlockScan;
    typedef cub::BlockScan<int, KALDI_CUDA_DECODER_LARGEST_1D_BLOCK> IntBlockScan;
    __shared__ typename IntBlockScan::TempStorage sh_temp_storage_int_scan;
    __shared__ typename Int2BlockScan::TempStorage sh_temp_storage_int2_scan;
  
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      ChannelCounters *channel_counters =
          cst_dev_params.d_channels_counters.channel(ichannel);
  
      int2 both = lane_counters->main_q_narcs_and_end;
      int32 main_q_narcs = both.x;
      int32 main_q_end = both.y;
      int32 main_q_local_offset = lane_counters->main_q_local_offset;
      const int32 main_q_global_offset = lane_counters->main_q_global_offset;
      // aux_q is empty when this kernel is called
      int32 aux_q_end = 0;
      IntegerCostType int_cutoff = lane_counters->int_cutoff;
      while (main_q_narcs > 0) {
        // Step 1 : ExpandArcs
        KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(offset, thread_idx,
                                                       main_q_narcs) {
          const int32 main_q_arc_idx = offset + thread_idx;
          // For details on how this code works, please refer to comments in
          // expand_arcs
          IntegerCostType total_int_cost = INT_MAX;
          int32 arc_idx;
          StateId arc_next_state;
          int32 main_q_idx;
          if (main_q_arc_idx < main_q_narcs) {
            main_q_idx = binsearch_maxle(
                cst_dev_params.d_main_q_degrees_prefix_sum.channel(ichannel),
                main_q_arc_idx, main_q_local_offset, main_q_end - 1);
  
            const int32 state_first_arc_idx_in_main_q =
                cst_dev_params.d_main_q_degrees_prefix_sum.channel(
                    ichannel)[main_q_idx];
            const int32 arc_offset_start =
                cst_dev_params.d_main_q_arc_offsets.channel(ichannel)[main_q_idx];
            arc_idx = arc_offset_start +
                      (main_q_arc_idx - state_first_arc_idx_in_main_q);
  
            arc_next_state = cst_dev_params.d_arc_nextstates[arc_idx];
            CostType arc_weight = cst_dev_params.d_arc_weights[arc_idx];
            CostType prev_token_cost =
                orderedIntToFloat(cst_dev_params.d_main_q_state_and_cost
                                      .channel(ichannel)[main_q_idx]
                                      .y);
            total_int_cost = floatToOrderedInt(arc_weight + prev_token_cost);
  	  if(total_int_cost < lane_counters->min_int_cost)
              atomicMin(&lane_counters->min_int_cost, total_int_cost);
            if (total_int_cost >= int_cutoff) {
              total_int_cost = INT_MAX;  // above cutoff
            }
          }
          const int32 has_successor = (total_int_cost < INT_MAX) ? 1 : 0;
  
          int32 local_aux_q_idx;
          int32 nsuccessors;
          IntBlockScan(sh_temp_storage_int_scan)
              .ExclusiveSum(has_successor, local_aux_q_idx,
                            nsuccessors);  // aggregate
  
          // Checking if we are overflowing the aux_q
          if ((aux_q_end + nsuccessors) >= cst_dev_params.aux_q_capacity) {
            lane_counters->q_overflow |= OVERFLOW_AUX_Q;
            // nothing to revert in global memory
            goto finalize_lane;
          }
  
          if (has_successor) {
            const int32 aux_q_idx = aux_q_end + local_aux_q_idx;
            const int32 prev_token_idx = main_q_global_offset + main_q_idx;
            cst_dev_params.d_aux_q_state_and_cost.lane(ilane)[aux_q_idx] = {
                arc_next_state, total_int_cost};
            cst_dev_params.d_aux_q_info.lane(ilane)[aux_q_idx] = {prev_token_idx,
                                                                  arc_idx};
          }
          aux_q_end += nsuccessors;
          // sync: reusing sh_temp_storage_scan_int
          __syncthreads();
        }
  
        // Step 2 : PreprocessAndContract
        // Reset for new iteration
        main_q_narcs = 0;
        main_q_local_offset = main_q_end;
        KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(offset, thread_idx,
                                                       aux_q_end) {
          const int32 aux_q_idx = offset + thread_idx;
          int32 degree = 0;
          int32 start = -1;
          StateId token_state;
          IntegerCostType token_int_cost;
          if (aux_q_idx < aux_q_end) {
            int2 both =
                cst_dev_params.d_aux_q_state_and_cost.lane(ilane)[aux_q_idx];
            token_state = both.x;
            token_int_cost = both.y;
            // beam may have changed since generation
            // We are non-emitting in this kernel, using ne offsets
            start = cst_dev_params.d_arc_ne_offsets[token_state];
            int32 end = cst_dev_params.d_arc_ne_offsets[token_state + 1];
            degree = end - start;
          }
          int has_valid_nonpruned_token = (start != -1) ? 1 : 0;
          int2 narcs_and_ntokens_prefix_sum = {degree, has_valid_nonpruned_token};
          int2 aggregate, zero2 = {0, 0};
          Int2BlockScan(sh_temp_storage_int2_scan)
              .ExclusiveScan(narcs_and_ntokens_prefix_sum,
                             narcs_and_ntokens_prefix_sum, zero2, PlusPlus(),
                             aggregate);
          // Checking if we are not overflowing the main_q
          const int32 total_ntokens = aggregate.y;
          if ((main_q_end + total_ntokens) >= cst_dev_params.main_q_capacity) {
            lane_counters->q_overflow |= OVERFLOW_MAIN_Q;
            goto finalize_lane;
          }
          const int32 degree_prefix_sum =
              main_q_narcs + narcs_and_ntokens_prefix_sum.x;
          const int32 degree_sum = aggregate.x;
          main_q_narcs += degree_sum;
          if (has_valid_nonpruned_token) {
            const int32 local_main_q_idx = narcs_and_ntokens_prefix_sum.y;
            const int32 main_q_idx = main_q_end + local_main_q_idx;
  
            cst_dev_params.d_main_q_arc_offsets.channel(ichannel)[main_q_idx] =
                start;
            cst_dev_params.d_main_q_degrees_prefix_sum.channel(
                ichannel)[main_q_idx] = degree_prefix_sum;
            cst_dev_params.d_main_q_state_and_cost.channel(
                ichannel)[main_q_idx] = {token_state, token_int_cost};
            cst_dev_params.d_main_q_info.lane(ilane)[main_q_idx] =
                cst_dev_params.d_aux_q_info.lane(ilane)[aux_q_idx];
            cst_dev_params.d_main_q_acoustic_cost.lane(ilane)[main_q_idx] =
                0.0f;  // we are always nonemitting in this kernel
          }
          main_q_end += total_ntokens;
          __syncthreads();
        }
        aux_q_end = 0;  // aux_q is now empty
      }
  
    finalize_lane:
      if (threadIdx.x == 0) {
        // This main_q is now final for that frame
        lane_counters->main_q_narcs_and_end = {0, main_q_end};
        cst_dev_params.h_lanes_counters.lane(ilane)->main_q_narcs_and_end = {
            0, main_q_end};  // pinned memory
      }
    }
  }
  
  // GetBestCost :
  // Finds all tokens with a cost in [min_cost;min_cost+lattice_beam[
  // Add the final_costs if use_final_probs
  // Does the computation in two steps
  //
  // Step 1: Find the value of min_cost, i.e. the minimum cost in the last token
  // queue
  // (the queue generated by the last frame computed)
  // We set both channel_counters->min_int_cost_and_arg_without_final
  // and channel_counters->min_int_cost_and_arg_with_final
  // One add the final_cost[token.state] before looking for the min
  __global__ void get_best_cost_step1_kernel(DeviceParams cst_dev_params,
                                             KernelParams params,
                                             bool use_final_probs,
                                             CostType fst_zero) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      ChannelCounters *channel_counters =
          cst_dev_params.d_channels_counters.channel(ichannel);
      const int32 main_q_end = channel_counters->prev_main_q_narcs_and_end.y;
      const int32 global_offset = channel_counters->prev_main_q_global_offset;
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(idx, main_q_end) {
        if (idx == 0)
          lane_counters->n_within_lattice_beam =
              0;  // will be used in the next kernel
        const int2 both =
            cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[idx];
        const int token_state = both.x;
        const int token_int_cost = both.y;
        CostType cost = orderedIntToFloat(token_int_cost);
        IntegerCostType int_cost = floatToOrderedInt(cost);
        int32 global_idx = global_offset + idx;
        // We know what is the min cost (without final costs)
        // we just need to have the index of one token with that min cost
  
        if (use_final_probs) {
          const CostType final_cost =
              cst_dev_params.d_fst_final_costs[token_state];
          IntegerCostType int_cost_with_final =
              floatToOrderedInt(cost + final_cost);
          if (final_cost != fst_zero) {
            int2 min_and_arg = {int_cost_with_final,
                                global_idx};  // sort by cost, put it first
            atomicMinI2(&channel_counters->min_int_cost_and_arg_with_final,
                        min_and_arg);
          }
        }
      }
    }
  }
  
  // Step2: Now that step1 found the min_cost (with and without final cost)
  // If at least one final token (token associated with a final fst state)
  // exists in the token queue, AND if use_final_probs is true,
  // We can detect all tokens with a cost within [min_cost;min_cost+lattice_beam]
  // and list them into d_list_final_tokens_in_main_q
  __global__ void get_best_cost_step2_kernel(DeviceParams cst_dev_params,
                                             KernelParams params,
                                             bool use_final_probs,
                                             CostType fst_zero) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      const ChannelCounters *channel_counters =
          cst_dev_params.d_channels_counters.channel(ichannel);
      const int32 main_q_end = channel_counters->prev_main_q_narcs_and_end.y;
      const int32 global_offset = channel_counters->prev_main_q_global_offset;
      const int2 min_int_cost_and_arg_with_final =
          channel_counters->min_int_cost_and_arg_with_final;
      const int2 min_int_cost_and_arg_without_final =
          channel_counters->min_int_cost_and_arg_without_final;
      bool has_reached_final = (min_int_cost_and_arg_with_final.x != INT_MAX);
      // Use final if we want to use final (use_final_probs is true) and if we
      // found a final state in the token list
      bool compute_final = use_final_probs && has_reached_final;
      IntegerCostType min_cost_to_use =
          compute_final ? min_int_cost_and_arg_with_final.x
                        : min_int_cost_and_arg_without_final.x;
  
      // if token.cost < lattice_cutoff, that token will belong in the output
      // lattice
      CostType lattice_cutoff =
          orderedIntToFloat(min_cost_to_use) + cst_dev_params.lattice_beam;
      IntegerCostType lattice_int_cutoff = floatToOrderedInt(lattice_cutoff);
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(idx, main_q_end) {
        // First thread of each lane will move the results into lane counters.
        // That's because we never move channel counters back to host,
        // so we move those values to the lane counters, and those lane counters
        // will be moved to host after this kernel
        if (idx == 0) {
          // The lane counters will be copied to host
          lane_counters->min_int_cost_and_arg =
              compute_final ? min_int_cost_and_arg_with_final
                            : min_int_cost_and_arg_without_final;
          lane_counters->has_reached_final = has_reached_final;
        }
        // Looking for a token with its int_cost < lattice_int_cutoff
        const int2 both =
            cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[idx];
        const int32 token_state = both.x;
        int32 token_int_cost = both.y;
        if (compute_final) {
          const CostType final_cost =
              cst_dev_params.d_fst_final_costs[token_state];
          const CostType token_cost = orderedIntToFloat(token_int_cost);
          // final_cost == fst_zero -> this state is not final
          token_int_cost = (final_cost != fst_zero)
                               ? floatToOrderedInt(token_cost + final_cost)
                               : INT_MAX;
        }
        if (token_int_cost < lattice_int_cutoff) {
          // That token will be included in the lattice (last frame)
          // save it
          int list_idx = atomicAdd(&lane_counters->n_within_lattice_beam, 1);
          cst_dev_params.h_list_final_tokens_in_main_q.lane(ilane)[list_idx] = {
              global_offset + idx, token_int_cost};
        }
      }
    }
  }
  __global__ void get_best_cost_step3_kernel(DeviceParams cst_dev_params,
                                             KernelParams params) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *d_lanes_counters =
          cst_dev_params.d_lanes_counters.lane(ilane);
      LaneCounters *h_lanes_counters =
          cst_dev_params.h_lanes_counters.lane(ilane);
      h_lanes_counters->min_int_cost_and_arg =
          d_lanes_counters->min_int_cost_and_arg;
      h_lanes_counters->has_reached_final = d_lanes_counters->has_reached_final;
      h_lanes_counters->n_within_lattice_beam =
          d_lanes_counters->n_within_lattice_beam;
    }
  }
  // compute_costs_histogram_kernel
  // Used in ApplyMaxActiveAndReduceBeam
  // Compute the histogram of the token.cost in the main_q
  __global__ void compute_costs_histogram_kernel(DeviceParams cst_dev_params,
                                                 KernelParams params,
                                                 bool use_aux_q) {
    const int nlanes = params.nlanes_used;
    typedef cub::BlockHistogram<BinId, KALDI_CUDA_DECODER_1D_BLOCK, 1,
                                KALDI_CUDA_DECODER_HISTO_NBINS + 1>
        BlockHistogram;
    __shared__ typename BlockHistogram::TempStorage temp_storage;
    __shared__ unsigned int smem_histogram[KALDI_CUDA_DECODER_HISTO_NBINS + 1];
  
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      const int32 q_end = use_aux_q ? lane_counters->post_expand_aux_q_end
                                    : lane_counters->main_q_narcs_and_end.y;
      bool compute_max_active = lane_counters->compute_max_active;
      if (!compute_max_active) {
        if (q_end <= cst_dev_params.max_active) continue;  // nothing to do
        // Otherwise let's turn max active on for this frame and lane
        lane_counters->compute_max_active = true;
      }
  
      // Reset local histogram for this lane
      BlockHistogram(temp_storage).InitHistogram(smem_histogram);
      CostType min_histo_cost = lane_counters->min_histo_cost;
      CostType max_histo_cost = lane_counters->max_histo_cost;
      CostType bin_width = lane_counters->histo_bin_width;
  
      // We have a sync inside the loop, keeping all threads alive
      KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(block_offset, thread_idx,
                                                     q_end) {
        const int32 q_idx = block_offset + thread_idx;
        // The last bin is for everything we don't want to count:
        // cost already above the beam, or non-valid tokens
        // It is the default bin
        BinId bin_id[1];
        bin_id[0] = KALDI_CUDA_DECODER_HISTO_NBINS;
        if (q_idx < q_end) {
          IntegerCostType int_cost =
              use_aux_q
                  ? cst_dev_params.d_aux_q_state_and_cost.lane(ilane)[q_idx].y
                  : cst_dev_params.d_main_q_state_and_cost
                        .channel(ichannel)[q_idx]
                        .y;
          CostType cost = orderedIntToFloat(int_cost);
          CostType extra = cost - min_histo_cost;
  	if(extra <= 0.0f) 
  		bin_id[0] = 0;
    	else if (extra < max_histo_cost) {
            bin_id[0] = (BinId)__fdiv_rd(extra, bin_width)+1; // +1 because first bin is cost < min_histo_cost
          }
        }
        BlockHistogram(temp_storage).Composite(bin_id, smem_histogram);  // sync
        __syncthreads();  // reusing temp_storage
      }
  
      // Not using the macros 1D_LOOP because that loop is only within a CTA
      for (int32 bin_id_w = threadIdx.x;
           bin_id_w < KALDI_CUDA_DECODER_HISTO_NBINS;
           bin_id_w += KALDI_CUDA_DECODER_1D_BLOCK) {
        // Writing the local histo to global
        // We don't care about the last bin (cf above)
        int32 s_count = (int32)smem_histogram[bin_id_w];
        atomicAdd(&cst_dev_params.d_histograms.lane(ilane)[bin_id_w], s_count);
      }
      // Making sure we're done reading from smem
      __syncthreads();
    }
  }
  
  // update_beam_using_histogram_kernel
  // used in ApplyMaxActiveAndReduceBeam
  // uses the histogram computed in compute_costs_histogram_kernel
  // to find where to cut (where to set the beam)
  // to keep only ~max_active_ tokens.
  // Important: use only one CTA per lane
  __global__ void update_beam_using_histogram_kernel(DeviceParams cst_dev_params,
                                                     KernelParams params,
                                                     bool use_aux_q) {
    typedef cub::BlockScan<int, KALDI_CUDA_DECODER_1D_BLOCK> BlockScan;
    __shared__ typename BlockScan::TempStorage temp_storage;
  
    const int nlanes = params.nlanes_used;
    const int max_active = cst_dev_params.max_active;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      bool compute_max_active = lane_counters->compute_max_active;
      if (!compute_max_active) continue;  // nothing to do
      CostType beam = orderedIntToFloat(lane_counters->int_beam);
      CostType min_histo_cost = lane_counters->min_histo_cost;
      CostType bin_width = lane_counters->histo_bin_width;
      // We now have our histogram of the token costs (computed in the previous
      // kernel)
      // Each thread i is responsible for a bin i, with that bin containing ni
      // tokens.
      // We compute the prefix sum of those ni, ending up for each thread with
      // si=sum[i=1..i](ni)
      // If the thread i detects that si < max_active_ and s[i+1] >= max_active_,
      // then we will cut the beam at
      // the cost of the bin [i+1]
      //
      // Assert : one thread in a CTA is responsible for at most one bin
      // we will not iterate over bins
      assert(KALDI_CUDA_DECODER_HISTO_NBINS < KALDI_CUDA_DECODER_1D_BLOCK);
      int bin_id = threadIdx.x;
      int val = 0;
      if (bin_id < KALDI_CUDA_DECODER_HISTO_NBINS) 
        val = cst_dev_params.d_histograms.lane(ilane)[bin_id];
      
      int prefix_sum;
      BlockScan(temp_storage).ExclusiveSum(val, prefix_sum);
  
      if (prefix_sum < max_active && (prefix_sum + val) >= max_active) {
        // We found our new beam regarding min_histo_cost
        // Howevever, the current min_cost could be lower than min_histo_cost
        // we need to add that diff to the new beam
        CostType new_beam_for_histo_min_cost = bin_width * bin_id;
        CostType current_min_cost = orderedIntToFloat(lane_counters->min_int_cost);
        CostType new_beam = (min_histo_cost - current_min_cost) + new_beam_for_histo_min_cost;
        IntegerCostType new_int_beam = floatToOrderedInt(new_beam);
        // Saving our new beam for this lane
        lane_counters->int_beam = new_int_beam;
        lane_counters->adaptive_int_beam_with_validity_index.x = new_int_beam;
        lane_counters->int_cutoff = floatToOrderedInt(current_min_cost + new_beam);
      }
    }
  }
  
  //
  // PostProcessingMainQueue kernels.
  // all the following kernels are called when postprocessing a frame
  //
  
  // Filling hashmap values with the tokens that we have in the main queue
  // We do that because multiple tokens associated with the same FST state
  // (but with different arc_idx) can exist in the main_q. We need to detect
  // that situation, count them, detect what the min_cost for that FST state is.
  // It is done using a hashmap
  __global__ void fill_hashmap_with_main_q_kernel(DeviceParams cst_dev_params,
                                                  KernelParams params) {
    // Operator for the prefix sum inside the CUDA block
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      ChannelCounters *channel_counters =
          cst_dev_params.d_channels_counters.channel(ichannel);
  
      const int32 main_q_end = lane_counters->main_q_narcs_and_end.y;
      int32 min_int_cost = lane_counters->min_int_cost;
      CostType min_cost = orderedIntToFloat(min_int_cost);
      const int32 global_offset = channel_counters->prev_main_q_global_offset;
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(main_q_idx, main_q_end) {
        // Position of considered token in the main_q
        if (main_q_idx < main_q_end) {
          int2 both = cst_dev_params.d_main_q_state_and_cost.channel(
              ichannel)[main_q_idx];
          StateId token_state = both.x;
          IntegerCostType token_int_cost = both.y;
          if (min_int_cost == token_int_cost) {
            // remove offset = min_cost, set it to 0 explicitely
            token_int_cost = floatToOrderedInt(0.0f);
            channel_counters->min_int_cost_and_arg_without_final = {
                token_int_cost, global_offset + main_q_idx};
            lane_counters->prev_arg_min_int_cost = main_q_idx;
          } else {
            // remove offset = min_cost
            CostType token_cost = orderedIntToFloat(token_int_cost) - min_cost;
            token_int_cost = floatToOrderedInt(token_cost);
          }
          int local_idx, hash_idx;
          hashmap_insert_or_aggregate(cst_dev_params.d_hashmap_values.lane(ilane),
                                      token_state, token_int_cost, main_q_idx,
                                      cst_dev_params.hashmap_capacity, &local_idx,
                                      &hash_idx);
          cst_dev_params.d_main_q_n_extra_prev_tokens_local_idx.lane(
              ilane)[main_q_idx] = local_idx;
          cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[main_q_idx].y =
              token_int_cost;
          // If we have the min, saving its index for get best cost and the min
          // cost estimate of the next frame
  
          // Saving where that token.state ended up in the hashmap
          // false = this token is not the representative of this state
          // We will update representing_state once we know more (in the next
          // kernel)
          // We first need to add all tokens to the hashmap. Which will be the
          // case when
          // this kernel returns.
          SetFSTStateHashIndex(
              hash_idx, false,
              &cst_dev_params.d_main_q_state_hash_idx.lane(ilane)[main_q_idx]);
        }
  
        if (main_q_idx == 0) {
          lane_counters->int_cutoff = floatToOrderedInt(
              orderedIntToFloat(lane_counters->int_cutoff) - min_cost);
        }
      }
    }
  }
  
  // preprocess_and_list_extra_prev_tokens_kernel_step[i] kernels
  // Called in PostProcessingMainQueue
  // They do two things:
  // - do the "emitting preprocessing". I.e. doing the preprocessing necessary for
  // the future ExpandArcsEmitting that may be done next (if the current frame is
  // not the last one)
  // It consists of filling the d_main_q_degrees_prefix_sum of the emitting arc
  // degrees of the tokens + setting d_main_q_arc_offsets
  // - when we have multiple tokens associated with the same FST state S, we will
  // list them in d_main_q_extra_prev_tokens. We need to know where to put them in
  // that array,
  // so we'll compute a prefix_sum also to compute those indexes. We'll then save
  // the location of each extra tokens list (its offset and size in
  // d_main_q_extra_prev_tokens),
  // and save it into d_main_q_info for later lattice processing
  //
  // First step : Reading the hashmap, detecting which token is representative for
  // each FST state, which is decided by fill_hashmap_with_main_q_kernel()
  // (we pick one of the best ones, with the best ones being the ones with the
  // lowest cost)
  // this representative will be responsible for K tokens, with K being the number
  // of tokens associated with that FST state. We only considers the cases where K
  // > 1,
  // because if K == 1, then we will not store that token in the special list
  // d_main_q_extra_prev_tokens
  // Each representative is also the only token that will propagate emitting arcs
  // for that FST state. Because a representative has the min_cost for that FST
  // state, it is enough to only propagate
  // that one
  // Each representative counts the number of emitting arcs it is responsible for,
  // and we will compute the prefix sum of the arc degrees
  __global__ void emitting_preprocess_and_list_extra_prev_tokens_step1_kernel(
      DeviceParams cst_dev_params, KernelParams params) {
    // Operator for the prefix sum inside the CUDA block
    typedef cub::BlockScan<int2, KALDI_CUDA_DECODER_1D_BLOCK> BlockScan;
    __shared__ typename BlockScan::TempStorage sh_temp_storage;
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      const LaneCounters *lane_counters =
          cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 main_q_end = lane_counters->main_q_narcs_and_end.y;
      // Final cutoff from last ExpandArc execution
      // The cutoff can have decreased since moving tokens to the main_q
      // min_cost cannot be lower than before (we only did non-emitting phases
      // since then)
      // but the adaptive beam may have lowered the beam
      const IntegerCostType int_cutoff = lane_counters->int_cutoff;
      // Keeping all threads in CTA alive
      // We'll __syncthreads()
      KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(block_offset, thread_idx,
                                                     main_q_end) {
        // We'll take care of the token at index main_q_idx
        const int32 main_q_idx = block_offset + thread_idx;
        const int32 ichannel = lane_counters->channel_to_compute;
        // If that token is the representative of its FST state (token.next_state)
        // The representative of a FST state is the token with the lowest
        // token.cost for that FST state
        // If multiple tokens have token1.cost == token2.cost ==
        // min_cost_for_that_state, then one is picked (first come first serve,
        // was done in fill_hashmap_with_main_q_kernel)
        bool representing_state = false;
        // Number of emitting arcs for that token
        // Only the token representative of that FST state can have degree > 0
        int32 degree = 0;
        // If that token is representative of a FST state S,
        // and if multiple tokens are associated with that state S,
        // then n_extra_prev_token will contain their count
        int32 n_extra_prev_token = 0;
        if (main_q_idx < main_q_end) {
          int2 both = cst_dev_params.d_main_q_state_and_cost.channel(
              ichannel)[main_q_idx];
          StateId token_state = both.x;
          IntegerCostType token_int_cost = both.y;
          // Loading info about token.next_state. Is there multiple tokens for
          // that state ?
          // How many ? What's the min token.cost for that state ?
          int32 hash_idx;    // we saved the hash_idx after inserting
          bool bool_buffer;  // will always be false. We just need it to call the
                             // function
          GetFSTStateHashIndex(
              cst_dev_params.d_main_q_state_hash_idx.lane(ilane)[main_q_idx],
              &hash_idx, &bool_buffer);
          HashmapValueT h_val =
              cst_dev_params.d_hashmap_values.lane(ilane)[hash_idx];
          // Token index of one of the token which the lowest token.cost for that
          // state
          uint32_t state_best_int_cost_argmin;
  	GetArgFromPackedArgminUInt64(h_val.min_and_argmin_int_cost_u64, &state_best_int_cost_argmin);
  
          // Checking if we're the representative of that state
          representing_state = (main_q_idx == state_best_int_cost_argmin);
          // Saving the hash_idx of that fst state + if we're responsible for that
          // state
          SetFSTStateHashIndex(
              hash_idx, representing_state,
              &cst_dev_params.d_main_q_state_hash_idx.lane(ilane)[main_q_idx]);
  
          // One of the best token for that state will represent that state in the
          // next frame
          if (representing_state) {
            if (token_int_cost < int_cutoff) {
              // Next step is emitting (next frame), using emitting offsets
              const int32 start = cst_dev_params.d_arc_e_offsets[token_state];
              const int32 end = cst_dev_params.d_arc_e_offsets[token_state + 1];
              degree = end - start;
              // Saving the start offset for the expand kernel
              // avoid a new random memory access
              cst_dev_params.d_main_q_arc_offsets.channel(ichannel)[main_q_idx] =
                  start;
            }
            // If that FST state has only one token associated to it, we store
            // that token directly in
            // d_main_q_info (its original place)
            // We only move it into the d_main_q_extra_prev_tokens list if
            // multiple tokens are associated to that state
            n_extra_prev_token = (h_val.count > 1) ? (h_val.count) : 0;
          }
        }
  
        // Computing a local prefix sum inside that CUDA block
        // Others kernels will take care of adding the necessary offset to those
        // local prefix sums
        int2 zeroi2 = {0, 0};
        int2 vali2 = {degree, n_extra_prev_token};
        int2 aggi2;
        BlockScan(sh_temp_storage)
            .ExclusiveScan(vali2, aggi2, zeroi2, PlusPlus());
        int32 degree_local_prefix_sum = aggi2.x;
        int32 n_extra_prev_token_prefix_sum = aggi2.y;
  
        if (main_q_idx < main_q_end) {
          // This is not the final global prefix sum
          // Other kernels will add the necessary offset
          cst_dev_params.d_main_q_degrees_prefix_sum.channel(
              ichannel)[main_q_idx] = degree_local_prefix_sum;
          cst_dev_params.d_main_q_extra_prev_tokens_prefix_sum.lane(
              ilane)[main_q_idx] = n_extra_prev_token_prefix_sum;
        }
  
        if (KALDI_CUDA_DECODER_IS_LAST_1D_THREAD()) {
          // Saving the local sum of degrees of that CUDA block
          // That's necessary to compute the global offset of that CUDA block,
          // and that offset is what we need to transform the local prefix sum
          // into a global prefix sum
          const int local_sum_index = block_offset / KALDI_CUDA_DECODER_1D_BLOCK;
          // the prefix sum was exclusive, adding missing value
          const int degree_inclusive_sum = degree_local_prefix_sum + degree;
          const int n_extra_prev_tokens_inclusive_sum =
              n_extra_prev_token_prefix_sum + n_extra_prev_token;
          cst_dev_params.d_main_q_block_sums_prefix_sum.lane(
              ilane)[local_sum_index] = {degree_inclusive_sum,
                                         n_extra_prev_tokens_inclusive_sum};
        }
  
        // Synchronization because:
        // - we may need to reuse sh_temp_storage if the for loop iterates (cf
        // CUB's doc)
        __syncthreads();
      }
    }
  }
  
  // In step1, we've computed the local (CTA-wide) prefix sums. We also have the
  // local sums of each individual CTAs
  // In this kernel, we will compute the offset of each CTA in the global prefix
  // sum. We will then add those offsets in step3
  // Only one CTA / lane
  __global__ void emitting_preprocess_and_list_extra_prev_tokens_step2_kernel(
      DeviceParams cst_dev_params, KernelParams params) {
    typedef cub::BlockScan<int2, KALDI_CUDA_DECODER_1D_BLOCK> BlockScan;
    __shared__ typename BlockScan::TempStorage sh_temp_storage;
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int main_q_end = lane_counters->main_q_narcs_and_end.y;
      const int ntiles = KALDI_CUDA_DECODER_DIV_ROUND_UP(
          main_q_end, KALDI_CUDA_DECODER_1D_BLOCK);
      // Using block_offset loop to keep entire CTA alive (we're going to use
      // __syncthreads in CUB)
      int2 sum_so_far = {0, 0};
      KALDI_CUDA_DECODER_1D_BLOCK_OFFSET_KERNEL_LOOP(offset, thread_idx, ntiles) {
        const int32 itile = offset + thread_idx;
        const int2 zeroi2 = {0, 0};
        const int2 val =
            (itile < ntiles)
                ? cst_dev_params.d_main_q_block_sums_prefix_sum.lane(ilane)[itile]
                : zeroi2;
  
        int2 prefix_sum, sum;
        BlockScan(sh_temp_storage)
            .ExclusiveScan(val, prefix_sum, zeroi2, PlusPlus(), sum);
        PlusPlus pp;
        prefix_sum = pp(prefix_sum, sum_so_far);
        sum_so_far = pp(sum_so_far, sum);
        if (itile < ntiles) {
          cst_dev_params.d_main_q_block_sums_prefix_sum.lane(ilane)[itile] =
              prefix_sum;
        }
        if (itile == (ntiles - 1)) {
          const int32 total_narcs = prefix_sum.x + val.x;
          const int32 total_n_extra_prev_tokens = prefix_sum.y + val.y;
          lane_counters->main_q_narcs_and_end.x = total_narcs;
          lane_counters->main_q_n_extra_prev_tokens = total_n_extra_prev_tokens;
          assert(total_n_extra_prev_tokens >= 0 &&
                 total_n_extra_prev_tokens <= main_q_end);
        }
      }
    }
  }
  
  // Step3: Uses the CTA offsets computed in step2 to transform the CTA-wide
  // prefix sums to global prefix sums
  // The representative of each FST states saves into the hashmap the location of
  // the extra_prev_tokens of that state
  // in d_main_q_extra_prev_tokens. That way each extra tokens will know where to
  // write itself in the next kernel.
  __global__ void emitting_preprocess_and_list_extra_prev_tokens_step3_kernel(
  		DeviceParams cst_dev_params, KernelParams params) {
  	const int nlanes = params.nlanes_used;
  	KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
  		const LaneCounters *lane_counters =
  			cst_dev_params.d_lanes_counters.lane(ilane);
  		const int32 ichannel = lane_counters->channel_to_compute;
  		const int main_q_end = lane_counters->main_q_narcs_and_end.y;
  		KALDI_CUDA_DECODER_1D_KERNEL_LOOP(main_q_idx, main_q_end) {
  			const int32 local_sum_idx = main_q_idx / KALDI_CUDA_DECODER_1D_BLOCK;
  			const int2 local_sum_offset =
  				cst_dev_params.d_main_q_block_sums_prefix_sum.lane(
  						ilane)[local_sum_idx];
  			cst_dev_params.d_main_q_degrees_prefix_sum.channel(
  					ichannel)[main_q_idx] += local_sum_offset.x;
  			int extra_prev_tokens_offset =
  				cst_dev_params.d_main_q_extra_prev_tokens_prefix_sum.lane(
  						ilane)[main_q_idx] +
  				local_sum_offset.y;
  			// Loading the hash index associate with token.state
  			// If representative, store the location of the extra prev tokens list for
  			// that state in the hashmap
  			bool is_representative;
  			int32 hash_idx;
  			GetFSTStateHashIndex(
  					cst_dev_params.d_main_q_state_hash_idx.lane(ilane)[main_q_idx],
  					&hash_idx, &is_representative);
                          if (is_representative) {
                            HashmapValueT &val =
                                cst_dev_params.d_hashmap_values.lane(
                                    ilane)[hash_idx];
                            uint32_t min;
                            GetMinFromPackedArgminUInt64(
                                val.min_and_argmin_int_cost_u64, &min);
                            unsigned long long new_pack;
                            PackArgminInUInt64(min, extra_prev_tokens_offset,
                                               &new_pack);
                            val.min_and_argmin_int_cost_u64 = new_pack;
                          }
  		}
  	}
  }
  
  // Step4: We now know where to store our extra prev tokens in
  // d_main_q_extra_prev_tokens.
  // We will now move the tokens that need to be moved (when multiple tokens are
  // associated to the same FST state)
  // into d_main_q_extra_prev_tokens. In d_main_q_info, we will store the location
  // of that list [offset,size]
  // so that when backtracking, when we read d_main_q_info[token_idx], we know
  // where to look to have the list
  // of the same-state tokens
  // It is the last step of the
  // emitting_preprocess_and_list_extra_prev_tokens_step[i]_kernel pipeline
  __global__ void emitting_preprocess_and_list_extra_prev_tokens_step4_kernel(
      DeviceParams cst_dev_params, KernelParams params) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      const LaneCounters *lane_counters =
          cst_dev_params.d_lanes_counters.lane(ilane);
      const int32 ichannel = lane_counters->channel_to_compute;
      const int main_q_end = lane_counters->main_q_narcs_and_end.y;
      // Previous frames have filled d_main_q_extra_prev_tokens.
      // d_main_q_extra_prev_tokens was then flushed to host. We want to set the
      // global
      // (global in the sense "for all frames") offset on where to read it the
      // h_all_tokens_extra_prev_tokens_ on host.
      // adding the main_q_extra_prev_tokens_global_offset for that
      const int prev_global_idx =
          lane_counters->main_q_extra_prev_tokens_global_offset;
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(main_q_idx, main_q_end) {
        // We'll take care of token at main_q_idx
        // Loading hashmap information about token.state
        bool is_representative;
        int32 hash_idx;
        GetFSTStateHashIndex(
            cst_dev_params.d_main_q_state_hash_idx.lane(ilane)[main_q_idx],
            &hash_idx, &is_representative);
  
        HashmapValueT val = cst_dev_params.d_hashmap_values.lane(ilane)[hash_idx];
        // How many tokens are associated with that fst state token.state
        int same_count = val.count;
        bool must_move_to_extra_prev_tokens = (same_count > 1);
        if (must_move_to_extra_prev_tokens) {
          // Moving to the extra_prev_tokens list.
          // Some of those tokens have an extra cost (compared to the best cost
          // for that FST state)
          // Generating and saving that extra cost. We will use it when generating
          // the lattice.
          CostType token_cost = orderedIntToFloat(
              cst_dev_params.d_main_q_state_and_cost.channel(ichannel)[main_q_idx]
                  .y);
  	uint32_t best_int_cost;
          // Where to write this state list in d_main_q_extra_prev_tokens
  	uint32_t extra_prev_tokens_offset;
  	unsigned long long pack = val.min_and_argmin_int_cost_u64;
  	GetMinFromPackedArgminUInt64(pack, &best_int_cost);
  	GetArgFromPackedArgminUInt64(pack, &extra_prev_tokens_offset);
          CostType best_cost = orderedIntToFloat((int)best_int_cost);
          CostType extra_cost = token_cost - best_cost;
  	assert(!is_representative || extra_cost == 0.0f);
          // Loading the token to be moved
          InfoToken inf_tok =
              cst_dev_params.d_main_q_info.lane(ilane)[main_q_idx];
          CostType acoustic_cost =
              cst_dev_params.d_main_q_acoustic_cost.lane(ilane)[main_q_idx];
          // Place of that specific token in the extra_prev_tokens sublist of that
          // specific FST state
          int32 local_idx =
              cst_dev_params.d_main_q_n_extra_prev_tokens_local_idx.lane(
                  ilane)[main_q_idx];
          // Saving the location of the extra prev tokens for that state into that
          // InfoToken
          SetSameFSTStateTokensList(
              prev_global_idx + extra_prev_tokens_offset, same_count,
              &cst_dev_params.d_main_q_info.lane(ilane)[main_q_idx]);
          // Where to write this token in d_main_q_extra_prev_tokens
          int32 list_idx = extra_prev_tokens_offset + local_idx;
          // Moving token. Also saving extra_cost
          cst_dev_params.d_main_q_extra_prev_tokens.lane(ilane)[list_idx] =
              inf_tok;
          cst_dev_params.d_main_q_extra_and_acoustic_cost.lane(
              ilane)[list_idx] = {extra_cost, acoustic_cost};
          assert(inf_tok.prev_token >= (lane_counters->main_q_global_offset -
                                        cst_dev_params.main_q_capacity) &&
                 inf_tok.prev_token <=
                     (lane_counters->main_q_global_offset + main_q_end));
        }
      }
    }
  }
  
  // Clear the hashmaps after use
  // Each element in the map has a representative in the main_q
  // Everyone of those representatives has the responsability to reset their
  // corresponding value in the hashmap
  // Once this kernel returns, the hashmaps are cleared
  __global__ void clear_hashmap_kernel(DeviceParams cst_dev_params,
                                       KernelParams params) {
    const int nlanes = params.nlanes_used;
    KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
      LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
      const int main_q_end = lane_counters->main_q_narcs_and_end.y;
      KALDI_CUDA_DECODER_1D_KERNEL_LOOP(main_q_idx, main_q_end) {
        bool is_representative;
        int32 hash_idx;
        GetFSTStateHashIndex(
            cst_dev_params.d_main_q_state_hash_idx.lane(ilane)[main_q_idx],
            &hash_idx, &is_representative);
        // Representative owns a state. Each representative resets its associated
        // token.state
        // in the hashmap
        if (is_representative) {
          cst_dev_params.d_hashmap_values.lane(ilane)[hash_idx] =
              KALDI_CUDA_DECODER_HASHMAP_NO_VAL;  // clear
        }
      }
    }
  }
  
  // Kernels wrappers
  
  void SaveChannelsStateFromLanesKernel(const dim3 &grid, const dim3 &block,
                                        const cudaStream_t &st,
                                        const DeviceParams &cst_dev_params,
                                        const KernelParams &kernel_params) {
    save_channels_state_from_lanes_kernel<<<grid, block, 0, st>>>(cst_dev_params,
                                                                  kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void LoadChannelsStateInLanesKernel(const dim3 &grid, const dim3 &block,
                                      const cudaStream_t &st,
                                      const DeviceParams &cst_dev_params,
                                      const KernelParams &kernel_params) {
    load_channels_state_in_lanes_kernel<<<grid, block, 0, st>>>(cst_dev_params,
                                                                kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void InitDecodingOnDeviceKernel(const dim3 &grid, const dim3 &block,
                                  const cudaStream_t &st,
                                  const DeviceParams &cst_dev_params,
                                  const KernelParams &kernel_params) {
    init_decoding_on_device_kernel<<<grid, block, 0, st>>>(cst_dev_params,
                                                           kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void InitializeInitialLaneKernel(const dim3 &grid, const dim3 &block,
                                   const cudaStream_t &st,
                                   const DeviceParams &cst_dev_params) {
    initialize_initial_lane_kernel<<<grid, block, 0, st>>>(cst_dev_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void ResetForFrameAndEstimateCutoffKernel(const dim3 &grid, const dim3 &block,
                                            const cudaStream_t &st,
                                            const DeviceParams &cst_dev_params,
                                            const KernelParams &kernel_params) {
    reset_for_frame_and_estimate_cutoff_kernel<<<grid, block, 0, st>>>(
        cst_dev_params, kernel_params);
  }
  
  template <bool IS_EMITTING>
  void ExpandArcsKernel(const dim3 &grid, const dim3 &block,
                        const cudaStream_t &st,
                        const DeviceParams &cst_dev_params,
                        const KernelParams &kernel_params) {
    expand_arcs_kernel<IS_EMITTING><<<grid, block, 0, st>>>(cst_dev_params,
                                                            kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  template <bool IS_EMITTING>
  void PostExpandKernel(const dim3 &grid, const dim3 &block,
                        const cudaStream_t &st,
                        const DeviceParams &cst_dev_params,
                        const KernelParams &kernel_params) {
    post_expand_kernel<IS_EMITTING><<<grid, block, 0, st>>>(cst_dev_params,
                                                            kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void PostContractAndPreprocessKernel(const dim3 &grid, const dim3 &block,
                                       const cudaStream_t &st,
                                       const DeviceParams &cst_dev_params,
                                       const KernelParams &kernel_params) {
    post_contract_and_preprocess_kernel<<<grid, block, 0, st>>>(cst_dev_params,
                                                                kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void NonEmittingPreprocessAndContractKernel(const dim3 &grid, const dim3 &block,
                                              const cudaStream_t &st,
                                              const DeviceParams &cst_dev_params,
                                              const KernelParams &kernel_params) {
    nonemitting_preprocess_and_contract_kernel<<<grid, block, 0, st>>>(
        cst_dev_params, kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void FillHashmapWithMainQKernel(const dim3 &grid, const dim3 &block,
                                  const cudaStream_t &st,
                                  const DeviceParams &cst_dev_params,
                                  const KernelParams &kernel_params) {
    fill_hashmap_with_main_q_kernel<<<grid, block, 0, st>>>(cst_dev_params,
                                                            kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void EmittingPreprocessAndListExtraPrevTokensStep1Kernel(
      const dim3 &grid, const dim3 &block, const cudaStream_t &st,
      const DeviceParams &cst_dev_params, const KernelParams &kernel_params) {
    emitting_preprocess_and_list_extra_prev_tokens_step1_kernel<<<grid, block, 0,
                                                                  st>>>(
        cst_dev_params, kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void EmittingPreprocessAndListExtraPrevTokensStep2Kernel(
      const dim3 &grid, const dim3 &block, const cudaStream_t &st,
      const DeviceParams &cst_dev_params, const KernelParams &kernel_params) {
    emitting_preprocess_and_list_extra_prev_tokens_step2_kernel<<<grid, block, 0,
                                                                  st>>>(
        cst_dev_params, kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void EmittingPreprocessAndListExtraPrevTokensStep3Kernel(
      const dim3 &grid, const dim3 &block, const cudaStream_t &st,
      const DeviceParams &cst_dev_params, const KernelParams &kernel_params) {
    emitting_preprocess_and_list_extra_prev_tokens_step3_kernel<<<grid, block, 0,
                                                                  st>>>(
        cst_dev_params, kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void EmittingPreprocessAndListExtraPrevTokensStep4Kernel(
      const dim3 &grid, const dim3 &block, const cudaStream_t &st,
      const DeviceParams &cst_dev_params, const KernelParams &kernel_params) {
    emitting_preprocess_and_list_extra_prev_tokens_step4_kernel<<<grid, block, 0,
                                                                  st>>>(
        cst_dev_params, kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void ComputeLaneOffsetsKernel(const dim3 &grid, const dim3 &block,
                                const cudaStream_t &st,
                                const DeviceParams &cst_dev_params,
                                const KernelParams &kernel_params) {
    compute_lane_offsets_kernel<<<grid, block, 0, st>>>(cst_dev_params,
                                                        kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  template <typename T>
  void ConcatenateLanesDataKernel(const dim3 &grid, const dim3 &block,
                                  const cudaStream_t &st,
                                  const DeviceParams &cst_dev_params,
                                  const KernelParams &kernel_params,
                                  const LaneMatrixView<T> &src, T *concat,
                                  int32 *lane_offsets) {
    concatenate_lanes_data_kernel<<<grid, block, 0, st>>>(
        cst_dev_params, kernel_params, src, concat, lane_offsets);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void InitHashmapKernel(const dim3 &grid, const dim3 &block,
                         const cudaStream_t &st,
                         const DeviceParams &cst_dev_params) {
    init_hashmap_kernel<<<grid, block, 0, st>>>(cst_dev_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void ClearHashmapKernel(const dim3 &grid, const dim3 &block,
                          const cudaStream_t &st,
                          const DeviceParams &cst_dev_params,
                          const KernelParams &kernel_params) {
    clear_hashmap_kernel<<<grid, block, 0, st>>>(cst_dev_params, kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void ComputeCostsHistogramKernel(const dim3 &grid, const dim3 &block,
                                   const cudaStream_t &st,
                                   const DeviceParams &cst_dev_params,
                                   const KernelParams &kernel_params,
                                   bool use_aux_q) {
    compute_costs_histogram_kernel<<<grid, block, 0, st>>>(
        cst_dev_params, kernel_params, use_aux_q);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void UpdateBeamUsingHistogramKernel(const dim3 &grid, const dim3 &block,
                                      const cudaStream_t &st,
                                      const DeviceParams &cst_dev_params,
                                      const KernelParams &kernel_params,
                                      bool use_aux_q) {
    update_beam_using_histogram_kernel<<<grid, block, 0, st>>>(
        cst_dev_params, kernel_params, use_aux_q);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void FinalizeProcessNonEmittingKernel(const dim3 &grid, const dim3 &block,
                                        const cudaStream_t &st,
                                        const DeviceParams &cst_dev_params,
                                        const KernelParams &kernel_params) {
    finalize_process_non_emitting_kernel<<<grid, block, 0, st>>>(cst_dev_params,
                                                                 kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void GetBestCostStep1Kernel(const dim3 &grid, const dim3 &block,
                              const cudaStream_t &st,
                              const DeviceParams &cst_dev_params,
                              const KernelParams &kernel_params, bool isfinal,
                              CostType fst_zero) {
    get_best_cost_step1_kernel<<<grid, block, 0, st>>>(
        cst_dev_params, kernel_params, isfinal, fst_zero);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void GetBestCostStep2Kernel(const dim3 &grid, const dim3 &block,
                              const cudaStream_t &st,
                              const DeviceParams &cst_dev_params,
                              const KernelParams &kernel_params, bool isfinal,
                              CostType fst_zero) {
    get_best_cost_step2_kernel<<<grid, block, 0, st>>>(
        cst_dev_params, kernel_params, isfinal, fst_zero);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  void GetBestCostStep3Kernel(const dim3 &grid, const dim3 &block,
                              const cudaStream_t &st,
                              const DeviceParams &cst_dev_params,
                              const KernelParams &kernel_params) {
    get_best_cost_step3_kernel<<<grid, block, 0, st>>>(cst_dev_params,
                                                       kernel_params);
    KALDI_DECODER_CUDA_CHECK_ERROR();
  }
  
  template void ExpandArcsKernel<true>(const dim3 &grid, const dim3 &block,
                                       const cudaStream_t &st,
                                       const DeviceParams &cst_dev_params,
                                       const KernelParams &params);
  template void ExpandArcsKernel<false>(const dim3 &grid, const dim3 &block,
                                        const cudaStream_t &st,
                                        const DeviceParams &cst_dev_params,
                                        const KernelParams &params);
  template void PostExpandKernel<true>(const dim3 &grid, const dim3 &block,
                                       const cudaStream_t &st,
                                       const DeviceParams &cst_dev_params,
                                       const KernelParams &params);
  template void PostExpandKernel<false>(const dim3 &grid, const dim3 &block,
                                        const cudaStream_t &st,
                                        const DeviceParams &cst_dev_params,
                                        const KernelParams &params);
  
  template void ConcatenateLanesDataKernel<InfoToken>(
      const dim3 &grid, const dim3 &block, const cudaStream_t &st,
      const DeviceParams &cst_dev_params, const KernelParams &params,
      const LaneMatrixView<InfoToken> &src, InfoToken *concat,
      int32 *lane_offsets);
  
  template void ConcatenateLanesDataKernel<CostType>(
      const dim3 &grid, const dim3 &block, const cudaStream_t &st,
      const DeviceParams &cst_dev_params, const KernelParams &params,
      const LaneMatrixView<CostType> &src, CostType *concat, int32 *lane_offsets);
  
  template void ConcatenateLanesDataKernel<float2>(
      const dim3 &grid, const dim3 &block, const cudaStream_t &st,
      const DeviceParams &cst_dev_params, const KernelParams &params,
      const LaneMatrixView<float2> &src, float2 *concat, int32 *lane_offsets);
  
  template void ConcatenateLanesDataKernel<int32>(
      const dim3 &grid, const dim3 &block, const cudaStream_t &st,
      const DeviceParams &cst_dev_params, const KernelParams &params,
      const LaneMatrixView<int32> &src, int32 *concat, int32 *lane_offsets);
  
  }  // end namespace cuda_decoder
  }  // end namespace kaldi