diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp index 82991011127..bcbaaf326b8 100644 --- a/runtime/executor/program.cpp +++ b/runtime/executor/program.cpp @@ -233,81 +233,16 @@ Result 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(constant_buffer->size()), - static_cast(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(constant_segment->segment_index()), - static_cast(segments->size())); - - const executorch_flatbuffer::DataSegment* data_segment = - segments->Get(constant_segment->segment_index()); - Result 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( loader, segment_base_offset, @@ -316,6 +251,48 @@ Result 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(constant_buffer->size()), + static_cast(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(constant_segment->segment_index()), + static_cast(segments->size())); + + const executorch_flatbuffer::DataSegment* data_segment = + segments->Get(constant_segment->segment_index()); + Result 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 { @@ -408,73 +385,49 @@ Result Program::get_constant_buffer_data( auto internal_program = static_cast(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( - (*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( - static_cast(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(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", + 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( + (*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( + static_cast(constant_segment_data_.data()) + + offset); } Result Program::get_named_data_map() const { diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index dc926184049..b5f7bed26ac 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -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"); @@ -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 = - 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 = programs_["add_mul_program"]->load_method( diff --git a/runtime/executor/test/program_test.cpp b/runtime/executor/test/program_test.cpp index e718901fa49..748040c6760 100644 --- a/runtime/executor/test/program_test.cpp +++ b/runtime/executor/test/program_test.cpp @@ -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 linear_loader = FileDataLoader::from(linear_path); ASSERT_EQ(linear_loader.error(), Error::Ok); - // This file should always be compatible. - Result 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::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) { @@ -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 loader = FileDataLoader::from(path); - ASSERT_EQ(loader.error(), Error::Ok); - - Result program = Program::load(&loader.get()); + Result 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 data = ProgramTestFriend::get_constant_buffer_data( &program.get(), /*buffer_index=*/0, /*nbytes=*/SIZE_MAX);