add integration test_add_vtcm to demo vtcm alloc

Author: adstraw

src/runtime/hexagon/hexagon/hexagon_buffer.cc

@@ -25,7 +25,7 @@
 
 #include "hexagon_common.h"
 
-#ifdef BUILD_FOR_HEXAGON
+#if defined(__hexagon__)
 #include "HAP_compute_res.h"
 #endif
 
@@ -68,19 +68,18 @@ struct DDRAllocation : public Allocation {
   }
 };
 
-#ifdef BUILD_FOR_HEXAGON
+#if defined(__hexagon__)
 struct VTCMAllocation : public Allocation {
   VTCMAllocation(size_t nbytes, size_t alignment) : Allocation(nbytes, alignment) {
-    // TODO(Straw): Alignment not used when allocating VTCM
     compute_res_attr_t res_info;
     HEXAGON_SAFE_CALL(HAP_compute_res_attr_init(&res_info));
-    // TODO(Straw): Magic number 1
-    HEXAGON_SAFE_CALL(HAP_compute_res_attr_set_vtcm_param(&res_info, nbytes, 1));
-    // TODO(Straw): HEXAGON_SAFE_CALL?
-    // TODO(Straw): Magic number 10000
-    context_id_ = HAP_compute_res_acquire(&res_info, 10000);
+
+    // allocate nbytes of vtcm on a single page
+    HEXAGON_SAFE_CALL(HAP_compute_res_attr_set_vtcm_param(&res_info, /*vtcm_size = */ nbytes,
+                                                          /*b_single_page = */ 1));
+    context_id_ = HAP_compute_res_acquire(&res_info, /*timeout = */ 10000);
+
     if (context_id_) {
-      // TODO(Straw): HEXAGON_SAFE_CALL?
       data_ = HAP_compute_res_attr_get_vtcm_ptr(&res_info);
       if (!data_) {
         HEXAGON_PRINT(ERROR, "ERROR: Allocated VTCM ptr is null.");
@@ -96,11 +95,8 @@ struct VTCMAllocation : public Allocation {
   ~VTCMAllocation() {
     // HEXAGON_PRINT(ALWAYS, "~VTCMAllocation() - Context ID: %u, VTCM ptr: %p", context_id_,
     // data_);
-    // TODO(Straw): Need to handle the else case(s) here
-    if (context_id_ && data_) {
-      HEXAGON_SAFE_CALL(HAP_compute_res_release(context_id_));
-      data_ = nullptr;
-    }
+    HEXAGON_SAFE_CALL(HAP_compute_res_release(context_id_));
+    data_ = nullptr;
   }
   unsigned int context_id_{0};
 };

src/runtime/hexagon/hexagon/hexagon_common.cc

@@ -87,7 +87,7 @@ PackedFunc WrapPackedFunc(TVMBackendPackedCFunc faddr, const ObjectPtr<Object>&
       if (args.type_codes[i] == kTVMDLTensorHandle) {
         DLTensor* tensor = static_cast<DLTensor*>(arg_values[i].v_handle);
         buffer_args.emplace_back(i, static_cast<HexagonBuffer*>(tensor->data));
-        tensor->data = buffer_args.back().second->GetPointer();
+        tensor->data = buffer_args.back().second->GetPointer()[0];
       }
     }
     int ret = (*faddr)(const_cast<TVMValue*>(args.values), const_cast<int*>(args.type_codes),

src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc

@@ -106,7 +106,7 @@ void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType typ
 void HexagonDeviceAPIv2::FreeWorkspace(Device dev, void* data) {
   CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
   auto it = workspace_allocations_.find(data);
-  ICHECK(it != workspace_allocations_.end())
+  CHECK(it != workspace_allocations_.end())
       << "Attempt made to free unknown or already freed workspace allocation";
   dmlc::ThreadLocalStore<HexagonWorkspacePool>::Get()->FreeWorkspace(dev, it->second);
   workspace_allocations_.erase(it);

tests/cpp/hexagon_test.cc

@@ -76,6 +76,54 @@ def test_add(tvm_tracker_host, tvm_tracker_port, android_serial_number):
     launcher.close()
 
 
+@requires_rpc_tracker
+@requires_hexagon_toolchain
+def test_add_vtcm(tvm_tracker_host, tvm_tracker_port, android_serial_number):
+    dtype = "int8"
+    A = tvm.te.placeholder((2,), dtype=dtype)
+    B = tvm.te.placeholder((1,), dtype=dtype)
+    C = tvm.te.compute(A.shape, lambda i: A[i] + B[0], name="C")
+    sched = tvm.te.create_schedule(C.op)
+
+    target_hexagon = tvm.target.hexagon("v68", link_params=True)
+    func = tvm.build(
+        sched, [A, B, C], tvm.target.Target(target_hexagon, host=target_hexagon), name="add"
+    )
+
+    temp = utils.tempdir()
+    dso_binary = "test_binary.so"
+    dso_binary_path = temp.relpath(dso_binary)
+    func.save(dso_binary_path)
+
+    launcher = HexagonLauncher(serial_number=android_serial_number)
+    launcher.android_run_rpc(rpc_tracker_host=tvm_tracker_host, rpc_tracker_port=tvm_tracker_port)
+    launcher.hexagon_setup()
+    remote_kw = {
+        "host": tvm_tracker_host,
+        "port": tvm_tracker_port,
+        "priority": 0,
+        "timeout": 60,
+    }
+    launcher.hexagon_session_setup(remote_kw)
+    launcher.upload(dso_binary_path, dso_binary)
+
+    with launcher.session as sess:
+        mod = launcher.get_module(dso_binary)
+        A_data = tvm.nd.empty(A.shape, A.dtype, sess.device, "global.vtcm")
+        A_data.copyfrom(np.array([2, 3]))
+
+        B_data = tvm.nd.empty(B.shape, B.dtype, sess.device, "global.vtcm")
+        B_data.copyfrom(np.array([4]))
+
+        C_data = tvm.nd.empty(C.shape, C.dtype, sess.device, "global.vtcm")
+        C_data.copyfrom(np.array([0, 0]))
+
+        mod["add"](A_data, B_data, C_data)
+        result = C_data.numpy()
+        assert (result == np.array([6, 7])).all()
+    launcher.close()
+
+
 class TestMatMul:
     M = tvm.testing.parameter(32)
     N = tvm.testing.parameter(32)