Bundled Program with **extension.module**

[Gasoonjia]

ET is updating the pybinding functions by migrating from pybinding.module into extension.module. BundledProgram has become a blocker for migration because:

• BundledProgram APIs take method module as input
• Extension.module does not expose method class to outside

This document focuses on solving the issue while having long-term benefits.

Targets:

We want to make some updates to replace pybinding.module with extension.module in pybinding functions and continue with the following requirements:

1. Supporting the current pybinding functionality of bundled program without BC issue
2. Supporting current Bundled Program users without BC issue
3. Don’t expose method class from extension.module

Proposals:

Create a new class called BundledModule, which extends the current Module class, specific for bundled program stuff. This way we wrap all method-related under BundledModule and not expose them:

class BundledModule : public Module {
   
   // create BundledModule from given bundled program file
   explict BundledModule(const std::string& bundled_program_path);

   // execute given method_name with its input_index input
   ET_NODISCARD inline runtime::Result<std::vector<runtime::Evalue>> execute(const std::string& method_name, const input_index);

   // compared the output of given method with output_index-th bundled expected output. Inside it invokes executorch::bundled_program::verify_method_outputs
   ET_NODISCARD inline runtime::Result<bool> verify_output_with_bundled(const std::string& method_name, const size_t bundled_output_index, std::vector<runtime::Evalue>> method_output; other_comparsion_configs);
   
  // extract Module class from BundledModule
  ET_NODISCARD inline runtime::Result<Module> extract_module();
}

Furthermore, in this scenario, we need to touch the method variable in the BundledModule, but the method variable is a private member of the Module class, so we should introduce an extra get function in Module’s protected scope to retrieve method module from method name while not exposing it to outside:

class Module {
    // Everything else keeps same

protected:
    // return method based on its name
    ET_NODISCARD inline runtime::Result<Method> get_method(const std::string& method_name);
}

In this way, we can make our final pybinding neat and concise:

py::list verify_result_with_bundled_expected_output(
    PyBundledModule& m, // here we can make PyBundledModule as a simple wrapper of BundledModule
    const std::string method_name,
    size_t testset_idx,
    double rtol = 1e-5,
    double atol = 1e-8) {

    output = m.execute(method_name, testset_idx);
    m.verify_output_with_bundled(method_name, testset_idx, output, rtol, atol);
}

@cccclai @tarun292 @byjlw

[zhenyan-zhang-meta]

Breaking down the discussion into issues:

• Add extension.BundledModule to Wrap extension.Module with Bundled Program Logic #10375
  • [ExecuTorch][#10375] Add extension.BundledModule to Wrap extension.Module with Bundled Program Logic #10449
• Extend PyBundledModule with extension.BundledModule #10447
  • [ExecuTorch][#10447] Extend PyBundledModule with extension.BundledModule #10450

We also got learnings from:

• Introduce APIs In Bundled Program to Not Take **Method** When Loading Input #10269
  • We were planning to break down part of bundled_program's load_bundled_input() to get_bundled_input(), and return the input for the usage of the existing Method::execute().
  • We realized that Method already have allocated memory for input of a method.
  • We can then reuse load_bundled_input(), then get the input written to that pre-allocated memory, so that eventually we can reuse Method::execute().
• Add Protected Method Getter in extension.Module #10364
  • [ExecuTorch][#10364] Add Protected Method Getter in extension.Module #10374
  • We realized that the method_ map is already protected, we can directly access it in child classes.