FIX: Buddy Allocator Free with Merge feature

paddle/memory/detail/buddy_allocator.cc

@@ -89,9 +89,8 @@ void BuddyAllocator::Free(void* p) {
                             block->index(cache_));
 
     // Invalidate GPU allocation from cache
-    if (system_allocator_->UseGpu()) {
     cache_.invalidate(block);
-    }
+
     return;
   }
 
@@ -104,12 +103,35 @@ void BuddyAllocator::Free(void* p) {
   if (block->has_right_buddy(cache_)) {
     DLOG(INFO) << "Merging this block " << block << " with its right buddy "
                << block->right_buddy(cache_);
+
+    auto right_buddy = block->right_buddy(cache_);
+
+    if (right_buddy->type(cache_) == MemoryBlock::FREE_CHUNK) {
+      // Take away right buddy from pool
+      pool_.erase({right_buddy->index(cache_), right_buddy->total_size(cache_),
+                   right_buddy});
+
+      // merge its right buddy to the block
+      block->merge(cache_, right_buddy);
+    }
   }
 
   // Trying to merge the left buddy
   if (block->has_left_buddy(cache_)) {
     DLOG(INFO) << "Merging this block " << block << " with its left buddy "
                << block->left_buddy(cache_);
+
+    auto left_buddy = block->left_buddy(cache_);
+
+    if (left_buddy->type(cache_) == MemoryBlock::FREE_CHUNK) {
+      // Take away right buddy from pool
+      pool_.erase({left_buddy->index(cache_), left_buddy->total_size(cache_),
+                   left_buddy});
+
+      // merge the block to its left buddy
+      left_buddy->merge(cache_, block);
+      block = left_buddy;
+    }
   }
 
   // Dumping this block into pool
@@ -167,13 +189,16 @@ BuddyAllocator::PoolSet::iterator BuddyAllocator::FindExistChunk(size_t size) {
 
   while (1) {
     auto it = pool_.lower_bound({index, size, nullptr});
+
+    // no match chunk memory
     if (it == pool_.end()) return it;
 
     if (std::get<0>(*it) > index) {
+      // find suitable one
       if (std::get<1>(*it) >= size) {
         return it;
       }
-
+      // update and continue
       index = std::get<0>(*it);
       continue;
     }

paddle/memory/detail/buddy_allocator.h

@@ -42,14 +42,14 @@ class BuddyAllocator {
   void Free(void*);
   size_t Used();
 
- public:
+ private:
   // Disable copy and assignment.
   BuddyAllocator(const BuddyAllocator&) = delete;
   BuddyAllocator& operator=(const BuddyAllocator&) = delete;
 
- private:
-  // Tuple type: allocator index, memory size, memory address
+  // Tuple (allocator index, memory size, memory address)
   using IndexSizeAddress = std::tuple<size_t, size_t, void*>;
+  // Each element in PoolSet is a free allocation
   using PoolSet = std::set<IndexSizeAddress>;
 
   /*! \brief Allocate fixed-size memory from system */
@@ -57,7 +57,6 @@ class BuddyAllocator {
 
   /*! \brief If existing chunks are not suitable, refill pool */
   PoolSet::iterator RefillPool();
-
   /**
    *  \brief Find the suitable chunk from existing pool
    *
@@ -77,13 +76,19 @@ class BuddyAllocator {
   size_t max_chunk_size_;  // the maximum size of each chunk
 
  private:
+  /**
+   * \brief A list of free allocation
+   *
+   * \note  Only store free chunk memory in pool
+   */
   PoolSet pool_;
 
  private:
-  // Unify the metadata format between GPU and CPU allocations
+  /*! Unify the metadata format between GPU and CPU allocations */
   MetadataCache cache_;
 
  private:
+  /*! Allocate CPU/GPU memory from system */
   SystemAllocator* system_allocator_;
   std::mutex mutex_;
 };

paddle/memory/detail/memory_block.cc

@@ -91,8 +91,8 @@ void MemoryBlock::split(MetadataCache& cache, size_t size) {
 
 void MemoryBlock::merge(MetadataCache& cache, MemoryBlock* right_buddy) {
   // only free blocks can be merged
-  PADDLE_ASSERT(type(cache) == FREE_MEMORY);
-  PADDLE_ASSERT(right_buddy->type(cache) == FREE_MEMORY);
+  PADDLE_ASSERT(type(cache) == FREE_CHUNK);
+  PADDLE_ASSERT(right_buddy->type(cache) == FREE_CHUNK);
 
   auto metadata = cache.load(this);
 

paddle/platform/CMakeLists.txt

@@ -6,4 +6,4 @@ nv_library(gpu_info SRCS gpu_info.cc DEPS gflags)
 cc_library(place SRCS place.cc)
 cc_test(place_test SRCS place_test.cc DEPS place glog gflags)
 
-cc_library(dynamic_loader SRCS dynload/dynamic_loader.cc)
+cc_library(dynamic_loader SRCS dynload/dynamic_loader.cc DEPS gflags)