[Hexagon] [runtime] Improve runtime resource management

python/tvm/contrib/hexagon/session.py

@@ -88,14 +88,25 @@ def __enter__(self):
                     self._rpc_receive_buffer_size_bytes,
                 ],
             )
+            func = self._rpc.get_function("device_api.hexagon.acquire_resources")
+            func()
             return self
 
         except RuntimeError as exception:
             raise exception
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
-        # close session to the tracker
-        del self._rpc
+        try:
+            func = self._rpc.get_function("device_api.hexagon.release_resources")
+            func()
+        except RuntimeError as exception:
+            print(
+                "Exception occurred while calling release_resources() during Session __exit__: ",
+                exception,
+            )
+        finally:
+            # close session to the tracker
+            del self._rpc
 
     @property
     def device(self):

src/runtime/hexagon/hexagon_buffer_manager.h

@@ -69,6 +69,9 @@ class HexagonBufferManager {
     return nullptr;
   }
 
+  //! \brief Returns whether the HexagonBufferManager has any allocations.
+  bool empty() { return hexagon_buffer_map_.empty(); }
+
  private:
   //! \brief Contains the HexagonBuffer objects managed by this class.
   std::unordered_map<void*, std::unique_ptr<HexagonBuffer>> hexagon_buffer_map_;

src/runtime/hexagon/hexagon_device_api.cc

@@ -90,18 +90,20 @@ void* HexagonDeviceAPI::AllocDataSpace(Device dev, int ndim, const int64_t* shap
 
   const size_t typesize = (dtype.bits / 8) * dtype.lanes;
 
+  HexagonBufferManager* mgr = runtime_hexbuffs ? runtime_hexbuffs.get() : &hexbuffs;
+
   if (ndim == 0) {
     // Allocate storage for a single scalar value.
-    return hexbuffs.AllocateHexagonBuffer(typesize, kHexagonAllocAlignment, mem_scope);
+    return mgr->AllocateHexagonBuffer(typesize, kHexagonAllocAlignment, mem_scope);
   } else if (ndim == 1) {
     // Allocate a single, contiguous memory region.
     size_t nbytes = shape[0] * typesize;
-    return hexbuffs.AllocateHexagonBuffer(nbytes, kHexagonAllocAlignment, mem_scope);
+    return mgr->AllocateHexagonBuffer(nbytes, kHexagonAllocAlignment, mem_scope);
   } else if (ndim == 2) {
     // Allocate the region(s) needed for Hexagon's indirect-tensor format.
     size_t nallocs = shape[0];
     size_t nbytes = shape[1] * typesize;
-    return hexbuffs.AllocateHexagonBuffer(nallocs, nbytes, kHexagonAllocAlignment, mem_scope);
+    return mgr->AllocateHexagonBuffer(nallocs, nbytes, kHexagonAllocAlignment, mem_scope);
   } else {
     return nullptr;  // unreachable
   }
@@ -115,13 +117,16 @@ void* HexagonDeviceAPI::AllocDataSpace(Device dev, size_t nbytes, size_t alignme
   if (alignment < kHexagonAllocAlignment) {
     alignment = kHexagonAllocAlignment;
   }
-  return hexbuffs.AllocateHexagonBuffer(nbytes, alignment, String("global"));
+  HexagonBufferManager* mgr = runtime_hexbuffs ? runtime_hexbuffs.get() : &hexbuffs;
+  return mgr->AllocateHexagonBuffer(nbytes, alignment, String("global"));
 }
 
 void HexagonDeviceAPI::FreeDataSpace(Device dev, void* ptr) {
   CHECK(ptr) << "buffer pointer is null";
   CHECK(IsValidDevice(dev)) << "dev.device_type: " << dev.device_type;
-  hexbuffs.FreeHexagonBuffer(ptr);
+  HexagonBufferManager* mgr =
+      (runtime_hexbuffs && runtime_hexbuffs->count(ptr) != 0) ? runtime_hexbuffs.get() : &hexbuffs;
+  mgr->FreeHexagonBuffer(ptr);
 }
 
 // WorkSpace: runtime allocations for Hexagon
@@ -137,7 +142,8 @@ void* HexagonDeviceAPI::AllocWorkspace(Device dev, size_t size, DLDataType type_
 
 void HexagonDeviceAPI::FreeWorkspace(Device dev, void* data) {
   CHECK(IsValidDevice(dev)) << "dev.device_type: " << dev.device_type;
-  CHECK(hexbuffs.count(data) != 0)
+  CHECK(hexbuffs.count(data) != 0 ||
+        (runtime_hexbuffs != nullptr && runtime_hexbuffs->count(data) != 0))
       << "Attempt made to free unknown or already freed workspace allocation";
   dmlc::ThreadLocalStore<HexagonWorkspacePool>::Get()->FreeWorkspace(dev, data);
 }
@@ -161,7 +167,10 @@ void HexagonDeviceAPI::CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHan
   CHECK_EQ(to->byte_offset, 0);
   CHECK_EQ(GetDataSize(*from), GetDataSize(*to));
 
-  auto lookup_hexagon_buffer = [this](void* ptr) -> HexagonBuffer* { return hexbuffs.find(ptr); };
+  auto lookup_hexagon_buffer = [this](void* ptr) -> HexagonBuffer* {
+    return (runtime_hexbuffs && runtime_hexbuffs->count(ptr) != 0) ? runtime_hexbuffs->find(ptr)
+                                                                   : hexbuffs.find(ptr);
+  };
 
   HexagonBuffer* hex_from_buf = lookup_hexagon_buffer(from->data);
   HexagonBuffer* hex_to_buf = lookup_hexagon_buffer(to->data);
@@ -235,6 +244,18 @@ TVM_REGISTER_GLOBAL("device_api.hexagon.free_nd").set_body([](TVMArgs args, TVMR
   *rv = static_cast<int32_t>(0);
 });
 
+TVM_REGISTER_GLOBAL("device_api.hexagon.acquire_resources")
+    .set_body([](TVMArgs args, TVMRetValue* rv) {
+      HexagonDeviceAPI* api = HexagonDeviceAPI::Global();
+      api->AcquireResources();
+    });
+
+TVM_REGISTER_GLOBAL("device_api.hexagon.release_resources")
+    .set_body([](TVMArgs args, TVMRetValue* rv) {
+      HexagonDeviceAPI* api = HexagonDeviceAPI::Global();
+      api->ReleaseResources();
+    });
+
 TVM_REGISTER_GLOBAL("device_api.hexagon").set_body([](TVMArgs args, TVMRetValue* rv) {
   DeviceAPI* ptr = HexagonDeviceAPI::Global();
   *rv = static_cast<void*>(ptr);

src/runtime/hexagon/hexagon_device_api.h

@@ -50,6 +50,22 @@ class HexagonDeviceAPI final : public DeviceAPI {
   //! \brief Destructor
   ~HexagonDeviceAPI() {}
 
+  //! \brief Ensures resource managers are in a good state for the runtime
+  void AcquireResources() {
+    CHECK_EQ(runtime_hexbuffs, nullptr);
+    runtime_hexbuffs = std::make_unique<HexagonBufferManager>();
+    LOG(INFO) << "runtime_hexbuffs created";
+  }
+
+  //! \brief Ensures all runtime resources are freed
+  void ReleaseResources() {
+    if (runtime_hexbuffs && !runtime_hexbuffs->empty()) {
+      LOG(INFO) << "runtime_hexbuffs was not empty in ReleaseResources";
+    }
+    LOG(INFO) << "runtime_hexbuffs reset";
+    runtime_hexbuffs.reset();
+  }
+
   /*! \brief Currently unimplemented interface to specify the active
    *  Hexagon device.
    */
@@ -138,7 +154,11 @@ class HexagonDeviceAPI final : public DeviceAPI {
   }
 
   //! \brief Manages underlying HexagonBuffer allocations
+  // runtime_hexbuffs is used for runtime allocations.  It is created
+  // with a call to AcquireResources, and destroyed on ReleaseResources.
+  // hexbuffs is used for all allocations outside of the session lifetime.
   HexagonBufferManager hexbuffs;
+  std::unique_ptr<HexagonBufferManager> runtime_hexbuffs;
 };
 }  // namespace hexagon
 }  // namespace runtime

tests/cpp-runtime/hexagon/hexagon_device_api_tests.cc

@@ -146,3 +146,21 @@ TEST_F(HexagonDeviceAPITest, DISABLED_alloc_free_diff_dev) {
   CHECK(buf != nullptr);
   EXPECT_THROW(hexapi->FreeDataSpace(cpu_dev, buf), InternalError);
 }
+
+// Alloc a non-runtime buffer
+// Alloc a runtime buffer
+// "Release" resources for runtime
+// Verify the runtime buffer cannot be freed, but the non-runtime buffer can
+// This test should be run last
+TEST_F(HexagonDeviceAPITest, leak_resources) {
+  hexapi->ReleaseResources();
+  void* pre_runtime_buf = hexapi->AllocDataSpace(hex_dev, nbytes, alignment, int8);
+  CHECK(pre_runtime_buf != nullptr);
+  hexapi->AcquireResources();
+  void* runtime_buf = hexapi->AllocDataSpace(hex_dev, nbytes, alignment, int8);
+  CHECK(runtime_buf != nullptr);
+  hexapi->ReleaseResources();
+  EXPECT_THROW(hexapi->FreeDataSpace(hex_dev, runtime_buf), InternalError);
+  hexapi->FreeDataSpace(hex_dev, pre_runtime_buf);
+  hexapi->AcquireResources();
+}