From 46cd48268f4066481a0e877b204465e4011fc19b Mon Sep 17 00:00:00 2001
From: Mehrdad Hessar <mhessar@octoml.ai>
Date: Thu, 26 May 2022 13:13:56 -0700
Subject: [PATCH] host workspace added

---
 python/tvm/contrib/hexagon/build.py         | 26 ++++++++++++++++---
 python/tvm/contrib/hexagon/pytest_plugin.py | 28 +++++++++++++++++++--
 2 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/python/tvm/contrib/hexagon/build.py b/python/tvm/contrib/hexagon/build.py
index 43856253cb18..ea51ecb62a80 100644
--- a/python/tvm/contrib/hexagon/build.py
+++ b/python/tvm/contrib/hexagon/build.py
@@ -318,7 +318,11 @@ class HexagonLauncherAndroid(HexagonLauncherRPC):
     ]
 
     def __init__(
-        self, serial_number: str, rpc_info: dict, workspace: Union[str, pathlib.Path] = None
+        self,
+        serial_number: str,
+        rpc_info: dict,
+        workspace: Union[str, pathlib.Path] = None,
+        host_workspace: Union[str, pathlib.Path] = None,
     ):
         """Configure a new HexagonLauncherAndroid
 
@@ -335,6 +339,7 @@ def __init__(
         """
         if not rpc_info.get("workspace_base"):
             rpc_info["workspace_base"] = self.ANDROID_HEXAGON_TEST_BASE_DIR
+        self._host_workspace = host_workspace
         self._serial_number = serial_number
         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]
@@ -472,6 +477,15 @@ def _cleanup_directory(self):
         # Remove workspace directory on remote target
         subprocess.Popen(self._adb_device_sub_cmd + ["shell", f"rm -rf {self._workspace}"])
 
+    def _copy_artifacts_to_host(self):
+        # Copy generated artifacts from device to host
+        if self._host_workspace:
+            for file in ["tvm_rpc_android.farf", "tvm_rpc_android.log"]:
+                subprocess.Popen(
+                    self._adb_device_sub_cmd
+                    + ["pull", str(self._workspace / file), str(self._host_workspace / file)]
+                )
+
     def start_server(self):
         """Abstract method implementation. See description in HexagonLauncherRPC."""
         self._copy_binaries()
@@ -481,6 +495,7 @@ def stop_server(self):
         """Abstract method implementation. See description in HexagonLauncherRPC."""
         self._cleanup_port_forwarding()
         self._terminate_remote()
+        self._copy_artifacts_to_host()
         self._cleanup_directory()
 
 
@@ -593,7 +608,12 @@ def _is_port_in_use(port: int) -> bool:
 
 
 # pylint: disable=invalid-name
-def HexagonLauncher(serial_number: str, rpc_info: dict, workspace: Union[str, pathlib.Path] = None):
+def HexagonLauncher(
+    serial_number: str,
+    rpc_info: dict,
+    workspace: Union[str, pathlib.Path] = None,
+    host_workspace: Union[str, pathlib.Path] = None,
+):
     if serial_number == "simulator":
         return HexagonLauncherSimulator(rpc_info, workspace)
-    return HexagonLauncherAndroid(serial_number, rpc_info, workspace)
+    return HexagonLauncherAndroid(serial_number, rpc_info, workspace, host_workspace)
diff --git a/python/tvm/contrib/hexagon/pytest_plugin.py b/python/tvm/contrib/hexagon/pytest_plugin.py
index 2c62a0a0b569..00efe170e621 100644
--- a/python/tvm/contrib/hexagon/pytest_plugin.py
+++ b/python/tvm/contrib/hexagon/pytest_plugin.py
@@ -22,6 +22,8 @@
 import os
 import random
 from typing import Optional, Union
+import pathlib
+import datetime
 
 import pytest
 
@@ -29,6 +31,7 @@
 import tvm.rpc.tracker
 from tvm.contrib.hexagon.build import HexagonLauncher, HexagonLauncherRPC
 from tvm.contrib.hexagon.session import Session
+from tvm.contrib.utils import tempdir
 
 HEXAGON_TOOLCHAIN = "HEXAGON_TOOLCHAIN"
 TVM_TRACKER_HOST = "TVM_TRACKER_HOST"
@@ -156,9 +159,28 @@ def adb_server_socket() -> str:
     return os.getenv(ADB_SERVER_SOCKET, default="tcp:5037")
 
 
+@pytest.fixture
+def temp_dir():
+    workspace = (
+        pathlib.Path(os.path.abspath("."))
+        / f"hexagon_workspace"
+        / datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+    )
+    workspace_base = str(workspace)
+    number = 1
+    while workspace.exists():
+        workspace = pathlib.Path(workspace_base + f"-{number}")
+        number += 1
+    if not os.path.exists(workspace.parent):
+        os.makedirs(workspace.parent)
+
+    # return tempdir(custom_path=board_workspace, keep_for_debug=True)
+    return tempdir(custom_path=workspace)
+
+
 @tvm.testing.fixture
 def hexagon_launcher(
-    request, android_serial_number, rpc_server_port, adb_server_socket
+    request, android_serial_number, rpc_server_port, adb_server_socket, temp_dir
 ) -> HexagonLauncherRPC:
     """Initials and returns hexagon launcher if ANDROID_SERIAL_NUMBER is defined"""
     if android_serial_number is None:
@@ -176,7 +198,9 @@ def hexagon_launcher(
             "rpc_server_port": rpc_server_port,
             "adb_server_socket": adb_server_socket,
         }
-        launcher = HexagonLauncher(serial_number=android_serial_number, rpc_info=rpc_info)
+        launcher = HexagonLauncher(
+            serial_number=android_serial_number, rpc_info=rpc_info, host_workspace=temp_dir
+        )
         launcher.start_server()
         try:
             yield launcher