[CustomDevice] Fix misalignment of memory allocated by buddy allocator (#59247)

ronny1996 · web-flow · commit 2c7368d4cc91 · 2023-11-29T13:53:59.000+08:00
* [CustomDevice] Fix misalignment of memory allocated by buddy allocator

* remove the resevered size in AvailableAllocSize
diff --git a/paddle/fluid/memory/allocation/buddy_allocator.cc b/paddle/fluid/memory/allocation/buddy_allocator.cc
@@ -52,6 +52,7 @@ BuddyAllocator::BuddyAllocator(
       return phi::DeviceManager::GetReallocSize(
           platform::PlaceHelper::CreatePlace(dev_type));
     };
+    use_custom_device_ = true;
   } else {
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
     init_allocate_size_func_ = &platform::GpuInitAllocSize;
@@ -89,6 +90,11 @@ void* BuddyAllocator::Alloc(size_t unaligned_size) {
   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_
@@ -116,22 +122,42 @@ void* BuddyAllocator::Alloc(size_t unaligned_size) {
       return nullptr;
     }
   } else {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+    if (use_custom_device_) {
+      VLOG(10) << "Allocation from existing memory block " << std::get<2>(*it)
+               << " at address " << std::get<2>(*it);
+    } else {
+      VLOG(10) << "Allocation from existing memory block " << std::get<2>(*it)
+               << " at address "
+               << reinterpret_cast<MemoryBlock*>(std::get<2>(*it))->Data();
+    }
+#else
     VLOG(10) << "Allocation from existing memory block " << std::get<2>(*it)
              << " at address "
              << reinterpret_cast<MemoryBlock*>(std::get<2>(*it))->Data();
+#endif
   }
 
   total_used_ += size;
   total_free_ -= size;
 
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  if (use_custom_device_) {
+    return SplitToAlloc(it, size);
+  }
+#endif
   // split the allocation and return data for use
   return reinterpret_cast<MemoryBlock*>(SplitToAlloc(it, size))->Data();
 }
 
 void BuddyAllocator::Free(void* p) {
   // Point back to metadata
   auto block = static_cast<MemoryBlock*>(p)->Metadata();
-
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  if (use_custom_device_) {
+    block = static_cast<MemoryBlock*>(p);
+  }
+#endif
   // Acquire the allocator lock
   std::lock_guard<std::mutex> lock(mutex_);
 
@@ -236,7 +262,11 @@ void* BuddyAllocator::SystemAlloc(size_t size) {
 
   static_cast<MemoryBlock*>(p)->Init(
       &cache_, MemoryBlock::HUGE_CHUNK, index, size, nullptr, nullptr);
-
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  if (use_custom_device_) {
+    return p;
+  }
+#endif
   return static_cast<MemoryBlock*>(p)->Data();
 }
 
diff --git a/paddle/fluid/memory/allocation/buddy_allocator.h b/paddle/fluid/memory/allocation/buddy_allocator.h
@@ -126,6 +126,7 @@ class BuddyAllocator {
   std::mutex mutex_;
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
   std::function<size_t()> init_allocate_size_func_, re_allocate_size_func_;
+  bool use_custom_device_ = false;
 #endif
 };
 
diff --git a/paddle/phi/backends/device_base.cc b/paddle/phi/backends/device_base.cc
@@ -229,15 +229,12 @@ size_t DeviceInterface::AvailableAllocSize(size_t dev_id) {
   size_t total = 0;
   size_t available = 0;
   MemoryStats(dev_id, &total, &available);
-  size_t reserving =
-      static_cast<size_t>(fraction_reserve_gpu_memory * available);
   // If available size is less than minimum chunk size, no usable memory exists
-  size_t available_to_alloc = available - reserving;
   size_t min_chunk_size = GetMinChunkSize(dev_id);
-  if (available_to_alloc < min_chunk_size) {
-    available_to_alloc = 0;
+  if (available < min_chunk_size) {
+    available = 0;
   }
-  return available_to_alloc;
+  return available;
 }
 
 size_t DeviceInterface::GetInitAllocSize(size_t dev_id) {
