25 #ifndef MXNET_COMMON_LAZY_ALLOC_ARRAY_H_ 26 #define MXNET_COMMON_LAZY_ALLOC_ARRAY_H_ 28 #include <dmlc/logging.h> 38 template<
typename TElem>
48 template<
typename FCreate>
49 inline std::shared_ptr<TElem>
Get(
int index, FCreate creator);
54 template<
typename FVisit>
55 inline void ForEach(FVisit fvisit);
60 template<
typename SyncObject>
63 explicit unique_unlock(std::unique_lock<SyncObject> *lock)
75 std::unique_lock<SyncObject> *lock_;
79 static constexpr std::size_t kInitSize = 16;
81 std::mutex create_mutex_;
83 std::array<std::shared_ptr<TElem>, kInitSize> head_;
85 std::vector<std::shared_ptr<TElem> > more_;
87 std::atomic<bool> is_clearing_;
90 template<
typename TElem>
92 : is_clearing_(false) {
96 template<
typename TElem>
97 template<
typename FCreate>
100 size_t idx =
static_cast<size_t>(index);
101 if (idx < kInitSize) {
102 std::shared_ptr<TElem> ptr = head_[idx];
106 std::lock_guard<std::mutex> lock(create_mutex_);
107 if (!is_clearing_.load()) {
108 std::shared_ptr<TElem> ptr = head_[idx];
112 ptr = head_[idx] = std::shared_ptr<TElem>(creator());
117 std::lock_guard<std::mutex> lock(create_mutex_);
118 if (!is_clearing_.load()) {
120 if (more_.size() <= idx) {
121 more_.reserve(idx + 1);
122 while (more_.size() <= idx) {
123 more_.push_back(std::shared_ptr<TElem>(
nullptr));
126 std::shared_ptr<TElem> ptr = more_[idx];
130 ptr = more_[idx] = std::shared_ptr<TElem>(creator());
137 template<
typename TElem>
139 std::unique_lock<std::mutex> lock(create_mutex_);
140 is_clearing_.store(
true);
144 for (
size_t i = 0; i < head_.size(); ++i) {
145 std::shared_ptr<TElem> p = head_[i];
146 head_[i] = std::shared_ptr<TElem>(
nullptr);
147 unique_unlock<std::mutex> unlocker(&lock);
148 p = std::shared_ptr<TElem>(
nullptr);
150 for (
size_t i = 0; i < more_.size(); ++i) {
151 std::shared_ptr<TElem> p = more_[i];
152 more_[i] = std::shared_ptr<TElem>(
nullptr);
153 unique_unlock<std::mutex> unlocker(&lock);
154 p = std::shared_ptr<TElem>(
nullptr);
157 is_clearing_.store(
false);
160 template<
typename TElem>
161 template<
typename FVisit>
163 std::lock_guard<std::mutex> lock(create_mutex_);
164 for (
size_t i = 0; i < head_.size(); ++i) {
165 if (head_[i].
get() !=
nullptr) {
166 fvisit(i, head_[i].
get());
169 for (
size_t i = 0; i < more_.size(); ++i) {
170 if (more_[i].
get() !=
nullptr) {
171 fvisit(i + kInitSize, more_[i].
get());
178 #endif // MXNET_COMMON_LAZY_ALLOC_ARRAY_H_ std::shared_ptr< TElem > Get(int index, FCreate creator)
Get element of corresponding index, if it is not created create by creator.
Definition: lazy_alloc_array.h:98
namespace of mxnet
Definition: api_registry.h:33
void Clear()
clear all the allocated elements in array
Definition: lazy_alloc_array.h:138
void ForEach(FVisit fvisit)
for each not null element of the array, call fvisit
Definition: lazy_alloc_array.h:162
Definition: lazy_alloc_array.h:39
LazyAllocArray()
Definition: lazy_alloc_array.h:91