1 /**
2 * Copyright 2019 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16 #ifndef MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_ALLOCATOR_H_
17 #define MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_ALLOCATOR_H_
18
19 #include <cstdlib>
20 #include <functional>
21 #include <memory>
22 #include <type_traits>
23 #include <utility>
24 #include "minddata/dataset/util/memory_pool.h"
25
26 namespace mindspore {
27 namespace dataset {
28 // The following conforms to the requirements of
29 // std::allocator. Do not rename/change any needed
30 // requirements, e.g. function names, typedef etc.
31 template <typename T>
32 class Allocator {
33 public:
34 template <typename U>
35 friend class Allocator;
36
37 using value_type = T;
38 using pointer = T *;
39 using const_pointer = const T *;
40 using reference = T &;
41 using const_reference = const T &;
42 using size_type = uint64_t;
43 using difference_type = std::ptrdiff_t;
44
45 template <typename U>
46 struct rebind {
47 using other = Allocator<U>;
48 };
49
50 using propagate_on_container_copy_assignment = std::true_type;
51 using propagate_on_container_move_assignment = std::true_type;
52 using propagate_on_container_swap = std::true_type;
53
Allocator(std::shared_ptr<MemoryPool> b)54 explicit Allocator(std::shared_ptr<MemoryPool> b) : pool_(std::move(b)) {}
55
56 ~Allocator() = default;
57
58 template <typename U>
Allocator(Allocator<U> const & rhs)59 explicit Allocator(Allocator<U> const &rhs) : pool_(rhs.pool_) {}
60
61 template <typename U>
62 bool operator==(Allocator<U> const &rhs) const {
63 return pool_ == rhs.pool_;
64 }
65
66 template <typename U>
67 bool operator!=(Allocator<U> const &rhs) const {
68 return pool_ != rhs.pool_;
69 }
70
allocate(std::size_t n)71 pointer allocate(std::size_t n) {
72 void *p = nullptr;
73 Status rc = pool_->Allocate(n * sizeof(T), &p);
74 if (rc.IsOk()) {
75 return reinterpret_cast<pointer>(p);
76 } else if (rc == StatusCode::kMDOutOfMemory) {
77 MS_LOG(ERROR) << rc.ToString();
78 return nullptr;
79 } else {
80 MS_LOG(ERROR) << rc.ToString();
81 return nullptr;
82 }
83 }
84
85 void deallocate(pointer p, std::size_t n = 0) noexcept { pool_->Deallocate(p); }
86
max_size()87 size_type max_size() { return pool_->get_max_size(); }
88
89 private:
90 std::shared_ptr<MemoryPool> pool_;
91 };
92
93 /// \brief It is a wrapper of unique_ptr with a custom Allocator class defined above
94 template <typename T, typename C = std::allocator<T>, typename... Args>
MakeUnique(std::unique_ptr<T[],std::function<void (T *)>> * out,C alloc,size_t n,Args &&...args)95 Status MakeUnique(std::unique_ptr<T[], std::function<void(T *)>> *out, C alloc, size_t n, Args &&... args) {
96 RETURN_UNEXPECTED_IF_NULL(out);
97 CHECK_FAIL_RETURN_UNEXPECTED(n > 0, "size must be positive");
98 T *data = nullptr;
99 try {
100 data = alloc.allocate(n);
101 // Some of our implementation of allocator (e.g. NumaAllocator) don't throw std::bad_alloc.
102 // So we have to catch for null ptr
103 if (data == nullptr) {
104 return Status(StatusCode::kMDOutOfMemory);
105 }
106 if (!std::is_arithmetic<T>::value) {
107 for (size_t i = 0; i < n; i++) {
108 std::allocator_traits<C>::construct(alloc, &(data[i]), std::forward<Args>(args)...);
109 }
110 }
111 auto deleter = [](T *p, C f_alloc, size_t f_n) {
112 if (!std::is_arithmetic<T>::value && std::is_destructible<T>::value) {
113 for (size_t i = 0; i < f_n; ++i) {
114 std::allocator_traits<C>::destroy(f_alloc, &p[i]);
115 }
116 }
117 f_alloc.deallocate(p, f_n);
118 };
119 *out = std::unique_ptr<T[], std::function<void(T *)>>(data, std::bind(deleter, std::placeholders::_1, alloc, n));
120 } catch (const std::bad_alloc &e) {
121 if (data != nullptr) {
122 alloc.deallocate(data, n);
123 }
124 return Status(StatusCode::kMDOutOfMemory);
125 } catch (const std::exception &e) {
126 if (data != nullptr) {
127 alloc.deallocate(data, n);
128 }
129 RETURN_STATUS_UNEXPECTED(e.what());
130 }
131 return Status::OK();
132 }
133
134 /// \brief It is a wrapper of the above custom unique_ptr with some additional methods
135 /// \tparam T The type of object to be allocated
136 /// \tparam C Allocator. Default to std::allocator
137 template <typename T, typename C = std::allocator<T>>
138 class MemGuard {
139 public:
140 using allocator = C;
MemGuard()141 MemGuard() : n_(0) {}
MemGuard(const allocator & a)142 explicit MemGuard(const allocator &a) : n_(0), alloc_(a) {}
143 // There is no copy constructor nor assignment operator because the memory is solely owned by this object.
144 MemGuard(const MemGuard &) = delete;
145 MemGuard &operator=(const MemGuard &) = delete;
146 // On the other hand, We can support move constructor
MemGuard(MemGuard && lhs)147 MemGuard(MemGuard &&lhs) noexcept : n_(lhs.n_), alloc_(std::move(lhs.alloc_)), ptr_(std::move(lhs.ptr_)) {}
148 MemGuard &operator=(MemGuard &&lhs) noexcept {
149 if (this != &lhs) {
150 this->deallocate();
151 n_ = lhs.n_;
152 alloc_ = std::move(lhs.alloc_);
153 ptr_ = std::move(lhs.ptr_);
154 }
155 return *this;
156 }
157 /// \brief Explicitly deallocate the memory if allocated
deallocate()158 void deallocate() {
159 if (ptr_) {
160 ptr_.reset();
161 }
162 }
163 /// \brief Allocate memory (with emplace feature). Previous one will be released. If size is 0, no new memory is
164 /// allocated.
165 /// \param n Number of objects of type T to be allocated
166 /// \tparam Args Extra arguments pass to the constructor of T
167 template <typename... Args>
allocate(size_t n,Args &&...args)168 Status allocate(size_t n, Args &&... args) noexcept {
169 deallocate();
170 n_ = n;
171 return MakeUnique(&ptr_, alloc_, n, std::forward<Args>(args)...);
172 }
~MemGuard()173 ~MemGuard() noexcept { deallocate(); }
174 /// \brief Getter function
175 /// \return The pointer to the memory allocated
GetPointer()176 T *GetPointer() const { return ptr_.get(); }
177 /// \brief Getter function
178 /// \return The pointer to the memory allocated
GetMutablePointer()179 T *GetMutablePointer() { return ptr_.get(); }
180 /// \brief Overload [] operator to access a particular element
181 /// \param x index to the element. Must be less than number of element allocated.
182 /// \return pointer to the x-th element
183 T *operator[](size_t x) { return GetMutablePointer() + x; }
184 /// \brief Overload [] operator to access a particular element
185 /// \param x index to the element. Must be less than number of element allocated.
186 /// \return pointer to the x-th element
187 T *operator[](size_t x) const { return GetPointer() + x; }
188 /// \brief Return how many bytes are allocated in total
189 /// \return Number of bytes allocated in total
GetSizeInBytes()190 size_t GetSizeInBytes() const { return n_ * sizeof(T); }
191
192 private:
193 size_t n_;
194 allocator alloc_;
195 std::unique_ptr<T[], std::function<void(T *)>> ptr_;
196 };
197 } // namespace dataset
198 } // namespace mindspore
199
200 #endif // MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_ALLOCATOR_H_
201