From e15ab78052ca9f3a63174f5fd1862d7027d4be07 Mon Sep 17 00:00:00 2001 From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com> Date: Thu, 11 May 2023 09:50:14 -0700 Subject: [PATCH] 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. --- .../cpu/transformers/beam_search_impl_base.h | 88 ++++++++----------- .../cpu/transformers/beam_search_impl_gpt.h | 52 ++++------- .../cpu/transformers/beam_search_impl_t5.h | 51 +++-------- .../cpu/transformers/beam_search_scorer.cc | 63 ++++++------- .../cpu/transformers/beam_search_scorer.h | 18 ++-- .../cpu/transformers/generation_shared.h | 2 - .../contrib_ops/cpu/transformers/sequences.cc | 19 ++-- .../contrib_ops/cpu/transformers/sequences.h | 2 - 8 files changed, 103 insertions(+), 192 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h index 43024cc5bb..75d53ad160 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h @@ -13,21 +13,14 @@ namespace contrib { namespace transformers { template -struct BeamSearchState : public IBeamSearchState { - 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(batch_size) * num_beams; +struct BeamSearchState : IBeamSearchState { + BeamSearchState(const IGenerationParameters& parameters, + AllocatorPtr allocator, + int has_decoder_masked_attention, + bool use_position) { + size_t batch_beam_size = SafeInt(parameters.batch_size) * parameters.num_beams; - size_t next_token_size = SafeInt(batch_beam_size) * vocab_size; + size_t next_token_size = SafeInt(batch_beam_size) * parameters.vocab_size; this->next_token_logits = AllocateBuffer(allocator, next_token_logits_buffer_, next_token_size); this->next_token_scores = AllocateBuffer(allocator, next_token_scores_buffer_, next_token_size); @@ -38,7 +31,7 @@ struct BeamSearchState : public IBeamSearchState { this->next_scores = AllocateBuffer(allocator, next_scores_buffer_, SafeInt(2) * batch_beam_size); constexpr size_t max_parts_of_vocab = 128; - size_t topk_buffer_size = SafeInt(batch_beam_size) * (max_parts_of_vocab + 1) * num_beams * 2 * 2; + size_t topk_buffer_size = SafeInt(batch_beam_size) * (max_parts_of_vocab + 1) * parameters.num_beams * 2 * 2; this->topk_buffer = AllocateBuffer(allocator, topk_temp_buffer_, topk_buffer_size); if (use_position) { @@ -47,8 +40,8 @@ struct BeamSearchState : public IBeamSearchState { this->beam_scores = AllocateBuffer(allocator, beam_scores_buffer_, batch_beam_size); - if (output_scores) { - size_t elements = SafeInt(max_length - sequence_length) * batch_size * num_beams * vocab_size; + if (parameters.output_scores) { + size_t elements = SafeInt(parameters.max_length - parameters.sequence_length) * parameters.batch_size * parameters.num_beams * parameters.vocab_size; this->scores = AllocateBuffer(allocator, scores_buffer_, elements); this->remaining_scores = this->scores; } @@ -56,7 +49,7 @@ struct BeamSearchState : public IBeamSearchState { 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(batch_beam_size), num_heads, max_length, head_size}; + TensorShape staging_for_past_state_reorder_buffer_shape = {static_cast(batch_beam_size), parameters.num_heads, parameters.max_length, parameters.head_size}; Tensor temp(DataTypeImpl::GetType(), staging_for_past_state_reorder_buffer_shape, allocator); @@ -82,57 +75,50 @@ struct BeamSearchState : public IBeamSearchState { 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(allocator, sequence_lengths_buffer_, batch_beam_size); + BeamSearchCpuState(const IGenerationParameters& parameters, AllocatorPtr allocator, bool is_cuda) + : parameters_{parameters} { + sequence_lengths = AllocateBuffer(allocator, sequence_lengths_buffer_, batch_beam_size_); - size_t sequences_bytes = SafeInt(2) * batch_beam_size * max_length; - this->sequences_space = AllocateBuffer(allocator, sequences_space_buffer_, sequences_bytes); - memset(this->sequences_space.data(), 0, this->sequences_space.size_bytes()); + size_t sequences_bytes = SafeInt(2) * batch_beam_size_ * parameters.max_length; + sequences_space = AllocateBuffer(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(allocator, topk_scores_buffer_, 2 * batch_beam_size); - this->topk_tokens = AllocateBuffer(allocator, topk_tokens_buffer_, 2 * batch_beam_size); - this->topk_indices = AllocateBuffer(allocator, topk_indices_buffer_, 2 * batch_beam_size); - this->final_beam_scores = AllocateBuffer(allocator, final_beam_scores_buffer_, batch_beam_size); + topk_scores = AllocateBuffer(allocator, topk_scores_buffer_, 2 * static_cast(batch_beam_size_)); + topk_tokens = AllocateBuffer(allocator, topk_tokens_buffer_, 2 * static_cast(batch_beam_size_)); + topk_indices = AllocateBuffer(allocator, topk_indices_buffer_, 2 * static_cast(batch_beam_size_)); + final_beam_scores = AllocateBuffer(allocator, final_beam_scores_buffer_, batch_beam_size_); } - - this->sequences.Init(this->sequences_space, static_cast(batch_beam_size), sequence_length, max_length); } - // Copy expanded input_ids to sequences[0] - void SetSequence(gsl::span input_ids_in_cpu, - size_t batch_beam_size, - int max_length, - int sequence_length) { - gsl::span 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(i) * max_length + j; - sequences_0[index] = input_ids_in_cpu[SafeInt(i) * sequence_length + j]; + // Copy expanded input_ids to sequences_space + void SetExpandedSequence(gsl::span 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(i) * parameters_.max_length + j; + sequences_space[index] = input_ids_in_cpu[SafeInt(i) * parameters_.sequence_length + j]; } } } - // Copy unexpanded input_ids to sequences[0] - void SetSequence(gsl::span input_ids_in_cpu, - size_t batch_beam_size, - int beam_size, - int max_length, - int sequence_length) { - gsl::span 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(i) * max_length + j; - sequences_0[index] = input_ids_in_cpu[SafeInt(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 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(i) * parameters_.max_length + j; + sequences_space[index] = input_ids_in_cpu[SafeInt(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_; diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h index 3ea4a4c242..e6cba6be1a 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_gpt.h @@ -204,24 +204,14 @@ Status BeamSearchGpt::Execute(const FeedsFetchesManager* init_run_feeds_fetch // Initialize resources onnxruntime::OrtStlAllocator hypothesis_score_allocator(this->cpu_allocator_); onnxruntime::OrtStlAllocator beam_hyps_allocator(this->cpu_allocator_); - this->beam_scorer_ = std::make_unique(static_cast(parameters->batch_size), - static_cast(parameters->num_beams), - static_cast(parameters->max_length), - parameters->length_penalty, - parameters->early_stopping, - static_cast(parameters->num_return_sequences), - parameters->pad_token_id, - parameters->eos_token_id, + this->beam_scorer_ = std::make_unique(*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(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 buffer; @@ -243,19 +233,11 @@ Status BeamSearchGpt::Execute(const FeedsFetchesManager* init_run_feeds_fetch } } - BeamSearchState 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 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::Execute(const FeedsFetchesManager* init_run_feeds_fetch this->ort_stream_); gsl::span input_ids = expanded_input_ids_in_cpu.Get().DataAsSpan(); - cpu_state.SetSequence(input_ids, - static_cast(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::Execute(const FeedsFetchesManager* init_run_feeds_fetch } } - gsl::span final_beam_scores(beam_state.beam_scores.data(), beam_state.beam_scores.size()); + gsl::span 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(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::Execute(const FeedsFetchesManager* init_run_feeds_fetch output_sequences_scores); // Output per token scores - if (output_scores != nullptr) { + if (output_scores) { gsl::span target = output_scores->MutableDataAsSpan(); - gsl::span source = gsl::span(beam_state.scores.data(), beam_state.scores.size()); + gsl::span source = beam_state.scores; assert(target.size() == source.size()); ORT_RETURN_IF_ERROR(this->device_copy_func_(target, source, nullptr, DeviceCopyDirection::deviceToDevice)); } diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h index 51e8ae7b13..f63da7bc34 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_t5.h @@ -131,12 +131,9 @@ Status BeamSearchT5::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(parameters->BatchBeamSize()), - parameters->max_length, - parameters->sequence_length, - this->IsCuda()); + BeamSearchCpuState cpu_state{*parameters, + this->cpu_allocator_, + this->IsCuda()}; IAllocatorUniquePtr 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::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().DataAsSpan(), - static_cast(parameters->BatchBeamSize()), - parameters->num_beams, - parameters->max_length, - parameters->sequence_length); + cpu_state.SetUnexpandedSequence(decoder_input_ids.Get().DataAsSpan()); onnxruntime::OrtStlAllocator hypothesis_score_allocator(this->cpu_allocator_); onnxruntime::OrtStlAllocator beam_hyps_allocator(this->cpu_allocator_); - this->beam_scorer_ = std::make_unique(static_cast(parameters->batch_size), - static_cast(parameters->num_beams), - static_cast(parameters->max_length), - parameters->length_penalty, - parameters->early_stopping, - static_cast(parameters->num_return_sequences), - parameters->pad_token_id, - parameters->eos_token_id, + this->beam_scorer_ = std::make_unique(*parameters, hypothesis_score_allocator, - beam_hyps_allocator); - this->beam_scorer_->Initialize(this->cpu_allocator_, parameters->sequence_length); + beam_hyps_allocator, + this->cpu_allocator_); - BeamSearchState 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 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::Execute(const FeedsFetchesManager& encoder_feeds_fetches } } - gsl::span final_beam_scores(beam_state.beam_scores.data(), beam_state.beam_scores.size()); + gsl::span 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(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::Execute(const FeedsFetchesManager& encoder_feeds_fetches // Output per token scores if (output_scores != nullptr) { gsl::span target = output_scores->MutableDataAsSpan(); - gsl::span source = gsl::span(beam_state.scores.data(), beam_state.scores.size()); + gsl::span source = beam_state.scores; assert(target.size() == source.size()); ORT_RETURN_IF_ERROR(this->device_copy_func_(target, source, nullptr, DeviceCopyDirection::deviceToDevice)); } diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc b/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc index e26d4245c8..8227db44e7 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.cc +++ b/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& hypothesis_score_allocator, - onnxruntime::OrtStlAllocator& 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& beam_hyps_allocator, + AllocatorPtr& allocator) + : batch_size_{static_cast(parameters.batch_size)}, + num_beams_{static_cast(parameters.num_beams)}, + max_length_{static_cast(parameters.max_length)}, + num_beam_hyps_to_keep_{static_cast(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(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(allocator, batch_size_, done_ptr_, no_fill); - std::fill_n(done_.data(), done_.size(), false); - next_beam_scores_ = Allocate(allocator, batch_beam_size, next_beam_scores_ptr_, no_fill); next_beam_tokens_ = Allocate(allocator, batch_beam_size, next_beam_tokens_ptr_, no_fill); next_beam_indices_ = Allocate(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(max_length_) * (max_length_ + 1) - (sequence_length - 1) * sequence_length) / 2; + size_t per_beam = (SafeInt(max_length_) * (max_length_ + 1) - (parameters.sequence_length - 1) * parameters.sequence_length) / 2; hypothesis_buffer_length_ = batch_beam_size * per_beam; hypothesis_buffer_ = Allocate(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& next_scores, gsl::span& next_tokens, @@ -248,7 +235,7 @@ void BeamSearchScorer::Finalize(ISequences* sequences, // Score of each sequence, with shape (batch_size * num_return_sequences). gsl::span sequence_scores; - if (output_sequence_scores != nullptr) { + if (output_sequence_scores) { sequence_scores = output_sequence_scores->MutableDataAsSpan(); } @@ -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); } diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h index fdd404efa0..237328d2db 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_scorer.h +++ b/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& hypothesis_score_allocator, - onnxruntime::OrtStlAllocator& beam_hyps_allocator); - - void Initialize(AllocatorPtr& allocator, int sequence_length) override; + onnxruntime::OrtStlAllocator& beam_hyps_allocator, + AllocatorPtr& allocator); void Process(ISequences* sequences, gsl::span& next_scores, @@ -118,8 +110,8 @@ class BeamSearchScorer : public IBeamScorer { IAllocatorUniquePtr hypothesis_buffer_ptr_; // Allocated buffer to hold all hypotheses gsl::span 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 beam_hyps_; }; diff --git a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h index 79e9f04fad..bb105942bc 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h +++ b/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& next_scores, gsl::span& next_tokens, diff --git a/onnxruntime/contrib_ops/cpu/transformers/sequences.cc b/onnxruntime/contrib_ops/cpu/transformers/sequences.cc index a42aa47bfd..35bb2a8959 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/sequences.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/sequences.cc @@ -23,11 +23,8 @@ void Sequences::Init(gsl::span buffer, int batch_beam_size, int sequenc } gsl::span Sequences::GetSequence(int beam_index) const { - gsl::span buffer(sequences[current_sequences_buffer].data(), - sequences[current_sequences_buffer].size()); - gsl::span sequence = buffer.subspan(SafeInt(beam_index) * max_length_, - static_cast(current_length_)); - return sequence; + gsl::span buffer = sequences[current_sequences_buffer]; + return buffer.subspan(SafeInt(beam_index) * max_length_, static_cast(current_length_)); } int Sequences::GetSequenceLength() const { @@ -47,9 +44,8 @@ void Sequences::PrintSequences(const IConsoleDumper* dumper) const { void Sequences::AppendNextTokenToSequences( gsl::span& beam_indices, gsl::span& beam_next_tokens) { - gsl::span input(sequences[current_sequences_buffer].data(), - sequences[current_sequences_buffer].size()); - gsl::span output = sequences[1 - current_sequences_buffer]; + gsl::span input = sequences[current_sequences_buffer]; + gsl::span 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& next_tokens) { - gsl::span output(sequences[current_sequences_buffer].data(), sequences[current_sequences_buffer].size()); +void Sequences::AppendNextTokenToSequences(gsl::span& next_tokens) { + auto output = sequences[current_sequences_buffer]; // Append next token to each sequence. for (int i = 0; i < batch_beam_size_; i++) { diff --git a/onnxruntime/contrib_ops/cpu/transformers/sequences.h b/onnxruntime/contrib_ops/cpu/transformers/sequences.h index cd7b714028..6a16f119ba 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/sequences.h +++ b/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 buffer, int batch_beam_size, int sequence_length, int max_length);