diff --git a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h index 9f372e5b3a673..6a325df7eb0ea 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h +++ b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h @@ -6,6 +6,7 @@ #include #include "contrib_ops/cpu/transformers/generation_shared.h" #include "contrib_ops/cpu/transformers/generate_impl_base.h" +#include "contrib_ops/cpu/transformers/sampling_buffer_element_count.h" namespace onnxruntime { namespace contrib { @@ -22,25 +23,32 @@ struct SamplingState : public ISamplingState { int seed, bool is_cuda, Stream* stream) { - int total_count = batch_size * vocab_size; + // Compute the product entirely in SafeInt's checked domain; an `int * int` multiply + // with model-controlled operands can silently positive-wrap before any SafeInt cast + // sees it, leading to under-allocated buffers (heap-buffer-overflow on the downstream + // memcpy in SamplingCpuHelper::Sample). `SafeMul` also rejects negative + // inputs (e.g. an unvalidated `vocab_size` of -1) up front, which a + // `static_cast` would silently turn into a huge value. `SamplingBufferElementCount` + // is the shared helper that regression tests exercise so a revert breaks them. + const size_t total_count = SamplingBufferElementCount(batch_size, vocab_size); - this->h_softmaxed_score = AllocateBuffer(cpu_allocator, h_softmaxed_score_buffer_, SafeInt(total_count), stream); + this->h_softmaxed_score = AllocateBuffer(cpu_allocator, h_softmaxed_score_buffer_, total_count, stream); this->generator = std::default_random_engine{gsl::narrow_cast(seed)}; if (is_cuda) { - this->d_index_in = AllocateBuffer(allocator, d_index_in_buffer_, SafeInt(total_count), stream); - this->d_index_out = AllocateBuffer(allocator, d_index_out_buffer_, SafeInt(total_count), stream); - this->d_offset = AllocateBuffer(allocator, d_offset_buffer_, SafeInt(batch_size + 1), stream); - this->d_sorted_score = AllocateBuffer(allocator, d_sorted_score_buffer_, SafeInt(total_count), stream); - this->d_sorted_softmaxed_score = AllocateBuffer(allocator, d_sorted_softmaxed_score_buffer_, SafeInt(total_count), stream); - this->d_softmaxed_score = AllocateBuffer(allocator, d_softmaxed_score_buffer_, SafeInt(total_count), stream); + this->d_index_in = AllocateBuffer(allocator, d_index_in_buffer_, total_count, stream); + this->d_index_out = AllocateBuffer(allocator, d_index_out_buffer_, total_count, stream); + this->d_offset = AllocateBuffer(allocator, d_offset_buffer_, SafeInt(batch_size) + 1, stream); + this->d_sorted_score = AllocateBuffer(allocator, d_sorted_score_buffer_, total_count, stream); + this->d_sorted_softmaxed_score = AllocateBuffer(allocator, d_sorted_softmaxed_score_buffer_, total_count, stream); + this->d_softmaxed_score = AllocateBuffer(allocator, d_softmaxed_score_buffer_, total_count, stream); this->d_sampled = AllocateBuffer(allocator, d_sampled_buffer_, SafeInt(batch_size), stream); - this->h_sampled_all = AllocateBuffer(cpu_allocator, h_sampled_all_buffer_, SafeInt(batch_size * max_iter), stream); + this->h_sampled_all = AllocateBuffer(cpu_allocator, h_sampled_all_buffer_, SamplingBufferElementCount(batch_size, max_iter), stream); this->d_indices = AllocateBuffer(allocator, d_indices_buffer_, SafeInt(batch_size), stream); this->temp_storage_bytes = 0; // TODO: Do not allocate this buffer if there's no presence_mask - this->d_presence_mask = AllocateBuffer(allocator, d_presence_mask_buffer_, SafeInt(total_count), stream); + this->d_presence_mask = AllocateBuffer(allocator, d_presence_mask_buffer_, total_count, stream); std::uniform_real_distribution distribution(0.0, 1.0); static_cast(distribution(this->generator)); @@ -49,8 +57,8 @@ struct SamplingState : public ISamplingState { } } else { // TODO: Some buffer can be reused for CPU - this->sorted_scores = AllocateBuffer(cpu_allocator, sorted_scores_buffer_, SafeInt(total_count), stream); - this->cumulative_probs = AllocateBuffer(cpu_allocator, cumulative_probs_buffer_, SafeInt(total_count), stream); + this->sorted_scores = AllocateBuffer(cpu_allocator, sorted_scores_buffer_, total_count, stream); + this->cumulative_probs = AllocateBuffer(cpu_allocator, cumulative_probs_buffer_, total_count, stream); } } diff --git a/onnxruntime/contrib_ops/cpu/transformers/sampling_buffer_element_count.h b/onnxruntime/contrib_ops/cpu/transformers/sampling_buffer_element_count.h new file mode 100644 index 0000000000000..b98f3a37b47ae --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/transformers/sampling_buffer_element_count.h @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/safeint.h" + +namespace onnxruntime { +namespace contrib { +namespace transformers { + +// Computes the element count for a `SamplingState` buffer using +// `SafeMul` so that overflow (from a plain `int * int` multiply) or +// negative-to-unsigned conversion (e.g. an unvalidated `-1`) of the +// model-controlled operands is detected up front (throws +// `OnnxRuntimeException`) rather than silently producing an under-sized or +// absurdly large allocation. Extracted into a self-contained header so that +// both `SamplingState::Init` and its regression tests call the exact same +// production code path; reverting this computation to `int * int` or +// `static_cast(...)` will fail the tests. +inline size_t SamplingBufferElementCount(int batch_size, int per_batch_count) { + return SafeMul(batch_size, per_batch_count); +} + +} // namespace transformers +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/sampling_state_test.cc b/onnxruntime/test/contrib_ops/sampling_state_test.cc new file mode 100644 index 0000000000000..8b8fc7514654b --- /dev/null +++ b/onnxruntime/test/contrib_ops/sampling_state_test.cc @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Regression tests for the buffer-size arithmetic in +// onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h +// (SamplingState::Init). The bug being guarded was that a plain `int * int` +// multiply of `batch_size * vocab_size`, or a `static_cast(...)` of a +// negative/unvalidated operand, could silently wrap and lead to under-allocated +// buffers (heap-buffer-overflow on the downstream memcpy in +// SamplingCpuHelper::Sample). +// +// The production computation lives in `sampling_buffer_element_count.h` as +// `SamplingBufferElementCount`, which `SamplingState::Init` calls directly. +// These tests call the same helper, so reverting the production code to +// `int * int` or `static_cast` on operands that may be negative will +// fail these tests. + +#include "gtest/gtest.h" + +#include "core/common/common.h" +#include "core/common/safeint.h" +#include "contrib_ops/cpu/transformers/sampling_buffer_element_count.h" + +namespace onnxruntime { +namespace test { + +using contrib::transformers::SamplingBufferElementCount; + +// Sanity check: well-formed inputs produce the expected element count. +TEST(SamplingStateArithmeticTest, ProducesExpectedTotalCountForValidInputs) { + EXPECT_EQ(SamplingBufferElementCount(4, 32), static_cast(4) * 32u); + EXPECT_EQ(SamplingBufferElementCount(1, 50257), static_cast(50257)); + // `h_sampled_all` uses the same helper with `max_iter` as the second operand. + EXPECT_EQ(SamplingBufferElementCount(8, 16), static_cast(8) * 16u); +} + +// A negative `vocab_size` (e.g. an unvalidated default of -1) used to be turned +// into SIZE_MAX by `static_cast(vocab_size)`, yielding a multiplication +// result that either silently wrapped or requested an absurdly large buffer. +// SafeInt rejects the negative-to-unsigned conversion up front. +TEST(SamplingStateArithmeticTest, ThrowsOnNegativeVocabSize) { + EXPECT_THROW(SamplingBufferElementCount(4, -1), OnnxRuntimeException); +} + +// Symmetric check for a negative `batch_size`. +TEST(SamplingStateArithmeticTest, ThrowsOnNegativeBatchSize) { + EXPECT_THROW(SamplingBufferElementCount(-1, 32), OnnxRuntimeException); +} + +// `max_iter` flows through the same helper for the `h_sampled_all` +// allocation, so a negative value in the second operand slot must also be +// rejected. Uses a valid `batch_size` and a negative `max_iter` to exercise +// the operand distinctly from `ThrowsOnNegativeVocabSize` above. +TEST(SamplingStateArithmeticTest, ThrowsOnNegativeMaxIter) { + constexpr int batch_size = 4; + constexpr int max_iter = -1; + EXPECT_THROW(SamplingBufferElementCount(batch_size, max_iter), OnnxRuntimeException); +} + +// The core bug this PR guards against: a plain `int * int` multiply of +// `batch_size * vocab_size` silently wraps once the mathematical product +// exceeds `INT_MAX`, producing an under-allocated buffer. `SafeMul` +// promotes both operands to `size_t` before multiplying, so the helper must +// return the correct non-wrapped product. Reverting the production code to +// `int * int` (or `static_cast(int_product)`) fails this test. +TEST(SamplingStateArithmeticTest, ReturnsCorrectProductWhenIntMultiplyWouldOverflow) { + // 50000 * 50000 = 2.5e9, which exceeds INT_MAX (~2.147e9) and would + // signed-overflow if computed in `int`, but still fits in a 32-bit + // `size_t` so the test is portable to 32-bit Windows builds where + // `size_t` is 32-bit (SIZE_MAX ~ 4.29e9). + constexpr int large_operand = 50000; + const size_t expected = static_cast(large_operand) * static_cast(large_operand); + EXPECT_EQ(SamplingBufferElementCount(large_operand, large_operand), expected); +} + +} // namespace test +} // namespace onnxruntime