add .ptd support to extension/module

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)

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/...",
         ],

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(

extension/module/module.cpp

@@ -10,6 +10,7 @@
 
 #include <executorch/extension/data_loader/file_data_loader.h>
 #include <executorch/extension/data_loader/mmap_data_loader.h>
+#include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
 #include <executorch/extension/memory_allocator/malloc_memory_allocator.h>
 #include <executorch/runtime/platform/runtime.h>
 
@@ -36,73 +37,125 @@
 namespace executorch {
 namespace extension {
 
+namespace {
+runtime::Result<std::unique_ptr<runtime::DataLoader>> load_file(
+    const std::string& file_path,
+    Module::LoadMode mode) {
+  std::unique_ptr<runtime::DataLoader> 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<runtime::EventTracer> event_tracer)
+    : file_path_(file_path),
+      load_mode_(load_mode),
+      memory_allocator_(std::make_unique<MallocMemoryAllocator>()),
+      temp_allocator_(std::make_unique<MallocMemoryAllocator>()),
+      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<runtime::EventTracer> event_tracer)
     : file_path_(file_path),
+      data_map_path_(data_map_path),
       load_mode_(load_mode),
       memory_allocator_(std::make_unique<MallocMemoryAllocator>()),
       temp_allocator_(std::make_unique<MallocMemoryAllocator>()),
-      event_tracer_(std::move(event_tracer)) {
+      event_tracer_(std::move(event_tracer)),
+      data_map_loader_(nullptr),
+      data_map_(nullptr) {
   runtime::runtime_init();
 }
 
 Module::Module(
     std::unique_ptr<runtime::DataLoader> data_loader,
     std::unique_ptr<runtime::MemoryAllocator> memory_allocator,
     std::unique_ptr<runtime::MemoryAllocator> temp_allocator,
-    std::unique_ptr<runtime::EventTracer> event_tracer)
+    std::unique_ptr<runtime::EventTracer> event_tracer,
+    std::unique_ptr<runtime::DataLoader> data_map_loader)
     : data_loader_(std::move(data_loader)),
       memory_allocator_(
           memory_allocator ? std::move(memory_allocator)
                            : std::make_unique<MallocMemoryAllocator>()),
       temp_allocator_(
           temp_allocator ? std::move(temp_allocator)
                          : std::make_unique<MallocMemoryAllocator>()),
-      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();
 }
 
 Module::Module(
     std::shared_ptr<runtime::Program> program,
     std::unique_ptr<runtime::MemoryAllocator> memory_allocator,
     std::unique_ptr<runtime::MemoryAllocator> temp_allocator,
-    std::unique_ptr<runtime::EventTracer> event_tracer)
+    std::unique_ptr<runtime::EventTracer> event_tracer,
+    std::unique_ptr<runtime::DataLoader> data_map_loader)
     : program_(std::move(program)),
       memory_allocator_(
           memory_allocator ? std::move(memory_allocator)
                            : std::make_unique<MallocMemoryAllocator>()),
       temp_allocator_(
           temp_allocator ? std::move(temp_allocator)
                          : std::make_unique<MallocMemoryAllocator>()),
-      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<runtime::Program>(
@@ -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));
   }

extension/module/module.h

@@ -51,6 +51,21 @@ class Module {
       const LoadMode load_mode = LoadMode::MmapUseMlock,
       std::unique_ptr<runtime::EventTracer> 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<runtime::EventTracer> 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<runtime::DataLoader> data_loader,
       std::unique_ptr<runtime::MemoryAllocator> memory_allocator = nullptr,
       std::unique_ptr<runtime::MemoryAllocator> temp_allocator = nullptr,
-      std::unique_ptr<runtime::EventTracer> event_tracer = nullptr);
+      std::unique_ptr<runtime::EventTracer> event_tracer = nullptr,
+      std::unique_ptr<runtime::DataLoader> 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<runtime::Program> program,
       std::unique_ptr<runtime::MemoryAllocator> memory_allocator = nullptr,
       std::unique_ptr<runtime::MemoryAllocator> temp_allocator = nullptr,
-      std::unique_ptr<runtime::EventTracer> event_tracer = nullptr);
+      std::unique_ptr<runtime::EventTracer> event_tracer = nullptr,
+      std::unique_ptr<runtime::DataLoader> data_map_loader = nullptr);
 
   Module(const Module&) = delete;
   Module& operator=(const Module&) = delete;
@@ -433,14 +452,16 @@ class Module {
     std::vector<runtime::EValue> inputs;
   };
 
- private:
   std::string file_path_;
+  std::string data_map_path_;
   LoadMode load_mode_{LoadMode::MmapUseMlock};
   std::shared_ptr<runtime::Program> program_;
   std::unique_ptr<runtime::DataLoader> data_loader_;
   std::unique_ptr<runtime::MemoryAllocator> memory_allocator_;
   std::unique_ptr<runtime::MemoryAllocator> temp_allocator_;
   std::unique_ptr<runtime::EventTracer> event_tracer_;
+  std::unique_ptr<runtime::DataLoader> data_map_loader_;
+  std::unique_ptr<runtime::NamedDataMap> data_map_;
 
  protected:
   std::unordered_map<std::string, MethodHolder> methods_;

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,

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);
+}

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
+    ```