Add Memdbginfo for Allocator (#67041)

wanghuancoder · web-flow · commit 255bed533f49 · 2024-08-07T14:24:32.000+08:00
* add mem dbg info

* refine

* refine
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -2490,7 +2490,15 @@ All parameter, weight, gradient are variables in Paddle.
   m.def("get_variable_tensor", framework::GetVariableTensor);
 
   m.def("_is_program_version_supported", IsProgramVersionSupported);
-
+#if defined(PADDLE_WITH_CUDA)
+  m.def("alloctor_dump", [](const phi::GPUPlace &place) {
+    auto allocator = std::dynamic_pointer_cast<
+        paddle::memory::allocation::AutoGrowthBestFitAllocator>(
+        paddle::memory::allocation::AllocatorFacade::Instance()
+            .GetAutoGrowthAllocator(place));
+    allocator->DumpInfo();
+  });
+#endif
   BindProgramDesc(&m);
   BindBlockDesc(&m);
   BindVarDesc(&m);
diff --git a/paddle/phi/core/memory/allocation/allocator_facade.cc b/paddle/phi/core/memory/allocation/allocator_facade.cc
@@ -250,6 +250,7 @@ class AllocatorFacadePrivate {
           InitAutoGrowthCUDAAllocator(phi::GPUPlace(dev_id),
                                       allow_free_idle_chunk_);
         }
+        auto_growth_allocators_ = allocators_;
 
         // Note(Ruibiao): For GPU multi-stream case without CUDA graph
         // capturing, the 'allocators_' map(place -> Allocator) hold the
@@ -363,6 +364,10 @@ class AllocatorFacadePrivate {
 #endif
   }
 
+  inline const AllocatorMap& GetAutoGrowthAllocatorMap() {
+    return auto_growth_allocators_;
+  }
+
   inline const std::shared_ptr<Allocator>& GetAllocator(const phi::Place& place,
                                                         size_t size) {
     const auto& allocators =
@@ -1596,6 +1601,7 @@ class AllocatorFacadePrivate {
 
   AllocatorStrategy strategy_;
   AllocatorMap allocators_;
+  AllocatorMap auto_growth_allocators_;
   static AllocatorMap zero_size_allocators_;
   static AllocatorMap system_allocators_;
   bool allow_free_idle_chunk_;
@@ -1643,6 +1649,11 @@ const std::shared_ptr<Allocator>& AllocatorFacade::GetAllocator(
       place, /* A non-zero num to choose allocator_ */ 1);
 }
 
+const std::shared_ptr<Allocator>& AllocatorFacade::GetAutoGrowthAllocator(
+    const phi::Place& place) {
+  return GetPrivate()->GetAutoGrowthAllocatorMap().at(place);
+}
+
 void* AllocatorFacade::GetBasePtr(
     const std::shared_ptr<phi::Allocation>& allocation) {
   PADDLE_ENFORCE_EQ(GetAllocatorStrategy(),
diff --git a/paddle/phi/core/memory/allocation/allocator_facade.h b/paddle/phi/core/memory/allocation/allocator_facade.h
@@ -56,6 +56,9 @@ class AllocatorFacade {
   TEST_API const std::shared_ptr<Allocator>& GetAllocator(
       const phi::Place& place);
 
+  TEST_API const std::shared_ptr<Allocator>& GetAutoGrowthAllocator(
+      const phi::Place& place);
+
   void* GetBasePtr(const std::shared_ptr<Allocation>& allocation);
 
   const std::shared_ptr<Allocator>& GetZeroAllocator(const phi::Place& place);
diff --git a/paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.cc b/paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.cc
@@ -42,6 +42,8 @@ PHI_DEFINE_EXPORTED_READONLY_bool(
 PHI_DEFINE_EXPORTED_READONLY_bool(print_allocator_trace_info,
                                   false,
                                   "print trace memory info");
+
+PHI_DEFINE_EXPORTED_READONLY_bool(dump_chunk_info, false, "dump chunk info");
 namespace paddle::memory::allocation {
 
 AutoGrowthBestFitAllocator::AutoGrowthBestFitAllocator(
@@ -62,6 +64,27 @@ AutoGrowthBestFitAllocator::AutoGrowthBestFitAllocator(
   VLOG(4) << "chunk_size_:" << chunk_size_;
 }
 
+void AutoGrowthBestFitAllocator::DumpInfo() const {
+  for (auto chunk_it = chunks_.begin(); chunk_it != chunks_.end(); ++chunk_it) {
+    std::cout << "Chunk\t";
+    std::ostringstream oss_used;
+    std::ostringstream oss_free;
+    size_t total = 0, free = 0, used = 0;
+    for (auto &b : chunk_it->blocks_) {
+      total += b.size_;
+      if (b.is_free_) {
+        free += b.size_;
+        oss_free << "(" << b.size_ << "," << b.ptr_ << ")";
+      } else {
+        used += b.size_;
+        oss_used << "(" << b.size_ << "," << b.ptr_ << ")";
+      }
+    }
+    std::cout << total << "\t" << used << "\t" << free << "\t";
+    std::cout << "[" << oss_used.str() << "]\t[" << oss_free.str() << "]"
+              << std::endl;
+  }
+}
 phi::Allocation *AutoGrowthBestFitAllocator::AllocateImpl(
     size_t unaligned_size) {
   phi::RecordEvent record("AutoGrowthBestFitAllocator::Allocate",
@@ -96,6 +119,12 @@ phi::Allocation *AutoGrowthBestFitAllocator::AllocateImpl(
       block_it->is_free_ = false;
     }
   } else {
+    if (FLAGS_dump_chunk_info) {
+      std::cout << "MemDbg memory not enough growth chunk, need size = " << size
+                << std::endl;
+      DumpInfo();
+    }
+
     if (FLAGS_free_when_no_cache_hit) {
       FreeIdleChunks();
     }
@@ -105,6 +134,10 @@ phi::Allocation *AutoGrowthBestFitAllocator::AllocateImpl(
       chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
           underlying_allocator_->Allocate(realloc_size)));
     } catch (BadAlloc &ex) {
+      if (FLAGS_dump_chunk_info) {
+        std::cout << "MemDbg OOM" << std::endl;
+        DumpInfo();
+      }
       if (FLAGS_free_when_no_cache_hit) throw ex;
       FreeIdleChunks();
       chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
@@ -125,6 +158,11 @@ phi::Allocation *AutoGrowthBestFitAllocator::AllocateImpl(
     block_it = --(blocks.end());
     VLOG(2) << "Not found and reallocate " << realloc_size << "("
             << static_cast<void *>(p) << "), and remaining " << remaining_size;
+    if (FLAGS_dump_chunk_info) {
+      std::cout << "MemDbg memory after growth chunk, realloc_size = "
+                << realloc_size << std::endl;
+      DumpInfo();
+    }
   }
   ++total_alloc_times_;
   total_alloc_size_ += size;
@@ -179,6 +217,9 @@ void AutoGrowthBestFitAllocator::FreeImpl(phi::Allocation *allocation) {
 }
 
 uint64_t AutoGrowthBestFitAllocator::FreeIdleChunks() {
+  if (FLAGS_dump_chunk_info) {
+    std::cout << "FreeIdleChunks called" << std::endl;
+  }
   if (!allow_free_idle_chunk_) {
     return 0;
   }
@@ -188,6 +229,10 @@ uint64_t AutoGrowthBestFitAllocator::FreeIdleChunks() {
     if (blocks.size() == 1 && blocks.begin()->is_free_) {
       auto &block = *blocks.begin();
       VLOG(2) << "Free chunk with size " << block.size_;
+      if (FLAGS_dump_chunk_info) {
+        std::cout << "FreeIdleChunks chunk is " << block.size_ << ", "
+                  << block.ptr_ << std::endl;
+      }
       bytes += block.size_;
       free_blocks_.erase(std::make_pair(block.size_, block.ptr_));
       chunk_it = chunks_.erase(chunk_it);
diff --git a/paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.h b/paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.h
@@ -37,6 +37,8 @@ class AutoGrowthBestFitAllocator : public Allocator {
 
   bool IsAllocThreadSafe() const override { return true; }
 
+  void DumpInfo() const;
+
  protected:
   phi::Allocation *AllocateImpl(size_t size) override;
 
