react-native-executorch 0.9.0 → 0.9.1

package/third-party/include/executorch/kernels/portable/NativeFunctions.h

@@ -20,8 +20,12 @@
 namespace torch {
 namespace executor {
 namespace native {
+torch::executor::Tensor & _adaptive_avg_pool2d_out(const torch::executor::Tensor & self, torch::executor::ArrayRef<int64_t> output_size, torch::executor::Tensor & out);
+torch::executor::Tensor & _adaptive_avg_pool2d_out(torch::executor::KernelRuntimeContext & context, const torch::executor::Tensor & self, torch::executor::ArrayRef<int64_t> output_size, torch::executor::Tensor & out);
 torch::executor::Tensor & _cdist_forward_out(const torch::executor::Tensor & x1, const torch::executor::Tensor & x2, double p, torch::executor::optional<int64_t> compute_mode, torch::executor::Tensor & out);
 torch::executor::Tensor & _cdist_forward_out(torch::executor::KernelRuntimeContext & context, const torch::executor::Tensor & x1, const torch::executor::Tensor & x2, double p, torch::executor::optional<int64_t> compute_mode, torch::executor::Tensor & out);
+torch::executor::Tensor & _conj_physical_out(const torch::executor::Tensor & self, torch::executor::Tensor & out);
+torch::executor::Tensor & _conj_physical_out(torch::executor::KernelRuntimeContext & context, const torch::executor::Tensor & self, torch::executor::Tensor & out);
 torch::executor::Tensor & log_softmax_out(const torch::executor::Tensor & self, int64_t dim, bool half_to_float, torch::executor::Tensor & out);
 torch::executor::Tensor & log_softmax_out(torch::executor::KernelRuntimeContext & context, const torch::executor::Tensor & self, int64_t dim, bool half_to_float, torch::executor::Tensor & out);
 ::std::tuple<torch::executor::Tensor &,torch::executor::Tensor &,torch::executor::Tensor &> _native_batch_norm_legit_out(const torch::executor::Tensor & input, const torch::executor::optional<torch::executor::Tensor> & weight, const torch::executor::optional<torch::executor::Tensor> & bias, torch::executor::Tensor & running_mean, torch::executor::Tensor & running_var, bool training, double momentum, double eps, torch::executor::Tensor & out, torch::executor::Tensor & save_mean, torch::executor::Tensor & save_invstd);
@@ -412,6 +416,8 @@ torch::executor::Tensor & upsample_nearest2d_vec_out(const torch::executor::Tens
 torch::executor::Tensor & upsample_nearest2d_vec_out(torch::executor::KernelRuntimeContext & context, const torch::executor::Tensor & input, torch::executor::optional<torch::executor::ArrayRef<int64_t>> output_size, torch::executor::optional<torch::executor::ArrayRef<double>> scale_factors, torch::executor::Tensor & out);
 torch::executor::Tensor & var_correction_out(const torch::executor::Tensor & self, torch::executor::optional<torch::executor::ArrayRef<int64_t>> dim, const torch::executor::optional<torch::executor::Scalar> & correction, bool keepdim, torch::executor::Tensor & out);
 torch::executor::Tensor & var_correction_out(torch::executor::KernelRuntimeContext & context, const torch::executor::Tensor & self, torch::executor::optional<torch::executor::ArrayRef<int64_t>> dim, const torch::executor::optional<torch::executor::Scalar> & correction, bool keepdim, torch::executor::Tensor & out);
+::std::tuple<torch::executor::Tensor &,torch::executor::Tensor &> var_mean_correction_out(const torch::executor::Tensor & self, torch::executor::optional<torch::executor::ArrayRef<int64_t>> dim, const torch::executor::optional<torch::executor::Scalar> & correction, bool keepdim, torch::executor::Tensor & out0, torch::executor::Tensor & out1);
+::std::tuple<torch::executor::Tensor &,torch::executor::Tensor &> var_mean_correction_out(torch::executor::KernelRuntimeContext & context, const torch::executor::Tensor & self, torch::executor::optional<torch::executor::ArrayRef<int64_t>> dim, const torch::executor::optional<torch::executor::Scalar> & correction, bool keepdim, torch::executor::Tensor & out0, torch::executor::Tensor & out1);
 torch::executor::Tensor & var_out(const torch::executor::Tensor & self, torch::executor::optional<torch::executor::ArrayRef<int64_t>> dim, bool unbiased, bool keepdim, torch::executor::Tensor & out);
 torch::executor::Tensor & var_out(torch::executor::KernelRuntimeContext & context, const torch::executor::Tensor & self, torch::executor::optional<torch::executor::ArrayRef<int64_t>> dim, bool unbiased, bool keepdim, torch::executor::Tensor & out);
 torch::executor::Tensor & view_as_real_copy_out(const torch::executor::Tensor & self, torch::executor::Tensor & out);

package/third-party/include/executorch/runtime/backend/backend_options_map.h

@@ -11,6 +11,7 @@
 #include <executorch/runtime/backend/options.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/span.h>
+#include <executorch/runtime/platform/assert.h>
 
 #include <cstring>
 
@@ -168,6 +169,42 @@ public:
    */
   size_t size() const { return size_; }
 
+  /**
+   * Non-owning view of a single (backend_id, options) entry, returned by
+   * entry_at(). The pointer / span are valid until the map is mutated or
+   * destroyed.
+   */
+  struct EntryView {
+    const char *backend_id = nullptr;
+    Span<const BackendOption> options;
+  };
+
+  /**
+   * Returns the (backend_id, options) entry at the given index for
+   * enumeration over the map's contents.
+   *
+   * @param index The entry index. Must be < size(); behavior is undefined
+   *     otherwise. Use this together with size() to walk every entry.
+   * @return EntryView referencing the entry's backend_id and options. The
+   *     view is valid until the next mutation of, or destruction of, this
+   *     map.
+   *
+   * Example:
+   * @code
+   *   for (size_t i = 0; i < map.size(); ++i) {
+   *     const auto entry = map.entry_at(i);
+   *     // use entry.backend_id and entry.options ...
+   *   }
+   * @endcode
+   */
+  EntryView entry_at(size_t index) const {
+    ET_DCHECK_MSG(index < size_, "entry_at index %zu out of bounds (size=%zu)",
+                  index, size_);
+    return EntryView{entries_[index].backend_id,
+                     Span<const BackendOption>(entries_[index].options.data(),
+                                               entries_[index].options.size())};
+  }
+
 private:
   static constexpr size_t kMaxBackends = 8;
   static constexpr size_t kMaxBackendIdLength = 64;

package/third-party/include/executorch/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 {
@@ -146,7 +147,8 @@ public:
   /// 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);
   }
 

package/third-party/include/executorch/runtime/core/error.h

@@ -152,6 +152,7 @@ constexpr const char *to_string(const Error error) {
   case Error::RegistrationAlreadyRegistered:
     return "Error::RegistrationAlreadyRegistered";
   }
+  return "Error::Unknown";
 }
 
 } // namespace runtime

package/third-party/include/executorch/runtime/core/evalue.h

@@ -8,6 +8,7 @@
 
 #pragma once
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/result.h>
 #include <executorch/runtime/core/tag.h>
 #include <executorch/runtime/platform/assert.h>
 
@@ -67,6 +68,29 @@ public:
    */
   executorch::aten::ArrayRef<T> get() const;
 
+  /**
+   * Result-returning counterpart of get(). Validates each wrapped EValue's
+   * tag before materializing; returns Error::InvalidType if any element's
+   * tag does not match T and Error::InvalidState if any element pointer is
+   * null. Use this when materializing lists from untrusted .pte data so that
+   * a malformed program cannot force a process abort inside to<T>() /
+   * ET_CHECK.
+   */
+  Result<executorch::aten::ArrayRef<T>> tryGet() const;
+
+  /**
+   * Destroys the unwrapped elements without re-dereferencing wrapped_vals_.
+   * This is safe to call during EValue destruction because it does not
+   * dereference wrapped_vals_, which may point to EValues mutated by
+   * MoveCall instructions.
+   */
+  void destroy_elements() {
+    for (typename executorch::aten::ArrayRef<T>::size_type i = 0;
+         i < wrapped_vals_.size(); i++) {
+      unwrapped_vals_[i].~T();
+    }
+  }
+
 private:
   static EValue **checkWrappedVals(EValue **wrapped_vals, int size) {
     ET_CHECK_MSG(wrapped_vals != nullptr, "wrapped_vals cannot be null");
@@ -89,6 +113,10 @@ template <>
 executorch::aten::ArrayRef<std::optional<executorch::aten::Tensor>>
 BoxedEvalueList<std::optional<executorch::aten::Tensor>>::get() const;
 
+template <>
+Result<executorch::aten::ArrayRef<std::optional<executorch::aten::Tensor>>>
+BoxedEvalueList<std::optional<executorch::aten::Tensor>>::tryGet() const;
+
 // Aggregate typing system similar to IValue only slimmed down with less
 // functionality, no dependencies on atomic, and fewer supported types to better
 // suit embedded systems (ie no intrusive ptr)
@@ -165,6 +193,13 @@ struct EValue {
     return payload.copyable_union.as_int;
   }
 
+  Result<int64_t> tryToInt() const {
+    if (!isInt()) {
+      return Error::InvalidType;
+    }
+    return payload.copyable_union.as_int;
+  }
+
   /****** Double Type ******/
   /*implicit*/ EValue(double d) : tag(Tag::Double) {
     payload.copyable_union.as_double = d;
@@ -177,6 +212,13 @@ struct EValue {
     return payload.copyable_union.as_double;
   }
 
+  Result<double> tryToDouble() const {
+    if (!isDouble()) {
+      return Error::InvalidType;
+    }
+    return payload.copyable_union.as_double;
+  }
+
   /****** Bool Type ******/
   /*implicit*/ EValue(bool b) : tag(Tag::Bool) {
     payload.copyable_union.as_bool = b;
@@ -189,6 +231,13 @@ struct EValue {
     return payload.copyable_union.as_bool;
   }
 
+  Result<bool> tryToBool() const {
+    if (!isBool()) {
+      return Error::InvalidType;
+    }
+    return payload.copyable_union.as_bool;
+  }
+
   /****** Scalar Type ******/
   /// Construct an EValue using the implicit value of a Scalar.
   /*implicit*/ EValue(executorch::aten::Scalar s) {
@@ -224,6 +273,19 @@ struct EValue {
     }
   }
 
+  Result<executorch::aten::Scalar> tryToScalar() const {
+    if (isDouble()) {
+      return executorch::aten::Scalar(payload.copyable_union.as_double);
+    }
+    if (isInt()) {
+      return executorch::aten::Scalar(payload.copyable_union.as_int);
+    }
+    if (isBool()) {
+      return executorch::aten::Scalar(payload.copyable_union.as_bool);
+    }
+    return Error::InvalidType;
+  }
+
   /****** Tensor Type ******/
   /*implicit*/ EValue(executorch::aten::Tensor t) : tag(Tag::Tensor) {
     // When built in aten mode, at::Tensor has a non trivial constructor
@@ -270,6 +332,16 @@ struct EValue {
     return payload.as_tensor;
   }
 
+  // Returns a copy of the Tensor handle (one intrusive_ptr refcount bump in
+  // ATen mode; free in lean mode). Unlike toTensor()'s const& / & overloads,
+  // tryToTensor() cannot return a reference — Result<T> wraps by value.
+  Result<executorch::aten::Tensor> tryToTensor() const {
+    if (!isTensor()) {
+      return Error::InvalidType;
+    }
+    return payload.as_tensor;
+  }
+
   /****** String Type ******/
   /*implicit*/ EValue(executorch::aten::ArrayRef<char> *s) : tag(Tag::String) {
     ET_CHECK_MSG(s != nullptr, "ArrayRef<char> pointer cannot be null");
@@ -286,6 +358,17 @@ struct EValue {
                             payload.copyable_union.as_string_ptr->size());
   }
 
+  Result<std::string_view> tryToString() const {
+    if (!isString()) {
+      return Error::InvalidType;
+    }
+    if (payload.copyable_union.as_string_ptr == nullptr) {
+      return Error::InvalidState;
+    }
+    return std::string_view(payload.copyable_union.as_string_ptr->data(),
+                            payload.copyable_union.as_string_ptr->size());
+  }
+
   /****** Int List Type ******/
   /*implicit*/ EValue(BoxedEvalueList<int64_t> *i) : tag(Tag::ListInt) {
     ET_CHECK_MSG(i != nullptr,
@@ -302,6 +385,16 @@ struct EValue {
     return (payload.copyable_union.as_int_list_ptr)->get();
   }
 
+  Result<executorch::aten::ArrayRef<int64_t>> tryToIntList() const {
+    if (!isIntList()) {
+      return Error::InvalidType;
+    }
+    if (payload.copyable_union.as_int_list_ptr == nullptr) {
+      return Error::InvalidState;
+    }
+    return (payload.copyable_union.as_int_list_ptr)->tryGet();
+  }
+
   /****** Bool List Type ******/
   /*implicit*/ EValue(executorch::aten::ArrayRef<bool> *b)
       : tag(Tag::ListBool) {
@@ -318,6 +411,16 @@ struct EValue {
     return *(payload.copyable_union.as_bool_list_ptr);
   }
 
+  Result<executorch::aten::ArrayRef<bool>> tryToBoolList() const {
+    if (!isBoolList()) {
+      return Error::InvalidType;
+    }
+    if (payload.copyable_union.as_bool_list_ptr == nullptr) {
+      return Error::InvalidState;
+    }
+    return *(payload.copyable_union.as_bool_list_ptr);
+  }
+
   /****** Double List Type ******/
   /*implicit*/ EValue(executorch::aten::ArrayRef<double> *d)
       : tag(Tag::ListDouble) {
@@ -334,6 +437,16 @@ struct EValue {
     return *(payload.copyable_union.as_double_list_ptr);
   }
 
+  Result<executorch::aten::ArrayRef<double>> tryToDoubleList() const {
+    if (!isDoubleList()) {
+      return Error::InvalidType;
+    }
+    if (payload.copyable_union.as_double_list_ptr == nullptr) {
+      return Error::InvalidState;
+    }
+    return *(payload.copyable_union.as_double_list_ptr);
+  }
+
   /****** Tensor List Type ******/
   /*implicit*/ EValue(BoxedEvalueList<executorch::aten::Tensor> *t)
       : tag(Tag::ListTensor) {
@@ -351,6 +464,17 @@ struct EValue {
     return payload.copyable_union.as_tensor_list_ptr->get();
   }
 
+  Result<executorch::aten::ArrayRef<executorch::aten::Tensor>>
+  tryToTensorList() const {
+    if (!isTensorList()) {
+      return Error::InvalidType;
+    }
+    if (payload.copyable_union.as_tensor_list_ptr == nullptr) {
+      return Error::InvalidState;
+    }
+    return payload.copyable_union.as_tensor_list_ptr->tryGet();
+  }
+
   /****** List Optional Tensor Type ******/
   /*implicit*/ EValue(
       BoxedEvalueList<std::optional<executorch::aten::Tensor>> *t)
@@ -371,6 +495,17 @@ struct EValue {
     return payload.copyable_union.as_list_optional_tensor_ptr->get();
   }
 
+  Result<executorch::aten::ArrayRef<std::optional<executorch::aten::Tensor>>>
+  tryToListOptionalTensor() const {
+    if (!isListOptionalTensor()) {
+      return Error::InvalidType;
+    }
+    if (payload.copyable_union.as_list_optional_tensor_ptr == nullptr) {
+      return Error::InvalidState;
+    }
+    return payload.copyable_union.as_list_optional_tensor_ptr->tryGet();
+  }
+
   /****** ScalarType Type ******/
   executorch::aten::ScalarType toScalarType() const {
     ET_CHECK_MSG(isInt(), "EValue is not a ScalarType.");
@@ -378,6 +513,14 @@ struct EValue {
         payload.copyable_union.as_int);
   }
 
+  Result<executorch::aten::ScalarType> tryToScalarType() const {
+    if (!isInt()) {
+      return Error::InvalidType;
+    }
+    return static_cast<executorch::aten::ScalarType>(
+        payload.copyable_union.as_int);
+  }
+
   /****** MemoryFormat Type ******/
   executorch::aten::MemoryFormat toMemoryFormat() const {
     ET_CHECK_MSG(isInt(), "EValue is not a MemoryFormat.");
@@ -385,12 +528,27 @@ struct EValue {
         payload.copyable_union.as_int);
   }
 
+  Result<executorch::aten::MemoryFormat> tryToMemoryFormat() const {
+    if (!isInt()) {
+      return Error::InvalidType;
+    }
+    return static_cast<executorch::aten::MemoryFormat>(
+        payload.copyable_union.as_int);
+  }
+
   /****** Layout Type ******/
   executorch::aten::Layout toLayout() const {
     ET_CHECK_MSG(isInt(), "EValue is not a Layout.");
     return static_cast<executorch::aten::Layout>(payload.copyable_union.as_int);
   }
 
+  Result<executorch::aten::Layout> tryToLayout() const {
+    if (!isInt()) {
+      return Error::InvalidType;
+    }
+    return static_cast<executorch::aten::Layout>(payload.copyable_union.as_int);
+  }
+
   /****** Device Type ******/
   executorch::aten::Device toDevice() const {
     ET_CHECK_MSG(isInt(), "EValue is not a Device.");
@@ -399,12 +557,29 @@ struct EValue {
                                     -1);
   }
 
+  Result<executorch::aten::Device> tryToDevice() const {
+    if (!isInt()) {
+      return Error::InvalidType;
+    }
+    return executorch::aten::Device(static_cast<executorch::aten::DeviceType>(
+                                        payload.copyable_union.as_int),
+                                    -1);
+  }
+
   template <typename T> T to() &&;
   template <typename T>
   typename internal::evalue_to_const_ref_overload_return<T>::type to() const &;
   template <typename T>
   typename internal::evalue_to_ref_overload_return<T>::type to() &;
 
+  /**
+   * Result-returning equivalent of `to<T>()`. Tag mismatch returns
+   * `Error::InvalidType`; a null list/string payload returns
+   * `Error::InvalidState`. Specializations are defined below via
+   * `EVALUE_DEFINE_TRY_TO`.
+   */
+  template <typename T> Result<T> tryTo() const;
+
   /**
    * Converts the EValue to an optional object that can represent both T and
    * an uninitialized state.
@@ -416,6 +591,22 @@ struct EValue {
     return this->to<T>();
   }
 
+  /**
+   * Result-returning equivalent of `toOptional<T>()`. None maps to an empty
+   * optional; any other tag that doesn't match T propagates `tryTo<T>()`'s
+   * error (`Error::InvalidType`).
+   */
+  template <typename T> inline Result<std::optional<T>> tryToOptional() const {
+    if (this->isNone()) {
+      return std::optional<T>(std::nullopt);
+    }
+    auto r = this->tryTo<T>();
+    if (!r.ok()) {
+      return r.error();
+    }
+    return std::optional<T>(std::move(r.get()));
+  }
+
 private:
   // Pre cond: the payload value has had its destructor called
   void clearToNone() noexcept {
@@ -446,17 +637,10 @@ private:
       payload.as_tensor.~Tensor();
     } else if (isTensorList() &&
                payload.copyable_union.as_tensor_list_ptr != nullptr) {
-      // for (auto& tensor : toTensorList()) {
-      for (auto &tensor : payload.copyable_union.as_tensor_list_ptr->get()) {
-        tensor.~Tensor();
-      }
+      payload.copyable_union.as_tensor_list_ptr->destroy_elements();
     } else if (isListOptionalTensor() &&
                payload.copyable_union.as_list_optional_tensor_ptr != nullptr) {
-      // for (auto& optional_tensor : toListOptionalTensor()) {
-      for (auto &optional_tensor :
-           payload.copyable_union.as_list_optional_tensor_ptr->get()) {
-        optional_tensor.~optional();
-      }
+      payload.copyable_union.as_list_optional_tensor_ptr->destroy_elements();
     }
   }
 
@@ -532,6 +716,53 @@ EVALUE_DEFINE_TO(
     toListOptionalTensor)
 #undef EVALUE_DEFINE_TO
 
+#define EVALUE_DEFINE_TRY_TO(T, method_name)                                   \
+  template <> inline Result<T> EValue::tryTo<T>() const {                      \
+    return this->method_name();                                                \
+  }
+
+EVALUE_DEFINE_TRY_TO(executorch::aten::Scalar, tryToScalar)
+EVALUE_DEFINE_TRY_TO(int64_t, tryToInt)
+EVALUE_DEFINE_TRY_TO(bool, tryToBool)
+EVALUE_DEFINE_TRY_TO(double, tryToDouble)
+EVALUE_DEFINE_TRY_TO(std::string_view, tryToString)
+EVALUE_DEFINE_TRY_TO(executorch::aten::ScalarType, tryToScalarType)
+EVALUE_DEFINE_TRY_TO(executorch::aten::MemoryFormat, tryToMemoryFormat)
+EVALUE_DEFINE_TRY_TO(executorch::aten::Layout, tryToLayout)
+EVALUE_DEFINE_TRY_TO(executorch::aten::Device, tryToDevice)
+// Tensor and Optional Tensor
+EVALUE_DEFINE_TRY_TO(executorch::aten::Tensor, tryToTensor)
+EVALUE_DEFINE_TRY_TO(std::optional<executorch::aten::Tensor>,
+                     tryToOptional<executorch::aten::Tensor>)
+
+// IntList and Optional IntList
+EVALUE_DEFINE_TRY_TO(executorch::aten::ArrayRef<int64_t>, tryToIntList)
+EVALUE_DEFINE_TRY_TO(std::optional<executorch::aten::ArrayRef<int64_t>>,
+                     tryToOptional<executorch::aten::ArrayRef<int64_t>>)
+
+// DoubleList and Optional DoubleList
+EVALUE_DEFINE_TRY_TO(executorch::aten::ArrayRef<double>, tryToDoubleList)
+EVALUE_DEFINE_TRY_TO(std::optional<executorch::aten::ArrayRef<double>>,
+                     tryToOptional<executorch::aten::ArrayRef<double>>)
+
+// BoolList and Optional BoolList
+EVALUE_DEFINE_TRY_TO(executorch::aten::ArrayRef<bool>, tryToBoolList)
+EVALUE_DEFINE_TRY_TO(std::optional<executorch::aten::ArrayRef<bool>>,
+                     tryToOptional<executorch::aten::ArrayRef<bool>>)
+
+// TensorList and Optional TensorList
+EVALUE_DEFINE_TRY_TO(executorch::aten::ArrayRef<executorch::aten::Tensor>,
+                     tryToTensorList)
+EVALUE_DEFINE_TRY_TO(
+    std::optional<executorch::aten::ArrayRef<executorch::aten::Tensor>>,
+    tryToOptional<executorch::aten::ArrayRef<executorch::aten::Tensor>>)
+
+// List of Optional Tensor
+EVALUE_DEFINE_TRY_TO(
+    executorch::aten::ArrayRef<std::optional<executorch::aten::Tensor>>,
+    tryToListOptionalTensor)
+#undef EVALUE_DEFINE_TRY_TO
+
 template <typename T>
 executorch::aten::ArrayRef<T> BoxedEvalueList<T>::get() const {
   for (typename executorch::aten::ArrayRef<T>::size_type i = 0;
@@ -542,6 +773,22 @@ executorch::aten::ArrayRef<T> BoxedEvalueList<T>::get() const {
   return executorch::aten::ArrayRef<T>{unwrapped_vals_, wrapped_vals_.size()};
 }
 
+template <typename T>
+Result<executorch::aten::ArrayRef<T>> BoxedEvalueList<T>::tryGet() const {
+  for (typename executorch::aten::ArrayRef<T>::size_type i = 0;
+       i < wrapped_vals_.size(); i++) {
+    if (wrapped_vals_[i] == nullptr) {
+      return Error::InvalidState;
+    }
+    auto r = wrapped_vals_[i]->template tryTo<T>();
+    if (!r.ok()) {
+      return r.error();
+    }
+    unwrapped_vals_[i] = std::move(r.get());
+  }
+  return executorch::aten::ArrayRef<T>{unwrapped_vals_, wrapped_vals_.size()};
+}
+
 } // namespace runtime
 } // namespace executorch
 

package/third-party/include/executorch/runtime/core/exec_aten/exec_aten.h

@@ -8,7 +8,10 @@
 
 #pragma once
 
+#include <executorch/runtime/core/error.h>                 // @manual
+#include <executorch/runtime/core/result.h>                // @manual
 #include <executorch/runtime/core/tensor_shape_dynamism.h> // @manual
+#include <executorch/runtime/platform/assert.h>            // @manual
 #include <executorch/runtime/platform/compiler.h>
 #ifdef USE_ATEN_LIB
 #include <ATen/Tensor.h> // @manual
@@ -28,6 +31,7 @@
 #include <c10/util/quint2x4.h>      // @manual
 #include <c10/util/quint4x2.h>      // @manual
 #include <c10/util/quint8.h>        // @manual
+#include <c10/util/safe_numerics.h> // @manual
 #include <c10/util/string_view.h>   // @manual
 #include <torch/torch.h>
 #else                                                            // use executor
@@ -107,6 +111,25 @@ inline ssize_t compute_numel(const SizesType *sizes, ssize_t dim) {
       c10::multiply_integers(c10::ArrayRef<SizesType>(sizes, dim)));
 }
 
+inline ::executorch::runtime::Result<ssize_t> safe_numel(const SizesType *sizes,
+                                                         ssize_t dim) {
+  ET_CHECK_OR_RETURN_ERROR(dim == 0 || sizes != nullptr, InvalidArgument,
+                           "Sizes must be provided for non-scalar tensors");
+  ssize_t numel = 1;
+  for (ssize_t i = 0; i < dim; i++) {
+    ET_CHECK_OR_RETURN_ERROR(
+        sizes[i] >= 0, InvalidArgument,
+        "Size must be non-negative, got %zd at dimension %zd",
+        static_cast<ssize_t>(sizes[i]), i);
+    ssize_t next_numel;
+    ET_CHECK_OR_RETURN_ERROR(
+        !c10::mul_overflows(numel, static_cast<ssize_t>(sizes[i]), &next_numel),
+        InvalidArgument, "Overflow computing numel at dimension %zd", i);
+    numel = next_numel;
+  }
+  return numel;
+}
+
 #undef ET_PRI_TENSOR_SIZE
 #define ET_PRI_TENSOR_SIZE PRId64
 
@@ -153,6 +176,7 @@ using OptionalArrayRef =
 using OptionalIntArrayRef = OptionalArrayRef<int64_t>;
 
 using torch::executor::compute_numel;
+using torch::executor::safe_numel;
 
 #endif // Use ExecuTorch types
 

package/third-party/include/executorch/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>
@@ -56,17 +57,19 @@ public:
                                                  size_t offset_bytes,
                                                  size_t size_bytes) {
     // Check for integer overflow in offset_bytes + size_bytes.
-    ET_CHECK_OR_RETURN_ERROR(size_bytes <= SIZE_MAX - offset_bytes,
-                             InvalidArgument,
-                             "Integer overflow in offset_bytes (%" ET_PRIsize_t
-                             ") + size_bytes (%" ET_PRIsize_t ")",
-                             offset_bytes, size_bytes);
+    size_t end_bytes = 0;
+    ET_CHECK_OR_RETURN_ERROR(
+        !c10::add_overflows(offset_bytes, size_bytes, &end_bytes),
+        InvalidArgument,
+        "Integer overflow in offset_bytes (%" ET_PRIsize_t
+        ") + size_bytes (%" ET_PRIsize_t ")",
+        offset_bytes, size_bytes);
     ET_CHECK_OR_RETURN_ERROR(memory_id < buffers_.size(), InvalidArgument,
                              "id %" PRIu32 " >= %" ET_PRIsize_t, memory_id,
                              buffers_.size());
     Span<uint8_t> buffer = buffers_[memory_id];
     ET_CHECK_OR_RETURN_ERROR(
-        offset_bytes + size_bytes <= buffer.size(), MemoryAllocationFailed,
+        end_bytes <= buffer.size(), MemoryAllocationFailed,
         "offset_bytes (%" ET_PRIsize_t ") + size_bytes (%" ET_PRIsize_t
         ") >= allocator size (%" ET_PRIsize_t ") "
         "for memory_id %" PRIu32,

package/third-party/include/executorch/runtime/core/portable_type/device.h

@@ -26,7 +26,6 @@ enum class DeviceType : int8_t {
 constexpr size_t kNumDeviceTypes = 2;
 
 /// An index representing a specific device; e.g. GPU 0 vs GPU 1.
-/// -1 means the default/unspecified device for that type.
 using DeviceIndex = int8_t;
 
 /**
@@ -41,7 +40,7 @@ struct Device final {
 
   /// Constructs a new `Device` from a `DeviceType` and an optional device
   /// index.
-  /* implicit */ Device(DeviceType type, DeviceIndex index = -1)
+  /* implicit */ Device(DeviceType type, DeviceIndex index = 0)
       : type_(type), index_(index) {}
 
   /// Returns the type of device the tensor data resides on.
@@ -50,7 +49,7 @@ struct Device final {
   /// Returns true if the device is of CPU type.
   bool is_cpu() const noexcept { return type_ == DeviceType::CPU; }
 
-  /// Returns the device index, or -1 if default/unspecified.
+  /// Returns the device index.
   DeviceIndex index() const noexcept { return index_; }
 
   bool operator==(const Device &other) const noexcept {
@@ -63,7 +62,7 @@ struct Device final {
 
 private:
   DeviceType type_;
-  DeviceIndex index_ = -1;
+  DeviceIndex index_ = 0;
 };
 
 } // namespace etensor