Fix extra_padding bugs. (#62373)

gongweibao · web-flow · commit a5181c549dab · 2024-03-05T11:02:12.000+08:00
diff --git a/paddle/fluid/memory/allocation/allocator_facade.cc b/paddle/fluid/memory/allocation/allocator_facade.cc
@@ -189,6 +189,7 @@ class AllocatorFacadePrivate {
     strategy_ = GetAllocatorStrategy();
     is_stream_safe_cuda_allocator_used_ = false;
     is_cuda_malloc_async_allocator_used_ = false;
+    VLOG(2) << "selected allocator strategy:" << int(strategy_) << std::endl;
     switch (strategy_) {
       case AllocatorStrategy::kNaiveBestFit: {
         InitNaiveBestFitCPUAllocator();
@@ -1289,7 +1290,11 @@ class AllocatorFacadePrivate {
     auto alignment = phi::DeviceManager::GetMinChunkSize(p);
     custom_device_allocators_[p][stream] =
         std::make_shared<AutoGrowthBestFitAllocator>(
-            custom_allocator, alignment, chunk_size, allow_free_idle_chunk_);
+            custom_allocator,
+            alignment,
+            chunk_size,
+            allow_free_idle_chunk_,
+            phi::DeviceManager::GetExtraPaddingSize(p));
   }
 
   void InitAutoGrowthCustomDeviceAllocator(platform::CustomPlace p,
@@ -1303,7 +1308,8 @@ class AllocatorFacadePrivate {
         custom_allocator,
         phi::DeviceManager::GetMinChunkSize(p),
         /*chunk_size=*/chunk_size,
-        allow_free_idle_chunk);
+        allow_free_idle_chunk,
+        phi::DeviceManager::GetExtraPaddingSize(p));
   }
 
   void WrapStreamSafeCustomDeviceAllocatorForDefault() {
diff --git a/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.cc b/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.cc
@@ -20,6 +20,7 @@
 #include "paddle/fluid/memory/allocation/aligned_allocator.h"
 #include "paddle/fluid/platform/flags.h"
 #include "paddle/fluid/platform/profiler/event_tracing.h"
+#include "paddle/phi/backends/device_manager.h"
 
 PADDLE_DEFINE_EXPORTED_READONLY_bool(
     free_idle_chunk,
@@ -40,7 +41,6 @@ PADDLE_DEFINE_EXPORTED_READONLY_bool(
 PADDLE_DEFINE_EXPORTED_READONLY_bool(print_allocator_trace_info,
                                      false,
                                      "print trace memory info");
-
 namespace paddle {
 namespace memory {
 namespace allocation {
@@ -49,11 +49,13 @@ AutoGrowthBestFitAllocator::AutoGrowthBestFitAllocator(
     const std::shared_ptr<Allocator> &underlying_allocator,
     size_t alignment,
     size_t chunk_size,
-    bool allow_free_idle_chunk)
+    bool allow_free_idle_chunk,
+    int extra_padding_size)
     : underlying_allocator_(underlying_allocator),
       alignment_(alignment),
       chunk_size_(std::max(AlignedSize(chunk_size, alignment), alignment)),
-      allow_free_idle_chunk_(allow_free_idle_chunk) {
+      allow_free_idle_chunk_(allow_free_idle_chunk),
+      extra_padding_size_(extra_padding_size) {
   total_alloc_times_ = 0;
   total_alloc_size_ = 0;
   total_free_times_ = 0;
@@ -66,8 +68,11 @@ phi::Allocation *AutoGrowthBestFitAllocator::AllocateImpl(
   platform::RecordEvent record("AutoGrowthBestFitAllocator::Allocate",
                                platform::TracerEventType::UserDefined,
                                9 /*level*/);
-  size_t size = AlignedSize(unaligned_size, alignment_);
-  VLOG(10) << "Allocate " << unaligned_size << " bytes, aligned to " << size;
+
+  size_t size = AlignedSize(unaligned_size + extra_padding_size_, alignment_);
+
+  VLOG(10) << "Allocate " << unaligned_size << " bytes, aligned to " << size
+           << ", extra size " << extra_padding_size_;
 
   std::lock_guard<SpinLock> guard(spinlock_);
   auto iter = free_blocks_.lower_bound(std::make_pair(size, nullptr));
diff --git a/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.h b/paddle/fluid/memory/allocation/auto_growth_best_fit_allocator.h
@@ -33,7 +33,8 @@ class AutoGrowthBestFitAllocator : public Allocator {
       const std::shared_ptr<Allocator> &underlying_allocator,
       size_t alignment,
       size_t chunk_size = 0,
-      bool allow_free_idle_chunk = true);
+      bool allow_free_idle_chunk = true,
+      int extra_padding_size = 0);
 
   bool IsAllocThreadSafe() const override { return true; }
 
@@ -93,6 +94,7 @@ class AutoGrowthBestFitAllocator : public Allocator {
   size_t alignment_;
   size_t chunk_size_;
   bool allow_free_idle_chunk_;
+  int extra_padding_size_;
 
   // stat info
   size_t total_alloc_times_;
diff --git a/paddle/fluid/memory/allocation/buddy_allocator.cc b/paddle/fluid/memory/allocation/buddy_allocator.cc
@@ -60,8 +60,10 @@ BuddyAllocator::BuddyAllocator(
 #endif
   }
 #endif
+
   VLOG(1) << "min_chunk_size_: " << min_chunk_size_
-          << ", max_chunk_size_:" << max_chunk_size_;
+          << ", max_chunk_size_:" << max_chunk_size_
+          << ", extra_padding_size_: " << extra_padding_size_;
 }
 
 BuddyAllocator::~BuddyAllocator() {
@@ -86,15 +88,9 @@ inline size_t align(size_t size, size_t alignment) {
 
 void* BuddyAllocator::Alloc(size_t unaligned_size) {
   // adjust allocation alignment
-
   size_t size =
       align(unaligned_size + sizeof(MemoryBlock::Desc) + extra_padding_size_,
             min_chunk_size_);
-#ifdef PADDLE_WITH_CUSTOM_DEVICE
-  if (use_custom_device_) {
-    size = align(unaligned_size + extra_padding_size_, min_chunk_size_);
-  }
-#endif
   VLOG(10) << "alloc: " << unaligned_size
            << ", padding for desc: " << sizeof(MemoryBlock::Desc)
            << ", extra padding: " << extra_padding_size_
diff --git a/paddle/fluid/memory/allocation/naive_best_fit_allocator.cc b/paddle/fluid/memory/allocation/naive_best_fit_allocator.cc
@@ -459,6 +459,9 @@ class BuddyAllocatorList {
       phi::DeviceManager::SetDevice(device_type_, dev_id);
       platform::CustomPlace place(device_type_, dev_id);
 
+      VLOG(10) << "Init BuddyAllocator on " << place
+               << " with GetExtraPaddingSize "
+               << phi::DeviceManager::GetExtraPaddingSize(place);
       allocators_[dev_id] = std::make_unique<BuddyAllocator>(
           std::unique_ptr<detail::SystemAllocator>(
               new detail::CustomAllocator(device_type_, dev_id)),
diff --git a/paddle/phi/backends/custom/custom_device.cc b/paddle/phi/backends/custom/custom_device.cc
@@ -534,8 +534,8 @@ class CustomDevice : public DeviceInterface {
     if (pimpl_->device_extra_padding_size) {
       PADDLE_ENFORCE_CUSTOM_DEVICE_SUCCESS(
           pimpl_->device_extra_padding_size(device, &padding_size));
-      VLOG(10) << Type() << " extra padding size " << (padding_size >> 20)
-               << "M";
+      VLOG(10) << Type() << " extra padding size:" << padding_size;
+      return padding_size;
     } else {
       return DeviceInterface::GetExtraPaddingSize(dev_id);
     }
