Skip to content

Fix #18654: Add safe_numel() #18698

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

Open
JiwaniZakir wants to merge 1 commit into
base: main
Choose a base branch
from
+55 −0
Open

Fix #18654: Add safe_numel() #18698

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
20 changes: 20 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 @@ -10,6 +10,7 @@

#include <algorithm>
#include <cstdint>
#include <limits>

#include <c10/util/irange.h>

Expand Down Expand Up @@ -68,6 +69,25 @@ TensorImpl::TensorImpl(
ET_CHECK_MSG(dim_ >= 0, "Dimension must be non-negative, got %zd", dim_);
}

Result<ssize_t> TensorImpl::safe_numel() const {
ssize_t numel = 1;
for (const auto i : c10::irange(dim_)) {
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);
Comment on lines +72 to +80
Copy link

Copilot AI Apr 4, 2026

Choose a reason for hiding this comment

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

safe_numel() duplicates much of compute_numel()’s logic (iteration + non-negative size checks). Consider factoring the shared computation into a single helper (e.g., a checked-multiply routine) so future changes to size validation/numel semantics don’t have to be kept in sync in multiple places.

Copilot uses AI. Check for mistakes.
ET_CHECK_OR_RETURN_ERROR(
sizes_[i] == 0 ||
numel <= std::numeric_limits<ssize_t>::max() / sizes_[i],
InvalidArgument,
"Tensor numel overflows ssize_t");
Copy link

Copilot AI Apr 4, 2026

Choose a reason for hiding this comment

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

The overflow failure path logs/returns "Tensor numel overflows ssize_t" without any context (which dimension triggered it, the offending size, or the partial product). Including at least the dimension index and size value (and optionally the partial numel) would make this error actionable when debugging shape issues.

Suggested change
"Tensor numel overflows ssize_t");
"Tensor numel overflows ssize_t at dimension %zd: size=%zd, partial_numel=%zd",
i,
static_cast<ssize_t>(sizes_[i]),
numel);

Copilot uses AI. Check for mistakes.
numel *= sizes_[i];
}
return numel;
}

size_t TensorImpl::nbytes() const {
return numel_ * elementSize(type_);
}
Expand Down
7 changes: 7 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 @@ -10,6 +10,7 @@

#include <executorch/runtime/core/array_ref.h>
#include <executorch/runtime/core/error.h>
#include <executorch/runtime/core/result.h>
#include <executorch/runtime/core/portable_type/device.h>
#include <executorch/runtime/core/portable_type/scalar_type.h>
#include <executorch/runtime/core/tensor_shape_dynamism.h>
Expand Down Expand Up @@ -149,6 +150,12 @@ class TensorImpl {
return numel_;
}

/**
* Returns the number of elements in the tensor, or an error if the result
* would overflow ssize_t.
Comment on lines +154 to +155
Copy link

Copilot AI Apr 4, 2026

Choose a reason for hiding this comment

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

The doc comment for safe_numel() says it errors only on overflow, but the implementation also returns InvalidArgument when any dimension size is negative. Please update the comment to reflect all error conditions (e.g., negative sizes and overflow) so callers know what to expect.

Suggested change
* Returns the number of elements in the tensor, or an error if the result
* would overflow ssize_t.
* Returns the number of elements in the tensor, or an error if any
* dimension size is negative or if computing the result would overflow
* ssize_t.

Copilot uses AI. Check for mistakes.
*/
Result<ssize_t> safe_numel() const;

/// Returns the type of the elements in the tensor (int32, float, bool, etc).
ScalarType scalar_type() const {
return type_;
Expand Down
28 changes: 28 additions & 0 deletions runtime/core/portable_type/test/tensor_impl_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include <executorch/runtime/core/portable_type/tensor_impl.h>

#include <gtest/gtest.h>
#include <limits>
#include <random>

#include <executorch/runtime/core/exec_aten/util/tensor_util.h>
Expand Down Expand Up @@ -40,6 +41,33 @@ class TensorImplTest : public ::testing::Test {
}
};

TEST_F(TensorImplTest, SafeNumelReturnsCorrectValue) {
SizesType sizes[2] = {3, 2};
TensorImpl t(ScalarType::Float, 2, sizes);
auto result = t.safe_numel();
ASSERT_TRUE(result.ok());
EXPECT_EQ(*result, 6);
}

TEST_F(TensorImplTest, SafeNumelScalar) {
TensorImpl t(ScalarType::Float, 0, nullptr);
auto result = t.safe_numel();
ASSERT_TRUE(result.ok());
EXPECT_EQ(*result, 1);
}

TEST_F(TensorImplTest, SafeNumelOverflowReturnsError) {
// Three large dimensions whose product overflows ssize_t on any platform:
// On 64-bit: INT32_MAX^2 * 3 > INT64_MAX; on 32-bit: INT32_MAX^2 > INT32_MAX.
SizesType sizes[3] = {
std::numeric_limits<SizesType>::max(),
std::numeric_limits<SizesType>::max(),
3};
Comment on lines +60 to +65
Copy link

Copilot AI Apr 4, 2026

Choose a reason for hiding this comment

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

This test assumes ssize_t is at most 64-bit (so INT32_MAX^2*3 always overflows). On platforms with wider ssize_t (or if the project ever targets such), safe_numel() could legitimately succeed and this test would fail. Consider constructing the shape to overflow based on std::numeric_limits<ssize_t>::max() (e.g., increasing the number of max-sized dimensions until overflow is guaranteed) or guarding the expectation with a static/runtime check of ssize_t width.

Copilot uses AI. Check for mistakes.
TensorImpl t(ScalarType::Float, 3, sizes);
auto result = t.safe_numel();
EXPECT_FALSE(result.ok());
}

TEST_F(TensorImplTest, TestCtorAndGetters) {
SizesType sizes[2] = {3, 2};
DimOrderType dim_order[2] = {0, 1};
Expand Down
Loading