diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index 80ad800b4fc..5b3dd671701 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -371,7 +371,7 @@ jobs: size=${arr[4]} # threshold=48120 on devserver with gcc11.4 # todo(lfq): update once binary size is below 50kb. - threshold="51408" + threshold="55504" if [[ "$size" -le "$threshold" ]]; then echo "Success $size <= $threshold" else @@ -406,7 +406,7 @@ jobs: output=$(ls -la cmake-out/test/size_test) arr=($output) size=${arr[4]} - threshold="47560" + threshold="51656" if [[ "$size" -le "$threshold" ]]; then echo "Success $size <= $threshold" else diff --git a/runtime/executor/method_meta.cpp b/runtime/executor/method_meta.cpp index e810d195370..fd980c64272 100644 --- a/runtime/executor/method_meta.cpp +++ b/runtime/executor/method_meta.cpp @@ -55,12 +55,24 @@ Result get_tag( size_t calculate_nbytes( Span sizes, executorch::aten::ScalarType scalar_type) { - ssize_t n = 1; + size_t n = 1; + size_t prev_n = 1; for (size_t i = 0; i < sizes.size(); i++) { + prev_n = n; n *= sizes[i]; + // Check for overflow + ET_CHECK(sizes[i] == 0 || n / sizes[i] == prev_n); } - // Use the full namespace to disambiguate from c10::elementSize. - return n * executorch::runtime::elementSize(scalar_type); + + size_t elem_size = executorch::runtime::elementSize(scalar_type); + + prev_n = n; + n = n * elem_size; + + // Check for overflow + ET_CHECK(elem_size == 0 || n / elem_size == prev_n); + + return n; } } // namespace diff --git a/runtime/executor/method_meta.h b/runtime/executor/method_meta.h index ec910f9f6e4..3e71ebf1420 100644 --- a/runtime/executor/method_meta.h +++ b/runtime/executor/method_meta.h @@ -21,6 +21,10 @@ struct ExecutionPlan; namespace executorch { namespace ET_RUNTIME_NAMESPACE { +namespace testing { +// Provides test access to private Program methods. +class TensorInfoTestFriend; +} // namespace testing /** * Metadata about a specific tensor of an ExecuTorch Program. @@ -71,6 +75,7 @@ class TensorInfo final { private: // Let MethodMeta create TensorInfo. friend class MethodMeta; + friend class testing::TensorInfoTestFriend; TensorInfo( Span sizes, diff --git a/runtime/executor/test/method_meta_test.cpp b/runtime/executor/test/method_meta_test.cpp index 211800d5dff..f5a07e352aa 100644 --- a/runtime/executor/test/method_meta_test.cpp +++ b/runtime/executor/test/method_meta_test.cpp @@ -9,11 +9,13 @@ #include #include -#include +#include +#include #include #include #include +#include #include using namespace ::testing; @@ -21,9 +23,34 @@ using executorch::runtime::Error; using executorch::runtime::MethodMeta; using executorch::runtime::Program; using executorch::runtime::Result; +using executorch::runtime::Span; using executorch::runtime::TensorInfo; using torch::executor::util::FileDataLoader; +namespace executorch { +namespace runtime { +namespace testing { +// Provides access to private TensorInfo methods. +class TensorInfoTestFriend final { + public: + ET_NODISCARD static TensorInfo get( + Span sizes, + Span dim_order, + executorch::aten::ScalarType scalar_type, + const bool is_memory_planned, + executorch::aten::string_view name) { + return TensorInfo( + Span(sizes.data(), sizes.size()), + Span(dim_order.data(), dim_order.size()), + scalar_type, + is_memory_planned, + name); + } +}; +} // namespace testing +} // namespace runtime +} // namespace executorch + class MethodMetaTest : public ::testing::Test { protected: void load_program(const char* path, const char* module_name) { @@ -163,3 +190,26 @@ TEST_F(MethodMetaTest, MethodMetaAttribute) { auto bad_access = method_meta->attribute_tensor_meta(1); ASSERT_EQ(bad_access.error(), Error::InvalidArgument); } + +TEST_F(MethodMetaTest, TensorInfoSizeOverflow) { + // Create sizes that will cause overflow when multiplied + std::vector overflow_sizes = { + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max(), + }; + + // Create a minimal dim_order + std::vector dim_order = {0, 1, 2, 3}; + + // Create a TensorInfo with the overflow sizes and expect it to fail. + ET_EXPECT_DEATH( + executorch::runtime::testing::TensorInfoTestFriend::get( + Span(overflow_sizes.data(), overflow_sizes.size()), + Span(dim_order.data(), dim_order.size()), + executorch::aten::ScalarType::Float, + false, // is_memory_planned + executorch::aten::string_view{nullptr, 0}), + ""); +}