diff --git a/CMakeLists.txt b/CMakeLists.txt index be0921a0b5a..01ad728c425 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -258,6 +258,11 @@ if(EXECUTORCH_BUILD_EXTENSION_TRAINING) set(EXECUTORCH_BUILD_EXTENSION_MODULE ON) endif() +if(EXECUTORCH_BUILD_EXTENSION_MODULE) + set(EXECUTORCH_BUILD_EXTENSION_DATA_LOADER ON) + set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON) +endif() + if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT) set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON) set(EXECUTORCH_BUILD_KERNELS_CUSTOM ON) diff --git a/extension/flat_tensor/targets.bzl b/extension/flat_tensor/targets.bzl index ed2adefc581..6f627492f24 100644 --- a/extension/flat_tensor/targets.bzl +++ b/extension/flat_tensor/targets.bzl @@ -9,13 +9,15 @@ def define_common_targets(): exported_headers = ["flat_tensor_data_map.h"], deps = [ "//executorch/extension/flat_tensor/serialize:generated_headers", - "//executorch/extension/flat_tensor/serialize:flat_tensor_header", "//executorch/runtime/core:core", "//executorch/runtime/core:evalue", "//executorch/runtime/core:named_data_map", "//executorch/runtime/core/exec_aten:lib", "//executorch/runtime/core/exec_aten/util:tensor_util", ], + exported_deps = [ + "//executorch/extension/flat_tensor/serialize:flat_tensor_header", + ], visibility = [ "//executorch/...", ], diff --git a/extension/module/CMakeLists.txt b/extension/module/CMakeLists.txt index 70441265c61..d144ce95356 100644 --- a/extension/module/CMakeLists.txt +++ b/extension/module/CMakeLists.txt @@ -27,7 +27,7 @@ if(CMAKE_TOOLCHAIN_IOS else() add_library(extension_module SHARED ${_extension_module__srcs}) endif() -target_link_libraries(extension_module PRIVATE executorch extension_data_loader) +target_link_libraries(extension_module PRIVATE executorch extension_data_loader extension_flat_tensor) target_include_directories(extension_module PUBLIC ${EXECUTORCH_ROOT}/..) target_compile_options( extension_module PUBLIC -Wno-deprecated-declarations -fPIC @@ -37,7 +37,7 @@ target_compile_options( # after cleaning up CMake targets. add_library(extension_module_static STATIC ${_extension_module__srcs}) target_link_libraries( - extension_module_static PRIVATE executorch extension_data_loader + extension_module_static PRIVATE executorch extension_data_loader extension_flat_tensor ) target_include_directories(extension_module_static PUBLIC ${EXECUTORCH_ROOT}/..) target_compile_options( diff --git a/extension/module/module.cpp b/extension/module/module.cpp index 99cc7e38bd6..aa750e2691e 100644 --- a/extension/module/module.cpp +++ b/extension/module/module.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -36,15 +37,59 @@ namespace executorch { namespace extension { +namespace { +runtime::Result> load_file( + const std::string& file_path, + Module::LoadMode mode) { + std::unique_ptr res = nullptr; + switch (mode) { + case Module::LoadMode::File: + res = ET_UNWRAP_UNIQUE(FileDataLoader::from(file_path.c_str())); + break; + case Module::LoadMode::Mmap: + res = ET_UNWRAP_UNIQUE(MmapDataLoader::from( + file_path.c_str(), MmapDataLoader::MlockConfig::NoMlock)); + break; + case Module::LoadMode::MmapUseMlock: + res = ET_UNWRAP_UNIQUE(MmapDataLoader::from(file_path.c_str())); + break; + case Module::LoadMode::MmapUseMlockIgnoreErrors: + res = ET_UNWRAP_UNIQUE(MmapDataLoader::from( + file_path.c_str(), + MmapDataLoader::MlockConfig::UseMlockIgnoreErrors)); + break; + } + return res; +} +} // namespace + +Module::Module( + const std::string& file_path, + const LoadMode load_mode, + std::unique_ptr event_tracer) + : file_path_(file_path), + load_mode_(load_mode), + memory_allocator_(std::make_unique()), + temp_allocator_(std::make_unique()), + event_tracer_(std::move(event_tracer)), + data_map_loader_(nullptr), + data_map_(nullptr) { + runtime::runtime_init(); +} + Module::Module( const std::string& file_path, + const std::string& data_map_path, const LoadMode load_mode, std::unique_ptr event_tracer) : file_path_(file_path), + data_map_path_(data_map_path), load_mode_(load_mode), memory_allocator_(std::make_unique()), temp_allocator_(std::make_unique()), - event_tracer_(std::move(event_tracer)) { + event_tracer_(std::move(event_tracer)), + data_map_loader_(nullptr), + data_map_(nullptr) { runtime::runtime_init(); } @@ -52,7 +97,8 @@ Module::Module( std::unique_ptr data_loader, std::unique_ptr memory_allocator, std::unique_ptr temp_allocator, - std::unique_ptr event_tracer) + std::unique_ptr event_tracer, + std::unique_ptr data_map_loader) : data_loader_(std::move(data_loader)), memory_allocator_( memory_allocator ? std::move(memory_allocator) @@ -60,7 +106,9 @@ Module::Module( temp_allocator_( temp_allocator ? std::move(temp_allocator) : std::make_unique()), - event_tracer_(std::move(event_tracer)) { + event_tracer_(std::move(event_tracer)), + data_map_loader_(std::move(data_map_loader)), + data_map_(nullptr) { runtime::runtime_init(); } @@ -68,7 +116,8 @@ Module::Module( std::shared_ptr program, std::unique_ptr memory_allocator, std::unique_ptr temp_allocator, - std::unique_ptr event_tracer) + std::unique_ptr event_tracer, + std::unique_ptr data_map_loader) : program_(std::move(program)), memory_allocator_( memory_allocator ? std::move(memory_allocator) @@ -76,33 +125,37 @@ Module::Module( temp_allocator_( temp_allocator ? std::move(temp_allocator) : std::make_unique()), - event_tracer_(std::move(event_tracer)) { + event_tracer_(std::move(event_tracer)), + data_map_loader_(std::move(data_map_loader)), + data_map_(nullptr) { runtime::runtime_init(); } runtime::Error Module::load(const runtime::Program::Verification verification) { if (!is_loaded()) { + // Load the program if (!data_loader_) { - switch (load_mode_) { - case LoadMode::File: - data_loader_ = - ET_UNWRAP_UNIQUE(FileDataLoader::from(file_path_.c_str())); - break; - case LoadMode::Mmap: - data_loader_ = ET_UNWRAP_UNIQUE(MmapDataLoader::from( - file_path_.c_str(), MmapDataLoader::MlockConfig::NoMlock)); - break; - case LoadMode::MmapUseMlock: - data_loader_ = - ET_UNWRAP_UNIQUE(MmapDataLoader::from(file_path_.c_str())); - break; - case LoadMode::MmapUseMlockIgnoreErrors: - data_loader_ = ET_UNWRAP_UNIQUE(MmapDataLoader::from( - file_path_.c_str(), - MmapDataLoader::MlockConfig::UseMlockIgnoreErrors)); - break; + auto res = load_file(file_path_, load_mode_); + if (!res.ok()) { + return res.error(); } - }; + data_loader_ = std::move(res.get()); + } + // If a .ptd path was given load it. + if (data_map_path_ != "") { + auto res = load_file(data_map_path_, load_mode_); + if (!res.ok()) { + return res.error(); + } + data_map_loader_ = std::move(res.get()); + } + // If we have a .ptd loader, then load the map. + if (data_map_loader_) { + data_map_ = + ET_UNWRAP_UNIQUE(FlatTensorDataMap::load(data_map_loader_.get())); + } + // else: either the map itself was provided or we have no data map, either + // way no work to do. auto program = ET_UNWRAP_UNIQUE( runtime::Program::load(data_loader_.get(), verification)); program_ = std::shared_ptr( @@ -130,6 +183,7 @@ runtime::Error Module::load_method( ET_CHECK_OK_OR_RETURN_ERROR(load()); MethodHolder method_holder; + const auto method_metadata = ET_UNWRAP(program_->method_meta(method_name.c_str())); const auto planned_buffersCount = @@ -155,7 +209,8 @@ runtime::Error Module::load_method( method_holder.method = ET_UNWRAP_UNIQUE(program_->load_method( method_name.c_str(), method_holder.memory_manager.get(), - event_tracer ? event_tracer : this->event_tracer())); + event_tracer ? event_tracer : this->event_tracer(), + data_map_.get())); method_holder.inputs.resize(method_holder.method->inputs_size()); methods_.emplace(method_name, std::move(method_holder)); } diff --git a/extension/module/module.h b/extension/module/module.h index 45ed38a7ff2..dc7c930d7c6 100644 --- a/extension/module/module.h +++ b/extension/module/module.h @@ -51,6 +51,21 @@ class Module { const LoadMode load_mode = LoadMode::MmapUseMlock, std::unique_ptr event_tracer = nullptr); + /** + * Constructs an instance by loading a program from a file with specified + * memory locking behavior. + * + * @param[in] file_path The path to the ExecuTorch program file to load. + * @param[in] data_map_path The path to a .ptd file + * @param[in] load_mode The loading mode to use. + * @param[in] event_tracer A EventTracer used for tracking and logging events. + */ + explicit Module( + const std::string& file_path, + const std::string& data_map_path, + const LoadMode load_mode = LoadMode::MmapUseMlock, + std::unique_ptr event_tracer = nullptr); + /** * Constructs an instance with the provided data loader and memory allocator. * @@ -59,12 +74,14 @@ class Module { * @param[in] temp_allocator A MemoryAllocator to use when allocating * temporary data during kernel or delegate execution. * @param[in] event_tracer A EventTracer used for tracking and logging events. + * @param[in] data_map_loader A DataLoader used for loading external weights. */ explicit Module( std::unique_ptr data_loader, std::unique_ptr memory_allocator = nullptr, std::unique_ptr temp_allocator = nullptr, - std::unique_ptr event_tracer = nullptr); + std::unique_ptr event_tracer = nullptr, + std::unique_ptr data_map_loader = nullptr); /** * Constructs an instance using an existing shared program. @@ -75,12 +92,14 @@ class Module { * @param[in] temp_allocator A MemoryAllocator to use when allocating * temporary data. * @param[in] event_tracer A EventTracer used for tracking and logging events. + * @param[in] data_map_loader A DataLoader used for loading external weights. */ explicit Module( std::shared_ptr program, std::unique_ptr memory_allocator = nullptr, std::unique_ptr temp_allocator = nullptr, - std::unique_ptr event_tracer = nullptr); + std::unique_ptr event_tracer = nullptr, + std::unique_ptr data_map_loader = nullptr); Module(const Module&) = delete; Module& operator=(const Module&) = delete; @@ -433,14 +452,16 @@ class Module { std::vector inputs; }; - private: std::string file_path_; + std::string data_map_path_; LoadMode load_mode_{LoadMode::MmapUseMlock}; std::shared_ptr program_; std::unique_ptr data_loader_; std::unique_ptr memory_allocator_; std::unique_ptr temp_allocator_; std::unique_ptr event_tracer_; + std::unique_ptr data_map_loader_; + std::unique_ptr data_map_; protected: std::unordered_map methods_; diff --git a/extension/module/targets.bzl b/extension/module/targets.bzl index 61251047dc8..4cbfa0ca0f5 100644 --- a/extension/module/targets.bzl +++ b/extension/module/targets.bzl @@ -25,6 +25,7 @@ def define_common_targets(): "//executorch/extension/memory_allocator:malloc_memory_allocator", "//executorch/extension/data_loader:file_data_loader", "//executorch/extension/data_loader:mmap_data_loader", + "//executorch/extension/flat_tensor:flat_tensor_data_map", ], exported_deps = [ "//executorch/runtime/executor:program" + aten_suffix, diff --git a/extension/module/test/module_test.cpp b/extension/module/test/module_test.cpp index 2dbb0fea936..ac7d4db13a9 100644 --- a/extension/module/test/module_test.cpp +++ b/extension/module/test/module_test.cpp @@ -22,14 +22,20 @@ using namespace ::executorch::runtime; class ModuleTest : public ::testing::Test { protected: static void SetUpTestSuite() { - model_path_ = std::getenv("RESOURCES_PATH") + std::string("/add.pte"); + std::string resources_path; + if (const char* env = std::getenv("RESOURCES_PATH")) { + resources_path = env; + } + model_path_ = resources_path + "/add.pte"; + linear_path_ = resources_path + "/linear.pte"; + linear_data_path_ = resources_path + "/linear.ptd"; } - static std::string model_path_; + static inline std::string model_path_; + static inline std::string linear_path_; + static inline std::string linear_data_path_; }; -std::string ModuleTest::model_path_; - TEST_F(ModuleTest, TestLoad) { Module module(model_path_); @@ -435,3 +441,14 @@ TEST_F(ModuleTest, TestSetOutputInvalidType) { EXPECT_NE(module.set_output(EValue()), Error::Ok); } + +TEST_F(ModuleTest, TestPTD) { + Module module(linear_path_, linear_data_path_); + + ASSERT_EQ(module.load_method("forward"), Error::Ok); + + auto tensor1 = + make_tensor_ptr({3, 3}, {2.f, 3.f, 4.f, 2.f, 3.f, 4.f, 2.f, 3.f, 4.f}); + + ASSERT_EQ(module.forward(tensor1).error(), Error::Ok); +} diff --git a/extension/module/test/resources/README.md b/extension/module/test/resources/README.md index e2b54633fae..ecbdd41c107 100644 --- a/extension/module/test/resources/README.md +++ b/extension/module/test/resources/README.md @@ -1,11 +1,23 @@ ## Resources -### model.pte +### add.pte, linear.pte, linear.ptd - Internally generated after D62209852, 2024-09-06 with: ``` buck2 run fbcode//executorch/examples/portable/scripts:export -- --model_name="add" ``` + + and + + ``` + buck2 run fbcode//executorch/examples/portable/scripts:export -- --model_name="linear" -examples + ``` - In OSS, the same file can be generated after [#5145](https://github.com/pytorch/executorch/pull/5145), 2024-09-06 with: ``` python -m examples.portable.scripts.export --model_name="add" ``` + + and + + ``` + python -m examples.portable.scripts.export --model_name="linear" -e + ``` diff --git a/extension/module/test/resources/linear.ptd b/extension/module/test/resources/linear.ptd new file mode 100644 index 00000000000..edab857bb3f Binary files /dev/null and b/extension/module/test/resources/linear.ptd differ diff --git a/extension/module/test/resources/linear.pte b/extension/module/test/resources/linear.pte new file mode 100644 index 00000000000..707815ad881 Binary files /dev/null and b/extension/module/test/resources/linear.pte differ