Skip to content

Commit

Permalink
[Hexagon] Support RPC execution of existing shared lib (#17162)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
quic-sanirudh authored Jul 16, 2024
1 parent 70c5308 commit 51d7c5e
Showing 1 changed file with 21 additions and 12 deletions.
33 changes: 21 additions & 12 deletions python/tvm/contrib/hexagon/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)}")
Expand Down Expand Up @@ -358,15 +358,17 @@ 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
---------
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
Expand All @@ -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))
Expand Down

0 comments on commit 51d7c5e

Please sign in to comment.