int32 nlanes, int32 nchannels)

: fst_(fst),

nlanes_(nlanes),

nchannels_(nchannels),

channel_lock_(nchannels + 1),

extra_cost_min_delta_(0.0f),

thread_pool_(NULL),

n_threads_used_(0),

n_h2h_task_not_done_(0),

n_init_decoding_h2h_task_not_done_(0),

h2h_threads_running_(true) {

ReadConfig(config);

// Static asserts on constants

CheckStaticAsserts();

// Runtime asserts

KALDI_ASSERT(nlanes > 0);

KALDI_ASSERT(nchannels > 0);

KALDI_ASSERT(nlanes_ <= nchannels_);

// All GPU work in decoder will be sent to compute_st_

cudaStreamCreate(&compute_st_);

// Copies D2H of tokens for storage on host are done on

// copy_st_, in parallel with compute_st_

cudaStreamCreate(&copy_st_);

// For all the allocating/initializing process

// We create a special channel

// containing the exact state a channel should have when starting a new decode

// It contains fst.Start(), the non-emitting tokens created by fst.Start(),

// and all the data used by the decoder.

// When calling InitDecoding() on a new channel, we simply clone this special

// channel into that new channel

++nchannels_; // adding the special initial channel

init_channel_id_ = nchannels_ - 1; // Using last one as init_channel_params

AllocateHostData();

AllocateDeviceData();

AllocateDeviceKernelParams();

InitDeviceParams();

InitHostData();

InitDeviceData();

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventCreate(&nnet3_done_evt_));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventCreate(&d2h_copy_acoustic_evt_));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventCreate(&d2h_copy_infotoken_evt_));

KALDI_DECODER_CUDA_API_CHECK_ERROR(

cudaEventCreate(&d2h_copy_extra_prev_tokens_evt_));

KALDI_DECODER_CUDA_API_CHECK_ERROR(

cudaEventCreate(&concatenated_data_ready_evt_));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaEventCreate(&lane_offsets_ready_evt_));

ComputeInitialChannel();

--nchannels_; // removing the special initial channel from the count

// Making sure that everything is ready to use

cudaStreamSynchronize(compute_st_);

KALDI_DECODER_CUDA_CHECK_ERROR();

CudaDecoderConfig config = cst_config; // deep copy

// Sets the missing values using other values

config.ComputeConfig();

default_beam_ = config.default_beam;

lattice_beam_ = config.lattice_beam;

ntokens_pre_allocated_ = config.ntokens_pre_allocated;

max_active_ = config.max_active;

aux_q_capacity_ = config.aux_q_capacity;

main_q_capacity_ = config.main_q_capacity;

KALDI_ASSERT(default_beam_ >= 0.0f);

KALDI_ASSERT(lattice_beam_ >= 0.0f);

KALDI_ASSERT(ntokens_pre_allocated_ >= 0);

KALDI_ASSERT(max_active_ > 0);

KALDI_ASSERT(main_q_capacity_ > 0);

KALDI_ASSERT(aux_q_capacity_ >= main_q_capacity_);

hashmap_capacity_ =

KALDI_CUDA_DECODER_HASHMAP_CAPACITY_FACTOR * main_q_capacity_;

d_channels_counters_.Resize(nchannels_, 1);

d_lanes_counters_.Resize(

nlanes_ + 1,

1); // +1 because we sometimes need last+1 value (for offsets)

d_main_q_state_and_cost_.Resize(nchannels_, main_q_capacity_);

d_main_q_info_.Resize(nlanes_, main_q_capacity_);

d_aux_q_state_and_cost_.Resize(nlanes_, aux_q_capacity_);

d_aux_q_info_.Resize(nlanes_, aux_q_capacity_);

d_main_q_degrees_prefix_sum_.Resize(nchannels_, main_q_capacity_);

d_histograms_.Resize(nlanes_, KALDI_CUDA_DECODER_HISTO_NBINS);

d_main_q_extra_prev_tokens_prefix_sum_.Resize(nlanes_, main_q_capacity_);

d_main_q_n_extra_prev_tokens_local_idx_.Resize(nlanes_, main_q_capacity_);

d_main_q_state_hash_idx_.Resize(nlanes_, main_q_capacity_);

d_main_q_extra_prev_tokens_.Resize(nlanes_, main_q_capacity_);

d_main_q_extra_and_acoustic_cost_.Resize(nlanes_, main_q_capacity_);

d_main_q_block_sums_prefix_sum_.Resize(

nlanes_, KALDI_CUDA_DECODER_DIV_ROUND_UP(main_q_capacity_,

KALDI_CUDA_DECODER_1D_BLOCK) +

1);

d_main_q_arc_offsets_.Resize(nchannels_, main_q_capacity_);

d_hashmap_values_.Resize(nlanes_, hashmap_capacity_);

d_main_q_acoustic_cost_.Resize(nlanes_, main_q_capacity_);

d_extra_and_acoustic_cost_concat_matrix_.Resize(nlanes_, main_q_capacity_);

d_acoustic_cost_concat_matrix_.Resize(nlanes_, main_q_capacity_);

d_infotoken_concat_matrix_.Resize(nlanes_, main_q_capacity_);

d_extra_prev_tokens_concat_matrix_.Resize(nlanes_, main_q_capacity_);

// Reusing data from aux_q. Those two are never used at the same time

// d_list_final_tokens_in_main_q_ is used in GetBestPath.

// the aux_q is used in AdvanceDecoding

h_list_final_tokens_in_main_q_.Resize(nlanes_, main_q_capacity_);

d_extra_prev_tokens_concat_ = d_extra_prev_tokens_concat_matrix_.lane(0);

d_extra_and_acoustic_cost_concat_ =

d_extra_and_acoustic_cost_concat_matrix_.lane(0);

d_acoustic_cost_concat_ = d_acoustic_cost_concat_matrix_.lane(0);

d_infotoken_concat_ = d_infotoken_concat_matrix_.lane(0);

channel_to_compute_.resize(nlanes_);

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(

&h_extra_and_acoustic_cost_concat_,

nlanes_ * main_q_capacity_ * sizeof(*h_extra_and_acoustic_cost_concat_)));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(

&h_acoustic_cost_concat_,

nlanes_ * main_q_capacity_ * sizeof(*h_acoustic_cost_concat_)));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(

&h_extra_prev_tokens_concat_,

nlanes_ * main_q_capacity_ * sizeof(*h_extra_prev_tokens_concat_)));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(

&h_infotoken_concat_,

nlanes_ * main_q_capacity_ * sizeof(*h_infotoken_concat_)));

KALDI_DECODER_CUDA_API_CHECK_ERROR(

cudaMallocHost(&h_extra_and_acoustic_cost_concat_tmp_,

nlanes_ * main_q_capacity_ *

sizeof(*h_extra_and_acoustic_cost_concat_tmp_)));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(

&h_acoustic_cost_concat_tmp_,

nlanes_ * main_q_capacity_ * sizeof(*h_acoustic_cost_concat_tmp_)));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(

&h_extra_prev_tokens_concat_tmp_,

nlanes_ * main_q_capacity_ * sizeof(*h_extra_prev_tokens_concat_tmp_)));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(

&h_infotoken_concat_tmp_,

nlanes_ * main_q_capacity_ * sizeof(*h_infotoken_concat_tmp_)));

h_lanes_counters_.Resize(

nlanes_ + 1,

1); // +1 because we sometimes need last+1 value (for offsets)

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMallocHost(

&h_channels_counters_, nchannels_ * sizeof(*h_channels_counters_)));

h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_.resize(nchannels_);

h_all_tokens_acoustic_cost_.resize(nchannels_);

h_all_tokens_extra_prev_tokens_.resize(nchannels_);

h_all_tokens_info_.resize(nchannels_);

for (int32 ichannel = 0; ichannel < nchannels_; ++ichannel) {

h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_[ichannel].reserve(

ntokens_pre_allocated_);

h_all_tokens_acoustic_cost_[ichannel].reserve(ntokens_pre_allocated_);

h_all_tokens_info_[ichannel].reserve(ntokens_pre_allocated_);

}

h_main_q_end_lane_offsets_.resize(nlanes_ + 1);

h_emitting_main_q_end_lane_offsets_.resize(nlanes_ + 1);

h_n_extra_prev_tokens_lane_offsets_.resize(nlanes_ + 1);

frame_offsets_.resize(nchannels_);

num_frames_decoded_.resize(nchannels_, -1);

lanes2channels_todo_.reserve(nlanes_);

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

h_all_final_tokens_list_.resize(nchannels_);

h_all_has_reached_final_.resize(nchannels_);

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemsetAsync(

d_channels_counters_.MutableData(), 0,

nchannels_ * sizeof(*d_channels_counters_.MutableData()), compute_st_));

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemsetAsync(

d_lanes_counters_.MutableData(), 0,

nlanes_ * sizeof(*d_lanes_counters_.MutableData()), compute_st_));

InitHashmapKernel(KaldiCudaDecoderNumBlocks(hashmap_capacity_, nlanes_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,

*h_device_params_);

KALDI_DECODER_CUDA_CHECK_ERROR();

// Setting Kernel Params

// Sent to cuda kernels by copy

// Making sure we'll be able to send it to the kernels

KALDI_ASSERT((sizeof(KernelParams) + sizeof(DeviceParams)) <

KALDI_CUDA_DECODER_MAX_KERNEL_ARGUMENTS_BYTE_SIZE);

h_device_params_->d_channels_counters = d_channels_counters_.GetView();

h_device_params_->d_lanes_counters = d_lanes_counters_.GetView();

h_device_params_->h_lanes_counters = h_lanes_counters_.GetView();

h_device_params_->d_main_q_state_and_cost =

d_main_q_state_and_cost_.GetView();

h_device_params_->d_main_q_info = d_main_q_info_.GetView();

h_device_params_->d_aux_q_state_and_cost = d_aux_q_state_and_cost_.GetView();

h_device_params_->d_main_q_extra_and_acoustic_cost =

d_main_q_extra_and_acoustic_cost_.GetView();

h_device_params_->d_main_q_acoustic_cost = d_main_q_acoustic_cost_.GetView();

h_device_params_->d_aux_q_info = d_aux_q_info_.GetView();

h_device_params_->d_main_q_degrees_prefix_sum =

d_main_q_degrees_prefix_sum_.GetView();

h_device_params_->d_main_q_block_sums_prefix_sum =

d_main_q_block_sums_prefix_sum_.GetView();

h_device_params_->d_main_q_state_hash_idx =

d_main_q_state_hash_idx_.GetView();

h_device_params_->d_main_q_extra_prev_tokens_prefix_sum =

d_main_q_extra_prev_tokens_prefix_sum_.GetView();

h_device_params_->d_main_q_n_extra_prev_tokens_local_idx =

d_main_q_n_extra_prev_tokens_local_idx_.GetView();

h_device_params_->d_main_q_extra_prev_tokens =

d_main_q_extra_prev_tokens_.GetView();

h_device_params_->d_main_q_arc_offsets = d_main_q_arc_offsets_.GetView();

h_device_params_->d_hashmap_values = d_hashmap_values_.GetView();

h_device_params_->d_histograms = d_histograms_.GetView();

h_device_params_->d_arc_e_offsets = fst_.d_e_offsets_;

h_device_params_->d_arc_ne_offsets = fst_.d_ne_offsets_;

h_device_params_->d_arc_pdf_ilabels = fst_.d_arc_pdf_ilabels_;

h_device_params_->d_arc_weights = fst_.d_arc_weights_;

h_device_params_->d_arc_nextstates = fst_.d_arc_nextstates_;

h_device_params_->d_fst_final_costs = fst_.d_final_;

h_device_params_->default_beam = default_beam_;

h_device_params_->lattice_beam = lattice_beam_;

h_device_params_->main_q_capacity = main_q_capacity_;

h_device_params_->aux_q_capacity = aux_q_capacity_;

h_device_params_->init_channel_id = init_channel_id_;

h_device_params_->max_nlanes = nlanes_;

h_device_params_->nstates = fst_.num_states_;

h_device_params_->init_state = fst_.Start();

KALDI_ASSERT(h_device_params_->init_state != fst::kNoStateId);

h_device_params_->init_cost = StdWeight::One().Value();

h_device_params_->hashmap_capacity = hashmap_capacity_;

h_device_params_->max_active = max_active_;

// For the first static_beam_q_length elements of the queue, we will keep the

// beam static

adaptive_beam_static_segment_ =

aux_q_capacity_ / KALDI_CUDA_DECODER_ADAPTIVE_BEAM_STATIC_SEGMENT;

// For the last adaptive_beam_q_length elements of the queue, we will decrease

// the beam, segment by segment

// For more information, please refer to the definition of GetAdaptiveBeam in

// cuda-decoder-kernels.cu

int32 adaptive_beam_q_length =

(aux_q_capacity_ - adaptive_beam_static_segment_);

int32 adaptive_beam_bin_width =

adaptive_beam_q_length / KALDI_CUDA_DECODER_ADAPTIVE_BEAM_NSTEPS;

h_device_params_->adaptive_beam_static_segment =

adaptive_beam_static_segment_;

h_device_params_->adaptive_beam_bin_width = adaptive_beam_bin_width;

// Reusing aux_q memory to list final states in GetLattice

// Those cannot be used at the same time

h_device_params_->h_list_final_tokens_in_main_q =

h_list_final_tokens_in_main_q_.GetView();

KALDI_ASSERT(nlanes_ > 0);

const int32 ilane = 0;

KALDI_ASSERT(ilane == 0);

// Following kernels working channel_id

std::vector<ChannelId> channels = {init_channel_id_};

SetChannelsInKernelParams(channels); // not calling LoadChannelsStateToLanes,

// init_channel_id_ is a special case

h_lanes_counters_.lane(ilane)->channel_to_compute = init_channel_id_;

cudaMemcpyAsync(d_lanes_counters_.MutableData(), h_lanes_counters_.lane(0),

1 * sizeof(*h_lanes_counters_.lane(0)),

cudaMemcpyHostToDevice, compute_st_);

h_lanes_counters_.lane(ilane)->main_q_narcs_and_end.y = 0;

// Adding the start state to the initial token queue

InitializeInitialLaneKernel(KaldiCudaDecoderNumBlocks(1, 1),

KALDI_CUDA_DECODER_ONE_THREAD_BLOCK, compute_st_,

*h_device_params_);

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

PruneAndPreprocess();

FinalizeProcessNonEmittingKernel(

KaldiCudaDecoderNumBlocks(1, 1), KALDI_CUDA_DECODER_LARGEST_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_);

CopyLaneCountersToHostSync();

PostProcessingMainQueue();

CopyLaneCountersToHostSync();

const int32 main_q_end =

h_lanes_counters_.lane(ilane)->main_q_narcs_and_end.y;

KALDI_ASSERT(main_q_end > 0);

// Moving all data linked to init_channel_id_ to host

// that data will be cloned to other channels when calling InitDecoding

CopyMainQueueDataToHost();

SaveChannelsStateFromLanes();

KALDI_ASSERT(

h_channels_counters_[init_channel_id_].prev_main_q_narcs_and_end.x > 0);

KALDI_ASSERT(

h_channels_counters_[init_channel_id_].prev_main_q_narcs_and_end.y > 0);

// Cloning the init_channel_id_ channel into all channels in the channels vec

const int nlanes_used = channels.size();

// Getting *h_kernel_params ready to use

LoadChannelsStateToLanes(channels);

cudaMemcpyAsync(d_lanes_counters_.MutableData(), h_lanes_counters_.lane(0),

nlanes_used_ * sizeof(*h_lanes_counters_.lane(0)),

cudaMemcpyHostToDevice, compute_st_);

// Size of the initial main_q

ChannelCounters &init_channel_counters =

h_channels_counters_[init_channel_id_];

const int32 init_main_q_size =

init_channel_counters.prev_main_q_narcs_and_end.y;

KALDI_ASSERT(init_main_q_size > 0);

// Getting the channels ready to compute new utterances

InitDecodingOnDeviceKernel(

KaldiCudaDecoderNumBlocks(init_main_q_size, nlanes_used),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_, *h_device_params_,

*h_kernel_params_);

{

std::lock_guard<std::mutex> n_h2h_not_done_lk(

n_init_decoding_h2h_task_not_done_mutex_);

n_init_decoding_h2h_task_not_done_ += channels.size();

}

for (ChannelId ichannel : channels) {

ChannelCounters &channel_counters = h_channels_counters_[ichannel];

channel_counters.prev_main_q_narcs_and_end =

init_channel_counters.prev_main_q_narcs_and_end;

channel_counters.prev_main_q_n_extra_prev_tokens =

init_channel_counters.prev_main_q_n_extra_prev_tokens;

channel_counters.prev_main_q_global_offset = 0;

channel_counters.prev_main_q_extra_prev_tokens_global_offset = 0;

channel_counters.prev_beam = default_beam_;

int32 n_initial_tokens = h_all_tokens_info_[init_channel_id_].size();

num_frames_decoded_[ichannel] = 0;

h_channels_counters_[ichannel] = h_channels_counters_[init_channel_id_];

h_all_argmin_cost_[ichannel] = {-1, 0.0f};

frame_offsets_[ichannel].clear();

frame_offsets_[ichannel].push_back(n_initial_tokens);

if (thread_pool_)

thread_pool_->enqueue(THREAD_POOL_HIGH_PRIORITY,

&CudaDecoder::InitDecodingH2HCopies, this,

ichannel);

else

InitDecodingH2HCopies(ichannel);

}

// Tokens from initial main_q needed on host

std::unique_lock<std::mutex> channel_lk(channel_lock_[ichannel]);

// Deep copy

h_all_tokens_info_[ichannel] = h_all_tokens_info_[init_channel_id_];

h_all_tokens_acoustic_cost_[ichannel] =

h_all_tokens_acoustic_cost_[init_channel_id_];

h_all_tokens_extra_prev_tokens_[ichannel] =

h_all_tokens_extra_prev_tokens_[init_channel_id_];

h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_[ichannel] =

h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_[init_channel_id_];

bool all_done;

{

std::lock_guard<std::mutex> lk_not_done(

n_init_decoding_h2h_task_not_done_mutex_);

all_done = (--n_init_decoding_h2h_task_not_done_ == 0);

}

if (all_done) {

init_decoding_h2h_done_.notify_all();

}

const std::vector<ChannelId> &channels) {

// Setting that channels configuration in kernel_params

SetChannelsInKernelParams(channels);

KALDI_ASSERT(nlanes_used_ > 0);

for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {

const ChannelId ichannel = channel_to_compute_[ilane];

ChannelCounters &channel_counters = h_channels_counters_[ichannel];

LaneCounters &lane_counters = *h_lanes_counters_.lane(ilane);

lane_counters.channel_to_compute = ichannel;

lane_counters.main_q_narcs_and_end =

channel_counters.prev_main_q_narcs_and_end;

lane_counters.main_q_n_extra_prev_tokens =

channel_counters.prev_main_q_n_extra_prev_tokens;

int32 int_beam = floatToOrderedIntHost(channel_counters.prev_beam);

lane_counters.int_beam = int_beam;

lane_counters.adaptive_int_beam_with_validity_index.x = int_beam;

lane_counters.adaptive_int_beam_with_validity_index.y =

adaptive_beam_static_segment_;

lane_counters.main_q_global_offset =

channel_counters.prev_main_q_global_offset;

lane_counters.main_q_extra_prev_tokens_global_offset =

channel_counters.prev_main_q_extra_prev_tokens_global_offset;

lane_counters.min_int_cost =

channel_counters.min_int_cost_and_arg_without_final.x;

lane_counters.prev_arg_min_int_cost =

channel_counters.min_int_cost_and_arg_without_final.y;

}

KALDI_ASSERT(nlanes_used_ > 0);

for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {

const ChannelId ichannel = channel_to_compute_[ilane];

ChannelCounters &channel_counters = h_channels_counters_[ichannel];

LaneCounters &lane_counters = *h_lanes_counters_.lane(ilane);

channel_counters.prev_main_q_narcs_and_end =

lane_counters.main_q_narcs_and_end;

channel_counters.prev_main_q_extra_prev_tokens_global_offset =

lane_counters.main_q_extra_prev_tokens_global_offset;

channel_counters.prev_main_q_global_offset =

lane_counters.main_q_global_offset;

channel_counters.prev_main_q_n_extra_prev_tokens =

lane_counters.main_q_n_extra_prev_tokens;

channel_counters.prev_beam = orderedIntToFloatHost(lane_counters.int_beam);

channel_counters.min_int_cost_and_arg_without_final = {

lane_counters.min_int_cost, lane_counters.prev_arg_min_int_cost};

}

SaveChannelsStateFromLanesKernel(KaldiCudaDecoderNumBlocks(1, nlanes_used_),

KALDI_CUDA_DECODER_ONE_THREAD_BLOCK,

compute_st_, *h_device_params_,

*h_kernel_params_);

ResetChannelsInKernelParams();

std::function<int32(const LaneCounters &)> func) {

int32 max_val = 0;

for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {

const int32 val = func(*h_lanes_counters_.lane(ilane));

max_val = std::max(max_val, val);

}

return max_val;

cudaMemcpyAsync(h_lanes_counters_.lane(0), d_lanes_counters_.MutableData(),

nlanes_used_ * sizeof(*h_lanes_counters_.lane(0)),

cudaMemcpyDeviceToHost, compute_st_);

const int32 ilane, const int32 ichannel,

const std::vector<int32> &lanes_offsets, T *h_concat,

std::vector<std::vector<T>> *vecvec) {

// Unpacking the concatenated vector into individual channel storage

int32 beg = lanes_offsets[ilane];

int32 end = lanes_offsets[ilane + 1];

auto &vec = (*vecvec)[ichannel];

vec.insert(vec.end(), h_concat + beg, h_concat + end);

// Checking if we should activate max active for the current frame

// once it is active, it is active for the whole frame (for all non emitting

// iterations)

// If at least one lane queue is bigger than max_active,

// we'll apply a topk on that queue (k=max_active_)

bool use_aux_q = (queue_id == AUX_Q);

ComputeCostsHistogramKernel(KaldiCudaDecoderNumBlocks(nlanes_used_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,

*h_device_params_, *h_kernel_params_, use_aux_q);

UpdateBeamUsingHistogramKernel(

KaldiCudaDecoderNumBlocks(1, nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_, use_aux_q);

const std::vector<ChannelId> &channels,

std::vector<CudaDecodableInterface *> &decodables, int32 max_num_frames) {

int32 nframes_to_decode = INT_MAX;

// std::vector<int> debug_ntokens;

// std::vector<int> debug_narcs;

for (int32 ilane = 0; ilane < nlanes_used_; ++ilane) {

const ChannelId ichannel = channels[ilane];

const int32 num_frames_decoded = num_frames_decoded_[ichannel];

KALDI_ASSERT(num_frames_decoded >= 0 &&

"You must call InitDecoding() before AdvanceDecoding()");

int32 num_frames_ready = decodables[ilane]->NumFramesReady();

// num_frames_ready must be >= num_frames_decoded, or else

// the number of frames ready must have decreased (which doesn't

// make sense) or the decodable object changed between calls

// (which isn't allowed).

KALDI_ASSERT(num_frames_ready >= num_frames_decoded);

int32 channel_nframes_to_decode = num_frames_ready - num_frames_decoded;

nframes_to_decode = std::min(nframes_to_decode, channel_nframes_to_decode);

}

if (max_num_frames >= 0)

nframes_to_decode = std::min(nframes_to_decode, max_num_frames);

return nframes_to_decode;

ExpandArcsKernel<true>(KaldiCudaDecoderNumBlocks(nlanes_used_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,

*h_device_params_, *h_kernel_params_);

// Updating a few counters, like resetting aux_q_end to 0...

// true is for IS_EMITTING

PostExpandKernel<true>(KaldiCudaDecoderNumBlocks(1, nlanes_used_),

KALDI_CUDA_DECODER_ONE_THREAD_BLOCK, compute_st_,

*h_device_params_, *h_kernel_params_);

// false is for non emitting

ExpandArcsKernel<false>(KaldiCudaDecoderNumBlocks(nlanes_used_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,

*h_device_params_, *h_kernel_params_);

// false is for non emitting

PostExpandKernel<false>(KaldiCudaDecoderNumBlocks(1, nlanes_used_),

KALDI_CUDA_DECODER_ONE_THREAD_BLOCK, compute_st_,

*h_device_params_, *h_kernel_params_);

NonEmittingPreprocessAndContractKernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_);

PostContractAndPreprocessKernel(KaldiCudaDecoderNumBlocks(1, nlanes_used_),

KALDI_CUDA_DECODER_ONE_THREAD_BLOCK,

compute_st_, *h_device_params_,

*h_kernel_params_);

ApplyMaxActiveAndReduceBeam(MAIN_Q);

FillHashmapWithMainQKernel(KaldiCudaDecoderNumBlocks(nlanes_used_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,

*h_device_params_, *h_kernel_params_);

EmittingPreprocessAndListExtraPrevTokensStep1Kernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_);

EmittingPreprocessAndListExtraPrevTokensStep2Kernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_);

// Step2 wrote main_q_n_extra_prev_tokens

// it was the last value missing to compute the lanes offsets

// doing it now

ComputeLaneOffsetsKernel(KaldiCudaDecoderNumBlocks(1, 1), // One CTA

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,

*h_device_params_, *h_kernel_params_);

cudaEventRecord(lane_offsets_ready_evt_, compute_st_);

EmittingPreprocessAndListExtraPrevTokensStep3Kernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_);

EmittingPreprocessAndListExtraPrevTokensStep4Kernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_);

ClearHashmapKernel(KaldiCudaDecoderNumBlocks(nlanes_used_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_,

*h_device_params_, *h_kernel_params_);

cudaEventRecord(concatenated_data_ready_evt_, compute_st_);

cudaStreamWaitEvent(copy_st_, concatenated_data_ready_evt_,

0); // the copies on copy_st will wait on compute_st_

cudaEventSynchronize(

lane_offsets_ready_evt_); // we need the total size of each segments

LaunchD2HCopies();

// Making sure the previous H2H copies are done

WaitForInitDecodingH2HCopies();

WaitForH2HCopies();

std::swap(h_extra_and_acoustic_cost_concat_tmp_,

h_extra_and_acoustic_cost_concat_);

std::swap(h_infotoken_concat_tmp_, h_infotoken_concat_);

std::swap(h_acoustic_cost_concat_tmp_, h_acoustic_cost_concat_);

std::swap(h_extra_prev_tokens_concat_tmp_, h_extra_prev_tokens_concat_);

// Saving the offsets computed previously

lanes2channels_todo_.clear();

for (int32 ilane = 0; ilane < (nlanes_used_ + 1); ++ilane) {

h_emitting_main_q_end_lane_offsets_[ilane] =

h_lanes_counters_.lane(ilane)->main_q_n_emitting_tokens_lane_offset;

h_main_q_end_lane_offsets_[ilane] =

h_lanes_counters_.lane(ilane)->main_q_end_lane_offset;

h_n_extra_prev_tokens_lane_offsets_[ilane] =

h_lanes_counters_.lane(ilane)->main_q_n_extra_prev_tokens_lane_offset;

lanes2channels_todo_.push_back(channel_to_compute_[ilane]);

}

LaunchH2HCopies();

// Last offset = total

int32 nelements_acoustic_costs = h_lanes_counters_.lane(nlanes_used_)

->main_q_n_emitting_tokens_lane_offset;

// Moving the d_concat to h_concat (host), async

if (nelements_acoustic_costs > 0) {

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpyAsync(

h_acoustic_cost_concat_tmp_, d_acoustic_cost_concat_,

nelements_acoustic_costs * sizeof(*d_acoustic_cost_concat_),

cudaMemcpyDeviceToHost, copy_st_));

}

cudaEventRecord(d2h_copy_acoustic_evt_, copy_st_);

int32 nelements_infotoken =

h_lanes_counters_.lane(nlanes_used_)->main_q_end_lane_offset;

if (nelements_infotoken > 0) {

KALDI_DECODER_CUDA_API_CHECK_ERROR(

cudaMemcpyAsync(h_infotoken_concat_tmp_, d_infotoken_concat_,

nelements_infotoken * sizeof(*d_infotoken_concat_),

cudaMemcpyDeviceToHost, copy_st_));

}

cudaEventRecord(d2h_copy_infotoken_evt_, copy_st_);

int32 nelements_extra_prev_tokens =

h_lanes_counters_.lane(nlanes_used_)

->main_q_n_extra_prev_tokens_lane_offset;

if (nelements_extra_prev_tokens > 0) {

KALDI_DECODER_CUDA_API_CHECK_ERROR(cudaMemcpyAsync(

h_extra_prev_tokens_concat_tmp_, d_extra_prev_tokens_concat_,

nelements_extra_prev_tokens * sizeof(*d_extra_prev_tokens_concat_),

cudaMemcpyDeviceToHost, copy_st_));

KALDI_DECODER_CUDA_API_CHECK_ERROR(

cudaMemcpyAsync(h_extra_and_acoustic_cost_concat_tmp_,

d_extra_and_acoustic_cost_concat_,

nelements_extra_prev_tokens *

sizeof(*d_extra_and_acoustic_cost_concat_),

cudaMemcpyDeviceToHost, copy_st_));

}

cudaEventRecord(d2h_copy_extra_prev_tokens_evt_, copy_st_);

ConcatenateLanesDataKernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_,

h_device_params_->d_main_q_acoustic_cost, d_acoustic_cost_concat_,

&d_lanes_counters_.lane(0)->main_q_n_emitting_tokens_lane_offset);

ConcatenateLanesDataKernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_,

h_device_params_->d_main_q_info, d_infotoken_concat_,

&d_lanes_counters_.lane(0)->main_q_end_lane_offset);

ConcatenateLanesDataKernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_,

h_device_params_->d_main_q_extra_prev_tokens, d_extra_prev_tokens_concat_,

&d_lanes_counters_.lane(0)->main_q_n_extra_prev_tokens_lane_offset);

ConcatenateLanesDataKernel(

KaldiCudaDecoderNumBlocks(nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_,

h_device_params_->d_main_q_extra_and_acoustic_cost,

d_extra_and_acoustic_cost_concat_,

&d_lanes_counters_.lane(0)->main_q_n_extra_prev_tokens_lane_offset);

const std::vector<ChannelId> &channels,

std::vector<CudaDecodableInterface *> &decodables, int32 max_num_frames) {

if (channels.size() == 0) return; // nothing to do

// Context switch : Loading the channels state in lanes

LoadChannelsStateToLanes(channels);

KALDI_ASSERT(nlanes_used_ > 0);

// We'll decode nframes_to_decode, such as all channels have at least that

// number

// of frames available

int32 nframes_to_decode =

NumFramesToDecode(channels, decodables, max_num_frames);

// Looping over the frames that we will compute

for (int32 iframe = 0; iframe < nframes_to_decode; ++iframe) {

// Loglikelihoods from the acoustic model

// Setting the loglikelihoods pointers for that frame

for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {

ChannelId ichannel = channel_to_compute_[ilane];

int32 frame = num_frames_decoded_[ichannel];

h_lanes_counters_.lane(ilane)->loglikelihoods =

decodables[ilane]->GetLogLikelihoodsCudaPointer(frame);

}

cudaMemcpyAsync(d_lanes_counters_.MutableData(), h_lanes_counters_.lane(0),

nlanes_used_ * sizeof(*h_lanes_counters_.lane(0)),

cudaMemcpyHostToDevice, compute_st_);

// compute_st_ will wait for nnet3 to complete

cudaEventRecord(nnet3_done_evt_, cudaStreamPerThread);

cudaStreamWaitEvent(compute_st_, nnet3_done_evt_, 0);

// Estimating cutoff using argmin from last frame

ResetForFrameAndEstimateCutoffKernel(

KaldiCudaDecoderNumBlocks(1, nlanes_used_), KALDI_CUDA_DECODER_1D_BLOCK,

compute_st_, *h_device_params_, *h_kernel_params_);

// Reset max active status. If necessary, ApplyMaxActiveAndReduceBeam will

// switch it back on

compute_max_active_ = false;

// Processing emitting arcs. We've done the preprocess stage at the end of

// the previous frame

ExpandArcsEmitting();

// We'll loop until we have a small enough number of non-emitting arcs

// in the token queue. We'll then break the loop

for (int i = 0; i < KALDI_CUDA_DECODER_N_NON_EMITTING_MAIN_ITERATIONS;

++i) {

// If one of the aux_q contains more than max_active_ tokens,

// we'll reduce the beam to only keep max_active_ tokens

ApplyMaxActiveAndReduceBeam(AUX_Q);

// Prune the aux_q. Apply the latest beam (using the one from

// ApplyMaxActiveAndReduceBeam if triggered)

// move the survival tokens to the main queue

// and do the preprocessing necessary for the next ExpandArcs

PruneAndPreprocess();

// "heavy duty" kernel for non-emitting. The long tail of small

// non-emitting iterations will be done in

// FinalizeProcessNonEmittingKernel

ExpandArcsNonEmitting();

}

ApplyMaxActiveAndReduceBeam(AUX_Q);

PruneAndPreprocess();

// Finalizing process non emitting. Takes care of the long tail,

// the final iterations with a small numbers of arcs. Do the work inside a

// single CTA (per lane),

FinalizeProcessNonEmittingKernel(KaldiCudaDecoderNumBlocks(1, nlanes_used_),

KALDI_CUDA_DECODER_LARGEST_1D_BLOCK,

compute_st_, *h_device_params_,

*h_kernel_params_);

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

// Post processing the tokens for that frame

// - do the preprocess necessary for the next emitting expand (will happen

// with next frame)

// - if a state S has more than one token associated to it, generate the

// list of those tokens

// It allows to backtrack efficiently in GetRawLattice

// - compute the extra costs

PostProcessingMainQueue();

// Waiting on previous d2h before writing on same device memory

cudaStreamWaitEvent(compute_st_, d2h_copy_extra_prev_tokens_evt_, 0);

// Concatenating the data that will be moved to host into large arrays

ConcatenateData();

// Copying the final lane counters for that frame

CopyLaneCountersToHostSync();

CheckOverflow();

// Moving the data necessary for GetRawLattice/GetBestPath back to host for

// storage

CopyMainQueueDataToHost();

for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {

const ChannelId ichannel = channel_to_compute_[ilane];

// We're done processing that frame

++num_frames_decoded_[ichannel];

const int32 main_q_end =

h_lanes_counters_.lane(ilane)->main_q_narcs_and_end.y;

// Saving frame offsets for GetRawLattice

frame_offsets_[ichannel].push_back(frame_offsets_[ichannel].back() +

main_q_end);

}

}

SaveChannelsStateFromLanes();

for (LaneId ilane = 0; ilane < nlanes_used_; ++ilane) {

LaneCounters *lane_counters = h_lanes_counters_.lane(ilane);

bool q_overflow = lane_counters->q_overflow;

if (q_overflow != OVERFLOW_NONE) {

// An overflow was prevented in a kernel

// The algorithm can still go on but quality of the result can be reduced

// (less tokens were generated)

if ((q_overflow & OVERFLOW_MAIN_Q) == OVERFLOW_MAIN_Q) {

// overflowed main_q

KALDI_WARN

<< "Preventing overflow of main_q. Continuing "

<< "execution but the quality of the output may be decreased. "

<< "To prevent this from happening, please increase the parameter "

"--main-q-capacity"

<< " and/or decrease --max-active";

}

if ((q_overflow & OVERFLOW_AUX_Q) == OVERFLOW_AUX_Q) {

// overflowed aux_q

KALDI_WARN

<< "Preventing overflow of aux_q. Continuing "

<< "execution but the quality of the output may be decreased. "

<< "To prevent this from happening, please increase the parameter "

"--aux-q-capacity"

<< " and/or decrease --beam";

}

KALDI_ASSERT(lane_counters->main_q_narcs_and_end.y < main_q_capacity_);

KALDI_ASSERT(lane_counters->main_q_narcs_and_end.x >= 0);

KALDI_ASSERT(lane_counters->main_q_narcs_and_end.y >= 0);

KALDI_ASSERT(lane_counters->post_expand_aux_q_end < aux_q_capacity_);

KALDI_ASSERT(lane_counters->post_expand_aux_q_end >= 0);

KALDI_ASSERT(lane_counters->aux_q_end < aux_q_capacity_);

KALDI_ASSERT(lane_counters->aux_q_end >= 0);

}

}

bool use_final_costs,

std::vector<std::pair<int32, CostType>> *argmins,

std::vector<std::vector<std::pair<int, float>>>

*list_finals_token_idx_and_cost,

std::vector<bool> *has_reached_final) {

if (channels.size() == 0) return;

// Getting the lanes ready to be used with those channels

LoadChannelsStateToLanes(channels);

cudaMemcpyAsync(d_lanes_counters_.MutableData(), h_lanes_counters_.lane(0),

nlanes_used_ * sizeof(*h_lanes_counters_.lane(0)),

cudaMemcpyHostToDevice, compute_st_);

auto func_main_q_end = [](const LaneCounters &c) {

return c.main_q_narcs_and_end.y;

};

int32 max_main_q_end = GetMaxForAllLanes(func_main_q_end);

// Step1 : Finding the best cost in the last token queue, with and without

// final costs.

// Also saving the indexes of those min.

GetBestCostStep1Kernel(

KaldiCudaDecoderNumBlocks(max_main_q_end, nlanes_used_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_, *h_device_params_,

*h_kernel_params_, use_final_costs, StdWeight::Zero().Value());

// Step2: Now that we now what the minimum cost is, we list all tokens within

// [min_cost; min_cost+lattice_beam]

// min_cost takes into account the final costs if use_final_costs is true,

// AND if a final state is is present in the last token queue

GetBestCostStep2Kernel(

KaldiCudaDecoderNumBlocks(max_main_q_end, nlanes_used_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_, *h_device_params_,

*h_kernel_params_, use_final_costs, StdWeight::Zero().Value());

// Step3 : Moves some data to host. We are moving the data that couldn't be

// moved

// directly in step 2, e.g. results of atomics (we don't know which one is

// last)

GetBestCostStep3Kernel(

KaldiCudaDecoderNumBlocks(max_main_q_end, nlanes_used_),

KALDI_CUDA_DECODER_1D_BLOCK, compute_st_, *h_device_params_,

*h_kernel_params_);

// Resetting the datastructures

argmins->clear();

has_reached_final->clear();

list_finals_token_idx_and_cost->clear();

// list_finals_token_idx_and_cost is a vector<vector<>>

// Each channel will have its own list of tokens within [best;

// best+lattice_beam]

list_finals_token_idx_and_cost->resize(nlanes_used_);

// Waiting for the copy

cudaStreamSynchronize(compute_st_);

for (int32 ilane = 0; ilane < nlanes_used_; ++ilane) {

int2 minarg = h_lanes_counters_.lane(ilane)->min_int_cost_and_arg;

// Min cost in that channel last token queue

CostType min_cost = orderedIntToFloatHost(minarg.x);

// index of that min cost

int32 arg = minarg.y;

// Saving both in output

argmins->push_back({arg, min_cost});

// Whether or not the last token queue contains at least one token

// associated with a final FST state

has_reached_final->push_back(

h_lanes_counters_.lane(ilane)->has_reached_final);

// Number of tokens within [min_cost; min_cost+lattice_beam]

int n_within_lattice_beam =

h_lanes_counters_.lane(ilane)->n_within_lattice_beam;

// Loading those tokens

(*list_finals_token_idx_and_cost)[ilane].resize(n_within_lattice_beam);

// Moving to output + int2float conversion

for (int i = 0; i < n_within_lattice_beam; ++i) {

int global_idx = h_list_final_tokens_in_main_q_.lane(ilane)[i].x;

float cost_with_final = orderedIntToFloatHost(

h_list_final_tokens_in_main_q_.lane(ilane)[i].y);

(*list_finals_token_idx_and_cost)[ilane][i].first = global_idx;

(*list_finals_token_idx_and_cost)[ilane][i].second = cost_with_final;

}

}

std::vector<Lattice *> &fst_out_vec,

bool use_final_probs) {

KALDI_ASSERT(channels.size() == fst_out_vec.size());

nvtxRangePushA("GetBestPath");

GetBestCost(channels, use_final_probs, &argmins_,

&list_finals_token_idx_and_cost_, &has_reached_final_);

std::vector<int32> reversed_path;

for (int32 ilane = 0; ilane < channels.size(); ++ilane) {

const ChannelId ichannel = channels[ilane];

const int32 token_with_best_cost = argmins_[ilane].first;

std::unique_lock<std::mutex> channel_lk(channel_lock_[ichannel]);