cuda_async_memory_resource.hpp
Go to the documentation of this file.
1 /*
2  * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
3  * SPDX-License-Identifier: Apache-2.0
4  */
5 #pragma once
6 
7 #include <rmm/aligned.hpp>
8 #include <rmm/detail/export.hpp>
9 #include <rmm/mr/detail/cuda_async_memory_resource_impl.hpp>
10 
11 #include <cuda/memory_resource>
12 #include <cuda_runtime_api.h>
13 
14 #include <cstddef>
15 #include <cstdint>
16 #include <optional>
17 
18 namespace RMM_NAMESPACE {
19 namespace mr {
30 class RMM_EXPORT cuda_async_memory_resource final
31  : public cuda::mr::shared_resource<detail::cuda_async_memory_resource_impl> {
32  using shared_base = cuda::mr::shared_resource<detail::cuda_async_memory_resource_impl>;
33 
34  public:
47  enum class allocation_handle_type : std::int32_t {
48  none = cudaMemHandleTypeNone,
49  posix_file_descriptor =
50  cudaMemHandleTypePosixFileDescriptor,
52  win32 =
53  cudaMemHandleTypeWin32,
54  win32_kmt = cudaMemHandleTypeWin32Kmt,
56  fabric = 0x8
57  };
58 
70  enum class mempool_usage : unsigned short {
71  hw_decompress = 0x2,
73  };
74 
78  RMM_CONSTEXPR_FRIEND void get_property(cuda_async_memory_resource const&,
79  cuda::mr::device_accessible) noexcept
80  {
81  }
82 
101  // NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
102  cuda_async_memory_resource(std::optional<std::size_t> initial_pool_size = {},
103  std::optional<std::size_t> release_threshold = {},
104  std::optional<allocation_handle_type> export_handle_type = {});
105 
111  [[nodiscard]] cudaMemPool_t pool_handle() const noexcept;
112 
113  ~cuda_async_memory_resource() = default;
115  default;
117  default;
119  default;
121  default;
122 };
123 
124 // static property checks
125 static_assert(cuda::mr::synchronous_resource<cuda_async_memory_resource>);
126 static_assert(cuda::mr::resource<cuda_async_memory_resource>);
127 static_assert(
128  cuda::mr::synchronous_resource_with<cuda_async_memory_resource, cuda::mr::device_accessible>);
129 static_assert(cuda::mr::resource_with<cuda_async_memory_resource, cuda::mr::device_accessible>);
130  // end of group
132 } // namespace mr
133 } // namespace RMM_NAMESPACE
Memory resource that uses cudaMallocAsync/cudaFreeAsync for allocation/deallocation.
Definition: cuda_async_memory_resource.hpp:31
allocation_handle_type
Flags for specifying memory allocation handle types.
Definition: cuda_async_memory_resource.hpp:47
friend void get_property(cuda_async_memory_resource const &, cuda::mr::device_accessible) noexcept
Enables the cuda::mr::device_accessible property.
Definition: cuda_async_memory_resource.hpp:78
cuda_async_memory_resource(std::optional< std::size_t > initial_pool_size={}, std::optional< std::size_t > release_threshold={}, std::optional< allocation_handle_type > export_handle_type={})
Constructs a cuda_async_memory_resource with the optionally specified initial pool size and release t...
mempool_usage
Flags for specifying memory pool usage.
Definition: cuda_async_memory_resource.hpp:70
cudaMemPool_t pool_handle() const noexcept
Returns the underlying native handle to the CUDA pool.