add .ptd support to extension/module (pytorch#8421)

Summary:
Pull Request resolved: pytorch#8421

Title

Reviewed By: lucylq

Differential Revision: D69478424

Author: JacobSzwejbka

CMakeLists.txt

@@ -254,6 +254,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

@@ -125,11 +125,24 @@ runtime::Result<std::unordered_set<std::string>> Module::method_names() {
 
 runtime::Error Module::load_method(
     const std::string& method_name,
-    torch::executor::EventTracer* event_tracer) {
+    torch::executor::EventTracer* event_tracer,
+    std::unique_ptr<runtime::DataLoader> data_map_data_loader) {
   if (!is_method_loaded(method_name)) {
     ET_CHECK_OK_OR_RETURN_ERROR(load());
 
     MethodHolder method_holder;
+
+    // If we have a .ptd load it.
+    const runtime::NamedDataMap* named_data_map = nullptr;
+    if (data_map_data_loader != nullptr) {
+      auto data_map =
+          ET_UNWRAP_UNIQUE(executorch::extension::FlatTensorDataMap::load(
+              data_map_data_loader.get()));
+      method_holder.data_map_loader = std::move(data_map_data_loader);
+      method_holder.data_map = std::move(data_map);
+    }
+    named_data_map = method_holder.data_map.get();
+
     const auto method_metadata =
         ET_UNWRAP(program_->method_meta(method_name.c_str()));
     const auto planned_buffersCount =
@@ -155,10 +168,30 @@ 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(),
+        named_data_map));
     method_holder.inputs.resize(method_holder.method->inputs_size());
     methods_.emplace(method_name, std::move(method_holder));
   }
+  return runtime::Error::Ok;
+    }
+
+runtime::Error Module::load_method(
+    const std::string& method_name,
+    const std::string& data_map_path,
+    torch::executor::EventTracer* event_tracer) {
+  if (!is_method_loaded(method_name)) {
+    ET_CHECK_OK_OR_RETURN_ERROR(load());
+
+    // If we have a .ptd get a dataloader for it.
+    std::unique_ptr<extension::FileDataLoader> data_map_data_loader;
+    if (!data_map_path.empty()) {
+      data_map_data_loader =
+          ET_UNWRAP_UNIQUE(FileDataLoader::from(data_map_path.c_str()));
+    }
+
+    return load_method(method_name, event_tracer, std::move(data_map_data_loader));
+  }
   return runtime::Error::Ok;
 }
 

extension/module/module.h

@@ -14,6 +14,7 @@
 #include <unordered_set>
 #include <vector>
 
+#include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
 #include <executorch/runtime/executor/program.h>
 
 namespace executorch {
@@ -133,6 +134,8 @@ class Module {
    * needed. The loaded method is cached to reuse the next time it's executed.
    *
    * @param[in] method_name The name of the method to load.
+   * @param[in] data_map_path Path to a .ptd file containing weights
+   * for this method.
    * @param[in] event_tracer Per-method event tracer to profile/trace methods
    * individually. When not given, the event tracer passed to the Module
    * constructor is used. Otherwise, this per-method event tracer takes
@@ -143,8 +146,29 @@ class Module {
   ET_NODISCARD
   runtime::Error load_method(
       const std::string& method_name,
+      const std::string& data_map_path,
       torch::executor::EventTracer* event_tracer = nullptr);
 
+  /**
+   * Load a specific method from the program and set up memory management if
+   * needed. The loaded method is cached to reuse the next time it's executed.
+   *
+   * @param[in] method_name The name of the method to load.
+   * @param[in] event_tracer Per-method event tracer to profile/trace methods
+   * individually. When not given, the event tracer passed to the Module
+   * constructor is used. Otherwise, this per-method event tracer takes
+   * precedence.
+   * @param[in] data_map_data_loader Optional data loader for the .ptd file
+   * for this method.
+   *
+   * @returns An Error to indicate success or failure.
+   */
+  ET_NODISCARD
+  runtime::Error load_method(
+      const std::string& method_name,
+      torch::executor::EventTracer* event_tracer = nullptr,
+      std::unique_ptr<runtime::DataLoader> data_map_data_loader = nullptr);
+
   /**
    * Load the 'forward' method from the program and set up memory management if
    * needed. The loaded method is cached to reuse the next time it's executed.
@@ -155,8 +179,9 @@ class Module {
    * @returns An Error to indicate success or failure.
    */
   ET_NODISCARD inline runtime::Error load_forward(
-      torch::executor::EventTracer* event_tracer = nullptr) {
-    return load_method("forward", event_tracer);
+      torch::executor::EventTracer* event_tracer = nullptr,
+      std::unique_ptr<runtime::DataLoader> data_map_data_loader = nullptr) {
+    return load_method("forward", event_tracer, std::move(data_map_data_loader));
   }
 
   /**
@@ -430,10 +455,11 @@ class Module {
     std::unique_ptr<runtime::HierarchicalAllocator> planned_memory;
     std::unique_ptr<runtime::MemoryManager> memory_manager;
     std::unique_ptr<runtime::Method> method;
+    std::unique_ptr<runtime::DataLoader> data_map_loader;
+    std::unique_ptr<extension::FlatTensorDataMap> data_map;
     std::vector<runtime::EValue> inputs;
   };
 
- private:
   std::string file_path_;
   LoadMode load_mode_{LoadMode::MmapUseMlock};
   std::shared_ptr<runtime::Program> program_;

extension/module/targets.bzl

@@ -28,5 +28,6 @@ def define_common_targets():
             ],
             exported_deps = [
                 "//executorch/runtime/executor:program" + aten_suffix,
+                "//executorch/extension/flat_tensor:flat_tensor_data_map",
             ],
         )

extension/module/test/module_test.cpp

@@ -23,12 +23,19 @@ class ModuleTest : public ::testing::Test {
  protected:
   static void SetUpTestSuite() {
     model_path_ = std::getenv("RESOURCES_PATH") + std::string("/add.pte");
+    linear_path_ = std::getenv("RESOURCES_PATH") + std::string("/linear.pte");
+    linear_data_path_ =
+        std::getenv("RESOURCES_PATH") + std::string("/linear.ptd");
   }
 
   static std::string model_path_;
+  static std::string linear_path_;
+  static std::string linear_data_path_;
 };
 
 std::string ModuleTest::model_path_;
+std::string ModuleTest::linear_path_;
+std::string ModuleTest::linear_data_path_;
 
 TEST_F(ModuleTest, TestLoad) {
   Module module(model_path_);
@@ -435,3 +442,15 @@ TEST_F(ModuleTest, TestSetOutputInvalidType) {
 
   EXPECT_NE(module.set_output(EValue()), Error::Ok);
 }
+
+TEST_F(ModuleTest, TestPTD) {
+  Module module(linear_path_);
+
+  ASSERT_EQ(
+      module.load_method("forward", linear_data_path_), 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
+    ```