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

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

#include "paddle/memory/memory.h"
16

L
liaogang 已提交
17
#include <algorithm>  // for transform
18
#include <cstring>    // for memcpy
L
liaogang 已提交
19
#include <memory>     // for unique_ptr
20 21
#include <mutex>      // for call_once

L
liaogang 已提交
22 23
#include "glog/logging.h"

L
liaogang 已提交
24 25
#include "paddle/memory/detail/buddy_allocator.h"
#include "paddle/memory/detail/system_allocator.h"
L
liaogang 已提交
26 27 28
#include "paddle/platform/gpu_info.h"

DECLARE_double(fraction_of_gpu_memory_to_use);
L
liaogang 已提交
29

30 31 32
namespace paddle {
namespace memory {

33 34
using BuddyAllocator = detail::BuddyAllocator;

L
liaogang 已提交
35 36
std::once_flag cpu_allocator_flag;
std::once_flag gpu_allocator_flag;
37 38

BuddyAllocator* GetCPUBuddyAllocator() {
L
liaogang 已提交
39
  static std::unique_ptr<BuddyAllocator> a{nullptr};
40

L
liaogang 已提交
41
  std::call_once(cpu_allocator_flag, [&]() {
42 43 44 45 46 47
    a.reset(new BuddyAllocator(new detail::CPUAllocator,
                               platform::CpuMinChunkSize(),
                               platform::CpuMaxChunkSize()));
  });

  return a.get();
L
liaogang 已提交
48 49
}

L
liaogang 已提交
50 51 52 53 54 55 56 57 58 59 60 61 62 63 64
template <>
void* Alloc<platform::CPUPlace>(platform::CPUPlace place, size_t size) {
  return GetCPUBuddyAllocator()->Alloc(size);
}

template <>
void Free<platform::CPUPlace>(platform::CPUPlace place, void* p) {
  GetCPUBuddyAllocator()->Free(p);
}

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

65
#ifdef PADDLE_WITH_CUDA
L
liaogang 已提交
66

67 68 69
BuddyAllocator* GetGPUBuddyAllocator(int gpu_id) {
  using BuddyAllocVec = std::vector<BuddyAllocator*>;
  static std::unique_ptr<BuddyAllocVec, void (*)(BuddyAllocVec * p)> as{
70
      new BuddyAllocVec, [](BuddyAllocVec* p) {
71 72 73 74
        std::for_each(p->begin(), p->end(),
                      [](BuddyAllocator* p) { delete p; });
      }};

L
liaogang 已提交
75 76
  // GPU buddy allocators
  auto& allocators = *as.get();
77 78

  // GPU buddy allocator initialization
L
liaogang 已提交
79
  std::call_once(gpu_allocator_flag, [&]() {
80
    int gpu_num = platform::GetCUDADeviceCount();
L
liaogang 已提交
81
    allocators.reserve(gpu_num);
L
liaogang 已提交
82 83
    for (int gpu = 0; gpu < gpu_num; gpu++) {
      platform::SetDeviceId(gpu);
L
liaogang 已提交
84 85 86
      allocators.emplace_back(new BuddyAllocator(new detail::GPUAllocator,
                                                 platform::GpuMinChunkSize(),
                                                 platform::GpuMaxChunkSize()));
L
liaogang 已提交
87
    }
L
liaogang 已提交
88 89 90 91 92
    VLOG(3) << "\n\nNOTE: each GPU device use "
            << FLAGS_fraction_of_gpu_memory_to_use * 100 << "% of GPU memory.\n"
            << "You can set environment variable '"
            << platform::kEnvFractionGpuMemoryToUse
            << "' to change the fraction of GPU usage.\n\n";
93 94
  });

Q
qijun 已提交
95
  platform::SetDeviceId(gpu_id);
L
liaogang 已提交
96
  return allocators[gpu_id];
L
liaogang 已提交
97 98
}

L
liaogang 已提交
99 100 101 102
template <>
void* Alloc<platform::GPUPlace>(platform::GPUPlace place, size_t size) {
  return GetGPUBuddyAllocator(place.device)->Alloc(size);
}
L
liaogang 已提交
103

L
liaogang 已提交
104 105 106
template <>
void Free<platform::GPUPlace>(platform::GPUPlace place, void* p) {
  GetGPUBuddyAllocator(place.device)->Free(p);
107 108
}

L
liaogang 已提交
109 110 111
template <>
size_t Used<platform::GPUPlace>(platform::GPUPlace place) {
  return GetGPUBuddyAllocator(place.device)->Used();
112 113 114
}

#endif  // PADDLE_ONLY_CPU
115 116 117

}  // namespace memory
}  // namespace paddle