All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
device_memory_resource.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2019-2025, 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 #pragma once
17 
18 #include <rmm/cuda_stream_view.hpp>
19 #include <rmm/detail/aligned.hpp>
20 #include <rmm/detail/export.hpp>
21 #include <rmm/detail/nvtx/ranges.hpp>
22 
23 #include <cuda/memory_resource>
24 
25 #include <cstddef>
26 
27 namespace RMM_NAMESPACE {
28 namespace mr {
94  public:
95  device_memory_resource() = default;
96  virtual ~device_memory_resource() = default;
99  default;
101  default;
103  default;
104 
120  void* allocate(std::size_t bytes, cuda_stream_view stream = cuda_stream_view{})
121  {
122  RMM_FUNC_RANGE();
123  return do_allocate(bytes, stream);
124  }
125 
142  void deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream = cuda_stream_view{})
143  {
144  RMM_FUNC_RANGE();
145  do_deallocate(ptr, bytes, stream);
146  }
147 
161  [[nodiscard]] bool is_equal(device_memory_resource const& other) const noexcept
162  {
163  return do_is_equal(other);
164  }
165 
178  void* allocate(std::size_t bytes, std::size_t alignment)
179  {
180  RMM_FUNC_RANGE();
181  return do_allocate(rmm::align_up(bytes, alignment), cuda_stream_view{});
182  }
183 
197  void deallocate(void* ptr, std::size_t bytes, std::size_t alignment)
198  {
199  RMM_FUNC_RANGE();
200  do_deallocate(ptr, rmm::align_up(bytes, alignment), cuda_stream_view{});
201  }
202 
216  void* allocate_async(std::size_t bytes, std::size_t alignment, cuda_stream_view stream)
217  {
218  RMM_FUNC_RANGE();
219  return do_allocate(rmm::align_up(bytes, alignment), stream);
220  }
221 
234  void* allocate_async(std::size_t bytes, cuda_stream_view stream)
235  {
236  RMM_FUNC_RANGE();
237  return do_allocate(bytes, stream);
238  }
239 
254  void deallocate_async(void* ptr,
255  std::size_t bytes,
256  std::size_t alignment,
257  cuda_stream_view stream)
258  {
259  RMM_FUNC_RANGE();
260  do_deallocate(ptr, rmm::align_up(bytes, alignment), stream);
261  }
262 
276  void deallocate_async(void* ptr, std::size_t bytes, cuda_stream_view stream)
277  {
278  RMM_FUNC_RANGE();
279  do_deallocate(ptr, bytes, stream);
280  }
281 
289  [[nodiscard]] bool operator==(device_memory_resource const& other) const noexcept
290  {
291  return do_is_equal(other);
292  }
293 
301  [[nodiscard]] bool operator!=(device_memory_resource const& other) const noexcept
302  {
303  return !do_is_equal(other);
304  }
305 
311  friend void get_property(device_memory_resource const&, cuda::mr::device_accessible) noexcept {}
312 
313  private:
326  virtual void* do_allocate(std::size_t bytes, cuda_stream_view stream) = 0;
327 
339  virtual void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) = 0;
340 
355  [[nodiscard]] virtual bool do_is_equal(device_memory_resource const& other) const noexcept
356  {
357  return this == &other;
358  }
359 };
360 static_assert(cuda::mr::async_resource_with<device_memory_resource, cuda::mr::device_accessible>); // end of group
362 } // namespace mr
363 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
Base class for all librmm device memory allocation.
Definition: device_memory_resource.hpp:93
void deallocate_async(void *ptr, std::size_t bytes, cuda_stream_view stream)
Deallocate memory pointed to by p.
Definition: device_memory_resource.hpp:276
friend void get_property(device_memory_resource const &, cuda::mr::device_accessible) noexcept
Enables the cuda::mr::device_accessible property.
Definition: device_memory_resource.hpp:311
device_memory_resource(device_memory_resource &&) noexcept=default
Default move constructor.
void deallocate(void *ptr, std::size_t bytes, std::size_t alignment)
Deallocate memory pointed to by p.
Definition: device_memory_resource.hpp:197
void * allocate(std::size_t bytes, std::size_t alignment)
Allocates memory of size at least bytes.
Definition: device_memory_resource.hpp:178
void * allocate_async(std::size_t bytes, cuda_stream_view stream)
Allocates memory of size at least bytes.
Definition: device_memory_resource.hpp:234
bool operator==(device_memory_resource const &other) const noexcept
Comparison operator with another device_memory_resource.
Definition: device_memory_resource.hpp:289
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:142
void deallocate_async(void *ptr, std::size_t bytes, std::size_t alignment, cuda_stream_view stream)
Deallocate memory pointed to by p.
Definition: device_memory_resource.hpp:254
bool operator!=(device_memory_resource const &other) const noexcept
Comparison operator with another device_memory_resource.
Definition: device_memory_resource.hpp:301
device_memory_resource(device_memory_resource const &)=default
Default copy constructor.
void * allocate_async(std::size_t bytes, std::size_t alignment, cuda_stream_view stream)
Allocates memory of size at least bytes.
Definition: device_memory_resource.hpp:216
bool is_equal(device_memory_resource const &other) const noexcept
Compare this resource to another.
Definition: device_memory_resource.hpp:161
constexpr std::size_t align_up(std::size_t value, std::size_t alignment) noexcept
Align up to nearest multiple of specified power of 2.
Definition: aligned.hpp:77