memory.h 2.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* 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. */

#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>
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>
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 60 61 62 63
struct Usage : public boost::static_visitor<size_t> {
  size_t operator()(const platform::CPUPlace& cpu) const;
  size_t operator()(const platform::CUDAPlace& gpu) const;
};

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

64 65 66 67 68 69 70 71
/**
 * \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.
 *
 */
72
template <typename T, typename Place>
73
class PODDeleter {
74 75
  static_assert(std::is_pod<T>::value, "T must be POD");

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

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

D
dzhwinter 已提交
84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101
/**
 * \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_;
};

102 103
}  // namespace memory
}  // namespace paddle