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

devtools/bundled_program/bundled_program.cpp

@@ -260,9 +260,16 @@ ET_NODISCARD Error load_bundled_input(
   if (!method_test.ok()) {
     return method_test.error();
   }
-
+  auto test_cases = method_test.get()->test_cases();
+  ET_CHECK_OR_RETURN_ERROR(
+      testset_idx < test_cases->size(),
+      InvalidArgument,
+      "testset_idx %zu is out of range [0, %u]",
+      testset_idx,
+      test_cases->size());
   auto bundled_inputs =
-      method_test.get()->test_cases()->Get(testset_idx)->inputs();
+      test_cases->Get(static_cast<flatbuffers::uoffset_t>(testset_idx))
+          ->inputs();
 
   for (size_t input_idx = 0; input_idx < method.inputs_size(); input_idx++) {
     auto bundled_input = bundled_inputs->GetMutableObject(input_idx);
@@ -359,8 +366,16 @@ ET_NODISCARD Error verify_method_outputs(
     return method_test.error();
   }
 
+  auto test_cases = method_test.get()->test_cases();
+  ET_CHECK_OR_RETURN_ERROR(
+      testset_idx < test_cases->size(),
+      InvalidArgument,
+      "testset_idx %zu is out of range [0, %u]",
+      testset_idx,
+      test_cases->size());
   auto bundled_expected_outputs =
-      method_test.get()->test_cases()->Get(testset_idx)->expected_outputs();
+      test_cases->Get(static_cast<flatbuffers::uoffset_t>(testset_idx))
+          ->expected_outputs();
 
   if (bundled_expected_outputs->size() == 0) {
     // No bundled expected outputs, so we can't verify the method outputs.

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,112 @@
+/*
+ * 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>
+
+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::unique_ptr<BundledModule>> BundledModule::from_file(
+    const std::string& file_path,
+    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) {
+  auto data_loader_result = FileDataLoader::from(file_path.c_str());
+  if (!data_loader_result.ok()) {
+    return data_loader_result.error();
+  }
+
+  auto file_size_result = data_loader_result->size();
+  if (!file_size_result.ok()) {
+    return file_size_result.error();
+  }
+
+  size_t file_size = file_size_result.get();
+  auto file_data = std::make_unique<uint8_t[]>(file_size);
+  auto buffer_result =
+      data_loader_result->load_into(0, file_size, {}, file_data.get());
+  if (buffer_result != runtime::Error::Ok) {
+    return buffer_result;
+  }
+
+  // Pass ownership of the data to BundledModule
+  auto bm = std::make_unique<BundledModule>(
+      file_data.release(),
+      std::move(memory_allocator),
+      std::move(temp_allocator),
+      std::move(event_tracer),
+      std::move(data_map_loader));
+
+  bm->is_loaded_from_file_ = true;
+
+  return bm;
+}
+
+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;
+
+  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->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,123 @@
+/*
+ * 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() {
+    if (is_loaded_from_file_) {
+      delete[] static_cast<const uint8_t*>(bundled_program_ptr_);
+    }
+  }
+
+  /**
+   * Constructs an instance by loading a bundled program from a file with
+   * specified memory locking behavior.
+   *
+   * @param[in] file_path The path to the ExecuTorch bundled program file to
+   * load.
+   * @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.
+   */
+  ET_NODISCARD static runtime::Result<std::unique_ptr<BundledModule>> from_file(
+      const std::string& file_path,
+      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);
+
+  using Module::execute;
+
+  /**
+   * 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_;
+  bool is_loaded_from_file_ = false;
+};
+
+} // 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

@@ -491,16 +491,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_no_prim_ops" + 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,
+            ],
+        )