memory.cc 3.1 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 "paddle/memory/detail/buddy_allocator.h"
#include "paddle/memory/detail/system_allocator.h"
L
liaogang 已提交
24

25 26 27
namespace paddle {
namespace memory {

28 29
using BuddyAllocator = detail::BuddyAllocator;

L
liaogang 已提交
30 31
std::once_flag cpu_allocator_flag;
std::once_flag gpu_allocator_flag;
32 33

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

L
liaogang 已提交
36
  std::call_once(cpu_allocator_flag, [&]() {
37 38 39 40 41 42
    a.reset(new BuddyAllocator(new detail::CPUAllocator,
                               platform::CpuMinChunkSize(),
                               platform::CpuMaxChunkSize()));
  });

  return a.get();
L
liaogang 已提交
43 44
}

L
liaogang 已提交
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
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();
}

#ifndef PADDLE_ONLY_CPU
L
liaogang 已提交
61

62 63 64
BuddyAllocator* GetGPUBuddyAllocator(int gpu_id) {
  using BuddyAllocVec = std::vector<BuddyAllocator*>;
  static std::unique_ptr<BuddyAllocVec, void (*)(BuddyAllocVec * p)> as{
65
      new BuddyAllocVec, [](BuddyAllocVec* p) {
66 67 68 69
        std::for_each(p->begin(), p->end(),
                      [](BuddyAllocator* p) { delete p; });
      }};

L
liaogang 已提交
70 71
  // GPU buddy allocators
  auto& allocators = *as.get();
72 73

  // GPU buddy allocator initialization
L
liaogang 已提交
74
  std::call_once(gpu_allocator_flag, [&]() {
L
liaogang 已提交
75
    int gpu_num = platform::GetDeviceCount();
L
liaogang 已提交
76
    allocators.reserve(gpu_num);
L
liaogang 已提交
77 78
    for (int gpu = 0; gpu < gpu_num; gpu++) {
      platform::SetDeviceId(gpu);
L
liaogang 已提交
79 80 81
      allocators.emplace_back(new BuddyAllocator(new detail::GPUAllocator,
                                                 platform::GpuMinChunkSize(),
                                                 platform::GpuMaxChunkSize()));
L
liaogang 已提交
82
    }
83 84
  });

Q
qijun 已提交
85
  platform::SetDeviceId(gpu_id);
L
liaogang 已提交
86
  return allocators[gpu_id];
L
liaogang 已提交
87 88
}

L
liaogang 已提交
89 90 91 92
template <>
void* Alloc<platform::GPUPlace>(platform::GPUPlace place, size_t size) {
  return GetGPUBuddyAllocator(place.device)->Alloc(size);
}
L
liaogang 已提交
93

L
liaogang 已提交
94 95 96
template <>
void Free<platform::GPUPlace>(platform::GPUPlace place, void* p) {
  GetGPUBuddyAllocator(place.device)->Free(p);
97 98
}

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

#endif  // PADDLE_ONLY_CPU
105 106 107

}  // namespace memory
}  // namespace paddle