diff --git a/paddle/fluid/memory/allocation/CMakeLists.txt b/paddle/fluid/memory/allocation/CMakeLists.txt index 8a8a7f9430e1ea099bf87aa618bf5ceb35016271..f3666438b60cd71b2b13dd54075803ef4f61b555 100644 --- a/paddle/fluid/memory/allocation/CMakeLists.txt +++ b/paddle/fluid/memory/allocation/CMakeLists.txt @@ -29,9 +29,6 @@ else() cpu_allocator) endif() - -cc_library(naive_managed_allocator SRCS naive_managed_allocator.cc DEPS allocator) -cc_test(naive_managed_allocator_test SRCS naive_managed_allocator_test.cc DEPS naive_managed_allocator) nv_library(pinned_allocator SRCS pinned_allocator.cc DEPS allocator) if (WITH_GPU) set(AllocatorFacadeDeps gpu_info cuda_allocator pinned_allocator cuda_device_guard) @@ -49,7 +46,6 @@ cc_library(allocator_facade SRCS allocator_facade.cc DEPS cpu_allocator locked_allocator best_fit_allocator - naive_managed_allocator aligned_allocator auto_increment_allocator zero_size_allocator @@ -61,6 +57,6 @@ cc_library(allocator_facade SRCS allocator_facade.cc DEPS nv_test(allocation_and_eigen_test SRCS allocation_and_eigen_test.cu DEPS allocator_facade) -cc_test(retry_allocator_test SRCS retry_allocator_test.cc DEPS retry_allocator naive_managed_allocator best_fit_allocator locked_allocator cpu_allocator) +cc_test(retry_allocator_test SRCS retry_allocator_test.cc DEPS retry_allocator best_fit_allocator locked_allocator cpu_allocator) cc_test(allocator_facade_test SRCS allocator_facade_test.cc DEPS allocator_facade) diff --git a/paddle/fluid/memory/allocation/aligned_allocator.cc b/paddle/fluid/memory/allocation/aligned_allocator.cc index ffaeadcbdc6a8c8e390934355ee82a00e80c098b..efae280dbd47a1db476f9c371ba73eac96c30df6 100644 --- a/paddle/fluid/memory/allocation/aligned_allocator.cc +++ b/paddle/fluid/memory/allocation/aligned_allocator.cc @@ -19,14 +19,9 @@ namespace memory { namespace allocation { ThinAlignedAllocator::ThinAlignedAllocator( - std::shared_ptr underlyning_allocator) + std::shared_ptr underlyning_allocator) : underlying_allocator_(std::move(underlyning_allocator)) {} -std::shared_ptr ThinAlignedAllocator::AllocateShared( - size_t size, Allocator::Attr attr) { - return std::shared_ptr(Allocate(size, attr).release()); -} - bool ThinAlignedAllocator::IsAllocThreadSafe() const { return underlying_allocator_->IsAllocThreadSafe(); } diff --git a/paddle/fluid/memory/allocation/aligned_allocator.h b/paddle/fluid/memory/allocation/aligned_allocator.h index 529943dc3da89187a0a360968f16292e62739f9a..835d6b5e5f7840fe86bf91f2ece4fbcc78c33b1e 100644 --- a/paddle/fluid/memory/allocation/aligned_allocator.h +++ b/paddle/fluid/memory/allocation/aligned_allocator.h @@ -70,17 +70,15 @@ class AlignedAllocation : public Allocation { // // NOTE(yy): This could be an over design. If it harms readability of code, it // could be removed later. -class ThinAlignedAllocator : public ManagedAllocator { +class ThinAlignedAllocator : public Allocator { public: explicit ThinAlignedAllocator( - std::shared_ptr underlyning_allocator); - - std::shared_ptr AllocateShared(size_t size, Attr attr) override; + std::shared_ptr underlyning_allocator); bool IsAllocThreadSafe() const; protected: - std::shared_ptr underlying_allocator_; + std::shared_ptr underlying_allocator_; }; // An aligned allocator will allocate `size+kAlignment` allocation and adjust diff --git a/paddle/fluid/memory/allocation/allocator.cc b/paddle/fluid/memory/allocation/allocator.cc index 8833b4e1cd6835b34f929e86fe82a19247f062c8..1aa4e878c4f72067609c26a607ce4eddbfe1c0ff 100644 --- a/paddle/fluid/memory/allocation/allocator.cc +++ b/paddle/fluid/memory/allocation/allocator.cc @@ -24,6 +24,11 @@ bool Allocator::IsAllocThreadSafe() const { return false; } const char* BadAlloc::what() const noexcept { return msg_.c_str(); } +MannualFreeAllocation::~MannualFreeAllocation() { allocator_->Free(this); } +std::unique_ptr MannualFreeAllocator::Allocate( + size_t size, Allocator::Attr attr) { + return std::unique_ptr(AllocateImpl(size, attr)); +} } // namespace allocation } // namespace memory } // namespace paddle diff --git a/paddle/fluid/memory/allocation/allocator.h b/paddle/fluid/memory/allocation/allocator.h index 9c838362d975b5b7c99f5196da7f757286126242..e283ee0616ec636b90a24f36b374fdd3aa997a39 100644 --- a/paddle/fluid/memory/allocation/allocator.h +++ b/paddle/fluid/memory/allocation/allocator.h @@ -121,19 +121,30 @@ class Allocator { virtual bool IsAllocThreadSafe() const; }; -// User need to invoke `Free` or `FreeUniquePtr` manually if allocated by -// a manally managed allocator. -class UnmanagedAllocator : public Allocator { +class MannualFreeAllocator; +class MannualFreeAllocation : public Allocation { public: - virtual void FreeUniquePtr(std::unique_ptr allocation) = 0; + MannualFreeAllocation(MannualFreeAllocator* allocator, void* ptr, size_t size, + platform::Place place) + : Allocation(ptr, size, place), allocator_(allocator) {} + + ~MannualFreeAllocation(); + + private: + MannualFreeAllocator* allocator_; }; -// The allocation will be managed by smart pointers. i.e., users do not need -// to free allocation manually. -class ManagedAllocator : public Allocator { +// User need to invoke `Free` or `FreeUniquePtr` manually if allocated by +// a manally managed allocator. +class MannualFreeAllocator : public Allocator { public: - virtual std::shared_ptr AllocateShared( - size_t size, Allocator::Attr attr = kDefault) = 0; + std::unique_ptr Allocate(size_t size, Attr attr) final; + + protected: + virtual void Free(MannualFreeAllocation* allocation) = 0; + virtual MannualFreeAllocation* AllocateImpl(size_t size, + Allocator::Attr attr) = 0; + friend class MannualFreeAllocation; }; } // namespace allocation diff --git a/paddle/fluid/memory/allocation/allocator_facade.cc b/paddle/fluid/memory/allocation/allocator_facade.cc index 4170e294301d655918744fc31029e0e439dbdf06..44b5ac2bb27d078aba3eb1b4961152cceb33e6b0 100644 --- a/paddle/fluid/memory/allocation/allocator_facade.cc +++ b/paddle/fluid/memory/allocation/allocator_facade.cc @@ -24,7 +24,6 @@ #include "paddle/fluid/memory/allocation/conditional_allocator.h" #include "paddle/fluid/memory/allocation/cpu_allocator.h" #include "paddle/fluid/memory/allocation/locked_allocator.h" -#include "paddle/fluid/memory/allocation/naive_managed_allocator.h" #include "paddle/fluid/memory/allocation/retry_allocator.h" #include "paddle/fluid/memory/allocation/zero_size_allocator.h" #include "paddle/fluid/platform/cpu_info.h" @@ -46,34 +45,28 @@ namespace memory { namespace allocation { // TODO(yy): Dirty code here. This class should be configurable in runtime. -class CPUManagedAllocator : public ManagedAllocator { +class CPUManagedAllocator : public Allocator { public: - CPUManagedAllocator() - : normal_allocator_(NaiveManagedAllocator::Create( - std::unique_ptr(new CPUAllocator()))) {} + CPUManagedAllocator() : normal_allocator_(new CPUAllocator()) {} std::unique_ptr Allocate(size_t size, Attr attr) override { return normal_allocator_->Allocate(size, attr); } - std::shared_ptr AllocateShared(size_t size, Attr attr) override { - return normal_allocator_->AllocateShared(size, attr); - } - bool IsAllocThreadSafe() const override { return true; } private: - std::shared_ptr normal_allocator_; + std::shared_ptr normal_allocator_; }; // TODO(yy): Dirty code here. This class should be configurable in runtime. -class ChunkedManagedAllocator : public ManagedAllocator { +class ChunkedManagedAllocator : public Allocator { public: explicit ChunkedManagedAllocator(std::unique_ptr system_allocator, size_t max_chunk_size, size_t capacity = 1, int64_t retry_time = -1) : max_chunk_size_(max_chunk_size), retry_time_(retry_time) { - raw_allocator_ = NaiveManagedAllocator::Create(std::move(system_allocator)); + raw_allocator_ = std::move(system_allocator); if (max_chunk_size_ == 0) { default_allocator_ = raw_allocator_; @@ -114,11 +107,7 @@ class ChunkedManagedAllocator : public ManagedAllocator { return default_allocator_->Allocate(size, attr); } - std::shared_ptr AllocateShared(size_t size, Attr attr) override { - return default_allocator_->AllocateShared(size, attr); - } - - std::shared_ptr BestFitAllocatorCreator() { + std::shared_ptr BestFitAllocatorCreator() { chunks_.emplace_back(raw_allocator_->Allocate(max_chunk_size_)); auto* allocation = chunks_.back().get(); std::unique_ptr unmanaged_allocator(new LockedAllocator( @@ -127,12 +116,13 @@ class ChunkedManagedAllocator : public ManagedAllocator { if (retry_time_ <= 0) { VLOG(10) << "Create NaiveManagedAllocator without retry"; return std::make_shared>( - NaiveManagedAllocator::Create(std::move(unmanaged_allocator))); + std::move(unmanaged_allocator)); } else { VLOG(10) << "Create RetryAllocator with retry_time " << retry_time_ << "ms"; - return std::make_shared>(RetryAllocator::Create( - std::move(unmanaged_allocator), static_cast(retry_time_))); + auto tmp = std::make_shared( + std::move(unmanaged_allocator), static_cast(retry_time_)); + return std::make_shared>(tmp); } } @@ -142,8 +132,8 @@ class ChunkedManagedAllocator : public ManagedAllocator { size_t max_chunk_size_; int64_t retry_time_; std::vector> chunks_; - std::shared_ptr raw_allocator_; - std::shared_ptr default_allocator_; + std::shared_ptr raw_allocator_; + std::shared_ptr default_allocator_; }; #ifdef PADDLE_WITH_CUDA @@ -193,7 +183,7 @@ class CUDAPinnedManagedAllocator : public ChunkedManagedAllocator { class AllocatorFacadePrivate { public: - std::map> allocators_; + std::map> allocators_; ~AllocatorFacadePrivate() = default; @@ -245,7 +235,8 @@ AllocatorFacade& AllocatorFacade::Instance() { std::shared_ptr AllocatorFacade::AllocShared( const platform::Place& place, size_t size, Allocator::Attr attr) { - return m_->allocators_.at(place)->AllocateShared(size, attr); + return std::shared_ptr( + m_->allocators_.at(place)->Allocate(size, attr).release()); } std::unique_ptr AllocatorFacade::Alloc(const platform::Place& place, diff --git a/paddle/fluid/memory/allocation/auto_increment_allocator.cc b/paddle/fluid/memory/allocation/auto_increment_allocator.cc index 1fac71b8321796417ac57f802e256c8b0df976b7..d198dce32ab489d2892168902f34486476d36f78 100644 --- a/paddle/fluid/memory/allocation/auto_increment_allocator.cc +++ b/paddle/fluid/memory/allocation/auto_increment_allocator.cc @@ -20,20 +20,61 @@ namespace allocation { std::unique_ptr AutoIncrementAllocator::Allocate( size_t size, Allocator::Attr attr) { - return InvokeOrCreateUnderlyingAllocator([&](ManagedAllocator& allocator) { - return allocator.Allocate(size, attr); - }); -} + auto cur = prev_success_allocator_.load(); + size_t retry_count = allocator_num_.load(); + size_t allocator_num = retry_count; + while (retry_count-- > 0) { // until there retry count is zero + try { + auto res = underlying_allocators_[cur]->Allocate(size, attr); + prev_success_allocator_ = cur; + return res; + } catch (BadAlloc&) { + if (++cur >= allocator_num) { + cur = 0; + } + } catch (...) { + // if there is another type of allocation, just rethrow it. + throw; + } + } -std::shared_ptr AutoIncrementAllocator::AllocateShared( - size_t size, Allocator::Attr attr) { - return InvokeOrCreateUnderlyingAllocator([&](ManagedAllocator& allocator) { - return allocator.AllocateShared(size, attr); - }); + // This happens when the first allocator is exhausted and + // there are more than 1 allocation requests + // In this situation, the first allocation request would success + // and the second allocation request would fail if we do not use + // the newly created allocator by the first allocation request. + for (cur = allocator_num; cur < allocator_num_; ++cur) { + try { + auto ret = underlying_allocators_[cur]->Allocate(size, attr); + prev_success_allocator_ = cur; + return ret; + } catch (BadAlloc&) { + } catch (...) { + throw; + } + } + // No suitable allocator + return CreateNewAllocator()->Allocate(size, attr); } bool AutoIncrementAllocator::IsAllocThreadSafe() const { return true; } +std::shared_ptr AutoIncrementAllocator::CreateNewAllocator() { + std::lock_guard guard(mtx_); + auto old_size = allocator_num_.load(); + PADDLE_ENFORCE_LT(old_size, underlying_allocators_.size(), + "Allocator number exceeds capacity %d", + underlying_allocators_.size()); + underlying_allocators_[old_size] = creator_(); + prev_success_allocator_ = old_size; + ++allocator_num_; + PADDLE_ENFORCE( + underlying_allocators_[old_size]->IsAllocThreadSafe(), + "the underlying allocator must be thread safe. This is a program " + "bug."); + return underlying_allocators_[old_size]; +} + } // namespace allocation } // namespace memory } // namespace paddle diff --git a/paddle/fluid/memory/allocation/auto_increment_allocator.h b/paddle/fluid/memory/allocation/auto_increment_allocator.h index f6e1677b4c4269008cc12f47ab039d2f85723e85..ffb5da5e1061aa6f6be07b8921e8d8917f8120c8 100644 --- a/paddle/fluid/memory/allocation/auto_increment_allocator.h +++ b/paddle/fluid/memory/allocation/auto_increment_allocator.h @@ -46,76 +46,20 @@ namespace allocation { // thread-safe std::vector with varying size is hard to implement. // Fortunately, we can get the total GPU memory and each chunk size. // Therefore, we can get the suitable capacity of AutoIncrementAllocator. -class AutoIncrementAllocator : public ManagedAllocator { +class AutoIncrementAllocator : public Allocator { public: // Creator is the method to create ManagedAllocator - using AllocatorCreator = std::function()>; + using AllocatorCreator = std::function()>; explicit AutoIncrementAllocator(AllocatorCreator&& creator, size_t capacity) : creator_(std::move(creator)), underlying_allocators_(capacity) {} + std::unique_ptr Allocate(size_t size, Attr attr) override; - std::shared_ptr AllocateShared(size_t size, Attr attr) override; + bool IsAllocThreadSafe() const override; private: - // NOTE: here use template Callback, it can be inlined when -O3 - template - inline typename std::result_of::type - InvokeOrCreateUnderlyingAllocator(Callback callback) { - auto cur = prev_success_allocator_.load(); - size_t retry_count = allocator_num_.load(); - size_t allocator_num = retry_count; - while (retry_count-- > 0) { // until there retry count is zero - try { - auto res = callback(*underlying_allocators_[cur]); - prev_success_allocator_ = cur; - return std::move(res); - } catch (BadAlloc&) { - if (++cur >= allocator_num) { - cur = 0; - } - } catch (...) { - // if there is another type of allocation, just rethrow it. - throw; - } - } - - // This happens when the first allocator is exhausted and - // there are more than 1 allocation requests - // In this situation, the first allocation request would success - // and the second allocation request would fail if we do not use - // the newly created allocator by the first allocation request. - for (cur = allocator_num; cur < allocator_num_; ++cur) { - try { - auto ret = callback(*underlying_allocators_[cur]); - prev_success_allocator_ = cur; - return std::move(ret); - } catch (BadAlloc&) { - } catch (...) { - throw; - } - } - // No suitable allocator - - ManagedAllocator* new_allocator; - { - std::lock_guard guard(mtx_); - auto old_size = allocator_num_.load(); - PADDLE_ENFORCE_LT(old_size, underlying_allocators_.size(), - "Allocator number exceeds capacity %d", - underlying_allocators_.size()); - underlying_allocators_[old_size] = creator_(); - new_allocator = underlying_allocators_[old_size].get(); - prev_success_allocator_ = old_size; - ++allocator_num_; - } - - PADDLE_ENFORCE( - new_allocator->IsAllocThreadSafe(), - "the underlying allocator must be thread safe. This is a program " - "bug."); - return callback(*new_allocator); - } + std::shared_ptr CreateNewAllocator(); AllocatorCreator creator_; diff --git a/paddle/fluid/memory/allocation/best_fit_allocator.cc b/paddle/fluid/memory/allocation/best_fit_allocator.cc index b903fa437bb5d4aae1a0d152263d2ec70c8034b3..4b17df399e6794c361501f843d95b43a6d45186b 100644 --- a/paddle/fluid/memory/allocation/best_fit_allocator.cc +++ b/paddle/fluid/memory/allocation/best_fit_allocator.cc @@ -45,23 +45,6 @@ BestFitAllocator::BestFitAllocator(Allocation* allocation) {chunk.size_, chunks_.begin()}); } -std::unique_ptr BestFitAllocator::Allocate(size_t size, Attr attr) { - auto highest_set_bit = static_cast(HighestBitPos(size)); - MapIt map_it; - for (; highest_set_bit < free_chunks_.size(); ++highest_set_bit) { - map_it = free_chunks_[highest_set_bit].lower_bound(size); - if (map_it != free_chunks_[highest_set_bit].end()) { - break; - } - } - if (UNLIKELY(highest_set_bit == free_chunks_.size())) { - throw BadAlloc(string::Sprintf( - "Cannot allocate %d, All fragments size is %d", size, FreeSize())); - } - auto chunk_it = SplitChunk(size, highest_set_bit, map_it); - return std::unique_ptr(new BestFitAllocation(this, chunk_it)); -} - size_t BestFitAllocator::FreeSize() const { size_t acc = 0; for (auto& array_item : free_chunks_) { @@ -104,8 +87,30 @@ BestFitAllocator::ListIt BestFitAllocator::SplitChunk(size_t request_size, return to_use_it; } -void BestFitAllocator::FreeUniquePtr(std::unique_ptr allocation) { - auto* bf_allocation = dynamic_cast(allocation.get()); +void BestFitAllocator::InsertFreeNode(const ListIt& it) { + auto pos = static_cast(HighestBitPos(it->size_)); + auto& free_map = free_chunks_[pos]; + free_map.insert({it->size_, it}); +} +void BestFitAllocator::EraseFreeNode(const ListIt& it) { + size_t pos = static_cast(HighestBitPos(it->size_)); + auto& free_map = free_chunks_[pos]; + auto map_it = free_map.find(it->size_); + while (map_it->second != it && map_it != free_map.end()) { + ++map_it; + } + PADDLE_ENFORCE(map_it != free_map.end()); + free_map.erase(map_it); +} +size_t BestFitAllocator::NumFreeChunks() const { + size_t num = 0; + for (auto& array_item : free_chunks_) { + num += array_item.size(); + } + return num; +} +void BestFitAllocator::Free(MannualFreeAllocation* allocation) { + auto* bf_allocation = dynamic_cast(allocation); auto chunk_it = bf_allocation->ChunkIterator(); PADDLE_ENFORCE(!chunk_it->is_free); chunk_it->is_free = true; @@ -132,38 +137,32 @@ void BestFitAllocator::FreeUniquePtr(std::unique_ptr allocation) { InsertFreeNode(chunk_it); } - -void BestFitAllocator::InsertFreeNode(const ListIt& it) { - auto pos = static_cast(HighestBitPos(it->size_)); - auto& free_map = free_chunks_[pos]; - free_map.insert({it->size_, it}); -} -void BestFitAllocator::EraseFreeNode(const ListIt& it) { - size_t pos = static_cast(HighestBitPos(it->size_)); - auto& free_map = free_chunks_[pos]; - auto map_it = free_map.find(it->size_); - while (map_it->second != it && map_it != free_map.end()) { - ++map_it; +MannualFreeAllocation* BestFitAllocator::AllocateImpl(size_t size, + Allocator::Attr attr) { + auto highest_set_bit = static_cast(HighestBitPos(size)); + MapIt map_it; + for (; highest_set_bit < free_chunks_.size(); ++highest_set_bit) { + map_it = free_chunks_[highest_set_bit].lower_bound(size); + if (map_it != free_chunks_[highest_set_bit].end()) { + break; + } } - PADDLE_ENFORCE(map_it != free_map.end()); - free_map.erase(map_it); -} -size_t BestFitAllocator::NumFreeChunks() const { - size_t num = 0; - for (auto& array_item : free_chunks_) { - num += array_item.size(); + if (UNLIKELY(highest_set_bit == free_chunks_.size())) { + throw BadAlloc(string::Sprintf( + "Cannot allocate %d, All fragments size is %d", size, FreeSize())); } - return num; + auto chunk_it = SplitChunk(size, highest_set_bit, map_it); + return new BestFitAllocation(this, chunk_it); } BestFitAllocation::BestFitAllocation( paddle::memory::allocation::BestFitAllocator* allocator, typename details::ChunkList::iterator chunk_it) - : Allocation(reinterpret_cast( - reinterpret_cast(allocator->BasePtr()) + - chunk_it->offset_), - chunk_it->size_, allocator->Place()), - allocator_(allocator), + : MannualFreeAllocation( + allocator, reinterpret_cast( + reinterpret_cast(allocator->BasePtr()) + + chunk_it->offset_), + chunk_it->size_, allocator->Place()), chunk_it_(chunk_it) {} } // namespace allocation } // namespace memory diff --git a/paddle/fluid/memory/allocation/best_fit_allocator.h b/paddle/fluid/memory/allocation/best_fit_allocator.h index 405306bba7b17abe7305740d1b700185076f8ee8..7e299fc4d317eede47471e41be2a7ba25c4c3917 100644 --- a/paddle/fluid/memory/allocation/best_fit_allocator.h +++ b/paddle/fluid/memory/allocation/best_fit_allocator.h @@ -71,7 +71,7 @@ using FreeChunkBin = class BestFitAllocator; // The BestFitAllocation maintain the List Node iterator. -class BestFitAllocation : public Allocation { +class BestFitAllocation : public MannualFreeAllocation { private: using ListIt = typename details::ChunkList::iterator; @@ -81,7 +81,6 @@ class BestFitAllocation : public Allocation { const ListIt& ChunkIterator() const { return chunk_it_; } private: - BestFitAllocator* allocator_; typename details::ChunkList::iterator chunk_it_; }; @@ -99,7 +98,7 @@ class BestFitAllocation : public Allocation { // // To free an allocation, it will set the chunk of allocation to free and merge // the prev-chunk and the next-chunk when possible. -class BestFitAllocator : public UnmanagedAllocator { +class BestFitAllocator : public MannualFreeAllocator { public: explicit BestFitAllocator(Allocation* allocation); @@ -107,9 +106,9 @@ class BestFitAllocator : public UnmanagedAllocator { const platform::Place& Place() const { return allocation_->place(); } - std::unique_ptr Allocate(size_t size, - Attr attr = kDefault) override; - void FreeUniquePtr(std::unique_ptr allocation) override; + // std::unique_ptr Allocate(size_t size, + // Attr attr = kDefault) override; + // void FreeUniquePtr(std::unique_ptr allocation) override; size_t NumFreeChunks() const; @@ -123,6 +122,12 @@ class BestFitAllocator : public UnmanagedAllocator { void EraseFreeNode(const ListIt& it); void InsertFreeNode(const ListIt& it); + protected: + void Free(MannualFreeAllocation* allocation) override; + MannualFreeAllocation* AllocateImpl(size_t size, + Allocator::Attr attr) override; + + private: Allocation* allocation_; // not owned details::ChunkList chunks_; details::FreeChunkBin free_chunks_; diff --git a/paddle/fluid/memory/allocation/buffered_allocator.cc b/paddle/fluid/memory/allocation/buffered_allocator.cc index 18d02f6f657569e61d6585fc6fe53e04fbb1c732..5d5ec710716f994d19b9069e08870d43a216aa40 100644 --- a/paddle/fluid/memory/allocation/buffered_allocator.cc +++ b/paddle/fluid/memory/allocation/buffered_allocator.cc @@ -16,14 +16,14 @@ #include #include #include +#include "paddle/fluid/memory/allocation/underlying_manual_allocation.h" namespace paddle { namespace memory { namespace allocation { -BufferedAllocator::BufferedAllocator(std::unique_ptr&& allocator) { - underlying_allocator_.reset( - dynamic_cast(allocator.release())); +BufferedAllocator::BufferedAllocator(std::unique_ptr &&allocator) + : underlying_allocator_(std::move(allocator)) { PADDLE_ENFORCE_NOT_NULL( underlying_allocator_, "Underlying allocator of BufferedAllocator must be unmanaged"); @@ -34,26 +34,6 @@ BufferedAllocator::BufferedAllocator(std::unique_ptr&& allocator) { BufferedAllocator::~BufferedAllocator() { FreeCache(-1UL); } -std::unique_ptr BufferedAllocator::Allocate(size_t size, - Allocator::Attr attr) { - { - platform::LockGuardPtr guard(mtx_); - auto it = allocations_.lower_bound(size); - if (it != allocations_.end() && it->first < size * 2) { - std::unique_ptr result(std::move(it->second)); - allocations_.erase(it); - return result; - } - } - - try { - return underlying_allocator_->Allocate(size, attr); - } catch (BadAlloc&) { - FreeCache(size); - return underlying_allocator_->Allocate(size, attr); - } -} - void BufferedAllocator::FreeCache(size_t size) { platform::LockGuardPtr guard(mtx_); if (UNLIKELY(size == 0)) return; @@ -61,19 +41,42 @@ void BufferedAllocator::FreeCache(size_t size) { while (!allocations_.empty()) { // free the largest auto it = --allocations_.end(); cur += it->second->size(); - underlying_allocator_->FreeUniquePtr(std::move(it->second)); allocations_.erase(it); if (cur >= size) return; } } -void BufferedAllocator::FreeUniquePtr(std::unique_ptr allocation) { +bool BufferedAllocator::IsAllocThreadSafe() const { + return this->underlying_allocator_->IsAllocThreadSafe(); +} +void BufferedAllocator::Free(MannualFreeAllocation *allocation) { platform::LockGuardPtr guard(mtx_); - allocations_.emplace(allocation->size(), std::move(allocation)); + + std::unique_ptr new_allocation(new UnderlyingManualAllocation( + this, std::move(reinterpret_cast(allocation) + ->allocation_))); + allocations_.emplace(allocation->size(), std::move(new_allocation)); } +MannualFreeAllocation *BufferedAllocator::AllocateImpl(size_t size, + Allocator::Attr attr) { + { + platform::LockGuardPtr guard(mtx_); + auto it = allocations_.lower_bound(size); + if (it != allocations_.end() && it->first < size * 2) { + std::unique_ptr result(std::move(it->second)); + allocations_.erase(it); + return new UnderlyingManualAllocation(this, std::move(result)); + } + } -bool BufferedAllocator::IsAllocThreadSafe() const { - return this->underlying_allocator_->IsAllocThreadSafe(); + try { + return new UnderlyingManualAllocation( + this, underlying_allocator_->Allocate(size, attr)); + } catch (BadAlloc &) { + FreeCache(size); + return new UnderlyingManualAllocation( + this, underlying_allocator_->Allocate(size, attr)); + } } } // namespace allocation diff --git a/paddle/fluid/memory/allocation/buffered_allocator.h b/paddle/fluid/memory/allocation/buffered_allocator.h index 1284661df1aba9ba756799372a6cb8cca1e2bc8a..67b95fe95a137d6076904a13e169856f3627bec9 100644 --- a/paddle/fluid/memory/allocation/buffered_allocator.h +++ b/paddle/fluid/memory/allocation/buffered_allocator.h @@ -29,16 +29,17 @@ namespace allocation { // memory allocation and reuse memory. // BufferedAllocator provides the same thread-safety level as // underlying_allocator_ -class BufferedAllocator : public UnmanagedAllocator { +class BufferedAllocator : public MannualFreeAllocator { public: - explicit BufferedAllocator(std::unique_ptr&& allocator); + explicit BufferedAllocator(std::unique_ptr &&allocator); ~BufferedAllocator(); - std::unique_ptr Allocate( - size_t size, Allocator::Attr attr = Allocator::Attr::kDefault) override; - - void FreeUniquePtr(std::unique_ptr allocation) override; + // std::unique_ptr Allocate( + // size_t size, Allocator::Attr attr = Allocator::Attr::kDefault) + // override; + // + // void FreeUniquePtr(std::unique_ptr allocation) override; bool IsAllocThreadSafe() const override; @@ -48,7 +49,13 @@ class BufferedAllocator : public UnmanagedAllocator { private: void FreeCache(size_t size); - std::unique_ptr underlying_allocator_; + protected: + void Free(MannualFreeAllocation *allocation) override; + MannualFreeAllocation *AllocateImpl(size_t size, + Allocator::Attr attr) override; + + private: + std::unique_ptr underlying_allocator_; std::multimap> allocations_; std::unique_ptr mtx_; }; diff --git a/paddle/fluid/memory/allocation/conditional_allocator.cc b/paddle/fluid/memory/allocation/conditional_allocator.cc index 2df10a89bc2d94920b286e4842f4b6a3d4918375..6a6437a7ff77efbce8140a8828af45cb1b67e195 100644 --- a/paddle/fluid/memory/allocation/conditional_allocator.cc +++ b/paddle/fluid/memory/allocation/conditional_allocator.cc @@ -20,23 +20,27 @@ namespace allocation { ConditionalAllocator& ConditionalAllocator::AddAllocator( std::function func, - std::shared_ptr allocator) { + std::shared_ptr allocator) { underlying_allocators_.emplace_back(std::move(func), std::move(allocator)); return *this; } std::unique_ptr ConditionalAllocator::Allocate( size_t size, Allocator::Attr attr) { - return SelectAndInvoke(size, attr, [&](ManagedAllocator& allocator) { - return allocator.Allocate(size, attr); - }); + for (auto& pair : underlying_allocators_) { + if (pair.first(size, attr)) { + return pair.second->Allocate(size, attr); + } + } + throw BadAlloc("No suitable allocator"); } -std::shared_ptr ConditionalAllocator::AllocateShared( - size_t size, Allocator::Attr attr) { - return SelectAndInvoke(size, attr, [&](ManagedAllocator& allocator) { - return allocator.AllocateShared(size, attr); - }); + +bool ConditionalAllocator::IsAllocThreadSafe() const { + return std::all_of(underlying_allocators_.begin(), + underlying_allocators_.end(), + [](const AllocatorWithCond& allocatorWithCond) { + return allocatorWithCond.second->IsAllocThreadSafe(); + }); } -bool ConditionalAllocator::IsAllocThreadSafe() const { return true; } } // namespace allocation } // namespace memory diff --git a/paddle/fluid/memory/allocation/conditional_allocator.h b/paddle/fluid/memory/allocation/conditional_allocator.h index 46af1099a5c9328aa1ec5e92e6655f277cdd8e93..942c125a4bba175c3ec1ee928194dc046b744652 100644 --- a/paddle/fluid/memory/allocation/conditional_allocator.h +++ b/paddle/fluid/memory/allocation/conditional_allocator.h @@ -38,32 +38,21 @@ namespace allocation { // // else // return true; // }, allocator_c); -class ConditionalAllocator : public ManagedAllocator { +class ConditionalAllocator : public Allocator { public: ConditionalAllocator() = default; - ConditionalAllocator& AddAllocator( - std::function func, - std::shared_ptr allocator); + ConditionalAllocator& AddAllocator(std::function func, + std::shared_ptr allocator); + std::unique_ptr Allocate(size_t size, Attr attr) override; - std::shared_ptr AllocateShared(size_t size, Attr attr) override; + bool IsAllocThreadSafe() const override; private: - template - inline typename std::result_of::type - SelectAndInvoke(size_t size, Attr attr, Callback callback) { - for (auto& pair : underlying_allocators_) { - if (pair.first(size, attr)) { - return callback(*pair.second); - } - } - PADDLE_THROW("No suitable allocator"); - } - - std::vector, - std::shared_ptr>> - underlying_allocators_; + using AllocatorWithCond = + std::pair, std::shared_ptr>; + std::vector underlying_allocators_; }; } // namespace allocation diff --git a/paddle/fluid/memory/allocation/cpu_allocator.cc b/paddle/fluid/memory/allocation/cpu_allocator.cc index 3714c0da746474db1a89f7cbb93a50ebdbb9ee02..35aca11664d9f43f34feaa005ecda877d49c6438 100644 --- a/paddle/fluid/memory/allocation/cpu_allocator.cc +++ b/paddle/fluid/memory/allocation/cpu_allocator.cc @@ -20,21 +20,27 @@ namespace paddle { namespace memory { namespace allocation { -std::unique_ptr CPUAllocator::Allocate(size_t size, Attr attr) { - void* ptr; +CPUAllocation::CPUAllocation( + paddle::memory::allocation::CPUAllocator *allocator, void *ptr, size_t size) + : MannualFreeAllocation(allocator, ptr, size, platform::CPUPlace()) {} + +bool CPUAllocator::IsAllocThreadSafe() const { return true; } + +void CPUAllocator::Free(MannualFreeAllocation *allocation) { + PADDLE_ENFORCE_NOT_NULL(dynamic_cast(allocation)); + free(allocation->ptr()); +} + +MannualFreeAllocation *CPUAllocator::AllocateImpl(size_t size, + Allocator::Attr attr) { + void *ptr; auto status = posix_memalign(&ptr, kAlignment, size); if (UNLIKELY(status) != 0) { throw BadAlloc(string::Sprintf("Cannot allocate cpu memory %d. Errno is %d", size, status)); } - return std::unique_ptr(new CPUAllocation(ptr, size)); -} -void CPUAllocator::FreeUniquePtr(std::unique_ptr allocation) { - PADDLE_ENFORCE_NOT_NULL(dynamic_cast(allocation.get())); - free(allocation->ptr()); + return new CPUAllocation(this, ptr, size); } - -bool CPUAllocator::IsAllocThreadSafe() const { return true; } } // namespace allocation } // namespace memory } // namespace paddle diff --git a/paddle/fluid/memory/allocation/cpu_allocator.h b/paddle/fluid/memory/allocation/cpu_allocator.h index 0852a58e577732a73a5bc6d924204533b6558c68..1c3610e5f344e2c0e6ef389d0d4b5006641295f1 100644 --- a/paddle/fluid/memory/allocation/cpu_allocator.h +++ b/paddle/fluid/memory/allocation/cpu_allocator.h @@ -25,19 +25,21 @@ namespace allocation { // // NOTE(yy): It is no need to use `BestFitAllocator` in CPU. We can import // an open-sourced allocator into Paddle. -class CPUAllocation : public Allocation { +class CPUAllocator; +class CPUAllocation : public MannualFreeAllocation { public: - CPUAllocation(void* ptr, size_t size) - : Allocation(ptr, size, platform::CPUPlace()) {} + CPUAllocation(CPUAllocator* allocator, void* ptr, size_t size); }; -class CPUAllocator : public UnmanagedAllocator { +class CPUAllocator : public MannualFreeAllocator { public: constexpr static size_t kAlignment = 64u; - std::unique_ptr Allocate(size_t size, - Attr attr = kDefault) override; - void FreeUniquePtr(std::unique_ptr allocation) override; bool IsAllocThreadSafe() const override; + + protected: + void Free(MannualFreeAllocation* allocation) override; + MannualFreeAllocation* AllocateImpl(size_t size, + Allocator::Attr attr) override; }; } // namespace allocation } // namespace memory diff --git a/paddle/fluid/memory/allocation/locked_allocator.cc b/paddle/fluid/memory/allocation/locked_allocator.cc index 0b9f1f753146e4c4e97a2a0402e8d0b1524324cd..a6931cff1c8086ae05160e9360ffb4f785f1f563 100644 --- a/paddle/fluid/memory/allocation/locked_allocator.cc +++ b/paddle/fluid/memory/allocation/locked_allocator.cc @@ -14,36 +14,32 @@ #include "paddle/fluid/memory/allocation/locked_allocator.h" #include // NOLINT - +#include "paddle/fluid/memory/allocation/underlying_manual_allocation.h" +#include "paddle/fluid/platform/lock_guard_ptr.h" namespace paddle { namespace memory { namespace allocation { -std::unique_ptr LockedAllocator::Allocate(size_t size, Attr attr) { - if (underlying_allocator_->IsAllocThreadSafe()) { - return underlying_allocator_->Allocate(size, attr); - } else { - std::lock_guard guard(mtx_); - return underlying_allocator_->Allocate(size, attr); - } -} -void LockedAllocator::FreeUniquePtr(std::unique_ptr allocation) { - if (underlying_allocator_->IsAllocThreadSafe()) { - return underlying_allocator_->FreeUniquePtr(std::move(allocation)); - } else { - std::lock_guard guard(mtx_); - return underlying_allocator_->FreeUniquePtr(std::move(allocation)); - } -} bool LockedAllocator::IsAllocThreadSafe() const { return true; } LockedAllocator::LockedAllocator( - std::unique_ptr &&underlying_allocator) { - auto *allocator = - dynamic_cast(underlying_allocator.get()); - PADDLE_ENFORCE_NOT_NULL(allocator); - underlying_allocator.release(); - underlying_allocator_.reset(allocator); + std::unique_ptr &&underlying_allocator) + : underlying_allocator_(std::move(underlying_allocator)) { + PADDLE_ENFORCE_NOT_NULL(underlying_allocator_); + if (!underlying_allocator_->IsAllocThreadSafe()) { + mtx_.reset(new std::mutex()); + } +} +void LockedAllocator::Free(MannualFreeAllocation *allocation) { + platform::LockGuardPtr guard(mtx_); + reinterpret_cast(allocation) + ->allocation_.reset(); +} +MannualFreeAllocation *LockedAllocator::AllocateImpl(size_t size, + Allocator::Attr attr) { + platform::LockGuardPtr guard(mtx_); + return new UnderlyingManualAllocation( + this, underlying_allocator_->Allocate(size, attr)); } } // namespace allocation } // namespace memory diff --git a/paddle/fluid/memory/allocation/locked_allocator.h b/paddle/fluid/memory/allocation/locked_allocator.h index 952622f534477dc01afc51788f58ba142481f7e9..35b151a801b25ae16b77eb450f33561b19ee367d 100644 --- a/paddle/fluid/memory/allocation/locked_allocator.h +++ b/paddle/fluid/memory/allocation/locked_allocator.h @@ -22,17 +22,19 @@ namespace memory { namespace allocation { // A allocator to make underlying allocator thread safe. -class LockedAllocator : public UnmanagedAllocator { +class LockedAllocator : public MannualFreeAllocator { public: - explicit LockedAllocator(std::unique_ptr&& underlying_allocator); - std::unique_ptr Allocate(size_t size, - Attr attr = kDefault) override; - void FreeUniquePtr(std::unique_ptr allocation) override; + explicit LockedAllocator(std::unique_ptr &&underlying_allocator); bool IsAllocThreadSafe() const override; + protected: + void Free(MannualFreeAllocation *allocation) override; + MannualFreeAllocation *AllocateImpl(size_t size, + Allocator::Attr attr) override; + private: - std::unique_ptr underlying_allocator_; - std::mutex mtx_; + std::unique_ptr underlying_allocator_; + std::unique_ptr mtx_; }; } // namespace allocation diff --git a/paddle/fluid/memory/allocation/naive_managed_allocator.cc b/paddle/fluid/memory/allocation/naive_managed_allocator.cc deleted file mode 100644 index 2a61aee84333fe229d4f002fd6b88235dd3859c4..0000000000000000000000000000000000000000 --- a/paddle/fluid/memory/allocation/naive_managed_allocator.cc +++ /dev/null @@ -1,69 +0,0 @@ -// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include "paddle/fluid/memory/allocation/naive_managed_allocator.h" - -namespace paddle { -namespace memory { -namespace allocation { - -NaiveManagedAllocator::NaiveManagedAllocator( - std::unique_ptr &&allocator) { - auto *underlying_allocator = - dynamic_cast(allocator.get()); - PADDLE_ENFORCE_NOT_NULL(underlying_allocator); - allocator.release(); - Init(std::unique_ptr(underlying_allocator)); -} - -NaiveManagedAllocator::NaiveManagedAllocator( - std::unique_ptr &&allocator) { - Init(std::move(allocator)); -} -void NaiveManagedAllocator::Init( - std::unique_ptr &&allocator) { - underlying_allocator_ = std::move(allocator); -} -bool NaiveManagedAllocator::IsAllocThreadSafe() const { - return underlying_allocator_->IsAllocThreadSafe(); -} -std::unique_ptr NaiveManagedAllocator::Allocate(size_t size, - Attr attr) { - std::unique_ptr allocation = - underlying_allocator_->Allocate(size, attr); - return std::unique_ptr( - new NaiveManagedAllocation(std::move(allocation), shared_from_this())); -} -std::shared_ptr NaiveManagedAllocator::AllocateShared(size_t size, - Attr attr) { - std::unique_ptr allocation = - underlying_allocator_->Allocate(size, attr); - return std::shared_ptr( - new NaiveManagedAllocation(std::move(allocation), shared_from_this())); -} - -NaiveManagedAllocation::~NaiveManagedAllocation() { - auto allocator = allocator_.lock(); - if (UNLIKELY(allocator == nullptr)) { - // the allocator is destructed before allocations. - // do nothing. - return; - } - // invoke Free - allocator->UnderlyingAllocator().FreeUniquePtr( - std::move(underlying_allocation_)); -} -} // namespace allocation -} // namespace memory -} // namespace paddle diff --git a/paddle/fluid/memory/allocation/naive_managed_allocator.h b/paddle/fluid/memory/allocation/naive_managed_allocator.h deleted file mode 100644 index 7a4cfdb662a3b6481b1f0071621588ac07471e84..0000000000000000000000000000000000000000 --- a/paddle/fluid/memory/allocation/naive_managed_allocator.h +++ /dev/null @@ -1,76 +0,0 @@ -// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#pragma once -#include -#include "paddle/fluid/memory/allocation/allocator.h" - -namespace paddle { -namespace memory { -namespace allocation { - -// An allocator to wrap an UnmanagedAllocator and make the allocation managed -// by C++ smart ptr. -// -// NOTE: if the NaiveManagedAllocator is destroyed before -// NaiveManagedAllocations, the allocation will never be released. -class NaiveManagedAllocator; -class NaiveManagedAllocation : public Allocation { - public: - NaiveManagedAllocation(std::unique_ptr&& underlying_allocation, - std::shared_ptr allocator) - : Allocation(underlying_allocation->ptr(), underlying_allocation->size(), - underlying_allocation->place()), - underlying_allocation_(std::move(underlying_allocation)), - allocator_(allocator) {} - - ~NaiveManagedAllocation() final; - - private: - std::unique_ptr underlying_allocation_; - std::weak_ptr allocator_; -}; - -class NaiveManagedAllocator - : public ManagedAllocator, - public std::enable_shared_from_this { - public: - template - static std::shared_ptr Create(ARGS... args) { - return std::static_pointer_cast( - std::shared_ptr( - new NaiveManagedAllocator(std::move(args)...))); - } - - inline UnmanagedAllocator& UnderlyingAllocator() { - return *underlying_allocator_; - } - - bool IsAllocThreadSafe() const override; - std::unique_ptr Allocate(size_t size, - Attr attr = kDefault) override; - std::shared_ptr AllocateShared(size_t size, - Attr attr = kDefault) override; - - private: - explicit NaiveManagedAllocator(std::unique_ptr&& allocator); - explicit NaiveManagedAllocator( - std::unique_ptr&& allocator); - void Init(std::unique_ptr&& allocator); - - std::unique_ptr underlying_allocator_; -}; -} // namespace allocation -} // namespace memory -} // namespace paddle diff --git a/paddle/fluid/memory/allocation/naive_managed_allocator_test.cc b/paddle/fluid/memory/allocation/naive_managed_allocator_test.cc deleted file mode 100644 index bb7440d394621f92c8a1bac97c1b5444caacd042..0000000000000000000000000000000000000000 --- a/paddle/fluid/memory/allocation/naive_managed_allocator_test.cc +++ /dev/null @@ -1,82 +0,0 @@ -// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include "paddle/fluid/memory/allocation/naive_managed_allocator.h" -#include // NOLINT -#include -#include // NOLINT -#include -#include "gtest/gtest.h" - -namespace paddle { -namespace memory { -namespace allocation { - -class StubAllocator : public UnmanagedAllocator { - public: - std::unique_ptr Allocate(size_t size, - Attr attr = kDefault) override { - counter_.fetch_add(1); - return std::unique_ptr( - new Allocation(nullptr, size, platform::CPUPlace())); - } - void FreeUniquePtr(std::unique_ptr allocation) override { - counter_.fetch_sub(1); - } - bool IsAllocThreadSafe() const override { return true; } - - std::atomic counter_{0}; -}; - -TEST(NaiveManagedAllocator, main) { - auto allocator = NaiveManagedAllocator::Create( - std::unique_ptr(new StubAllocator())); - - auto th_main = [=] { - std::random_device dev; - std::default_random_engine engine(dev()); - std::uniform_int_distribution dist(0, 1); - - std::vector> allocations; - - for (int j = 0; j < 1024; ++j) { - bool to_insert = static_cast(dist(engine)); - if (to_insert) { - allocations.emplace_back(allocator->AllocateShared(10)); - } else { - if (!allocations.empty()) { - allocations.pop_back(); - } - } - } - }; - - { - std::vector threads; - for (size_t i = 0; i < 1024; ++i) { - threads.emplace_back(th_main); - } - for (auto& th : threads) { - th.join(); - } - } - ASSERT_EQ(reinterpret_cast( - std::dynamic_pointer_cast(allocator) - ->UnderlyingAllocator()) - .counter_, - 0); -} -} // namespace allocation -} // namespace memory -} // namespace paddle diff --git a/paddle/fluid/memory/allocation/retry_allocator.cc b/paddle/fluid/memory/allocation/retry_allocator.cc index 9dc568ef2ab6d4ce79e2c923311a53cf0cede278..68c983c63afaefa86f4d22740e7df366dc642f65 100644 --- a/paddle/fluid/memory/allocation/retry_allocator.cc +++ b/paddle/fluid/memory/allocation/retry_allocator.cc @@ -18,29 +18,25 @@ namespace paddle { namespace memory { namespace allocation { -RetryAllocation::~RetryAllocation() { - auto allocator = retry_allocator_.lock(); - // Allocator is destroyed before allocation. Should not happened usually. - if (UNLIKELY(allocator == nullptr)) return; - allocator->FreeUnderlyingAllocation(std::move(underlying_allocation_)); +bool RetryAllocator::IsAllocThreadSafe() const { + return underlying_allocator_->IsAllocThreadSafe(); } -bool RetryAllocator::IsAllocThreadSafe() const { return true; } - -std::shared_ptr RetryAllocator::AllocateShared( - size_t size, Allocator::Attr attr) { - return std::shared_ptr(AllocateImpl(size, attr)); -} - -std::unique_ptr RetryAllocator::Allocate(size_t size, - Allocator::Attr attr) { - return std::unique_ptr(AllocateImpl(size, attr)); +void RetryAllocator::Free(MannualFreeAllocation* allocation) { + reinterpret_cast(allocation) + ->underlying_allocation_.reset(); + { + // notify all waited allocators, they can try to allocate memory after free. + std::lock_guard lock(mutex_); + cv_.notify_all(); + } } -Allocation* RetryAllocator::AllocateImpl(size_t size, Allocator::Attr attr) { +MannualFreeAllocation* RetryAllocator::AllocateImpl(size_t size, + Allocator::Attr attr) { auto alloc_func = [&, this]() { return new RetryAllocation(underlying_allocator_->Allocate(size, attr), - this->shared_from_this()); + this); }; // In fact, we can unify the code of allocation success and failure // But it would add lock even when allocation success at the first time @@ -73,15 +69,6 @@ Allocation* RetryAllocator::AllocateImpl(size_t size, Allocator::Attr attr) { throw; } } -void RetryAllocator::FreeUnderlyingAllocation( - std::unique_ptr&& allocation) { - underlying_allocator_->FreeUniquePtr(std::move(allocation)); - { - // notify all waited allocators, they can try to allocate memory after free. - std::lock_guard lock(mutex_); - cv_.notify_all(); - } -} } // namespace allocation } // namespace memory diff --git a/paddle/fluid/memory/allocation/retry_allocator.h b/paddle/fluid/memory/allocation/retry_allocator.h index 25461e5423a20425d9d191b5edb68063bb8a796c..3dc4855333699ba60c03001b3c88b1393bf2d186 100644 --- a/paddle/fluid/memory/allocation/retry_allocator.h +++ b/paddle/fluid/memory/allocation/retry_allocator.h @@ -26,52 +26,27 @@ namespace allocation { class RetryAllocator; -class RetryAllocation : public Allocation { +class RetryAllocation : public MannualFreeAllocation { public: RetryAllocation(std::unique_ptr&& underlying_allocation, - const std::shared_ptr& retry_allocator) - : Allocation(underlying_allocation->ptr(), underlying_allocation->size(), - underlying_allocation->place()), - underlying_allocation_(std::move(underlying_allocation)), - retry_allocator_(retry_allocator) {} - - ~RetryAllocation() final; - - private: + MannualFreeAllocator* allocator) + : MannualFreeAllocation(allocator, underlying_allocation->ptr(), + underlying_allocation->size(), + underlying_allocation->place()), + underlying_allocation_(std::move(underlying_allocation)) {} std::unique_ptr underlying_allocation_; - std::weak_ptr retry_allocator_; }; -class RetryAllocator : public ManagedAllocator, - public std::enable_shared_from_this { - private: +class RetryAllocator : public MannualFreeAllocator { + public: RetryAllocator(std::unique_ptr&& allocator, size_t retry_ms) - : underlying_allocator_( - dynamic_cast(allocator.release())), - retry_time_(retry_ms) { + : underlying_allocator_(std::move(allocator)), retry_time_(retry_ms) { EnforceCheck(); } - public: - template - static std::shared_ptr Create(Args... args) { - return std::shared_ptr( - new RetryAllocator(std::forward(args)...)); - } - bool IsAllocThreadSafe() const override; - std::unique_ptr Allocate(size_t size, - Allocator::Attr attr) override; - - std::shared_ptr AllocateShared(size_t size, - Allocator::Attr attr) override; - - void FreeUnderlyingAllocation(std::unique_ptr&& allocation); - private: - Allocation* AllocateImpl(size_t size, Allocator::Attr attr); - void EnforceCheck() { PADDLE_ENFORCE_NOT_NULL( underlying_allocator_.get(), @@ -80,7 +55,13 @@ class RetryAllocator : public ManagedAllocator, "UnderlyingAllocator of RetryAllocator must be thread-safe"); } - std::unique_ptr underlying_allocator_; + protected: + void Free(MannualFreeAllocation* allocation) override; + MannualFreeAllocation* AllocateImpl(size_t size, + Allocator::Attr attr) override; + + private: + std::unique_ptr underlying_allocator_; std::chrono::milliseconds retry_time_; std::mutex mutex_; std::condition_variable cv_; diff --git a/paddle/fluid/memory/allocation/underlying_manual_allocation.h b/paddle/fluid/memory/allocation/underlying_manual_allocation.h new file mode 100644 index 0000000000000000000000000000000000000000..a54aee71a871f35c83c37e97b8c2945abad60d61 --- /dev/null +++ b/paddle/fluid/memory/allocation/underlying_manual_allocation.h @@ -0,0 +1,35 @@ +// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include "paddle/fluid/memory/allocation/allocator.h" + +namespace paddle { +namespace memory { +namespace allocation { + +class UnderlyingManualAllocation : public MannualFreeAllocation { + public: + UnderlyingManualAllocation(MannualFreeAllocator* allocator, + std::unique_ptr allocation) + : MannualFreeAllocation(allocator, allocation->ptr(), allocation->size(), + allocation->place()), + allocation_(std::move(allocation)) {} + std::unique_ptr allocation_; +}; + +} // namespace allocation +} // namespace memory +} // namespace paddle diff --git a/paddle/fluid/memory/allocation/zero_size_allocator.cc b/paddle/fluid/memory/allocation/zero_size_allocator.cc index e6cf754a4698db4a51630e72e2ad4d8960e09a7d..663688e94c304f520521b332eb6145f28c7823d9 100644 --- a/paddle/fluid/memory/allocation/zero_size_allocator.cc +++ b/paddle/fluid/memory/allocation/zero_size_allocator.cc @@ -26,15 +26,10 @@ std::unique_ptr ZeroSizeAllocator::Allocate(size_t size, return underlying_allocator_->Allocate(size, attr); } } -std::shared_ptr ZeroSizeAllocator::AllocateShared( - size_t size, Allocator::Attr attr) { - if (size == 0) { - return std::shared_ptr(new ZeroSizeAllocation(place_)); - } else { - return underlying_allocator_->AllocateShared(size, attr); - } + +bool ZeroSizeAllocator::IsAllocThreadSafe() const { + return underlying_allocator_->IsAllocThreadSafe(); } -bool ZeroSizeAllocator::IsAllocThreadSafe() const { return true; } } // namespace allocation } // namespace memory } // namespace paddle diff --git a/paddle/fluid/memory/allocation/zero_size_allocator.h b/paddle/fluid/memory/allocation/zero_size_allocator.h index 35a4552469f2d3c50fc0df8fab50d99c6bdc08f5..4046c783e798732ffcf440ebfa130caf5302258d 100644 --- a/paddle/fluid/memory/allocation/zero_size_allocator.h +++ b/paddle/fluid/memory/allocation/zero_size_allocator.h @@ -12,10 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include - #pragma once - +#include #include "paddle/fluid/memory/allocation/allocator.h" namespace paddle { @@ -31,18 +29,17 @@ class ZeroSizeAllocation : public Allocation { : Allocation(nullptr, 0, p) {} }; -class ZeroSizeAllocator : public ManagedAllocator { +class ZeroSizeAllocator : public Allocator { public: - ZeroSizeAllocator( - const std::shared_ptr& underlying_allocator, - const platform::Place& p) - : underlying_allocator_(underlying_allocator), place_(p) {} + ZeroSizeAllocator(std::shared_ptr underlying_allocator, + const platform::Place& p) + : underlying_allocator_(std::move(underlying_allocator)), place_(p) {} std::unique_ptr Allocate(size_t size, Attr attr) override; - std::shared_ptr AllocateShared(size_t size, Attr attr) override; + bool IsAllocThreadSafe() const override; private: - std::shared_ptr underlying_allocator_; + std::shared_ptr underlying_allocator_; const platform::Place& place_; };