diff --git a/paddle/memory/detail/buddy_allocator.cc b/paddle/memory/detail/buddy_allocator.cc index aa5b6b557c5ddd10cb1d080bc71c36483d8f915c..9f334a7048fc688da53d007c8ca2af023ca61614 100644 --- a/paddle/memory/detail/buddy_allocator.cc +++ b/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); - } + 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; } diff --git a/paddle/memory/detail/buddy_allocator.h b/paddle/memory/detail/buddy_allocator.h index ecf23b77ae8ff5f49d9a6871a722aecb34c13e67..2fd9c8162a98aab51ca6cdb706050aefc0978ef9 100644 --- a/paddle/memory/detail/buddy_allocator.h +++ b/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; + // Each element in PoolSet is a free allocation using PoolSet = std::set; /*! \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_; }; diff --git a/paddle/memory/detail/memory_block.cc b/paddle/memory/detail/memory_block.cc index eaa97e7b4ad3d96d4d3256edb272b58e9d17423f..bc67bcef0fdf9358460d72da5f443b48744c6838 100644 --- a/paddle/memory/detail/memory_block.cc +++ b/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); diff --git a/paddle/platform/CMakeLists.txt b/paddle/platform/CMakeLists.txt index 2f3d1c061e37ccdd7aec2b9c4b34ae5f5b8b08c1..0ad11f1b106be2a40a0143eb53e4f7a1317d3486 100644 --- a/paddle/platform/CMakeLists.txt +++ b/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)