Merge pull request #16424 from sneaxiy/fix_allocator_bug

Fix allocator bug

Author: sneaxiy

paddle/fluid/framework/operator.h

@@ -365,9 +365,6 @@ class ExecutionContext {
     auto shared_allocation = std::shared_ptr<memory::allocation::Allocation>(
         allocation_ptr, deleter);
 
-    PADDLE_ENFORCE(
-        dynamic_cast<platform::TemporaryAllocation*>(allocation_ptr) != nullptr,
-        "The AllocationPtr must be TemporaryAllocation.");
     PADDLE_ENFORCE_GE(allocation_ptr->size(),
                       framework::product(dim) * sizeof(T));
 

paddle/fluid/memory/allocation/CMakeLists.txt

@@ -4,6 +4,7 @@ cc_library(best_fit_allocator SRCS best_fit_allocator.cc DEPS allocator)
 cc_library(locked_allocator SRCS locked_allocator.cc DEPS allocator)
 cc_library(buffered_allocator SRCS buffered_allocator.cc DEPS allocator)
 cc_library(legacy_allocator SRCS legacy_allocator.cc DEPS allocator buddy_allocator profiler)
+cc_library(zero_size_allocator SRCS zero_size_allocator.cc DEPS allocator)
 cc_test(buffered_allocator_test SRCS buffered_allocator_test.cc DEPS best_fit_allocator locked_allocator buffered_allocator cpu_allocator)
 
 if (WITH_GPU)
@@ -37,30 +38,20 @@ else ()
     set(AllocatorFacadeDeps)
 endif()
 
+list(APPEND AllocatorFacadeDeps cpu_allocator locked_allocator best_fit_allocator aligned_allocator auto_increment_allocator conditional_allocator retry_allocator buffered_allocator legacy_allocator zero_size_allocator)
+
 cc_library(aligned_allocator SRCS aligned_allocator.cc DEPS allocator)
 cc_library(auto_increment_allocator SRCS auto_increment_allocator.cc DEPS allocator)
-cc_library(zero_size_allocator SRCS zero_size_allocator.cc DEPS allocator)
 cc_library(conditional_allocator SRCS conditional_allocator.cc DEPS allocator)
-cc_library(allocator_strategy SRCS allocator_strategy.cc DEPS gflags)
-cc_library(allocator_facade SRCS allocator_facade.cc DEPS
-        ${AllocatorFacadeDeps}
-        cpu_allocator
-        locked_allocator
-        best_fit_allocator
-        aligned_allocator
-        auto_increment_allocator
-        zero_size_allocator
-        conditional_allocator
-        retry_allocator
-        buffered_allocator
-        allocator_strategy
-        legacy_allocator
-        )
+cc_library(allocator_strategy SRCS allocator_strategy.cc DEPS gflags ${AllocatorFacadeDeps})
+cc_library(allocator_facade SRCS allocator_facade.cc DEPS allocator_strategy)
 
 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 best_fit_allocator locked_allocator cpu_allocator)
 
+cc_test(naive_best_fit_allocator_facade_test SRCS naive_best_fit_allocator_facade_test.cc DEPS allocator_facade)
+
 cc_test(allocator_facade_abs_flags_test SRCS allocator_facade_abs_flags_test.cc DEPS allocator_facade)
 
 cc_test(allocator_facade_frac_flags_test SRCS allocator_facade_frac_flags_test.cc DEPS allocator_facade)

paddle/fluid/memory/allocation/aligned_allocator.h

@@ -14,6 +14,7 @@
 
 #pragma once
 #include <memory>
+#include <utility>
 #include "paddle/fluid/memory/allocation/allocator.h"
 
 namespace paddle {
@@ -93,6 +94,8 @@ class AlignedAllocator : public ThinAlignedAllocator {
         underlying_allocator_->Allocate(size + kAlignment, attr);
     return new AlignedAllocation<kAlignment>(std::move(raw_allocation), size);
   }
+
+  void FreeImpl(Allocation* allocation) override { delete allocation; }
 };
 
 }  // namespace allocation

paddle/fluid/memory/allocation/allocator.cc

@@ -27,16 +27,24 @@ bool Allocator::IsAllocThreadSafe() const { return false; }
 
 AllocationPtr Allocator::Allocate(size_t size, Allocator::Attr attr) {
   auto ptr = AllocateImpl(size, attr);
-  ptr->set_allocator(this);
+  ptr->RegisterDecoratedAllocator(this);
   return AllocationPtr(ptr);
 }
 
-void Allocator::Free(Allocation* allocation) { delete allocation; }
+void Allocator::FreeImpl(Allocation* allocation) {
+  Allocator* allocator = allocation->TopDecoratedAllocator();
+  allocator->Free(allocation);
+}
+
+void Allocator::Free(Allocation* allocation) {
+  allocation->PopDecoratedAllocator();
+  FreeImpl(allocation);
+}
 
 const char* BadAlloc::what() const noexcept { return msg_.c_str(); }
 
 void AllocationDeleter::operator()(Allocation* allocation) const {
-  auto* allocator = allocation->allocator();
+  Allocator* allocator = allocation->TopDecoratedAllocator();
   allocator->Free(allocation);
 }
 

paddle/fluid/memory/allocation/allocator.h

@@ -15,6 +15,8 @@
 #pragma once
 #include <memory>
 #include <string>
+#include <utility>
+#include <vector>
 #include "paddle/fluid/platform/place.h"
 
 namespace paddle {
@@ -44,13 +46,56 @@ class Allocator;
 // NOTE: this is the base class of Allocation. Each allocator can use its own
 //       allocation object.
 // NOTE: the `Allocation::ptr()` could be nullptr, if the allocation size is 0
+
+/**
+ * Allocation is returned by Allocator::Allocate() method.
+ *
+ * An allocator may be decorated by another allocator. For example, we can
+ * decorate
+ * a RetryAllocator to any allocator to perform allocation retry when first
+ * allocation request fails.
+ *
+ * Explanations of Allocator design is as follows:
+ *
+ * Suppose we have an allocator which is decorated by several allocators:
+ *
+ *   A(1) <- A(2) <- A(3) <- ... <- A(n)
+ *
+ * , and the public allocator is A(1).
+ *
+ * The allocation process would be:
+ *
+ *   A(n).Allocate() -> ... -> A(2).Allocate() -> A(1).Allocate()
+ *
+ * , and the free process would be:
+ *
+ *   A(1).Free() -> A(2).Free() -> ... -> A(n).Free()
+ *
+ * Therefore, we should record the allocator chain when allocating, so
+ * that we can free the allocation in the reverse order of allocator chain.
+ * The field `decorated_allocators_` is used to record this chain.
+ *
+ * Another example is that we want to add additional fields in Allocation,
+ * e.g., something what is done in AlignedAllocator, etc.
+ * In this case, we should declare a derived class of Allocation, which
+ * contains an underlying Allocation allocated by the underlying allocator.
+ * Therefore, `decorated_allocators_` of the new Allocation object would
+ * be a new chain, differing from the underlying Allocation object.
+ */
 class Allocation {
  public:
   Allocation(void* ptr, size_t size, platform::Place place)
-      : allocator_(nullptr), ptr_(ptr), size_(size), place_(place) {}
+      : ptr_(ptr), size_(size), place_(place) {
+    // NOTE(zjl): Since decorated_allocators_ is usually a small vector
+    // We reserve a small buffer to it to prevent frequent heap allocation
+    // Not quite sure whether we need something like gtl vector.
+    decorated_allocators_.reserve(8);
+  }
 
   Allocation(const Allocation& o) = delete;
   Allocation& operator=(const Allocation& o) = delete;
+  Allocation(Allocation&& o) = delete;
+  Allocation& operator=(Allocation&& o) = delete;
 
   // Returns the holding pointer.
   // NOTE: For performance consideration, it is better not to make this method
@@ -72,17 +117,31 @@ class Allocation {
 
   const platform::Place& place() const { return place_; }
 
-  Allocator* allocator() { return allocator_; }
+  virtual ~Allocation();
 
-  void set_allocator(Allocator* allocator) { allocator_ = allocator; }
+ private:
+  const std::vector<Allocator*>& DecoratedAllocators() const {
+    return decorated_allocators_;
+  }
 
-  virtual ~Allocation();
+  inline void RegisterDecoratedAllocator(Allocator* allocator) {
+    decorated_allocators_.push_back(allocator);
+  }
+
+  inline void PopDecoratedAllocator() { decorated_allocators_.pop_back(); }
+
+  inline Allocator* TopDecoratedAllocator() {
+    return decorated_allocators_.back();
+  }
 
  private:
-  Allocator* allocator_;
   void* ptr_;
   size_t size_;
   platform::Place place_;
+  std::vector<Allocator*> decorated_allocators_;
+
+  friend class Allocator;
+  friend class AllocationDeleter;
 };
 
 using AllocationPtr = std::unique_ptr<Allocation, AllocationDeleter>;
@@ -132,9 +191,12 @@ class Allocator {
   // True if the `Allocate` is thread safe.
   virtual bool IsAllocThreadSafe() const;
 
+  // This function should not be called outside
+  void Free(Allocation* allocation);
+
  protected:
-  virtual void Free(Allocation* allocation);
   virtual Allocation* AllocateImpl(size_t size, Allocator::Attr attr) = 0;
+  virtual void FreeImpl(Allocation* allocation);
 
  private:
   friend class AllocationDeleter;

paddle/fluid/memory/allocation/allocator_facade.cc

@@ -17,6 +17,7 @@
 #include <map>
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 #include "paddle/fluid/memory/allocation/aligned_allocator.h"
 #include "paddle/fluid/memory/allocation/allocator_facade.h"
@@ -30,6 +31,7 @@
 #include "paddle/fluid/memory/allocation/retry_allocator.h"
 #include "paddle/fluid/memory/allocation/zero_size_allocator.h"
 #include "paddle/fluid/platform/cpu_info.h"
+#include "paddle/fluid/platform/enforce.h"
 #include "paddle/fluid/platform/place.h"
 #ifdef PADDLE_WITH_CUDA
 #include "paddle/fluid/memory/allocation/cuda_allocator.h"
@@ -47,6 +49,17 @@ namespace paddle {
 namespace memory {
 namespace allocation {
 
+static inline std::shared_ptr<Allocator> WrapRetryAllocator(
+    std::shared_ptr<Allocator> allocator, int64_t retry_time) {
+  if (retry_time > 0) {
+    auto* retry_allocator =
+        new RetryAllocator(std::move(allocator), retry_time);
+    allocator.reset(retry_allocator);
+  }
+
+  return allocator;
+}
+
 // TODO(yy): Dirty code here. This class should be configurable in runtime.
 class CPUManagedAllocator : public Allocator {
  public:
@@ -110,14 +123,10 @@ class ChunkedAllocator : public Allocator {
   std::shared_ptr<Allocator> CreateAllocatorWithChunk() {
     chunks_.emplace_back(raw_allocator_->Allocate(max_chunk_size_));
     auto* allocation = chunks_.back().get();
-    std::unique_ptr<Allocator> allocator(new LockedAllocator(
-        std::unique_ptr<Allocator>(new BestFitAllocator(allocation))));
+    std::shared_ptr<Allocator> allocator(new LockedAllocator(
+        std::shared_ptr<Allocator>(new BestFitAllocator(allocation))));
 
-    if (retry_time_ > 0) {
-      auto* retry_allocator =
-          new RetryAllocator(std::move(allocator), retry_time_);
-      allocator.reset(retry_allocator);
-    }
+    allocator = WrapRetryAllocator(allocator, retry_time_);
 
     return std::make_shared<AlignedAllocator<64u>>(std::move(allocator));
   }
@@ -188,13 +197,23 @@ class AllocatorFacadePrivate {
   ~AllocatorFacadePrivate() = default;
 
   AllocatorFacadePrivate() {
-    if (GetAllocatorStrategy() == AllocatorStrategy::kLegacy) {
-      InitLegacyAllocator();
-    } else {
-      InitCPUAllocator();
-      InitCUDAAllocator();
-      InitCUDAPinnedAllocator();
-      WrapZeroSizeAllocator();
+    auto strategy = GetAllocatorStrategy();
+    switch (strategy) {
+      case AllocatorStrategy::kLegacy: {
+        InitLegacyAllocator();
+        break;
+      }
+      case AllocatorStrategy::kNaiveBestFit: {
+        InitCPUAllocator();
+        InitCUDAAllocator();
+        InitCUDAPinnedAllocator();
+        WrapZeroSizeAllocator();
+        break;
+      }
+      default: {
+        PADDLE_THROW("Unsupported allocator strategy: %d",
+                     static_cast<int>(strategy));
+      }
     }
   }
 
@@ -252,8 +271,7 @@ AllocatorFacade& AllocatorFacade::Instance() {
 
 std::shared_ptr<Allocation> AllocatorFacade::AllocShared(
     const platform::Place& place, size_t size, Allocator::Attr attr) {
-  return std::shared_ptr<Allocation>(Alloc(place, size, attr).release(),
-                                     AllocationDeleter());
+  return std::shared_ptr<Allocation>(Alloc(place, size, attr));
 }
 
 AllocationPtr AllocatorFacade::Alloc(const platform::Place& place, size_t size,

paddle/fluid/memory/allocation/allocator_strategy.cc

@@ -14,20 +14,27 @@
 
 #include "paddle/fluid/memory/allocation/allocator_strategy.h"
 #include "gflags/gflags.h"
+#include "paddle/fluid/platform/enforce.h"
 
 DEFINE_string(
     allocator_strategy, "legacy",
     "The allocation strategy. Legacy means the original allocator of Fluid."
-    "New means the experimental allocators of Fluid. in [legacy, new]");
+    "naive_best_fit means the experimental best fit allocator. "
+    "allocator. Enum in [legacy, naive_best_fit].");
 
 namespace paddle {
 namespace memory {
 namespace allocation {
 
 static AllocatorStrategy GetStrategyFromFlag() {
-  return FLAGS_allocator_strategy == "legacy"
-             ? AllocatorStrategy::kLegacy
-             : AllocatorStrategy::kNaiveBestFit;
+  if (FLAGS_allocator_strategy == "legacy") {
+    return AllocatorStrategy::kLegacy;
+  } else if (FLAGS_allocator_strategy == "naive_best_fit") {
+    return AllocatorStrategy::kNaiveBestFit;
+  } else {
+    PADDLE_THROW("Unsupported allocator strategy: %s",
+                 FLAGS_allocator_strategy);
+  }
 }
 
 AllocatorStrategy GetAllocatorStrategy() {

paddle/fluid/memory/allocation/best_fit_allocator.cc

@@ -109,7 +109,7 @@ size_t BestFitAllocator::NumFreeChunks() const {
   }
   return num;
 }
-void BestFitAllocator::Free(Allocation* allocation) {
+void BestFitAllocator::FreeImpl(Allocation* allocation) {
   auto* bf_allocation = dynamic_cast<BestFitAllocation*>(allocation);
   PADDLE_ENFORCE_NOT_NULL(bf_allocation,
                           "The input allocation is not BestFitAllocation.");

paddle/fluid/memory/allocation/best_fit_allocator.h

@@ -119,7 +119,7 @@ class BestFitAllocator : public Allocator {
   void InsertFreeNode(const ListIt& it);
 
  protected:
-  void Free(Allocation* allocation) override;
+  void FreeImpl(Allocation* allocation) override;
   Allocation* AllocateImpl(size_t size, Allocator::Attr attr) override;
 
  private: