From 1b89c86b3560b9c3ede9ac3b61e6c896d153efde Mon Sep 17 00:00:00 2001 From: Anirudh Sundar Date: Fri, 12 Jul 2024 04:58:51 -0500 Subject: [PATCH] [Hexagon] Support RPC execution of existing shared lib This patch modifies the `get_executor_from_factory` for relax to support accepting a string that points to an already exported shared library. This allows us to run models that were already compiled through the RPC executor. --- python/tvm/contrib/hexagon/session.py | 33 +++++++++++++++++---------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py index 9f1166823423..50064e42ba08 100644 --- a/python/tvm/contrib/hexagon/session.py +++ b/python/tvm/contrib/hexagon/session.py @@ -287,14 +287,14 @@ def get_graph_debug_executor( ) def get_executor_from_factory( - self, module: Union[ExecutorFactoryModule, relax.Executable], hexagon_arch: str = "v68" + self, module: Union[ExecutorFactoryModule, relax.Executable, str], hexagon_arch: str = "v68" ): """Create a local GraphModule which consumes a remote libmod. Parameters ---------- - module : Union[ExecutorFactoryModule, relax.Executable] + module : Union[ExecutorFactoryModule, relax.Executable, str] The module to upload to the remote session and load. @@ -305,7 +305,7 @@ def get_executor_from_factory( return self._aot_executor_from_factory(module) if isinstance(module, GraphExecutorFactoryModule): return self._graph_executor_from_factory(module) - if isinstance(module, relax.Executable): + if isinstance(module, (relax.Executable, str)): return self._relax_vm_executable_executor(module, hexagon_arch=hexagon_arch) raise TypeError(f"Unsupported executor type: {type(module)}") @@ -358,7 +358,9 @@ def _graph_executor_from_factory( """ return self.get_graph_executor(module.get_graph_json(), module.get_lib()) - def _relax_vm_executable_executor(self, vm_exec: relax.Executable, hexagon_arch: str): + def _relax_vm_executable_executor( + self, vm_exec: Union[relax.Executable, str], hexagon_arch: str + ): """Create a local TVM module which consumes a remote vm executable. Paramters @@ -366,7 +368,7 @@ def _relax_vm_executable_executor(self, vm_exec: relax.Executable, hexagon_arch: vm_exec : relax.Executable The Relax VM Executable to upload to the remote and load. This will typically be the - output of `relax.build`. + output of `relax.build` or the path to an already built and exported shared library hexagon_arch : str The hexagon arch to be used Returns @@ -376,14 +378,21 @@ def _relax_vm_executable_executor(self, vm_exec: relax.Executable, hexagon_arch: """ assert self._rpc is not None, "Hexagon session must be started using __enter__ prior to use" - temp_dir = utils.tempdir() - path_exec = temp_dir.relpath("exec.so") + if isinstance(vm_exec, relax.Executable): + temp_dir = utils.tempdir() + path_exec = temp_dir.relpath("exec.so") - vm_exec.mod.export_library( - path_exec, - fcompile=hexagon.create_aot_shared, - hexagon_arch=hexagon_arch, - ) + vm_exec.mod.export_library( + path_exec, + fcompile=hexagon.create_aot_shared, + hexagon_arch=hexagon_arch, + ) + + path = self.upload(path_exec, "exec.so") + elif isinstance(vm_exec, str): + path_exec = vm_exec + else: + raise TypeError(f"Unsupported executor type: {type(vm_exec)}") path = self.upload(path_exec, "exec.so") return self._rpc.get_function("tvm.hexagon.load_module")(str(path))