Hexagon compilation on MacOS system (#14308)

Short desc
This changes allow my to compile and tune models for hexagon directly from my macOS laptop without full switching to linux environment.

List of changes

Replace local linker call with call from docker container with Hexagon SDK. Yes, that is the only SDK tool used by TVM during compilation.
Enhanced search of ADB. Not only in PATH, but also in ANDROID_HOME, ANDROID_SDK_ROOT and default sdk installation directory. Mac OS doesn't allow to easily change default PATH env var for UI application launched from dock bar. So adb is not available for IDE by default.
Motivation
Some engineers would like to continue work with comfortable macOS environment even if they have to play with hexagon devices. At this moment there is no official Hexagon SDK for macOS system. Alternatives are next: fully switch to remote linux, use local linux virtual machine or try to introduce required hexagon SDK functionality for macOS. The last option is more preferable to me.

Signed-off-by: Alexander Peskov <[email protected]>

python/tvm/contrib/hexagon/build.py

@@ -29,6 +29,7 @@
 import random
 import string
 import subprocess
+import sys
 import tempfile
 from typing import Union
 
@@ -89,6 +90,67 @@ def _get_test_directory_name() -> str:
     return f"{date_str}-{random_str}"
 
 
+def _get_adb_path() -> str:
+    """Define path to adb
+
+    Order of search:
+      1. From PATH
+      2. From ANDROID_SDK_ROOT
+      3. From ANDROID_HOME
+      3. From default android sdk installation directory (platform specific)
+    """
+
+    def check_execution(exe_path):
+        try:
+            ret_code = subprocess.call(
+                [exe_path, "--version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+        except FileNotFoundError:
+            ret_code = -1
+
+        return ret_code == 0
+
+    # Check if adb available via PATH
+    if check_execution("adb"):
+        return "adb"
+
+    # Check if adb available via env vars or default directories
+    list_of_paths = [
+        os.environ.get("ANDROID_SDK_ROOT", default=""),
+        os.environ.get("ANDROID_HOME", default=""),
+    ]
+
+    if sys.platform == "darwin":
+        list_of_paths += [
+            os.path.join(pathlib.Path.home(), "Library", "Android", "sdk", "platform-tools")
+        ]
+    if sys.platform == "win32":
+        list_of_paths += [
+            os.path.join(
+                pathlib.Path.home(), "AppData", "Local", "Android", "sdk", "platform-tools"
+            )
+        ]
+    if sys.platform == "linux":
+        list_of_paths += [os.path.join(pathlib.Path.home(), "Android", "Sdk", "platform-tools")]
+
+    list_of_paths = [path for path in list_of_paths if path != ""]
+
+    found_path = None
+    for candidate_path in list_of_paths:
+        adb_path = os.path.join(candidate_path, "adb")
+        if os.path.isfile(adb_path) and check_execution(adb_path):
+            found_path = adb_path
+            break
+
+    if found_path is None:
+        raise RuntimeError(
+            "ADB was not found. It should be available via PATH, ANDROID_SDK_ROOT "
+            "or ANDROID_HOME env var."
+        )
+
+    return found_path
+
+
 class HexagonLauncherRPC(metaclass=abc.ABCMeta):
     """Base class for RPC-based launchers.
 
@@ -301,7 +363,8 @@ def __init__(
         assert self._serial_number != "", "Android serial number is not set."
 
         adb_socket = rpc_info["adb_server_socket"] if rpc_info["adb_server_socket"] else "tcp:5037"
-        self._adb_device_sub_cmd = ["adb", "-L", adb_socket, "-s", self._serial_number]
+        adb_exe = _get_adb_path()
+        self._adb_device_sub_cmd = [adb_exe, "-L", adb_socket, "-s", self._serial_number]
         self.forwarded_ports_ = []
         self._hexagon_debug = hexagon_debug
         self._clear_logcat = clear_logcat

python/tvm/contrib/hexagon/session.py

@@ -403,6 +403,7 @@ def _aot_executor_from_factory(
             elif target_type == "llvm":
                 module.export_library(
                     str(binary_path),
+                    fcompile=hexagon.create_shared,
                     cc=hexagon.hexagon_clang_plus(),
                 )
             else:

python/tvm/contrib/hexagon/tools.py

@@ -20,6 +20,9 @@
 import os
 import pathlib
 from typing import Union
+import sys
+import tarfile
+import io
 import numpy
 
 import tvm
@@ -43,6 +46,9 @@
 
 HEXAGON_TOOLCHAIN = os.environ.get("HEXAGON_TOOLCHAIN", default="")  # pylint: disable=invalid-name
 HEXAGON_SDK_ROOT = os.environ.get("HEXAGON_SDK_ROOT", default="")  # pylint: disable=invalid-name
+HEXAGON_SDK_DOCKER_IMAGE = os.environ.get(
+    "HEXAGON_SDK_DOCKER_IMAGE", default=""
+)  # pylint: disable=invalid-name
 HEXAGON_LINK_MAIN = (
     pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-link"
 )  # pylint: disable=invalid-name
@@ -145,6 +151,74 @@ def to_str(s):
     return 0
 
 
+def link_shared_macos(so_name, objs, extra_args=None):
+    """Link Hexagon shared library using docker container with proper tooling.
+
+    Parameters
+    ----------
+    so_name : str
+        Name of the shared library file.
+    objs : list[str,StringImm]
+    extra_args : dict (str->str) or Map<String,String>
+        Additional arguments:
+            'hex_arch' - Hexagon architecture, e.g. v66
+
+    Returns
+    -------
+    ret_val : int
+        This function returns 0 at the moment.
+    """
+    # The list of object files can be passed as built-in Python strings,
+    # or as tvm.tir.StringImm's.
+    def to_str(s):
+        if isinstance(s, tvm.tir.StringImm):
+            return s.value
+        assert isinstance(s, str), 'argument "' + str(s) + '" should be a string or StrImm'
+        return s
+
+    objs = [to_str(s) for s in objs]
+
+    if not extra_args:
+        extra_args = {}
+    hex_arch = extra_args.get("hex_arch") or "v66"
+
+    ses = ContainerSession(HEXAGON_SDK_DOCKER_IMAGE)
+
+    hexagon_sdk_tools_path = ses.get_env("HEXAGON_TOOLCHAIN")
+    libpath = os.path.join(hexagon_sdk_tools_path, "target", "hexagon", "lib", hex_arch, "G0")
+    linker = os.path.join(hexagon_sdk_tools_path, "bin", "hexagon-link")
+
+    # Copy input data to docker container
+    docker_objs = [ses.copy_to(obj) for obj in objs]
+    docker_so_name = ses.tmp_dir + "/" + os.path.basename(so_name)
+
+    link_cmd = [linker, "-shared", "-fPIC", "-o", docker_so_name]
+    link_cmd += docker_objs
+    link_cmd += [
+        "-Bdynamic",
+        "-export-dynamic",
+        "-L" + os.path.join(libpath, "pic"),
+        "-lgcc",
+    ]
+    ses.exec(link_cmd)
+
+    # Copy result back to host
+    ses.copy_from(docker_so_name, so_name)
+    return 0
+
+
+if sys.platform == "darwin":
+
+    def __create_shared_mac(so_name, objs, **kwargs):
+        return link_shared_macos(so_name, objs, kwargs)
+
+    create_shared = __create_shared_mac
+    register_func("tvm.contrib.hexagon.link_shared", f=link_shared_macos, override=True)
+else:  # Linux and Win32
+    create_shared = cc.create_shared
+    register_func("tvm.contrib.hexagon.link_shared", f=link_shared, override=True)
+
+
 def create_aot_shared(so_name: Union[str, pathlib.Path], files, hexagon_arch: str, options=None):
     """Export Hexagon AOT module."""
     options = options or []
@@ -242,3 +316,127 @@ def allocate_hexagon_array(
         arr.copyfrom(data.reshape(physical_shape))
 
     return arr._create_view(tensor_shape)
+
+
+class ContainerSession:
+    """Docker container session
+
+    Parameters
+    ----------
+    base_image_name : str
+        Docker image name to use. Empty string means to use default "tlcpack/ci-hexagon"
+        base image.
+    """
+
+    def __init__(self, base_image_name: str = ""):
+        self._client = None
+        self._container = None
+        self.tmp_dir = None
+
+        self._client = ContainerSession._get_docker_client()
+
+        if base_image_name == "":
+            base_image_name = ContainerSession._get_latest_ci_image(self._client)
+
+        self._container = ContainerSession._find_container_or_create(self._client, base_image_name)
+
+        exit_code, tmp_dir_b = self._container.exec_run("mktemp -d -t tvm-toolbox-XXXXXXXXXX")
+        assert exit_code == 0
+
+        self.tmp_dir = tmp_dir_b.decode("utf-8").rstrip()
+
+    def __del__(self):
+        self.close()
+
+    @staticmethod
+    def _get_latest_ci_image(client) -> str:
+        ci_images = client.images.list(name="tlcpack/ci-hexagon")
+        ci_images.sort(reverse=True, key=lambda img: img.tags[0])
+        return ci_images[0].tags[0]
+
+    @staticmethod
+    def _get_docker_client():
+        try:
+            # pylint: disable=import-outside-toplevel
+            from docker import from_env
+            from docker.errors import DockerException
+        except (ModuleNotFoundError, ImportError):
+            raise Exception("Docker SDK module is not installed. Please install it.")
+
+        try:
+            client = from_env()
+        except DockerException:
+            raise Exception(
+                "Docker server is not available. Please verify the docker is installed, "
+                "launched and available via command line ('dokcer ps' should works)."
+            )
+
+        return client
+
+    @staticmethod
+    def _find_container_or_create(client, image_name: str):
+        all_containers = client.containers.list(all=True)
+
+        filtered_containers = []
+        for container in all_containers:
+            tags: list = container.image.tags
+            img_name: str = tags[0]
+            if img_name.startswith(image_name) and container.name.startswith("tvm-hex-toolbox"):
+                filtered_containers.append(container)
+
+        if len(filtered_containers) == 0:
+            container = client.containers.run(
+                image=image_name, detach=True, tty=True, name="tvm-hex-toolbox"
+            )
+        else:
+            container = filtered_containers[0]
+
+        if container.status != "running":
+            container.start()
+
+        return container
+
+    def exec(self, cmd) -> str:
+        """Execute command inside docker container"""
+        exit_code, res = self._container.exec_run(cmd)
+        assert exit_code == 0
+        return res.decode("utf-8")
+
+    def get_env(self, key: str) -> str:
+        """Return env var value from docker container"""
+        res: str = self.exec(f"bash -c 'echo \"${key}\"'")
+        return res.rstrip(" \n")
+
+    def copy_to(self, host_file_path: str) -> str:
+        """Upload file to docker container"""
+        file_name = os.path.basename(host_file_path)
+
+        byte_stream = io.BytesIO()
+        with tarfile.open(fileobj=byte_stream, mode="w:gz") as tar:
+            tar.add(host_file_path, arcname=file_name)
+
+        self._container.put_archive(path=self.tmp_dir, data=byte_stream.getvalue())
+
+        return f"{self.tmp_dir}/{file_name}"
+
+    def copy_from(self, container_file_path: str, host_file_path: str):
+        """Download file from docker container"""
+        tar_bytes_gen, _ = self._container.get_archive(container_file_path)
+
+        # convert to bytes
+        tar_bytes = bytes()
+        for chunk in tar_bytes_gen:
+            tar_bytes += chunk
+
+        tar = tarfile.open(fileobj=io.BytesIO(initial_bytes=tar_bytes))
+        assert len(tar.getmembers()) == 1
+        tar_element_reader = tar.extractfile(tar.getmembers()[0])
+        with open(host_file_path, "wb") as host_file:
+            for chunk in tar_element_reader:
+                host_file.write(chunk)
+
+    def close(self):
+        """Close docker container session"""
+        if self.tmp_dir is not None:
+            exit_code, _ = self._container.exec_run(f"rm -rf {self.tmp_dir}")
+            assert exit_code == 0