// cudadecoder/batched-threaded-nnet3-cuda-pipeline.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_BATCHED_THREADED_CUDA_DECODER_H_
#define KALDI_CUDA_DECODER_BATCHED_THREADED_CUDA_DECODER_H_

#include <atomic>
#include <thread>

#include "cudadecoder/cuda-decoder.h"
#include "decodable-cumatrix.h"
#include "feat/wave-reader.h"
#include "lat/determinize-lattice-pruned.h"
#include "nnet3/nnet-batch-compute.h"
#include "online2/online-nnet2-feature-pipeline.h"
#include "cudafeat/online-cuda-feature-pipeline.h"
#include "thread-pool.h"

// If num_channels sets to automatic,
// num_channels = [this define] * max_batch_size
#define KALDI_CUDA_DECODER_CHANNELS_BATCH_SIZE_RATIO 1.3

namespace kaldi {
namespace cuda_decoder {

/* BatchedThreadedNnet3CudaPipelineConfig
 * This class is a common configuration class for the various components
 * of a batched cuda multi-threaded pipeline.  It defines a single place
 * to control all operations and ensures that the various componets
 * match configurations
 */
// configuration options common to the BatchedThreadedNnet3CudaPipeline and
// BatchedThreadedNnet3CudaPipeline
struct BatchedThreadedNnet3CudaPipelineConfig {
  BatchedThreadedNnet3CudaPipelineConfig()
      : max_batch_size(200),
        num_channels(-1),
        batch_drain_size(10),
        num_control_threads(2),
        num_worker_threads(20),
        determinize_lattice(true),
        max_pending_tasks(4000),
        num_decoder_copy_threads(2),
        gpu_feature_extract(true) {};
  void Register(OptionsItf *po) {
    po->Register("max-batch-size", &max_batch_size,
                 "The maximum batch size to be used by the decoder. "
                 "This is also the number of lanes in the CudaDecoder. "
                 "Larger = Faster and more GPU memory used.");
    std::ostringstream num_channels_desc;
    num_channels_desc
        << "The number of channels "
           "allocated to the cuda decoder.  This should be larger "
           "than max_batch_size.  Each channel consumes a small "
           "amount of memory but also allows us to better overlap "
           "computation"
           " (-1 = set to "
        << KALDI_CUDA_DECODER_CHANNELS_BATCH_SIZE_RATIO << "*max-batch-size).";
    po->Register("num-channels", &num_channels, num_channels_desc.str());
    po->Register("batch-drain-size", &batch_drain_size,
                 "How far to drain the batch before refilling work. This "
                 "batches pre/post decode work.");
    po->Register("cuda-control-threads", &num_control_threads,
                 "The number of pipeline control threads for the CUDA work. "
                 "e.g. 2 control threads -> 2 independent CUDA pipeline (nnet3 "
                 "and decoder).");
    po->Register(
        "cuda-worker-threads", &num_worker_threads,
        "The total number of CPU threads launched to process CPU tasks.");
    po->Register("determinize-lattice", &determinize_lattice,
                 "Determinize the lattice before output.");
    po->Register("max-outstanding-queue-length", &max_pending_tasks,
                 "Number of files to allow to be outstanding at a time. When "
                 "the number of files is larger than this handles will be "
                 "closed before opening new ones in FIFO order.");
    po->Register("cuda-decoder-copy-threads", &num_decoder_copy_threads,
                 "Advanced - Number of worker threads used in the decoder for "
                 "the host to host copies.");
    po->Register("gpu-feature-extract", &gpu_feature_extract,
                 "Extract features on the GPU.  This reduces CPU overhead "
                 "leading to better scalability but may reduce overall "
                 "performance for a single GPU.");

    feature_opts.Register(po);
    decoder_opts.Register(po);
    det_opts.Register(po);
    compute_opts.Register(po);
  }
  int max_batch_size;
  int num_channels;
  int batch_drain_size;
  int num_control_threads;
  int num_worker_threads;
  bool determinize_lattice;
  int max_pending_tasks;
  int num_decoder_copy_threads;
  bool gpu_feature_extract;

  void ComputeConfig() {
    if (num_channels == -1)
      num_channels =
          max_batch_size * KALDI_CUDA_DECODER_CHANNELS_BATCH_SIZE_RATIO;
  }

  OnlineNnet2FeaturePipelineConfig feature_opts;      // constant readonly
  CudaDecoderConfig decoder_opts;                     // constant readonly
  fst::DeterminizeLatticePhonePrunedOptions det_opts; // constant readonly
  nnet3::NnetBatchComputerOptions compute_opts;       // constant readonly
};

/*
 * BatchedThreadedNnet3CudaPipeline uses multiple levels of parallelism in order to
 * decode quickly on CUDA GPUs. This is the primary interface for cuda decoding.
 * For examples of how to use this decoder see cudadecoder/README and
 * cudadecoderbin/batched-wav-nnet3-cuda.cc
 */
class BatchedThreadedNnet3CudaPipeline {
public:
 BatchedThreadedNnet3CudaPipeline(
     const BatchedThreadedNnet3CudaPipelineConfig &config)
     : config_(config), all_group_tasks_not_done_(0) {
   config_.ComputeConfig();
 };

 // allocates reusable objects that are common across all decodings
 void Initialize(const fst::Fst<fst::StdArc> &decode_fst,
                 const nnet3::AmNnetSimple &nnet,
                 const TransitionModel &trans_model);

 // deallocates reusable objects
 void Finalize();

 // query a specific key to see if compute on it is complete
 bool isFinished(const std::string &key);

 // remove an audio file from the decoding and clean up resources
 void CloseDecodeHandle(const std::string &key);
 void CloseAllDecodeHandlesForGroup(const std::string &group);
 void CloseAllDecodeHandles();

 // Adds a decoding task to the decoder
 // When passing in a vector of data, the caller must ensure the data exists
 // until the CloseDecodeHandle/WaitForAllTasks is called
 // callback is called once task is done and we pass it the final lattice
 // callback can be used to compute lattice rescoring, find best path in
 // lattice, writing lattice to disk, etc.
 // Important: callback is launched in the threadpool. It must be threadsafe.
 // For instance, if writing to disk, or to stdout,
 // use a lock:
 // e.g. :
 // {
 // 	std::lock_guard<std::mutex> lock(global_mutex);
 // 	// write lattice to disk
 //    // lock is released in the destructor of lock_guard<>
 // }
 void OpenDecodeHandle(
     const std::string &key, const WaveData &wave_data,
     const std::string &group = std::string(),
     const std::function<void(CompactLattice &clat)> &callback =
         std::function<void(CompactLattice &clat)>());
 // When passing in a vector of data, the caller must ensure the data exists
 // until the CloseDecodeHandle is called
 void OpenDecodeHandle(
     const std::string &key, const VectorBase<BaseFloat> &wave_data,
     float sample_rate, const std::string &group = std::string(),
     const std::function<void(CompactLattice &clat)> &callback =
         std::function<void(CompactLattice &clat)>());

 // Copies the raw lattice for decoded handle "key" into lat
 bool GetRawLattice(const std::string &key, Lattice *lat);
 // Determinizes raw lattice and returns a compact lattice
 bool GetLattice(const std::string &key, CompactLattice *lat);

 int32 GetNumberOfTasksPending();

 // Wait for all tasks to complete
 void WaitForAllTasks();
 // Wait for all tasks in the group to complete
 void WaitForGroup(const std::string &group);
 // Check if a group is available. Returns if not.
 bool IsGroupCompleted(const std::string &group);
 // Wait for any group to complete, then returns which group completed
 std::string WaitForAnyGroup();
 // Check if any group is available. If one is available, set its name in *group
 bool IsAnyGroupCompleted(std::string *group);
 inline int NumPendingTasks() {
   return (tasks_back_ - tasks_front_ + config_.max_pending_tasks + 1) %
          (config_.max_pending_tasks + 1);
  };

private:
 // Task data used during computation
 // Is cleared when task is completed
 struct TaskData {
   Vector<BaseFloat> raw_data;  // Wave input data when wave_reader passed
   std::shared_ptr<SubVector<BaseFloat>>
       wave_samples;  // Used as a pointer to either the raw
                      // data or the samples passed
   float sample_frequency;
   Vector<BaseFloat> ivector_features_cpu;
   Matrix<BaseFloat> input_features_cpu;
   CuVector<BaseFloat> ivector_features;
   CuMatrix<BaseFloat> input_features;
   CuMatrix<BaseFloat> posteriors;

   TaskData(const WaveData &wave_data_in)
       : wave_samples(NULL), sample_frequency(0) {
     raw_data.Resize(
         wave_data_in.Data().NumRows() * wave_data_in.Data().NumCols(),
         kUndefined);
     memcpy(raw_data.Data(), wave_data_in.Data().Data(),
            raw_data.Dim() * sizeof(BaseFloat));
     wave_samples =
         std::make_shared<SubVector<BaseFloat>>(raw_data, 0, raw_data.Dim());
     sample_frequency = wave_data_in.SampFreq();
   };

   // Init when raw data is passed in.  This data is shallow copied.
   TaskData(const VectorBase<BaseFloat> &wave_data_in, float sample_rate) {
     wave_samples = std::make_shared<SubVector<BaseFloat>>(wave_data_in, 0,
                                                           wave_data_in.Dim());
     sample_frequency = sample_rate;
   }
 };

 // State needed for each decode task.
 // This state can be passed around by reference or pointer safely
 // and provides a convieniet way to store all decoding state.
 struct TaskState {
   std::string key;
   std::string group;  // group for that task. "" is default
   bool error;
   std::string error_string;

   std::shared_ptr<TaskData> task_data;

   int32 ichannel;              // associated CudaDecoder channel
   Lattice lat;                 // Raw Lattice output
   CompactLattice dlat;         // Determinized lattice output.  Only set if
                                // determinize-lattice=true
   std::atomic<bool> finished;  // Tells master thread if task has finished
                                // execution

   bool determinized;

   // (optional) callback is called task is finished and we have a lattice
   // ready
   // that way we can compute all CPU tasks in the threadpool (lattice
   // rescoring, find best path in lattice, etc.)
   std::function<void(CompactLattice &clat)> callback;

   TaskState() : error(false), finished(false), determinized(false) {}

   // Init when wave data is passed directly in.  This data is deep copied.
   void Init(const std::string &key_in, const WaveData &wave_data_in) {
     task_data = std::make_shared<TaskData>(wave_data_in);
     key = key_in;
   };
   // Init when raw data is passed in.  This data is shallow copied.
   void Init(const std::string &key_in,
             const VectorBase<BaseFloat> &wave_data_in, float sample_rate) {
     task_data = std::make_shared<TaskData>(wave_data_in, sample_rate);
     key = key_in;
   }
  };

  // Creating a new task in the hashmaps
  TaskState *AddTask(const std::string &key, const std::string &group);

  // Holds the current channel state for a worker
  struct ChannelState {
    std::vector<ChannelId> channels;
    std::vector<ChannelId> free_channels;
    std::vector<ChannelId> completed_channels;
    std::mutex free_channels_mutex;
  };

  // Adds task to the PendingTaskQueue
  void AddTaskToPendingTaskQueue(TaskState *task);

  // Attempts to fill the batch from the task queue.  May not fully fill the
  // batch.
  void AquireAdditionalTasks(CudaDecoder &cuda_decoder,
                             ChannelState &channel_state,
                             std::vector<TaskState *> &tasks);

  // Computes Features for a single decode instance.
  void ComputeOneFeatureCPU(TaskState *task);

  // Computes features across the tasks[first,tasks.size()
  void ComputeBatchFeatures(int32 first,
                            std::vector<TaskState *> &tasks,
                            OnlineCudaFeaturePipeline &feature_pipeline);

  // Computes Nnet across the current decode batch
  void ComputeBatchNnet(nnet3::NnetBatchComputer &computer, int32 first,
                        std::vector<TaskState *> &tasks);

  // Allocates decodables for tasks in the range of
  // dstates[first,dstates.size())
  void AllocateDecodables(int32 first, std::vector<TaskState *> &tasks,
                          std::vector<CudaDecodableInterface *> &decodables);

  // Removes all completed channels from the channel list.
  // Also enqueues up work for post processing
  void
  RemoveCompletedChannels(CudaDecoder &cuda_decoder,
                          ChannelState &channel_state,
                          std::vector<CudaDecodableInterface *> &decodables,
                          std::vector<TaskState *> &tasks);

  // For each completed decode perform post processing work and clean up
  void PostDecodeProcessing(CudaDecoder &cuda_decoder,
                            ChannelState &channel_state,
                            std::vector<CudaDecodableInterface *> &decodables,
                            std::vector<TaskState *> &tasks);

  // Calls ConcurrentGetRawLatticeSingleChannel and Determinize
  // on a dedicated CPU worker thread at the end of the decode
  void CompleteTask(CudaDecoder *cuda_decoder, ChannelState *channel_state,
                    TaskState *state);

  // Determinize one lattice
  void DeterminizeOneLattice(TaskState *task);
  // Thread execution function.  This is a single worker thread which processes
  // input.
  void ExecuteWorker(int threadId);

  BatchedThreadedNnet3CudaPipelineConfig config_;

  CudaFst cuda_fst_;
  const TransitionModel *trans_model_;
  const nnet3::AmNnetSimple *am_nnet_;
  nnet3::DecodableNnetSimpleLoopedInfo *decodable_info_;
  OnlineNnet2FeaturePipelineInfo *feature_info_;

  std::mutex tasks_mutex_; // protects tasks_front_ and pending_task_queue_ for
                           // workers
  std::mutex tasks_add_mutex_; // protect OpenDecodeHandle if multiple threads
                               // access
  std::mutex tasks_lookup_mutex_; // protext tasks_lookup map
  std::condition_variable tasks_lookup_cv_;
  std::atomic<int> tasks_front_, tasks_back_;
  TaskState **pending_task_queue_;

  std::atomic<bool> exit_;      // signals threads to exit
  std::atomic<int> numStarted_; // signals master how many threads have started

  ThreadPool *work_pool_; // thread pool for CPU work
  std::map<std::string, int32> group_tasks_not_done_;
  int32 all_group_tasks_not_done_;
  std::mutex group_tasks_mutex_;
  std::condition_variable group_done_cv_;
  std::unordered_multimap<std::string, TaskState *>
      tasks_group_lookup_;  // group -> list of tasks
  std::unordered_map<std::string, TaskState>
      tasks_lookup_;                              // Contains a map of
                                                  // utterance to TaskState
  std::vector<std::thread> thread_contexts_;      // A list of thread contexts
};

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

#endif  // KALDI_CUDA_DECODER_BATCHED_THREADED_CUDA_DECODER_H_