[Hexagon] [runtime] Improve runtime resource management

python/tvm/contrib/hexagon/build.py

@@ -149,6 +149,18 @@ def stop_server(self):
         """Stop the RPC server"""
         ...
 
+    def _acquire_device_resources(self):
+        """Call into device to initialize objects to acquire resources"""
+        with self.start_session() as session:
+            func = session._rpc.get_function("device_api.hexagon.acquire_resources")
+            func()
+
+    def _release_device_resources(self):
+        """Call into device to release resources"""
+        with self.start_session() as session:
+            func = session._rpc.get_function("device_api.hexagon.release_resources")
+            func()
+
     @abc.abstractmethod
     def cleanup_directory(self):
         """Cleanup working directory"""
@@ -508,9 +520,11 @@ def start_server(self):
         """Abstract method implementation. See description in HexagonLauncherRPC."""
         self._copy_binaries()
         self._run_server_script()
+        super()._acquire_device_resources()
 
     def stop_server(self):
         """Abstract method implementation. See description in HexagonLauncherRPC."""
+        super()._release_device_resources()
         self._cleanup_port_forwarding()
         self._terminate_remote()
         self.cleanup_directory()
@@ -613,12 +627,14 @@ def _start(self):
 
         self._server_process = mp.Process(target=lambda *a: _start(self, *a))
         self._server_process.start()
+        super()._acquire_device_resources()
 
     def cleanup_directory(self):
         """Abstract method implementation. See description in HexagonLauncherRPC."""
 
     def stop_server(self):
         """Abstract method implementation. See description in HexagonLauncherRPC."""
+        super()._release_device_resources()
         self._server_process.terminate()
 
 

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

@@ -92,16 +92,16 @@ void* HexagonDeviceAPI::AllocDataSpace(Device dev, int ndim, const int64_t* shap
 
   if (ndim == 0) {
     // Allocate storage for a single scalar value.
-    return hexbuffs.AllocateHexagonBuffer(typesize, kHexagonAllocAlignment, mem_scope);
+    return hexbuffs->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 hexbuffs->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 hexbuffs->AllocateHexagonBuffer(nallocs, nbytes, kHexagonAllocAlignment, mem_scope);
   } else {
     return nullptr;  // unreachable
   }
@@ -115,13 +115,13 @@ void* HexagonDeviceAPI::AllocDataSpace(Device dev, size_t nbytes, size_t alignme
   if (alignment < kHexagonAllocAlignment) {
     alignment = kHexagonAllocAlignment;
   }
-  return hexbuffs.AllocateHexagonBuffer(nbytes, alignment, String("global"));
+  return hexbuffs->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);
+  hexbuffs->FreeHexagonBuffer(ptr);
 }
 
 // WorkSpace: runtime allocations for Hexagon
@@ -137,7 +137,7 @@ 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)
       << "Attempt made to free unknown or already freed workspace allocation";
   dmlc::ThreadLocalStore<HexagonWorkspacePool>::Get()->FreeWorkspace(dev, data);
 }
@@ -161,7 +161,7 @@ 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 hexbuffs->find(ptr); };
 
   HexagonBuffer* hex_from_buf = lookup_hexagon_buffer(from->data);
   HexagonBuffer* hex_to_buf = lookup_hexagon_buffer(to->data);
@@ -235,6 +235,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

@@ -45,11 +45,27 @@ class HexagonDeviceAPI final : public DeviceAPI {
   static HexagonDeviceAPI* Global();
 
   //! \brief Constructor
-  HexagonDeviceAPI() {}
+  HexagonDeviceAPI() { hexbuffs = std::make_unique<HexagonBufferManager>(); }
 
   //! \brief Destructor
   ~HexagonDeviceAPI() {}
 
+  //! \brief Creates resource managers for the runtime
+  void AcquireResources() {
+    if (!hexbuffs->empty()) {
+      LOG(INFO) << "hexbuffs was not empty in AcquireResources";
+    }
+  }
+
+  //! \brief Ensures we have freed all resources when we end the runtime
+  void ReleaseResources() {
+    if (!hexbuffs->empty()) {
+      LOG(INFO) << "hexbuffs was not empty in ReleaseResources, resetting";
+      hexbuffs.reset();
+      hexbuffs = std::make_unique<HexagonBufferManager>();
+    }
+  }
+
   /*! \brief Currently unimplemented interface to specify the active
    *  Hexagon device.
    */
@@ -138,7 +154,7 @@ class HexagonDeviceAPI final : public DeviceAPI {
   }
 
   //! \brief Manages underlying HexagonBuffer allocations
-  HexagonBufferManager hexbuffs;
+  std::unique_ptr<HexagonBufferManager> hexbuffs;
 };
 }  // namespace hexagon
 }  // namespace runtime

tests/cpp-runtime/hexagon/hexagon_device_api_tests.cc

@@ -146,3 +146,13 @@ TEST_F(HexagonDeviceAPITest, DISABLED_alloc_free_diff_dev) {
   CHECK(buf != nullptr);
   EXPECT_THROW(hexapi->FreeDataSpace(cpu_dev, buf), InternalError);
 }
+
+// Alloc a buffer, but do not free it.
+// "Release" resources, and verify it cannot be freed.
+TEST_F(HexagonDeviceAPITest, leak_resources) {
+  hexapi->AcquireResources();
+  void* buf = hexapi->AllocDataSpace(hex_dev, nbytes, alignment, int8);
+  CHECK(buf != nullptr);
+  hexapi->ReleaseResources();
+  EXPECT_THROW(hexapi->FreeDataSpace(hex_dev, buf), InternalError);
+}