diff --git a/extension/pybindings/TARGETS b/extension/pybindings/TARGETS index ecf23e4658..17ccbb2477 100644 --- a/extension/pybindings/TARGETS +++ b/extension/pybindings/TARGETS @@ -67,6 +67,7 @@ runtime.python_library( srcs = ["portable_lib.py"], visibility = [ "//executorch/exir/...", + "//executorch/runtime/...", "@EXECUTORCH_CLIENTS", ], deps = [":_portable_lib"], diff --git a/extension/pybindings/portable_lib.py b/extension/pybindings/portable_lib.py index 1348cc7a3f..25624ad60c 100644 --- a/extension/pybindings/portable_lib.py +++ b/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 ) diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 6db42dbf15..62687c6b9a 100644 --- a/extension/pybindings/pybindings.cpp +++ b/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 method_names() const { + std::vector names; + for (const auto& method : methods_) { + names.push_back(method.first); + } + return names; + } + bool has_etdump() { return static_cast(event_tracer_); } @@ -905,6 +914,10 @@ struct PyModule final { return std::make_unique(module_, method.method_meta()); } + std::vector method_names() { + return module_->method_names(); + } + private: std::shared_ptr 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, diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index 5264985c59..818df1f760 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -58,6 +58,7 @@ class ExecuTorchModule: self, path: str, debug_buffer_path: Optional[str] = None ) -> None: ... def method_meta(self, method_name: str) -> MethodMeta: ... + def method_names(self) -> List[str]: ... @experimental("This API is experimental and subject to change without notice.") class BundledModule: diff --git a/extension/pybindings/test/TARGETS b/extension/pybindings/test/TARGETS index 335bd68ed1..73063deb65 100644 --- a/extension/pybindings/test/TARGETS +++ b/extension/pybindings/test/TARGETS @@ -10,7 +10,10 @@ runtime.python_library( srcs = [ "make_test.py", ], - visibility = ["//executorch/extension/pybindings/..."], + visibility = [ + "//executorch/extension/pybindings/...", + "//executorch/runtime/...", + ], deps = [ "//caffe2:torch", "//caffe2:torch_fx", diff --git a/extension/pybindings/test/make_test.py b/extension/pybindings/test/make_test.py index 73acdd4f45..90cb21d1a1 100644 --- a/extension/pybindings/test/make_test.py +++ b/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( + 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 ######### @@ -298,7 +304,6 @@ def test_constant_output_not_memory_planned(tester): tester.assertEqual(str(torch.ones(2, 2)), str(executorch_output[1])) def test_method_meta(tester) -> None: - # pyre-fixme[16]: Callable `make_test` has no attribute `wrapper`. exported_program, inputs = create_program(ModuleAdd()) # Use pybindings to load the program and query its metadata. @@ -345,7 +350,6 @@ def test_method_meta(tester) -> None: def test_bad_name(tester) -> None: # Create an ExecuTorch program from ModuleAdd. - # pyre-fixme[16]: Callable `make_test` has no attribute `wrapper`. exported_program, inputs = create_program(ModuleAdd()) # Use pybindings to load and execute the program. @@ -356,7 +360,6 @@ def test_bad_name(tester) -> None: def test_verification_config(tester) -> None: # Create an ExecuTorch program from ModuleAdd. - # pyre-fixme[16]: Callable `make_test` has no attribute `wrapper`. exported_program, inputs = create_program(ModuleAdd()) Verification = runtime.Verification diff --git a/pytest.ini b/pytest.ini index ecd58ea07e..166890bd25 100644 --- a/pytest.ini +++ b/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 diff --git a/runtime/TARGETS b/runtime/TARGETS new file mode 100644 index 0000000000..b9b0fc2c30 --- /dev/null +++ b/runtime/TARGETS @@ -0,0 +1,14 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +runtime.python_library( + name = "runtime", + srcs = ["__init__.py"], + deps = [ + "//executorch/extension/pybindings:portable_lib", + ], + visibility = [ + "//executorch/runtime/...", + ], +) diff --git a/runtime/__init__.py b/runtime/__init__.py new file mode 100644 index 0000000000..9f319995f7 --- /dev/null +++ b/runtime/__init__.py @@ -0,0 +1,199 @@ +# 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. + +"""Interface to the native C++ ExecuTorch runtime. + +Example usage: +.. code-block:: text + + from pathlib import Path + + import torch + from executorch.runtime import Verification, Runtime + + et_runtime: Runtime = Runtime.get() + program: Program = et_runtime.load_program( + Path("/tmp/program.pte"), + verification=Verification.Minimal, + ) + print("Program methods:", program.method_names) + forward: Method = program.load_method("forward") + + inputs = (torch.ones(2, 2), torch.ones(2, 2)) + outputs = forward.execute(inputs) + print(f"Ran forward({inputs})") + print(f" outputs: {outputs}") + +Example output: +.. code-block:: text + + Program methods: ('forward', 'forward2') + Ran forward((tensor([[1., 1.], + [1., 1.]]), tensor([[1., 1.], + [1., 1.]]))) + outputs: [tensor([[1., 1.], + [1., 1.]])] +""" + +import functools +from pathlib import Path +from types import ModuleType +from typing import Any, BinaryIO, Dict, Optional, Set, Union + +try: + from executorch.extension.pybindings.portable_lib import ( + ExecuTorchModule, + MethodMeta, + Verification, + ) +except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "Prebuilt /extension/pybindings/_portable_lib.so " + "is not found. Please reinstall ExecuTorch from pip." + ) from e + + +class Method: + """An ExecuTorch method, loaded from a Program. + This can be used to execute the method with inputs. + """ + + def __init__(self, method_name: str, module: ExecuTorchModule) -> None: + # TODO: This class should be pybind to the C++ counterpart instead of hosting ExecuTorchModule. + self._method_name = method_name + self._module = module + + def execute(self, inputs: Sequence[Any]) -> Sequence[Any]: + """Executes the method with the given inputs. + + Args: + inputs: The inputs to the method. + + Returns: + The outputs of the method. + """ + return self._module.run_method(self._method_name, inputs) + + @property + def metadata(self) -> MethodMeta: + """Gets the metadata for the method. + + Returns: + The metadata for the method. + """ + return self._module.method_meta(self._method_name) + + +class Program: + """An ExecuTorch program, loaded from binary PTE data. + + This can be used to load the methods/models defined by the program. + """ + + def __init__(self, module: ExecuTorchModule, data: Optional[bytes]) -> None: + # Hold the data so the program is not freed. + self._data = data + self._module = module + self._methods: Dict[str, Method] = {} + # ExecuTorchModule already pre-loads all Methods when created, so this + # doesn't do any extra work. TODO: Don't load a given Method until + # load_method() is called. Create a separate Method instance each time, + # to allow multiple independent instances of the same model. + for method_name in self._module.method_names(): + self._methods[method_name] = Method(method_name, self._module) + + @property + def method_names(self) -> Set[str]: + return set(self._methods.keys()) + + def load_method(self, name: str) -> Optional[Method]: + """Loads a method from the program. + + Args: + name: The name of the method to load. + + Returns: + The loaded method. + """ + return self._methods.get(name, None) + + +class OperatorRegistry: + """The registry of operators that are available to the runtime.""" + + def __init__(self, legacy_module: ModuleType) -> None: + # TODO: Expose the kernel callables to Python. + self._legacy_module = legacy_module + + @property + def operator_names(self) -> Set[str]: + """The names of all registered operators.""" + # pyre-fixme[7]: Expected `Set[str]` but got `Set[typing.Any]`. + return set(self._legacy_module._get_operator_names()) + + +class Runtime: + """An instance of the ExecuTorch runtime environment. + + This can be used to concurrently load and execute any number of ExecuTorch + programs and methods. + """ + + @staticmethod + @functools.lru_cache(maxsize=1) + def get() -> "Runtime": + """Gets the Runtime singleton.""" + import executorch.extension.pybindings.portable_lib as legacy_module + + return Runtime(legacy_module=legacy_module) + + def __init__(self, *, legacy_module: ModuleType) -> None: + # Public attributes. + self.operator_registry = OperatorRegistry(legacy_module) + # Private attributes. + self._legacy_module = legacy_module + + def load_program( + self, + data: Union[bytes, bytearray, BinaryIO, Path, str], + *, + verification: Verification = Verification.InternalConsistency, + ) -> Program: + """Loads an ExecuTorch program from a PTE binary. + + Args: + data: The binary program data to load; typically PTE data. + verification: level of program verification to perform. + + Returns: + The loaded program. + """ + if isinstance(data, (Path, str)): + m = self._legacy_module._load_for_executorch( + str(data), + enable_etdump=False, + debug_buffer_size=0, + program_verification=verification, + ) + return Program(m, data=None) + elif isinstance(data, BinaryIO): + data_bytes = data.read() + elif isinstance(data, bytearray): + data_bytes = bytes(data) + elif isinstance(data, bytes): + data_bytes = data + else: + raise TypeError( + f"Expected data to be bytes, bytearray, a path to a .pte file, or a file-like object, but got {type(data).__name__}." + ) + m = self._legacy_module._load_for_executorch_from_buffer( + data_bytes, + enable_etdump=False, + debug_buffer_size=0, + program_verification=verification, + ) + + return Program(m, data=data_bytes) diff --git a/runtime/test/TARGETS b/runtime/test/TARGETS new file mode 100644 index 0000000000..728de01b01 --- /dev/null +++ b/runtime/test/TARGETS @@ -0,0 +1,12 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +runtime.python_test( + name = "test_runtime", + srcs = ["test_runtime.py"], + deps = [ + "//executorch/extension/pybindings/test:make_test", + "//executorch/runtime:runtime", + ], +) diff --git a/runtime/test/test_runtime.py b/runtime/test/test_runtime.py new file mode 100644 index 0000000000..f0722f357e --- /dev/null +++ b/runtime/test/test_runtime.py @@ -0,0 +1,78 @@ +# 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. + +import tempfile +import unittest +from pathlib import Path + +import torch + +from executorch.extension.pybindings.test.make_test import ( + create_program, + ModuleAdd, + ModuleMulti, +) +from executorch.runtime import Runtime, Verification + + +class RuntimeTest(unittest.TestCase): + def test_smoke(self): + ep, inputs = create_program(ModuleAdd()) + runtime = Runtime.get() + # Demonstrate that get() returns a singleton. + runtime2 = Runtime.get() + self.assertTrue(runtime is runtime2) + program = runtime.load_program(ep.buffer, verification=Verification.Minimal) + method = program.load_method("forward") + outputs = method.execute(inputs) + self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1])) + + def test_module_with_multiple_method_names(self): + ep, inputs = create_program(ModuleMulti()) + runtime = Runtime.get() + + program = runtime.load_program(ep.buffer, verification=Verification.Minimal) + self.assertEqual(program.method_names, set({"forward", "forward2"})) + method = program.load_method("forward") + outputs = method.execute(inputs) + self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1])) + + method = program.load_method("forward2") + outputs = method.execute(inputs) + self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1] + 1)) + + def test_print_operator_names(self): + ep, inputs = create_program(ModuleAdd()) + runtime = Runtime.get() + + operator_names = runtime.operator_registry.operator_names + self.assertGreater(len(operator_names), 0) + + self.assertIn("aten::add.out", operator_names) + + def test_load_program_with_path(self): + ep, inputs = create_program(ModuleAdd()) + runtime = Runtime.get() + + def test_add(program): + method = program.load_method("forward") + outputs = method.execute(inputs) + self.assertTrue(torch.allclose(outputs[0], inputs[0] + inputs[1])) + + with tempfile.NamedTemporaryFile() as f: + f.write(ep.buffer) + f.flush() + # filename + program = runtime.load_program(f.name) + test_add(program) + # pathlib.Path + path = Path(f.name) + program = runtime.load_program(path) + test_add(program) + # BytesIO + with open(f.name, "rb") as f: + program = runtime.load_program(f.read()) + test_add(program) diff --git a/setup.py b/setup.py index 02ceed8744..ef027b451c 100644 --- a/setup.py +++ b/setup.py @@ -708,6 +708,7 @@ def get_ext_modules() -> List[Extension]: "executorch/schema": "schema", "executorch/devtools": "devtools", "executorch/devtools/bundled_program": "devtools/bundled_program", + "executorch/runtime": "runtime", "executorch/util": "util", # Note: This will install a top-level module called "serializer", # which seems too generic and might conflict with other pip packages.