mxnet
lazy_alloc_array.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
26 #ifndef MXNET_COMMON_LAZY_ALLOC_ARRAY_H_
27 #define MXNET_COMMON_LAZY_ALLOC_ARRAY_H_
28 
29 #include <dmlc/logging.h>
30 #include <memory>
31 #include <mutex>
32 #include <array>
33 #include <vector>
34 #include <atomic>
35 
36 namespace mxnet {
37 namespace common {
38 
39 template<typename TElem>
41  public:
49  template<typename FCreate>
50  inline std::shared_ptr<TElem> Get(int index, FCreate creator);
55  template<typename FVisit>
56  inline void ForEach(FVisit fvisit);
58  inline void Clear();
59 
60  private:
61  template<typename SyncObject>
62  class unique_unlock {
63  public:
64  explicit unique_unlock(std::unique_lock<SyncObject> *lock)
65  : lock_(lock) {
66  if (lock_) {
67  lock_->unlock();
68  }
69  }
70  ~unique_unlock() {
71  if (lock_) {
72  lock_->lock();
73  }
74  }
75  private:
76  std::unique_lock<SyncObject> *lock_;
77  };
78 
80  static constexpr std::size_t kInitSize = 16;
82  std::mutex create_mutex_;
84  std::array<std::shared_ptr<TElem>, kInitSize> head_;
86  std::vector<std::shared_ptr<TElem> > more_;
88  std::atomic<bool> is_clearing_;
89 };
90 
91 template<typename TElem>
93  : is_clearing_(false) {
94 }
95 
96 // implementations
97 template<typename TElem>
98 template<typename FCreate>
99 inline std::shared_ptr<TElem> LazyAllocArray<TElem>::Get(int index, FCreate creator) {
100  CHECK_GE(index, 0);
101  size_t idx = static_cast<size_t>(index);
102  if (idx < kInitSize) {
103  std::shared_ptr<TElem> ptr = head_[idx];
104  if (ptr) {
105  return ptr;
106  } else {
107  std::lock_guard<std::mutex> lock(create_mutex_);
108  if (!is_clearing_.load()) {
109  std::shared_ptr<TElem> ptr = head_[idx];
110  if (ptr) {
111  return ptr;
112  }
113  ptr = head_[idx] = std::shared_ptr<TElem>(creator());
114  return ptr;
115  }
116  }
117  } else {
118  std::lock_guard<std::mutex> lock(create_mutex_);
119  if (!is_clearing_.load()) {
120  idx -= kInitSize;
121  if (more_.size() <= idx) {
122  more_.reserve(idx + 1);
123  while (more_.size() <= idx) {
124  more_.push_back(std::shared_ptr<TElem>(nullptr));
125  }
126  }
127  std::shared_ptr<TElem> ptr = more_[idx];
128  if (ptr) {
129  return ptr;
130  }
131  ptr = more_[idx] = std::shared_ptr<TElem>(creator());
132  return ptr;
133  }
134  }
135  return nullptr;
136 }
137 
138 template<typename TElem>
140  std::unique_lock<std::mutex> lock(create_mutex_);
141  is_clearing_.store(true);
142  // Currently, head_ and more_ never get smaller, so it's safe to
143  // iterate them outside of the lock. The loops should catch
144  // any growth which might happen when create_mutex_ is unlocked
145  for (size_t i = 0; i < head_.size(); ++i) {
146  std::shared_ptr<TElem> p = head_[i];
147  head_[i] = std::shared_ptr<TElem>(nullptr);
148  unique_unlock<std::mutex> unlocker(&lock);
149  p = std::shared_ptr<TElem>(nullptr);
150  }
151  for (size_t i = 0; i < more_.size(); ++i) {
152  std::shared_ptr<TElem> p = more_[i];
153  more_[i] = std::shared_ptr<TElem>(nullptr);
154  unique_unlock<std::mutex> unlocker(&lock);
155  p = std::shared_ptr<TElem>(nullptr);
156  }
157  more_.clear();
158  is_clearing_.store(false);
159 }
160 
161 template<typename TElem>
162 template<typename FVisit>
163 inline void LazyAllocArray<TElem>::ForEach(FVisit fvisit) {
164  std::lock_guard<std::mutex> lock(create_mutex_);
165  for (size_t i = 0; i < head_.size(); ++i) {
166  if (head_[i].get() != nullptr) {
167  fvisit(i, head_[i].get());
168  }
169  }
170  for (size_t i = 0; i < more_.size(); ++i) {
171  if (more_[i].get() != nullptr) {
172  fvisit(i + kInitSize, more_[i].get());
173  }
174  }
175 }
176 
177 } // namespace common
178 } // namespace mxnet
179 #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:99
namespace of mxnet
Definition: base.h:118
void Clear()
clear all the allocated elements in array
Definition: lazy_alloc_array.h:139
void ForEach(FVisit fvisit)
for each not null element of the array, call fvisit
Definition: lazy_alloc_array.h:163
Definition: lazy_alloc_array.h:40
LazyAllocArray()
Definition: lazy_alloc_array.h:92