polymorphic_allocator.hpp
1 /*
2  * Copyright (c) 2020, NVIDIA CORPORATION.
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 
17 #pragma once
18 
19 #include <memory>
20 #include <type_traits>
21 
22 #include <rmm/cuda_stream_view.hpp>
23 #include <rmm/mr/device/device_memory_resource.hpp>
25 
26 namespace rmm {
27 namespace mr {
28 
43 template <typename T>
45  public:
46  using value_type = T;
47 
53  polymorphic_allocator() = default;
54 
63 
72 
80  template <typename U>
81  polymorphic_allocator(polymorphic_allocator<U> const& other) noexcept : mr_{other.resource()}
82  {
83  }
84 
92  value_type* allocate(std::size_t n, cuda_stream_view stream)
93  {
94  return static_cast<value_type*>(resource()->allocate(n * sizeof(T), stream));
95  }
96 
107  void deallocate(value_type* p, std::size_t n, cuda_stream_view stream)
108  {
109  resource()->deallocate(p, n * sizeof(T), stream);
110  }
111 
117  device_memory_resource* resource() const noexcept { return mr_; }
118 
119  private:
122 };
123 
124 template <typename T, typename U>
125 bool operator==(polymorphic_allocator<T> const& lhs, polymorphic_allocator<U> const& rhs)
126 {
127  return lhs.resource()->is_equal(*rhs.resource());
128 }
129 
130 template <typename T, typename U>
131 bool operator!=(polymorphic_allocator<T> const& lhs, polymorphic_allocator<U> const& rhs)
132 {
133  return not(lhs == rhs);
134 }
135 
158 template <typename Allocator>
160  public:
161  using value_type = typename std::allocator_traits<Allocator>::value_type;
162 
163  stream_allocator_adaptor() = delete;
164 
174  stream_allocator_adaptor(Allocator const& a, cuda_stream_view stream) : alloc_{a}, stream_{stream}
175  {
176  }
177 
186 
195  template <typename OtherAllocator>
198  {
199  }
200 
206  template <typename T>
207  struct rebind {
208  using other =
210  };
211 
219  value_type* allocate(std::size_t n) { return alloc_.allocate(n, stream()); }
220 
230  void deallocate(value_type* p, std::size_t n) { alloc_.deallocate(p, n, stream()); }
231 
236  cuda_stream_view stream() const noexcept { return stream_; }
237 
242  Allocator underlying_allocator() const noexcept { return alloc_; }
243 
244  private:
245  Allocator alloc_;
246  cuda_stream_view stream_;
247 };
248 
249 template <typename A, typename O>
250 bool operator==(stream_allocator_adaptor<A> const& lhs, stream_allocator_adaptor<O> const& rhs)
251 {
252  return lhs.underlying_allocator() == rhs.underlying_allocator();
253 }
254 
255 template <typename A, typename O>
256 bool operator!=(stream_allocator_adaptor<A> const& lhs, stream_allocator_adaptor<O> const& rhs)
257 {
258  return not(lhs == rhs);
259 }
260 
271 template <typename Allocator>
272 auto make_stream_allocator_adaptor(Allocator const& allocator, cuda_stream_view s)
273 {
274  return stream_allocator_adaptor<Allocator>{allocator, s};
275 }
276 
277 } // namespace mr
278 } // namespace rmm
rmm::mr::polymorphic_allocator::polymorphic_allocator
polymorphic_allocator(device_memory_resource *mr)
Construct a polymorphic_allocator using the provided memory resource.
Definition: polymorphic_allocator.hpp:71
rmm::mr::device_memory_resource::allocate
void * allocate(std::size_t bytes, cuda_stream_view stream=cuda_stream_view{})
Allocates memory of size at least bytes.
Definition: device_memory_resource.hpp:102
per_device_resource.hpp
Management of per-device device_memory_resources.
rmm::mr::stream_allocator_adaptor
Adapts a stream ordered allocator to provide a standard Allocator interface.
Definition: polymorphic_allocator.hpp:159
rmm::mr::polymorphic_allocator::allocate
value_type * allocate(std::size_t n, cuda_stream_view stream)
Allocates storage for n objects of type T using the underlying memory resource.
Definition: polymorphic_allocator.hpp:92
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::polymorphic_allocator::polymorphic_allocator
polymorphic_allocator()=default
Construct a polymorphic_allocator using the return value of rmm::mr::get_current_device_resource() as...
rmm::mr::device_memory_resource::deallocate
void deallocate(void *p, std::size_t bytes, cuda_stream_view stream=cuda_stream_view{})
Deallocate memory pointed to by p.
Definition: device_memory_resource.hpp:125
rmm::mr::stream_allocator_adaptor::allocate
value_type * allocate(std::size_t n)
Allocates storage for n objects of type T using the underlying allocator on stream().
Definition: polymorphic_allocator.hpp:219
rmm::mr::polymorphic_allocator
A stream ordered Allocator using a rmm::mr::device_memory_resource to satisfy (de)allocations.
Definition: polymorphic_allocator.hpp:44
rmm::mr::polymorphic_allocator::polymorphic_allocator
polymorphic_allocator(polymorphic_allocator< U > const &other) noexcept
Construct a polymorphic_allocator using other.resource() as the underlying memory resource.
Definition: polymorphic_allocator.hpp:81
rmm::mr::polymorphic_allocator::deallocate
void deallocate(value_type *p, std::size_t n, cuda_stream_view stream)
Deallocates storage pointed to by p.
Definition: polymorphic_allocator.hpp:107
rmm::mr::stream_allocator_adaptor::deallocate
void deallocate(value_type *p, std::size_t n)
Deallocates storage pointed to by p using the underlying allocator on stream().
Definition: polymorphic_allocator.hpp:230
rmm::mr::stream_allocator_adaptor::rebind
Rebinds the allocator to the specified type.
Definition: polymorphic_allocator.hpp:207
rmm::mr::stream_allocator_adaptor::stream_allocator_adaptor
stream_allocator_adaptor(stream_allocator_adaptor< Allocator > const &other)=default
Construct a stream_allocator_adaptor using other.underlying_allocator() and other....
rmm::mr::stream_allocator_adaptor::underlying_allocator
Allocator underlying_allocator() const noexcept
Returns the underlying stream-ordered allocator.
Definition: polymorphic_allocator.hpp:242
rmm::mr::polymorphic_allocator::resource
device_memory_resource * resource() const noexcept
Returns pointer to the underlying rmm::mr::device_memory_resource.
Definition: polymorphic_allocator.hpp:117
rmm::mr::stream_allocator_adaptor::stream_allocator_adaptor
stream_allocator_adaptor(stream_allocator_adaptor< OtherAllocator > const &other)
Construct a stream_allocator_adaptor using other.underlying_allocator() and other....
Definition: polymorphic_allocator.hpp:196
rmm::mr::stream_allocator_adaptor::stream_allocator_adaptor
stream_allocator_adaptor(Allocator const &a, cuda_stream_view stream)
Construct a stream_allocator_adaptor using a as the underlying allocator.
Definition: polymorphic_allocator.hpp:174
rmm::mr::stream_allocator_adaptor::stream
cuda_stream_view stream() const noexcept
Returns the underlying stream on which calls to the underlying allocator are made.
Definition: polymorphic_allocator.hpp:236
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:83
rmm::mr::get_current_device_resource
device_memory_resource * get_current_device_resource()
Get the memory resource for the current device.
Definition: per_device_resource.hpp:229
rmm::mr::polymorphic_allocator::polymorphic_allocator
polymorphic_allocator(polymorphic_allocator< T > const &other)=default
Construct a polymorphic_allocator using other.resource() as the underlying memory resource.