[runtime][Hexagon] AOTExecutor implementation for C Codegen (#10311)

* Hexagon AOT tests work

* fix and address comments

python/tvm/contrib/hexagon/build.py

@@ -241,6 +241,25 @@ def get_graph_executor(
         graph_mod = self.load_module(module_name, session)
         return tvm.contrib.graph_executor.create(graph_json, graph_mod, session.device)
 
+    def get_aot_executor(self, module_name: Union[str, pathlib.Path], session: Session):
+        """Create a local AoTModule which consumes a remote libmod.
+
+        Parameters
+        ----------
+        module_name : str or pathlib.Path
+            Remote module filename. Same restrictions apply as in load_module().
+        session : Session
+            Remote session. The session must be established (via __enter__)
+            prior to calling this function.
+
+        Returns
+        -------
+        aot_module : AotModule
+            Runtime AOT module that can be used to execute.
+        """
+        aot_mod = self.load_module(module_name, session)
+        return tvm.runtime.executor.AotModule(aot_mod["default"](session.device))
+
 
 class HexagonLauncherAndroid(HexagonLauncherRPC):
     """Hexagon Launcher for Android."""
@@ -284,7 +303,7 @@ def _copy_to_remote(
 
     def _create_remote_directory(self, remote_path: Union[str, pathlib.Path]):
         """Abstract method implementation. See description in HexagonLauncherRPC."""
-        subprocess.check_call(self._adb_device_sub_cmd + ["shell", "mkdir", "-p", str(path)])
+        subprocess.check_call(self._adb_device_sub_cmd + ["shell", "mkdir", "-p", str(remote_path)])
 
     def _copy_binaries(self):
         """Upload Android server binaries."""

python/tvm/contrib/hexagon/hexagon.py

@@ -19,6 +19,9 @@
 
 import functools as ft
 import os
+import pathlib
+from typing import Union
+
 import tvm
 import tvm.ir
 import tvm.contrib.cc as cc
@@ -39,10 +42,18 @@
 #
 # Subsequent calls to 'link_shared' will use the newly registered linker.
 
-hexagon_toolchain_root = os.environ.get("HEXAGON_TOOLCHAIN") or ""  # pylint: disable=invalid-name
-hexagon_link_main = os.path.join(  # pylint: disable=invalid-name
-    hexagon_toolchain_root, "bin", "hexagon-link"
-)
+HEXAGON_TOOLCHAIN = os.environ.get("HEXAGON_TOOLCHAIN", default="")  # pylint: disable=invalid-name
+HEXAGON_SDK_PATH = os.environ.get("HEXAGON_SDK_PATH", default="")  # pylint: disable=invalid-name
+HEXAGON_LINK_MAIN = (
+    pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-link"
+)  # pylint: disable=invalid-name
+HEXAGON_CLANG_PLUS = (
+    pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-clang++"
+)  # pylint: disable=invalid-name
+HEXAGON_SDK_INCLUDE_DIRS = [  # pylint: disable=invalid-name
+    pathlib.Path(HEXAGON_SDK_PATH) / "incs",
+    pathlib.Path(HEXAGON_SDK_PATH) / "incs" / "stddef",
+]
 
 
 def register_linker(f):
@@ -51,9 +62,14 @@ def register_linker(f):
 
 
 @register_func("tvm.contrib.hexagon.hexagon.hexagon_link")
-def hexagon_link():
+def hexagon_link() -> str:
     """Return path to the Hexagon linker."""
-    return hexagon_link_main
+    return str(HEXAGON_LINK_MAIN)
+
+
+def hexagon_clang_plus() -> str:
+    """Return path to the Hexagon clang++."""
+    return str(HEXAGON_CLANG_PLUS)
 
 
 @register_func("tvm.contrib.hexagon.hexagon.link_shared")
@@ -101,12 +117,12 @@ def to_str(s):
             message += (
                 " Please verify the value of the HEXAGON_LINKER environment variable "
                 + '(currently set to "'
-                + hexagon_toolchain_root
+                + HEXAGON_TOOLCHAIN
                 + '").'
             )
         raise Exception(message)
 
-    libpath = os.path.join(hexagon_toolchain_root, "target", "hexagon", "lib", "v66", "G0")
+    libpath = os.path.join(HEXAGON_TOOLCHAIN, "target", "hexagon", "lib", "v66", "G0")
     cc.create_shared(
         so_name,
         objs,
@@ -248,3 +264,42 @@ def transform(func, mod, ctx):
 
 def ir_lower_vtcm_pass():
     return [(3, ir_lower_vtcm())]
+
+
+def create_aot_shared(so_name: Union[str, pathlib.Path], files, hexagon_arch: str, options=None):
+    """Export Hexagon AOT module."""
+    if not os.access(str(HEXAGON_CLANG_PLUS), os.X_OK):
+        raise Exception(
+            'The Clang++ "' + str(HEXAGON_CLANG_PLUS) + '" does not exist or is not executable.'
+        )
+    if not HEXAGON_TOOLCHAIN:
+        raise Exception(
+            " The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
+            + "HEXAGON_TOOLCHAIN in your environment."
+        )
+    if not HEXAGON_SDK_PATH:
+        raise Exception(
+            " The environment variable HEXAGON_SDK_PATH is unset. Please export "
+            + "HEXAGON_SDK_PATH in your environment."
+        )
+
+    tvm_dir = pathlib.Path(os.path.dirname(os.path.realpath(__file__))) / ".." / ".." / ".." / ".."
+    compute_arch = f"compute{hexagon_arch}"
+    compile_options = [
+        f"-I{tvm_dir / 'include'}",
+        f"-I{tvm_dir / '3rdparty' / 'dlpack' / 'include'}",
+        f"-I{tvm_dir / '3rdparty' / 'dmlc-core' / 'include'}",
+        f"-I{pathlib.Path(HEXAGON_SDK_PATH) / 'rtos' / 'qurt' / compute_arch / 'include'/ 'posix'}",
+        f"-I{pathlib.Path(HEXAGON_SDK_PATH) / 'rtos' / 'qurt' / compute_arch / 'include' / 'qurt'}",
+        f"-DDMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>",
+        f"-D_MACH_I32=int",
+    ]
+
+    # For debugging
+    for path in HEXAGON_SDK_INCLUDE_DIRS:
+        compile_options.append(f"-I{str(path)}")
+
+    cross_compile = cc.cross_compiler(compile_func=hexagon_clang_plus())
+    cross_compile.output_format = "o"
+    c_files = [str(file) for file in files]
+    cross_compile(str(so_name), c_files, options=compile_options)

python/tvm/contrib/hexagon/session.py

@@ -22,6 +22,7 @@
 from typing import Union
 from tvm import rpc as _rpc
 
+
 class Session:
     """Hexagon Device Session
 

src/relay/backend/aot_executor_codegen.cc

@@ -1167,6 +1167,9 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
       if (!target_host.defined() && it.second->kind->device_type == kDLCPU) {
         target_host = it.second;
       }
+      if (!target_host.defined() && it.second->kind->device_type == kDLHexagon) {
+        target_host = *(new Target("c"));
+      }
       ICHECK(dev_type);
       targets[static_cast<DLDeviceType>(dev_type->value)] = it.second;
     }

src/runtime/aot_executor/aot_executor.cc

@@ -45,8 +45,11 @@ AotExecutor::AotExecutor(tvm::runtime::Module module, const std::vector<Device>&
   DLDevice expected_device{kDLCPU, 0};
   ICHECK_EQ(devices_[0].device_id, expected_device.device_id)
       << "At this time, AOTExecutor supports only execution on kDLCPU 0";
-  ICHECK_EQ(devices_[0].device_type, expected_device.device_type)
-      << "At this time, AOTExecutor supports only execution on kDLCPU 0";
+  // TODO(tvm-team): Temporary hack since Hexagon is defined different than kDLCPU.
+  bool is_valid_device = (TVMDeviceExtType(devices_[0].device_type) == kDLHexagon) ||
+                         (DLDeviceType(devices_[0].device_type) == kDLCPU);
+  CHECK(is_valid_device)
+      << "At this time, AOTExecutor supports only execution on kDLCPU 0 or kDLHexagon 0";
 
   for (auto input : metadata_->inputs()) {
     // TODO(areusch): Encode device information in Metadata.

src/runtime/crt/microtvm_rpc_server/rpc_server.cc

@@ -31,7 +31,7 @@
 // NOTE: dmlc/base.h contains some declarations that are incompatible with some C embedded
 // toolchains. Just pull the bits we need for this file.
 #define DMLC_CMAKE_LITTLE_ENDIAN DMLC_IO_USE_LITTLE_ENDIAN
-#define DMLC_LITTLE_ENDIAN true
+#define DMLC_LITTLE_ENDIAN 1
 #include <tvm/runtime/c_runtime_api.h>
 #include <tvm/runtime/crt/crt.h>
 #include <tvm/runtime/crt/logging.h>

src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc

@@ -56,7 +56,7 @@ void HexagonDeviceAPIv2::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* r
 // DataSpace: static allocations for Hexagon
 void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, int ndim, const int64_t* shape,
                                          DLDataType dtype, Optional<String> mem_scope) {
-  CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
+  CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type;
 
   // Forcing contiguous allocation, for now
   // TODO(Straw): Enable discontiguous allocation after RFC 39 lands
@@ -84,7 +84,9 @@ void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, size_t nbytes, size_t align
 }
 
 void HexagonDeviceAPIv2::FreeDataSpace(Device dev, void* ptr) {
-  CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
+  bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) ||
+                         (DLDeviceType(dev.device_type) == kDLCPU);
+  CHECK(is_valid_device) << "dev.device_type: " << dev.device_type;
   auto* hexbuf = static_cast<HexagonBuffer*>(ptr);
   CHECK(hexbuf != nullptr);
   delete hexbuf;
@@ -97,7 +99,7 @@ struct HexagonWorkspacePool : public WorkspacePool {
 };
 
 void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType type_hint) {
-  CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
+  CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type;
   auto* hexbuf = static_cast<HexagonBuffer*>(
       dmlc::ThreadLocalStore<HexagonWorkspacePool>::Get()->AllocWorkspace(dev, size));
 
@@ -109,7 +111,7 @@ void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType typ
 }
 
 void HexagonDeviceAPIv2::FreeWorkspace(Device dev, void* data) {
-  CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
+  CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type;
   auto it = workspace_allocations_.find(data);
   CHECK(it != workspace_allocations_.end())
       << "Attempt made to free unknown or already freed workspace allocation";

src/runtime/minrpc/minrpc_server.h

@@ -28,7 +28,7 @@
 #ifndef TVM_RUNTIME_MINRPC_MINRPC_SERVER_H_
 #define TVM_RUNTIME_MINRPC_MINRPC_SERVER_H_
 
-#define DMLC_LITTLE_ENDIAN true
+#define DMLC_LITTLE_ENDIAN 1
 #include <string.h>
 #include <tvm/runtime/c_runtime_api.h>
 

src/runtime/vm/vm.cc

@@ -249,7 +249,7 @@ void VirtualMachine::SetOneInput(std::string func_name, const TVMArgValue& tag,
   const auto& vm_func = CheckAndGetVMFunction(func_name);
   size_t params_num = vm_func.params.size();
 
-  int inp_index;
+  int inp_index = 0;
   if (tag.type_code() == kTVMArgInt) {
     inp_index = tag;
   } else if (tag.type_code() == kTVMStr) {

tests/python/contrib/test_hexagon/test_cache_read_write.py

@@ -112,9 +112,9 @@ def test_cache_read_write(android_serial_number, tvm_tracker_host, tvm_tracker_p
         pytest.skip("Skip hardware test since ANDROID_SERIAL_NUMBER is not set.")
 
     rpc_info = {
-      "rpc_tracker_host" : tvm_tracker_host,
-      "rpc_tracker_port" : tvm_tracker_port,
-      "rpc_server_port" : 7070,
+        "rpc_tracker_host": tvm_tracker_host,
+        "rpc_tracker_port": tvm_tracker_port,
+        "rpc_server_port": 7070,
     }
     launcher = HexagonLauncher(serial_number=android_serial_number, rpc_info=rpc_info)
     launcher.upload(dso_binary_path, dso_binary)