fix bs allocator bug (#892)

* fix bs allocator bug * minor change
DeepLink-org · Jul 12, 2024 · ee58ddc · ee58ddc
1 parent 4d62b53
commit ee58ddc
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBSCachingAllocator.cpp b/dipu/torch_dipu/csrc_dipu/runtime/core/allocator/DIPUBSCachingAllocator.cpp
@@ -138,6 +138,7 @@ class BSCachingAllocator : public CacheAllocator {
   }
 
   void empty_resource_pool() const {
+    std::lock_guard<mutex_t> lk(mutex_);
     DIPU_DEBUG_ALLOCATOR(
         8, "BSCachingAllocator::empty_resource_pool ,allocator:" << this);
     while (!async_mem_pool()->empty()) {
@@ -180,6 +181,7 @@ class BSCachingAllocator : public CacheAllocator {
   void release_all_memory() const override { release_all_memory_impl(); }
 
   void flush_mem_pool() const {
+    std::lock_guard<mutex_t> lk(mutex_);
     DIPU_DEBUG_ALLOCATOR(
         8, "BSCachingAllocator::flush_mem_pool allocator:" << this);
     while (async_mem_pool()->ready()) {