system_allocator.cc 4.2 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/detail/system_allocator.h"
L
liaogang 已提交
16
#include "paddle/platform/assert.h"
L
liaogang 已提交
17
#include "paddle/platform/enforce.h"
L
liaogang 已提交
18
#include "paddle/platform/gpu_info.h"
19 20 21 22 23 24 25 26 27 28 29

#include <stdlib.h>    // for malloc and free
#include <sys/mman.h>  // for mlock and munlock

#include "gflags/gflags.h"

// If use_pinned_memory is true, CPUAllocator calls mlock, which
// returns pinned and locked memory as staging areas for data exchange
// between host and device.  Allocates too much would reduce the amount
// of memory available to the system for paging.  So, by default, we
// should set false to use_pinned_memory.
L
liaogang 已提交
30
DEFINE_bool(use_pinned_memory, false, "If set, allocate cpu pinned memory.");
31 32 33 34 35

namespace paddle {
namespace memory {
namespace detail {

L
liaogang 已提交
36
void* CPUAllocator::Alloc(size_t& index, size_t size) {
37 38 39 40 41
  // According to http://www.cplusplus.com/reference/cstdlib/malloc/,
  // malloc might not return nullptr if size is zero, but the returned
  // pointer shall not be dereferenced -- so we make it nullptr.
  if (size <= 0) return nullptr;

42
  index = 0;  // unlock memory
L
liaogang 已提交
43

44
  void* p = malloc(size);
45 46 47 48 49 50

  if (p != nullptr) {
    if (FLAGS_use_pinned_memory) {
      index = 1;
      mlock(p, size);  // lock memory
    }
51
  }
52

53 54 55
  return p;
}

L
liaogang 已提交
56
void CPUAllocator::Free(void* p, size_t size, size_t index) {
57
  if (p != nullptr && index == 1) {
58 59 60 61 62
    munlock(p, size);
  }
  free(p);
}

L
liaogang 已提交
63
bool CPUAllocator::UseGpu() const { return false; }
L
liaogang 已提交
64

65 66
#ifndef PADDLE_ONLY_CPU

L
liaogang 已提交
67
void* GPUAllocator::Alloc(size_t& index, size_t size) {
68 69
  // CUDA documentation doesn't explain if cudaMalloc returns nullptr
  // if size is 0.  We just make sure it does.
L
liaogang 已提交
70
  if (size <= 0) return nullptr;
71

L
liaogang 已提交
72 73 74 75 76 77
  size_t available = 0;
  size_t capacity = 0;
  paddle::platform::GpuMemoryUsage(available, capacity);

  // Reserve memory for page tables, etc.
  size_t reserving = capacity - paddle::platform::GpuMaxAllocSize();
78
  size_t usable = available > reserving ? available - reserving : 0;
L
liaogang 已提交
79 80 81

  // If remaining size no less than expected size, using general
  // cudaMalloc to allocate GPU memory.
82
  void* p = 0;
83
  if (size <= usable) {
L
liaogang 已提交
84 85 86
    cudaError_t result = cudaMalloc(&p, size);
    if (result == cudaSuccess) {
      index = 0;
87
      gpu_alloc_size_ += size;
L
liaogang 已提交
88 89
      return p;
    }
90
  }
L
liaogang 已提交
91 92 93

  // If remaining size less than expected size or cudaMalloc failed,
  // cudaMallocHost will be considered as a fallback allocator.
94 95 96 97 98 99 100 101
  //
  // NOTE: here, we use GpuMaxAllocSize() as the maximum memory size
  // of host fallback allocation. Allocates too much would reduce
  // the amount of memory available to the underlying system for paging.
  usable = paddle::platform::GpuMaxAllocSize() - fallback_alloc_size_;

  if (size > usable) return nullptr;

L
liaogang 已提交
102 103 104
  cudaError_t result = cudaMallocHost(&p, size);
  if (result == cudaSuccess) {
    index = 1;
105
    fallback_alloc_size_ += size;
L
liaogang 已提交
106 107 108 109
    return p;
  }

  return nullptr;
110 111
}

L
liaogang 已提交
112
void GPUAllocator::Free(void* p, size_t size, size_t index) {
113 114 115 116 117 118 119 120 121 122 123 124
  cudaError_t err;

  if (index == 0) {
    PADDLE_ASSERT(gpu_alloc_size_ >= size);
    gpu_alloc_size_ -= size;
    err = cudaFree(p);
  } else {
    PADDLE_ASSERT(fallback_alloc_size_ >= size);
    fallback_alloc_size_ -= size;
    err = cudaFreeHost(p);
  }

125 126 127 128 129 130
  // Purposefully allow cudaErrorCudartUnloading, because
  // that is returned if you ever call cudaFree after the
  // driver has already shutdown. This happens only if the
  // process is terminating, in which case we don't care if
  // cudaFree succeeds.
  if (err != cudaErrorCudartUnloading) {
L
liaogang 已提交
131
    PADDLE_ENFORCE(err, "cudaFree{Host} failed in GPUAllocator::Free.");
132 133 134
  }
}

L
liaogang 已提交
135
bool GPUAllocator::UseGpu() const { return true; }
L
liaogang 已提交
136

137 138 139 140 141
#endif  // PADDLE_ONLY_CPU

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