mxnet
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros
lazy_alloc_array.h
Go to the documentation of this file.
1 
7 #ifndef MXNET_COMMON_LAZY_ALLOC_ARRAY_H_
8 #define MXNET_COMMON_LAZY_ALLOC_ARRAY_H_
9 
10 #include <dmlc/logging.h>
11 #include <memory>
12 #include <mutex>
13 #include <array>
14 #include <vector>
15 
16 namespace mxnet {
17 namespace common {
18 
19 template<typename TElem>
21  public:
28  template<typename FCreate>
29  inline TElem* Get(int index, FCreate creator);
34  template<typename FVisit>
35  inline void ForEach(FVisit fvisit);
37  inline void Clear();
38 
39  private:
41  static constexpr std::size_t kInitSize = 16;
43  std::mutex create_mutex_;
45  std::array<std::unique_ptr<TElem>, kInitSize> head_;
47  std::vector<std::unique_ptr<TElem> > more_;
48 };
49 
50 // implementations
51 template<typename TElem>
52 template<typename FCreate>
53 inline TElem* LazyAllocArray<TElem>::Get(int index, FCreate creator) {
54  CHECK_GE(index, 0);
55  size_t idx = static_cast<size_t>(index);
56  if (idx < kInitSize) {
57  TElem *ptr = head_[idx].get();
58  if (ptr != nullptr) {
59  return ptr;
60  } else {
61  std::lock_guard<std::mutex> lock(create_mutex_);
62  TElem *ptr = head_[idx].get();
63  if (ptr != nullptr) return ptr;
64  head_[idx].reset(ptr = creator());
65  return ptr;
66  }
67  } else {
68  std::lock_guard<std::mutex> lock(create_mutex_);
69  idx -= kInitSize;
70  if (more_.size() <= idx) more_.resize(idx + 1);
71  TElem *ptr = more_[idx].get();
72  if (ptr != nullptr) return ptr;
73  more_[idx].reset(ptr = creator());
74  return ptr;
75  }
76 }
77 
78 template<typename TElem>
80  std::lock_guard<std::mutex> lock(create_mutex_);
81  for (size_t i = 0; i < head_.size(); ++i) {
82  head_[i].reset(nullptr);
83  }
84  for (size_t i = 0; i < more_.size(); ++i) {
85  more_[i].reset(nullptr);
86  }
87 }
88 
89 template<typename TElem>
90 template<typename FVisit>
91 inline void LazyAllocArray<TElem>::ForEach(FVisit fvisit) {
92  std::lock_guard<std::mutex> lock(create_mutex_);
93  for (size_t i = 0; i < head_.size(); ++i) {
94  if (head_[i].get() != nullptr) {
95  fvisit(i, head_[i].get());
96  }
97  }
98  for (size_t i = 0; i < more_.size(); ++i) {
99  if (more_[i].get() != nullptr) {
100  fvisit(i + kInitSize, more_[i].get());
101  }
102  }
103 }
104 } // namespace common
105 } // namespace mxnet
106 #endif // MXNET_COMMON_LAZY_ALLOC_ARRAY_H_
TElem * Get(int index, FCreate creator)
Get element of corresponding index, if it is not created create by creator.
Definition: lazy_alloc_array.h:53
void Clear()
clear all the allocated elements in array
Definition: lazy_alloc_array.h:79
void ForEach(FVisit fvisit)
for each not null element of the array, call fvisit
Definition: lazy_alloc_array.h:91
Definition: lazy_alloc_array.h:20