um_allocator.h
Go to the documentation of this file.
1 /*
2  * This file is a part of TiledArray.
3  * Copyright (C) 2018 Virginia Tech
4  *
5  * This program is free software: you can redistribute it and/or modify
6  * it under the terms of the GNU General Public License as published by
7  * the Free Software Foundation, either version 3 of the License, or
8  * (at your option) any later version.
9  *
10  * This program is distributed in the hope that it will be useful,
11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13  * GNU General Public License for more details.
14  *
15  * You should have received a copy of the GNU General Public License
16  * along with this program. If not, see <http://www.gnu.org/licenses/>.
17  *
18  * Eduard Valeyev
19  * Department of Chemistry, Virginia Tech
20  * Jan 31, 2018
21  *
22  */
23 
24 #ifndef TILEDARRAY_CUDA_UM_ALLOCATOR_H___INCLUDED
25 #define TILEDARRAY_CUDA_UM_ALLOCATOR_H___INCLUDED
26 
27 #include <TiledArray/config.h>
28 
29 #ifdef TILEDARRAY_HAS_CUDA
30 
32 
33 #include <memory>
34 #include <stdexcept>
35 
36 namespace TiledArray {
37 
40 template <class T>
41 class cuda_um_allocator_impl {
42  public:
43  using value_type = T;
44  using pointer = T*;
45  using reference = T&;
46  using const_reference = const T&;
47 
48  cuda_um_allocator_impl() noexcept
49  : um_dynamic_pool_(&cudaEnv::instance()->um_dynamic_pool()) {}
50 
51  template <class U>
52  cuda_um_allocator_impl(const cuda_um_allocator_impl<U>& rhs) noexcept
53  : um_dynamic_pool_(rhs.um_dynamic_pool_) {}
54 
56  pointer allocate(size_t n) {
57  pointer result = nullptr;
58 
59  TA_ASSERT(um_dynamic_pool_);
60 
61  result = static_cast<pointer>(um_dynamic_pool_->allocate(n * sizeof(T)));
62 
63  return result;
64  }
65 
67  void deallocate(value_type* ptr, size_t) {
68  TA_ASSERT(um_dynamic_pool_);
69  um_dynamic_pool_->deallocate(ptr);
70  }
71 
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;
75 
76  private:
77  umpire::Allocator* um_dynamic_pool_;
78 }; // class cuda_um_allocator
79 
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_;
84 }
85 
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 {
89  return !(lhs == rhs);
90 }
91 
94 template <typename T, typename A>
95 class default_init_allocator : public A {
96  using a_t = std::allocator_traits<A>;
97 
98  public:
99  using reference = typename A::reference; // std::allocator<T>::reference
100  // deprecated in C++17, but thrust
101  // still relying on this
102  using const_reference = typename A::const_reference; // ditto
103 
104  template <typename U>
105  struct rebind {
106  using other =
107  default_init_allocator<U, typename a_t::template rebind_alloc<U>>;
108  };
109 
110  using A::A;
111 
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;
116  }
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)...);
120  }
121 };
122 
123 template <typename T>
124 using cuda_um_allocator = default_init_allocator<T, cuda_um_allocator_impl<T>>;
125 
126 } // namespace TiledArray
127 
128 #endif // TILEDARRAY_HAS_CUDA
129 
130 #endif // TILEDARRAY_CUDA_UM_ALLOCATOR_H___INCLUDED
bool operator==(const BlockRange &r1, const BlockRange &r2)
BlockRange equality comparison.
Definition: block_range.h:433
constexpr bool operator!=(const DenseShape &a, const DenseShape &b)
Definition: dense_shape.h:382
#define TA_ASSERT(EXPR,...)
Definition: error.h:39