From 66f8b3ce93d89c173f922b0d584e50efafb1269a Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Tue, 20 Jan 2026 08:18:29 -0800 Subject: [PATCH 1/5] overflow --- runtime/core/portable_type/tensor_impl.cpp | 7 ++++++- .../portable_type/test/tensor_impl_test.cpp | 17 +++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/runtime/core/portable_type/tensor_impl.cpp b/runtime/core/portable_type/tensor_impl.cpp index ede5a3d4101..74f391e8c33 100644 --- a/runtime/core/portable_type/tensor_impl.cpp +++ b/runtime/core/portable_type/tensor_impl.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -66,7 +67,11 @@ TensorImpl::TensorImpl( } size_t TensorImpl::nbytes() const { - return numel_ * elementSize(type_); + size_t result; + bool overflow = + c10::mul_overflows(static_cast(numel_), elementSize(type_), &result); + ET_CHECK_MSG(!overflow, "nbytes overflowed"); + return result; } // Return the size of one element of the tensor diff --git a/runtime/core/portable_type/test/tensor_impl_test.cpp b/runtime/core/portable_type/test/tensor_impl_test.cpp index 0b8ae05f4da..3e3b8f74a76 100644 --- a/runtime/core/portable_type/test/tensor_impl_test.cpp +++ b/runtime/core/portable_type/test/tensor_impl_test.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -449,3 +450,19 @@ TEST_F(TensorImplTest, TestResizingTensorToZeroAndBack) { EXPECT_GT(t.numel(), 0); EXPECT_EQ(t.data(), data); } + +TEST_F(TensorImplTest, TestNbytesOverflow) { + // Create a tensor with sizes that would cause overflow when computing nbytes. + // SIZE_MAX / sizeof(float) will give us a numel that, when multiplied by + // sizeof(float), would overflow size_t. + constexpr size_t max_safe_numel = SIZE_MAX / sizeof(float); + // Use sizes that when multiplied together exceed max_safe_numel + // For a 2D tensor, we need size[0] * size[1] > max_safe_numel + SizesType sizes[2] = { + static_cast(max_safe_numel / 2 + 1), + static_cast(3)}; + TensorImpl t(ScalarType::Float, 2, sizes); + + // Calling nbytes() should trigger an overflow check and die + ET_EXPECT_DEATH(t.nbytes(), ""); +} From d2b3148b403bc32fc0741df9a78a5fcec50db3a8 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Tue, 20 Jan 2026 09:13:59 -0800 Subject: [PATCH 2/5] test and numel --- runtime/core/portable_type/tensor_impl.cpp | 4 +++- .../core/portable_type/test/tensor_impl_test.cpp | 13 ++----------- 2 files changed, 5 insertions(+), 12 deletions(-) diff --git a/runtime/core/portable_type/tensor_impl.cpp b/runtime/core/portable_type/tensor_impl.cpp index 74f391e8c33..595e1e97373 100644 --- a/runtime/core/portable_type/tensor_impl.cpp +++ b/runtime/core/portable_type/tensor_impl.cpp @@ -39,7 +39,9 @@ ssize_t compute_numel(const TensorImpl::SizesType* sizes, ssize_t dim) { "Size must be non-negative, got %zd at dimension %zd", static_cast(sizes[i]), i); - numel *= sizes[i]; + bool overflow = + c10::mul_overflows(numel, static_cast(sizes[i]), &numel); + ET_CHECK_MSG(!overflow, "numel overflowed"); } return numel; } diff --git a/runtime/core/portable_type/test/tensor_impl_test.cpp b/runtime/core/portable_type/test/tensor_impl_test.cpp index 3e3b8f74a76..84a1055358c 100644 --- a/runtime/core/portable_type/test/tensor_impl_test.cpp +++ b/runtime/core/portable_type/test/tensor_impl_test.cpp @@ -452,17 +452,8 @@ TEST_F(TensorImplTest, TestResizingTensorToZeroAndBack) { } TEST_F(TensorImplTest, TestNbytesOverflow) { - // Create a tensor with sizes that would cause overflow when computing nbytes. - // SIZE_MAX / sizeof(float) will give us a numel that, when multiplied by - // sizeof(float), would overflow size_t. constexpr size_t max_safe_numel = SIZE_MAX / sizeof(float); - // Use sizes that when multiplied together exceed max_safe_numel - // For a 2D tensor, we need size[0] * size[1] > max_safe_numel - SizesType sizes[2] = { - static_cast(max_safe_numel / 2 + 1), - static_cast(3)}; - TensorImpl t(ScalarType::Float, 2, sizes); - - // Calling nbytes() should trigger an overflow check and die + SizesType sizes[1] = {static_cast(max_safe_numel + 1)}; + TensorImpl t(ScalarType::Float, 1, sizes); ET_EXPECT_DEATH(t.nbytes(), ""); } From 421c0afb7be14a4fac153a1904f765989bb1ae8a Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Tue, 20 Jan 2026 09:15:13 -0800 Subject: [PATCH 3/5] lint --- runtime/core/portable_type/tensor_impl.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/runtime/core/portable_type/tensor_impl.cpp b/runtime/core/portable_type/tensor_impl.cpp index 595e1e97373..64ecb73a726 100644 --- a/runtime/core/portable_type/tensor_impl.cpp +++ b/runtime/core/portable_type/tensor_impl.cpp @@ -70,8 +70,8 @@ TensorImpl::TensorImpl( size_t TensorImpl::nbytes() const { size_t result; - bool overflow = - c10::mul_overflows(static_cast(numel_), elementSize(type_), &result); + bool overflow = c10::mul_overflows( + static_cast(numel_), elementSize(type_), &result); ET_CHECK_MSG(!overflow, "nbytes overflowed"); return result; } From f264fd81df927e838afafb6ca0b55964181e9a0e Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Tue, 20 Jan 2026 09:46:56 -0800 Subject: [PATCH 4/5] forgot about int32 need multiple --- .../portable_type/test/tensor_impl_test.cpp | 32 ++++++------------- 1 file changed, 10 insertions(+), 22 deletions(-) diff --git a/runtime/core/portable_type/test/tensor_impl_test.cpp b/runtime/core/portable_type/test/tensor_impl_test.cpp index 84a1055358c..6940ae6971b 100644 --- a/runtime/core/portable_type/test/tensor_impl_test.cpp +++ b/runtime/core/portable_type/test/tensor_impl_test.cpp @@ -434,26 +434,14 @@ TEST_F(TensorImplTest, TestResizingTensorToZeroAndBack) { nullptr, nullptr, TensorShapeDynamism::DYNAMIC_BOUND); - - float data[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; - t.set_data(data); - EXPECT_GT(t.numel(), 0); - EXPECT_EQ(t.data(), data); - - SizesType zero_sizes[2] = {0, 0}; - t.set_sizes_contiguous({zero_sizes, 2}); - EXPECT_EQ(t.numel(), 0); - EXPECT_EQ(t.data(), data); - - SizesType new_sizes[2] = {3, 2}; - t.set_sizes_contiguous({new_sizes, 2}); - EXPECT_GT(t.numel(), 0); - EXPECT_EQ(t.data(), data); -} - -TEST_F(TensorImplTest, TestNbytesOverflow) { - constexpr size_t max_safe_numel = SIZE_MAX / sizeof(float); - SizesType sizes[1] = {static_cast(max_safe_numel + 1)}; - TensorImpl t(ScalarType::Float, 1, sizes); - ET_EXPECT_DEATH(t.nbytes(), ""); +TEST_F(TensorImplTest, TestNumelOverflow) { + // Create a tensor with sizes that would cause overflow when computing numel. + // SizesType is int32_t, so each dimension can be at most ~2^31. + // Using 3 dimensions: 2^21 * 2^21 * 2^21 = 2^63 which overflows ssize_t. + SizesType sizes[3] = { + static_cast(1 << 21), + static_cast(1 << 21), + static_cast(1 << 21)}; + // This should die during construction because numel overflows + ET_EXPECT_DEATH(TensorImpl t(ScalarType::Float, 3, sizes), ""); } From 5977020ee789c14e509c28c6a5efa5f1e484ebb6 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Tue, 20 Jan 2026 10:30:02 -0800 Subject: [PATCH 5/5] fix --- .../portable_type/test/tensor_impl_test.cpp | 23 +++++++++++++++---- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/runtime/core/portable_type/test/tensor_impl_test.cpp b/runtime/core/portable_type/test/tensor_impl_test.cpp index 6940ae6971b..08b26c64e09 100644 --- a/runtime/core/portable_type/test/tensor_impl_test.cpp +++ b/runtime/core/portable_type/test/tensor_impl_test.cpp @@ -434,14 +434,27 @@ TEST_F(TensorImplTest, TestResizingTensorToZeroAndBack) { nullptr, nullptr, TensorShapeDynamism::DYNAMIC_BOUND); -TEST_F(TensorImplTest, TestNumelOverflow) { - // Create a tensor with sizes that would cause overflow when computing numel. - // SizesType is int32_t, so each dimension can be at most ~2^31. - // Using 3 dimensions: 2^21 * 2^21 * 2^21 = 2^63 which overflows ssize_t. + + float data[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + t.set_data(data); + EXPECT_GT(t.numel(), 0); + EXPECT_EQ(t.data(), data); + + SizesType zero_sizes[2] = {0, 0}; + t.set_sizes_contiguous({zero_sizes, 2}); + EXPECT_EQ(t.numel(), 0); + EXPECT_EQ(t.data(), data); + + SizesType new_sizes[2] = {3, 2}; + t.set_sizes_contiguous({new_sizes, 2}); + EXPECT_GT(t.numel(), 0); + EXPECT_EQ(t.data(), data); +} + +TEST_F(TensorImplTest, TestNbytesOverflow) { SizesType sizes[3] = { static_cast(1 << 21), static_cast(1 << 21), static_cast(1 << 21)}; - // This should die during construction because numel overflows ET_EXPECT_DEATH(TensorImpl t(ScalarType::Float, 3, sizes), ""); }