[DIPU] refactor: rewrite AsyncResourcePool

dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBFCachingAllocator.cpp

@@ -8,7 +8,6 @@
 #include <vector>
 
 #include "DIPUCachingAllocator.h"
-#include "DIPUSpinMutex.h"
 
 namespace dipu {
 
@@ -131,7 +130,7 @@ class BFCachingAllocatorImpl {
   using StreamSetHandle = std::unique_ptr<StreamSet>;
   std::vector<StreamSetHandle> streamSets_;
 
-  using mutex_t = SpinMutex;
+  using mutex_t = std::mutex;
   mutable mutex_t mut_;
 
   static size_t roundBytes(size_t nbytes) {
@@ -409,35 +408,37 @@ class BFCachingAllocator : public CacheAllocator {
  private:
   void restore() const {
     std::lock_guard<mutex_t> lk(resource_pool_mutex_);
-    while (async_mem_pool()->ready()) {
-      const auto block = async_mem_pool()->get();
-      void* ptr = std::get<0>(block);
-      int id = static_cast<int>(std::get<1>(block));
+
+    auto& pool = *async_mem_pool();
+    for (auto item = pool.pop(); item; item = pool.pop()) {
+      auto [ptr, id] = item.value();
       DIPU_DEBUG_ALLOCATOR(
           8, "BFCachingAllocator: "
                  << __FUNCTION__ << " ,ptr:" << ptr << " ,id:" << id
                  << " ,allocator:" << this << ", device:" << device()
                  << ", async_pool.size:" << async_mem_pool()->size());
-      impl->releaseRaw(ptr, id);
+      impl->releaseRaw(ptr, static_cast<int>(id));
     }
+
     set_memory_reserved(impl->memory_reserved());
   }
 
   void empty_resource_pool() const {
     std::lock_guard<mutex_t> lk(resource_pool_mutex_);
-    while (!async_mem_pool()->empty()) {
-      if (!async_mem_pool()->ready()) {
+    auto& pool = *async_mem_pool();
+    while (not pool.empty()) {
+      auto item = pool.pop();
+      if (not item) {
         std::this_thread::yield();
         continue;
       }
-      const auto block = async_mem_pool()->get();
-      void* ptr = std::get<0>(block);
-      int id = static_cast<int>(std::get<1>(block));
+
+      auto [ptr, id] = item.value();
       DIPU_DEBUG_ALLOCATOR(
           8, "BFCachingAllocator: " << __FUNCTION__ << " ,ptr:" << ptr
                                     << " ,id:" << id << " ,allocator:" << this
                                     << ", device:" << device());
-      impl->releaseRaw(ptr, id);
+      impl->releaseRaw(ptr, static_cast<int>(id));
     }
   }
 
@@ -481,15 +482,8 @@ class BFCachingAllocator : public CacheAllocator {
                                   << ", device:" << allocator_->device());
       if (allocator_->impl) {
         if (ptr()) {
-          std::deque<DIPUEvent> events;
-          for (auto const& stream : streams()) {
-            events.emplace_back();
-            DIPU_DEBUG_ALLOCATOR(8, "BFCachingAllocator: record to stream:"
-                                        << stream.rawstream());
-            events.back().record(stream);
-          }
-          allocator_->async_mem_pool()->add(std::make_tuple(ptr(), id_),
-                                            events);
+          allocator_->async_mem_pool()->put(std::make_tuple(ptr(), id_),
+                                            streams_to_events());
           allocator_->set_memory_allocated(allocator_->memory_allocated() -
                                            nbytes_);
         }

dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBSCachingAllocator.cpp

@@ -140,10 +140,9 @@ class BSCachingAllocator : public CacheAllocator {
   void empty_resource_pool() const {
     DIPU_DEBUG_ALLOCATOR(
         8, "BSCachingAllocator::empty_resource_pool ,allocator:" << this);
-    while (!async_mem_pool()->empty()) {
-      if (async_mem_pool()->ready()) {
-        flush_mem_pool();
-      } else {
+
+    while (not async_mem_pool()->empty()) {
+      if (not flush_mem_pool()) {
         std::this_thread::yield();
       }
     }
@@ -179,13 +178,17 @@ class BSCachingAllocator : public CacheAllocator {
 
   void release_all_memory() const override { release_all_memory_impl(); }
 
-  void flush_mem_pool() const {
+  bool flush_mem_pool() const {
     DIPU_DEBUG_ALLOCATOR(
         8, "BSCachingAllocator::flush_mem_pool allocator:" << this);
-    while (async_mem_pool()->ready()) {
-      auto mem = async_mem_pool()->get();
-      restore(std::get<1>(mem), std::get<0>(mem));
+
+    auto& pool = *async_mem_pool();
+    auto done = false;
+    for (auto item = pool.pop(); item; item = pool.pop(), done = true) {
+      auto [ptr, size] = item.value();
+      restore(size, ptr);
     }
+    return done;
   }
 
   struct Context : public DataPtrContextBase {
@@ -199,14 +202,8 @@ class BSCachingAllocator : public CacheAllocator {
                                            << ", ptr:" << ptr()
                                            << ", size_:" << size());
       if (allocator_->impl) {
-        std::deque<DIPUEvent> events;
-        for (const auto& item : streams()) {
-          events.emplace_back();
-          events.back().record(item);
-        }
-
-        allocator_->async_mem_pool()->add(std::make_tuple(ptr(), size()),
-                                          events);
+        allocator_->async_mem_pool()->put(std::make_tuple(ptr(), size()),
+                                          streams_to_events());
         allocator_->set_memory_allocated(allocator_->memory_allocated() -
                                          real_size_);
         allocator_->flush_mem_pool();

dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUCachingAllocator.h

@@ -5,9 +5,9 @@
 #include <c10/core/Device.h>
 #include <c10/util/flat_hash_map.h>
 
-#include "DIPUAsyncResourcePool.h"
 #include "DIPUCachingAllocatorUtils.h"
 #include "DIPURawAllocator.h"
+#include "async_resource_pool.h"
 
 namespace dipu {
 
@@ -42,7 +42,7 @@ const MemoryAlignmentStrategy* getMemoryAlignmentStrategy();
 void setMemoryAlignmentStrategy(
     const MemoryAlignmentStrategy* memoryAlignStrategy);
 
-using AsyncMemPool = AsyncResourcePool<std::tuple<void*, size_t>>;
+using AsyncMemPool = AsyncResourceQueue<std::tuple<void*, size_t>>;
 
 class MemStats {
  private:
@@ -148,6 +148,15 @@ class DIPU_API CacheAllocator : public c10::Allocator, public MemStats {
 
     ska::flat_hash_set<DIPUStream>& streams() { return streams_; }
 
+    std::vector<DIPUEvent> streams_to_events() const {
+      auto index = std::size_t{};
+      auto events = std::vector<DIPUEvent>(streams_.size());
+      for (auto& stream : streams_) {
+        events[index++].record(stream);
+      }
+      return events;
+    }
+
     const CacheAllocator* allocator() { return allocator_; }
 
     void* ptr() { return ptr_; }
@@ -233,9 +242,7 @@ c10::Allocator* get_allocator(int device_id, c10::Allocator* raw_allocator) {
   namespace name##device_type {                                                \
     static allocator_details::RawAllocator<at::DeviceType::device_type>::type  \
         raw_allocator;                                                         \
-    using AsyncMemPool =                                                       \
-        AsyncResourcePoolImpl<std::tuple<void*, size_t>,                       \
-                              at::DeviceType::device_type, priority>;          \
+    using AsyncMemPool = AsyncResourceQueue<std::tuple<void*, size_t>>;        \
     static const std::function<c10::Allocator*(int)> allocator_get_fn =        \
         std::bind(                                                             \
             allocator_details::get_allocator<CachingAllocator, AsyncMemPool>,  \

dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPURawCachingAllocator.cpp

@@ -17,17 +17,13 @@ class RawCachingAllocator : public CacheAllocator {
     Context(const CacheAllocator* allocator, void* ptr, size_t size,
             size_t real_size)
         : DataPtrContextBase(allocator, ptr, size), real_size_(real_size) {}
+
     ~Context() {
-      std::deque<DIPUEvent> events;
-      for (const auto& item : streams()) {
-        events.emplace_back();
-        events.back().record(item);
-      }
-      auto allocator_ = static_cast<const RawCachingAllocator*>(allocator());
-      allocator_->async_mem_pool()->add(std::make_tuple(ptr(), size()), events);
-      allocator_->set_memory_allocated(allocator_->memory_allocated() -
-                                       real_size_);
-      allocator_->empty_cache();
+      auto alloc = static_cast<const RawCachingAllocator*>(allocator());
+      alloc->async_mem_pool()->put(std::make_tuple(ptr(), size()),
+                                   streams_to_events());
+      alloc->set_memory_allocated(alloc->memory_allocated() - real_size_);
+      alloc->empty_cache();
     }
     size_t real_size_ = 0;
   };
@@ -51,17 +47,13 @@ class RawCachingAllocator : public CacheAllocator {
 
   void empty_cache() const override {
     DIPU_DEBUG_ALLOCATOR(8, "RawCachingAllocator: empty_cache");
-    while (!async_mem_pool()->empty()) {
-      if (async_mem_pool()->ready()) {
-        auto mem = async_mem_pool()->get();
-        void* ptr = std::get<0>(mem);
-        size_t size = std::get<1>(mem);
-        size_t nbytes = getAllocateSize(size);
-        raw_allocator()->raw_deallocate(ptr);
-        set_memory_reserved(memory_reserved() - nbytes);
-      } else {
-        std::this_thread::yield();
-      }
+
+    auto& pool = *async_mem_pool();
+    for (auto item = pool.pop(); item; item = pool.pop()) {
+      auto [ptr, size] = item.value();
+      auto nbytes = getAllocateSize(size);
+      raw_allocator()->raw_deallocate(ptr);
+      set_memory_reserved(memory_reserved() - nbytes);
     }
   }