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

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

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

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

17
#include "paddle/memory/detail/meta_cache.h"
L
liaogang 已提交
18
#include "paddle/memory/detail/meta_data.h"
L
liaogang 已提交
19
#include "paddle/memory/detail/system_allocator.h"
L
liaogang 已提交
20 21 22
#include "paddle/platform/assert.h"
#include "paddle/platform/cpu_info.h"
#include "paddle/platform/gpu_info.h"
L
liaogang 已提交
23

L
liaogang 已提交
24
#include <mutex>
L
liaogang 已提交
25
#include <set>
L
liaogang 已提交
26
#include <unordered_map>
L
liaogang 已提交
27
#include <vector>
L
liaogang 已提交
28

L
liaogang 已提交
29 30 31 32 33
namespace paddle {
namespace memory {
namespace detail {

class BuddyAllocator {
34
 public:
L
liaogang 已提交
35 36 37
  BuddyAllocator(SystemAllocator* system_allocator, size_t min_chunk_size,
                 size_t max_chunk_size);

38 39
  ~BuddyAllocator();

L
liaogang 已提交
40 41
 public:
  void* Alloc(size_t unaligned_size);
L
liaogang 已提交
42
  void Free(void* ptr);
43 44
  size_t Used();

L
liaogang 已提交
45
 public:
L
liaogang 已提交
46
  // Disable copy and assignment
L
liaogang 已提交
47 48 49
  BuddyAllocator(const BuddyAllocator&) = delete;
  BuddyAllocator& operator=(const BuddyAllocator&) = delete;

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

L
liaogang 已提交
56 57
  /*! \brief Allocate fixed-size memory from system */
  void* SystemAlloc(size_t size);
58

L
liaogang 已提交
59 60
  /*! \brief If existing chunks are not suitable, refill pool */
  PoolSet::iterator RefillPool();
L
liaogang 已提交
61

L
liaogang 已提交
62
  /**
L
liaogang 已提交
63 64 65 66 67
   *  \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 已提交
68
   *
L
liaogang 已提交
69
   *  \return  the left buddy address
L
liaogang 已提交
70 71
   */
  void* SplitToAlloc(PoolSet::iterator it, size_t size);
L
liaogang 已提交
72

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

L
liaogang 已提交
76 77 78 79 80
  /*! \brief Clean idle fallback allocation */
  void CleanIdleFallBackAlloc();

  /*! \brief Clean idle normal allocation */
  void CleanIdleNormalAlloc();
L
liaogang 已提交
81

L
liaogang 已提交
82 83 84 85 86 87 88 89
 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

 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
  /*! Record fallback allocation count for auto-scaling */
  size_t fallback_alloc_count_ = 0;

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

 private:
105
  /*! Allocate CPU/GPU memory from system */
L
liaogang 已提交
106 107
  SystemAllocator* system_allocator_;
  std::mutex mutex_;
L
liaogang 已提交
108 109 110 111 112
};

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