malloc.h 2.8 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 15 16

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

#pragma once

Y
Yi Wang 已提交
17
#include "paddle/fluid/platform/place.h"
18 19 20 21

namespace paddle {
namespace memory {

22 23 24 25 26 27 28 29 30 31 32 33 34
/**
 * \brief   Allocate memory block in one place.
 *
 * \param[in]  place  Allocation place (CPU or GPU).
 * \param[in]  size   Allocation size.
 *
 * \return  Allocated memory block address.
 *
 * \note    If return nullptr, it indicates memory allocation failed
 *          because insufficient memory in current system. When Alloc
 *          function is invoked, you must check the returned memory
 *          address is valid or not.
 */
L
liaogang 已提交
35
template <typename Place>
C
chengduoZH 已提交
36
void* Alloc(Place place, size_t size);
L
liaogang 已提交
37

38 39 40 41 42 43 44
/**
 * \brief   Free memory block in one place.
 *
 * \param[in]  place  Allocation place (CPU or GPU).
 * \param[in]  ptr    Memory block address to free.
 *
 */
L
liaogang 已提交
45
template <typename Place>
C
chengduoZH 已提交
46
void Free(Place place, void* ptr);
L
liaogang 已提交
47

48 49 50 51 52 53
/**
 * \brief   Total size of used memory in one place.
 *
 * \param[in]  place  Allocation place (CPU or GPU).
 *
 */
L
liaogang 已提交
54
template <typename Place>
55
size_t Used(Place place);
56

57 58 59
struct Usage : public boost::static_visitor<size_t> {
  size_t operator()(const platform::CPUPlace& cpu) const;
  size_t operator()(const platform::CUDAPlace& gpu) const;
C
chengduoZH 已提交
60
  size_t operator()(const platform::CUDAPinnedPlace& cuda_pinned) const;
61 62 63 64
};

size_t memory_usage(const platform::Place& p);

65 66 67 68 69 70 71 72
/**
 * \brief   Free memory block in one place.
 *
 * \note    In some cases, custom deleter is used to
 *          deallocate the memory automatically for
 *          std::unique_ptr<T> in tensor.h.
 *
 */
73
template <typename T, typename Place>
74
class PODDeleter {
75 76
  static_assert(std::is_pod<T>::value, "T must be POD");

F
fengjiayi 已提交
77
 public:
C
chengduoZH 已提交
78 79
  explicit PODDeleter(Place place) : place_(place) {}
  void operator()(T* ptr) { Free(place_, static_cast<void*>(ptr)); }
F
fengjiayi 已提交
80 81

 private:
82
  Place place_;
F
fengjiayi 已提交
83 84
};

D
dzhwinter 已提交
85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
/**
 * \brief   Free memory block in one place does not meet POD
 *
 * \note    In some cases, custom deleter is used to
 *          deallocate the memory automatically for
 *          std::unique_ptr<T> in tensor.h.
 *
 */
template <typename T, typename Place>
class PlainDeleter {
 public:
  explicit PlainDeleter(Place place) : place_(place) {}
  void operator()(T* ptr) { Free(place_, reinterpret_cast<void*>(ptr)); }

 private:
  Place place_;
};

103 104
}  // namespace memory
}  // namespace paddle