// cudadecoder/cuda-decoder.h
  //
  // Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
  // Hugo Braun, Justin Luitjens, Ryan Leary
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //     http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing, software
  // distributed under the License is distributed on an "AS IS" BASIS,
  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  // See the License for the specific language governing permissions and
  // limitations under the License.
  
  #ifndef KALDI_CUDA_DECODER_CUDA_DECODER_H_
  #define KALDI_CUDA_DECODER_CUDA_DECODER_H_
  
  #include "cudadecoder/cuda-decodable-itf.h"
  #include "cudadecoder/cuda-decoder-common.h"
  #include "cudadecoder/cuda-fst.h"
  #include "nnet3/decodable-online-looped.h"
  #include "thread-pool.h"
  
  #include <cuda_runtime_api.h>
  #include <mutex>
  #include <tuple>
  #include <vector>
  namespace kaldi {
  namespace cuda_decoder {
  
  struct CudaDecoderConfig {
    BaseFloat default_beam;
    BaseFloat lattice_beam;
    int32 ntokens_pre_allocated;
    int32 main_q_capacity, aux_q_capacity;
    int32 max_active;
  
    CudaDecoderConfig()
        : default_beam(15.0),
          lattice_beam(10.0),
          ntokens_pre_allocated(2000000),
          main_q_capacity(-1),
          aux_q_capacity(-1),
          max_active(10000) {}
  
    void Register(OptionsItf *opts) {
      opts->Register("beam", &default_beam,
                     "Decoding beam. Larger->slower, more accurate. If "
                     "aux-q-capacity is too small, we may decrease the beam "
                     "dynamically to avoid overflow (adaptive beam, see "
                     "aux-q-capacity parameter)");
      opts->Register("lattice-beam", &lattice_beam,
                     "The width of the lattice beam");
      opts->Register("max-active", &max_active,
                     "At the end of each frame computation, we keep only its "
                     "best max-active tokens. One token is the instantiation of "
                     "a single arc. Typical values are within the 5k-10k range.");
      opts->Register("ntokens-pre-allocated", &ntokens_pre_allocated,
                     "Advanced - Number of tokens pre-allocated in host buffers. "
                     "If this size is exceeded the buffer will reallocate, "
                     "reducing performance.");
      std::ostringstream main_q_capacity_desc;
      main_q_capacity_desc
          << "Advanced - Capacity of the main queue : Maximum number of "
             "tokens that can be stored *after* pruning for each frame. "
             "Lower -> less memory usage, Higher -> More accurate. "
             "Tokens stored in the main queue were already selected "
             "through a max-active pre-selection. It means that for each "
             "emitting/non-emitting iteration, we can add at most "
             "~max-active tokens to the main queue. Typically only the "
             "emitting iteration creates a large number of tokens. Using "
             "main-q-capacity=k*max-active with k=4..10 should be safe. "
             "If main-q-capacity is too small, we will print a warning "
             "but prevent the overflow. The computation can safely "
             "continue, but the quality of the output may decrease "
             "(-1 = set to "
          << KALDI_CUDA_DECODER_MAX_ACTIVE_MAIN_Q_CAPACITY_FACTOR
          << "*max-active).";
      opts->Register("main-q-capacity", &main_q_capacity,
                     main_q_capacity_desc.str());
      std::ostringstream aux_q_capacity_desc;
      aux_q_capacity_desc
          << "Advanced - Capacity of the auxiliary queue : Maximum "
             "number of raw tokens that can be stored *before* pruning "
             "for each frame. Lower -> less memory usage, Higher -> More "
             "accurate. During the tokens generation, if we detect that "
             "we are getting close to saturating that capacity, we will "
             "reduce the beam dynamically (adaptive beam) to keep only "
             "the best tokens in the remaining space. If the aux queue "
             "is still too small, we will print an overflow warning, but "
             "prevent the overflow. The computation can safely continue, "
             "but the quality of the output may decrease. We strongly "
             "recommend keeping aux-q-capacity large (>400k), to avoid "
             "triggering the adaptive beam and/or the overflow "
             "(-1 = set to "
          << KALDI_CUDA_DECODER_AUX_Q_MAIN_Q_CAPACITIES_FACTOR
          << "*main-q-capacity).";
      opts->Register("aux-q-capacity", &aux_q_capacity,
                     aux_q_capacity_desc.str());
    }
  
    void Check() const {
      KALDI_ASSERT(default_beam > 0.0 && ntokens_pre_allocated >= 0 &&
                   lattice_beam >= 0.0f && max_active > 0);
    }
  
    void ComputeConfig() {
      if (main_q_capacity == -1)
        main_q_capacity =
            max_active * KALDI_CUDA_DECODER_MAX_ACTIVE_MAIN_Q_CAPACITY_FACTOR;
      if (aux_q_capacity == -1)
        aux_q_capacity =
            main_q_capacity * KALDI_CUDA_DECODER_AUX_Q_MAIN_Q_CAPACITIES_FACTOR;
    }
  };
  
  // Forward declaration.
  // Those contains CUDA code. We don't want to include their definition
  // in this header
  class DeviceParams;
  class KernelParams;
  
  class CudaDecoder {
   public:
    // Creating a new CudaDecoder, associated to the FST fst
    // nlanes and nchannels are defined as follow
  
    // A decoder channel is linked to one utterance.
    // When we need to perform decoding on an utterance,
    // we pick an available channel, call InitDecoding on that channel
    // (with that ChannelId in the channels vector in the arguments)
    // then call AdvanceDecoding whenever frames are ready for the decoder
    // for that utterance (also passing the same ChannelId to AdvanceDecoding)
    //
    // A decoder lane is where the computation actually happens
    // a decoder lane is channel, and perform the actual decoding
    // of that channel.
    // If we have 200 lanes, we can compute 200 utterances (channels)
    // at the same time. We need many lanes in parallel to saturate the big GPUs
    //
    // An analogy would be lane -> a CPU core, channel -> a software thread
    // A channel saves the current state of the decoding for a given utterance.
    // It can be kept idle until more frames are ready to be processed
    //
    // We will use as many lanes as necessary to saturate the GPU, but not more.
    // A lane has an higher memory usage than a channel. If you just want to be
    // able to
    // keep more audio channels open at the same time (when I/O is the bottleneck
    // for instance,
    // typically in the context of online decoding), you should instead use more
    // channels.
    //
    // A channel is typically way smaller in term of memory usage, and can be used
    // to oversubsribe lanes in the context of online decoding
    // For instance, we could choose nlanes=200 because it gives us good
    // performance
    // on a given GPU. It gives us an end-to-end performance of 3000 XRTF. We are
    // doing online,
    // so we only get audio at realtime speed for a given utterance/channel.
    // We then decide to receive audio from 2500 audio channels at the same time
    // (each at realtime speed),
    // and as soon as we have frames ready for nlanes=200 channels, we call
    // AdvanceDecoding on those channels
    // In that configuration, we have nlanes=200 (for performance), and
    // nchannels=2500 (to have enough audio
    // available at a given time).
    // Using nlanes=2500 in that configuration would first not be possible (out of
    // memory), but also not necessary.
    // Increasing the number of lanes is only useful if it increases performance.
    // If the GPU is saturated at nlanes=200,
    // you should not increase that number
    CudaDecoder(const CudaFst &fst, const CudaDecoderConfig &config, int32 nlanes,
                int32 nchannels);
  
    // Reads the config from config
    void ReadConfig(const CudaDecoderConfig &config);
    // Special constructor for nlanes = nchannels. Here for the non-advanced user
    // Here we can consider nchannels = batch size. If we want to decode 10
    // utterances at a time,
    // we can use nchannels = 10
    CudaDecoder(const CudaFst &fst, const CudaDecoderConfig &config,
                int32 nchannels)
        : CudaDecoder(fst, config, nchannels, nchannels) {}
    ~CudaDecoder();
  
    // InitDecoding initializes the decoding, and should only be used if you
    // intend to call AdvanceDecoding() on the channels listed in channels
    void InitDecoding(const std::vector<ChannelId> &channels);
    // Computes the heavy H2H copies of InitDecoding. Usually launched on the
    // threadpool
    void InitDecodingH2HCopies(ChannelId ichannel);
    // AdvanceDecoding on a given batch
    // a batch is defined by the channels vector
    // We can compute N channels at the same time (in the same batch)
    // where N = number of lanes, as defined in the constructor
    // AdvanceDecoding will compute as many frames as possible while running the
    // full batch
    // when at least one channel has no more frames ready to be computed,
    // AdvanceDecoding returns
    // The user then decides what to do, i.e.:
    //
    // 1) either remove the empty channel from the channels list
    // and call again AdvanceDecoding
    // 2) or swap the empty channel with another one that has frames ready
    // and call again AdvanceDecoding
    //
    // Solution 2) should be preferred because we need to run full, big batches to
    // saturate the GPU
    //
    // If max_num_frames is >= 0 it will decode no more than
    // that many frames.
    void AdvanceDecoding(const std::vector<ChannelId> &channels,
                         std::vector<CudaDecodableInterface *> &decodables,
                         int32 max_num_frames = -1);
  
    // Returns the number of frames already decoded in a given channel
    int32 NumFramesDecoded(ChannelId ichannel) const;
    // GetBestPath gets the one-best decoding traceback. If "use_final_probs" is
    // true
    // AND we reached a final state, it limits itself to final states;
    // otherwise it gets the most likely token not taking into account
    // final-probs.
    void GetBestPath(const std::vector<ChannelId> &channels,
                     std::vector<Lattice *> &fst_out_vec,
                     bool use_final_probs = true);
    // It is possible to use a threadsafe version of GetRawLattice, which is
    // ConcurrentGetRawLatticeSingleChannel()
    // Which will do the heavy CPU work associated with GetRawLattice
    // It is necessary to first call PrepareForGetRawLattice *on the main thread*
    // on the channels.
    // The main thread is the one we use to call all other functions, like
    // InitDecoding or AdvanceDecoding
    // We usually call it "cuda control thread", but it is a CPU thread
    // For example:
    // on main cpu thread : Call PrepareForGetRawLattice on channel 8,6,3
    // then:
    // on some cpu thread : Call ConcurrentGetRawLatticeSingleChannel on channel 3
    // on some cpu thread : Call ConcurrentGetRawLatticeSingleChannel on channel 8
    // on some cpu thread : Call ConcurrentGetRawLatticeSingleChannel on channel 6
    void PrepareForGetRawLattice(const std::vector<ChannelId> &channels,
                                 bool use_final_probs);
    void ConcurrentGetRawLatticeSingleChannel(ChannelId ichannel,
                                              Lattice *fst_out);
  
    // GetRawLattice gets the lattice decoding traceback (using the lattice-beam
    // in the CudaConfig parameters).
    // If "use_final_probs" is true
    // AND we reached a final state, it limits itself to final states;
    // otherwise it gets the most likely token not taking into account
    // final-probs.
    void GetRawLattice(const std::vector<ChannelId> &channels,
                       std::vector<Lattice *> &fst_out_vec, bool use_final_probs);
    // GetBestCost finds the best cost in the last tokens queue
    // for each channel in channels. If isfinal is true,
    // we also add the final cost to the token costs before
    // finding the minimum cost
    // We list all tokens that have a cost within [best; best+lattice_beam]
    // in list_lattice_tokens.
    // We alsos set has_reached_final[ichannel] to true if token associated to a
    // final state
    // exists in the last token queue of that channel
    void GetBestCost(
        const std::vector<ChannelId> &channels, bool isfinal,
        std::vector<std::pair<int32, CostType>> *argmins,
        std::vector<std::vector<std::pair<int, float>>> *list_lattice_tokens,
        std::vector<bool> *has_reached_final);
    // (optional) Giving the decoder access to the cpu thread pool
    // We will use it to compute specific CPU work, such as InitDecodingH2HCopies
    // For recurrent CPU work, such as ComputeH2HCopies, we will use dedicated CPU
    // threads
    // We will launch nworkers of those threads
    void SetThreadPoolAndStartCPUWorkers(ThreadPool *thread_pool, int32 nworkers);
  
   private:
    // Data allocation. Called in constructor
    void AllocateDeviceData();
    void AllocateHostData();
    void AllocateDeviceKernelParams();
    // Data initialization. Called in constructor
    void InitDeviceData();
    void InitHostData();
    void InitDeviceParams();
    // Computes the initial channel
    // The initial channel is used to initialize a channel
    // when a new utterance starts (we clone it into the given channel)
    void ComputeInitialChannel();
    // Updates *h_kernel_params using channels
    void SetChannelsInKernelParams(const std::vector<ChannelId> &channels);
    void ResetChannelsInKernelParams();
    // Context-switch functions
    // Used to perform the context-switch of load/saving the state of a channels
    // into a lane. When a channel will be executed on a lane, we load that
    // channel into that lane (same idea than when we load a software threads into
    // the registers of a CPU)
    void LoadChannelsStateToLanes(const std::vector<ChannelId> &channels);
    void SaveChannelsStateFromLanes();
    // We compute the decodes by batch. Each decodable in the batch has a
    // different number of frames ready
    // We compute the min number of frames ready (so that the full batch is
    // executing). If max_num_frames
    // is > 0, we apply that ceiling to the NumFramesToDecode.
    int32 NumFramesToDecode(const std::vector<ChannelId> &channels,
                            std::vector<CudaDecodableInterface *> &decodables,
                            int32 max_num_frames);
    // Expand the arcs, emitting stage. Must be called after
    // a preprocess_in_place, which happens in PostProcessingMainQueue.
    // ExpandArcsEmitting is called first when decoding a frame,
    // using the preprocessing that happened at the end of the previous frame,
    // in PostProcessingMainQueue
    void ExpandArcsEmitting();
    // ExpandArcs, non-emitting stage. Must be called after PruneAndPreprocess.
    void ExpandArcsNonEmitting();
    // If we have more than max_active_ tokens in the queue (either after an
    // expand, or at the end of the frame)
    // we will compute a new beam that will only keep a number of tokens as close
    // as possible to max_active_ tokens
    // (that number is >= max_active_) (soft topk)
    // All ApplyMaxActiveAndReduceBeam is find the right beam for that topk and
    // set it.
    // We need to then call PruneAndPreprocess (explicitly pruning tokens with
    // cost > beam)
    // Or PostProcessingMainQueue (ignoring tokens with cost > beam in the next
    // frame)
    void ApplyMaxActiveAndReduceBeam(enum QUEUE_ID queue_id);
    // Called after an ExpandArcs. Prune the aux_q (output of the ExpandArcs),
    // move the survival tokens to the main_q, do the preprocessing at the same
    // time
    // We don't need it after the last ExpandArcsNonEmitting.
    void PruneAndPreprocess();
    // Once the non-emitting is done, the main_q is final for that frame.
    // We now generate all the data associated with that main_q, such as listing
    // the different tokens sharing the same token.next_state
    // we also preprocess for the ExpandArcsEmitting of the next frame
    // Once PostProcessingMainQueue, all working data is back to its original
    // state, to make sure we're ready for the next context switch
    void PostProcessingMainQueue();
    // Moving the relevant data to host, ie the data that will be needed in
    // GetBestPath/GetRawLattice.
    // Happens when PostProcessingMainQueue is done generating that data
    void CopyMainQueueDataToHost();
    // CheckOverflow
    // If a kernel sets the flag h_q_overflow, we send a warning to stderr
    // Overflows are detected and prevented on the device. It only means
    // that we've discarded the tokens that were created after the queue was full
    // That's why we only send a warning. It is not a fatal error
    void CheckOverflow();
    // Evaluates the function func for each lane, returning the max of all return
    // values
    // (func returns int32)
    // Used for instance to ge the max number of arcs for all lanes
    // func is called with h_lanes_counters_[ilane] for each lane.
    // h_lanes_counters_
    // must be ready to be used when calling GetMaxForAllLanes (you might want to
    // call
    // CopyLaneCountersToHost[A|]sync to make sure everything is ready first)
    int32 GetMaxForAllLanes(std::function<int32(const LaneCounters &)> func);
    // Copy the lane counters back to host, async or sync
    // The lanes counters contain all the information such as main_q_end (number
    // of tokens in the main_q)
    // main_q_narcs (number of arcs) during the computation. That's why we
    // frequently copy it back to host
    // to know what to do next
    void CopyLaneCountersToHostAsync();
    void CopyLaneCountersToHostSync();
    // The selected tokens for each frame will be copied back to host. We will
    // store them on host memory, and we wil use them to create the final lattice
    // once we've reached the last frame
    // We will also copy information on those tokens that we've generated on the
    // device, such as which tokens are associated to the same FST state in the
    // same frame, or their extra cost.
    // We cannot call individuals Device2Host copies for each channel, because it
    // would lead to a lot of small copies, reducing performance. Instead we
    // concatenate all channels data into a single
    // continuous array, copy that array to host, then unpack it to the individual
    // channel vectors
    // The first step (pack then copy to host, async) is done in
    // ConcatenateData
    // The second step is done in LaunchD2H and sLaunchH2HCopies
    // A sync on cudaStream st has to happen between the two functions to make
    // sure that the copy is done
    //
    // Each lane contains X elements to be copied, where X = func(ilane)
    // That data is contained in the array (pointer, X), with pointer = src[ilane]
    // It will be concatenated in d_concat on device, then copied async into
    // h_concat
    // That copy is launched on stream st
    // The offset of the data of each lane in the concatenate array is saved in
    // *lanes_offsets_ptr
    // it will be used for unpacking in MoveConcatenatedCopyToVector
    //
    // func is called with h_lanes_counters_[ilane] for each lane.
    // h_lanes_counters_
    // must be ready to be used when calling GetMaxForAllLanes (you might want to
    // call
    // CopyLaneCountersToHost[A|]sync to make sure everything is ready first)
    // Concatenate data on device before calling the D2H copies
    void ConcatenateData();
    // Start the D2H copies used to send data back to host at the end of each
    // frames
    void LaunchD2HCopies();
    // ComputeH2HCopies
    // At the end of each frame, we copy data back to host
    // That data was concatenated into a single continous array
    // We then have to unpack it and move it inside host memory
    // This is done by ComputeH2HCopies
    void ComputeH2HCopies();
    // Takes care of preparing the data for ComputeH2HCopies
    // and check whether we can use the threadpool or we have to do the work on
    // the current thread
    void LaunchH2HCopies();
    // Function called by the CPU worker threads
    // Calls ComputeH2HCopies when triggered
    void ComputeH2HCopiesCPUWorker();
  
    template <typename T>
    void MoveConcatenatedCopyToVector(const LaneId ilane,
                                      const ChannelId ichannel,
                                      const std::vector<int32> &lanes_offsets,
                                      T *h_concat,
                                      std::vector<std::vector<T>> *vecvec);
    void WaitForH2HCopies();
    void WaitForInitDecodingH2HCopies();
    // Computes a set of static asserts on the static values
    // In theory we should do them at compile time
    void CheckStaticAsserts();
    // Can be called in GetRawLattice to do a bunch of deep asserts on the data
    // Slow, so disabled by default
    void DebugValidateLattice();
  
    //
    // Data members
    //
  
    // The CudaFst data structure contains the FST graph
    // in the CSR format, on both the GPU and CPU memory
    const CudaFst fst_;
    // Counters used by a decoder lane
    // Contains all the single values generated during computation,
    // such as the current size of the main_q, the number of arcs currently in
    // that queue
    // We load data from the channel state during context-switch (for instance the
    // size of the last token queue for that channel)
    HostLaneMatrix<LaneCounters> h_lanes_counters_;
    // Counters of channels
    // Contains all the single values saved to remember the state of a channel
    // not used during computation. Those values are loaded/saved into/from a lane
    // during context switching
    ChannelCounters *h_channels_counters_;
    // Contain the various counters used by lanes/channels, such as main_q_end,
    // main_q_narcs. On device memory (equivalent of h_channels_counters on
    // device)
    DeviceChannelMatrix<ChannelCounters> d_channels_counters_;
    DeviceLaneMatrix<LaneCounters> d_lanes_counters_;
    // Number of lanes and channels, as defined in the constructor arguments
    int32 nlanes_, nchannels_;
  
    // We will now define the data used on the GPU
    // The data is mainly linked to two token queues
    // - the main queue
    // - the auxiliary queue
    //
    // The auxiliary queue is used to store the raw output of ExpandArcs.
    // We then prune that aux queue (and apply max-active) and move the survival
    // tokens in the main queue.
    // Tokens stored in the main q can then be used to generate new tokens (using
    // ExpandArcs)
    // We also generate more information about what's in the main_q at the end of
    // a frame (in PostProcessingMainQueue)
    //
    // As a reminder, here's the data structure of a token :
    //
    // struct Token { state, cost, prev_token, arc_idx }
    //
    // Please keep in mind that this structure is also used in the context
    // of lattice decoding. We are not storing a list of forward links like in the
    // CPU decoder. A token stays an instanciation of an single arc.
    //
    // For performance reasons, we split the tokens in three parts :
    // { state } , { cost }, { prev_token, arc_idx }
    // Each part has its associated queue
    // For instance, d_main_q_state[i], d_main_q_cost[i], d_main_q_info[i]
    // all refer to the same token (at index i)
    // The data structure InfoToken contains { prev_token, arc_idx }
    // We also store the acoustic costs independently in d_main_q_acoustic_cost_
    //
    // The data is eiher linked to a channel, or to a lane.
    //
    // Channel data (DeviceChannelMatrix):
    //
    // The data linked with a channel contains the data of frame i we need to
    // remember
    // to compute frame i+1. It is the list of tokens from frame i, with some
    // additional info
    // (ie the prefix sum of the emitting arcs degrees from those tokens).
    // We are only storing d_main_q_state_and_cost_ as channel data because that's
    // all we need in a token to compute
    // frame i+1. We don't need token.arc_idx or token.prev_token.
    // The reason why we also store that prefix sum is because we do the emitting
    // preprocessing
    // at the end of frame i. The reason for that is that we need infos from the
    // hashmap to do that preprocessing.
    // The hashmap is always cleared at the end of a frame. So we need to do the
    // preprocessing at the end of frame i,
    // and then save d_main_q_degrees_prefix_sum_. d_main_q_arc_offsets is
    // generated also during preprocessing.
    //
    // Lane data (DeviceLaneMatrix):
    //
    // The lane data is everything we use during computation, but which we reset
    // at the end of each frame.
    // For instance we use a hashmap at some point during the computation, but at
    // the end of each frame we reset it. That
    // way that hashmap is able to compute whichever channel the next time
    // AdvanceDecoding is called. The reasons why we do that is :
    //
    // - We use context switching. Before and after every frames, we can do a
    // context switching. Which means that a lane cannot save a channel's state
    // in any way once AdvanceDecoding returns. e.g., during a call of
    // AdvanceDecoding, ilane=2 may compute 5 frames from channel=57 (as defined
    // in the std::vector<ChannelId> channels).
    // In the next call, the same ilane=2 may compute 10 frames from channel=231.
    // A lane data has to be reset to its original state at the end of each
    // AdvanceDecoding call.
    // If somehow some data has to be saved, it needs to be declared as channel
    // data.
    //
    // - The reason why we make the distinction between lane and channel data (in
    // theory everything could be consider channel data), is because
    // a lane uses more memory than a channel. In the context of online decoding,
    // we need to create a lot channels, and we need them to be as small as
    // possible in memory.
    // Everything that can be reused between channels is stored as lane data.
  
    //
    // Channel data members:
    //
  
    DeviceChannelMatrix<int2> d_main_q_state_and_cost_;
    // Prefix sum of the arc's degrees in the main_q. Used by ExpandArcs,
    // set in the preprocess stages (either PruneAndPreprocess or
    // preprocess_in_place in PostProcessingMainQueue)
    DeviceChannelMatrix<int32> d_main_q_degrees_prefix_sum_;
    // d_main_q_arc_offsets[i] = fst_.arc_offsets[d_main_q_state[i]]
    // we pay the price for the random memory accesses of fst_.arc_offsets in the
    // preprocess kernel
    // we cache the results in d_main_q_arc_offsets which will be read in a
    // coalesced fashion in expand
    DeviceChannelMatrix<int32> d_main_q_arc_offsets_;
  
    //
    // Lane data members:
    //
  
    // InfoToken
    // Usually contains {prev_token, arc_idx}
    // If more than one token is associated to a fst_state,
    // it will contain where to find the list of those tokens in
    // d_main_q_extra_prev_tokens
    // ie {offset,size} in that list. We differentiate the two situations by
    // calling InfoToken.IsUniqueTokenForStateAndFrame()
    DeviceLaneMatrix<InfoToken> d_main_q_info_;
    // Acoustic cost of a given token
    DeviceLaneMatrix<CostType> d_main_q_acoustic_cost_;
    // At the end of a frame, we use a hashmap to detect the tokens that are
    // associated with the same FST state S
    // We do it that the very end, to only use the hashmap on post-prune, post-max
    // active tokens
    DeviceLaneMatrix<HashmapValueT> d_hashmap_values_;
    // Reminder: in the GPU lattice decoder, a token is always associated
    // to a single arc. Which means that multiple tokens in the same frame
    // can be associated with the same FST state.
    //
    // We are NOT listing those duplicates as ForwardLinks in an unique meta-token
    // like in the CPU lattice decoder
    //
    // When more than one token is associated to a single FST state,
    // we will list those tokens into another list : d_main_q_extra_prev_tokens
    // we will also save data useful in such a case, such as the extra_cost of a
    // token compared to the best for that state
    DeviceLaneMatrix<InfoToken> d_main_q_extra_prev_tokens_;
    DeviceLaneMatrix<float2> d_main_q_extra_and_acoustic_cost_;
    // Histogram. Used to perform the histogram of the token costs
    // in the main_q. Used to perform a soft topk of the main_q (max-active)
    DeviceLaneMatrix<int32> d_histograms_;
    // When filling the hashmap in PostProcessingMainQueue, we create a hashmap
    // value for each FST state
    // presents in the main_q (if at least one token is associated with that
    // state)
    // d_main_q_state_hash_idx_[token_idx] is the index of the state token.state
    // in the hashmap
    // Stored into a FSTStateHashIndex, which is actually a int32.
    // FSTStateHashIndex should only
    // be accessed through [Get|Set]FSTStateHashIndex, because it uses the bit
    // sign to also remember if that token is the representative of that state.
    // If only one token is associated with S, its representative will be itself
    DeviceLaneMatrix<FSTStateHashIndex> d_main_q_state_hash_idx_;
    // local_idx of the extra cost list for a state
    // For a given state S, first token associated with S will have local_idx=0
    // the second one local_idx=1, etc. The order of the local_idxs is random
    DeviceLaneMatrix<int32> d_main_q_n_extra_prev_tokens_local_idx_;
    // Where to write the extra_prev_tokens in the d_main_q_extra_prev_tokens_
    // queue
    DeviceLaneMatrix<int32> d_main_q_extra_prev_tokens_prefix_sum_;
    // Used when computing the prefix_sums in preprocess_in_place. Stores
    // the local_sums per CTA
    DeviceLaneMatrix<int2> d_main_q_block_sums_prefix_sum_;
    // Defining the aux_q. Filled by ExpandArcs.
    // The tokens are moved to the main_q by PruneAndPreprocess
    DeviceLaneMatrix<int2> d_aux_q_state_and_cost_;
    DeviceLaneMatrix<InfoToken> d_aux_q_info_;
    // Dedicated space for the concat of extra_cost. We should reuse memory
    DeviceLaneMatrix<float2> d_extra_and_acoustic_cost_concat_matrix_;
    DeviceLaneMatrix<InfoToken> d_extra_prev_tokens_concat_matrix_;
    DeviceLaneMatrix<CostType> d_acoustic_cost_concat_matrix_;
    DeviceLaneMatrix<InfoToken> d_infotoken_concat_matrix_;
    // We will list in d_list_final_tokens_in_main_q all tokens within [min_cost;
    // min_cost+lattice_beam]
    // It is used when calling GetBestCost
    // We only use an interface here because we will actually reuse data from
    // d_aux_q_state_and_cost
    // We are done using the aux_q when GetBestCost is called, so we can reuse
    // that memory
    HostLaneMatrix<int2> h_list_final_tokens_in_main_q_;
    // Parameters used by the kernels
    // DeviceParams contains all the parameters that won't change
    // i.e. memory address of the main_q for instance
    // KernelParams contains information that can change.
    // For instance which channel is executing on which lane
    DeviceParams *h_device_params_;
    KernelParams *h_kernel_params_;
    std::vector<ChannelId> channel_to_compute_;
    int32 nlanes_used_;  // number of lanes used in h_kernel_params_
    // Initial lane
    // When starting a new utterance,
    // init_channel_id is used to initialize a channel
    int32 init_channel_id_;
    // CUDA streams used by the decoder
    cudaStream_t compute_st_, copy_st_;
    // Parameters extracted from CudaDecoderConfig
    // Those are defined in CudaDecoderConfig
    CostType default_beam_;
    CostType lattice_beam_;
    int32 ntokens_pre_allocated_;
    int32 max_active_;  // Target value from the parameters
    int32 aux_q_capacity_;
    int32 main_q_capacity_;
    // Hashmap capacity. Multiple of max_tokens_per_frame
    int32 hashmap_capacity_;
    // Static segment of the adaptive beam. Cf InitDeviceParams
    int32 adaptive_beam_static_segment_;
    // The first index of all the following vectors (or vector<vector>)
    // is the ChannelId. e.g., to get the number of frames decoded in channel 2,
    // look into num_frames_decoded_[2].
  
    // Keep track of the number of frames decoded in the current file.
    std::vector<int32> num_frames_decoded_;
    // Offsets of each frame in h_all_tokens_info_
    std::vector<std::vector<int32>> frame_offsets_;
    // Data storage. We store on host what we will need in
    // GetRawLattice/GetBestPath
    std::vector<std::vector<InfoToken>> h_all_tokens_info_;
    std::vector<std::vector<CostType>> h_all_tokens_acoustic_cost_;
    std::vector<std::vector<InfoToken>> h_all_tokens_extra_prev_tokens_;
    std::vector<std::vector<float2>>
        h_all_tokens_extra_prev_tokens_extra_and_acoustic_cost_;
    std::vector<std::mutex> channel_lock_;  // at some point we should switch to a
                                            // shared_lock (to be able to compute
                                            // partial lattices while still
                                            // streaming new data for this
                                            // channel)
    bool worker_threads_running_;
    // For each channel, set by PrepareForGetRawLattice
    // argmin cost, list of the tokens within [best_cost;best_cost+lattice_beam]
    // and if we've reached a final token. Set by PrepareForGetRawLattice.
    std::vector<std::pair<int32, CostType>> h_all_argmin_cost_;
    std::vector<std::vector<std::pair<int, float>>> h_all_final_tokens_list_;
    std::vector<bool> h_all_has_reached_final_;
  
    // Pinned memory arrays. Used for the DeviceToHost copies
    float2 *h_extra_and_acoustic_cost_concat_, *d_extra_and_acoustic_cost_concat_;
    InfoToken *h_infotoken_concat_, *d_infotoken_concat_;
    CostType *h_acoustic_cost_concat_, *d_acoustic_cost_concat_;
    InfoToken *h_extra_prev_tokens_concat_, *d_extra_prev_tokens_concat_;
    // second memory space used for double buffering
    float2 *h_extra_and_acoustic_cost_concat_tmp_;
    InfoToken *h_infotoken_concat_tmp_;
    CostType *h_acoustic_cost_concat_tmp_;
    InfoToken *h_extra_prev_tokens_concat_tmp_;
    // Offsets used in MoveConcatenatedCopyToVector
    std::vector<int32> h_main_q_end_lane_offsets_;
    std::vector<int32> h_emitting_main_q_end_lane_offsets_;
    std::vector<int32> h_n_extra_prev_tokens_lane_offsets_;
    // Used when calling GetBestCost
    std::vector<std::pair<int32, CostType>> argmins_;
    std::vector<bool> has_reached_final_;
    std::vector<std::vector<std::pair<int32, CostType>>>
        list_finals_token_idx_and_cost_;
    bool compute_max_active_;
    cudaEvent_t nnet3_done_evt_;
    cudaEvent_t d2h_copy_acoustic_evt_;
    cudaEvent_t d2h_copy_infotoken_evt_;
    cudaEvent_t d2h_copy_extra_prev_tokens_evt_;
    cudaEvent_t concatenated_data_ready_evt_;
    cudaEvent_t lane_offsets_ready_evt_;
    // GetRawLattice helper
    // Data used when building the lattice in GetRawLattice
  
    // few typedef to make GetRawLattice easier to understand
    // Returns a unique id for each (iframe, fst_state) pair
    // We need to be able to quickly identity a (iframe, fst_state) ID
    //
    // A lattice state is defined by the pair (iframe, fst_state)
    // A token is associated to a lattice state (iframe, token.next_state)
    // Multiple token in the same frame can be associated to the same lattice
    // state
    // (they all go to the same token.next_state)
    // We need to quickly identify what is the lattice state of a token.
    // We are able to do that through GetLatticeStateInternalId(token),
    // which returns the internal unique ID for each lattice state for a token
    //
    // When we build the output lattice, we a get new lattice state
    // output_lattice_state = fst_out->AddState()
    // We call this one OutputLatticeState
    // The conversion between the two is done through maps
    // [curr|prev]_f_raw_lattice_state_
    typedef int32 LatticeStateInternalId;
    typedef StateId OutputLatticeState;
    typedef int32 TokenId;
    LatticeStateInternalId GetLatticeStateInternalId(int32 total_ntokens,
                                                     TokenId token_idx,
                                                     InfoToken token);
    // Keeping track of a variety of info about states in the lattice
    // - token_extra_cost. A path going from the current lattice_state to the
    // end has an extra cost
    // compared to the best path (which has an extra cost of 0).
    // token_extra_cost is the minimum of the extra_cost of all paths going from
    // the current lattice_state
    // to the final frame.
    // - fst_lattice_state is the StateId of the lattice_state in fst_out (in
    // the output lattice). lattice_state is an internal state used in
    // GetRawLattice.
    // - is_state_closed is true if the token_extra_cost has been read by
    // another token. It means that the
    // token_extra_cost value has been used, and if we modify token_extra_cost
    // again, we may need to recompute the current frame (so that everyone uses
    // the latest
    // token_extra_cost value)
    struct RawLatticeState {
      CostType token_extra_cost;
      OutputLatticeState fst_lattice_state;
      bool is_state_closed;
    };
    // extra_cost_min_delta_ used in the must_replay_frame situation. Please read
    // comments
    // associated with must_replay_frame in GetRawLattice to understand what it
    // does
    CostType extra_cost_min_delta_;
    ThreadPool *thread_pool_;
    std::vector<std::thread> cpu_dedicated_threads_;
    int32 n_threads_used_;
    std::vector<ChannelId> lanes2channels_todo_;
    std::atomic<int> n_acoustic_h2h_copies_todo_;
    std::atomic<int> n_extra_prev_tokens_h2h_copies_todo_;
    std::atomic<int> n_d2h_copies_ready_;
    std::atomic<int> n_infotoken_h2h_copies_todo_;
    int32 n_h2h_task_not_done_;
    int32 n_init_decoding_h2h_task_not_done_;
    std::atomic<int> n_h2h_main_task_todo_;
    std::mutex n_h2h_task_not_done_mutex_;
    std::mutex n_init_decoding_h2h_task_not_done_mutex_;
    std::mutex n_h2h_main_task_todo_mutex_;
    std::condition_variable n_h2h_main_task_todo_cv_;
    std::condition_variable h2h_done_;
    std::condition_variable init_decoding_h2h_done_;
    std::atomic<bool> active_wait_;
    bool h2h_threads_running_;
    // Using the output from GetBestPath, we add the best tokens (as selected in
    // GetBestCost)
    // from the final frame to the output lattice. We also fill the data
    // structures
    // (such as q_curr_frame_todo_, or curr_f_raw_lattice_state_) accordingly
    void AddFinalTokensToLattice(
        ChannelId ichannel,
        std::vector<std::pair<TokenId, InfoToken>> *q_curr_frame_todo,
        std::unordered_map<LatticeStateInternalId, RawLatticeState>
            *curr_f_raw_lattice_state,
        Lattice *fst_out);
    // Check if a token should be added to the lattice. If it should, then
    // keep_arc will be true
    void ConsiderTokenForLattice(
        ChannelId ichannel, int32 iprev, int32 total_ntokens, TokenId token_idx,
        OutputLatticeState fst_lattice_start, InfoToken *tok_beg,
        float2 *arc_extra_cost_beg, CostType token_extra_cost,
        TokenId list_prev_token_idx, int32 list_arc_idx,
        InfoToken *list_prev_token, CostType *this_arc_prev_token_extra_cost,
        CostType *acoustic_cost, OutputLatticeState *lattice_src_state,
        bool *keep_arc, bool *dbg_found_zero);
    // Add the arc to the lattice. Also updates what needs to be updated in the
    // GetRawLattice datastructures.
    void AddArcToLattice(
        int32 list_arc_idx, TokenId list_prev_token_idx,
        InfoToken list_prev_token, int32 curr_frame_offset,
        CostType acoustic_cost, CostType this_arc_prev_token_extra_cost,
        LatticeStateInternalId src_state_internal_id,
        OutputLatticeState fst_lattice_start,
        OutputLatticeState to_fst_lattice_state,
        std::vector<std::pair<TokenId, InfoToken>> *q_curr_frame_todo,
        std::vector<std::pair<TokenId, InfoToken>> *q_prev_frame_todo,
        std::unordered_map<LatticeStateInternalId, RawLatticeState>
            *curr_f_raw_lattice_state,
        std::unordered_map<LatticeStateInternalId, RawLatticeState>
            *prev_f_raw_lattice_state,
        std::unordered_set<int32> *f_arc_idx_added, Lattice *fst_out,
        bool *must_replay_frame);
    // Read a token information
    void GetTokenRawLatticeData(
        TokenId token_idx, InfoToken token, int32 total_ntokens,
        std::unordered_map<LatticeStateInternalId, RawLatticeState>
            *curr_f_raw_lattice_state,
        CostType *token_extra_cost, OutputLatticeState *to_fst_lattice_state);
  
    // A token is an instance of an arc. It goes to a FST state (token.next_state)
    // Multiple token in the same frame can go to the same FST state.
    // GetSameFSTStateTokenList
    // returns that list
    void GetSameFSTStateTokenList(ChannelId ichannel, InfoToken token,
                                  InfoToken **tok_beg,
                                  float2 **arc_extra_cost_beg, int32 *nprevs);
  
    // Swap datastructures at the end of a frame. prev becomes curr (we go
    // backward)
    //
    void SwapPrevAndCurrLatticeMap(
        int32 iframe, bool dbg_found_best_path,
        std::vector<std::pair<TokenId, InfoToken>> *q_curr_frame_todo,
        std::vector<std::pair<TokenId, InfoToken>> *q_prev_frame_todo,
        std::unordered_map<LatticeStateInternalId, RawLatticeState>
            *curr_f_raw_lattice_state,
        std::unordered_map<LatticeStateInternalId, RawLatticeState>
            *prev_f_raw_lattice_state,
        std::unordered_set<int32> *f_arc_idx_added);
    KALDI_DISALLOW_COPY_AND_ASSIGN(CudaDecoder);
  };
  
  }  // end namespace cuda_decoder
  }  // end namespace kaldi
  
  #endif  // KALDI_CUDA_DECODER_CUDA_DECODER_H_