Ryanunderhill/beamsearch simplify (#15883)

### Description
Simplify some sections of code by removing some extra gsl::span
conversions and passing parameter packs by an existing structure vs
directly.

### Motivation and Context
While stepping through the code, I noticed parts that could be
simplified. Simplifying then helped me understand it further.

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h

@@ -13,21 +13,14 @@ namespace contrib {
 namespace transformers {
 
 template <typename T>
-struct BeamSearchState : public IBeamSearchState<T> {
-  void Init(AllocatorPtr allocator,
-            int batch_size,
-            int num_beams,
-            int vocab_size,
-            int sequence_length,
-            int max_length,
-            int num_heads,
-            int head_size,
-            int has_decoder_masked_attention,
-            bool output_scores,
-            bool use_position) {
-    size_t batch_beam_size = SafeInt<size_t>(batch_size) * num_beams;
+struct BeamSearchState : IBeamSearchState<T> {
+  BeamSearchState(const IGenerationParameters& parameters,
+                  AllocatorPtr allocator,
+                  int has_decoder_masked_attention,
+                  bool use_position) {
+    size_t batch_beam_size = SafeInt<size_t>(parameters.batch_size) * parameters.num_beams;
 
-    size_t next_token_size = SafeInt<size_t>(batch_beam_size) * vocab_size;
+    size_t next_token_size = SafeInt<size_t>(batch_beam_size) * parameters.vocab_size;
     this->next_token_logits = AllocateBuffer<T>(allocator, next_token_logits_buffer_, next_token_size);
     this->next_token_scores = AllocateBuffer<float>(allocator, next_token_scores_buffer_, next_token_size);
 
@@ -38,7 +31,7 @@ struct BeamSearchState : public IBeamSearchState<T> {
     this->next_scores = AllocateBuffer<float>(allocator, next_scores_buffer_, SafeInt<size_t>(2) * batch_beam_size);
 
     constexpr size_t max_parts_of_vocab = 128;
-    size_t topk_buffer_size = SafeInt<size_t>(batch_beam_size) * (max_parts_of_vocab + 1) * num_beams * 2 * 2;
+    size_t topk_buffer_size = SafeInt<size_t>(batch_beam_size) * (max_parts_of_vocab + 1) * parameters.num_beams * 2 * 2;
     this->topk_buffer = AllocateBuffer<float>(allocator, topk_temp_buffer_, topk_buffer_size);
 
     if (use_position) {
@@ -47,8 +40,8 @@ struct BeamSearchState : public IBeamSearchState<T> {
 
     this->beam_scores = AllocateBuffer<float>(allocator, beam_scores_buffer_, batch_beam_size);
 
-    if (output_scores) {
-      size_t elements = SafeInt<size_t>(max_length - sequence_length) * batch_size * num_beams * vocab_size;
+    if (parameters.output_scores) {
+      size_t elements = SafeInt<size_t>(parameters.max_length - parameters.sequence_length) * parameters.batch_size * parameters.num_beams * parameters.vocab_size;
       this->scores = AllocateBuffer<float>(allocator, scores_buffer_, elements);
       this->remaining_scores = this->scores;
     }
@@ -56,7 +49,7 @@ struct BeamSearchState : public IBeamSearchState<T> {
     if (has_decoder_masked_attention) {
       // We need a temp staging buffer to do the past 'K' state re-ordering that is needed
       // when using DecoderMaskedSelfAttention
-      TensorShape staging_for_past_state_reorder_buffer_shape = {static_cast<int64_t>(batch_beam_size), num_heads, max_length, head_size};
+      TensorShape staging_for_past_state_reorder_buffer_shape = {static_cast<int64_t>(batch_beam_size), parameters.num_heads, parameters.max_length, parameters.head_size};
 
       Tensor temp(DataTypeImpl::GetType<T>(), staging_for_past_state_reorder_buffer_shape, allocator);
 
@@ -82,57 +75,50 @@ struct BeamSearchState : public IBeamSearchState<T> {
   BufferUniquePtr chosen_indices_buffer_;
 };
 
-struct BeamSearchCpuState : public IBeamSearchCpuState {
+struct BeamSearchCpuState : IBeamSearchCpuState {
   Sequences sequences;
 
-  void Init(AllocatorPtr allocator, size_t batch_beam_size, int max_length, int sequence_length, bool is_cuda) {
-    this->sequence_lengths = AllocateBuffer<int32_t>(allocator, sequence_lengths_buffer_, batch_beam_size);
+  BeamSearchCpuState(const IGenerationParameters& parameters, AllocatorPtr allocator, bool is_cuda)
+      : parameters_{parameters} {
+    sequence_lengths = AllocateBuffer<int32_t>(allocator, sequence_lengths_buffer_, batch_beam_size_);
 
-    size_t sequences_bytes = SafeInt<size_t>(2) * batch_beam_size * max_length;
-    this->sequences_space = AllocateBuffer<int32_t>(allocator, sequences_space_buffer_, sequences_bytes);
-    memset(this->sequences_space.data(), 0, this->sequences_space.size_bytes());
+    size_t sequences_bytes = SafeInt<size_t>(2) * batch_beam_size_ * parameters.max_length;
+    sequences_space = AllocateBuffer<int32_t>(allocator, sequences_space_buffer_, sequences_bytes, true /* fill */);
+    sequences.Init(sequences_space, batch_beam_size_, parameters.sequence_length, parameters.max_length);
 
     if (is_cuda) {
       // buffers used by CUDA operator but not by CPU operator.
-      this->topk_scores = AllocateBuffer<float>(allocator, topk_scores_buffer_, 2 * batch_beam_size);
-      this->topk_tokens = AllocateBuffer<int32_t>(allocator, topk_tokens_buffer_, 2 * batch_beam_size);
-      this->topk_indices = AllocateBuffer<int32_t>(allocator, topk_indices_buffer_, 2 * batch_beam_size);
-      this->final_beam_scores = AllocateBuffer<float>(allocator, final_beam_scores_buffer_, batch_beam_size);
+      topk_scores = AllocateBuffer<float>(allocator, topk_scores_buffer_, 2 * static_cast<size_t>(batch_beam_size_));
+      topk_tokens = AllocateBuffer<int32_t>(allocator, topk_tokens_buffer_, 2 * static_cast<size_t>(batch_beam_size_));
+      topk_indices = AllocateBuffer<int32_t>(allocator, topk_indices_buffer_, 2 * static_cast<size_t>(batch_beam_size_));
+      final_beam_scores = AllocateBuffer<float>(allocator, final_beam_scores_buffer_, batch_beam_size_);
     }
-
-    this->sequences.Init(this->sequences_space, static_cast<int>(batch_beam_size), sequence_length, max_length);
   }
 
-  // Copy expanded input_ids to sequences[0]
-  void SetSequence(gsl::span<const int32_t> input_ids_in_cpu,
-                   size_t batch_beam_size,
-                   int max_length,
-                   int sequence_length) {
-    gsl::span<int32_t> sequences_0 = sequences_space;
-    for (size_t i = 0; i < batch_beam_size; i++) {
-      for (int j = 0; j < sequence_length; j++) {
-        const size_t index = SafeInt<gsl::index>(i) * max_length + j;
-        sequences_0[index] = input_ids_in_cpu[SafeInt<gsl::index>(i) * sequence_length + j];
+  // Copy expanded input_ids to sequences_space
+  void SetExpandedSequence(gsl::span<const int32_t> input_ids_in_cpu) {
+    for (int i = 0; i < batch_beam_size_; i++) {
+      for (int j = 0; j < parameters_.sequence_length; j++) {
+        const size_t index = SafeInt<gsl::index>(i) * parameters_.max_length + j;
+        sequences_space[index] = input_ids_in_cpu[SafeInt<gsl::index>(i) * parameters_.sequence_length + j];
       }
     }
   }
 
-  // Copy unexpanded input_ids to sequences[0]
-  void SetSequence(gsl::span<const int32_t> input_ids_in_cpu,
-                   size_t batch_beam_size,
-                   int beam_size,
-                   int max_length,
-                   int sequence_length) {
-    gsl::span<int32_t> sequences_0 = sequences_space;
-    for (size_t i = 0; i < batch_beam_size; i++) {
-      for (int j = 0; j < sequence_length; j++) {
-        const size_t index = SafeInt<gsl::index>(i) * max_length + j;
-        sequences_0[index] = input_ids_in_cpu[SafeInt<gsl::index>(i / beam_size) * sequence_length + j];
+  // Copy unexpanded input_ids to sequences_space (only difference from SetExpandedSequence is i is divided by parameters_.num_beams
+  void SetUnexpandedSequence(gsl::span<const int32_t> input_ids_in_cpu) {
+    for (int i = 0; i < batch_beam_size_; i++) {
+      for (int j = 0; j < parameters_.sequence_length; j++) {
+        const size_t index = SafeInt<gsl::index>(i) * parameters_.max_length + j;
+        sequences_space[index] = input_ids_in_cpu[SafeInt<gsl::index>(i / parameters_.num_beams) * parameters_.sequence_length + j];
       }
     }
   }
 
  private:
+  const IGenerationParameters& parameters_;
+  const int batch_beam_size_{parameters_.batch_size * parameters_.num_beams};
+
   BufferUniquePtr final_beam_scores_buffer_;
   BufferUniquePtr sequence_lengths_buffer_;
   BufferUniquePtr topk_scores_buffer_;

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h

@@ -204,24 +204,14 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
   // Initialize resources
   onnxruntime::OrtStlAllocator<HypothesisScore> hypothesis_score_allocator(this->cpu_allocator_);
   onnxruntime::OrtStlAllocator<BeamHypotheses> beam_hyps_allocator(this->cpu_allocator_);
-  this->beam_scorer_ = std::make_unique<BeamSearchScorer>(static_cast<size_t>(parameters->batch_size),
-                                                          static_cast<size_t>(parameters->num_beams),
-                                                          static_cast<size_t>(parameters->max_length),
-                                                          parameters->length_penalty,
-                                                          parameters->early_stopping,
-                                                          static_cast<size_t>(parameters->num_return_sequences),
-                                                          parameters->pad_token_id,
-                                                          parameters->eos_token_id,
+  this->beam_scorer_ = std::make_unique<BeamSearchScorer>(*parameters,
                                                           hypothesis_score_allocator,
-                                                          beam_hyps_allocator);
-  this->beam_scorer_->Initialize(this->cpu_allocator_, parameters->sequence_length);
+                                                          beam_hyps_allocator,
+                                                          this->cpu_allocator_);
 
-  BeamSearchCpuState cpu_state;
-  cpu_state.Init(this->cpu_allocator_,
-                 static_cast<size_t>(parameters->BatchBeamSize()),
-                 parameters->max_length,
-                 parameters->sequence_length,
-                 this->IsCuda());
+  BeamSearchCpuState cpu_state{*parameters,
+                               this->cpu_allocator_,
+                               this->IsCuda()};
 
   // buffer in GPU for input_ids, position_ids and attention_mask
   IAllocatorUniquePtr<char> buffer;
@@ -243,19 +233,11 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
     }
   }
 
-  BeamSearchState<T> beam_state;
   constexpr bool use_position = true;
-  beam_state.Init(this->temp_space_allocator_,
-                  parameters->batch_size,
-                  parameters->num_beams,
-                  parameters->vocab_size,
-                  parameters->sequence_length,
-                  parameters->max_length,
-                  parameters->num_heads,
-                  parameters->head_size,
-                  gpt_subgraph_.has_decoder_masked_attention_,
-                  parameters->output_scores,
-                  use_position);
+  BeamSearchState<T> beam_state{*parameters,
+                                this->temp_space_allocator_,
+                                gpt_subgraph_.has_decoder_masked_attention_,
+                                use_position};
 
   init_beam_state_func_(&beam_state,
                         cpu_state.sequence_lengths,
@@ -264,10 +246,7 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
                         this->ort_stream_);
 
   gsl::span<const int32_t> input_ids = expanded_input_ids_in_cpu.Get<Tensor>().DataAsSpan<int32_t>();
-  cpu_state.SetSequence(input_ids,
-                        static_cast<size_t>(parameters->BatchBeamSize()),
-                        parameters->max_length,
-                        parameters->sequence_length);
+  cpu_state.SetExpandedSequence(input_ids);
 
 #ifdef DEBUG_GENERATION
   const IConsoleDumper* dumper = this->GetConsoleDumper();
@@ -394,14 +373,13 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
     }
   }
 
-  gsl::span<const float> final_beam_scores(beam_state.beam_scores.data(), beam_state.beam_scores.size());
+  gsl::span<const float> final_beam_scores = beam_state.beam_scores;
   if (this->IsCuda()) {
     ORT_RETURN_IF_ERROR(this->device_copy_func_(cpu_state.final_beam_scores,
                                                 final_beam_scores,
                                                 nullptr,
                                                 DeviceCopyDirection::deviceToHost));
-    final_beam_scores = gsl::make_span<const float>(cpu_state.final_beam_scores.data(),
-                                                    cpu_state.final_beam_scores.size());
+    final_beam_scores = cpu_state.final_beam_scores;
   }
 
   this->beam_scorer_->Finalize(&(cpu_state.sequences),
@@ -410,9 +388,9 @@ Status BeamSearchGpt<T>::Execute(const FeedsFetchesManager* init_run_feeds_fetch
                                output_sequences_scores);
 
   // Output per token scores
-  if (output_scores != nullptr) {
+  if (output_scores) {
     gsl::span<float> target = output_scores->MutableDataAsSpan<float>();
-    gsl::span<const float> source = gsl::span<const float>(beam_state.scores.data(), beam_state.scores.size());
+    gsl::span<const float> source = beam_state.scores;
     assert(target.size() == source.size());
     ORT_RETURN_IF_ERROR(this->device_copy_func_(target, source, nullptr, DeviceCopyDirection::deviceToDevice));
   }

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h

@@ -131,12 +131,9 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
 
   const OrtValue* encoder_attn_mask_value = this->context_.GetInputOrtValue(9);
 
-  BeamSearchCpuState cpu_state;
-  cpu_state.Init(this->cpu_allocator_,
-                 static_cast<size_t>(parameters->BatchBeamSize()),
-                 parameters->max_length,
-                 parameters->sequence_length,
-                 this->IsCuda());
+  BeamSearchCpuState cpu_state{*parameters,
+                               this->cpu_allocator_,
+                               this->IsCuda()};
 
   IAllocatorUniquePtr<char> buffer;
   OrtValue decoder_input_ids;  // Tensor in CPU, and it will be used to initialize sequence in cpu_state
@@ -184,39 +181,20 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
   // ------------------------------------
 
   // Copy decoder_input_ids (in CPU) to sequence. It contains decoder_start_token_id for each beam.
-  cpu_state.SetSequence(decoder_input_ids.Get<Tensor>().DataAsSpan<int32_t>(),
-                        static_cast<size_t>(parameters->BatchBeamSize()),
-                        parameters->num_beams,
-                        parameters->max_length,
-                        parameters->sequence_length);
+  cpu_state.SetUnexpandedSequence(decoder_input_ids.Get<Tensor>().DataAsSpan<int32_t>());
 
   onnxruntime::OrtStlAllocator<HypothesisScore> hypothesis_score_allocator(this->cpu_allocator_);
   onnxruntime::OrtStlAllocator<BeamHypotheses> beam_hyps_allocator(this->cpu_allocator_);
-  this->beam_scorer_ = std::make_unique<BeamSearchScorer>(static_cast<size_t>(parameters->batch_size),
-                                                          static_cast<size_t>(parameters->num_beams),
-                                                          static_cast<size_t>(parameters->max_length),
-                                                          parameters->length_penalty,
-                                                          parameters->early_stopping,
-                                                          static_cast<size_t>(parameters->num_return_sequences),
-                                                          parameters->pad_token_id,
-                                                          parameters->eos_token_id,
+  this->beam_scorer_ = std::make_unique<BeamSearchScorer>(*parameters,
                                                           hypothesis_score_allocator,
-                                                          beam_hyps_allocator);
-  this->beam_scorer_->Initialize(this->cpu_allocator_, parameters->sequence_length);
+                                                          beam_hyps_allocator,
+                                                          this->cpu_allocator_);
 
-  BeamSearchState<T> beam_state;
   constexpr bool use_position = false;
-  beam_state.Init(this->temp_space_allocator_,
-                  parameters->batch_size,
-                  parameters->num_beams,
-                  parameters->vocab_size,
-                  parameters->sequence_length,
-                  parameters->max_length,
-                  parameters->num_heads,
-                  parameters->head_size,
-                  decoder_subgraph_.has_decoder_masked_attention_,
-                  parameters->output_scores,
-                  use_position);
+  BeamSearchState<T> beam_state{*parameters,
+                                this->temp_space_allocator_,
+                                decoder_subgraph_.has_decoder_masked_attention_,
+                                use_position};
 
   init_beam_state_func_(&beam_state,
                         cpu_state.sequence_lengths,
@@ -386,14 +364,13 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
     }
   }
 
-  gsl::span<const float> final_beam_scores(beam_state.beam_scores.data(), beam_state.beam_scores.size());
+  gsl::span<const float> final_beam_scores = beam_state.beam_scores;
   if (this->IsCuda()) {
     ORT_RETURN_IF_ERROR(this->device_copy_func_(cpu_state.final_beam_scores,
                                                 final_beam_scores,
                                                 nullptr,
                                                 DeviceCopyDirection::deviceToHost));
-    final_beam_scores = gsl::make_span<const float>(cpu_state.final_beam_scores.data(),
-                                                    cpu_state.final_beam_scores.size());
+    final_beam_scores = cpu_state.final_beam_scores;
   }
 
   this->beam_scorer_->Finalize(&(cpu_state.sequences),
@@ -404,7 +381,7 @@ Status BeamSearchT5<T>::Execute(const FeedsFetchesManager& encoder_feeds_fetches
   // Output per token scores
   if (output_scores != nullptr) {
     gsl::span<float> target = output_scores->MutableDataAsSpan<float>();
-    gsl::span<const float> source = gsl::span<const float>(beam_state.scores.data(), beam_state.scores.size());
+    gsl::span<const float> source = beam_state.scores;
     assert(target.size() == source.size());
     ORT_RETURN_IF_ERROR(this->device_copy_func_(target, source, nullptr, DeviceCopyDirection::deviceToDevice));
   }

onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc

@@ -89,57 +89,44 @@ void BeamHypotheses::Output(
   }
 }
 
-BeamSearchScorer::BeamSearchScorer(size_t batch_size,
-                                   size_t num_beams,
-                                   size_t max_length,
-                                   float length_penalty,
-                                   bool early_stopping,
-                                   size_t num_return_sequences,
-                                   int pad_token_id,
-                                   int eos_token_id,
+BeamSearchScorer::BeamSearchScorer(const IGenerationParameters& parameters,
                                    onnxruntime::OrtStlAllocator<HypothesisScore>& hypothesis_score_allocator,
-                                   onnxruntime::OrtStlAllocator<BeamHypotheses>& beam_hyps_allocator)
-    : batch_size_(batch_size),
-      num_beams_(num_beams),
-      max_length_(max_length),
-      num_beam_hyps_to_keep_(num_return_sequences),
-      pad_token_id_(pad_token_id),
-      eos_token_id_(eos_token_id),
-      hypothesis_buffer_length_(0),
-      hypothesis_buffer_offset_(0),
+                                   onnxruntime::OrtStlAllocator<BeamHypotheses>& beam_hyps_allocator,
+                                   AllocatorPtr& allocator)
+    : batch_size_{static_cast<size_t>(parameters.batch_size)},
+      num_beams_{static_cast<size_t>(parameters.num_beams)},
+      max_length_{static_cast<size_t>(parameters.max_length)},
+      num_beam_hyps_to_keep_{static_cast<size_t>(parameters.num_return_sequences)},
+      pad_token_id_{parameters.pad_token_id},
+      eos_token_id_{parameters.eos_token_id},
       beam_hyps_(beam_hyps_allocator) {
-  for (size_t i = 0; i < batch_size; i++) {
-    beam_hyps_.push_back(BeamHypotheses(num_beams, length_penalty, early_stopping, hypothesis_score_allocator));
+  for (size_t i = 0; i < batch_size_; i++) {
+    beam_hyps_.push_back(BeamHypotheses(num_beams_, parameters.length_penalty, parameters.early_stopping, hypothesis_score_allocator));
   }
-}
-
-bool BeamSearchScorer::IsDone() {
-  for (size_t batch = 0; batch < batch_size_; batch++) {
-    if (!done_[batch])
-      return false;
-  }
-  return true;
-}
-
-void BeamSearchScorer::Initialize(AllocatorPtr& allocator, int sequence_length) {
-  ORT_ENFORCE(next_beam_scores_.empty());  // Make sure this is called only once.
 
   size_t batch_beam_size = batch_size_ * num_beams_;
+
+  done_ = Allocate<bool>(allocator, batch_size_, done_ptr_, true /* fill allocated array */, false /* fill with false */);
+
   constexpr bool no_fill = false;  // Do not fill values after allocation
-
-  done_ = Allocate<bool>(allocator, batch_size_, done_ptr_, no_fill);
-  std::fill_n(done_.data(), done_.size(), false);
-
   next_beam_scores_ = Allocate<float>(allocator, batch_beam_size, next_beam_scores_ptr_, no_fill);
   next_beam_tokens_ = Allocate<int32_t>(allocator, batch_beam_size, next_beam_tokens_ptr_, no_fill);
   next_beam_indices_ = Allocate<int32_t>(allocator, batch_beam_size, next_beam_indices_ptr_, no_fill);
 
   // Space to store intermediate sequence with length sequence_length, sequence_length + 1, ..., max_sequence_length.
-  size_t per_beam = (SafeInt<size_t>(max_length_) * (max_length_ + 1) - (sequence_length - 1) * sequence_length) / 2;
+  size_t per_beam = (SafeInt<size_t>(max_length_) * (max_length_ + 1) - (parameters.sequence_length - 1) * parameters.sequence_length) / 2;
   hypothesis_buffer_length_ = batch_beam_size * per_beam;
   hypothesis_buffer_ = Allocate<int32_t>(allocator, hypothesis_buffer_length_, hypothesis_buffer_ptr_, no_fill);
 }
 
+bool BeamSearchScorer::IsDone() {
+  for (auto done : done_) {
+    if (!done)
+      return false;
+  }
+  return true;
+}
+
 void BeamSearchScorer::Process(ISequences* sequences,
                                gsl::span<const float>& next_scores,
                                gsl::span<const int32_t>& next_tokens,
@@ -248,7 +235,7 @@ void BeamSearchScorer::Finalize(ISequences* sequences,
 
   // Score of each sequence, with shape (batch_size * num_return_sequences).
   gsl::span<float> sequence_scores;
-  if (output_sequence_scores != nullptr) {
+  if (output_sequence_scores) {
     sequence_scores = output_sequence_scores->MutableDataAsSpan<float>();
   }
 
@@ -263,7 +250,7 @@ void BeamSearchScorer::Finalize(ISequences* sequences,
     auto batch_output = output.subspan(batch_index * num_return_sequences * max_length_,
                                        num_return_sequences * max_length_);
 
-    if (output_sequence_scores != nullptr) {
+    if (output_sequence_scores) {
       batch_sequence_score = sequence_scores.subspan(batch_index * num_return_sequences, num_return_sequences);
     }
 

onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h

@@ -67,18 +67,10 @@ class BeamHypotheses {
 
 class BeamSearchScorer : public IBeamScorer {
  public:
-  BeamSearchScorer(size_t batch_size,
-                   size_t num_beams,
-                   size_t max_length,
-                   float length_penalty,
-                   bool early_stopping,
-                   size_t num_return_sequences,
-                   int pad_token_id,
-                   int eos_token_id,
+  BeamSearchScorer(const IGenerationParameters& parameters,
                    onnxruntime::OrtStlAllocator<HypothesisScore>& hypothesis_score_allocator,
-                   onnxruntime::OrtStlAllocator<BeamHypotheses>& beam_hyps_allocator);
-
-  void Initialize(AllocatorPtr& allocator, int sequence_length) override;
+                   onnxruntime::OrtStlAllocator<BeamHypotheses>& beam_hyps_allocator,
+                   AllocatorPtr& allocator);
 
   void Process(ISequences* sequences,
                gsl::span<const float>& next_scores,
@@ -118,8 +110,8 @@ class BeamSearchScorer : public IBeamScorer {
 
   IAllocatorUniquePtr<int32_t> hypothesis_buffer_ptr_;  // Allocated buffer to hold all hypotheses
   gsl::span<int32_t> hypothesis_buffer_;                // Span of the allocated buffer
-  size_t hypothesis_buffer_length_;                     // Total number of elements
-  size_t hypothesis_buffer_offset_;                     // Offset of available buffer, or length of used buffer.
+  size_t hypothesis_buffer_length_{};                   // Total number of elements
+  size_t hypothesis_buffer_offset_{};                   // Offset of available buffer, or length of used buffer.
 
   onnxruntime::FastAllocVector<BeamHypotheses> beam_hyps_;
 };

onnxruntime/contrib_ops/cpu/transformers/generation_shared.h

@@ -111,8 +111,6 @@ class IBeamScorer {
  public:
   virtual ~IBeamScorer() {}
 
-  virtual void Initialize(AllocatorPtr& allocator, int sequence_length) = 0;
-
   virtual void Process(ISequences* sequences,
                        gsl::span<const float>& next_scores,
                        gsl::span<const int32_t>& next_tokens,

onnxruntime/contrib_ops/cpu/transformers/sequences.cc

@@ -23,11 +23,8 @@ void Sequences::Init(gsl::span<int32_t> buffer, int batch_beam_size, int sequenc
 }
 
 gsl::span<const int32_t> Sequences::GetSequence(int beam_index) const {
-  gsl::span<const int32_t> buffer(sequences[current_sequences_buffer].data(),
-                                  sequences[current_sequences_buffer].size());
-  gsl::span<const int32_t> sequence = buffer.subspan(SafeInt<size_t>(beam_index) * max_length_,
-                                                     static_cast<gsl::index>(current_length_));
-  return sequence;
+  gsl::span<const int32_t> buffer = sequences[current_sequences_buffer];
+  return buffer.subspan(SafeInt<size_t>(beam_index) * max_length_, static_cast<gsl::index>(current_length_));
 }
 
 int Sequences::GetSequenceLength() const {
@@ -47,9 +44,8 @@ void Sequences::PrintSequences(const IConsoleDumper* dumper) const {
 void Sequences::AppendNextTokenToSequences(
     gsl::span<int32_t>& beam_indices,
     gsl::span<int32_t>& beam_next_tokens) {
-  gsl::span<const int32_t> input(sequences[current_sequences_buffer].data(),
-                                 sequences[current_sequences_buffer].size());
-  gsl::span<int32_t> output = sequences[1 - current_sequences_buffer];
+  gsl::span<const int32_t> input = sequences[current_sequences_buffer];
+  gsl::span<int32_t> output = sequences[current_sequences_buffer ^ 1];
 
   for (int i = 0; i < batch_beam_size_; i++) {
     int beam_index = beam_indices[i];
@@ -68,12 +64,11 @@ void Sequences::AppendNextTokenToSequences(
   ++current_length_;
 
   // Rotate buffer for next round.
-  current_sequences_buffer = 1 - current_sequences_buffer;
+  current_sequences_buffer ^= 1;
 }
 
-void Sequences::AppendNextTokenToSequences(
-    gsl::span<int32_t>& next_tokens) {
-  gsl::span<int32_t> output(sequences[current_sequences_buffer].data(), sequences[current_sequences_buffer].size());
+void Sequences::AppendNextTokenToSequences(gsl::span<int32_t>& next_tokens) {
+  auto output = sequences[current_sequences_buffer];
 
   // Append next token to each sequence.
   for (int i = 0; i < batch_beam_size_; i++) {

onnxruntime/contrib_ops/cpu/transformers/sequences.h

@@ -13,8 +13,6 @@ namespace transformers {
 // This class keeps track of sequences generated.
 class Sequences : public ISequences {
  public:
-  Sequences() {}
-
   // Initialize the sequence.
   void Init(gsl::span<int32_t> buffer, int batch_beam_size, int sequence_length, int max_length);