[Hexagon] Add HexagonThreadManager

src/runtime/hexagon/hexagon_buffer_manager.h

@@ -0,0 +1,81 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#ifndef TVM_RUNTIME_HEXAGON_HEXAGON_BUFFER_MANAGER_H_
+#define TVM_RUNTIME_HEXAGON_HEXAGON_BUFFER_MANAGER_H_
+
+#include <tvm/runtime/logging.h>
+
+#include <memory>
+#include <unordered_map>
+#include <utility>
+
+#include "hexagon_buffer.h"
+
+namespace tvm {
+namespace runtime {
+namespace hexagon {
+
+class HexagonBufferManager {
+ public:
+  /*!
+   * \brief Free a HexagonBuffer.
+   * \param ptr Address of the HexagonBuffer as returned by `AllocateHexagonBuffer`.
+   */
+  void FreeHexagonBuffer(void* ptr) {
+    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);
+    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) {
+    auto buf = std::make_unique<HexagonBuffer>(std::forward<Args>(args)...);
+    void* ptr = buf->GetPointer();
+    hexagon_buffer_map_.insert({ptr, std::move(buf)});
+    return ptr;
+  }
+
+  //! \brief Returns whether the HexagonBuffer is in the map.
+  size_t count(void* ptr) { return hexagon_buffer_map_.count(ptr); }
+
+  //! \brief Returns an iterator to the HexagonBuffer within the map.
+  HexagonBuffer* find(void* ptr) {
+    auto it = hexagon_buffer_map_.find(ptr);
+    if (it != hexagon_buffer_map_.end()) {
+      return it->second.get();
+    }
+    return nullptr;
+  }
+
+ private:
+  //! \brief Contains the HexagonBuffer objects managed by this class.
+  std::unordered_map<void*, std::unique_ptr<HexagonBuffer>> hexagon_buffer_map_;
+};
+
+}  // namespace hexagon
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_HEXAGON_HEXAGON_BUFFER_MANAGER_H_

src/runtime/hexagon/hexagon_device_api.cc

@@ -32,7 +32,6 @@
 #include <cstring>
 
 #include "../workspace_pool.h"
-#include "hexagon_buffer.h"
 #include "hexagon_common.h"
 
 namespace tvm {
@@ -74,14 +73,14 @@ void* HexagonDeviceAPI::AllocDataSpace(Device dev, int ndim, const int64_t* shap
   }
 
   if (ndim == 0) {
-    return AllocateHexagonBuffer(typesize, alignment, mem_scope);
+    return hexbuffs.AllocateHexagonBuffer(typesize, alignment, mem_scope);
   } else if (ndim == 1) {
     size_t nbytes = shape[0] * typesize;
-    return AllocateHexagonBuffer(nbytes, alignment, mem_scope);
+    return hexbuffs.AllocateHexagonBuffer(nbytes, alignment, mem_scope);
   } else if (ndim == 2) {
     size_t nallocs = shape[0];
     size_t nbytes = shape[1] * typesize;
-    return AllocateHexagonBuffer(nallocs, nbytes, alignment, mem_scope);
+    return hexbuffs.AllocateHexagonBuffer(nallocs, nbytes, alignment, mem_scope);
   } else {
     LOG(FATAL) << "Hexagon Device API supports only 1d and 2d allocations, but received ndim = "
                << ndim;
@@ -97,13 +96,13 @@ void* HexagonDeviceAPI::AllocDataSpace(Device dev, size_t nbytes, size_t alignme
   if (alignment < kHexagonAllocAlignment) {
     alignment = kHexagonAllocAlignment;
   }
-  return 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;
-  FreeHexagonBuffer(ptr);
+  hexbuffs.FreeHexagonBuffer(ptr);
 }
 
 // WorkSpace: runtime allocations for Hexagon
@@ -119,7 +118,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(hexagon_buffer_map_.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);
 }
@@ -143,13 +142,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* {
-    auto it = this->hexagon_buffer_map_.find(ptr);
-    if (it != this->hexagon_buffer_map_.end()) {
-      return it->second.get();
-    }
-    return nullptr;
-  };
+  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);
@@ -172,14 +165,6 @@ void HexagonDeviceAPI::CopyDataFromTo(const void* from, size_t from_offset, void
   memcpy(static_cast<char*>(to) + to_offset, static_cast<const char*>(from) + from_offset, size);
 }
 
-void HexagonDeviceAPI::FreeHexagonBuffer(void* ptr) {
-  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);
-  hexagon_buffer_map_.erase(it);
-}
-
 TVM_REGISTER_GLOBAL("device_api.hexagon.mem_copy").set_body([](TVMArgs args, TVMRetValue* rv) {
   void* dst = args[0];
   void* src = args[1];

src/runtime/hexagon/hexagon_device_api.h

@@ -30,6 +30,7 @@
 #include <vector>
 
 #include "hexagon_buffer.h"
+#include "hexagon_buffer_manager.h"
 
 namespace tvm {
 namespace runtime {
@@ -72,7 +73,7 @@ class HexagonDeviceAPI final : public DeviceAPI {
    */
   void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final;
 
-  //! Erase from tracked hexagon_buffer_map and free
+  //! Erase from HexagonBufferManager and free
   void FreeWorkspace(Device dev, void* data) final;
 
   /*!
@@ -127,18 +128,6 @@ class HexagonDeviceAPI final : public DeviceAPI {
                       TVMStreamHandle stream) final;
 
  private:
-  /*! \brief Helper to allocate a HexagonBuffer and register the result
-   *  in the owned buffer map.
-   *  \return Raw data storage managed by the hexagon buffer
-   */
-  template <typename... Args>
-  void* AllocateHexagonBuffer(Args&&... args) {
-    auto buf = std::make_unique<HexagonBuffer>(std::forward<Args>(args)...);
-    void* ptr = buf->GetPointer();
-    hexagon_buffer_map_.insert({ptr, std::move(buf)});
-    return ptr;
-  }
-
   /*! \brief Helper to check if the device type is valid for the Hexagon Device API
    *  \return Boolean indicating whether the device type is valid
    */
@@ -148,12 +137,8 @@ class HexagonDeviceAPI final : public DeviceAPI {
            (DLDeviceType(dev.device_type) == kDLCPU);
   }
 
-  /*! \brief Helper to free a HexagonBuffer and unregister the result
-   *  from the owned buffer map.
-   */
-  void FreeHexagonBuffer(void* ptr);
-  //! Lookup table for the HexagonBuffer managing an allocation.
-  std::unordered_map<void*, std::unique_ptr<HexagonBuffer>> hexagon_buffer_map_;
+  //! \brief Manages underlying HexagonBuffer allocations
+  HexagonBufferManager hexbuffs;
 };
 }  // namespace hexagon
 }  // namespace runtime