1 // Copyright (c) 2017-2021, Lawrence Livermore National Security, LLC and
2 // other Axom Project Developers. See the top-level LICENSE file for details.
3 //
4 // SPDX-License-Identifier: (BSD-3-Clause)
5
6 #ifndef AXOM_MEMORYMANAGEMENT_HPP_
7 #define AXOM_MEMORYMANAGEMENT_HPP_
8
9 // Axom includes
10 #include "axom/config.hpp" // for AXOM compile-time definitions
11 #include "axom/core/Macros.hpp"
12
13 // Umpire includes
14 #ifdef AXOM_USE_UMPIRE
15 #include "umpire/config.hpp"
16 #include "umpire/ResourceManager.hpp"
17 #include "umpire/op/MemoryOperationRegistry.hpp"
18 #include "umpire/resource/MemoryResourceTypes.hpp"
19 #include "umpire/strategy/QuickPool.hpp"
20 #else
21 #include <cstring> // for std::memcpy
22 #include <cstdlib> // for std::malloc, std::realloc, std::free
23 #endif
24
25 namespace axom
26 {
27 constexpr int INVALID_ALLOCATOR_ID = -1;
28
29 /*!
30 * \brief Memory spaces supported by Array-like types
31 *
32 * This abstraction is not implemented using Umpire's MemoryResourceType enum
33 * in order to also include a "Dynamic" option as a default template parameter
34 * for Array-like types
35 */
36 enum class MemorySpace
37 {
38 Dynamic,
39 #ifdef AXOM_USE_UMPIRE
40 Host,
41 Device,
42 Unified,
43 Pinned,
44 Constant
45 #endif
46 };
47
48 /// \name Memory Management Routines
49 /// @{
50
51 #ifdef AXOM_USE_UMPIRE
52
53 /*!
54 * \brief Returns the ID of the predefined allocator for a given resource.
55 * \param [in] resource_type the Umpire resource type
56 * \return ID the id of the predefined umpire allocator.
57 */
getUmpireResourceAllocatorID(umpire::resource::MemoryResourceType resource_type)58 inline int getUmpireResourceAllocatorID(
59 umpire::resource::MemoryResourceType resource_type)
60 {
61 umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
62 umpire::Allocator alloc = rm.getAllocator(resource_type);
63 return alloc.getId();
64 }
65
66 #endif
67
68 /*!
69 * \brief Sets the default memory allocator to use.
70 * \param [in] allocatorID ID of the Umpire allocator to use.
71 *
72 * \note This function has no effect when Axom is not compiled with Umpire.
73 */
setDefaultAllocator(int allocatorID)74 inline void setDefaultAllocator(int allocatorID)
75 {
76 #ifdef AXOM_USE_UMPIRE
77 umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
78 umpire::Allocator allocator = rm.getAllocator(allocatorID);
79 rm.setDefaultAllocator(allocator);
80 #else
81 AXOM_UNUSED_VAR(allocatorID);
82 #endif
83 }
84
85 /*!
86 * \brief Returns the ID of the current default allocator.
87 * \return ID the ID of the current default allocator.
88 * \post ID != INVALID_ALLOCATOR_ID
89 */
getDefaultAllocatorID()90 inline int getDefaultAllocatorID()
91 {
92 #ifdef AXOM_USE_UMPIRE
93 return umpire::ResourceManager::getInstance().getDefaultAllocator().getId();
94 #else
95 return 0;
96 #endif
97 }
98
99 /*!
100 * \brief Allocates a chunk of memory of type T.
101 *
102 * \param [in] n the number of elements to allocate.
103 * \param [in] allocID the Umpire allocator to use (optional)
104 *
105 * \tparam T the type of pointer returned.
106 *
107 * \note By default allocate() will use the current default allocator. The
108 * caller may explicitly specify a different allocator to use by supplying the
109 * second, optional argument, or change the default allocator by calling
110 * axom::setDefaultAllocator().
111 *
112 * \return p pointer to the new allocation or a nullptr if allocation failed.
113 */
114 template <typename T>
115 inline T* allocate(std::size_t n, int allocID = getDefaultAllocatorID()) noexcept;
116
117 /*!
118 * \brief Frees the chunk of memory pointed to by the supplied pointer, p.
119 * \param [in/out] p a pointer to memory allocated with allocate/reallocate or a
120 * nullptr.
121 * \post p == nullptr
122 */
123 template <typename T>
124 inline void deallocate(T*& p) noexcept;
125
126 /*!
127 * \brief Reallocates the chunk of memory pointed to by the supplied pointer.
128 *
129 * \param [in] p pointer to memory allocated with allocate/reallocate, or a
130 * nullptr.
131 * \param [in] n the number of elements to allocate.
132 * \param [in] allocID the ID of the allocator to use if pointer is null
133 * (optional)
134 *
135 * \tparam T the type pointer p points to.
136 *
137 * \return p pointer to the new allocation or a nullptr if allocation failed.
138 *
139 * \note When n == 0, this function returns a valid pointer (of size 0) in the
140 * current allocator's memory space. This follows the semantics of
141 * Umpire's reallocate function.
142 * \note When p is a null pointer, allocID is used to allocate the data.
143 * Otherwise, it is unused.
144 */
145 template <typename T>
146 inline T* reallocate(T* p,
147 std::size_t n,
148 int allocID = getDefaultAllocatorID()) noexcept;
149
150 /*!
151 * \brief Copies memory from the source to the destination.
152 *
153 * \param [in/out] dst the destination to copy to.
154 * \param [in] src the source to copy from.
155 * \param [in] numbytes the number of bytes to copy.
156 *
157 * \note When using Umpire if either src or dst is not registered with the
158 * ResourceManager then the default host allocation strategy is assumed for
159 * that pointer.
160 */
161 inline void copy(void* dst, const void* src, std::size_t numbytes) noexcept;
162
163 /// @}
164
165 //------------------------------------------------------------------------------
166 // IMPLEMENTATION
167 //------------------------------------------------------------------------------
168
169 template <typename T>
allocate(std::size_t n,int allocID)170 inline T* allocate(std::size_t n, int allocID) noexcept
171 {
172 const std::size_t numbytes = n * sizeof(T);
173
174 #ifdef AXOM_USE_UMPIRE
175
176 umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
177 umpire::Allocator allocator = rm.getAllocator(allocID);
178 return static_cast<T*>(allocator.allocate(numbytes));
179
180 #else
181 AXOM_UNUSED_VAR(allocID);
182 return static_cast<T*>(std::malloc(numbytes));
183 #endif
184 }
185 //------------------------------------------------------------------------------
186 template <typename T>
deallocate(T * & pointer)187 inline void deallocate(T*& pointer) noexcept
188 {
189 if(pointer == nullptr) return;
190
191 #ifdef AXOM_USE_UMPIRE
192
193 umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
194 rm.deallocate(pointer);
195
196 #else
197
198 std::free(pointer);
199
200 #endif
201
202 pointer = nullptr;
203 }
204
205 //------------------------------------------------------------------------------
206 template <typename T>
reallocate(T * pointer,std::size_t n,int allocID)207 inline T* reallocate(T* pointer, std::size_t n, int allocID) noexcept
208 {
209 const std::size_t numbytes = n * sizeof(T);
210
211 #if defined(AXOM_USE_UMPIRE)
212
213 umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
214 if(pointer == nullptr)
215 {
216 pointer = axom::allocate<T>(n, allocID);
217 }
218 else
219 {
220 pointer = static_cast<T*>(rm.reallocate(pointer, numbytes));
221 }
222
223 #else
224
225 pointer = static_cast<T*>(std::realloc(pointer, numbytes));
226
227 // Consistently handle realloc(0) for std::realloc to match Umpire's behavior
228 if(n == 0 && pointer == nullptr)
229 {
230 pointer = axom::allocate<T>(0);
231 }
232
233 AXOM_UNUSED_VAR(allocID);
234 #endif
235
236 return pointer;
237 }
238
239 //------------------------------------------------------------------------------
copy(void * dst,const void * src,std::size_t numbytes)240 inline void copy(void* dst, const void* src, std::size_t numbytes) noexcept
241 {
242 #ifdef AXOM_USE_UMPIRE
243 umpire::ResourceManager& rm = umpire::ResourceManager::getInstance();
244 umpire::op::MemoryOperationRegistry& op_registry =
245 umpire::op::MemoryOperationRegistry::getInstance();
246
247 auto dstStrategy = rm.getAllocator("HOST").getAllocationStrategy();
248 auto srcStrategy = dstStrategy;
249
250 using AllocationRecord = umpire::util::AllocationRecord;
251 AllocationRecord* dstRecord = nullptr;
252 AllocationRecord* srcRecord = nullptr;
253
254 if(rm.hasAllocator(dst))
255 {
256 dstRecord = const_cast<AllocationRecord*>(rm.findAllocationRecord(dst));
257 dstStrategy = dstRecord->strategy;
258 }
259
260 if(rm.hasAllocator(const_cast<void*>(src)))
261 {
262 srcRecord = const_cast<AllocationRecord*>(
263 rm.findAllocationRecord(const_cast<void*>(src)));
264 srcStrategy = srcRecord->strategy;
265 }
266
267 auto op = op_registry.find("COPY", srcStrategy, dstStrategy);
268 op->transform(const_cast<void*>(src), &dst, srcRecord, dstRecord, numbytes);
269 #else
270 std::memcpy(dst, src, numbytes);
271 #endif
272 }
273
274 namespace detail
275 {
276 /// \brief Translates between the MemorySpace enum and Umpire allocator IDs
277 template <MemorySpace SPACE>
278 inline int getAllocatorID();
279
280 template <>
getAllocatorID()281 inline int getAllocatorID<MemorySpace::Dynamic>()
282 {
283 return axom::getDefaultAllocatorID();
284 }
285
286 #ifdef AXOM_USE_UMPIRE
287
288 template <>
getAllocatorID()289 inline int getAllocatorID<MemorySpace::Host>()
290 {
291 return axom::getUmpireResourceAllocatorID(
292 umpire::resource::MemoryResourceType::Host);
293 }
294
295 template <>
getAllocatorID()296 inline int getAllocatorID<MemorySpace::Device>()
297 {
298 return axom::getUmpireResourceAllocatorID(
299 umpire::resource::MemoryResourceType::Device);
300 }
301
302 template <>
getAllocatorID()303 inline int getAllocatorID<MemorySpace::Unified>()
304 {
305 return axom::getUmpireResourceAllocatorID(
306 umpire::resource::MemoryResourceType::Unified);
307 }
308
309 template <>
getAllocatorID()310 inline int getAllocatorID<MemorySpace::Pinned>()
311 {
312 return axom::getUmpireResourceAllocatorID(
313 umpire::resource::MemoryResourceType::Pinned);
314 }
315
316 template <>
getAllocatorID()317 inline int getAllocatorID<MemorySpace::Constant>()
318 {
319 return axom::getUmpireResourceAllocatorID(
320 umpire::resource::MemoryResourceType::Constant);
321 }
322
323 #endif
324
325 } // namespace detail
326
327 } // namespace axom
328
329 #endif /* AXOM_MEMORYMANAGEMENT_HPP_ */
330