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 
102  // NOLINTNEXTLINE(bugprone-easily-swappable-parameters)
103  cuda_async_memory_resource(std::optional<std::size_t> initial_pool_size = {},
104  std::optional<std::size_t> release_threshold = {},
105  std::optional<allocation_handle_type> export_handle_type = {});
106 
112  [[nodiscard]] cudaMemPool_t pool_handle() const noexcept;
113 
114  ~cuda_async_memory_resource() = default;
116  default;
118  default;
120  default;
122  default;
123 };
124 
125 // static property checks
126 static_assert(cuda::mr::synchronous_resource<cuda_async_memory_resource>);
127 static_assert(cuda::mr::resource<cuda_async_memory_resource>);
128 static_assert(
129  cuda::mr::synchronous_resource_with<cuda_async_memory_resource, cuda::mr::device_accessible>);
130 static_assert(cuda::mr::resource_with<cuda_async_memory_resource, cuda::mr::device_accessible>);
131  // end of group
133 } // namespace mr
134 } // 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.