diff --git a/backends/apple/metal/runtime/metal_backend.cpp b/backends/apple/metal/runtime/metal_backend.cpp index 2975a2d2375..c0d996df62b 100644 --- a/backends/apple/metal/runtime/metal_backend.cpp +++ b/backends/apple/metal/runtime/metal_backend.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include @@ -459,8 +460,10 @@ class ET_EXPERIMENTAL MetalBackend final ET_LOG(Debug, "MetalBackend n_outputs %zd generated", n_outputs); + size_t n_io_sum = 0; ET_CHECK_OR_RETURN_ERROR( - n_inputs + n_outputs == args.size(), + !c10::add_overflows(n_inputs, n_outputs, &n_io_sum) && + n_io_sum == args.size(), InvalidArgument, "number of user input %zd and output %zd generated from AOT Inductor does not match ET runner's %zd. Exit.", n_inputs, diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index eb0a07b8d8f..4cd6b3acf51 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include @@ -550,8 +551,10 @@ class ET_EXPERIMENTAL CudaBackend final setCurrentCUDAStream(handle->get_cuda_stream(), 0); + size_t n_io_sum = 0; ET_CHECK_OR_RETURN_ERROR( - n_inputs + n_outputs == args.size(), + !c10::add_overflows(n_inputs, n_outputs, &n_io_sum) && + n_io_sum == args.size(), InvalidArgument, "number of user input %zd and output %zd generated from AOT Inductor does not match ET runner's %zd. Exit.", n_inputs, diff --git a/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp b/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp index 916931226ea..59744d488bd 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include using executorch::aten::TensorImpl; @@ -248,9 +249,12 @@ Result PromptProcessor::prefill( ET_CHECK_MSG( start_pos == 0, "Bert model doesn't support multi-turn conversation."); } else if (!enable_attention_sink) { + int64_t end_pos = 0; ET_CHECK_MSG( - (start_pos + num_prompt_tokens) <= - (metadata_.context_len - metadata_.ar_len), + !c10::add_overflows( + start_pos, static_cast(num_prompt_tokens), &end_pos) && + end_pos <= static_cast(metadata_.context_len) - + static_cast(metadata_.ar_len), "The sequence length exceeds the maximum limit that the prompt processor can handle."); } diff --git a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp index dccdda17fa2..0e9b7860dbd 100644 --- a/examples/qualcomm/oss_scripts/llama/runner/runner.cpp +++ b/examples/qualcomm/oss_scripts/llama/runner/runner.cpp @@ -9,6 +9,7 @@ // A llama 3.2 runner that includes preprocessing and post processing // logic. The module takes in a string as input and emits a string as output. +#include #include #include #include @@ -433,8 +434,11 @@ Error Runner::generate_from_prompt_or_file( } int num_prompt_tokens = prompt_tokens.size(); ET_CHECK_MSG(num_prompt_tokens >= 1, "Expected at least 1 prompt token"); + int64_t end_pos = 0; ET_CHECK_MSG( - cur_pos_ + num_prompt_tokens < seq_len, + !c10::add_overflows( + cur_pos_, static_cast(num_prompt_tokens), &end_pos) && + end_pos < static_cast(seq_len), "sequence length exceeded - please increase the seq_len value"); // Prompt Processor first diff --git a/extension/flat_tensor/flat_tensor_data_map.cpp b/extension/flat_tensor/flat_tensor_data_map.cpp index 327d3d231ff..48684da1239 100644 --- a/extension/flat_tensor/flat_tensor_data_map.cpp +++ b/extension/flat_tensor/flat_tensor_data_map.cpp @@ -8,6 +8,8 @@ #include +#include + #include #include @@ -73,9 +75,13 @@ Result get_named_data( key.data(), segments->size()); // Validate the segment. + uint64_t seg_end = 0; ET_CHECK_OR_RETURN_ERROR( - (segments->Get(segment_index)->offset() + - segments->Get(segment_index)->size()) <= segment_end_offset, + !c10::add_overflows( + static_cast(segments->Get(segment_index)->offset()), + static_cast(segments->Get(segment_index)->size()), + &seg_end) && + seg_end <= static_cast(segment_end_offset), InvalidExternalData, "Invalid segment offset %" PRIu64 " is larger than the segment_base_offset + segment_data_size %" PRIu64 diff --git a/runtime/core/array_ref.h b/runtime/core/array_ref.h index a23509e8698..d920761c9fc 100644 --- a/runtime/core/array_ref.h +++ b/runtime/core/array_ref.h @@ -30,6 +30,7 @@ #include #include +#include #include namespace executorch { @@ -161,7 +162,8 @@ class ArrayRef final { /// slice(n, m) - Take M elements of the array starting at element N ArrayRef slice(size_t N, size_t M) const { // cant slice longer then the array - ET_CHECK(N + M <= size()); + size_t end = 0; + ET_CHECK(!c10::add_overflows(N, M, &end) && end <= size()); return ArrayRef(data() + N, M); } diff --git a/runtime/core/hierarchical_allocator.h b/runtime/core/hierarchical_allocator.h index c5dd6243f3b..d41b98f69d5 100644 --- a/runtime/core/hierarchical_allocator.h +++ b/runtime/core/hierarchical_allocator.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #include @@ -58,8 +59,9 @@ class HierarchicalAllocator final { size_t offset_bytes, size_t size_bytes) { // Check for integer overflow in offset_bytes + size_bytes. + size_t end_bytes = 0; ET_CHECK_OR_RETURN_ERROR( - size_bytes <= SIZE_MAX - offset_bytes, + !c10::add_overflows(offset_bytes, size_bytes, &end_bytes), InvalidArgument, "Integer overflow in offset_bytes (%" ET_PRIsize_t ") + size_bytes (%" ET_PRIsize_t ")", @@ -73,7 +75,7 @@ class HierarchicalAllocator final { buffers_.size()); Span buffer = buffers_[memory_id]; ET_CHECK_OR_RETURN_ERROR( - offset_bytes + size_bytes <= buffer.size(), + end_bytes <= buffer.size(), MemoryAllocationFailed, "offset_bytes (%" ET_PRIsize_t ") + size_bytes (%" ET_PRIsize_t ") >= allocator size (%" ET_PRIsize_t