Address comments

python/tvm/contrib/hexagon/build.py

@@ -182,9 +182,14 @@ def upload(self, local_path: Union[str, pathlib.Path], remote_filename: str):
         assert self._workspace
         self._copy_to_remote(local_path, os.path.join(str(self._workspace), remote_filename))
 
-    def start_session(self, name="hexagon-rpc") -> Session:
+    def start_session(self, session_name: str = "hexagon-rpc") -> Session:
         """Connect to the RPC server.
 
+        Parameters
+        ----------
+        session_name : str
+            RPC session name.
+
         Returns
         -------
         Session :
@@ -197,7 +202,7 @@ def start_session(self, name="hexagon-rpc") -> Session:
             "timeout": 0,
             "key": self._device_key,
         }
-        return Session(self, hexagon_remote_kw, session_name=name)
+        return Session(self, hexagon_remote_kw, session_name=session_name)
 
     def load_module(self, module: Union[str, pathlib.Path, tvm.runtime.Module], session: Session):
         """Load TVM module.

python/tvm/contrib/hexagon/session.py

@@ -44,7 +44,7 @@ class Session:
         Remote configs for RPC tracker.
 
     session_name : str
-        Hexagon RPC session name.
+        Hexagon RPC session name. Options are [hexagon-rpc, cpu-rpc]
 
     remote_stack_size_bytes : int
         The stack size of the remote device, to be passed to
@@ -91,7 +91,9 @@ def __enter__(self):
             elif self._session_name == "hexagon-rpc":
                 self.device = self._rpc.hexagon(0)
             else:
-                raise RuntimeError(f"Incorrect session name: {self._session_name}")
+                raise RuntimeError(
+                    f"Incorrect session name: {self._session_name}. Options for session name are [hexagon-rpc, cpu-rpc]"
+                )
             return self
 
         except RuntimeError as exception:
@@ -317,7 +319,9 @@ def _aot_executor_from_factory(
                     cc=hexagon.hexagon_clang_plus(),
                 )
             else:
-                raise ValueError("Incorrect Target kind.")
+                raise ValueError(
+                    "Incorrect Target kind. Target kind should be from these options: [hexagon, llvm]"
+                )
 
             self.upload(binary_path, binary_name)
 

python/tvm/script/tir/__init__.pyi

@@ -226,6 +226,7 @@ def alloc_buffer(
 """
 special_stmt - Reads/Writes
 """
+
 @overload
 def reads(read_regions: List[BufferSlice]) -> None: ...
 @overload
@@ -337,6 +338,7 @@ def Assert(condition: Union[PrimExpr, builtins.bool], message: str) -> PrimExpr:
 """
 Scope handler - Loops
 """
+
 @overload
 def serial(
     begin: Union[PrimExpr, int],

src/relay/backend/aot_executor_codegen.cc

@@ -1234,6 +1234,8 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
     Target target_host;
     for (const auto& it : tmp) {
       auto dev_type = it.first.as<tir::IntImmNode>();
+      // TODO(tvm-team): AoT only works with kDLCPU device type. We can remove kDLHexagon
+      // here once we refactored kDLHexagon to kDLCPU.
       if (!target_host.defined() && ((it.second->kind->device_type == kDLCPU) ||
                                      (it.second->kind->device_type == kDLHexagon))) {
         target_host = it.second;

src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc

@@ -84,6 +84,7 @@ void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, int ndim, const int64_t* sh
 
 void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, size_t nbytes, size_t alignment,
                                          DLDataType type_hint) {
+  // Added kDLCPU since we use hexagon as a sub-target of LLVM which by default maps to kDLCPU;
   bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) ||
                          (DLDeviceType(dev.device_type) == kDLCPU);
   CHECK(is_valid_device) << "dev.device_type: " << dev.device_type;
@@ -94,6 +95,7 @@ void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, size_t nbytes, size_t align
 }
 
 void HexagonDeviceAPIv2::FreeDataSpace(Device dev, void* ptr) {
+  // Added kDLCPU since we use hexagon as a sub-target of LLVM which by default maps to kDLCPU;
   bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) ||
                          (DLDeviceType(dev.device_type) == kDLCPU);
   CHECK(is_valid_device) << "dev.device_type: " << dev.device_type;
@@ -107,13 +109,15 @@ struct HexagonWorkspacePool : public WorkspacePool {
 };
 
 void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType type_hint) {
+  // Added kDLCPU since we use hexagon as a sub-target of LLVM which by default maps to kDLCPU;
   bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) ||
                          (DLDeviceType(dev.device_type) == kDLCPU);
   CHECK(is_valid_device) << "dev.device_type: " << dev.device_type;
   return dmlc::ThreadLocalStore<HexagonWorkspacePool>::Get()->AllocWorkspace(dev, size);
 }
 
 void HexagonDeviceAPIv2::FreeWorkspace(Device dev, void* data) {
+  // Added kDLCPU since we use hexagon as a sub-target of LLVM which by default maps to kDLCPU;
   bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) ||
                          (DLDeviceType(dev.device_type) == kDLCPU);
   CHECK(is_valid_device) << "dev.device_type: " << dev.device_type;

src/runtime/hexagon/rpc/hexagon/rpc_server.cc

@@ -58,7 +58,7 @@ class HexagonIOHandler {
         read_buffer_size_bytes_{read_buffer_size_bytes},
         write_buffer_available_length_{0} {}
 
-  void MessageStart(size_t message_size_bytes) { LOG(INFO) << "MessageStart called."; }
+  void MessageStart(size_t message_size_bytes) {}
 
   ssize_t PosixWrite(const uint8_t* buf, size_t write_len_bytes) {
     LOG(INFO) << "HexagonIOHandler PosixWrite called, write_len_bytes(" << write_len_bytes << ")";
@@ -159,8 +159,9 @@ class HexagonRPCServer {
    * Otherwise, returns -1;
    */
   int64_t Write(const uint8_t* data, size_t data_size_bytes) {
-    if (io_.SetReadBuffer(data, data_size_bytes) != AEE_SUCCESS) {
-      LOG(ERROR) << "ERROR: SetReadBuffer failed";
+    AEEResult rc = io_.SetReadBuffer(data, data_size_bytes);
+    if (rc != AEE_SUCCESS) {
+      LOG(ERROR) << "ERROR: SetReadBuffer failed: " << rc;
       return -1;
     }
 
@@ -216,6 +217,7 @@ const tvm::runtime::PackedFunc get_runtime_func(const std::string& name) {
 void reset_device_api() {
   const tvm::runtime::PackedFunc api = get_runtime_func("device_api.hexagon.v2");
   tvm::runtime::Registry::Register("device_api.hexagon", true).set_body(api);
+  // Registering device_api.cpu as device_api.hexagon.v2 since we use hexagon as sub-target of LLVM.
   tvm::runtime::Registry::Register("device_api.cpu", true).set_body(api);
 }
 

tests/python/contrib/test_hexagon/test_launcher.py

@@ -274,11 +274,13 @@ def _workaround_create_aot_shared():
 def get_target_and_session(target_kind: str):
     if target_kind == "c":
         target_hexagon = tvm.target.hexagon("v68")
-        session_key = "hexagon-rpc"
+        session_name = "hexagon-rpc"
     elif target_kind.startswith("llvm"):
         target_hexagon = target_kind
-        session_key = "cpu-rpc"
-    return target_hexagon, session_key
+        session_name = "cpu-rpc"
+    else:
+        assert False, "Incorrect target_kind: {target_kind}. Options are [c, llvm]."
+    return target_hexagon, session_name
 
 
 @requires_hexagon_toolchain
@@ -301,7 +303,7 @@ def test_aot_executor(hexagon_launcher, aot_target_kind):
     relay_mod = tvm.IRModule.from_expr(f)
     relay_mod = relay.transform.InferType()(relay_mod)
 
-    target_hexagon, session_key = get_target_and_session(aot_target_kind)
+    target_hexagon, session_name = get_target_and_session(aot_target_kind)
 
     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(
@@ -320,12 +322,11 @@ def test_aot_executor(hexagon_launcher, aot_target_kind):
             executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}),
         )
 
-    hexagon_session = hexagon_launcher.start_session(name=session_key)
-    hexagon_session.__enter__()
-    aot_mod = hexagon_session.get_executor_from_factory(lowered)
-    aot_mod.set_input(**inputs)
-    aot_mod.run()
-    hexagon_output = aot_mod.get_output(0).numpy()
+    with hexagon_launcher.start_session(session_name=session_name) as hexagon_session:
+        aot_mod = hexagon_session.get_executor_from_factory(lowered)
+        aot_mod.set_input(**inputs)
+        aot_mod.run()
+        hexagon_output = aot_mod.get_output(0).numpy()
 
     target_llvm = tvm.target.Target("llvm")
     with tvm.transform.PassContext(opt_level=3):
@@ -375,7 +376,7 @@ def test_aot_executor_multiple_conv2d(hexagon_launcher, aot_target_kind):
     relay_mod = tvm.IRModule.from_expr(f)
     relay_mod = relay.transform.InferType()(relay_mod)
 
-    target_hexagon, session_key = get_target_and_session(aot_target_kind)
+    target_hexagon, session_name = get_target_and_session(aot_target_kind)
 
     weight1_data = np.random.rand(w1_shape[0], w1_shape[1], w1_shape[2], w1_shape[3]).astype(
         dtype=dtype
@@ -399,12 +400,11 @@ def test_aot_executor_multiple_conv2d(hexagon_launcher, aot_target_kind):
             executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}),
         )
 
-    hexagon_session = hexagon_launcher.start_session(name=session_key)
-    hexagon_session.__enter__()
-    aot_mod = hexagon_session.get_executor_from_factory(lowered)
-    aot_mod.set_input(**inputs)
-    aot_mod.run()
-    hexagon_output = aot_mod.get_output(0).numpy()
+    with hexagon_launcher.start_session(session_name=session_name) as hexagon_session:
+        aot_mod = hexagon_session.get_executor_from_factory(lowered)
+        aot_mod.set_input(**inputs)
+        aot_mod.run()
+        hexagon_output = aot_mod.get_output(0).numpy()
 
     target_llvm = tvm.target.Target("llvm")
     with tvm.transform.PassContext(opt_level=3):