[Hexagon] [runtime] Improve runtime resource management (#12727)

* First pass at improving runtime resource management

* Add unit test

* Fix lint and clang format errors

* Disable resource reset for simulator

* Moved acquire/release calls to session object, separate buffer managers for non-runtime (static) and runtime (dynamic).

* Fix lint errors

* Fix lint errors

* Improve robustness of session shutdown

* Fix lint

* Address feedback

* Only allow call to Acquire in a clean state

* Use a pointer to indicate the "active" manager

Author: janetsc

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_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 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 +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 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);
+  mgr->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(mgr->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 mgr->find(ptr); };
 
   HexagonBuffer* hex_from_buf = lookup_hexagon_buffer(from->data);
   HexagonBuffer* hex_to_buf = lookup_hexagon_buffer(to->data);
@@ -246,6 +246,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,29 @@ class HexagonDeviceAPI final : public DeviceAPI {
   static HexagonDeviceAPI* Global();
 
   //! \brief Constructor
-  HexagonDeviceAPI() {}
+  HexagonDeviceAPI() { mgr = &hexbuffs; }
 
   //! \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";
+    mgr = runtime_hexbuffs.get();
+  }
+
+  //! \brief Ensures all runtime resources are freed
+  void ReleaseResources() {
+    if (runtime_hexbuffs && !runtime_hexbuffs->empty()) {
+      LOG(INFO) << "runtime_hexbuffs was not empty in ReleaseResources";
+    }
+    mgr = &hexbuffs;
+    LOG(INFO) << "runtime_hexbuffs reset";
+    runtime_hexbuffs.reset();
+  }
+
   /*! \brief Currently unimplemented interface to specify the active
    *  Hexagon device.
    */
@@ -138,7 +156,14 @@ 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;
+
+  //! \brief Current buffer manager
+  HexagonBufferManager* mgr;
 };
 }  // 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();
+}