system_allocator.cc 3.8 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

#include <stdlib.h>    // for malloc and free
#include <sys/mman.h>  // for mlock and munlock
22
#include <algorithm>   // for std::max
23 24 25 26 27 28 29 30

#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.
31
DEFINE_bool(use_pinned_memory, true, "If set, allocate cpu pinned memory.");
32
DECLARE_double(fraction_of_gpu_memory_to_use);
33 34 35 36
namespace paddle {
namespace memory {
namespace detail {

L
liaogang 已提交
37
void* CPUAllocator::Alloc(size_t& index, size_t size) {
38 39 40 41 42
  // 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;

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

45 46
  void* p;

T
tensor-tang 已提交
47
#ifdef PADDLE_WITH_MKLDNN
48 49 50 51 52 53 54
  // refer to https://github.com/01org/mkl-dnn/blob/master/include/mkldnn.hpp
  // memory alignment
  PADDLE_ENFORCE_EQ(posix_memalign(&p, 4096ul, size), 0);
#else
  PADDLE_ENFORCE_EQ(posix_memalign(&p, 32ul, size), 0);
#endif
  PADDLE_ENFORCE(p, "Fail to allocate CPU memory: size = %d .", size);
55 56 57 58 59 60

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

63 64 65
  return p;
}

L
liaogang 已提交
66
void CPUAllocator::Free(void* p, size_t size, size_t index) {
67
  if (p != nullptr && index == 1) {
68 69 70 71 72
    munlock(p, size);
  }
  free(p);
}

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

75
#ifdef PADDLE_WITH_CUDA
76

L
liaogang 已提交
77
void* GPUAllocator::Alloc(size_t& index, size_t size) {
78 79
  // CUDA documentation doesn't explain if cudaMalloc returns nullptr
  // if size is 0.  We just make sure it does.
L
liaogang 已提交
80
  if (size <= 0) return nullptr;
81 82
  void* p;
  cudaError_t result = cudaMalloc(&p, size);
L
liaogang 已提交
83
  if (result == cudaSuccess) {
84 85
    index = 0;
    gpu_alloc_size_ += size;
L
liaogang 已提交
86
    return p;
87 88 89 90 91 92 93
  } else {
    LOG(WARNING)
        << "Cannot malloc " << size / 1024.0 / 1024.0
        << " MB GPU memory. Please shrink FLAGS_fraction_of_gpu_memory_to_use "
           "environment variable to a lower value. Current value is "
        << FLAGS_fraction_of_gpu_memory_to_use;
    return nullptr;
L
liaogang 已提交
94
  }
95 96
}

L
liaogang 已提交
97
void GPUAllocator::Free(void* p, size_t size, size_t index) {
98 99 100 101 102 103 104 105 106 107 108 109
  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);
  }

110 111 112 113 114 115
  // 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 已提交
116
    PADDLE_ENFORCE(err, "cudaFree{Host} failed in GPUAllocator::Free.");
117 118 119
  }
}

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

L
Luo Tao 已提交
122
#endif
123 124 125 126

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