Fix overflows in et (pytorch#19057)

### Summary
Check various overflows

Author: lucylq

backends/apple/metal/runtime/metal_backend.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <c10/util/safe_numerics.h>
 #include <dlfcn.h>
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/error.h>
@@ -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,

backends/cuda/runtime/cuda_backend.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <c10/util/safe_numerics.h>
 #include <cuda_runtime.h>
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/backend/options.h>
@@ -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,

examples/qualcomm/oss_scripts/llama/runner/prompt_processor.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <c10/util/safe_numerics.h>
 #include <executorch/examples/qualcomm/oss_scripts/llama/runner/prompt_processor.h>
 #include <numeric>
 using executorch::aten::TensorImpl;
@@ -248,9 +249,12 @@ Result<uint64_t> PromptProcessor<T>::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<int64_t>(num_prompt_tokens), &end_pos) &&
+            end_pos <= static_cast<int64_t>(metadata_.context_len) -
+                    static_cast<int64_t>(metadata_.ar_len),
         "The sequence length exceeds the maximum limit that the prompt processor can handle.");
   }
 

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 <c10/util/safe_numerics.h>
 #include <executorch/examples/models/llama/runner/runner.h>
 #include <executorch/examples/models/llama/tokenizer/llama_tiktoken.h>
 #include <executorch/examples/qualcomm/oss_scripts/llama/runner/client_mem.h>
@@ -433,8 +434,11 @@ Error Runner<T>::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<int64_t>(num_prompt_tokens), &end_pos) &&
+          end_pos < static_cast<int64_t>(seq_len),
       "sequence length exceeded - please increase the seq_len value");
 
   // Prompt Processor first

extension/flat_tensor/flat_tensor_data_map.cpp

@@ -8,6 +8,8 @@
 
 #include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
 
+#include <c10/util/safe_numerics.h>
+
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_generated.h>
 #include <executorch/extension/flat_tensor/serialize/flat_tensor_header.h>
 
@@ -73,9 +75,13 @@ Result<const flat_tensor_flatbuffer::NamedData*> 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<uint64_t>(segments->Get(segment_index)->offset()),
+              static_cast<uint64_t>(segments->Get(segment_index)->size()),
+              &seg_end) &&
+              seg_end <= static_cast<uint64_t>(segment_end_offset),
           InvalidExternalData,
           "Invalid segment offset %" PRIu64
           " is larger than the segment_base_offset + segment_data_size %" PRIu64

runtime/core/array_ref.h

@@ -30,6 +30,7 @@
 #include <cstdint>
 
 #include <c10/util/irange.h>
+#include <c10/util/safe_numerics.h>
 #include <executorch/runtime/platform/assert.h>
 
 namespace executorch {
@@ -161,7 +162,8 @@ class ArrayRef final {
   /// slice(n, m) - Take M elements of the array starting at element N
   ArrayRef<T> 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<T>(data() + N, M);
   }
 

runtime/core/hierarchical_allocator.h

@@ -9,6 +9,7 @@
 #pragma once
 
 #include <c10/util/irange.h>
+#include <c10/util/safe_numerics.h>
 
 #include <executorch/runtime/core/memory_allocator.h>
 #include <executorch/runtime/core/result.h>
@@ -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<uint8_t> 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