buddy_allocator.h 3.4 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);
42 43 44
  void Free(void*);
  size_t Used();

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

50
 private:
L
liaogang 已提交
51 52 53
  // Tuple type: allocator index, memory size, memory address
  using IndexSizeAddress = std::tuple<size_t, size_t, void*>;
  using PoolSet = std::set<IndexSizeAddress>;
54

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

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

L
liaogang 已提交
61
  /**
L
liaogang 已提交
62
   *  \brief Find the suitable chunk from existing pool
L
liaogang 已提交
63
   *
L
liaogang 已提交
64 65 66 67
   *  \param it   pool iterator which contains suitable block.
   *  \param size the size of allocation.
   */
  void* SplitToAlloc(PoolSet::iterator it, size_t size);
L
liaogang 已提交
68

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

L
liaogang 已提交
72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88
 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:
  PoolSet pool_;

 private:
  // Unify the metadata format between GPU and CPU allocations
  MetadataCache cache_;

 private:
  SystemAllocator* system_allocator_;
  std::mutex mutex_;
L
liaogang 已提交
89 90
};

L
liaogang 已提交
91 92
BuddyAllocator* GetCPUBuddyAllocator() {
  static BuddyAllocator* a = nullptr;
L
liaogang 已提交
93
  if (a == nullptr) {
L
liaogang 已提交
94 95
    a = new BuddyAllocator(new CPUAllocator, platform::CpuMinChunkSize(),
                           platform::CpuMaxChunkSize());
L
liaogang 已提交
96 97 98 99 100 101
  }
  return a;
}

#ifndef PADDLE_ONLY_CPU  // The following code are for CUDA.

L
liaogang 已提交
102 103
BuddyAllocator* GetGPUBuddyAllocator(int gpu_id) {
  static BuddyAllocator** as = NULL;
L
liaogang 已提交
104
  if (as == NULL) {
L
liaogang 已提交
105 106
    int gpu_num = platform::GpuDeviceCount();
    as = new BuddyAllocator*[gpu_num];
L
liaogang 已提交
107
    for (int gpu = 0; gpu < gpu_num; gpu++) {
L
liaogang 已提交
108 109 110
      as[gpu] =
          new BuddyAllocator(new GPUAllocator, platform::GpuMinChunkSize(),
                             platform::GpuMaxChunkSize());
L
liaogang 已提交
111 112 113 114 115
    }
  }
  return as[gpu_id];
}

116
#endif  // PADDLE_ONLY_CPU
L
liaogang 已提交
117 118 119 120

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