|
11 | 11 |
|
12 | 12 | #include <cstdint> |
13 | 13 | #include <cstring> |
| 14 | +#include <memory> |
14 | 15 |
|
15 | 16 | #include "coreneuron/utils/nrn_assert.h" |
16 | 17 | #include "coreneuron/nrniv/nrniv_decl.h" |
|
21 | 22 | #define NRN_SOA_BYTE_ALIGN (8 * sizeof(double)) |
22 | 23 | #endif |
23 | 24 |
|
| 25 | +namespace coreneuron { |
| 26 | +/** @brief Check if allocate_unified will return a unified memory address. |
| 27 | + * |
| 28 | + * If false, [de]allocate_unified simply forward to new/delete. It is |
| 29 | + * convenient to include this method here to avoid having to access |
| 30 | + * corenrn_param directly. |
| 31 | + */ |
| 32 | +bool unified_memory_enabled(); |
| 33 | + |
| 34 | +/** @brief Allocate unified memory in GPU builds iff GPU enabled, otherwise new |
| 35 | + */ |
| 36 | +void* allocate_unified(std::size_t num_bytes); |
| 37 | + |
| 38 | +/** @brief Deallocate memory allocated by `allocate_unified`. |
| 39 | + */ |
| 40 | +void deallocate_unified(void* ptr, std::size_t num_bytes); |
| 41 | + |
| 42 | +/** @brief C++ allocator that uses [de]allocate_unified. |
| 43 | + */ |
| 44 | +template <typename T> |
| 45 | +struct unified_allocator { |
| 46 | + using value_type = T; |
| 47 | + |
| 48 | + unified_allocator() = default; |
| 49 | + |
| 50 | + template <typename U> |
| 51 | + unified_allocator(unified_allocator<U> const&) noexcept {} |
| 52 | + |
| 53 | + value_type* allocate(std::size_t n) { |
| 54 | + return static_cast<value_type*>(allocate_unified(n * sizeof(value_type))); |
| 55 | + } |
| 56 | + |
| 57 | + void deallocate(value_type* p, std::size_t n) noexcept { |
| 58 | + deallocate_unified(p, n * sizeof(value_type)); |
| 59 | + } |
| 60 | +}; |
| 61 | + |
| 62 | +template <typename T, typename U> |
| 63 | +bool operator==(unified_allocator<T> const&, unified_allocator<U> const&) noexcept { |
| 64 | + return true; |
| 65 | +} |
| 66 | + |
| 67 | +template <typename T, typename U> |
| 68 | +bool operator!=(unified_allocator<T> const& x, unified_allocator<U> const& y) noexcept { |
| 69 | + return !(x == y); |
| 70 | +} |
| 71 | + |
| 72 | +/** @brief Allocator-aware deleter for use with std::unique_ptr. |
| 73 | + * |
| 74 | + * This is copied from https://stackoverflow.com/a/23132307. See also |
| 75 | + * http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2017/p0316r0.html, |
| 76 | + * http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2020/p0211r3.html, and |
| 77 | + * boost::allocate_unique<...>. |
| 78 | + * Hopefully std::allocate_unique will be included in C++23. |
| 79 | + */ |
| 80 | +template <typename Alloc> |
| 81 | +struct alloc_deleter { |
| 82 | + alloc_deleter() = default; // OL210813 addition |
| 83 | + alloc_deleter(const Alloc& a) |
| 84 | + : a(a) {} |
| 85 | + |
| 86 | + typedef typename std::allocator_traits<Alloc>::pointer pointer; |
| 87 | + |
| 88 | + void operator()(pointer p) const { |
| 89 | + Alloc aa(a); |
| 90 | + std::allocator_traits<Alloc>::destroy(aa, std::addressof(*p)); |
| 91 | + std::allocator_traits<Alloc>::deallocate(aa, p, 1); |
| 92 | + } |
| 93 | + |
| 94 | + private: |
| 95 | + Alloc a; |
| 96 | +}; |
| 97 | + |
| 98 | +template <typename T, typename Alloc, typename... Args> |
| 99 | +auto allocate_unique(const Alloc& alloc, Args&&... args) { |
| 100 | + using AT = std::allocator_traits<Alloc>; |
| 101 | + static_assert(std::is_same<typename AT::value_type, std::remove_cv_t<T>>{}(), |
| 102 | + "Allocator has the wrong value_type"); |
| 103 | + |
| 104 | + Alloc a(alloc); |
| 105 | + auto p = AT::allocate(a, 1); |
| 106 | + try { |
| 107 | + AT::construct(a, std::addressof(*p), std::forward<Args>(args)...); |
| 108 | + using D = alloc_deleter<Alloc>; |
| 109 | + return std::unique_ptr<T, D>(p, D(a)); |
| 110 | + } catch (...) { |
| 111 | + AT::deallocate(a, p, 1); |
| 112 | + throw; |
| 113 | + } |
| 114 | +} |
| 115 | +} // namespace coreneuron |
| 116 | + |
24 | 117 | /// for gpu builds with unified memory support |
25 | 118 | /// OL210812: why do we include __CUDACC__ here? |
26 | 119 | #if (defined(__CUDACC__) || defined(CORENEURON_UNIFIED_MEMORY)) |
|
0 commit comments