Skip to content

Add safe_numel() #19074

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
meta-codesync merged 13 commits into from
Apr 25, 2026
+75 −1
Merged

Add safe_numel() #19074

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
25e8f81
Add safe_numel() and compute_numel_overflow()
Apr 23, 2026
e9450de
Update base for Update on "Add safe_numel() and compute_numel_overflo…
Apr 24, 2026
86633f6
Update on "Add safe_numel() and compute_numel_overflow()"
Apr 24, 2026
b3db241
Update base for Update on "Add safe_numel() and compute_numel_overflo…
Apr 24, 2026
f09383c
Update on "Add safe_numel() and compute_numel_overflow()"
Apr 24, 2026
e2b7cdb
Update base for Update on "Add safe_numel()"
Apr 24, 2026
1302631
Update on "Add safe_numel()"
Apr 24, 2026
0fb439e
Update base for Update on "Add safe_numel()"
Apr 24, 2026
d17950c
Update on "Add safe_numel()"
Apr 24, 2026
34b3e45
Update base for Update on "Add safe_numel()"
Apr 24, 2026
1303417
Update on "Add safe_numel()"
Apr 24, 2026
788683f
Update base for Update on "Add safe_numel()"
Apr 24, 2026
ae06ce8
Update on "Add safe_numel()"
Apr 24, 2026
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
31 changes: 31 additions & 0 deletions runtime/core/exec_aten/exec_aten.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -110,6 +114,32 @@ 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++) {
Comment on lines +120 to +125
Copy link

Copilot AI Apr 24, 2026

Choose a reason for hiding this comment

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

exec_aten::safe_numel() doesn’t reject negative dim. For malformed metadata, a negative dim will skip the loop and return 1, which is indistinguishable from a scalar. Consider explicitly returning InvalidArgument when dim < 0 (in addition to the existing sizes!=nullptr check for dim>0).

Copilot uses AI. Check for mistakes.
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,
Comment on lines +126 to +135
Copy link

Copilot AI Apr 24, 2026

Choose a reason for hiding this comment

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

In ATen mode, SizesType is int64_t but safe_numel() casts each size to ssize_t before checking overflow. On platforms where ssize_t is narrower than int64_t (or if a size is larger than SSIZE_MAX), this cast is implementation-defined and can truncate/flip sign, potentially letting invalid shapes slip through or producing incorrect results. Add an explicit range check that sizes[i] <= SSIZE_MAX (and >= 0) before converting, or do the multiplication in a wider unsigned type and validate the final result fits in ssize_t.

Copilot uses AI. Check for mistakes.
"Overflow computing numel at dimension %zd",
i);
numel = next_numel;
}
return numel;
}
Comment thread
lucylq marked this conversation as resolved.

#undef ET_PRI_TENSOR_SIZE
#define ET_PRI_TENSOR_SIZE PRId64

Expand Down Expand Up @@ -158,6 +188,7 @@ using OptionalArrayRef =
using OptionalIntArrayRef = OptionalArrayRef<int64_t>;

using torch::executor::compute_numel;
using torch::executor::safe_numel;

#endif // Use ExecuTorch types

Expand Down
5 changes: 4 additions & 1 deletion runtime/core/exec_aten/targets.bzl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@ def define_common_targets():
exported_headers = ["exec_aten.h"],
exported_preprocessor_flags = ["-DUSE_ATEN_LIB"] if aten_mode else [],
visibility = ["PUBLIC"],
exported_deps = ["//executorch/runtime/core:tensor_shape_dynamism"] + ([] if aten_mode else ["//executorch/runtime/core/portable_type:portable_type"]),
exported_deps = [
"//executorch/runtime/core:core",
"//executorch/runtime/core:tensor_shape_dynamism",
] + ([] if aten_mode else ["//executorch/runtime/core/portable_type:portable_type"]),
exported_external_deps = ["libtorch"] if aten_mode else [],
)
27 changes: 27 additions & 0 deletions runtime/core/portable_type/tensor_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
#include <cstdint>

#include <c10/util/irange.h>
#include <c10/util/safe_numerics.h>

#include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
Expand Down Expand Up @@ -43,6 +44,32 @@ ssize_t compute_numel(const TensorImpl::SizesType* sizes, ssize_t dim) {
return numel;
}

::executorch::runtime::Result<ssize_t> safe_numel(
const TensorImpl::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 (const auto i : c10::irange(dim)) {
Comment on lines +50 to +55
Copy link

Copilot AI Apr 24, 2026

Choose a reason for hiding this comment

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

safe_numel() doesn’t reject negative dim. With c10::irange(dim), a negative dim produces an empty range, so the function returns 1 instead of an error. Since this helper is intended for validating potentially malformed/untrusted shape metadata, it should return InvalidArgument when dim < 0 (before the loop).

Copilot uses AI. Check for mistakes.
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;
Comment on lines +62 to +68
Copy link

Copilot AI Apr 24, 2026

Choose a reason for hiding this comment

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

safe_numel() uses c10::mul_overflows() with signed ssize_t. In our portable_type c10 implementation, the MSVC fallback path computes volatile T tmp = a * b for signed T, which can itself trigger signed-overflow UB when the product doesn’t fit. This undermines the overflow checking and can be exploitable for attacker-controlled shapes on MSVC builds. Prefer an overflow check that avoids signed multiplication (e.g., since sizes are validated non-negative, use a size_t/uint64_t accumulator with if (a > SSIZE_MAX / b) bounds checking, or use an unsigned safe-multiply helper and then range-check back to ssize_t).

Copilot uses AI. Check for mistakes.
}
return numel;
}
Comment on lines +47 to +71
Copy link

Copilot AI Apr 24, 2026

Choose a reason for hiding this comment

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

There are existing portable_type tensor_impl unit tests, but safe_numel() introduces new behavior (overflow detection, negative sizes, null sizes for dim>0). Please add focused tests covering: scalar (dim==0, sizes==nullptr), negative dim, negative size, and an overflow case near SSIZE_MAX to ensure the new checks behave as expected.

Copilot uses AI. Check for mistakes.

TensorImpl::TensorImpl(
ScalarType type,
ssize_t dim,
Expand Down
13 changes: 13 additions & 0 deletions runtime/core/portable_type/tensor_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,9 @@
#include <executorch/runtime/core/error.h>
#include <executorch/runtime/core/portable_type/device.h>
#include <executorch/runtime/core/portable_type/scalar_type.h>
#include <executorch/runtime/core/result.h>
#include <executorch/runtime/core/tensor_shape_dynamism.h>
#include <executorch/runtime/platform/compiler.h>

// Forward declaration of a helper that provides access to internal resizing
// methods of TensorImpl. Real definition is in
Expand Down Expand Up @@ -293,6 +295,16 @@ ssize_t compute_numel(
const ::executorch::runtime::etensor::TensorImpl::SizesType* sizes,
ssize_t dim);

/**
* Compute the number of elements based on the sizes of a tensor.
* Returns Error::InvalidArgument if any intermediate multiplication would
* overflow ssize_t, or if a size is negative. Prefer this over compute_numel()
* for paths that can propagate an Error upward.
*/
::executorch::runtime::Result<ssize_t> safe_numel(
Comment thread
lucylq marked this conversation as resolved.
const ::executorch::runtime::etensor::TensorImpl::SizesType* sizes,
ssize_t dim);
Comment thread
lucylq marked this conversation as resolved.

/// Appropriate format specifier for the result of calling
/// size(). Must be used instead of using zd directly to support ATen
/// mode.
Expand Down Expand Up @@ -322,6 +334,7 @@ namespace executor {
// TODO(T197294990): Remove these deprecated aliases once all users have moved
// to the new `::executorch` namespaces.
using ::executorch::runtime::etensor::compute_numel;
using ::executorch::runtime::etensor::safe_numel;
using ::executorch::runtime::etensor::TensorImpl;
} // namespace executor
} // namespace torch
Loading