[ExecuTorch][#10375] Add extension.BundledModule to Wrap extension.Module with Bundled Program Logic

Pull Request resolved: #10449

#10375

# Context

This issue is a step of #9638.

In #9638, we want to have `extension.Module` as the single source of implementation in `pybindings`, which means that `pybindings.PyModule` should use `extension.Module` rather than its own `pybindings.Module`.

The issue is that `pybindings.PyModule` is dependent on the `method` getter from `pybindings.Module`, which `extension.Module` do not have. Since we don't want to expose `method` getter in `extension.Module`, we have to protect the getter, wrap the functions that is dependent on it and use the protected getter there, ultimately decouple `pybindings` from a `method` getter.

# Proposal

Now that we have a protected `method` getter, we can introduce a `extension.BundledModule`, a child class inheriting `extension.Module` which wraps up bundled program logic that is dependent on the `method` getter.
ghstack-source-id: 280825735

Differential Revision: [D73564125](https://our.internmc.facebook.com/intern/diff/D73564125/)

Author: zhenyan-zhang-meta

devtools/bundled_program/schema/targets.bzl

@@ -74,6 +74,7 @@ def define_common_targets():
         visibility = [
             "//executorch/devtools/bundled_program/...",
             "//executorch/extension/pybindings/...",
+            "//executorch/extension/module/...",
         ],
         exported_headers = {
             OUTPUT_BUNDLED_HEADER: ":{}[{}]".format(BUNDLED_GEN_RULE_NAME, OUTPUT_BUNDLED_HEADER),

extension/module/bundled_module.cpp

@@ -0,0 +1,97 @@
+/*
+ * 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.
+ */
+
+#include <executorch/extension/module/bundled_module.h>
+
+#include <executorch/devtools/bundled_program/bundled_program.h>
+#include <executorch/devtools/bundled_program/schema/bundled_program_schema_generated.h>
+#include <executorch/extension/data_loader/buffer_data_loader.h>
+#include <executorch/extension/data_loader/file_data_loader.h>
+
+/**
+ * Unwrap a Result to obtain its value (direct object, not a pointer).
+ * If the Result contains an error, propagate the error via trivial function
+ * return. The macro wraps the object into a unique_ptr.
+ *
+ * Note: A function using ET_UNWRAP_UNIQUE should itself return a Result or
+ * Error.
+ *
+ * @param[in] result__ Expression yielding the result to unwrap.
+ */
+#define ET_UNWRAP_UNIQUE(result__)                                     \
+  ({                                                                   \
+    auto et_result__ = (result__);                                     \
+    if (!et_result__.ok()) {                                           \
+      return et_result__.error();                                      \
+    }                                                                  \
+    std::make_unique<std::remove_reference_t<decltype(*et_result__)>>( \
+        std::move(*et_result__));                                      \
+  })
+
+namespace executorch {
+namespace extension {
+
+namespace {
+std::unique_ptr<BufferDataLoader> program_data_loader(
+    const void* bundled_program_ptr) {
+  auto bundled_program =
+      bundled_program_flatbuffer::GetBundledProgram(bundled_program_ptr);
+  // the program inside the bundled program
+  auto program = bundled_program->program();
+  return std::make_unique<BufferDataLoader>(program->data(), program->size());
+}
+} // namespace
+
+BundledModule::BundledModule(
+    const void* bundled_program_ptr,
+    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)
+    : Module(
+          program_data_loader(bundled_program_ptr),
+          std::move(memory_allocator),
+          std::move(temp_allocator),
+          std::move(event_tracer),
+          std::move(data_map_loader)),
+      bundled_program_ptr_(bundled_program_ptr) {}
+
+runtime::Result<std::vector<runtime::EValue>> BundledModule::execute(
+    const std::string& method_name,
+    const size_t testset_idx) {
+  ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));
+  auto& method = methods_.at(method_name).method;
+  auto& inputs = methods_.at(method_name).inputs;
+
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      executorch::BUNDLED_PROGRAM_NAMESPACE::load_bundled_input(
+          *method, bundled_program_ptr_, testset_idx));
+  ET_CHECK_OK_OR_RETURN_ERROR(method->get_inputs(inputs.data(), inputs.size()));
+  ET_CHECK_OK_OR_RETURN_ERROR(method->execute());
+
+  const auto outputs_size = method->outputs_size();
+  std::vector<runtime::EValue> outputs(outputs_size);
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      method->get_outputs(outputs.data(), outputs_size));
+
+  return outputs;
+}
+
+runtime::Error BundledModule::verify_method_outputs(
+    const std::string& method_name,
+    const size_t testset_idx,
+    double rtol,
+    double atol) {
+  ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));
+  auto& method = methods_.at(method_name).method;
+  return executorch::BUNDLED_PROGRAM_NAMESPACE::verify_method_outputs(
+      *method, bundled_program_ptr_, testset_idx, rtol, atol);
+}
+
+} // namespace extension
+} // namespace executorch

extension/module/bundled_module.h

@@ -0,0 +1,97 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <executorch/extension/module/module.h>
+
+namespace executorch {
+namespace extension {
+
+/**
+ * A facade class for loading bundled programs and executing methods within
+ * them.
+ */
+class BundledModule : public Module {
+ public:
+  /**
+   * Constructs an instance with the bundled program buffer pointer.
+   *
+   * This constructor reads the program from bundled program buffer to load the
+   * module with data loader. The bundled program pointer is preserved so that
+   * the portion outside of program is accessible.
+   *
+   * @param[in] bundled_program_ptr A DataLoader used for loading program data.
+   * @param[in] memory_allocator A MemoryAllocator used for memory management.
+   * @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 BundledModule(
+      const void* bundled_program_ptr,
+      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::DataLoader> data_map_loader = nullptr);
+
+  // Disallow copying
+  BundledModule(const BundledModule&) = delete;
+  BundledModule& operator=(const BundledModule&) = delete;
+  // Disallow copying
+  BundledModule(BundledModule&&) = delete;
+  BundledModule& operator=(BundledModule&&) = delete;
+  // Default destructor
+  ~BundledModule() = default;
+
+  /**
+   * Execute a specific method with the input value at the given `testset_idx`
+   * from the bundle to the method. Loads the program and method before
+   * executing if needed.
+   *
+   * This function is a wrapper of `load_bundled_input` in `bundled_program`.
+   *
+   * @param[in] method_name The name of the method to execute.
+   * @param[in] testset_idx The index of the input value to be passed to the
+   * method.
+   *
+   * @returns Return Error::Ok on a successful load, or the error happens during
+   * execution.
+   */
+  ET_NODISCARD
+  runtime::Result<std::vector<runtime::EValue>> execute(
+      const std::string& method_name,
+      const size_t testset_idx);
+
+  /**
+   * Verify the output of a specific method with the expected output from the
+   * program bundle at the given `testset_idx`.
+   *
+   * This function is a wrapper of `verify_method_outputs` in `bundled_program`.
+   *
+   * @param[in] method_name The name of the method to extract outputs from.
+   * @param[in] testset_idx  The index of expected output needs to be compared.
+   * @param[in] rtol Relative tolerance used for data comparsion.
+   * @param[in] atol Absolute tolerance used for data comparsion.
+   *
+   * @returns Return Error::Ok if two outputs match, or the error happens during
+   * execution.
+   */
+  ET_NODISCARD
+  runtime::Error verify_method_outputs(
+      const std::string& method_name,
+      const size_t testset_idx,
+      double rtol = 1e-5,
+      double atol = 1e-8);
+
+ private:
+  const void* bundled_program_ptr_;
+};
+
+} // namespace extension
+} // namespace executorch

extension/module/module.cpp

@@ -302,15 +302,5 @@ runtime::Error Module::set_output(
       output_tensor.mutable_data_ptr(), output_tensor.nbytes(), output_index);
 }
 
-ET_NODISCARD inline runtime::Result<Method*> Module::get_method(
-    const std::string& method_name) {
-  ET_CHECK_OR_RETURN_ERROR(
-      methods_.count(method_name) > 0,
-      InvalidArgument,
-      "no such method in program: %s",
-      method_name.c_str());
-  return methods_[method_name].method.get();
-}
-
 } // namespace extension
 } // namespace executorch

extension/module/module.h

@@ -493,16 +493,6 @@ class Module {
   std::unique_ptr<NamedDataMap> data_map_;
 
  protected:
-  /**
-   * Get a method by method name.
-   *
-   * @param[in] method_name The name of the method to get.
-   *
-   * @returns A Result object containing either a pointer to the requested
-   *          method or an error to indicate failure.
-   */
-  ET_NODISCARD inline runtime::Result<Method*> get_method(
-      const std::string& method_name);
   std::unordered_map<std::string, MethodHolder> methods_;
 
   friend class ExecuTorchJni;

extension/module/targets.bzl

@@ -31,3 +31,25 @@ def define_common_targets():
                 "//executorch/runtime/executor:program" + aten_suffix,
             ],
         )
+
+        runtime.cxx_library(
+            name = "bundled_module" + aten_suffix,
+            srcs = [
+                "bundled_module.cpp",
+            ],
+            exported_headers = [
+                "bundled_module.h",
+            ],
+            visibility = [
+                "@EXECUTORCH_CLIENTS",
+            ],
+            deps = [
+                "//executorch/extension/data_loader:buffer_data_loader",
+                "//executorch/extension/data_loader:file_data_loader",
+                "//executorch/devtools/bundled_program:runtime" + aten_suffix,
+                "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs",
+            ],
+            exported_deps = [
+                "//executorch/extension/module:module" + aten_suffix,
+            ],
+        )

extension/module/test/bundled_module_test.cpp

@@ -0,0 +1,54 @@
+/*
+ * 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.
+ */
+
+#include <executorch/extension/module/bundled_module.h>
+
+#include <gtest/gtest.h>
+
+#include <executorch/extension/data_loader/file_data_loader.h>
+
+using namespace ::executorch::extension;
+using namespace ::executorch::runtime;
+
+class BundledModuleTest : public ::testing::Test {
+ protected:
+  static void SetUpTestSuite() {
+    std::string resources_path;
+    if (const char* env = std::getenv("RESOURCES_PATH")) {
+      resources_path = env;
+    }
+    bpte_path_ = resources_path + "/bundled_program.bpte";
+  }
+
+  static inline std::string bpte_path_;
+};
+
+#include <fstream>
+
+std::vector<uint8_t> load_file_or_die(const char* path) {
+  std::ifstream file(path, std::ios::binary | std::ios::ate);
+  const size_t nbytes = file.tellg();
+  file.seekg(0, std::ios::beg);
+  auto file_data = std::vector<uint8_t>(nbytes);
+  ET_CHECK_MSG(
+      file.read(reinterpret_cast<char*>(file_data.data()), nbytes),
+      "Could not load contents of file '%s'",
+      path);
+  return file_data;
+}
+
+TEST_F(BundledModuleTest, TestExecute) {
+  std::vector<uint8_t> file_data = load_file_or_die(bpte_path_.c_str());
+  BundledModule bundled_module(reinterpret_cast<void*>(file_data.data()));
+
+  auto outputs = bundled_module.execute("forward", 0);
+  EXPECT_EQ(outputs.error(), Error::Ok);
+  auto status = bundled_module.verify_method_outputs(
+      "forward", 0, 1e-3, 1e-5);
+  EXPECT_EQ(status, Error::Ok);
+}

extension/module/test/resources/README.md

@@ -21,3 +21,9 @@
     ```
     python -m examples.portable.scripts.export --model_name="linear" -e
     ```
+
+### bundled_program.bpte
+
+   ```
+   python3 extension/module/test/resources/gen_bundled_program.py
+   ```