[pybind] New Runtime pybind API

Summary:
Based on this proposal: https://docs.google.com/document/d/10Q4-pt97inQQtFf-FjjwhMaDXXCfk1zGy6V6EkygNUY/edit#heading=h.fcrpnrtb6cud

Historically our pybinding APIs are not following the same C++ modeling
(Program, Method etc) and hence it's hard to use and easy to hit
footguns - for example, if we load the program and return it from a
python method, it goes out of the scope and releases the memory.

This effort is to create Pybind APIs that resembles C++ objects so it's
less confusing to the users.

Add the following python classes:
* `Runtime`: a singleton object hosting methods like `load_program`.
  Returns a `Program` object when calling `load_program`. Also exposes
  the operator registry
* `Program`: each pte file should have one `Program` object. Most
  important method is `load_method` which returns a `Method` object. It
  has a property `method_names` where we can inspect what methods are
  inside this .pte file.
* `Method`: one object per method name in a given `Program`. Exposes
  `execute` which takes in pytree flattened torch tensors as input and
  return pytree flattened output. It also exposes `MethodMeta` for users
  to inspect more information regarding input/output of this method.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: larryliu0820

extension/pybindings/portable_lib.py

@@ -45,6 +45,7 @@
     _reset_profile_results,  # noqa: F401
     BundledModule,  # noqa: F401
     ExecuTorchModule,  # noqa: F401
+    MethodMeta,  # noqa: F401
     Verification,  # noqa: F401
 )
 

extension/pybindings/pybindings.cpp

@@ -311,6 +311,15 @@ class Module final {
     return *methods_[method_name].get();
   }
 
+  /// Returns the names of all methods in the program.
+  std::vector<std::string> method_names() const {
+    std::vector<std::string> names;
+    for (const auto& method : methods_) {
+      names.push_back(method.first);
+    }
+    return names;
+  }
+
   bool has_etdump() {
     return static_cast<bool>(event_tracer_);
   }
@@ -905,6 +914,10 @@ struct PyModule final {
     return std::make_unique<PyMethodMeta>(module_, method.method_meta());
   }
 
+  std::vector<std::string> method_names() {
+    return module_->method_names();
+  }
+
  private:
   std::shared_ptr<Module> module_;
   // Need to keep-alive output storages until they can be compared in case of
@@ -1043,6 +1056,7 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
           &PyModule::method_meta,
           py::arg("method_name"),
           call_guard)
+      .def("method_names", &PyModule::method_names, call_guard)
       .def(
           "run_method",
           &PyModule::run_method,

extension/pybindings/test/make_test.py

@@ -16,118 +16,124 @@
 from torch.export import export
 
 
-def make_test(  # noqa: C901
-    tester: unittest.TestCase,
-    runtime: ModuleType,
-) -> Callable[[unittest.TestCase], None]:
-    """
-    Returns a function that operates as a test case within a unittest.TestCase class.
+class ModuleAdd(torch.nn.Module):
+    """The module to serialize and execute."""
 
-    Used to allow the test code for pybindings to be shared across different pybinding libs
-    which will all have different load functions. In this case each individual test case is a
-    subfunction of wrapper.
-    """
-    load_fn: Callable = runtime._load_for_executorch_from_buffer
+    def __init__(self):
+        super(ModuleAdd, self).__init__()
 
-    def wrapper(tester: unittest.TestCase) -> None:
-        class ModuleAdd(torch.nn.Module):
-            """The module to serialize and execute."""
+    def forward(self, x, y):
+        return x + y
 
-            def __init__(self):
-                super(ModuleAdd, self).__init__()
+    def get_methods_to_export(self):
+        return ("forward",)
 
-            def forward(self, x, y):
-                return x + y
+    def get_inputs(self):
+        return (torch.ones(2, 2), torch.ones(2, 2))
 
-            def get_methods_to_export(self):
-                return ("forward",)
 
-            def get_inputs(self):
-                return (torch.ones(2, 2), torch.ones(2, 2))
+class ModuleMulti(torch.nn.Module):
+    """The module to serialize and execute."""
 
-        class ModuleMulti(torch.nn.Module):
-            """The module to serialize and execute."""
+    def __init__(self):
+        super(ModuleMulti, self).__init__()
 
-            def __init__(self):
-                super(ModuleMulti, self).__init__()
+    def forward(self, x, y):
+        return x + y
 
-            def forward(self, x, y):
-                return x + y
+    def forward2(self, x, y):
+        return x + y + 1
 
-            def forward2(self, x, y):
-                return x + y + 1
+    def get_methods_to_export(self):
+        return ("forward", "forward2")
 
-            def get_methods_to_export(self):
-                return ("forward", "forward2")
+    def get_inputs(self):
+        return (torch.ones(2, 2), torch.ones(2, 2))
 
-            def get_inputs(self):
-                return (torch.ones(2, 2), torch.ones(2, 2))
 
-        class ModuleAddSingleInput(torch.nn.Module):
-            """The module to serialize and execute."""
+class ModuleAddSingleInput(torch.nn.Module):
+    """The module to serialize and execute."""
 
-            def __init__(self):
-                super(ModuleAddSingleInput, self).__init__()
+    def __init__(self):
+        super(ModuleAddSingleInput, self).__init__()
 
-            def forward(self, x):
-                return x + x
+    def forward(self, x):
+        return x + x
 
-            def get_methods_to_export(self):
-                return ("forward",)
+    def get_methods_to_export(self):
+        return ("forward",)
 
-            def get_inputs(self):
-                return (torch.ones(2, 2),)
+    def get_inputs(self):
+        return (torch.ones(2, 2),)
 
-        class ModuleAddConstReturn(torch.nn.Module):
-            """The module to serialize and execute."""
 
-            def __init__(self):
-                super(ModuleAddConstReturn, self).__init__()
-                self.state = torch.ones(2, 2)
+class ModuleAddConstReturn(torch.nn.Module):
+    """The module to serialize and execute."""
 
-            def forward(self, x):
-                return x + self.state, self.state
+    def __init__(self):
+        super(ModuleAddConstReturn, self).__init__()
+        self.state = torch.ones(2, 2)
 
-            def get_methods_to_export(self):
-                return ("forward",)
+    def forward(self, x):
+        return x + self.state, self.state
 
-            def get_inputs(self):
-                return (torch.ones(2, 2),)
+    def get_methods_to_export(self):
+        return ("forward",)
 
-        def create_program(
-            eager_module: torch.nn.Module,
-            et_config: Optional[ExecutorchBackendConfig] = None,
-        ) -> Tuple[ExecutorchProgramManager, Tuple[Any, ...]]:
-            """Returns an executorch program based on ModuleAdd, along with inputs."""
+    def get_inputs(self):
+        return (torch.ones(2, 2),)
 
-            # Trace the test module and create a serialized ExecuTorch program.
-            inputs = eager_module.get_inputs()
-            input_map = {}
-            for method in eager_module.get_methods_to_export():
-                input_map[method] = inputs
 
-            class WrapperModule(torch.nn.Module):
-                def __init__(self, fn):
-                    super().__init__()
-                    self.fn = fn
+def create_program(
+    eager_module: torch.nn.Module,
+    et_config: Optional[ExecutorchBackendConfig] = None,
+) -> Tuple[ExecutorchProgramManager, Tuple[Any, ...]]:
+    """Returns an executorch program based on ModuleAdd, along with inputs."""
 
-                def forward(self, *args, **kwargs):
-                    return self.fn(*args, **kwargs)
+    # Trace the test module and create a serialized ExecuTorch program.
+    inputs = eager_module.get_inputs()
+    input_map = {}
+    for method in eager_module.get_methods_to_export():
+        input_map[method] = inputs
 
-            exported_methods = {}
-            # These cleanup passes are required to convert the `add` op to its out
-            # variant, along with some other transformations.
-            for method_name, method_input in input_map.items():
-                wrapped_mod = WrapperModule(  # pyre-ignore[16]
-                    getattr(eager_module, method_name)
-                )
-                exported_methods[method_name] = export(wrapped_mod, method_input)
+    class WrapperModule(torch.nn.Module):
+        def __init__(self, fn):
+            super().__init__()
+            self.fn = fn
+
+        def forward(self, *args, **kwargs):
+            return self.fn(*args, **kwargs)
+
+    exported_methods = {}
+    # These cleanup passes are required to convert the `add` op to its out
+    # variant, along with some other transformations.
+    for method_name, method_input in input_map.items():
+        wrapped_mod = WrapperModule(  # pyre-ignore[16]
+            getattr(eager_module, method_name)
+        )
+        exported_methods[method_name] = export(wrapped_mod, method_input)
+
+    exec_prog = to_edge(exported_methods).to_executorch(config=et_config)
 
-            exec_prog = to_edge(exported_methods).to_executorch(config=et_config)
+    # Create the ExecuTorch program from the graph.
+    exec_prog.dump_executorch_program(verbose=True)
+    return (exec_prog, inputs)
 
-            # Create the ExecuTorch program from the graph.
-            exec_prog.dump_executorch_program(verbose=True)
-            return (exec_prog, inputs)
+
+def make_test(  # noqa: C901
+    tester: unittest.TestCase,
+    runtime: ModuleType,
+) -> Callable[[unittest.TestCase], None]:
+    """
+    Returns a function that operates as a test case within a unittest.TestCase class.
+
+    Used to allow the test code for pybindings to be shared across different pybinding libs
+    which will all have different load functions. In this case each individual test case is a
+    subfunction of wrapper.
+    """
+    load_fn: Callable = runtime._load_for_executorch_from_buffer
+
+    def wrapper(tester: unittest.TestCase) -> None:
 
         ######### TEST CASES #########
 

pytest.ini

@@ -38,6 +38,8 @@ addopts =
     backends/xnnpack/test
     # extension/
     extension/pybindings/test
+    # Runtime
+    runtime
     # test
     test/end2end/test_end2end.py
     --ignore=backends/xnnpack/test/ops/linear.py