bugfix: crash when competition (PaddlePaddle#3)

m3ngyang · Mar 24, 2022 · cad1183 · cad1183
1 parent f31494a
commit cad1183
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 4 deletions.
diff --git a/paddle/fluid/memory/allocation/mixed_mem_best_fit_allocator.cc b/paddle/fluid/memory/allocation/mixed_mem_best_fit_allocator.cc
@@ -41,6 +41,7 @@ Allocation* MixedMemBestFitAllocator::AllocateImpl(size_t size) {
       platform::errors::InvalidArgument("Underlying host allocator of "
                                         "MixedMemBestFitAllocator is nullptr"));
 
+  std::lock_guard<std::mutex> lock(mtx_);
   if (!device_allocator_->ReachLimit()) {
     try {
       void* ptr = device_allocator_->Alloc(size);
@@ -83,9 +84,10 @@ Allocation* MixedMemBestFitAllocator::AllocateImpl(size_t size) {
 }
 
 void MixedMemBestFitAllocator::FreeImpl(Allocation* allocation) {
-  auto place = allocation->place();
+  const auto place = allocation->place();
   bool succ = false;
 
+  std::lock_guard<std::mutex> lock(mtx_);
   auto it = devptr2hostptr_.find(allocation->ptr());
   if (it == devptr2hostptr_.end()) {
     device_allocator_->Free(allocation->ptr());

diff --git a/paddle/fluid/memory/allocation/mixed_mem_best_fit_allocator.h b/paddle/fluid/memory/allocation/mixed_mem_best_fit_allocator.h
@@ -74,6 +74,7 @@ class MixedMemBestFitAllocator : public Allocator {
   std::shared_ptr<detail::BuddyAllocator> device_allocator_;
   std::shared_ptr<detail::BuddyAllocator> host_allocator_;
 
+  std::mutex mtx_;
   std::unordered_map<void*, MappedAddr> devptr2hostptr_;
 };
 

diff --git a/paddle/fluid/platform/gpu_info.cc b/paddle/fluid/platform/gpu_info.cc
@@ -510,9 +510,9 @@ class RecordedCudaMallocHelper {
  private:
   explicit RecordedCudaMallocHelper(int dev_id, uint64_t limit_size = 0)
       : dev_id_(dev_id), limit_size_(limit_size) {
-    // if (NeedRecord()) {
-    mtx_.reset(new std::mutex());
-    // }
+    if (NeedRecord()) {
+      mtx_.reset(new std::mutex());
+    }
   }
 
   DISABLE_COPY_AND_ASSIGN(RecordedCudaMallocHelper);