Skip to content

Numel and Nbytes Validation (#19121) #19121

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
JacobSzwejbka merged 1 commit into from
Apr 24, 2026
Merged

Numel and Nbytes Validation (#19121) #19121

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
28 changes: 28 additions & 0 deletions runtime/executor/method.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include <executorch/runtime/executor/method.h>

#include <c10/util/irange.h>
#include <c10/util/safe_numerics.h>
#include <array>
#include <cinttypes> // @donotremove
#include <cstdint>
Expand Down Expand Up @@ -1194,6 +1195,33 @@ Method::set_input(const EValue& input_evalue, size_t input_idx) {
input_idx,
executorch::runtime::toString(t_dst.scalar_type()),
executorch::runtime::toString(t_src.scalar_type()));

ssize_t numel = 1;
for (ssize_t i = 0; i < t_src.dim(); i++) {
bool overflow = c10::mul_overflows(
numel, static_cast<ssize_t>(t_src.size(i)), &numel);
ET_CHECK_OR_RETURN_ERROR(
!overflow,
InvalidArgument,
"Input %" ET_PRIsize_t
": numel overflowed at dimension %zd with size %zd",
input_idx,
(size_t)i,
(size_t)t_src.size(i));
}
size_t nbytes;
bool nbytes_overflow = c10::mul_overflows(
static_cast<size_t>(numel),
executorch::runtime::elementSize(t_src.scalar_type()),
&nbytes);
ET_CHECK_OR_RETURN_ERROR(
!nbytes_overflow,
InvalidArgument,
"Input %" ET_PRIsize_t
": nbytes overflowed: numel %zd with element size %zu",
input_idx,
numel,
executorch::runtime::elementSize(t_src.scalar_type()));
// Reset the shape for the Method's input as the size of forwarded input
// tensor for shape dynamism. Also is a safety check if need memcpy.
ET_CHECK_OK_OR_RETURN_ERROR(
Expand Down
90 changes: 59 additions & 31 deletions runtime/executor/program_validation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
#include <executorch/runtime/platform/log.h>
#include <executorch/schema/program_generated.h>

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

namespace executorch {
namespace runtime {
Expand All @@ -32,7 +32,8 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) {
return Error::InvalidProgram;
}

// ssize_t numel = 1;
ssize_t numel = 1;
bool numel_overflowed = false;
for (flatbuffers::uoffset_t i = 0; i < sizes->size(); i++) {
int32_t size = sizes->Get(i);

Expand All @@ -45,16 +46,10 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) {
return Error::InvalidProgram;
}

// bool overflow =
// c10::mul_overflows(numel, static_cast<ssize_t>(size), &numel);
// if (overflow) {
// ET_LOG(
// Error,
// "numel overflowed at dimension %u with size %d",
// static_cast<unsigned>(i),
// size);
// return Error::InvalidProgram;
// }
if (!numel_overflowed) {
numel_overflowed =
c10::mul_overflows(numel, static_cast<ssize_t>(size), &numel);
}
}

auto scalar_type =
Expand All @@ -64,19 +59,18 @@ validate_tensor(const executorch_flatbuffer::Tensor* tensor) {
return Error::InvalidProgram;
}

// size_t nbytes;
// bool nbytes_overflow = c10::mul_overflows(
// static_cast<size_t>(numel),
// executorch::runtime::elementSize(scalar_type),
// &nbytes);
// if (nbytes_overflow) {
// ET_LOG(
// Error,
// "nbytes overflowed: numel %zd with element size %zu",
// numel,
// executorch::runtime::elementSize(scalar_type));
// return Error::InvalidProgram;
// }
if (numel_overflowed) {
return Error::InvalidProgram;
}

size_t nbytes;
bool nbytes_overflow = c10::mul_overflows(
static_cast<size_t>(numel),
executorch::runtime::elementSize(scalar_type),
&nbytes);
if (nbytes_overflow) {
return Error::InvalidProgram;
}

return Error::Ok;
}
Expand Down Expand Up @@ -114,6 +108,27 @@ validate_program(const executorch_flatbuffer::Program* program) {
return Error::InvalidProgram;
}

const auto* inputs = plan->inputs();
auto is_dynamic_input = [&](flatbuffers::uoffset_t idx) -> bool {
if (inputs == nullptr) {
return false;
}
for (flatbuffers::uoffset_t i = 0; i < inputs->size(); i++) {
if (inputs->Get(i) == static_cast<int32_t>(idx)) {
const auto* value = values->Get(idx);
if (value == nullptr) {
return false;
}
const auto* tensor =
static_cast<const executorch_flatbuffer::Tensor*>(value->val());
return tensor != nullptr &&
tensor->shape_dynamism() !=
executorch_flatbuffer::TensorShapeDynamism::STATIC;
}
}
return false;
};

for (flatbuffers::uoffset_t value_idx = 0; value_idx < values->size();
value_idx++) {
const auto* value = values->Get(value_idx);
Expand All @@ -128,12 +143,25 @@ validate_program(const executorch_flatbuffer::Program* program) {

Error err = validate_tensor(tensor);
if (err != Error::Ok) {
ET_LOG(
Error,
"Tensor validation failed for value %u in execution plan %u",
static_cast<unsigned>(value_idx),
static_cast<unsigned>(plan_idx));
return err;
// Dynamic input tensors may have upper-bound sizes serialized for
// 64-bit machines that would overflow on 32-bit. Since their actual
// sizes are provided at set_input time, we defer overflow checks
// for those to Method::set_input.
if (is_dynamic_input(value_idx)) {
ET_LOG(
Info,
"Skipping validation failure for dynamic input tensor "
"at value %u in execution plan %u",
static_cast<unsigned>(value_idx),
static_cast<unsigned>(plan_idx));
} else {
ET_LOG(
Error,
"Tensor validation failed for value %u in execution plan %u",
static_cast<unsigned>(value_idx),
static_cast<unsigned>(plan_idx));
return err;
}
}
}

Expand Down
9 changes: 4 additions & 5 deletions runtime/executor/program_validation.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,13 +22,12 @@ namespace executorch {
namespace runtime {

/**
* Validates that computing numel (number of elements) from the tensor's sizes
* will not overflow. This check should be performed before creating TensorImpl
* objects to prevent undefined behavior from integer overflow.
* Validates that a tensor's metadata is semantically valid: sizes are
* non-negative, scalar type is valid, and computing numel/nbytes will not
* overflow.
*
* @param[in] tensor The flatbuffer Tensor to validate.
* @return Error::Ok if the numel calculation is safe, Error::InvalidProgram
* if computing numel would overflow.
* @return Error::Ok if validation passes, Error::InvalidProgram otherwise.
*/
ET_NODISCARD Error validate_tensor(const executorch_flatbuffer::Tensor* tensor);

Expand Down
25 changes: 25 additions & 0 deletions runtime/executor/test/method_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,31 @@ TEST_F(MethodTest, AliasedIOTest) {
}
}

TEST_F(MethodTest, SetInputRejectsOverflowingSizes) {
// The "cat" model (ModuleDynamicCatUnallocatedIO) has a 2D input.
// set_input validates numel/nbytes overflow before resize_tensor.
ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
Result<Method> method = programs_["cat"]->load_method("forward", &mmm.get());
ASSERT_EQ(method.error(), Error::Ok);

// Create a 2D tensor with enormous sizes. On 32-bit platforms the numel
// multiplication overflows; on 64-bit the resize bounds check rejects it.
int32_t sizes[2] = {2000000000, 2000000000};
uint8_t dim_order[2] = {0, 1};
int32_t strides[2] = {1, 1};
executorch::aten::TensorImpl impl(
executorch::aten::ScalarType::Float,
2,
sizes,
nullptr,
dim_order,
strides);

auto input_err =
method->set_input(EValue(executorch::aten::Tensor(&impl)), 0);
EXPECT_NE(input_err, Error::Ok);
}

TEST_F(MethodTest, ConstantSegmentTest) {
// Execute model with constants stored in segment.
ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
Expand Down
126 changes: 98 additions & 28 deletions runtime/executor/test/program_validation_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,12 +64,15 @@ struct EValueConfig {
EValueType type;
std::vector<int32_t> tensor_sizes; // For Tensor type.
std::vector<int32_t> tensor_list_items; // For TensorList type (indices).
bool is_dynamic = false; // For Tensor type: if true, uses DYNAMIC_BOUND.
};

// Unified helper to create a minimal valid PTE flatbuffer with configurable
// evalues. Returns a buffer containing the flatbuffer data.
// evalues. Returns a buffer containing the flatbuffer data. input_indices
// specifies which value indices appear in the execution plan's inputs list.
std::vector<uint8_t> CreateTestProgram(
const std::vector<EValueConfig>& configs) {
const std::vector<EValueConfig>& configs,
const std::vector<int32_t>& input_indices = {}) {
flatbuffers::FlatBufferBuilder builder(1024);

std::vector<flatbuffers::Offset<executorch_flatbuffer::EValue>> evalues;
Expand All @@ -93,7 +96,9 @@ std::vector<uint8_t> CreateTestProgram(
/*data_buffer_idx=*/0,
/*allocation_info=*/0,
/*layout=*/0,
executorch_flatbuffer::TensorShapeDynamism::STATIC,
config.is_dynamic
? executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_BOUND
: executorch_flatbuffer::TensorShapeDynamism::STATIC,
/*extra_tensor_info=*/0);
evalues.push_back(executorch_flatbuffer::CreateEValue(
builder,
Expand Down Expand Up @@ -124,6 +129,7 @@ std::vector<uint8_t> CreateTestProgram(

auto values_vec = builder.CreateVector(evalues);
auto plan_name = builder.CreateString("forward");
auto inputs_vec = builder.CreateVector(input_indices);
auto empty_int_vec = builder.CreateVector(std::vector<int32_t>{});
auto empty_int64_vec = builder.CreateVector(std::vector<int64_t>{0});
auto empty_chain_vec = builder.CreateVector(
Expand All @@ -139,8 +145,8 @@ std::vector<uint8_t> CreateTestProgram(
plan_name,
/*container_meta_type=*/0,
values_vec,
empty_int_vec,
empty_int_vec,
/*inputs=*/inputs_vec,
/*outputs=*/empty_int_vec,
empty_chain_vec,
empty_operators_vec,
empty_delegates_vec,
Expand Down Expand Up @@ -206,29 +212,93 @@ TEST_F(ProgramValidationTest, InternalConsistencyDetectsTruncatedData) {
ASSERT_EQ(program.error(), Error::InvalidProgram);
}

// TEST_F(ProgramValidationTest, TensorNumelOverflowDetected) {
// std::vector<EValueConfig> configs = {
// {EValueType::Tensor, {2000000000, 2000000000, 2000000000}, {}}};
//
// AlignedBuffer buf(CreateTestProgram(configs));
// auto loader = buf.loader();
//
// Result<Program> program =
// Program::load(&loader, Program::Verification::InternalConsistency);
// EXPECT_EQ(program.error(), Error::InvalidProgram);
// }

// TEST_F(ProgramValidationTest, TensorNumelOverflowNotDetectedWithMinimal) {
// std::vector<EValueConfig> configs = {
// {EValueType::Tensor, {2000000000, 2000000000, 2000000000}, {}}};
//
// AlignedBuffer buf(CreateTestProgram(configs));
// auto loader = buf.loader();
//
// // Minimal verification doesn't run program validation.
// Result<Program> program =
// Program::load(&loader, Program::Verification::Minimal);
// }
TEST_F(ProgramValidationTest, TensorNumelOverflowDetectedForStaticTensor) {
// Static tensors always have their overflow checked at validation time.
std::vector<EValueConfig> configs = {
{EValueType::Tensor,
{2000000000, 2000000000, 2000000000},
{},
/*is_dynamic=*/false}};

AlignedBuffer buf(CreateTestProgram(configs));
auto loader = buf.loader();

Result<Program> program =
Program::load(&loader, Program::Verification::InternalConsistency);
EXPECT_EQ(program.error(), Error::InvalidProgram);
}

TEST_F(ProgramValidationTest, TensorNumelOverflowNotDetectedWithMinimal) {
std::vector<EValueConfig> configs = {
{EValueType::Tensor,
{2000000000, 2000000000, 2000000000},
{},
/*is_dynamic=*/false}};

AlignedBuffer buf(CreateTestProgram(configs));
auto loader = buf.loader();

// Minimal verification doesn't run program validation.
Result<Program> program =
Program::load(&loader, Program::Verification::Minimal);
EXPECT_EQ(program.error(), Error::Ok);
}

TEST_F(ProgramValidationTest, TensorNumelOverflowSkippedForDynamicInput) {
// Dynamic input tensors skip overflow checks at validation time; the check
// is deferred to set_input where actual sizes are known.
std::vector<EValueConfig> configs = {
{EValueType::Tensor,
{2000000000, 2000000000, 2000000000},
{},
/*is_dynamic=*/true}};

// Mark value index 0 as a plan input.
AlignedBuffer buf(CreateTestProgram(configs, /*input_indices=*/{0}));
auto loader = buf.loader();

Result<Program> program =
Program::load(&loader, Program::Verification::InternalConsistency);
EXPECT_EQ(program.error(), Error::Ok);
}

TEST_F(
ProgramValidationTest,
TensorNumelOverflowDetectedForDynamicNonInputTensor) {
// A dynamic tensor that is NOT in the inputs list should still have its
// overflow checked at validation time.
std::vector<EValueConfig> configs = {
{EValueType::Tensor,
{2000000000, 2000000000, 2000000000},
{},
/*is_dynamic=*/true}};

// No input indices — the tensor is not a plan input.
AlignedBuffer buf(CreateTestProgram(configs));
auto loader = buf.loader();

Result<Program> program =
Program::load(&loader, Program::Verification::InternalConsistency);
EXPECT_EQ(program.error(), Error::InvalidProgram);
}

TEST_F(ProgramValidationTest, TensorNumelOverflowDetectedForStaticInputTensor) {
// A static input tensor should still have its overflow checked at
// validation time since its sizes cannot change.
std::vector<EValueConfig> configs = {
{EValueType::Tensor,
{2000000000, 2000000000, 2000000000},
{},
/*is_dynamic=*/false}};

// Mark value index 0 as a plan input.
AlignedBuffer buf(CreateTestProgram(configs, /*input_indices=*/{0}));
auto loader = buf.loader();

Result<Program> program =
Program::load(&loader, Program::Verification::InternalConsistency);
EXPECT_EQ(program.error(), Error::InvalidProgram);
}

TEST_F(ProgramValidationTest, NegativeSizeDetected) {
std::vector<EValueConfig> configs = {{EValueType::Tensor, {10, -5, 10}, {}}};
Expand Down
Loading