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))