[Hexagon] Refactor HexagonBufferManager

src/runtime/hexagon/hexagon_buffer_manager.h

@@ -40,21 +40,20 @@ class HexagonBufferManager {
    * \param ptr Address of the HexagonBuffer as returned by `AllocateHexagonBuffer`.
    */
   void FreeHexagonBuffer(void* ptr) {
+    std::lock_guard<std::mutex> lock(map_mutex_);
     auto it = hexagon_buffer_map_.find(ptr);
     CHECK(it != hexagon_buffer_map_.end())
         << "Attempt made to free unknown or already freed dataspace allocation";
     CHECK(it->second != nullptr);
-    {
-      std::lock_guard<std::mutex> lock(map_mutex_);
-      hexagon_buffer_map_.erase(it);
-    }
+    hexagon_buffer_map_.erase(it);
   }
   /*!
    * \brief Allocate a HexagonBuffer.
    * \param args Templated arguments to pass through to HexagonBuffer constructor.
    */
   template <typename... Args>
   void* AllocateHexagonBuffer(Args&&... args) {
+    CHECK(enabled_);
     auto buf = std::make_unique<HexagonBuffer>(std::forward<Args>(args)...);
     void* ptr = buf->GetPointer();
     {
@@ -64,14 +63,9 @@ class HexagonBufferManager {
     return ptr;
   }
 
-  //! \brief Returns whether the HexagonBuffer is in the map.
-  size_t count(void* ptr) {
-    std::lock_guard<std::mutex> lock(map_mutex_);
-    return hexagon_buffer_map_.count(ptr);
-  }
-
   //! \brief Returns an iterator to the HexagonBuffer within the map.
-  HexagonBuffer* find(void* ptr) {
+  HexagonBuffer* FindHexagonBuffer(void* ptr) {
+    CHECK(enabled_);
     std::lock_guard<std::mutex> lock(map_mutex_);
     auto it = hexagon_buffer_map_.find(ptr);
     if (it != hexagon_buffer_map_.end()) {
@@ -80,30 +74,23 @@ class HexagonBufferManager {
     return nullptr;
   }
 
-  //! \brief Returns whether the HexagonBufferManager has any allocations.
-  bool empty() {
+  void Enable() {
     std::lock_guard<std::mutex> lock(map_mutex_);
-    return hexagon_buffer_map_.empty();
+    CHECK(hexagon_buffer_map_.empty());
+    enabled_ = true;
   }
 
-  //! \brief Returns a vector of currently allocated pointers, owned by the manager.
-  // Note - this should only be used by the device API to keep track of what
-  // was in the manager when HexagonDeviceAPI::ReleaseResources is called.
-  std::vector<void*> current_allocations() {
-    std::vector<void*> allocated;
-    std::lock_guard<std::mutex> lock(map_mutex_);
-    for (const auto& [data_ptr, buffer] : hexagon_buffer_map_) {
-      allocated.push_back(data_ptr);
-    }
-    return allocated;
-  }
+  void Disable() { enabled_ = false; }
 
  private:
   //! \brief Contains the HexagonBuffer objects managed by this class.
   std::unordered_map<void*, std::unique_ptr<HexagonBuffer>> hexagon_buffer_map_;
 
   //! \brief Protects updates to the map.
   std::mutex map_mutex_;
+
+  //! \brief Tracks whether this HexagonBufferManager is enabled
+  bool enabled_ = true;
 };
 
 }  // namespace hexagon

src/runtime/hexagon/hexagon_device_api.cc

@@ -90,23 +90,19 @@ void* HexagonDeviceAPI::AllocDataSpace(Device dev, int ndim, const int64_t* shap
 
   const size_t typesize = (dtype.bits / 8) * dtype.lanes;
 
-  CHECK(runtime_hexbuffs) << "Attempted to allocate Hexagon data with "
-                          << "HexagonDeviceAPI::AllocDataSpace before initializing resources.  "
-                          << "Please call HexagonDeviceAPI::AcquireResources";
-
   if (ndim == 0) {
     // Allocate storage for a single scalar value.
-    return runtime_hexbuffs->AllocateHexagonBuffer(typesize, kHexagonAllocAlignment, mem_scope);
+    return runtime_hexbuffs.AllocateHexagonBuffer(typesize, kHexagonAllocAlignment, mem_scope);
   } else if (ndim == 1) {
     // Allocate a single, contiguous memory region.
     size_t nbytes = shape[0] * typesize;
-    return runtime_hexbuffs->AllocateHexagonBuffer(nbytes, kHexagonAllocAlignment, mem_scope);
+    return runtime_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 runtime_hexbuffs->AllocateHexagonBuffer(nallocs, nbytes, kHexagonAllocAlignment,
-                                                   mem_scope);
+    return runtime_hexbuffs.AllocateHexagonBuffer(nallocs, nbytes, kHexagonAllocAlignment,
+                                                  mem_scope);
   } else {
     return nullptr;  // unreachable
   }
@@ -120,27 +116,13 @@ void* HexagonDeviceAPI::AllocDataSpace(Device dev, size_t nbytes, size_t alignme
   if (alignment < kHexagonAllocAlignment) {
     alignment = kHexagonAllocAlignment;
   }
-  CHECK(runtime_hexbuffs) << "Attempted to allocate Hexagon data with "
-                          << "HexagonDeviceAPI::AllocDataSpace before initializing resources.  "
-                          << "Please call HexagonDeviceAPI::AcquireResources";
-  return runtime_hexbuffs->AllocateHexagonBuffer(nbytes, alignment, String("global"));
+  return runtime_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;
-  if (runtime_hexbuffs) {
-    runtime_hexbuffs->FreeHexagonBuffer(ptr);
-  } else {
-    // Either AcquireResources was never called, or ReleaseResources was called.  Check the
-    // list of buffers that were still allocated at the time of release.  If this buffer is
-    // in that list, this is a no-op as it is being freed as part of another object teardown.
-    // If the pointer isn't in that list, we raise an exception as this is an unexpected Free.
-    auto it = std::find(released_runtime_buffers.begin(), released_runtime_buffers.end(), ptr);
-    CHECK(it != released_runtime_buffers.end()) << "Attempted to free Hexagon data with "
-                                                << "HexagonDeviceAPI::FreeDataSpace that was not "
-                                                << "allocated during the session.";
-  }
+  runtime_hexbuffs.FreeHexagonBuffer(ptr);
 }
 
 void* HexagonDeviceAPI::AllocRpcBuffer(size_t nbytes, size_t alignment) {
@@ -167,11 +149,6 @@ 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(runtime_hexbuffs) << "Attempted to free Hexagon workspace with "
-                          << "HexagonDeviceAPI::FreeWorkspace outside of a session.  "
-                          << "Please call HexagonDeviceAPI::AcquireResources";
-  CHECK(runtime_hexbuffs->count(data) != 0)
-      << "Attempt made to free unknown or already freed workspace allocation";
   dmlc::ThreadLocalStore<HexagonWorkspacePool>::Get()->FreeWorkspace(dev, data);
 }
 
@@ -193,12 +170,9 @@ void HexagonDeviceAPI::CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHan
   CHECK_EQ(from->byte_offset, 0);
   CHECK_EQ(to->byte_offset, 0);
   CHECK_EQ(GetDataSize(*from), GetDataSize(*to));
-  CHECK(runtime_hexbuffs) << "Attempted to copy Hexagon data with "
-                          << "HexagonDeviceAPI::CopyDataFromTo before initializing resources.  "
-                          << "Please call HexagonDeviceAPI::AcquireResources";
 
   auto lookup_hexagon_buffer = [this](void* ptr) -> HexagonBuffer* {
-    return runtime_hexbuffs->find(ptr);
+    return runtime_hexbuffs.FindHexagonBuffer(ptr);
   };
 
   HexagonBuffer* hex_from_buf = lookup_hexagon_buffer(from->data);

src/runtime/hexagon/hexagon_device_api.h

@@ -55,13 +55,11 @@ class HexagonDeviceAPI final : public DeviceAPI {
 
   //! \brief Ensures resource managers are in a good state for the runtime
   void AcquireResources() {
+    runtime_hexbuffs.Enable();
+
     CHECK_EQ(runtime_vtcm, nullptr);
     runtime_vtcm = std::make_unique<HexagonVtcmPool>();
 
-    CHECK_EQ(runtime_hexbuffs, nullptr);
-    runtime_hexbuffs = std::make_unique<HexagonBufferManager>();
-    released_runtime_buffers.clear();
-
     CHECK_EQ(runtime_threads, nullptr);
     runtime_threads = std::make_unique<HexagonThreadManager>(threads, stack_size, pipe_size);
 
@@ -77,15 +75,10 @@ class HexagonDeviceAPI final : public DeviceAPI {
     CHECK(runtime_threads) << "runtime_threads was not created in AcquireResources";
     runtime_threads.reset();
 
-    CHECK(runtime_hexbuffs) << "runtime_hexbuffs was not created in AcquireResources";
-    if (!runtime_hexbuffs->empty()) {
-      DLOG(INFO) << "runtime_hexbuffs was not empty in ReleaseResources";
-      released_runtime_buffers = runtime_hexbuffs->current_allocations();
-    }
-    runtime_hexbuffs.reset();
-
     CHECK(runtime_vtcm) << "runtime_vtcm was not created in AcquireResources";
     runtime_vtcm.reset();
+
+    runtime_hexbuffs.Disable();
   }
 
   /*! \brief Currently unimplemented interface to specify the active
@@ -202,15 +195,9 @@ class HexagonDeviceAPI final : public DeviceAPI {
   HexagonBufferManager rpc_hexbuffs;
 
   //! \brief Manages runtime HexagonBuffer allocations
-  // runtime_hexbuffs is used for runtime allocations, separate from rpc_hexbuffs.  It is created
-  // with a call to AcquireResources, and destroyed on ReleaseResources.  The buffers in this
-  // manager are scoped to the lifetime of a user application session.
-  std::unique_ptr<HexagonBufferManager> runtime_hexbuffs;
-
-  //! \brief Keeps a list of released runtime HexagonBuffer allocations
-  // ReleaseResources can be called when there are still buffers in runtime_hexbuffs.  This list
-  // stores the buffers that were released.
-  std::vector<void*> released_runtime_buffers;
+  // runtime_hexbuffs is used for runtime allocations, separate from rpc_hexbuffs. The buffers
+  // in this manager are scoped to the lifetime of a user application session.
+  HexagonBufferManager runtime_hexbuffs;
 
   //! \brief Thread manager
   std::unique_ptr<HexagonThreadManager> runtime_threads;

tests/cpp-runtime/hexagon/hexagon_vtcm_pool_tests.cc

@@ -166,36 +166,34 @@ TEST_F(HexagonVtcmPoolTest, free_alloc_combinations) {
 
 // Test alignment edge cases allocating through HexagonBuffer
 TEST_F(HexagonVtcmPoolTest, vtcm_alignment) {
-  std::unique_ptr<HexagonBufferManager> test_hexbuffs = std::make_unique<HexagonBufferManager>();
+  HexagonBufferManager test_hexbuffs;
   void* ptr;
 
   // Invalid alignments
-  EXPECT_THROW(test_hexbuffs->AllocateHexagonBuffer(min_bytes, 128 + 1, String("global")),
+  EXPECT_THROW(test_hexbuffs.AllocateHexagonBuffer(min_bytes, 128 + 1, String("global")),
                InternalError);
-  EXPECT_THROW(test_hexbuffs->AllocateHexagonBuffer(min_bytes, 2048 + 1, String("global")),
+  EXPECT_THROW(test_hexbuffs.AllocateHexagonBuffer(min_bytes, 2048 + 1, String("global")),
                InternalError);
 
   // Valid alignments, sizes need to be adjusted
-  ptr = test_hexbuffs->AllocateHexagonBuffer(1, 128, String("global"));
+  ptr = test_hexbuffs.AllocateHexagonBuffer(1, 128, String("global"));
   CHECK((reinterpret_cast<uintptr_t>(ptr) & 0x7F) == 0) << "Must be multiple of 128 " << ptr;
 
-  ptr = test_hexbuffs->AllocateHexagonBuffer(127, 128, String("global"));
+  ptr = test_hexbuffs.AllocateHexagonBuffer(127, 128, String("global"));
   CHECK((reinterpret_cast<uintptr_t>(ptr) & 0x7F) == 0) << "Must be multiple of 128 " << ptr;
 
-  ptr = test_hexbuffs->AllocateHexagonBuffer(129, 128, String("global"));
+  ptr = test_hexbuffs.AllocateHexagonBuffer(129, 128, String("global"));
   CHECK((reinterpret_cast<uintptr_t>(ptr) & 0x7F) == 0) << "Must be multiple of 128 " << ptr;
 
-  ptr = test_hexbuffs->AllocateHexagonBuffer(1, 2048, String("global"));
+  ptr = test_hexbuffs.AllocateHexagonBuffer(1, 2048, String("global"));
   CHECK((reinterpret_cast<uintptr_t>(ptr) & 0x7FF) == 0) << "Must be multiple of 2k " << ptr;
 
-  ptr = test_hexbuffs->AllocateHexagonBuffer(2047, 2048, String("global"));
+  ptr = test_hexbuffs.AllocateHexagonBuffer(2047, 2048, String("global"));
   CHECK((reinterpret_cast<uintptr_t>(ptr) & 0x7FF) == 0) << "Must be multiple of 2k " << ptr;
 
-  ptr = test_hexbuffs->AllocateHexagonBuffer(2049, 2048, String("global"));
+  ptr = test_hexbuffs.AllocateHexagonBuffer(2049, 2048, String("global"));
   CHECK((reinterpret_cast<uintptr_t>(ptr) & 0x7FF) == 0) << "Must be multiple of 2k " << ptr;
 
-  test_hexbuffs.reset();
-
   // Make sure at the end we have the full amount available again
   ptr = vtcm_pool->Allocate(max_bytes);
   vtcm_pool->Free(ptr, max_bytes);