Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
235 changes: 94 additions & 141 deletions runtime/executor/program.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -233,81 +233,16 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
pte_data_map.emplace(std::move(pte_data_map_result.get()));
}

// Constant data may live inside the flatbuffer data (constant_buffer) or in a
// separate segment (constant_segment). It should not be in both.
// Check constant_segment->offsets()->size() > 1, as the offsets list will
// always contain a placeholder value 0 for non-const tensors. If this is the
// only offset, the constant segment is empty and does not need to be loaded.
// Constant data lives in a separate segment (constant_segment).
const auto* constant_segment = flatbuffer_program->constant_segment();
if (constant_segment != nullptr && constant_segment->offsets() != nullptr &&
constant_segment->offsets()->size() > 0) {
if (constant_segment->offsets()->size() == 1) {
// No constants; the constant segment is empty and does not
// need to be loaded.
return Program(
loader,
segment_base_offset,
std::move(program_data.get()),
flatbuffer_program,
/*constant_segment_data=*/FreeableBuffer{},
std::move(pte_data_map));
}
// The constant data is inside a separate segment.
const auto* constant_buffer = flatbuffer_program->constant_buffer();
ET_CHECK_OR_RETURN_ERROR(
constant_buffer == nullptr || constant_buffer->size() == 0,
InvalidProgram,
"constant_buffer contains %zu items, "
"constant_segment.offsets contains %zu items. Only one should be used.",
static_cast<size_t>(constant_buffer->size()),
static_cast<size_t>(constant_segment->offsets()->size()));
const auto* segments = flatbuffer_program->segments();
ET_CHECK_OR_RETURN_ERROR(
segments != nullptr, InvalidProgram, "No segments in program");

// Load constant segment.
// TODO(T171839323): Add test for segment_index > num available segments.
ET_CHECK_OR_RETURN_ERROR(
constant_segment->segment_index() < segments->size(),
InvalidProgram,
"Constant segment index %zu invalid for program segments range %zu",
static_cast<size_t>(constant_segment->segment_index()),
static_cast<size_t>(segments->size()));

const executorch_flatbuffer::DataSegment* data_segment =
segments->Get(constant_segment->segment_index());
Result<FreeableBuffer> constant_segment_data = loader->load(
segment_base_offset + data_segment->offset(),
data_segment->size(),
DataLoader::SegmentInfo(
DataLoader::SegmentInfo::Type::Constant,
constant_segment->segment_index()));
if (!constant_segment_data.ok()) {
return constant_segment_data.error();
}
// The FreeableBuffer owns the data that flatbuffer_program points into.
// Also keep a pointer to the loader so it can load more segments when
// necessary.
return Program(
loader,
segment_base_offset,
std::move(program_data.get()),
flatbuffer_program,
std::move(constant_segment_data.get()),
std::move(pte_data_map));
} else {
// The constant data is stored inside the flatbuffer, so this program does
// not contain a separate segment for it.

// NOTE: This branch is deprecated from ExecuTorch 0.7 onwards.
// Please regenerate your PTE file to ensure newer ExecuTorch runtimes can
// support it. ExecuTorch deprecation policy:
// https://docs.pytorch.org/executorch/stable/api-life-cycle.html#deprecation-policy.
// For support, contact the PyTorch Edge team or make an issue in:
// https://github.com/pytorch/executorch/issues.
ET_LOG(
Error,
"!!DEPRECATED!! This branch is deprecated from ExecuTorch 0.7; re-export this PTE file to ensure support on newer runtimes.");
ET_CHECK_OR_RETURN_ERROR(
constant_segment != nullptr && constant_segment->offsets() != nullptr,
InvalidProgram,
"Program constant segment is nullptr. Expect at least one constant placeholder.");
// The offsets list will always contain a placeholder value 0 for non-const
// tensors. If this is the only offset, the constant segment is empty and
// does not need to be loaded.
if (constant_segment->offsets()->size() == 1) {
return Program(
Comment on lines +236 to 246
Copy link

Copilot AI Apr 20, 2026

Choose a reason for hiding this comment

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

In Program::load(), the early return for constant_segment->offsets()->size() == 1 happens before validating Program.constant_buffer is empty. This means a malformed/corrupt PTE could still include deprecated inline constant_buffer data and not be rejected. Also consider explicitly rejecting offsets()->size() == 0 (no placeholder) since the loader currently treats it as “has constants” and attempts to load a segment, which contradicts the placeholder invariant used elsewhere/tests.

Copilot uses AI. Check for mistakes.
loader,
segment_base_offset,
Expand All @@ -316,6 +251,48 @@ Result<executorch_flatbuffer::ExecutionPlan*> get_execution_plan(
/*constant_segment_data=*/FreeableBuffer{},
std::move(pte_data_map));
}
const auto* constant_buffer = flatbuffer_program->constant_buffer();
ET_CHECK_OR_RETURN_ERROR(
constant_buffer == nullptr || constant_buffer->size() == 0,
InvalidProgram,
"constant_buffer contains %zu items, "
"constant_segment.offsets contains %zu items. Only one should be used.",
static_cast<size_t>(constant_buffer->size()),
static_cast<size_t>(constant_segment->offsets()->size()));
const auto* segments = flatbuffer_program->segments();
ET_CHECK_OR_RETURN_ERROR(
segments != nullptr, InvalidProgram, "No segments in program");

// Load constant segment.
// TODO(T171839323): Add test for segment_index > num available segments.
ET_CHECK_OR_RETURN_ERROR(
constant_segment->segment_index() < segments->size(),
InvalidProgram,
"Constant segment index %zu invalid for program segments range %zu",
static_cast<size_t>(constant_segment->segment_index()),
static_cast<size_t>(segments->size()));

const executorch_flatbuffer::DataSegment* data_segment =
segments->Get(constant_segment->segment_index());
Result<FreeableBuffer> constant_segment_data = loader->load(
segment_base_offset + data_segment->offset(),
data_segment->size(),
DataLoader::SegmentInfo(
DataLoader::SegmentInfo::Type::Constant,
constant_segment->segment_index()));
if (!constant_segment_data.ok()) {
return constant_segment_data.error();
}
// The FreeableBuffer owns the data that flatbuffer_program points into.
// Also keep a pointer to the loader so it can load more segments when
// necessary.
return Program(
loader,
segment_base_offset,
std::move(program_data.get()),
flatbuffer_program,
std::move(constant_segment_data.get()),
std::move(pte_data_map));
}

size_t Program::num_methods() const {
Expand Down Expand Up @@ -408,73 +385,49 @@ Result<const void*> Program::get_constant_buffer_data(
auto internal_program =
static_cast<const executorch_flatbuffer::Program*>(internal_program_);

// Constant data is either in a separate segment (constant_segment_data) and
// loaded during Program::load, or stored inside the flatbuffer data
// (constant_buffer).
if (constant_segment_data_.data() != nullptr) {
const auto* constant_segment = internal_program->constant_segment();
size_t num_elems = constant_segment == nullptr
? 0
: (constant_segment->offsets() == nullptr
? 0
: constant_segment->offsets()->size());
ET_CHECK_OR_RETURN_ERROR(
buffer_index < num_elems,
InvalidArgument,
"Constant segment buffer index %zu invalid for program constant segment range %zu",
buffer_index,
num_elems);

// All constant data is stored in one segment, with each tensor aligned to
// @executorch_tensor_alignment. Tensor offsets are stored in the flatbuffer
// data in Program.constant_segment.offsets.
// The constant data at buffer_index is located at: base address of the
// constant segment + offset for tensor at buffer_index.
uint64_t offset = static_cast<uint64_t>(
(*internal_program->constant_segment()->offsets())[buffer_index]);

size_t size = constant_segment_data_.size();
ET_CHECK_OR_RETURN_ERROR(
offset <= size && nbytes <= size - offset,
InvalidArgument,
"Constant segment offset %" PRIu64
" + size_bytes %zu invalid for program constant segment size %zu",
offset,
nbytes,
size);

// Offset is wrt the beginning of the constant segment.
return static_cast<const void*>(
static_cast<const unsigned char*>(constant_segment_data_.data()) +
offset);
} else {
// Otherwise, the constant data is stored inside Program.constant_buffer.
const auto* constant_buffer_ptr = internal_program->constant_buffer();
size_t num_elems =
constant_buffer_ptr == nullptr ? 0 : constant_buffer_ptr->size();
ET_CHECK_OR_RETURN_ERROR(
buffer_index < num_elems,
InvalidArgument,
"Constant buffer index %zu invalid for program constant buffer range %zu",
buffer_index,
num_elems);

const auto& constant_buffer = *constant_buffer_ptr;
const auto* storage = constant_buffer[buffer_index]->storage();
auto storage_size = storage == nullptr ? 0 : storage->size();
// nbytes (requested from the program) should be less than storage_size
// (size of the constant buffer from PTE), to prevent reading out of bounds.
// in some cases storage size may be larger than nbytes because of padding;
// executorch-tensor-alignment, or 16 by default.
ET_CHECK_OR_RETURN_ERROR(
nbytes <= storage_size,
InvalidArgument,
"Requested nbytes %zu exceeds constant buffer storage size %zu",
nbytes,
static_cast<size_t>(storage_size));

return storage->data();
}
// Constant data must live in a separate segment (constant_segment_data) that
// was loaded during Program::load. The deprecated inline
// Program.constant_buffer path is rejected at load time.
ET_CHECK_OR_RETURN_ERROR(
constant_segment_data_.data() != nullptr,
InvalidProgram,
"Program has no constant segment; cannot retrieve constant data.");

const auto* constant_segment = internal_program->constant_segment();
size_t num_elems = constant_segment == nullptr
? 0
: (constant_segment->offsets() == nullptr
? 0
: constant_segment->offsets()->size());
ET_CHECK_OR_RETURN_ERROR(
buffer_index < num_elems,
InvalidArgument,
"Constant segment buffer index %zu invalid for program constant segment range %zu",
Comment on lines +388 to +405
Copy link

Copilot AI Apr 20, 2026

Choose a reason for hiding this comment

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

Program::get_constant_buffer_data() now hard-fails with InvalidProgram when constant_segment_data_.data() == nullptr. However, valid programs with an empty constant segment (segment size 0 / offsets only placeholder) will naturally have constant_segment_data_ empty/null (and Program::load() intentionally doesn’t load it). Consider using constant_segment.offsets()->size() to distinguish “no constants / invalid buffer_index” (should be InvalidArgument) from “constants expected but segment data missing” (could be InvalidProgram).

Copilot uses AI. Check for mistakes.
buffer_index,
num_elems);

// All constant data is stored in one segment, with each tensor aligned to
// @executorch_tensor_alignment. Tensor offsets are stored in the flatbuffer
// data in Program.constant_segment.offsets.
// The constant data at buffer_index is located at: base address of the
// constant segment + offset for tensor at buffer_index.
uint64_t offset = static_cast<uint64_t>(
(*internal_program->constant_segment()->offsets())[buffer_index]);

size_t size = constant_segment_data_.size();
ET_CHECK_OR_RETURN_ERROR(
offset <= size && nbytes <= size - offset,
InvalidArgument,
"Constant segment offset %" PRIu64
" + size_bytes %zu invalid for program constant segment size %zu",
offset,
nbytes,
size);

// Offset is wrt the beginning of the constant segment.
return static_cast<const void*>(
static_cast<const unsigned char*>(constant_segment_data_.data()) +
offset);
}

Result<const NamedDataMap*> Program::get_named_data_map() const {
Expand Down
30 changes: 0 additions & 30 deletions runtime/executor/test/method_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -81,9 +81,6 @@ class MethodTest : public ::testing::Test {
std::getenv("ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH"), "cat");
load_program(std::getenv("ET_MODULE_ADD_MUL_PATH"), "add_mul");
load_program(std::getenv("ET_MODULE_STATEFUL_PATH"), "stateful");
load_program(
std::getenv("DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH"),
"linear_constant_buffer");

load_program(
std::getenv("ET_MODULE_ADD_MUL_PROGRAM_PATH"), "add_mul_program");
Expand Down Expand Up @@ -337,33 +334,6 @@ TEST_F(MethodTest, ConstantSegmentTest) {
ASSERT_EQ(err, Error::Ok);
}

TEST_F(MethodTest, ConstantBufferTest) {
// Execute model with constants stored in the program flatbuffer.
ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
Result<Method> method =
programs_["linear_constant_buffer"]->load_method("forward", &mmm.get());
ASSERT_EQ(method.error(), Error::Ok);

// Set a dummy input.
int32_t sizes[2] = {2, 2};
uint8_t dim_order[2] = {0, 1};
int32_t strides[2] = {2, 1};
executorch::aten::TensorImpl impl(
executorch::aten::ScalarType::Float,
2,
sizes,
nullptr,
dim_order,
strides);
auto input_err = method->set_input(
executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 0);
ASSERT_EQ(input_err, Error::Ok);

// Can execute the method.
Error err = method->execute();
ASSERT_EQ(err, Error::Ok);
}

TEST_F(MethodTest, ProgramDataSeparationTest) {
ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
Result<Method> method = programs_["add_mul_program"]->load_method(
Expand Down
38 changes: 7 additions & 31 deletions runtime/executor/test/program_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -502,36 +502,17 @@ TEST_F(ProgramTest, LoadConstantSegment) {
EXPECT_GE(flatbuffer_program->constant_segment()->offsets()->size(), 1);
}

TEST_F(ProgramTest, LoadConstantSegmentWhenConstantBufferExists) {
// Load the serialized ModuleAddMul data, with constants in the flatbuffer and
// no constants in the segment.
TEST_F(ProgramTest, RejectsDeprecatedInlineConstantBuffer) {
// PTE files that store constants inline in the flatbuffer
// (Program.constant_buffer, deprecated since ExecuTorch 0.7) must be
// rejected at load.
const char* linear_path =
std::getenv("DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH");
Result<FileDataLoader> linear_loader = FileDataLoader::from(linear_path);
ASSERT_EQ(linear_loader.error(), Error::Ok);

// This file should always be compatible.
Result<FreeableBuffer> linear_header = linear_loader->load(
/*offset=*/0,
Program::kMinHeadBytes,
/*segment_info=*/
DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program));
ASSERT_EQ(linear_header.error(), Error::Ok);
EXPECT_EQ(
Program::check_header(linear_header->data(), linear_header->size()),
Program::HeaderStatus::CompatibleVersion);

Result<Program> program = Program::load(&linear_loader.get());
ASSERT_EQ(program.error(), Error::Ok);

const executorch_flatbuffer::Program* flatbuffer_program =
ProgramTestFriend::GetInternalProgram(&program.get());

// Expect no segments.
EXPECT_EQ(flatbuffer_program->segments()->size(), 0);

// The constant buffer should exist.
EXPECT_GE(flatbuffer_program->constant_buffer()->size(), 1);
EXPECT_EQ(program.error(), Error::InvalidProgram);
}

TEST_F(ProgramTest, LoadFromMutableSegment) {
Expand Down Expand Up @@ -867,15 +848,10 @@ TEST_F(ProgramTest, NullPlanNameDoesNotCrash) {
}

TEST_F(ProgramTest, GetConstantBufferDataRejectsOversizedRequest) {
const char* path =
std::getenv("DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH");
Result<FileDataLoader> loader = FileDataLoader::from(path);
ASSERT_EQ(loader.error(), Error::Ok);

Result<Program> program = Program::load(&loader.get());
Result<Program> program = Program::load(add_loader_.get());
ASSERT_EQ(program.error(), Error::Ok);

// Request way more bytes than any buffer could contain
// Request way more bytes than the constant segment could contain.
Result<const void*> data = ProgramTestFriend::get_constant_buffer_data(
&program.get(), /*buffer_index=*/0, /*nbytes=*/SIZE_MAX);

Comment on lines 850 to 857
Copy link

Copilot AI Apr 20, 2026

Choose a reason for hiding this comment

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

This test now loads add_loader_ (ET_MODULE_ADD), which earlier tests indicate has no constants (constant segment is only a placeholder / size 0). With the new runtime behavior, get_constant_buffer_data() will return InvalidProgram (no segment data) rather than exercising the oversized-request bounds check. Use a module that actually has constant-segment data (e.g., ET_MODULE_ADD_MUL_PATH) and a non-placeholder buffer_index (typically 1) to keep this test asserting the intended behavior.

Copilot uses AI. Check for mistakes.
Expand Down
Loading