allocator.h 5.4 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>
S
sneaxiy 已提交
18
#include <vector>
S
sneaxiy 已提交
19
#include "paddle/fluid/framework/inlined_stack.h"
20 21 22 23 24 25
#include "paddle/fluid/platform/place.h"

namespace paddle {
namespace memory {
namespace allocation {

Y
Yu Yang 已提交
26
// Exception when `Alloc`/`AllocShared` failed
27 28
class BadAlloc : public std::exception {
 public:
Y
Yu Yang 已提交
29
  explicit BadAlloc(std::string msg) : msg_(std::move(msg)) {}
30 31 32 33 34 35
  const char* what() const noexcept override;

 private:
  std::string msg_;
};

Y
Yu Yang 已提交
36
class Allocation;
Y
Yu Yang 已提交
37 38
class AllocationDeleter {
 public:
Y
Yu Yang 已提交
39 40 41
  void operator()(Allocation* allocation) const;
};

Y
Yu Yang 已提交
42
class Allocator;
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
49 50 51
class Allocation {
 public:
  Allocation(void* ptr, size_t size, platform::Place place)
S
sneaxiy 已提交
52
      : ptr_(ptr), size_(size), place_(place) {}
53 54 55

  Allocation(const Allocation& o) = delete;
  Allocation& operator=(const Allocation& o) = delete;
S
sneaxiy 已提交
56 57
  Allocation(Allocation&& o) = delete;
  Allocation& operator=(Allocation&& o) = delete;
58

Y
Yu Yang 已提交
59 60 61 62 63
  // 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_`.
64 65
  void* ptr() const { return ptr_; }

Y
Yu Yang 已提交
66 67 68 69 70 71 72 73 74
  // 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`.
75 76 77 78
  size_t size() const { return size_; }

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

S
sneaxiy 已提交
79
  virtual ~Allocation();
80

S
sneaxiy 已提交
81 82 83 84
  // This function should only be used in unittest
  std::vector<Allocator*> GetAllocatorChain() const {
    std::vector<Allocator*> allocators;
    for (size_t i = 0; i < allocator_chain_.size(); ++i) {
S
sneaxiy 已提交
85
      allocators.push_back(allocator_chain_[i]);
S
sneaxiy 已提交
86 87 88
    }
    return allocators;
  }
Y
Yu Yang 已提交
89

S
sneaxiy 已提交
90 91 92 93 94 95 96 97
 private:
  inline void RegisterAllocatorChain(Allocator* allocator) {
    allocator_chain_.push(allocator);
  }

  inline void PopAllocator() { allocator_chain_.pop(); }

  inline Allocator* TopAllocator() { return allocator_chain_.top(); }
Y
Yu Yang 已提交
98

99 100 101 102
 private:
  void* ptr_;
  size_t size_;
  platform::Place place_;
S
sneaxiy 已提交
103
  framework::InlinedStack<Allocator*, 8> allocator_chain_;
S
sneaxiy 已提交
104 105 106

  friend class Allocator;
  friend class AllocationDeleter;
107 108
};

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

Y
Yu Yang 已提交
111 112 113 114 115 116
// Base interface class of memory Allocator.
// To allocate a memory, allocator needs two parameters:
//    1. size of bytes.
//    2. Attribute of memory.
// NOTE: the attribute of memory might be ignored if the allocator does not
// care it.
117 118 119
class Allocator {
 public:
  enum Attr {
Y
Yu Yang 已提交
120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145
    kDefault = 0,  // Default attribute. Uses the fast or stablest allocation
                   // algorithm.

    kFixedHuge = 1,  // The allocation may not be freed until the program
                     // ends. e.g., `Parameters` and `Momentum`.

    kFluxHuge = 2,  // The allocation may create and freed frequently and the
                    // allocation is considerable huge. Like `activations`
                    // and gradients.

    kScratchpad =
        3,  // The `Scratchpad` memory is allocated and freed very soon,
            // usually within an operator or aux memory.
            // Like CUDNN workspace, AUX memory in batch norm, etc.
            //
            // https://en.wikipedia.org/wiki/Scratchpad_memory

    kCrossDevice =
        4,  // The memory used cross-device memory copy/communication.
            // For example:
            // 1. it can use an `pinned` memory for CPU-GPU
            //    communication.
            // 2. it can use an `registered` memory for RDMA
            //    communication.

    NumOfAttrs = 5  // The number of all attributes. It is used internally.
146 147 148
  };

  virtual ~Allocator();
Y
Yu Yang 已提交
149

Y
Yu Yang 已提交
150 151
  // Allocate an allocation.
  AllocationPtr Allocate(size_t size, Allocator::Attr attr = kDefault);
152

Y
Yu Yang 已提交
153
  // True if the `Allocate` is thread safe.
154
  virtual bool IsAllocThreadSafe() const;
Y
Yu Yang 已提交
155

S
sneaxiy 已提交
156 157 158
  // This function should not be called outside
  void Free(Allocation* allocation);

Y
Yu Yang 已提交
159
 protected:
Y
Yu Yang 已提交
160
  virtual Allocation* AllocateImpl(size_t size, Allocator::Attr attr) = 0;
S
sneaxiy 已提交
161
  virtual void FreeImpl(Allocation* allocation);
Y
Yu Yang 已提交
162 163 164

 private:
  friend class AllocationDeleter;
165 166 167 168 169
};

}  // namespace allocation
}  // namespace memory
}  // namespace paddle