[src] CudaDecoder endpointing (#4146, #4101)

 * Partial hypotheses
 * PR comments
 * Endpointing
 * PR comments
 * Neg non-em partial traceback bug fix

Author: hugovbraun

src/cudadecoder/batched-threaded-nnet3-cuda-online-pipeline.cc

@@ -85,6 +85,11 @@ void BatchedThreadedNnet3CudaOnlinePipeline::AllocateAndInitializeData(
     cuda_decoder_->SetThreadPoolAndStartCPUWorkers(
         thread_pool_.get(), config_.num_decoder_copy_threads);
   }
+
+  cuda_decoder_->SetOutputFrameShiftInSeconds(
+      feature_info_->FrameShiftInSeconds() *
+      config_.compute_opts.frame_subsampling_factor);
+
   n_samples_valid_.resize(max_batch_size_);
   n_input_frames_valid_.resize(max_batch_size_);
   n_lattice_callbacks_not_done_.store(0);
@@ -217,7 +222,9 @@ void BatchedThreadedNnet3CudaOnlinePipeline::DecodeBatch(
     const std::vector<CorrelationID> &corr_ids,
     const std::vector<SubVector<BaseFloat>> &wave_samples,
     const std::vector<bool> &is_first_chunk,
-    const std::vector<bool> &is_last_chunk) {
+    const std::vector<bool> &is_last_chunk,
+    std::vector<const std::string *> *partial_hypotheses,
+    std::vector<bool> *end_points) {
   nvtxRangePushA("DecodeBatch");
   KALDI_ASSERT(corr_ids.size() > 0);
   KALDI_ASSERT(corr_ids.size() == wave_samples.size());
@@ -242,9 +249,37 @@ void BatchedThreadedNnet3CudaOnlinePipeline::DecodeBatch(
     }
   }
   int features_frame_stride = d_all_features_.Stride();
+  if (partial_hypotheses) {
+    // We're going to have to generate the partial hypotheses
+    if (word_syms_ == nullptr) {
+      KALDI_ERR << "You need to set --word-symbol-table to use "
+                << "partial hypotheses";
+    }
+    cuda_decoder_->AllowPartialHypotheses();
+  }
+  if (end_points) cuda_decoder_->AllowEndpointing();
+
   DecodeBatch(corr_ids, d_features_ptrs_, features_frame_stride,
               n_input_frames_valid_, d_ivectors_ptrs_, is_first_chunk,
               is_last_chunk, &channels_);
+
+  if (partial_hypotheses) {
+    partial_hypotheses->resize(channels_.size());
+    for (size_t i = 0; i < channels_.size(); ++i) {
+      PartialHypothesis *partial_hypothesis;
+      ChannelId ichannel = channels_[i];
+      cuda_decoder_->GetPartialHypothesis(ichannel, &partial_hypothesis);
+      (*partial_hypotheses)[i] = &partial_hypothesis->out_str;
+    }
+  }
+
+  if (end_points) {
+    end_points->resize(channels_.size());
+    for (size_t i = 0; i < channels_.size(); ++i) {
+      ChannelId ichannel = channels_[i];
+      (*end_points)[i] = cuda_decoder_->EndpointDetected(ichannel);
+    }
+  }
 }
 
 void BatchedThreadedNnet3CudaOnlinePipeline::DecodeBatch(
@@ -393,10 +428,11 @@ void BatchedThreadedNnet3CudaOnlinePipeline::ReadParametersFromModel() {
   input_dim_ = feature.InputFeature()->Dim();
   if (use_ivectors_) ivector_dim_ = feature.IvectorFeature()->Dim();
   model_frequency_ = feature_info_->GetSamplingFrequency();
-  BaseFloat frame_shift = feature_info_->FrameShiftInSeconds();
+  BaseFloat frame_shift_seconds = feature_info_->FrameShiftInSeconds();
   input_frames_per_chunk_ = config_.compute_opts.frames_per_chunk;
-  seconds_per_chunk_ = input_frames_per_chunk_ * frame_shift;
-  int32 samp_per_frame = static_cast<int>(model_frequency_ * frame_shift);
+  seconds_per_chunk_ = input_frames_per_chunk_ * frame_shift_seconds;
+  int32 samp_per_frame =
+      static_cast<int>(model_frequency_ * frame_shift_seconds);
   samples_per_chunk_ = input_frames_per_chunk_ * samp_per_frame;
   BatchedStaticNnet3Config nnet3_config;
   nnet3_config.compute_opts = config_.compute_opts;
@@ -433,7 +469,8 @@ void BatchedThreadedNnet3CudaOnlinePipeline::FinalizeDecoding(
   }
 
   if (dlat.NumStates() > 0) {
-    if (word_syms_) {
+    // Used for debugging
+    if (false && word_syms_) {
       CompactLattice best_path_clat;
       CompactLatticeShortestPath(dlat, &best_path_clat);
 

src/cudadecoder/batched-threaded-nnet3-cuda-online-pipeline.h

@@ -169,10 +169,18 @@ class BatchedThreadedNnet3CudaOnlinePipeline {
   // If it contains some last chunks for given utterances, it will call
   // FinalizeDecoding (building the final lattice, determinize it, etc.)
   // asynchronously. The callback for that utterance will then be called
-  void DecodeBatch(const std::vector<CorrelationID> &corr_ids,
-                   const std::vector<SubVector<BaseFloat>> &wave_samples,
-                   const std::vector<bool> &is_first_chunk,
-                   const std::vector<bool> &is_last_chunk);
+  //
+  // If partial_hypotheses is not null, generate and set the current partial
+  // hypotheses in partial_hypotheses The pointers in partial_hypotheses are
+  // only valid until the next DecodeBatch call - perform a deep copy if
+  // necessary
+  void DecodeBatch(
+      const std::vector<CorrelationID> &corr_ids,
+      const std::vector<SubVector<BaseFloat>> &wave_samples,
+      const std::vector<bool> &is_first_chunk,
+      const std::vector<bool> &is_last_chunk,
+      std::vector<const std::string *> *partial_hypotheses = nullptr,
+      std::vector<bool> *end_point = nullptr);
 
   // Version providing directly the features. Only runs nnet3 & decoder
   // Used when we want to provide the final ivectors (offline case)
@@ -207,8 +215,12 @@ class BatchedThreadedNnet3CudaOnlinePipeline {
   // Maximum number of seconds per chunk
   BaseFloat GetSecondsPerChunk() { return seconds_per_chunk_; }
 
-  // Used when debugging. Used to Print the text when a decoding is done
-  void SetSymbolTable(fst::SymbolTable *word_syms) { word_syms_ = word_syms; }
+  // Used for partial hypotheses
+  void SetSymbolTable(const fst::SymbolTable &word_syms) {
+    word_syms_ = &word_syms;
+    KALDI_ASSERT(cuda_decoder_);
+    cuda_decoder_->SetSymbolTable(word_syms);
+  }
 
   // Wait for all lattice callbacks to complete
   // Can be called after DecodeBatch
@@ -355,7 +367,7 @@ class BatchedThreadedNnet3CudaOnlinePipeline {
   std::unique_ptr<ThreadPoolLight> thread_pool_;
 
   // Used for debugging
-  fst::SymbolTable *word_syms_;
+  const fst::SymbolTable *word_syms_;
   // Used when printing to stdout for debugging purposes
   std::mutex stdout_m_;
 };

src/cudadecoder/batched-threaded-nnet3-cuda-pipeline2.h

@@ -245,7 +245,7 @@ class BatchedThreadedNnet3CudaPipeline2 {
   void WaitForAllTasks();
 
   // Used for debug
-  void SetSymbolTable(fst::SymbolTable *word_syms) {
+  void SetSymbolTable(const fst::SymbolTable &word_syms) {
     cuda_online_pipeline_.SetSymbolTable(word_syms);
   }
 

src/cudadecoder/cuda-decoder-common.h

@@ -345,6 +345,56 @@ class DeviceChannelMatrix : public DeviceMatrix<T> {
   }
 };
 
+// InfoToken contains data that needs to be saved for the backtrack
+// in GetBestPath/GetRawLattice
+// We don't need the token.cost or token.next_state.
+struct __align__(8) InfoToken {
+  int32 prev_token;
+  int32 arc_idx;
+  bool IsUniqueTokenForStateAndFrame() {
+    // This is a trick used to save space and PCI-E bandwidth (cf
+    // preprocess_in_place kernel)
+    // This token is associated with a next_state s, created during the
+    // processing of frame f.
+    // If we have multiple tokens associated with the state s in the frame f,
+    // arc_idx < 0 and -arc_idx is the
+    // count of such tokens. We will then have to look at another list to read
+    // the actually arc_idx and prev_token values
+    // If the current token is the only one, prev_token and arc_idx are valid
+    // and can be used directly
+    return (arc_idx >= 0);
+  }
+
+  // Called if this token is linked to others tokens in the same frame (cf
+  // comments for IsUniqueTokenForStateAndFrame)
+  // return the {offset,size} pair necessary to list those tokens in the
+  // extra_prev_tokens list
+  // They are stored at offset "offset", and we have "size" of those
+  std::pair<int32, int32> GetSameFSTStateTokensList() {
+    KALDI_ASSERT(!IsUniqueTokenForStateAndFrame());
+
+    return {prev_token, -arc_idx};
+  }
+};
+
+// Device function, used to set a in an InfoToken the [offset,size] related to
+// InfoToken.GetSameFSTStateTokensList
+__device__ __inline__ void SetSameFSTStateTokensList(int32 offset, int32 size,
+                                                     InfoToken *info_token) {
+  // We always have size > 0
+  *info_token = {offset, -size};
+}
+
+// Information about the best path head
+// Used by partial hypotheses and endpoiting
+struct BestPathTracebackHead {
+  int index;
+  CostType relative_cost;
+
+  void Reset() { index = -1; }
+  bool IsSet() { return (index != -1); }
+};
+
 // LaneCounters/ChannelCounters
 // The counters are all the singular values associated to a lane/channel
 // For instance  the main queue size. Or the min_cost of all tokens in that
@@ -402,6 +452,7 @@ struct LaneCounters {
   int32 main_q_extra_prev_tokens_global_offset;
   // Minimum token for that frame
   IntegerCostType min_int_cost;
+  IntegerCostType int_relative_cost;
   // Current beam. Can be different from default_beam,
   // because of the AdaptiveBeam process, or because of
   // ApplyMaxActiveAndReduceBeam
@@ -471,46 +522,6 @@ class CudaDecoderException : public std::exception {
   const bool recoverable;
 };
 
-// InfoToken contains data that needs to be saved for the backtrack
-// in GetBestPath/GetRawLattice
-// We don't need the token.cost or token.next_state.
-struct __align__(8) InfoToken {
-  int32 prev_token;
-  int32 arc_idx;
-  bool IsUniqueTokenForStateAndFrame() {
-    // This is a trick used to save space and PCI-E bandwidth (cf
-    // preprocess_in_place kernel)
-    // This token is associated with a next_state s, created during the
-    // processing of frame f.
-    // If we have multiple tokens associated with the state s in the frame f,
-    // arc_idx < 0 and -arc_idx is the
-    // count of such tokens. We will then have to look at another list to read
-    // the actually arc_idx and prev_token values
-    // If the current token is the only one, prev_token and arc_idx are valid
-    // and can be used directly
-    return (arc_idx >= 0);
-  }
-
-  // Called if this token is linked to others tokens in the same frame (cf
-  // comments for IsUniqueTokenForStateAndFrame)
-  // return the {offset,size} pair necessary to list those tokens in the
-  // extra_prev_tokens list
-  // They are stored at offset "offset", and we have "size" of those
-  std::pair<int32, int32> GetSameFSTStateTokensList() {
-    KALDI_ASSERT(!IsUniqueTokenForStateAndFrame());
-
-    return {prev_token, -arc_idx};
-  }
-};
-
-// Device function, used to set a in an InfoToken the [offset,size] related to
-// InfoToken.GetSameFSTStateTokensList
-__device__ __inline__ void SetSameFSTStateTokensList(int32 offset, int32 size,
-                                                     InfoToken *info_token) {
-  // We always have size > 0
-  *info_token = {offset, -size};
-}
-
 // Used to store the index in the GPU hashmap of that FST state
 // The hashmap is only generated with the final main queue (post max_active_) of
 // each frame
@@ -558,6 +569,25 @@ enum OVERFLOW_TYPE {
 
 enum QUEUE_ID { MAIN_Q = 0, AUX_Q = 1 };
 
+// Used internally to generate partial paths
+struct PartialPathArc {
+  int32 token_idx;
+  int32 arc_idx;
+};
+
+// Partial hypothesis formatted and meant to be used by user
+struct PartialHypothesis {
+  std::vector<int> arc_idx;
+  std::vector<int> olabel;
+  std::string out_str;
+
+  void clear() {
+    arc_idx.clear();
+    olabel.clear();
+    out_str.clear();
+  }
+};
+
 }  // end namespace cuda_decoder
 }  // end namespace kaldi
 

src/cudadecoder/cuda-decoder-kernels.cu

@@ -505,6 +505,7 @@ __global__ void reset_for_frame_and_estimate_cutoff_kernel(
       lane_counters->int_cutoff = INT_MAX;
       lane_counters->min_int_cost = INT_MAX;
       lane_counters->q_overflow = OVERFLOW_NONE;
+      lane_counters->int_relative_cost = INT_MAX;
       lane_counters->aux_q_requested = 0;
       lane_counters->main_q_requested = 0;
       lane_counters->main_q_local_offset = 0;
@@ -1122,8 +1123,7 @@ __launch_bounds__(KALDI_CUDA_DECODER_LARGEST_1D_BLOCK, 1) __global__
 // One add the final_cost[token.state] before looking for the min
 __global__ void get_best_cost_step1_kernel(DeviceParams cst_dev_params,
                                            KernelParams params,
-                                           bool use_final_probs,
-                                           CostType fst_zero) {
+                                           bool use_final_probs) {
   const int nlanes = params.nlanes_used;
   KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
     LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
@@ -1151,7 +1151,7 @@ __global__ void get_best_cost_step1_kernel(DeviceParams cst_dev_params,
             cst_dev_params.d_fst_final_costs[token_state];
         IntegerCostType int_cost_with_final =
             floatToOrderedInt(cost + final_cost);
-        if (final_cost != fst_zero) {
+        if (final_cost != cst_dev_params.fst_zero) {
           int2 min_and_arg = {int_cost_with_final,
                               global_idx};  // sort by cost, put it first
           atomicMinI2(&channel_counters->min_int_cost_and_arg_with_final,
@@ -1169,8 +1169,7 @@ __global__ void get_best_cost_step1_kernel(DeviceParams cst_dev_params,
 // and list them into d_list_final_tokens_in_main_q
 __global__ void get_best_cost_step2_kernel(DeviceParams cst_dev_params,
                                            KernelParams params,
-                                           bool use_final_probs,
-                                           CostType fst_zero) {
+                                           bool use_final_probs) {
   const int nlanes = params.nlanes_used;
   KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
     LaneCounters *lane_counters = cst_dev_params.d_lanes_counters.lane(ilane);
@@ -1218,7 +1217,7 @@ __global__ void get_best_cost_step2_kernel(DeviceParams cst_dev_params,
             cst_dev_params.d_fst_final_costs[token_state];
         const CostType token_cost = orderedIntToFloat(token_int_cost);
         // final_cost == fst_zero -> this state is not final
-        token_int_cost = (final_cost != fst_zero)
+        token_int_cost = (final_cost != cst_dev_params.fst_zero)
                              ? floatToOrderedInt(token_cost + final_cost)
                              : INT_MAX;
       }
@@ -1401,7 +1400,7 @@ __global__ void fill_hashmap_with_main_q_kernel(DeviceParams cst_dev_params,
     const int32 main_q_end = lane_counters->main_q_narcs_and_end.y;
     int32 min_int_cost = lane_counters->min_int_cost;
     CostType min_cost = orderedIntToFloat(min_int_cost);
-    const int32 global_offset = channel_counters->prev_main_q_global_offset;
+    const int32 global_offset = lane_counters->main_q_global_offset;
     KALDI_CUDA_DECODER_1D_KERNEL_LOOP(main_q_idx, main_q_end) {
       // Position of considered token in the main_q
       if (main_q_idx < main_q_end) {
@@ -1490,7 +1489,7 @@ __global__ void emitting_preprocess_and_list_extra_prev_tokens_step1_kernel(
   __shared__ typename BlockScan::TempStorage sh_temp_storage;
   const int nlanes = params.nlanes_used;
   KALDI_CUDA_DECODER_BATCH_KERNEL_LOOP(ilane, nlanes) {
-    const LaneCounters *lane_counters =
+    LaneCounters *lane_counters =
         cst_dev_params.d_lanes_counters.lane(ilane);
     const int32 main_q_end = lane_counters->main_q_narcs_and_end.y;
     // Final cutoff from last ExpandArc execution
@@ -1561,6 +1560,17 @@ __global__ void emitting_preprocess_and_list_extra_prev_tokens_step1_kernel(
             // avoid a new random memory access
             cst_dev_params.d_main_q_arc_offsets.channel(ichannel)[main_q_idx] =
                 start;
+
+	    // Saving best cost with final cost, to compute the final_extra_cost
+	    // It seems like ~5% of all states are final, so the following atomic may be fine
+	    // if necessary, we could first reduce locally at the CTA level
+	    const CostType final_cost =
+		    cst_dev_params.d_fst_final_costs[token_state];
+	    if(final_cost != cst_dev_params.fst_zero) {
+		    IntegerCostType token_int_cost_with_final = floatToOrderedInt(orderedIntToFloat(token_int_cost) + final_cost);
+		    IntegerCostType int_relative_cost = token_int_cost_with_final; // - 0.0f, the min_cost was reset to 0.0f
+		    atomicMin(&lane_counters->int_relative_cost, int_relative_cost);
+	    }
           }
           // If that FST state has only one token associated to it, we store
           // that token directly in
@@ -2023,20 +2033,18 @@ void FinalizeProcessNonEmittingKernel(const dim3 &grid, const dim3 &block,
 void GetBestCostStep1Kernel(const dim3 &grid, const dim3 &block,
                             const cudaStream_t &st,
                             const DeviceParams &cst_dev_params,
-                            const KernelParams &kernel_params, bool isfinal,
-                            CostType fst_zero) {
+                            const KernelParams &kernel_params, bool isfinal) {
   get_best_cost_step1_kernel<<<grid, block, 0, st>>>(
-      cst_dev_params, kernel_params, isfinal, fst_zero);
+      cst_dev_params, kernel_params, isfinal);
   KALDI_DECODER_CUDA_CHECK_ERROR();
 }
 
 void GetBestCostStep2Kernel(const dim3 &grid, const dim3 &block,
                             const cudaStream_t &st,
                             const DeviceParams &cst_dev_params,
-                            const KernelParams &kernel_params, bool isfinal,
-                            CostType fst_zero) {
+                            const KernelParams &kernel_params, bool isfinal) {
   get_best_cost_step2_kernel<<<grid, block, 0, st>>>(
-      cst_dev_params, kernel_params, isfinal, fst_zero);
+      cst_dev_params, kernel_params, isfinal);
   KALDI_DECODER_CUDA_CHECK_ERROR();
 }