Skip to content

Fix integer overflows in tensor byte-size computations #19055

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lucylq merged 1 commit into from
Apr 23, 2026
+38 −7
Merged

Fix integer overflows in tensor byte-size computations #19055

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion extension/tensor/tensor_ptr.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#include <vector>

#include <c10/macros/Macros.h>
#include <c10/util/safe_numerics.h>
#include <executorch/runtime/core/error.h>
#include <executorch/runtime/core/exec_aten/exec_aten.h>
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
Expand Down Expand Up @@ -117,7 +118,16 @@ inline TensorPtr make_tensor_ptr(
ET_CHECK_MSG(
runtime::canCast(deduced_type, type),
"Cannot cast deduced type to specified type.");
std::vector<uint8_t> casted_data(data.size() * aten::elementSize(type));
size_t casted_bytes = 0;
ET_CHECK_MSG(
!c10::mul_overflows(
data.size(),
static_cast<size_t>(aten::elementSize(type)),
&casted_bytes),
"casted_data size overflow: %zu elements * %zu bytes/element",
data.size(),
static_cast<size_t>(aten::elementSize(type)));
std::vector<uint8_t> casted_data(casted_bytes);

// Create a minimal context for error handling in ET_SWITCH
struct {
Expand Down
16 changes: 13 additions & 3 deletions extension/tensor/tensor_ptr_maker.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@

#include <random>

#include <c10/util/safe_numerics.h>

namespace executorch {
namespace extension {
namespace {
Expand Down Expand Up @@ -111,9 +113,17 @@ TensorPtr empty_strided(
std::vector<executorch::aten::StridesType> strides,
executorch::aten::ScalarType type,
executorch::aten::TensorShapeDynamism dynamism) {
std::vector<uint8_t> data(
executorch::aten::compute_numel(sizes.data(), sizes.size()) *
executorch::aten::elementSize(type));
const auto numel = static_cast<size_t>(
executorch::aten::compute_numel(sizes.data(), sizes.size()));
Comment on lines +116 to +117
Copy link

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

empty_strided() still relies on executorch::aten::compute_numel() (which returns ssize_t and, in executor mode, multiplies sizes using signed arithmetic without overflow checks). If the size product overflows inside compute_numel, that's undefined behavior and may yield a wrapped numel before you even reach the mul_overflows(numel, elem_size, ...) guard.

Consider computing numel locally using size_t + c10::mul_overflows over the sizes vector (and validating non-negative sizes) so the overflow is caught deterministically before any UB occurs, then multiply by elem_size as you do now.

Suggested change
const auto numel = static_cast<size_t>(
executorch::aten::compute_numel(sizes.data(), sizes.size()));
size_t numel = 1;
for (const auto dim : sizes) {
ET_CHECK_MSG(
dim >= 0,
"empty_strided requires non-negative sizes, got %zd",
static_cast<ssize_t>(dim));
size_t next_numel = 0;
ET_CHECK_MSG(
!c10::mul_overflows(
numel, static_cast<size_t>(dim), &next_numel),
"empty_strided size overflow while computing numel: %zu * %zu",
numel,
static_cast<size_t>(dim));
numel = next_numel;
}

Copilot uses AI. Check for mistakes.
Copy link
Copy Markdown
Contributor Author

@lucylq lucylq Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

need to use safe_numel here.

const auto elem_size =
static_cast<size_t>(executorch::aten::elementSize(type));
size_t nbytes = 0;
ET_CHECK_MSG(
!c10::mul_overflows(numel, elem_size, &nbytes),
"empty_strided size overflow: numel %zu * element size %zu",
numel,
elem_size);
std::vector<uint8_t> data(nbytes);
return make_tensor_ptr(
std::move(sizes),
std::move(data),
Expand Down
17 changes: 14 additions & 3 deletions runtime/core/tensor_layout.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
*/

#include <c10/util/irange.h>
#include <c10/util/safe_numerics.h>
#include <executorch/runtime/core/exec_aten/exec_aten.h>
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
#include <executorch/runtime/core/span.h>
Expand All @@ -19,15 +20,25 @@ namespace {
Result<size_t> calculate_nbytes(
const Span<const int32_t>& sizes,
const executorch::aten::ScalarType& scalar_type) {
ssize_t n = 1;
size_t n = 1;
for (const auto i : c10::irange(sizes.size())) {
if (sizes[i] < 0) {
return Error::InvalidArgument;
}
n *= sizes[i];
size_t next = 0;
if (c10::mul_overflows(n, static_cast<size_t>(sizes[i]), &next)) {
return Error::InvalidArgument;
}
n = next;
Comment on lines +28 to +32
Copy link

Copilot AI Apr 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new overflow handling in calculate_nbytes() changes behavior to return Error::InvalidArgument on overflow, but runtime/core/test/tensor_layout_test.cpp currently has no coverage for this path. Add a unit test that uses sizes whose product (or product * elementSize) overflows size_t (e.g., {INT32_MAX, INT32_MAX} with ScalarType::Double on 64-bit) and assert TensorLayout::create() returns InvalidArgument.

Copilot uses AI. Check for mistakes.
}
// Use the full namespace to disambiguate from c10::elementSize.
return n * executorch::runtime::elementSize(scalar_type);
const size_t elem_size =
static_cast<size_t>(executorch::runtime::elementSize(scalar_type));
size_t total = 0;
if (c10::mul_overflows(n, elem_size, &total)) {
return Error::InvalidArgument;
}
return total;
}
} // namespace

Expand Down
Loading