malloc.cc 5.3 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14

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. */

Y
Yi Wang 已提交
15
#include "paddle/fluid/memory/memory.h"
16

L
liaogang 已提交
17 18
#include "glog/logging.h"

Y
Yi Wang 已提交
19 20 21
#include "paddle/fluid/memory/detail/buddy_allocator.h"
#include "paddle/fluid/memory/detail/system_allocator.h"
#include "paddle/fluid/platform/gpu_info.h"
L
liaogang 已提交
22 23

DECLARE_double(fraction_of_gpu_memory_to_use);
L
liaogang 已提交
24

25 26 27
namespace paddle {
namespace memory {

28 29 30
using BuddyAllocator = detail::BuddyAllocator;

BuddyAllocator* GetCPUBuddyAllocator() {
31 32 33 34 35 36 37
  static detail::BuddyAllocator* a = nullptr;
  if (a == nullptr) {
    a = new detail::BuddyAllocator(new detail::CPUAllocator,
                                   platform::CpuMinChunkSize(),
                                   platform::CpuMaxChunkSize());
  }
  return a;
L
liaogang 已提交
38 39
}

L
liaogang 已提交
40
template <>
C
chengduoZH 已提交
41
void* Alloc<platform::CPUPlace>(platform::CPUPlace place, size_t size) {
42
  VLOG(10) << "Allocate " << size << " bytes on " << platform::Place(place);
43
  void* p = GetCPUBuddyAllocator()->Alloc(size);
44
  VLOG(10) << "  pointer=" << p;
45
  return p;
L
liaogang 已提交
46 47 48
}

template <>
C
chengduoZH 已提交
49
void Free<platform::CPUPlace>(platform::CPUPlace place, void* p) {
50
  VLOG(10) << "Free pointer=" << p << " on " << platform::Place(place);
L
liaogang 已提交
51 52 53 54 55 56 57 58
  GetCPUBuddyAllocator()->Free(p);
}

template <>
size_t Used<platform::CPUPlace>(platform::CPUPlace place) {
  return GetCPUBuddyAllocator()->Used();
}

59
#ifdef PADDLE_WITH_CUDA
L
liaogang 已提交
60

61
BuddyAllocator* GetGPUBuddyAllocator(int gpu_id) {
62 63
  static BuddyAllocator** as = NULL;
  if (as == NULL) {
64
    int gpu_num = platform::GetCUDADeviceCount();
65
    as = new BuddyAllocator*[gpu_num];
L
liaogang 已提交
66
    for (int gpu = 0; gpu < gpu_num; gpu++) {
67
      as[gpu] = nullptr;
L
liaogang 已提交
68
    }
69 70 71
  }
  platform::SetDeviceId(gpu_id);
  if (!as[gpu_id]) {
Y
Yu Yang 已提交
72
    as[gpu_id] = new BuddyAllocator(new detail::GPUAllocator(gpu_id),
73 74
                                    platform::GpuMinChunkSize(),
                                    platform::GpuMaxChunkSize());
75 76 77
    VLOG(10) << "\n\nNOTE: each GPU device use "
             << FLAGS_fraction_of_gpu_memory_to_use * 100
             << "% of GPU memory.\n"
78 79
             << "You can set GFlags environment variable '"
             << "FLAGS_fraction_of_gpu_memory_to_use"
80
             << "' to change the fraction of GPU usage.\n\n";
81 82
  }
  return as[gpu_id];
L
liaogang 已提交
83 84
}

L
liaogang 已提交
85
template <>
D
dzhwinter 已提交
86
size_t Used<platform::CUDAPlace>(platform::CUDAPlace place) {
87
  return GetGPUBuddyAllocator(place.device)->Used();
L
liaogang 已提交
88
}
L
liaogang 已提交
89

L
liaogang 已提交
90
template <>
C
chengduoZH 已提交
91 92
void* Alloc<platform::CUDAPlace>(platform::CUDAPlace place, size_t size) {
  auto* buddy_allocator = GetGPUBuddyAllocator(place.device);
C
chengduoZH 已提交
93
  auto* ptr = buddy_allocator->Alloc(size);
94 95 96 97
  if (ptr == nullptr) {
    int cur_dev = platform::GetCurrentDeviceId();
    platform::SetDeviceId(place.device);
    size_t avail, total;
98
    platform::GpuMemoryUsage(&avail, &total);
99 100 101 102 103
    LOG(WARNING) << "Cannot allocate " << size << " bytes in GPU "
                 << place.device << ", available " << avail << " bytes";
    LOG(WARNING) << "total " << total;
    LOG(WARNING) << "GpuMinChunkSize " << platform::GpuMinChunkSize();
    LOG(WARNING) << "GpuMaxChunkSize " << platform::GpuMaxChunkSize();
D
dzhwinter 已提交
104
    LOG(WARNING) << "GPU memory used: " << Used<platform::CUDAPlace>(place);
105 106 107
    platform::SetDeviceId(cur_dev);
  }
  return ptr;
108 109
}

L
liaogang 已提交
110
template <>
C
chengduoZH 已提交
111 112 113 114
void Free<platform::CUDAPlace>(platform::CUDAPlace place, void* p) {
  GetGPUBuddyAllocator(place.device)->Free(p);
}

C
chengduoZH 已提交
115 116 117 118
BuddyAllocator* GetCUDAPinnedBuddyAllocator() {
  static BuddyAllocator* ba = NULL;
  if (ba == NULL) {
    ba = new BuddyAllocator(new detail::CUDAPinnedAllocator,
119 120
                            platform::CUDAPinnedMinChunkSize(),
                            platform::CUDAPinnedMaxChunkSize());
C
chengduoZH 已提交
121 122 123 124 125
  }
  return ba;
}

template <>
C
chengduoZH 已提交
126
size_t Used<platform::CUDAPinnedPlace>(platform::CUDAPinnedPlace place) {
C
chengduoZH 已提交
127
  return GetCUDAPinnedBuddyAllocator()->Used();
C
chengduoZH 已提交
128 129 130 131 132
}

template <>
void* Alloc<platform::CUDAPinnedPlace>(platform::CUDAPinnedPlace place,
                                       size_t size) {
C
chengduoZH 已提交
133
  auto* buddy_allocator = GetCUDAPinnedBuddyAllocator();
C
chengduoZH 已提交
134 135
  void* ptr = buddy_allocator->Alloc(size);

C
chengduoZH 已提交
136 137 138 139
  if (ptr == nullptr) {
    LOG(WARNING) << "cudaMallocHost Cannot allocate " << size
                 << " bytes in CUDAPinnedPlace";
  }
C
chengduoZH 已提交
140 141 142 143 144
  return ptr;
}

template <>
void Free<platform::CUDAPinnedPlace>(platform::CUDAPinnedPlace place, void* p) {
C
chengduoZH 已提交
145
  GetCUDAPinnedBuddyAllocator()->Free(p);
146
}
L
Luo Tao 已提交
147
#endif
148

149 150 151 152 153 154 155 156 157 158 159 160
size_t Usage::operator()(const platform::CPUPlace& cpu) const {
  return Used(cpu);
}

size_t Usage::operator()(const platform::CUDAPlace& gpu) const {
#ifdef PADDLE_WITH_CUDA
  return Used(gpu);
#else
  PADDLE_THROW("'CUDAPlace' is not supported in CPU only device.");
#endif
}

C
chengduoZH 已提交
161
size_t Usage::operator()(const platform::CUDAPinnedPlace& cuda_pinned) const {
C
chengduoZH 已提交
162
#ifdef PADDLE_WITH_CUDA
C
chengduoZH 已提交
163
  return Used(cuda_pinned);
C
chengduoZH 已提交
164 165 166
#else
  PADDLE_THROW("'CUDAPinnedPlace' is not supported in CPU only device.");
#endif
C
chengduoZH 已提交
167 168
}

169 170 171 172
size_t memory_usage(const platform::Place& p) {
  return boost::apply_visitor(Usage(), p);
}

173 174
}  // namespace memory
}  // namespace paddle