Refine code

Author: reyoung

paddle/fluid/memory/allocation/aligned_allocator.h

@@ -33,8 +33,7 @@ class AlignedAllocation : public Allocation {
                 "kAlignment must be 2^N");
 
  public:
-  AlignedAllocation(std::unique_ptr<Allocation>&& underlying_allocation,
-                    size_t size)
+  AlignedAllocation(AllocationPtr&& underlying_allocation, size_t size)
       : Allocation(AlignedPtr(underlying_allocation->ptr()),
                    size + kAlignment - Offset(underlying_allocation->ptr()),
                    underlying_allocation->place()),
@@ -59,7 +58,7 @@ class AlignedAllocation : public Allocation {
     }
   }
 
-  std::unique_ptr<Allocation> underlying_allocation_;
+  AllocationPtr underlying_allocation_;
 };
 
 // Thin aligned allocator is trivial and used to generate a small size binary.
@@ -87,10 +86,10 @@ template <size_t kAlignment>
 class AlignedAllocator : public ThinAlignedAllocator {
  public:
   using ThinAlignedAllocator::ThinAlignedAllocator;
-  std::unique_ptr<Allocation> Allocate(size_t size, Attr attr) override {
+  AllocationPtr Allocate(size_t size, Attr attr) override {
     auto raw_allocation =
         underlying_allocator_->Allocate(size + kAlignment, attr);
-    return std::unique_ptr<Allocation>(
+    return AllocationPtr(
         new AlignedAllocation<kAlignment>(std::move(raw_allocation), size));
   }
 };

paddle/fluid/memory/allocation/allocator.cc

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 #include "paddle/fluid/memory/allocation/allocator.h"
+#include <functional>
+
 namespace paddle {
 namespace memory {
 namespace allocation {
@@ -24,10 +26,20 @@ bool Allocator::IsAllocThreadSafe() const { return false; }
 
 const char* BadAlloc::what() const noexcept { return msg_.c_str(); }
 
-MannualFreeAllocation::~MannualFreeAllocation() { allocator_->Free(this); }
-std::unique_ptr<Allocation> MannualFreeAllocator::Allocate(
-    size_t size, Allocator::Attr attr) {
-  return std::unique_ptr<Allocation>(AllocateImpl(size, attr));
+AllocationPtr MannualFreeAllocator::Allocate(size_t size,
+                                             Allocator::Attr attr) {
+  auto allocation = AllocateImpl(size, attr);
+  allocation->Deleter =
+      std::bind1st(std::mem_fn(&MannualFreeAllocator::Free), this);
+  return AllocationPtr(allocation);
+}
+void AllocationDeleter::operator()(Allocation* allocation) const {
+  if (allocation->Deleter) {
+    auto deleter = std::move(allocation->Deleter);
+    deleter(allocation);
+  } else {
+    delete allocation;
+  }
 }
 }  // namespace allocation
 }  // namespace memory

paddle/fluid/memory/allocation/allocator.h

@@ -31,6 +31,11 @@ class BadAlloc : public std::exception {
   std::string msg_;
 };
 
+class Allocation;
+struct AllocationDeleter {
+  void operator()(Allocation* allocation) const;
+};
+
 // Allocation is the object holding the actually pointer. Use
 // `Allocation::ptr()` will returns the pointer that allocated.
 //
@@ -67,12 +72,16 @@ class Allocation {
 
   virtual ~Allocation();
 
+  std::function<void(Allocation*)> Deleter;
+
  private:
   void* ptr_;
   size_t size_;
   platform::Place place_;
 };
 
+using AllocationPtr = std::unique_ptr<Allocation, AllocationDeleter>;
+
 // Base interface class of memory Allocator.
 // To allocate a memory, allocator needs two parameters:
 //    1. size of bytes.
@@ -114,36 +123,22 @@ class Allocator {
 
   // Allocate an allocation. Note the return allocation might need to be freed
   // manually if the Allocator is an `UnmanagedAllocator`.
-  virtual std::unique_ptr<Allocation> Allocate(
-      size_t size, Allocator::Attr attr = kDefault) = 0;
+  virtual AllocationPtr Allocate(size_t size,
+                                 Allocator::Attr attr = kDefault) = 0;
 
   // True if the `Allocate` is thread safe.
   virtual bool IsAllocThreadSafe() const;
 };
 
-class MannualFreeAllocator;
-class MannualFreeAllocation : public Allocation {
- public:
-  MannualFreeAllocation(MannualFreeAllocator* allocator, void* ptr, size_t size,
-                        platform::Place place)
-      : Allocation(ptr, size, place), allocator_(allocator) {}
-
-  ~MannualFreeAllocation();
-
- private:
-  MannualFreeAllocator* allocator_;
-};
-
 // User need to invoke `Free` or `FreeUniquePtr` manually if allocated by
 // a manally managed allocator.
 class MannualFreeAllocator : public Allocator {
  public:
-  std::unique_ptr<Allocation> Allocate(size_t size, Attr attr) final;
+  AllocationPtr Allocate(size_t size, Attr attr) final;
 
  protected:
-  virtual void Free(MannualFreeAllocation* allocation) = 0;
-  virtual MannualFreeAllocation* AllocateImpl(size_t size,
-                                              Allocator::Attr attr) = 0;
+  virtual void Free(Allocation* allocation) = 0;
+  virtual Allocation* AllocateImpl(size_t size, Allocator::Attr attr) = 0;
   friend class MannualFreeAllocation;
 };
 

paddle/fluid/memory/allocation/allocator_facade.cc

@@ -49,7 +49,7 @@ class CPUManagedAllocator : public Allocator {
  public:
   CPUManagedAllocator() : normal_allocator_(new CPUAllocator()) {}
 
-  std::unique_ptr<Allocation> Allocate(size_t size, Attr attr) override {
+  AllocationPtr Allocate(size_t size, Attr attr) override {
     return normal_allocator_->Allocate(size, attr);
   }
 
@@ -103,7 +103,7 @@ class ChunkedManagedAllocator : public Allocator {
     raw_allocator_.reset();
   }
 
-  std::unique_ptr<Allocation> Allocate(size_t size, Attr attr) override {
+  AllocationPtr Allocate(size_t size, Attr attr) override {
     return default_allocator_->Allocate(size, attr);
   }
 
@@ -131,7 +131,7 @@ class ChunkedManagedAllocator : public Allocator {
  protected:
   size_t max_chunk_size_;
   int64_t retry_time_;
-  std::vector<std::unique_ptr<Allocation>> chunks_;
+  std::vector<AllocationPtr> chunks_;
   std::shared_ptr<Allocator> raw_allocator_;
   std::shared_ptr<Allocator> default_allocator_;
 };
@@ -236,12 +236,12 @@ AllocatorFacade& AllocatorFacade::Instance() {
 std::shared_ptr<Allocation> AllocatorFacade::AllocShared(
     const platform::Place& place, size_t size, Allocator::Attr attr) {
   return std::shared_ptr<Allocation>(
-      m_->allocators_.at(place)->Allocate(size, attr).release());
+      m_->allocators_.at(place)->Allocate(size, attr).release(),
+      AllocationDeleter());
 }
 
-std::unique_ptr<Allocation> AllocatorFacade::Alloc(const platform::Place& place,
-                                                   size_t size,
-                                                   Allocator::Attr attr) {
+AllocationPtr AllocatorFacade::Alloc(const platform::Place& place, size_t size,
+                                     Allocator::Attr attr) {
   return m_->allocators_.at(place)->Allocate(size, attr);
 }
 

paddle/fluid/memory/allocation/allocator_facade.h

@@ -43,8 +43,8 @@ class AllocatorFacade {
       Allocator::Attr attr = Allocator::kDefault);
 
   // Allocate a unique allocation.
-  std::unique_ptr<Allocation> Alloc(const platform::Place& place, size_t size,
-                                    Allocator::Attr attr = Allocator::kDefault);
+  AllocationPtr Alloc(const platform::Place& place, size_t size,
+                      Allocator::Attr attr = Allocator::kDefault);
 
   // TODO(yy): Allocate a Copy-On-Write allocation?
  private:

paddle/fluid/memory/allocation/auto_increment_allocator.cc

@@ -18,8 +18,8 @@ namespace paddle {
 namespace memory {
 namespace allocation {
 
-std::unique_ptr<Allocation> AutoIncrementAllocator::Allocate(
-    size_t size, Allocator::Attr attr) {
+AllocationPtr AutoIncrementAllocator::Allocate(size_t size,
+                                               Allocator::Attr attr) {
   auto cur = prev_success_allocator_.load();
   size_t retry_count = allocator_num_.load();
   size_t allocator_num = retry_count;

paddle/fluid/memory/allocation/auto_increment_allocator.h

@@ -54,7 +54,7 @@ class AutoIncrementAllocator : public Allocator {
   explicit AutoIncrementAllocator(AllocatorCreator&& creator, size_t capacity)
       : creator_(std::move(creator)), underlying_allocators_(capacity) {}
 
-  std::unique_ptr<Allocation> Allocate(size_t size, Attr attr) override;
+  AllocationPtr Allocate(size_t size, Attr attr) override;
 
   bool IsAllocThreadSafe() const override;
 

paddle/fluid/memory/allocation/best_fit_allocator.cc

@@ -109,7 +109,7 @@ size_t BestFitAllocator::NumFreeChunks() const {
   }
   return num;
 }
-void BestFitAllocator::Free(MannualFreeAllocation* allocation) {
+void BestFitAllocator::Free(Allocation* allocation) {
   auto* bf_allocation = dynamic_cast<BestFitAllocation*>(allocation);
   auto chunk_it = bf_allocation->ChunkIterator();
   PADDLE_ENFORCE(!chunk_it->is_free);
@@ -136,9 +136,9 @@ void BestFitAllocator::Free(MannualFreeAllocation* allocation) {
   }
 
   InsertFreeNode(chunk_it);
+  delete allocation;
 }
-MannualFreeAllocation* BestFitAllocator::AllocateImpl(size_t size,
-                                                      Allocator::Attr attr) {
+Allocation* BestFitAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
   auto highest_set_bit = static_cast<size_t>(HighestBitPos(size));
   MapIt map_it;
   for (; highest_set_bit < free_chunks_.size(); ++highest_set_bit) {
@@ -158,11 +158,10 @@ MannualFreeAllocation* BestFitAllocator::AllocateImpl(size_t size,
 BestFitAllocation::BestFitAllocation(
     paddle::memory::allocation::BestFitAllocator* allocator,
     typename details::ChunkList::iterator chunk_it)
-    : MannualFreeAllocation(
-          allocator, reinterpret_cast<void*>(
-                         reinterpret_cast<uintptr_t>(allocator->BasePtr()) +
-                         chunk_it->offset_),
-          chunk_it->size_, allocator->Place()),
+    : Allocation(reinterpret_cast<void*>(
+                     reinterpret_cast<uintptr_t>(allocator->BasePtr()) +
+                     chunk_it->offset_),
+                 chunk_it->size_, allocator->Place()),
       chunk_it_(chunk_it) {}
 }  // namespace allocation
 }  // namespace memory

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 MannualFreeAllocation {
+class BestFitAllocation : public Allocation {
  private:
   using ListIt = typename details::ChunkList::iterator;
 
@@ -123,9 +123,8 @@ class BestFitAllocator : public MannualFreeAllocator {
   void InsertFreeNode(const ListIt& it);
 
  protected:
-  void Free(MannualFreeAllocation* allocation) override;
-  MannualFreeAllocation* AllocateImpl(size_t size,
-                                      Allocator::Attr attr) override;
+  void Free(Allocation* allocation) override;
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
 
  private:
   Allocation* allocation_;  // not owned

paddle/fluid/memory/allocation/buffered_allocator.cc

@@ -49,33 +49,28 @@ void BufferedAllocator::FreeCache(size_t size) {
 bool BufferedAllocator::IsAllocThreadSafe() const {
   return this->underlying_allocator_->IsAllocThreadSafe();
 }
-void BufferedAllocator::Free(MannualFreeAllocation *allocation) {
+void BufferedAllocator::Free(Allocation *allocation) {
   platform::LockGuardPtr<std::mutex> guard(mtx_);
-
-  std::unique_ptr<Allocation> new_allocation(new UnderlyingManualAllocation(
-      this, std::move(reinterpret_cast<UnderlyingManualAllocation *>(allocation)
-                          ->allocation_)));
-  allocations_.emplace(allocation->size(), std::move(new_allocation));
+  allocations_.emplace(allocation->size(), AllocationPtr(allocation));
 }
-MannualFreeAllocation *BufferedAllocator::AllocateImpl(size_t size,
-                                                       Allocator::Attr attr) {
+Allocation *BufferedAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
   {
     platform::LockGuardPtr<std::mutex> guard(mtx_);
     auto it = allocations_.lower_bound(size);
     if (it != allocations_.end() && it->first < size * 2) {
-      std::unique_ptr<Allocation> result(std::move(it->second));
+      AllocationPtr result(std::move(it->second));
       allocations_.erase(it);
-      return new UnderlyingManualAllocation(this, std::move(result));
+      return new UnderlyingManualAllocation(std::move(result));
     }
   }
 
   try {
     return new UnderlyingManualAllocation(
-        this, underlying_allocator_->Allocate(size, attr));
+        underlying_allocator_->Allocate(size, attr));
   } catch (BadAlloc &) {
     FreeCache(size);
     return new UnderlyingManualAllocation(
-        this, underlying_allocator_->Allocate(size, attr));
+        underlying_allocator_->Allocate(size, attr));
   }
 }