24 #ifndef TILEDARRAY_CUDA_UM_ALLOCATOR_H___INCLUDED
25 #define TILEDARRAY_CUDA_UM_ALLOCATOR_H___INCLUDED
27 #include <TiledArray/config.h>
29 #ifdef TILEDARRAY_HAS_CUDA
41 class cuda_um_allocator_impl {
46 using const_reference =
const T&;
48 cuda_um_allocator_impl() noexcept
49 : um_dynamic_pool_(&cudaEnv::instance()->um_dynamic_pool()) {}
52 cuda_um_allocator_impl(
const cuda_um_allocator_impl<U>& rhs) noexcept
53 : um_dynamic_pool_(rhs.um_dynamic_pool_) {}
56 pointer allocate(
size_t n) {
57 pointer result =
nullptr;
61 result =
static_cast<pointer
>(um_dynamic_pool_->allocate(n *
sizeof(T)));
67 void deallocate(value_type* ptr,
size_t) {
69 um_dynamic_pool_->deallocate(ptr);
72 template <
typename T1,
typename T2>
73 friend bool operator==(
const cuda_um_allocator_impl<T1>& lhs,
74 const cuda_um_allocator_impl<T2>& rhs) noexcept;
77 umpire::Allocator* um_dynamic_pool_;
80 template <
class T1,
class T2>
81 bool operator==(
const cuda_um_allocator_impl<T1>& lhs,
82 const cuda_um_allocator_impl<T2>& rhs) noexcept {
83 return lhs.um_dynamic_pool_ == rhs.um_dynamic_pool_;
86 template <
class T1,
class T2>
87 bool operator!=(
const cuda_um_allocator_impl<T1>& lhs,
88 const cuda_um_allocator_impl<T2>& rhs) noexcept {
94 template <
typename T,
typename A>
95 class default_init_allocator :
public A {
96 using a_t = std::allocator_traits<A>;
99 using reference =
typename A::reference;
102 using const_reference =
typename A::const_reference;
104 template <
typename U>
107 default_init_allocator<U, typename a_t::template rebind_alloc<U>>;
112 template <
typename U>
113 void construct(U* ptr) noexcept(
114 std::is_nothrow_default_constructible<U>::value) {
115 ::new (
static_cast<void*
>(ptr)) U;
117 template <
typename U,
typename... Args>
118 void construct(U* ptr, Args&&... args) {
119 a_t::construct(
static_cast<A&
>(*
this), ptr, std::forward<Args>(args)...);
123 template <
typename T>
124 using cuda_um_allocator = default_init_allocator<T, cuda_um_allocator_impl<T>>;
128 #endif // TILEDARRAY_HAS_CUDA
130 #endif // TILEDARRAY_CUDA_UM_ALLOCATOR_H___INCLUDED