buddy_allocator.h 3.2 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
L
liaogang 已提交
2

L
Luo Tao 已提交
3 4 5
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
L
liaogang 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
L
liaogang 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
L
liaogang 已提交
14 15 16

#pragma once

G
gongweibao 已提交
17
#include <memory>
Y
Yi Wang 已提交
18 19 20 21 22 23 24
#include <mutex>  // NOLINT
#include <set>
#include <tuple>
#include <unordered_map>
#include <vector>

#include "paddle/fluid/memory/detail/memory_block.h"
Y
Yi Wang 已提交
25 26 27
#include "paddle/fluid/memory/detail/system_allocator.h"
#include "paddle/fluid/platform/cpu_info.h"
#include "paddle/fluid/platform/gpu_info.h"
L
liaogang 已提交
28 29 30 31 32 33

namespace paddle {
namespace memory {
namespace detail {

class BuddyAllocator {
34
 public:
G
gongweibao 已提交
35 36
  BuddyAllocator(std::unique_ptr<SystemAllocator> system_allocator,
                 size_t min_chunk_size, size_t max_chunk_size);
L
liaogang 已提交
37

38 39
  ~BuddyAllocator();

L
liaogang 已提交
40 41
 public:
  void* Alloc(size_t unaligned_size);
L
liaogang 已提交
42
  void Free(void* ptr);
43
  // Release the unused memory pool, a real free operation for the OS.
W
Wilber 已提交
44
  uint64_t Release();
45
  size_t Used();
46 47
  size_t GetMinChunkSize();
  size_t GetMaxChunkSize();
48

L
liaogang 已提交
49
 public:
L
liaogang 已提交
50
  // Disable copy and assignment
L
liaogang 已提交
51 52 53
  BuddyAllocator(const BuddyAllocator&) = delete;
  BuddyAllocator& operator=(const BuddyAllocator&) = delete;

54
 private:
55
  // Tuple (allocator index, memory size, memory address)
L
liaogang 已提交
56
  using IndexSizeAddress = std::tuple<size_t, size_t, void*>;
57
  // Each element in PoolSet is a free allocation
L
liaogang 已提交
58
  using PoolSet = std::set<IndexSizeAddress>;
59

L
liaogang 已提交
60 61
  /*! \brief Allocate fixed-size memory from system */
  void* SystemAlloc(size_t size);
62

L
liaogang 已提交
63
  /*! \brief If existing chunks are not suitable, refill pool */
64
  PoolSet::iterator RefillPool(size_t request_bytes);
L
liaogang 已提交
65

L
liaogang 已提交
66
  /**
L
liaogang 已提交
67 68 69 70 71
   *  \brief   Find the suitable chunk from existing pool and split
   *           it to left and right buddies
   *
   *  \param   it     the iterator of pool list
   *  \param   size   the size of allocation
L
liaogang 已提交
72
   *
L
liaogang 已提交
73
   *  \return  the left buddy address
L
liaogang 已提交
74 75
   */
  void* SplitToAlloc(PoolSet::iterator it, size_t size);
L
liaogang 已提交
76

L
liaogang 已提交
77
  /*! \brief Find the existing chunk which used to allocation */
L
liaogang 已提交
78
  PoolSet::iterator FindExistChunk(size_t size);
L
liaogang 已提交
79

L
liaogang 已提交
80 81 82 83 84 85 86
 private:
  size_t total_used_ = 0;  // the total size of used memory
  size_t total_free_ = 0;  // the total size of free memory

  size_t min_chunk_size_;  // the minimum size of each chunk
  size_t max_chunk_size_;  // the maximum size of each chunk

Z
zhhsplendid 已提交
87 88
  size_t realloc_size_ = 0;  // the size of re-allocated chunk

L
liaogang 已提交
89
 private:
90 91 92 93 94
  /**
   * \brief A list of free allocation
   *
   * \note  Only store free chunk memory in pool
   */
L
liaogang 已提交
95 96
  PoolSet pool_;

97 98 99 100 101
  /**
   * \brief Record the allocated chunks when Refill pool.
   */
  PoolSet chunks_;

L
liaogang 已提交
102
 private:
103
  /*! Unify the metadata format between GPU and CPU allocations */
L
liaogang 已提交
104 105 106
  MetadataCache cache_;

 private:
107
  /*! Allocate CPU/GPU memory from system */
G
gongweibao 已提交
108
  std::unique_ptr<SystemAllocator> system_allocator_;
L
liaogang 已提交
109
  std::mutex mutex_;
L
liaogang 已提交
110 111 112 113 114
};

}  // namespace detail
}  // namespace memory
}  // namespace paddle