polymorphic_allocator.hpp
1 /*
2  * Copyright (c) 2020-2021, 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 <rmm/cuda_stream_view.hpp>
20 #include <rmm/mr/device/device_memory_resource.hpp>
22 
23 #include <cstddef>
24 #include <memory>
25 #include <type_traits>
26 
27 namespace rmm::mr {
28 
43 template <typename T>
45  public:
46  using value_type = T;
52  polymorphic_allocator() = default;
53 
62 
70  template <typename U>
71  polymorphic_allocator(polymorphic_allocator<U> const& other) noexcept : mr_{other.resource()}
72  {
73  }
74 
82  value_type* allocate(std::size_t num, cuda_stream_view stream)
83  {
84  return static_cast<value_type*>(resource()->allocate(num * sizeof(T), stream));
85  }
86 
97  void deallocate(value_type* ptr, std::size_t num, cuda_stream_view stream)
98  {
99  resource()->deallocate(ptr, num * sizeof(T), stream);
100  }
101 
107  [[nodiscard]] device_memory_resource* resource() const noexcept { return mr_; }
108 
109  private:
112 };
113 
114 template <typename T, typename U>
115 bool operator==(polymorphic_allocator<T> const& lhs, polymorphic_allocator<U> const& rhs)
116 {
117  return lhs.resource()->is_equal(*rhs.resource());
118 }
119 
120 template <typename T, typename U>
121 bool operator!=(polymorphic_allocator<T> const& lhs, polymorphic_allocator<U> const& rhs)
122 {
123  return not(lhs == rhs);
124 }
125 
148 template <typename Allocator>
150  public:
151  using value_type = typename std::allocator_traits<Allocator>::value_type;
152 
153  stream_allocator_adaptor() = delete;
154 
164  stream_allocator_adaptor(Allocator const& allocator, cuda_stream_view stream)
165  : alloc_{allocator}, stream_{stream}
166  {
167  }
168 
177  template <typename OtherAllocator>
180  {
181  }
182 
188  template <typename T>
189  struct rebind {
190  using other =
192  };
193 
201  value_type* allocate(std::size_t num) { return alloc_.allocate(num, stream()); }
202 
212  void deallocate(value_type* ptr, std::size_t num) { alloc_.deallocate(ptr, num, stream()); }
213 
218  [[nodiscard]] cuda_stream_view stream() const noexcept { return stream_; }
219 
224  [[nodiscard]] Allocator underlying_allocator() const noexcept { return alloc_; }
225 
226  private:
227  Allocator alloc_;
228  cuda_stream_view stream_;
229 };
230 
231 template <typename A, typename O>
232 bool operator==(stream_allocator_adaptor<A> const& lhs, stream_allocator_adaptor<O> const& rhs)
233 {
234  return lhs.underlying_allocator() == rhs.underlying_allocator();
235 }
236 
237 template <typename A, typename O>
238 bool operator!=(stream_allocator_adaptor<A> const& lhs, stream_allocator_adaptor<O> const& rhs)
239 {
240  return not(lhs == rhs);
241 }
242 
253 template <typename Allocator>
254 auto make_stream_allocator_adaptor(Allocator const& allocator, cuda_stream_view stream)
255 {
256  return stream_allocator_adaptor<Allocator>{allocator, stream};
257 }
258 
259 } // namespace rmm::mr
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:61
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:106
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:149
rmm::mr::polymorphic_allocator::deallocate
void deallocate(value_type *ptr, std::size_t num, cuda_stream_view stream)
Deallocates storage pointed to by ptr.
Definition: polymorphic_allocator.hpp:97
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::stream_allocator_adaptor::stream_allocator_adaptor
stream_allocator_adaptor(Allocator const &allocator, cuda_stream_view stream)
Construct a stream_allocator_adaptor using a as the underlying allocator.
Definition: polymorphic_allocator.hpp:164
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::stream_allocator_adaptor::deallocate
void deallocate(value_type *ptr, std::size_t num)
Deallocates storage pointed to by ptr using the underlying allocator on stream().
Definition: polymorphic_allocator.hpp:212
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:71
rmm::mr::device_memory_resource::deallocate
void deallocate(void *ptr, std::size_t bytes, cuda_stream_view stream=cuda_stream_view{})
Deallocate memory pointed to by p.
Definition: device_memory_resource.hpp:129
rmm::mr::stream_allocator_adaptor::rebind
Rebinds the allocator to the specified type.
Definition: polymorphic_allocator.hpp:189
rmm::mr::stream_allocator_adaptor::allocate
value_type * allocate(std::size_t num)
Allocates storage for num objects of type T using the underlying allocator on stream().
Definition: polymorphic_allocator.hpp:201
rmm::mr::stream_allocator_adaptor::underlying_allocator
Allocator underlying_allocator() const noexcept
Returns the underlying stream-ordered allocator.
Definition: polymorphic_allocator.hpp:224
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:107
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:178
rmm::mr::polymorphic_allocator::allocate
value_type * allocate(std::size_t num, cuda_stream_view stream)
Allocates storage for num objects of type T using the underlying memory resource.
Definition: polymorphic_allocator.hpp:82
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:218
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82
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:197