text: // cudadecoderbin/batched-wav-nnet3-cuda-online.cc
  //
  // Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
  // Hugo Braun
  //
  // 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 <queue>
  #if HAVE_CUDA == 1
  
  #include <cuda.h>
  #include <cuda_profiler_api.h>
  #include <nvToolsExt.h>
  #include <iomanip>
  #include <sstream>
  #include "cudadecoder/cuda-online-pipeline-dynamic-batcher.h"
  #include "cudadecoderbin/cuda-bin-tools.h"
  #include "cudamatrix/cu-allocator.h"
  #include "fstext/fstext-lib.h"
  #include "lat/lattice-functions.h"
  #include "nnet3/am-nnet-simple.h"
  #include "nnet3/nnet-utils.h"
  #include "util/kaldi-thread.h"
  
  using namespace kaldi;
  using namespace cuda_decoder;
  
  //
  // Binary for the online pipeline BatchedThreadedNnet3CudaOnlinePipeline
  // Can serve both as a benchmarking tool and an example on how to call
  // BatchedThreadedNnet3CudaOnlinePipeline
  //
  
  struct Stream {
    std::shared_ptr<WaveData> wav;
    BatchedThreadedNnet3CudaOnlinePipeline::CorrelationID corr_id;
    int offset;
    double send_next_chunk_at;
    double *latency_ptr;
  
    Stream(const std::shared_ptr<WaveData> &_wav,
           BatchedThreadedNnet3CudaOnlinePipeline::CorrelationID _corr_id,
           double _send_next_chunk_at, double *_latency_ptr)
        : wav(_wav),
          corr_id(_corr_id),
          offset(0),
          send_next_chunk_at(_send_next_chunk_at),
          latency_ptr(_latency_ptr) {}
  
    bool operator<(const Stream &other) const {
      return (send_next_chunk_at > other.send_next_chunk_at);
    }
  };
  
  int main(int argc, char *argv[]) {
    try {
      using namespace kaldi;
      using namespace fst;
  
      typedef kaldi::int32 int32;
      typedef kaldi::int64 int64;
  
      CudaOnlineBinaryOptions opts;
      int errno;
      if ((errno = SetUpAndReadCmdLineOptions(argc, argv, &opts))) return errno;
  
      TransitionModel trans_model;
      nnet3::AmNnetSimple am_nnet;
      fst::Fst<fst::StdArc> *decode_fst;
      fst::SymbolTable *word_syms;
      ReadModels(opts, &trans_model, &am_nnet, &decode_fst, &word_syms);
      BatchedThreadedNnet3CudaOnlinePipeline cuda_pipeline(
          opts.batched_decoder_config, *decode_fst, am_nnet, trans_model);
      delete decode_fst;
      if (word_syms) cuda_pipeline.SetSymbolTable(*word_syms);
  
      CompactLatticeWriter clat_writer(opts.clat_wspecifier);
      std::mutex clat_writer_m;
      if (!opts.write_lattice) {
        KALDI_LOG << "If you want to write lattices to disk, please set "
                     "--write-lattice=true";
      }
  
      int chunk_length = cuda_pipeline.GetNSampsPerChunk();
      double chunk_seconds = cuda_pipeline.GetSecondsPerChunk();
      double seconds_per_sample = chunk_seconds / chunk_length;
  
      std::vector<std::shared_ptr<WaveData>> all_wav;
      std::vector<std::string> all_wav_keys;
      ReadDataset(opts, &all_wav, &all_wav_keys);
  
      BatchedThreadedNnet3CudaOnlinePipeline::CorrelationID correlation_id_cnt =
          0;
  
      CudaOnlinePipelineDynamicBatcherConfig dynamic_batcher_config;
      CudaOnlinePipelineDynamicBatcher dynamic_batcher(dynamic_batcher_config,
                                                       cuda_pipeline);
  
      // Streaming code
      // Wav reader index
      size_t all_wav_i = 0;
      size_t all_wav_max = all_wav.size() * opts.niterations;
      std::vector<double> latencies(all_wav_max);
  
      // We will start all utterances at a random position within the first second
      std::random_device rd;
      std::mt19937 gen(rd());
      std::uniform_real_distribution<> dis(0.0, 1.0);
  
      std::priority_queue<Stream> streams;
      nvtxRangePush("Global Timer");
      Timer timer;
  
      // Initial set of streams will start randomly within the first second of
      // streaming That's to simulate something more realistic than a large spike
      // of all channels starting at the same time
      bool add_random_offset = true;
      KALDI_LOG << "Inferencing...";
      while (true) {
        while (streams.size() < opts.num_streaming_channels) {
          int wav_i = all_wav_i++;
          if (wav_i >= all_wav_max) {
            break;
          }
          uint64_t corr_id = correlation_id_cnt++;
          size_t all_wav_i_modulo = wav_i % (all_wav.size());
          double *latency_ptr = &latencies[corr_id];
  
          // The "speaker" start talking at stream_will_start_at
          double stream_will_start_at = timer.Elapsed();  // "now"
          if (add_random_offset) stream_will_start_at += dis(gen);
          // The "speaker" will speak for stream_duration seconds
          double stream_duration = all_wav[all_wav_i_modulo]->Duration();
          // The first chunk will be available whenever the first chunk was
          // "spoken" e.g. if the first chunk is made of 0.5s for audio, we have
          // to wait 0.5s after stream_will_start_at
          double first_chunk_available_at =
              stream_will_start_at + std::min(stream_duration, chunk_seconds);
  
          // stream_will_stop_at is used for latency computation
          // We started streaming at t0=stream_will_start_at, so the audio will be
          // done streaming at t1=stream_will_start_at+duration We will get
          // the result at t2, we then have latency=t2-t1
          double stream_will_stop_at = stream_will_start_at + stream_duration;
          // tmp storing in lat_ptr
          // we'll do lat_ptr = now - lat_ptr in callback
          *latency_ptr = stream_will_stop_at;
  
          // Defining the callback for results
          cuda_pipeline.SetBestPathCallback(
              corr_id,
              [latency_ptr, &timer, &opts, corr_id](
                  const std::string &str, bool partial, bool endpoint_detected) {
                if (partial && opts.print_partial_hypotheses)
                  KALDI_LOG << "corr_id #" << corr_id << " [partial] : " << str;
  
                if (endpoint_detected && opts.print_endpoints)
                  KALDI_LOG << "corr_id #" << corr_id << " [endpoint detected]";
  
                if (!partial) {
                  // *latency_ptr currently contains t1="stream_will_start_at +
                  // duration" where stream_will_start_at is when this stream
                  // started and duration is the duration of this audio file so t1
                  // is the time when the virtual user is done talking
                  // timer.Elapsed() now contains t2, i.e. when the result is
                  // ready latency = t2 - t1
                  if (!opts.generate_lattice) {
                    // If we need to gen a lattice, latency will take the lattice
                    // gen into account
                    *latency_ptr = timer.Elapsed() - *latency_ptr;
                  }
                  if (opts.print_hypotheses)
                    KALDI_LOG << "corr_id #" << corr_id << " : " << str;
                }
              });
  
          if (opts.generate_lattice) {
            // Setting a callback will indicate the pipeline to generate a
            // lattice
            int iter = all_wav_i / all_wav.size();
            std::string key = all_wav_keys[all_wav_i_modulo];
            if (opts.niterations > 1) key += std::to_string(iter);
            cuda_pipeline.SetLatticeCallback(
                corr_id, [&opts, &clat_writer, &clat_writer_m, key, latency_ptr,
                          &timer](CompactLattice &clat) {
                  *latency_ptr = timer.Elapsed() - *latency_ptr;
                  if (opts.write_lattice) {
                    std::lock_guard<std::mutex> lk(clat_writer_m);
                    clat_writer.Write(key, clat);
                  }
                });
          }
  
          // Adding that stream to our simulation stream pool
          streams.emplace(all_wav[all_wav_i_modulo], corr_id,
                          first_chunk_available_at, latency_ptr);
        }
        add_random_offset = false;  // the next streams will just start whenever a
                                    // spot is available
  
        // If we reach this, we're just done with all streams
        if (streams.empty()) break;
  
        auto chunk = streams.top();
        streams.pop();
        double wait_for = chunk.send_next_chunk_at - timer.Elapsed();
        if (wait_for > 0) usleep(wait_for * 1e6);
  
        SubVector<BaseFloat> data(chunk.wav->Data(), 0);
  
        // Current chunk
        int32 total_num_samp = data.Dim();
        int this_chunk_num_samp =
            std::min(total_num_samp - chunk.offset, chunk_length);
        bool is_last_chunk =
            ((chunk.offset + this_chunk_num_samp) == total_num_samp);
        bool is_first_chunk = (chunk.offset == 0);
        SubVector<BaseFloat> wave_part(data, chunk.offset, this_chunk_num_samp);
  
        // Giving current chunk to the dynamic batcher for processing
        dynamic_batcher.Push(chunk.corr_id, is_first_chunk, is_last_chunk,
                             wave_part);
  
        // Current chunk was sent, done
        chunk.offset += this_chunk_num_samp;
  
        // Streaming simulation:
        // We need to know the duration of the next chunk
        // The next time we will "wake up" the stream will be at
        // send_next_chunk_at with send_next_chunk_at = send_current_chunk_at +
        // next_chunk_duration
        int next_chunk_num_samp =
            std::min(total_num_samp - chunk.offset, chunk_length);
        double next_chunk_seconds = next_chunk_num_samp * seconds_per_sample;
        chunk.send_next_chunk_at += next_chunk_seconds;
  
        // If there is a next chunk, add it to the list of streams tasks
        if (!is_last_chunk) streams.push(chunk);
      }
  
      dynamic_batcher.WaitForCompletion();
      KALDI_LOG << "Done.";
      nvtxRangePop();
  
      KALDI_LOG << "\tLatency stats:";
      PrintLatencyStats(latencies);
      delete word_syms;  // will delete if non-NULL.
  
      clat_writer.Close();
      cudaDeviceSynchronize();
  
      return 0;
    } catch (const std::exception &e) {
      std::cerr << e.what();
      return -1;
    }
  }  // main()
  
  #endif  // if HAVE_CUDA == 1
  