Polish allocation

Clean allocation->Deleter

test=develop

Author: reyoung

paddle/fluid/memory/allocation/aligned_allocator.h

@@ -86,11 +86,12 @@ template <size_t kAlignment>
 class AlignedAllocator : public ThinAlignedAllocator {
  public:
   using ThinAlignedAllocator::ThinAlignedAllocator;
-  AllocationPtr Allocate(size_t size, Attr attr) override {
+
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override {
     auto raw_allocation =
         underlying_allocator_->Allocate(size + kAlignment, attr);
-    return AllocationPtr(
-        new AlignedAllocation<kAlignment>(std::move(raw_allocation), size));
+    return new AlignedAllocation<kAlignment>(std::move(raw_allocation), size);
   }
 };
 

paddle/fluid/memory/allocation/underlying_manual_allocation.h renamed to paddle/fluid/memory/allocation/allocation_with_underlying.h

@@ -20,9 +20,9 @@ namespace paddle {
 namespace memory {
 namespace allocation {
 
-class UnderlyingManualAllocation : public Allocation {
+class AllocationWithUnderlying : public Allocation {
  public:
-  explicit UnderlyingManualAllocation(AllocationPtr allocation)
+  explicit AllocationWithUnderlying(AllocationPtr allocation)
       : Allocation(allocation->ptr(), allocation->size(), allocation->place()),
         allocation_(std::move(allocation)) {}
   AllocationPtr allocation_;

paddle/fluid/memory/allocation/allocator.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/fluid/memory/allocation/allocator.h"
+
 #include <functional>
 
 namespace paddle {
@@ -24,23 +25,20 @@ Allocator::~Allocator() {}
 
 bool Allocator::IsAllocThreadSafe() const { return false; }
 
+AllocationPtr Allocator::Allocate(size_t size, Allocator::Attr attr) {
+  auto ptr = AllocateImpl(size, attr);
+  ptr->set_allocator(this);
+  return AllocationPtr(ptr);
+}
+
+void Allocator::Free(Allocation* allocation) { delete allocation; }
+
 const char* BadAlloc::what() const noexcept { return msg_.c_str(); }
 
-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;
-  }
+  allocation->allocator()->Free(allocation);
 }
+
 }  // namespace allocation
 }  // namespace memory
 }  // namespace paddle

paddle/fluid/memory/allocation/allocator.h

@@ -32,10 +32,12 @@ class BadAlloc : public std::exception {
 };
 
 class Allocation;
-struct AllocationDeleter {
+class AllocationDeleter {
+ public:
   void operator()(Allocation* allocation) const;
 };
 
+class Allocator;
 // Allocation is the object holding the actually pointer. Use
 // `Allocation::ptr()` will returns the pointer that allocated.
 //
@@ -45,7 +47,7 @@ struct AllocationDeleter {
 class Allocation {
  public:
   Allocation(void* ptr, size_t size, platform::Place place)
-      : ptr_(ptr), size_(size), place_(place) {}
+      : allocator_(nullptr), ptr_(ptr), size_(size), place_(place) {}
 
   Allocation(const Allocation& o) = delete;
   Allocation& operator=(const Allocation& o) = delete;
@@ -70,11 +72,14 @@ class Allocation {
 
   const platform::Place& place() const { return place_; }
 
-  virtual ~Allocation();
+  Allocator* allocator() { return allocator_; }
 
-  std::function<void(Allocation*)> Deleter;
+  void set_allocator(Allocator* allocator) { allocator_ = allocator; }
+
+  virtual ~Allocation();
 
  private:
+  Allocator* allocator_;
   void* ptr_;
   size_t size_;
   platform::Place place_;
@@ -121,25 +126,18 @@ class Allocator {
 
   virtual ~Allocator();
 
-  // Allocate an allocation. Note the return allocation might need to be freed
-  // manually if the Allocator is an `UnmanagedAllocator`.
-  virtual AllocationPtr Allocate(size_t size,
-                                 Allocator::Attr attr = kDefault) = 0;
+  // Allocate an allocation.
+  AllocationPtr Allocate(size_t size, Allocator::Attr attr = kDefault);
 
   // True if the `Allocate` is thread safe.
   virtual bool IsAllocThreadSafe() const;
-};
-
-// User need to invoke `Free` or `FreeUniquePtr` manually if allocated by
-// a manally managed allocator.
-class MannualFreeAllocator : public Allocator {
- public:
-  AllocationPtr Allocate(size_t size, Attr attr) final;
 
  protected:
-  virtual void Free(Allocation* allocation) = 0;
+  virtual void Free(Allocation* allocation);
   virtual Allocation* AllocateImpl(size_t size, Allocator::Attr attr) = 0;
-  friend class MannualFreeAllocation;
+
+ private:
+  friend class AllocationDeleter;
 };
 
 }  // namespace allocation

paddle/fluid/memory/allocation/allocator_facade.cc

@@ -49,12 +49,13 @@ class CPUManagedAllocator : public Allocator {
  public:
   CPUManagedAllocator() : normal_allocator_(new CPUAllocator()) {}
 
-  AllocationPtr Allocate(size_t size, Attr attr) override {
-    return normal_allocator_->Allocate(size, attr);
-  }
-
   bool IsAllocThreadSafe() const override { return true; }
 
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override {
+    return normal_allocator_->Allocate(size, attr).release();
+  }
+
  private:
   std::shared_ptr<Allocator> normal_allocator_;
 };
@@ -103,10 +104,6 @@ class ChunkedManagedAllocator : public Allocator {
     raw_allocator_.reset();
   }
 
-  AllocationPtr Allocate(size_t size, Attr attr) override {
-    return default_allocator_->Allocate(size, attr);
-  }
-
   std::shared_ptr<Allocator> BestFitAllocatorCreator() {
     chunks_.emplace_back(raw_allocator_->Allocate(max_chunk_size_));
     auto* allocation = chunks_.back().get();
@@ -128,6 +125,11 @@ class ChunkedManagedAllocator : public Allocator {
 
   bool IsAllocThreadSafe() const override { return true; }
 
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override {
+    return default_allocator_->Allocate(size, attr).release();
+  }
+
  protected:
   size_t max_chunk_size_;
   int64_t retry_time_;

paddle/fluid/memory/allocation/auto_increment_allocator.cc

@@ -17,18 +17,34 @@
 namespace paddle {
 namespace memory {
 namespace allocation {
+bool AutoIncrementAllocator::IsAllocThreadSafe() const { return true; }
 
-AllocationPtr AutoIncrementAllocator::Allocate(size_t size,
-                                               Allocator::Attr attr) {
+std::shared_ptr<Allocator> AutoIncrementAllocator::CreateNewAllocator() {
+  std::lock_guard<std::mutex> 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];
+}
+Allocation *AutoIncrementAllocator::AllocateImpl(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;
   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&) {
+      return res.release();
+    } catch (BadAlloc &) {
       if (++cur >= allocator_num) {
         cur = 0;
       }
@@ -47,32 +63,14 @@ AllocationPtr AutoIncrementAllocator::Allocate(size_t size,
     try {
       auto ret = underlying_allocators_[cur]->Allocate(size, attr);
       prev_success_allocator_ = cur;
-      return ret;
-    } catch (BadAlloc&) {
+      return ret.release();
+    } catch (BadAlloc &) {
     } catch (...) {
       throw;
     }
   }
   // No suitable allocator
-  return CreateNewAllocator()->Allocate(size, attr);
-}
-
-bool AutoIncrementAllocator::IsAllocThreadSafe() const { return true; }
-
-std::shared_ptr<Allocator> AutoIncrementAllocator::CreateNewAllocator() {
-  std::lock_guard<std::mutex> 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];
+  return CreateNewAllocator()->Allocate(size, attr).release();
 }
 
 }  // namespace allocation

paddle/fluid/memory/allocation/auto_increment_allocator.h

@@ -54,13 +54,15 @@ class AutoIncrementAllocator : public Allocator {
   explicit AutoIncrementAllocator(AllocatorCreator&& creator, size_t capacity)
       : creator_(std::move(creator)), underlying_allocators_(capacity) {}
 
-  AllocationPtr Allocate(size_t size, Attr attr) override;
-
   bool IsAllocThreadSafe() const override;
 
  private:
   std::shared_ptr<Allocator> CreateNewAllocator();
 
+ protected:
+  Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
+
+ private:
   AllocatorCreator creator_;
 
   std::vector<AllocatorCreator::result_type> underlying_allocators_;

paddle/fluid/memory/allocation/best_fit_allocator.h

@@ -98,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 MannualFreeAllocator {
+class BestFitAllocator : public Allocator {
  public:
   explicit BestFitAllocator(Allocation* allocation);
 

paddle/fluid/memory/allocation/buffered_allocator.cc

@@ -16,7 +16,7 @@
 #include <algorithm>
 #include <limits>
 #include <utility>
-#include "paddle/fluid/memory/allocation/underlying_manual_allocation.h"
+#include "paddle/fluid/memory/allocation/allocation_with_underlying.h"
 
 namespace paddle {
 namespace memory {
@@ -60,16 +60,16 @@ Allocation *BufferedAllocator::AllocateImpl(size_t size, Allocator::Attr attr) {
     if (it != allocations_.end() && it->first < size * 2) {
       AllocationPtr result(std::move(it->second));
       allocations_.erase(it);
-      return new UnderlyingManualAllocation(std::move(result));
+      return new AllocationWithUnderlying(std::move(result));
     }
   }
 
   try {
-    return new UnderlyingManualAllocation(
+    return new AllocationWithUnderlying(
         underlying_allocator_->Allocate(size, attr));
   } catch (BadAlloc &) {
     FreeCache(size);
-    return new UnderlyingManualAllocation(
+    return new AllocationWithUnderlying(
         underlying_allocator_->Allocate(size, attr));
   }
 }

paddle/fluid/memory/allocation/buffered_allocator.h

@@ -29,7 +29,7 @@ namespace allocation {
 // memory allocation and reuse memory.
 // BufferedAllocator provides the same thread-safety level as
 // underlying_allocator_
-class BufferedAllocator : public MannualFreeAllocator {
+class BufferedAllocator : public Allocator {
  public:
   explicit BufferedAllocator(std::unique_ptr<Allocator> &&allocator);