[Hexagon] Move aot/graph_executor interactions into launcher

python/tvm/contrib/hexagon/session.py

@@ -20,10 +20,16 @@
 import os
 import pathlib
 import tempfile
-from typing import Union
+from typing import Union, Optional
 
 import tvm
 from tvm import rpc as _rpc
+import tvm.contrib.hexagon as hexagon
+from tvm.relay.backend.executor_factory import (
+    ExecutorFactoryModule,
+    AOTExecutorFactoryModule,
+    GraphExecutorFactoryModule,
+)
 
 
 class Session:
@@ -136,3 +142,156 @@ def load_module(self, module: Union[str, pathlib.Path, tvm.runtime.Module]):
 
         assert isinstance(module, (str, pathlib.Path)), "Invalid path type:" + str(type(module))
         return self._rpc.get_function("tvm.hexagon.load_module")(str(module))
+
+    def get_graph_executor(
+        self,
+        module: Union[str, pathlib.Path, GraphExecutorFactoryModule],
+        graph_json: Optional[str] = None,
+    ):
+        """Create a local GraphModule which consumes a remote libmod.
+
+        The session must be established (via __enter__) prior to
+        calling this function.
+
+        Parameters
+        ----------
+
+        module : Union[str, pathlib.Path, GraphExecutorFactoryModule]
+
+            The module to load.  If `module` is a
+            `GraphExecutorFactoryModule`, it will be uploaded to the remote
+            session and loaded.  In this case, `graph_json` should be
+            None.
+
+            If `module` is a string or `pathlib.Path`, it must be
+            either a bare file name (without any path components), or
+            a full path in the remote system.  In this case, the file
+            must already have been uploaded to the remote and placed
+            in the remote workspace, and `graph_json` must be
+            provided.
+
+        graph_json : Optional[str]
+
+            The string with the graph JSON.  This should be passed if
+            and only if `module` is a file path to a binary that has
+            been uploaded to the remote.
+
+        Returns
+        -------
+        GraphModule :
+            Runtime graph module that can be used to execute the graph.
+
+        """
+
+        assert (
+            self.device is not None
+        ), "Hexagon session must be started using __enter__ prior to use"
+
+        if isinstance(module, GraphExecutorFactoryModule):
+            assert graph_json is None, "May not specify graph_json if full module is provided"
+            graph_json = module.get_graph_json()
+            graph_mod = self.load_module(module.get_lib())
+        else:
+            assert (
+                graph_json is not None
+            ), "Must specify graph_json if binary module has already been uploaded"
+            graph_mod = self.load_module(module)
+
+        return tvm.contrib.graph_executor.create(graph_json, graph_mod, self.device)
+
+    def get_aot_executor(
+        self,
+        module: Union[str, pathlib.Path, AOTExecutorFactoryModule],
+    ):
+        """Create a local GraphModule which consumes a remote libmod.
+
+        The session must be established (via __enter__) prior to
+        calling this function.
+
+        Parameters
+        ----------
+
+        module : Union[str, pathlib.Path, AOTExecutorFactoryModule]
+
+            The module to load.  If `module` is a
+            `AOTExecutorFactoryModule`, it will be uploaded to the remote
+            session and loaded.
+
+            If `module` is a string or `pathlib.Path`, it must be
+            either a bare file name (without any path components), or
+            a full path in the remote system.  In this case, the file
+            must already have been uploaded to the remote and placed
+            in the remote workspace, and `graph_json` must be
+            provided.
+
+        Returns
+        -------
+        GraphModule :
+            Runtime graph module that can be used to execute the graph.
+
+        """
+
+        assert (
+            self.device is not None
+        ), "Hexagon session must be started using __enter__ prior to use"
+
+        def _workaround_create_aot_shared():
+            # The C codegen uses TVM/RT functions directly. On Hexagon it should use
+            # functions pointers via __TVMxyz variables. This workaround makes the
+            # runtime symbols visible to the compiled shared library.
+            extra_link_flags = os.environ.get("HEXAGON_SHARED_LINK_FLAGS")
+            extra_options = str(extra_link_flags).split() if extra_link_flags else []
+            return lambda so_name, files, hexagon_arch, options: hexagon.create_aot_shared(
+                so_name, files, hexagon_arch, options=extra_options + options
+            )
+
+        if isinstance(module, ExecutorFactoryModule):
+            hexagon_arch = set(
+                target.mcpu.replace("hexagon", "")
+                for target in module.target.values()
+                if "hexagon" in target.keys
+            )
+            assert hexagon_arch, "No hexagon target architecture found"
+            assert (
+                len(hexagon_arch) == 1
+            ), f"Inconsistent hexagon architecture found, {hexagon_arch}"
+            hexagon_arch = hexagon_arch.pop()
+
+            with tempfile.TemporaryDirectory() as temp_dir:
+                temp_dir = pathlib.Path(temp_dir)
+                binary_name = "test_binary.so"
+                binary_path = temp_dir / binary_name
+
+                # Uncomment this once the workaround is not needed.
+                # module.export_library(
+                #     str(binary_path), fcompile=hexagon.create_aot_shared, hexagon_arch=hexagon_arch
+                # )
+                module.export_library(
+                    str(binary_path),
+                    fcompile=_workaround_create_aot_shared(),
+                    hexagon_arch=hexagon_arch,
+                )
+
+                self.upload(binary_path, binary_name)
+                module = binary_name
+
+        aot_mod = self.load_module(module)
+        return tvm.runtime.executor.AotModule(aot_mod["default"](self.device))
+
+    def get_executor(self, module: ExecutorFactoryModule):
+        """Create a local GraphModule which consumes a remote libmod.
+
+        Parameters
+        ----------
+
+        module : ExecutorFactoryModule
+
+            The module to upload to the remote
+            session and load.
+        """
+        if isinstance(module, AOTExecutorFactoryModule):
+            return self.get_aot_executor(module)
+        elif isinstance(module, GraphExecutorFactoryModule):
+            return self.get_graph_executor(module)
+        else:
+            raise TypeError(f"Unsupported executor type: {type(module)}")

tests/python/contrib/test_hexagon/test_launcher.py

@@ -144,7 +144,7 @@ def test_matmul(self, hexagon_session, M, N, K):
 
 
 @requires_hexagon_toolchain
-def test_graph_executor(hexagon_launcher, hexagon_session):
+def test_graph_executor(hexagon_session):
     dtype = "float32"
     data = relay.var("data", relay.TensorType((1, 64, 64, 3), dtype))
     weight = relay.var("weight", relay.TensorType((5, 5, 3, 8), dtype))
@@ -170,27 +170,18 @@ def test_graph_executor(hexagon_launcher, hexagon_session):
     params = {"weight": weight_in}
     inputs = {"data": data_in}
 
-    temp = utils.tempdir()
-    dso_binary = "test_binary.so"
-    dso_binary_path = temp.relpath(dso_binary)
-
     with tvm.transform.PassContext(opt_level=3):
         lowered = tvm.relay.build(
             relay_mod,
             tvm.target.Target(target_hexagon, host=target_hexagon),
             runtime=runtime,
             executor=executor,
         )
-        lowered.get_lib().save(dso_binary_path)
 
     if hexagon_session is None:
         pytest.skip(msg="Skip hardware test since ANDROID_SERIAL_NUMBER is not set.")
 
-    hexagon_launcher.upload(dso_binary_path, dso_binary)
-
-    graph_mod = hexagon_launcher.get_graph_executor(
-        lowered.get_graph_json(), dso_binary, hexagon_session
-    )
+    graph_mod = hexagon_session.get_executor(lowered)
     graph_mod.set_input(**params)
     graph_mod.run(**inputs)
     hexagon_output = graph_mod.get_output(0).numpy()
@@ -212,7 +203,7 @@ def test_graph_executor(hexagon_launcher, hexagon_session):
 
 
 @requires_hexagon_toolchain
-def test_graph_executor_multiple_conv2d(hexagon_launcher, hexagon_session):
+def test_graph_executor_multiple_conv2d(hexagon_session):
     dtype = "float32"
     input_shape = (1, 8, 8, 3)
     w1_shape = (5, 5, 3, 1)
@@ -246,24 +237,17 @@ def test_graph_executor_multiple_conv2d(hexagon_launcher, hexagon_session):
     runtime = Runtime("cpp")
     executor = Executor("graph")
 
-    temp = utils.tempdir()
-    dso_binary = "test_binary.so"
-    dso_binary_path = temp.relpath(dso_binary)
-
     with tvm.transform.PassContext(opt_level=3):
         lowered = tvm.relay.build(
             relay_mod,
             tvm.target.Target(target_hexagon, host=target_hexagon),
             runtime=runtime,
             executor=executor,
         )
-        lowered.get_lib().save(dso_binary_path)
 
     if hexagon_session is None:
         pytest.skip(msg="Skip hardware test since ANDROID_SERIAL_NUMBER is not set.")
 
-    hexagon_launcher.upload(dso_binary_path, dso_binary)
-
     weight1_data = np.random.rand(w1_shape[0], w1_shape[1], w1_shape[2], w1_shape[3]).astype(
         dtype=dtype
     )
@@ -277,9 +261,7 @@ def test_graph_executor_multiple_conv2d(hexagon_launcher, hexagon_session):
     params = {"weight1": weight1_data, "weight2": weight2_data}
     inputs = {"data": input_data}
 
-    graph_mod = hexagon_launcher.get_graph_executor(
-        lowered.get_graph_json(), dso_binary, hexagon_session
-    )
+    graph_mod = hexagon_session.get_executor(lowered)
     graph_mod.set_input(**params)
     graph_mod.run(**inputs)
     hexagon_output = graph_mod.get_output(0).numpy()
@@ -312,7 +294,7 @@ def _workaround_create_aot_shared():
 
 
 @requires_hexagon_toolchain
-def test_aot_executor(hexagon_launcher, hexagon_session):
+def test_aot_executor(hexagon_session):
     dtype = "float32"
     input_shape = (1, 128, 128, 3)
     w_shape = (5, 5, 3, 8)
@@ -332,9 +314,6 @@ def test_aot_executor(hexagon_launcher, hexagon_session):
     relay_mod = relay.transform.InferType()(relay_mod)
 
     target_hexagon = tvm.target.hexagon("v68")
-    temp = utils.tempdir()
-    dso_binary = "test_binary.so"
-    dso_binary_path = temp / dso_binary
 
     weight_data = np.random.rand(w_shape[0], w_shape[1], w_shape[2], w_shape[3]).astype(dtype=dtype)
     input_data = np.random.rand(
@@ -352,20 +331,11 @@ def test_aot_executor(hexagon_launcher, hexagon_session):
             runtime=Runtime("cpp"),
             executor=Executor("aot", {"unpacked-api": False, "interface-api": "c"}),
         )
-        # Uncomment this once the workaround is not needed.
-        # lowered.export_library(
-        #     dso_binary_path, fcompile=hexagon.create_aot_shared, hexagon_arch="v68"
-        # )
-        lowered.export_library(
-            dso_binary_path, fcompile=_workaround_create_aot_shared(), hexagon_arch="v68"
-        )
 
     if hexagon_session is None:
         pytest.skip(msg="Skip hardware test, ANDROID_SERIAL_NUMBER is not set.")
 
-    hexagon_launcher.upload(dso_binary_path, dso_binary)
-
-    aot_mod = hexagon_launcher.get_aot_executor(dso_binary, hexagon_session)
+    aot_mod = hexagon_session.get_executor(lowered)
     aot_mod.set_input(**inputs)
     aot_mod.run()
     hexagon_output = aot_mod.get_output(0).numpy()
@@ -388,7 +358,7 @@ def test_aot_executor(hexagon_launcher, hexagon_session):
 
 
 @requires_hexagon_toolchain
-def test_aot_executor_multiple_conv2d(hexagon_launcher, hexagon_session):
+def test_aot_executor_multiple_conv2d(hexagon_session):
     dtype = "float32"
     input_shape = (1, 8, 8, 3)
     w1_shape = (5, 5, 3, 1)
@@ -419,9 +389,6 @@ def test_aot_executor_multiple_conv2d(hexagon_launcher, hexagon_session):
     relay_mod = relay.transform.InferType()(relay_mod)
 
     target_hexagon = tvm.target.hexagon("v68")
-    temp = utils.tempdir()
-    dso_binary = "test_binary.so"
-    dso_binary_path = temp / dso_binary
 
     weight1_data = np.random.rand(w1_shape[0], w1_shape[1], w1_shape[2], w1_shape[3]).astype(
         dtype=dtype
@@ -444,20 +411,11 @@ def test_aot_executor_multiple_conv2d(hexagon_launcher, hexagon_session):
             runtime=Runtime("cpp"),
             executor=Executor("aot", {"unpacked-api": False, "interface-api": "c"}),
         )
-        # Uncomment this once the workaround is not needed.
-        # lowered.export_library(
-        #     dso_binary_path, fcompile=hexagon.create_aot_shared, hexagon_arch="v68"
-        # )
-        lowered.export_library(
-            dso_binary_path, fcompile=_workaround_create_aot_shared(), hexagon_arch="v68"
-        )
 
     if hexagon_session is None:
         pytest.skip(msg="Skip hardware test, ANDROID_SERIAL_NUMBER is not set.")
 
-    hexagon_launcher.upload(dso_binary_path, dso_binary)
-
-    aot_mod = hexagon_launcher.get_aot_executor(dso_binary, hexagon_session)
+    aot_mod = hexagon_session.get_executor(lowered)
     aot_mod.set_input(**inputs)
     aot_mod.run()
     hexagon_output = aot_mod.get_output(0).numpy()