allocator.h 6.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
// 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 <memory>
#include <string>
Z
Zeng Jinle 已提交
18
#include <type_traits>
S
sneaxiy 已提交
19
#include <utility>
S
sneaxiy 已提交
20
#include <vector>
W
wanghuancoder 已提交
21

Z
Zeng Jinle 已提交
22
#include "paddle/fluid/framework/inlined_vector.h"
23
#include "paddle/fluid/platform/enforce.h"
24 25 26 27 28 29
#include "paddle/fluid/platform/place.h"

namespace paddle {
namespace memory {
namespace allocation {

Y
Yu Yang 已提交
30
// Exception when `Alloc`/`AllocShared` failed
31 32 33 34
struct BadAlloc : public std::exception {
  inline explicit BadAlloc(std::string err_msg, const char* file, int line)
      : err_str_(platform::GetTraceBackString(std::move(err_msg), file, line)) {
  }
Z
Zeng Jinle 已提交
35

36
  const char* what() const noexcept override { return err_str_.c_str(); }
37

38
  std::string err_str_;
39 40
};

Y
Yu Yang 已提交
41
class Allocator;
Z
Zeng Jinle 已提交
42

Y
Yu Yang 已提交
43 44 45 46 47 48
// Allocation is the object holding the actually pointer. Use
// `Allocation::ptr()` will returns the pointer that allocated.
//
// 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
Z
Zeng Jinle 已提交
49 50 51 52 53 54 55 56

/**
 * 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.
 *
Z
Zeng Jinle 已提交
57
 * Explanations of Allocator design are as follows:
Z
Zeng Jinle 已提交
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83
 *
 * 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.
 */
84 85
class Allocation {
 public:
Z
Zeng Jinle 已提交
86 87
  inline Allocation(void* ptr, size_t size, platform::Place place)
      : ptr_(ptr), size_(size), place_(place) {}
88 89 90

  Allocation(const Allocation& o) = delete;
  Allocation& operator=(const Allocation& o) = delete;
Z
Zeng Jinle 已提交
91 92
  Allocation(Allocation&& o) = delete;
  Allocation& operator=(Allocation&& o) = delete;
93

Y
Yu Yang 已提交
94 95 96 97 98
  // Returns the holding pointer.
  // NOTE: For performance consideration, it is better not to make this method
  // as a virtual method. If we want to implement a `defragmentation` later,
  // we might need to make `ptr_` field as a protected field, and add a virtual
  // method like `defragmentation` to change `ptr_`.
Z
Zeng Jinle 已提交
99
  inline void* ptr() const { return ptr_; }
100

Y
Yu Yang 已提交
101 102 103 104 105 106 107 108 109
  // Returns the size of this memory buffer, i.e., ptr() + size() - 1 is the
  // last valid element.
  //
  // NOTE: Some allocator might alloc more memory than request. The size
  // could larger than its request. For example,
  //    the AlignedAllocator will always allocate memory as size + kAlignment.
  //    The raw pointer might not aligned, so an offset might be added to raw
  //    the pointer. The size of this allocation will be
  //    `size + kAlignemnt - offset`.
Z
Zeng Jinle 已提交
110 111 112
  inline size_t size() const { return size_; }

  inline const platform::Place& place() const { return place_; }
113

Z
Zeng Jinle 已提交
114
  virtual ~Allocation() {}
115

Z
Zeng Jinle 已提交
116 117 118 119
 private:
  inline void RegisterDecoratedAllocator(Allocator* allocator) {
    decorated_allocators_.emplace_back(allocator);
  }
S
sneaxiy 已提交
120

Z
Zeng Jinle 已提交
121
  inline void PopDecoratedAllocator() { decorated_allocators_.pop_back(); }
S
sneaxiy 已提交
122

Z
Zeng Jinle 已提交
123 124 125
  inline Allocator* TopDecoratedAllocator() {
    return decorated_allocators_.back();
  }
Y
Yu Yang 已提交
126

127 128 129 130 131
 private:
  void* ptr_;
  size_t size_;
  platform::Place place_;

Z
Zeng Jinle 已提交
132 133 134 135 136 137 138 139 140
  /**
   * NOTE(zjl): Since decorated_allocators_ is usually a small vector.
   * We reserve a small buffer to it to prevent frequent heap allocation
   *
   * Instead, we can use a std::vector<Allocator *> here, and reserve
   * kReserveAllocatorNum in constructor of Allocation.
   * But using std::vector<Allocator *> would make ocr recognition model
   * fail in CE. The train duration is 8% slower than KPI.
   */
Z
Zeng Jinle 已提交
141 142 143 144 145 146 147 148
  static constexpr size_t kReserveAllocatorNum = 8;
  using DecoratedAllocatorStack =
      framework::InlinedVector<Allocator*, kReserveAllocatorNum>;

  DecoratedAllocatorStack decorated_allocators_;

  friend class Allocator;
};
Y
Yu Yang 已提交
149

Y
Yu Yang 已提交
150
// Base interface class of memory Allocator.
151 152
class Allocator {
 public:
Z
Zeng Jinle 已提交
153 154 155 156 157 158 159 160 161 162 163
  virtual ~Allocator() {}

  class AllocationDeleter {
   public:
    inline void operator()(Allocation* allocation) const {
      Allocator* allocator = allocation->TopDecoratedAllocator();
      allocator->Free(allocation);
    }
  };

  using AllocationPtr = std::unique_ptr<Allocation, AllocationDeleter>;
Y
Yu Yang 已提交
164

Y
Yu Yang 已提交
165
  // Allocate an allocation.
166 167 168
  // size may be 0, but it would be too complex if we handle size == 0
  // in each Allocator. So we handle size == 0 inside AllocatorFacade
  // in our design.
169 170
  inline AllocationPtr Allocate(size_t size) {
    auto ptr = AllocateImpl(size);
Z
Zeng Jinle 已提交
171 172 173 174 175 176 177 178 179
    ptr->RegisterDecoratedAllocator(this);
    return AllocationPtr(ptr);
  }

  // This function should not be called outside Allocator class
  inline void Free(Allocation* allocation) {
    allocation->PopDecoratedAllocator();
    FreeImpl(allocation);
  }
180

181 182
  inline void Release(const platform::Place& place) { ReleaseImpl(place); }

Y
Yu Yang 已提交
183
  // True if the `Allocate` is thread safe.
184
  virtual bool IsAllocThreadSafe() const;
Y
Yu Yang 已提交
185 186

 protected:
187
  virtual Allocation* AllocateImpl(size_t size) = 0;
Z
Zeng Jinle 已提交
188
  virtual void FreeImpl(Allocation* allocation);
189
  virtual void ReleaseImpl(const platform::Place& place) {}
190 191
};

Z
Zeng Jinle 已提交
192 193 194
using AllocationDeleter = Allocator::AllocationDeleter;
using AllocationPtr = Allocator::AllocationPtr;

195 196 197 198 199 200 201 202 203 204 205
inline size_t AlignedSize(size_t size, size_t alignment) {
  auto remaining = size % alignment;
  return remaining == 0 ? size : size + alignment - remaining;
}

inline size_t AlignedPtrOffset(const void* ptr, size_t alignment) {
  auto ptr_addr = reinterpret_cast<uintptr_t>(ptr);
  auto diff = ptr_addr % alignment;
  return diff == 0 ? 0 : alignment - diff;
}

206 207 208
}  // namespace allocation
}  // namespace memory
}  // namespace paddle