diff --git a/CMakeLists.txt b/CMakeLists.txt index ca8d1bbbcf2..f3b3390c0ed 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -182,6 +182,10 @@ option(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER "Build the Data Loader extension" OFF ) +option(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR "Build the Flat Tensor extension" + OFF +) + option(EXECUTORCH_BUILD_EXTENSION_MODULE "Build the Module extension" OFF) option(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL "Build the Runner Util extension" @@ -694,6 +698,11 @@ if(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/data_loader) endif() +if(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/flat_tensor) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/flat_tensor/serialize) +endif() + if(EXECUTORCH_BUILD_EXTENSION_MODULE) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/module) endif() diff --git a/build/cmake_deps.toml b/build/cmake_deps.toml index c7680708688..287db04045e 100644 --- a/build/cmake_deps.toml +++ b/build/cmake_deps.toml @@ -148,6 +148,27 @@ deps = [ "executorch", ] +[targets.extension_flat_tensor_schema] +buck_targets = [ + "//extension/flat_tensor/serialize:generated_headers", +] +filters = [ + ".fbs$", +] + +[targets.extension_flat_tensor] +buck_targets = [ + "//extension/flat_tensor:flat_tensor_data_map", +] +filters = [ + ".cpp$", +] +deps = [ + "extension_flat_tensor_schema", + "executorch_core", + "executorch", +] + [targets.extension_module] buck_targets = [ "//extension/module:module", diff --git a/extension/flat_tensor/CMakeLists.txt b/extension/flat_tensor/CMakeLists.txt new file mode 100644 index 00000000000..c0bfa516a4d --- /dev/null +++ b/extension/flat_tensor/CMakeLists.txt @@ -0,0 +1,39 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Please this file formatted by running: +# ~~~ +# cmake-format -i CMakeLists.txt +# ~~~ + +cmake_minimum_required(VERSION 3.19) + +# Source root directory for executorch. +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..) +endif() + +list(TRANSFORM _extension_flat_tensor__srcs PREPEND "${EXECUTORCH_ROOT}/") +add_library(extension_flat_tensor ${_extension_flat_tensor__srcs}) +target_link_libraries(extension_flat_tensor executorch extension_data_loader) +target_include_directories(extension_flat_tensor PUBLIC + ${EXECUTORCH_ROOT}/.. + "${CMAKE_BINARY_DIR}/extension/flat_tensor/include" + "${EXECUTORCH_ROOT}/third-party/flatbuffers/include" + ${_common_include_directories}) +target_compile_options(extension_flat_tensor PUBLIC ${_common_compile_options}) + +# Install libraries +install( + TARGETS extension_flat_tensor + DESTINATION lib + INCLUDES + DESTINATION ${_common_include_directories} +) + +if(BUILD_TESTING) + add_subdirectory(test) +endif() diff --git a/extension/flat_tensor/flat_tensor_data_map.cpp b/extension/flat_tensor/flat_tensor_data_map.cpp index b6d03f88148..20ebc99994a 100644 --- a/extension/flat_tensor/flat_tensor_data_map.cpp +++ b/extension/flat_tensor/flat_tensor_data_map.cpp @@ -8,9 +8,10 @@ #include +#include #include -#include +#include #include #include #include diff --git a/extension/flat_tensor/serialize/CMakeLists.txt b/extension/flat_tensor/serialize/CMakeLists.txt new file mode 100644 index 00000000000..d8541b95cf3 --- /dev/null +++ b/extension/flat_tensor/serialize/CMakeLists.txt @@ -0,0 +1,69 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Flatbuffer schema header lib. Please this file formatted by running: +# ~~~ +# cmake-format -i CMakeLists.txt +# ~~~ + +if(NOT FLATC_EXECUTABLE) + set(FLATC_EXECUTABLE flatc) +endif() + +# The include directory that will contain the generated schema headers. +set(_flat_tensor_schema__include_dir "${CMAKE_BINARY_DIR}/extension/flat_tensor/include") +set(_flat_tensor_schema__output_dir "${_flat_tensor_schema__include_dir}/executorch/extension/flat_tensor/serialize") +# Source root directory for executorch. +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/..) +endif() + +function(generate_flat_tensor_schema _schema_srcs _schema_name) + set(_schema_outputs) + foreach(fbs_file ${_schema_srcs}) + string(REGEX REPLACE "[.]fbs$" "_generated.h" generated "${fbs_file}") + list(APPEND _schema_outputs + "${_flat_tensor_schema__output_dir}/${generated}" + ) + endforeach() + + # Generate the headers from the .fbs files. + add_custom_command( + OUTPUT ${_schema_outputs} + COMMAND + ${FLATC_EXECUTABLE} --cpp --cpp-std c++11 --gen-mutable --scoped-enums -o + "${_flat_tensor_schema__output_dir}" ${_schema_srcs} + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + DEPENDS ${FLATC_EXECUTABLE} ${_schema_srcs} + COMMENT "Generating ${_schema_name} headers" + VERBATIM + ) + + add_library(${_schema_name} INTERFACE ${_schema_outputs}) + set_target_properties(${_schema_name} PROPERTIES LINKER_LANGUAGE CXX) + + # exir lets users set the alignment of tensor data embedded in the flatbuffer, + # and some users need an alignment larger than the default, which is typically + # 32. + target_compile_definitions( + ${_schema_name} INTERFACE FLATBUFFERS_MAX_ALIGNMENT=1024 + ) + + target_include_directories( + ${_schema_name} + INTERFACE ${_flat_tensor_schema__include_dir} + ${EXECUTORCH_ROOT}/third-party/flatbuffers/include + ) +endfunction() + +# Generate common schema +set(scalar_type_schema_srcs scalar_type.fbs) +generate_flat_tensor_schema("${scalar_type_schema_srcs}" "scalar_type_schema") + +# For the other schemas +set(flat_tensor_schema_srcs flat_tensor.fbs) +generate_flat_tensor_schema("${flat_tensor_schema_srcs}" "flat_tensor_schema") +add_dependencies(flat_tensor_schema scalar_type_schema) diff --git a/extension/flat_tensor/serialize/serialize.cpp b/extension/flat_tensor/serialize/serialize.cpp index 11f67c9b16e..720e104ab7f 100644 --- a/extension/flat_tensor/serialize/serialize.cpp +++ b/extension/flat_tensor/serialize/serialize.cpp @@ -8,9 +8,9 @@ #include +#include #include #include -#include #include #include diff --git a/extension/flat_tensor/serialize/targets.bzl b/extension/flat_tensor/serialize/targets.bzl index e93638b5e4a..78054af30e9 100644 --- a/extension/flat_tensor/serialize/targets.bzl +++ b/extension/flat_tensor/serialize/targets.bzl @@ -8,7 +8,7 @@ def define_common_targets(): "scalar_type.fbs", ], outs = { - "schema_generated.h": ["flat_tensor_generated.h"], + "flat_tensor_generated.h": ["flat_tensor_generated.h"], "scalar_type_generated.h": ["scalar_type_generated.h"] }, cmd = " ".join([ @@ -29,7 +29,7 @@ def define_common_targets(): "//executorch/...", ], exported_headers = { - "schema_generated.h": ":gen_schema[schema_generated.h]", + "flat_tensor_generated.h": ":gen_schema[flat_tensor_generated.h]", "scalar_type_generated.h": ":gen_schema[scalar_type_generated.h]", }, exported_external_deps = ["flatbuffers-api"], diff --git a/extension/flat_tensor/test/CMakeLists.txt b/extension/flat_tensor/test/CMakeLists.txt new file mode 100644 index 00000000000..dbda1144fae --- /dev/null +++ b/extension/flat_tensor/test/CMakeLists.txt @@ -0,0 +1,51 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# @generated by test/utils/generate_gtest_cmakelists.py +# +# This file should be formatted with +# ~~~ +# cmake-format -i CMakeLists.txt +# ~~~ +# It should also be cmake-lint clean. +# + +cmake_minimum_required(VERSION 3.19) + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) + +include(${EXECUTORCH_ROOT}/build/Test.cmake) + +add_custom_command( + COMMENT "Lucy: Generating from extension/flat_tensor ${CMAKE_CURRENT_BINARY_DIR}" + OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte" + "${CMAKE_CURRENT_BINARY_DIR}/_default_external_constant.ptd" + + COMMAND + python -m test.models.export_program --modules "ModuleLinear" --external-constants --outdir "${CMAKE_CURRENT_BINARY_DIR}" 2> /dev/null + + WORKING_DIRECTORY ${EXECUTORCH_ROOT} +) + +add_custom_target( + program_data_files + DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte" + "${CMAKE_CURRENT_BINARY_DIR}/_default_external_constant.ptd" +) + +set(test_env_ + "ET_MODULE_LINEAR_PROGRAM_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModuleLinearProgram.pte" + "ET_MODULE_LINEAR_DATA_PATH=${CMAKE_CURRENT_BINARY_DIR}/_default_external_constant.ptd" +) + +# set(ENV{RESOURCES_PATH} "${CMAKE_BINARY_DIR}") + +set(_test_srcs flat_tensor_data_map_test.cpp flat_tensor_header_test.cpp) + +et_cxx_test(extension_flat_tensor_test SOURCES ${_test_srcs} EXTRA_LIBS extension_flat_tensor extension_data_loader) + +add_dependencies(extension_flat_tensor_test program_data_files) +set_property(TEST extension_flat_tensor_test PROPERTY ENVIRONMENT ${test_env_}) diff --git a/extension/flat_tensor/test/flat_tensor_data_map_test.cpp b/extension/flat_tensor/test/flat_tensor_data_map_test.cpp index 9673f5c850c..fd64c92119e 100644 --- a/extension/flat_tensor/test/flat_tensor_data_map_test.cpp +++ b/extension/flat_tensor/test/flat_tensor_data_map_test.cpp @@ -8,8 +8,8 @@ #include #include +#include #include -#include #include #include #include @@ -35,7 +35,8 @@ class FlatTensorDataMapTest : public ::testing::Test { // Load data map. The eager linear model is defined at: // //executorch/test/models/linear_model.py - const char* path = std::getenv("ET_MODULE_LINEAR_DATA"); + const char* path = std::getenv("ET_MODULE_LINEAR_DATA_PATH"); + ET_LOG(Info, "Lucy: %s", path); Result loader = FileDataLoader::from(path); ASSERT_EQ(loader.error(), Error::Ok); @@ -46,6 +47,8 @@ class FlatTensorDataMapTest : public ::testing::Test { }; TEST_F(FlatTensorDataMapTest, LoadFlatTensorDataMap) { + const char* path = std::getenv("ET_MODULE_LINEAR_DATA_PATH"); + ET_LOG(Info, "Lucy: %s", path); Result data_map = FlatTensorDataMap::load(data_map_loader_.get()); EXPECT_EQ(data_map.error(), Error::Ok); diff --git a/extension/flat_tensor/test/targets.bzl b/extension/flat_tensor/test/targets.bzl index 813a1b56d98..bc04edfbe1e 100644 --- a/extension/flat_tensor/test/targets.bzl +++ b/extension/flat_tensor/test/targets.bzl @@ -35,8 +35,8 @@ def define_common_targets(is_fbcode=False): # The tests use this var to find the program file to load. This uses # an fbcode target path because the authoring/export tools # intentionally don't work in xplat (since they're host-only tools). - "ET_MODULE_LINEAR_PROGRAM": "$(location fbcode//executorch/test/models:exported_programs_with_data_separated[ModuleLinear.pte])", - "ET_MODULE_LINEAR_DATA": "$(location fbcode//executorch/test/models:exported_programs_with_data_separated[ModuleLinear.ptd])", + "ET_MODULE_LINEAR_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.pte])", + "ET_MODULE_LINEAR_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.ptd])", } runtime.cxx_test( diff --git a/extension/flat_tensor/test/test_serialize.cpp b/extension/flat_tensor/test/test_serialize.cpp index 5013d75622e..82f47684c71 100644 --- a/extension/flat_tensor/test/test_serialize.cpp +++ b/extension/flat_tensor/test/test_serialize.cpp @@ -8,9 +8,9 @@ #include +#include #include #include -#include #include #include diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index c6fe98abcc8..d435678ca2b 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -288,7 +289,7 @@ Result parse_cond_value(const EValue& cond_value) { } // namespace -Error Method::parse_values() { +Error Method::parse_values(const NamedDataMap* named_data_map) { auto flatbuffer_values = serialization_plan_->values(); ET_CHECK_OR_RETURN_ERROR( flatbuffer_values != nullptr, InvalidProgram, "Missing values"); @@ -414,7 +415,8 @@ Error Method::parse_values() { auto t = deserialization::parseTensor( program_, memory_manager_, - static_cast(val)); + static_cast(val), + named_data_map); if (!t.ok()) { ET_LOG( Error, @@ -607,7 +609,8 @@ Result Method::load( executorch_flatbuffer::ExecutionPlan* s_plan, const Program* program, MemoryManager* memory_manager, - EventTracer* event_tracer) { + EventTracer* event_tracer, + const NamedDataMap* named_data_map) { MemoryAllocator* temp_allocator = memory_manager->temp_allocator(); if (temp_allocator == nullptr) { PlatformMemoryAllocator* platform_allocator = @@ -621,7 +624,7 @@ Result Method::load( } Method method(program, memory_manager, event_tracer, temp_allocator); - Error err = method.init(s_plan); + Error err = method.init(s_plan, named_data_map); if (err != Error::Ok) { return err; } else { @@ -630,7 +633,9 @@ Result Method::load( } } -Error Method::init(executorch_flatbuffer::ExecutionPlan* s_plan) { +Error Method::init( + executorch_flatbuffer::ExecutionPlan* s_plan, + const NamedDataMap* named_data_map) { EXECUTORCH_SCOPE_PROF("Method::init"); internal::EventTracerProfileMethodScope event_tracer_profile_scope = internal::EventTracerProfileMethodScope(event_tracer_, "Method::init"); @@ -647,7 +652,7 @@ Error Method::init(executorch_flatbuffer::ExecutionPlan* s_plan) { { // Parse the elements of the values_ array. - Error err = parse_values(); + Error err = parse_values(named_data_map); if (err != Error::Ok) { return err; } diff --git a/runtime/executor/method.h b/runtime/executor/method.h index ac6bdd8728b..dff4e818f9f 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -294,14 +295,17 @@ class Method final { executorch_flatbuffer::ExecutionPlan* s_plan, const Program* program, MemoryManager* memory_manager, - EventTracer* event_tracer); + EventTracer* event_tracer, + const NamedDataMap* named_data_map); /** * Initialize the method from its serialized representation. * * @returns Error::Ok on success, non-Ok on failure. */ - ET_NODISCARD Error init(executorch_flatbuffer::ExecutionPlan* s_plan); + ET_NODISCARD Error init( + executorch_flatbuffer::ExecutionPlan* s_plan, + const NamedDataMap* named_data_map); /// Returns true if the Method was successfully initialized. inline bool initialized() const { @@ -339,7 +343,7 @@ class Method final { * the number of successfully-initialized entries so that ~Method doesn't try * to clean up uninitialized entries. */ - ET_NODISCARD Error parse_values(); + ET_NODISCARD Error parse_values(const NamedDataMap* named_data_map); ET_NODISCARD Error resolve_operator( int32_t op_index, diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp index b32832d7718..3c970fe169d 100644 --- a/runtime/executor/program.cpp +++ b/runtime/executor/program.cpp @@ -240,7 +240,8 @@ Result Program::get_method_name(size_t plan_index) const { Result Program::load_method( const char* method_name, MemoryManager* memory_manager, - EventTracer* event_tracer) const { + EventTracer* event_tracer, + const NamedDataMap* named_data_map) const { EXECUTORCH_SCOPE_PROF("Program::load_method"); internal::event_tracer_create_event_block(event_tracer, "Default"); internal::EventTracerProfileMethodScope event_tracer_scope = @@ -257,7 +258,8 @@ Result Program::load_method( if (!plan.ok()) { return plan.error(); } - return Method::load(plan.get(), this, memory_manager, event_tracer); + return Method::load( + plan.get(), this, memory_manager, event_tracer, named_data_map); } Result Program::method_meta(const char* method_name) const { diff --git a/runtime/executor/program.h b/runtime/executor/program.h index aac3d1cdfe9..7313b19d66d 100644 --- a/runtime/executor/program.h +++ b/runtime/executor/program.h @@ -129,13 +129,16 @@ class Program final { * execution of the loaded method. If `memory_manager.temp_allocator()` is * null, the runtime will allocate temp memory using `et_pal_allocate()`. * @param[in] event_tracer The event tracer to use for this method run. + * @param[in] named_data_map An optional map of {name, blob} used to resolve + * data that is external to the PTE, if any. * * @returns The loaded method on success, or an error on failure. */ Result load_method( const char* method_name, MemoryManager* memory_manager, - EventTracer* event_tracer = nullptr) const; + EventTracer* event_tracer = nullptr, + const NamedDataMap* named_data_map = nullptr) const; /** * Gathers metadata for the named method. diff --git a/runtime/executor/targets.bzl b/runtime/executor/targets.bzl index 158da5d1087..67163ed8789 100644 --- a/runtime/executor/targets.bzl +++ b/runtime/executor/targets.bzl @@ -79,6 +79,7 @@ def define_common_targets(): ":memory_manager", "//executorch/runtime/backend:interface", "//executorch/runtime/core:core", + "//executorch/runtime/core:named_data_map", "//executorch/runtime/core:evalue" + aten_suffix, "//executorch/runtime/core:event_tracer" + aten_suffix, "//executorch/runtime/core/exec_aten:lib" + aten_suffix, diff --git a/runtime/executor/tensor_parser.h b/runtime/executor/tensor_parser.h index 6b593afe7c0..2ffb473544d 100644 --- a/runtime/executor/tensor_parser.h +++ b/runtime/executor/tensor_parser.h @@ -7,6 +7,9 @@ */ #pragma once +// Disable -Wdeprecated-declarations, as some builds use 'Werror'. +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" #include #include @@ -21,7 +24,8 @@ namespace deserialization { ET_NODISCARD Result parseTensor( const Program* program, MemoryManager* memory_manager, - const executorch_flatbuffer::Tensor* s_tensor); + const executorch_flatbuffer::Tensor* s_tensor, + const NamedDataMap* named_data_map = nullptr); ET_NODISCARD Result> parseTensorList( const flatbuffers::Vector* tensor_indices, @@ -100,6 +104,8 @@ parseListOptionalType( * @param[in] program The Program to use for constant buffer data. * @param[in] nbytes The amount of memory to get from the allocator. * @param[in] allocator The source of memory for non-constant tensors. + * @param[in] named_data_map An optional map of {name, blob} used to resolve + * data that is external to the PTE, if any. * * @returns On success, the data pointer to use for the tensor. On failure, a * non-Ok Error. @@ -108,7 +114,8 @@ ET_NODISCARD Result getTensorDataPtr( const executorch_flatbuffer::Tensor* s_tensor, const Program* program, size_t nbytes, - HierarchicalAllocator* allocator); + HierarchicalAllocator* allocator, + const NamedDataMap* named_data_map = nullptr); } // namespace deserialization } // namespace runtime @@ -126,3 +133,4 @@ using ::executorch::runtime::deserialization::parseTensorList; } // namespace deserialization } // namespace executor } // namespace torch +#pragma GCC diagnostic pop diff --git a/runtime/executor/tensor_parser_aten.cpp b/runtime/executor/tensor_parser_aten.cpp index d92e7e6eb90..ab9af3d0399 100644 --- a/runtime/executor/tensor_parser_aten.cpp +++ b/runtime/executor/tensor_parser_aten.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -31,7 +32,8 @@ void deleteNothing(void*) {} Result parseTensor( const Program* program, MemoryManager* memory_manager, - const executorch_flatbuffer::Tensor* s_tensor) { + const executorch_flatbuffer::Tensor* s_tensor, + const NamedDataMap* named_data_map) { EXECUTORCH_SCOPE_PROF("TensorParser::parseTensor"); ET_CHECK_OR_RETURN_ERROR( @@ -102,7 +104,11 @@ Result parseTensor( } else { // Now that we know how big the tensor is, find and assign its memory. Result data_ptr = getTensorDataPtr( - s_tensor, program, tensor.nbytes(), memory_manager->planned_memory()); + s_tensor, + program, + tensor.nbytes(), + memory_manager->planned_memory(), + named_data_map); if (!data_ptr.ok()) { ET_LOG( Error, diff --git a/runtime/executor/tensor_parser_exec_aten.cpp b/runtime/executor/tensor_parser_exec_aten.cpp index 4feae452995..83310ff680c 100644 --- a/runtime/executor/tensor_parser_exec_aten.cpp +++ b/runtime/executor/tensor_parser_exec_aten.cpp @@ -19,6 +19,8 @@ namespace executorch { namespace runtime { namespace deserialization { +using executorch::aten::ScalarType; +using executorch::runtime::TensorLayout; // Provides access to private Program methods. class TensorParser final { public: @@ -113,7 +115,8 @@ ET_NODISCARD Result getTensorDataPtr( const executorch_flatbuffer::Tensor* s_tensor, const Program* program, size_t nbytes, - HierarchicalAllocator* allocator) { + HierarchicalAllocator* allocator, + const NamedDataMap* named_data_map) { auto data_buffer_idx = s_tensor->data_buffer_idx(); const executorch_flatbuffer::AllocationDetails* allocation_info = s_tensor->allocation_info(); @@ -131,9 +134,103 @@ ET_NODISCARD Result getTensorDataPtr( return err; } return planned_ptr; + } + + // External tensors. + else if ( + s_tensor->extra_tensor_info() != nullptr && + s_tensor->extra_tensor_info()->location() == + executorch_flatbuffer::TensorDataLocation::EXTERNAL) { + // Check that fqn is not null. + ET_CHECK_OR_RETURN_ERROR( + s_tensor->extra_tensor_info()->fully_qualified_name() != nullptr, + InvalidExternalData, + "Fully qualified name of external tensor is null"); + // Look up tensor in named data map. + Result tensor_layout_res = named_data_map->get_metadata( + s_tensor->extra_tensor_info()->fully_qualified_name()->c_str()); + if (!tensor_layout_res.ok()) { + return tensor_layout_res.error(); + } + const TensorLayout& tensor_layout = tensor_layout_res.get(); + + // Compatibility checking. + ET_CHECK_OR_RETURN_ERROR( + static_cast(s_tensor->scalar_type()) == + tensor_layout.scalar_type(), + InvalidExternalData, + "Scalar type mismatch. Expected %hhd, got %hhd.", + static_cast(s_tensor->scalar_type()), + static_cast(tensor_layout.scalar_type())); + ET_CHECK_OR_RETURN_ERROR( + nbytes == tensor_layout.nbytes(), + InvalidExternalData, + "Nbytes mismatch. Expected %zu, got %zu.", + nbytes, + tensor_layout.nbytes()); + int dim = s_tensor->sizes()->size(); + ET_CHECK_OR_RETURN_ERROR( + dim == tensor_layout.sizes().size(), + InvalidExternalData, + "Dim mismatch. Expected %d, got %zu.", + dim, + tensor_layout.sizes().size()); + for (int i = 0; i < dim; i++) { + ET_CHECK_OR_RETURN_ERROR( + s_tensor->sizes()->Get(i) == tensor_layout.sizes()[i], + InvalidExternalData, + "Sizes mismatch. Expected %d, got %d for size at index %d.", + s_tensor->sizes()->Get(i), + tensor_layout.sizes()[i], + i); + ET_CHECK_OR_RETURN_ERROR( + s_tensor->dim_order()->Get(i) == tensor_layout.dim_order()[i], + InvalidExternalData, + "Dim order mismatch. Expected %d, got %d for dim at index %d.", + s_tensor->dim_order()->Get(i), + tensor_layout.dim_order()[i], + i); + } + + // Constant value. + if (allocation_info == nullptr) { + Result data_res = named_data_map->get_data( + s_tensor->extra_tensor_info()->fully_qualified_name()->c_str()); + if (!data_res.ok()) { + return data_res.error(); + } + // The const_cast is 'ok' here because program and runtime should + // guarantee that this data is never modified. Temporary until runtime + // takes ownership of FreeableBuffers in TODO(T214294528). + return const_cast(data_res.get().data()); + } + + // Mutable value. + else { + // Call load_into. + auto planned_ptr = getMemPlannedPtr(allocation_info, nbytes, allocator); + if (!planned_ptr.ok()) { + return planned_ptr.error(); + } + auto size = named_data_map->load_data_into( + s_tensor->extra_tensor_info()->fully_qualified_name()->c_str(), + planned_ptr.get(), + nbytes); + if (size.error() != Error::Ok) { + return size.error(); + } + ET_CHECK_OR_RETURN_ERROR( + size.get() == nbytes, + InvalidExternalData, + "Expected to load %zu bytes, actually loaded %u bytes", + nbytes, + static_cast(size.get())); + return planned_ptr; + } + } - // Constant - } else if (data_buffer_idx > 0 && allocation_info == nullptr) { + // Constant, stored in PTE file. + else if (data_buffer_idx > 0 && allocation_info == nullptr) { auto const_data = program->get_constant_buffer_data(data_buffer_idx, nbytes); if (!const_data.ok()) { diff --git a/runtime/executor/tensor_parser_portable.cpp b/runtime/executor/tensor_parser_portable.cpp index 414961e0ff3..a53295470fc 100644 --- a/runtime/executor/tensor_parser_portable.cpp +++ b/runtime/executor/tensor_parser_portable.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -27,7 +28,8 @@ using torch::executor::TensorImpl; Result parseTensor( const Program* program, MemoryManager* memory_manager, - const executorch_flatbuffer::Tensor* s_tensor) { + const executorch_flatbuffer::Tensor* s_tensor, + const NamedDataMap* named_data_map) { EXECUTORCH_SCOPE_PROF("TensorParser::parseTensor"); auto method_allocator = memory_manager->method_allocator(); @@ -146,7 +148,8 @@ Result parseTensor( s_tensor, program, tensor_impl->nbytes(), - memory_manager->planned_memory()); + memory_manager->planned_memory(), + named_data_map); if (!data_ptr.ok()) { ET_LOG( Error, diff --git a/runtime/executor/test/CMakeLists.txt b/runtime/executor/test/CMakeLists.txt index 298a796b7b7..1eceb8f51d9 100644 --- a/runtime/executor/test/CMakeLists.txt +++ b/runtime/executor/test/CMakeLists.txt @@ -25,12 +25,18 @@ add_custom_command( "${CMAKE_BINARY_DIR}/ModuleLinear.pte" "${CMAKE_BINARY_DIR}/ModuleMultipleEntry.pte" "${CMAKE_BINARY_DIR}/ModuleSimpleTrain.pte" + "${CMAKE_BINARY_DIR}/ModuleLinearProgram.pte" + "${CMAKE_BINARY_DIR}/_default_external_constant.ptd" COMMAND - python3 -m test.models.export_program --modules + python -m test.models.export_program --modules "ModuleAdd,ModuleAddHalf,ModuleDynamicCatUnallocatedIO,ModuleIndex,ModuleLinear,ModuleMultipleEntry,ModuleSimpleTrain" --outdir "${CMAKE_BINARY_DIR}" 2> /dev/null COMMAND - python3 -m test.models.export_delegated_program --modules "ModuleAddMul" + python -m test.models.export_program --modules + "ModuleLinear" --external-constants + --outdir "${CMAKE_BINARY_DIR}" 2> /dev/null + COMMAND + python -m test.models.export_delegated_program --modules "ModuleAddMul" --backend_id "StubBackend" --outdir "${CMAKE_BINARY_DIR}" || true WORKING_DIRECTORY ${EXECUTORCH_ROOT} ) @@ -44,6 +50,8 @@ add_custom_target( "${CMAKE_BINARY_DIR}/ModuleLinear.pte" "${CMAKE_BINARY_DIR}/ModuleMultipleEntry.pte" "${CMAKE_BINARY_DIR}/ModuleSimpleTrain.pte" + "${CMAKE_BINARY_DIR}/ModuleLinearProgram.pte" + "${CMAKE_BINARY_DIR}/_default_external_constant.ptd" ) set(test_env @@ -53,6 +61,8 @@ set(test_env "ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH=${CMAKE_BINARY_DIR}/ModuleDynamicCatUnallocatedIO.pte" "ET_MODULE_INDEX_PATH=${CMAKE_BINARY_DIR}/ModuleIndex.pte" "ET_MODULE_LINEAR_PATH=${CMAKE_BINARY_DIR}/ModuleLinear.pte" + "ET_MODULE_LINEAR_PROGRAM_PATH=${CMAKE_BINARY_DIR}/ModuleLinearProgram.pte" + "ET_MODULE_LINEAR_DATA_PATH=${CMAKE_BINARY_DIR}/_default_external_constant.ptd" "ET_MODULE_MULTI_ENTRY_PATH=${CMAKE_BINARY_DIR}/ModuleMultipleEntry.pte" "ET_MODULE_SIMPLE_TRAIN_PATH=${CMAKE_BINARY_DIR}/ModuleSimpleTrain.pte" ) @@ -87,6 +97,7 @@ et_cxx_test( portable_ops_lib portable_kernels extension_data_loader + extension_flat_tensor extension_runner_util ) add_dependencies(method_test generated_pte_files) diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index 8ef4cfcb369..0c6a2db94b7 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include #include @@ -21,6 +22,7 @@ using namespace ::testing; using executorch::aten::ArrayRef; +using executorch::extension::FlatTensorDataMap; using executorch::extension::prepare_input_tensors; using executorch::runtime::Error; using executorch::runtime::EValue; @@ -52,6 +54,23 @@ class MethodTest : public ::testing::Test { {module_name, std::make_unique(std::move(program.get()))}); } + void load_data_map(const char* path, const char* module_name) { + // Create a loader for the serialized data map. + Result loader = FileDataLoader::from(path); + ASSERT_EQ(loader.error(), Error::Ok); + loaders_.insert( + {module_name, + std::make_unique(std::move(loader.get()))}); + + Result data_map = + FlatTensorDataMap::load(loaders_[module_name].get()); + EXPECT_EQ(data_map.error(), Error::Ok); + + data_maps_.insert( + {module_name, + std::make_unique(std::move(data_map.get()))}); + } + void SetUp() override { executorch::runtime::runtime_init(); @@ -63,6 +82,10 @@ class MethodTest : public ::testing::Test { load_program( std::getenv("DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH"), "linear_constant_buffer"); + + load_program( + std::getenv("ET_MODULE_LINEAR_PROGRAM_PATH"), "linear_program"); + load_data_map(std::getenv("ET_MODULE_LINEAR_DATA_PATH"), "linear_data"); } private: @@ -71,6 +94,8 @@ class MethodTest : public ::testing::Test { protected: std::unordered_map> programs_; + std::unordered_map> + data_maps_; }; TEST_F(MethodTest, MoveTest) { @@ -303,6 +328,17 @@ TEST_F(MethodTest, ConstantBufferTest) { ASSERT_EQ(err, Error::Ok); } +TEST_F(MethodTest, ProgramDataSeparationTest) { + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["linear_program"]->load_method( + "forward", &mmm.get(), nullptr, data_maps_["linear_data"].get()); + ASSERT_EQ(method.error(), Error::Ok); + + // Can execute the method. + Error err = method->execute(); + ASSERT_EQ(err, Error::Ok); +} + /* * TODO(T161163608): Test is disabled due to a resize bug in tensor_index_out of * the portable op lib diff --git a/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl index 72923e9868f..922fa17ba75 100644 --- a/runtime/executor/test/targets.bzl +++ b/runtime/executor/test/targets.bzl @@ -109,6 +109,8 @@ def define_common_targets(is_fbcode = False): "ET_MODULE_LINEAR_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleLinear.pte])", "ET_MODULE_MULTI_ENTRY_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleMultipleEntry.pte])", "ET_MODULE_SIMPLE_TRAIN_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleSimpleTrain.pte])", + "ET_MODULE_LINEAR_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.pte])", + "ET_MODULE_LINEAR_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleLinear.ptd])", } runtime.cxx_test( @@ -135,6 +137,7 @@ def define_common_targets(is_fbcode = False): ":managed_memory_manager", "//executorch/runtime/executor:program", "//executorch/extension/data_loader:file_data_loader", + "//executorch/extension/flat_tensor:flat_tensor_data_map", "//executorch/extension/runner_util:inputs", "//executorch/kernels/portable:generated_lib", ], diff --git a/test/models/export_program.py b/test/models/export_program.py index 7d5d5606830..ada80ff342f 100644 --- a/test/models/export_program.py +++ b/test/models/export_program.py @@ -264,6 +264,8 @@ def main() -> None: # Skip type promotion to keep the model in fp16. # Type promotion will convert to fp32. skip_type_promotion = True + if args.external_constants: + module_name = f"{module_name}Program" outfile = os.path.join(args.outdir, f"{module_name}.pte") prog = export_module_to_program( module_class, diff --git a/test/models/targets.bzl b/test/models/targets.bzl index e2cd0f264b3..51ad10f0f01 100644 --- a/test/models/targets.bzl +++ b/test/models/targets.bzl @@ -92,7 +92,7 @@ def define_common_targets(): ) runtime.genrule( - name = "exported_programs_with_data_separated", + name = "exported_program_and_data", cmd = "$(exe :export_program) --modules ModuleLinear --external-constants --outdir $OUT", outs = { "ModuleLinear.pte": ["ModuleLinear.pte"], diff --git a/test/run_oss_cpp_tests.sh b/test/run_oss_cpp_tests.sh index 2c8685ea5b7..e138693b49c 100755 --- a/test/run_oss_cpp_tests.sh +++ b/test/run_oss_cpp_tests.sh @@ -33,11 +33,11 @@ build_executorch() { -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON \ -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ -DEXECUTORCH_BUILD_DEVTOOLS=ON \ - -DEXECUTORCH_BUILD_VULKAN=$BUILD_VULKAN \ -DEXECUTORCH_BUILD_XNNPACK=ON \ -DEXECUTORCH_BUILD_TESTS=ON \ -Bcmake-out diff --git a/test/utils/OSSTestConfig.json b/test/utils/OSSTestConfig.json index ce8c9723d23..a43d285180d 100644 --- a/test/utils/OSSTestConfig.json +++ b/test/utils/OSSTestConfig.json @@ -18,6 +18,13 @@ "../print_evalue.cpp" ] }, + { + "directory": "extension/flat_tensor/test", + "sources": [ + "flat_tensor_data_map_test.cpp", + "flat_tensor_header_test.cpp" + ] + }, { "directory": "extension/kernel_util/test", "sources": [