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

#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.

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

ghstack-source-id: 279969988
Pull Request resolved: #10449

Author: zhenyanzhang

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/module.cpp

@@ -8,6 +8,9 @@
 
 #include <executorch/extension/module/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>
 #include <executorch/extension/data_loader/mmap_data_loader.h>
 #include <executorch/extension/flat_tensor/flat_tensor_data_map.h>
@@ -302,14 +305,57 @@ 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 {
+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::Error BundledModule::load_bundled_input(
+    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;
+
+  auto status = executorch::BUNDLED_PROGRAM_NAMESPACE::load_bundled_input(
+      *method, bundled_program_ptr_, testset_idx);
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      status,
+      "Bundled Program's load_bundled_input failed with status 0x%" PRIx32,
+      static_cast<uint32_t>(status));
+
+  return method->get_inputs(inputs.data(), inputs.size());
+}
+
+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

extension/module/module.h

@@ -493,19 +493,85 @@ class Module {
   std::unique_ptr<NamedDataMap> data_map_;
 
  protected:
+  std::unordered_map<std::string, MethodHolder> methods_;
+
+  friend class ExecuTorchJni;
+};
+
+/**
+ * A facade class for loading bundled programs and executing methods within
+ * them.
+ */
+class BundledModule : public Module {
+ public:
   /**
-   * Get a method by method name.
+   * Constructs an instance with the bundled program buffer pointer.
    *
-   * @param[in] method_name The name of the method to get.
+   * 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.
    *
-   * @returns A Result object containing either a pointer to the requested
-   *          method or an error to indicate failure.
+   * @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.
    */
-  ET_NODISCARD inline runtime::Result<Method*> get_method(
-      const std::string& method_name);
-  std::unordered_map<std::string, MethodHolder> methods_;
+  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);
 
-  friend class ExecuTorchJni;
+  BundledModule(const BundledModule&) = delete;
+  BundledModule& operator=(const BundledModule&) = delete;
+  BundledModule(Module&&) = delete;
+  BundledModule& operator=(BundledModule&&) = delete;
+
+  /**
+   * Execute a specific method with the input value at the given testset_idx
+   * from the program bundle.
+   *
+   * Before execution, this function loads the program and method with
+   * 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 A Result object containing either a vector of output values
+   *          from the method or an error to indicate failure.
+   */
+  ET_NODISCARD
+  runtime::Error load_bundled_input(
+      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

extension/module/targets.bzl

@@ -22,6 +22,9 @@ def define_common_targets():
                 "@EXECUTORCH_CLIENTS",
             ],
             deps = [
+                "//executorch/extension/data_loader:buffer_data_loader",
+                "//executorch/devtools/bundled_program:runtime",
+                "//executorch/devtools/bundled_program/schema:bundled_program_schema_fbs",
                 "//executorch/extension/memory_allocator:malloc_memory_allocator",
                 "//executorch/extension/data_loader:file_data_loader",
                 "//executorch/extension/data_loader:mmap_data_loader",