From f043b0c9bd3788fc2b657385059b529a336a5fd3 Mon Sep 17 00:00:00 2001 From: lucylq Date: Thu, 2 Oct 2025 19:14:26 -0700 Subject: [PATCH] Use merged data map in module Differential Revision: [D83799869](https://our.internmc.facebook.com/intern/diff/D83799869/) [ghstack-poisoned] --- extension/module/module.cpp | 24 +++++++++++++----------- extension/module/targets.bzl | 1 + extension/module/test/module_test.cpp | 23 ++++++++++++++++------- extension/module/test/targets.bzl | 2 ++ 4 files changed, 32 insertions(+), 18 deletions(-) diff --git a/extension/module/module.cpp b/extension/module/module.cpp index 4b1c30ae6b5..e7c537b18a2 100644 --- a/extension/module/module.cpp +++ b/extension/module/module.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include /** @@ -155,10 +156,6 @@ runtime::Error Module::load(const Program::Verification verification) { data_loader_ = ET_UNWRAP(make_data_loader(file_path_, load_mode_)); } if (data_files_.size() > 0) { - ET_CHECK_OR_RETURN_ERROR( - data_files_.size() == 1, - NotImplemented, - "Multiple named data map paths are not supported yet."); for (const auto& data_file : data_files_) { data_map_loaders_.push_back( ET_UNWRAP(make_data_loader(data_file, load_mode_))); @@ -166,13 +163,18 @@ runtime::Error Module::load(const Program::Verification verification) { } if (data_map_loaders_.size() > 0) { - ET_CHECK_OR_RETURN_ERROR( - data_map_loaders_.size() == 1 && merged_data_map_ == nullptr, - NotImplemented, - "Multiple named data map loaders are not supported yet."); - // TODO(lfq): support multiple named data map loaders. - merged_data_map_ = - ET_UNWRAP_UNIQUE(FlatTensorDataMap::load(data_map_loaders_[0].get())); + for (auto i = 0; i < data_map_loaders_.size(); ++i) { + named_data_maps_.push_back(ET_UNWRAP_UNIQUE( + FlatTensorDataMap::load(data_map_loaders_[i].get()))); + } + + // Extract raw pointers from unique_ptrs to pass to MergedDataMap::load() + std::vector raw_data_maps; + raw_data_maps.reserve(named_data_maps_.size()); + for (const auto& data_map : named_data_maps_) { + raw_data_maps.push_back(data_map.get()); + } + merged_data_map_ = ET_UNWRAP_UNIQUE(MergedDataMap::load(raw_data_maps)); } auto program = diff --git a/extension/module/targets.bzl b/extension/module/targets.bzl index 3e449da5e14..0db909ce053 100644 --- a/extension/module/targets.bzl +++ b/extension/module/targets.bzl @@ -26,6 +26,7 @@ def define_common_targets(): "//executorch/extension/data_loader:file_data_loader", "//executorch/extension/data_loader:mmap_data_loader", "//executorch/extension/flat_tensor:flat_tensor_data_map" + aten_suffix, + "//executorch/extension/named_data_map:merged_data_map" + aten_suffix, ], exported_deps = [ "//executorch/runtime/executor:program_no_prim_ops" + aten_suffix, diff --git a/extension/module/test/module_test.cpp b/extension/module/test/module_test.cpp index 6f7e8a44558..1ff51eaf9e4 100644 --- a/extension/module/test/module_test.cpp +++ b/extension/module/test/module_test.cpp @@ -26,11 +26,15 @@ class ModuleTest : public ::testing::Test { model_path_ = std::getenv("ET_MODULE_ADD_PATH"); add_mul_path_ = std::getenv("ET_MODULE_ADD_MUL_PROGRAM_PATH"); add_mul_data_path_ = std::getenv("ET_MODULE_ADD_MUL_DATA_PATH"); + linear_path_ = std::getenv("ET_MODULE_LINEAR_PROGRAM_PATH"); + linear_data_path_ = std::getenv("ET_MODULE_LINEAR_DATA_PATH"); } static inline std::string model_path_; static inline std::string add_mul_path_; static inline std::string add_mul_data_path_; + static inline std::string linear_path_; + static inline std::string linear_data_path_; }; TEST_F(ModuleTest, TestLoad) { @@ -532,16 +536,21 @@ TEST_F(ModuleTest, TestPTD) { } TEST_F(ModuleTest, TestPTD_Multiple) { - std::vector data_files = {add_mul_data_path_}; - Module module(add_mul_path_, data_files); - - ASSERT_EQ(module.load_method("forward"), Error::Ok); - + std::vector data_files = {add_mul_data_path_, linear_data_path_}; + + // Create module with add mul. + Module module_add_mul(add_mul_path_, data_files); + ASSERT_EQ(module_add_mul.load_method("forward"), Error::Ok); auto tensor = make_tensor_ptr({2, 2}, {2.f, 3.f, 4.f, 2.f}); - ASSERT_EQ(module.forward(tensor).error(), Error::Ok); + ASSERT_EQ(module_add_mul.forward(tensor).error(), Error::Ok); // Confirm that the data_file is not std::move'd away. ASSERT_EQ(std::strcmp(data_files[0].c_str(), add_mul_data_path_.c_str()), 0); + ASSERT_EQ(std::strcmp(data_files[1].c_str(), linear_data_path_.c_str()), 0); - // TODO(lfq): add test when merge capability is supported. + // Create module with linear. + Module module_linear(linear_path_, data_files); + ASSERT_EQ(module_linear.load_method("forward"), Error::Ok); + auto tensor2 = make_tensor_ptr({3}, {2.f, 3.f, 4.f}); + ASSERT_EQ(module_linear.forward(tensor2).error(), Error::Ok); } diff --git a/extension/module/test/targets.bzl b/extension/module/test/targets.bzl index d1aa73f6789..da7f1cc91bd 100644 --- a/extension/module/test/targets.bzl +++ b/extension/module/test/targets.bzl @@ -19,6 +19,8 @@ def define_common_targets(is_fbcode=False): "ET_MODULE_ADD_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleAdd.pte])", "ET_MODULE_ADD_MUL_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleAddMul.pte])", "ET_MODULE_ADD_MUL_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleAddMul.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])", "ET_MODULE_SHARED_STATE": "$(location fbcode//executorch/test/models:exported_programs[ModuleSharedState.pte])", }